nanocode/skills/context-hub/chub_wrapper.py

#!/usr/bin/env python3
"""
Context Hub (chub) 包装器
为nanobot提供更友好的chub接口
"""

import subprocess
import json
import sys
import os
from typing import List, Dict, Any, Optional

class ChubWrapper:
    """chub CLI的Python包装器"""
    
    def __init__(self):
        """初始化chub包装器"""
        self.chub_path = self._find_chub()
        
    def _find_chub(self) -> str:
        """查找chub可执行文件路径"""
        try:
            result = subprocess.run(['which', 'chub'], 
                                  capture_output=True, 
                                  text=True, 
                                  check=True)
            return result.stdout.strip()
        except subprocess.CalledProcessError:
            raise RuntimeError("chub未安装。请运行: npm install -g @aisuite/chub")
    
    def _run_chub(self, args: List[str], capture_output: bool = True) -> str:
        """运行chub命令"""
        cmd = [self.chub_path] + args
        try:
            result = subprocess.run(cmd, 
                                  capture_output=capture_output, 
                                  text=True, 
                                  check=True)
            return result.stdout
        except subprocess.CalledProcessError as e:
            error_msg = f"chub命令失败: {' '.join(cmd)}\n"
            if e.stderr:
                error_msg += f"错误: {e.stderr}"
            raise RuntimeError(error_msg)
    
    def search(self, query: str, json_output: bool = True) -> List[Dict[str, Any]]:
        """
        搜索文档
        
        Args:
            query: 搜索查询
            json_output: 是否返回JSON格式
            
        Returns:
            搜索结果列表
        """
        args = ['search', query]
        if json_output:
            args.append('--json')
        
        output = self._run_chub(args)
        
        if json_output:
            try:
                data = json.loads(output)
                return data.get('results', [])
            except json.JSONDecodeError:
                # 如果不是JSON格式，返回原始输出
                return [{'raw_output': output}]
        else:
            return [{'raw_output': output}]
    
    def get_doc(self, doc_id: str, language: Optional[str] = None, 
                output_file: Optional[str] = None) -> str:
        """
        获取文档内容
        
        Args:
            doc_id: 文档ID (如 'openai/chat')
            language: 语言 (如 'py', 'js', 'ts')
            output_file: 输出文件路径
            
        Returns:
            文档内容
        """
        args = ['get', doc_id]
        
        if language:
            args.extend(['--lang', language])
        
        if output_file:
            args.extend(['-o', output_file])
        
        return self._run_chub(args, capture_output=not output_file)
    
    def annotate(self, doc_id: str, note: str) -> str:
        """
        添加注解
        
        Args:
            doc_id: 文档ID
            note: 注解内容
            
        Returns:
            执行结果
        """
        args = ['annotate', doc_id, note]
        return self._run_chub(args)
    
    def list_annotations(self) -> str:
        """列出所有注解"""
        args = ['annotate', '--list']
        return self._run_chub(args)
    
    def feedback(self, doc_id: str, rating: str, label: Optional[str] = None) -> str:
        """
        提供反馈
        
        Args:
            doc_id: 文档ID
            rating: 评分 ('up' 或 'down')
            label: 标签 (如 'outdated', 'inaccurate')
            
        Returns:
            执行结果
        """
        args = ['feedback', doc_id, rating]
        if label:
            args.extend(['--label', label])
        
        return self._run_chub(args)
    
    def search_and_get(self, query: str, language: Optional[str] = None) -> Dict[str, Any]:
        """
        搜索并获取文档（一步完成）
        
        Args:
            query: 搜索查询
            language: 语言
            
        Returns:
            包含搜索结果和文档内容的字典
        """
        # 搜索
        results = self.search(query)
        
        if not results:
            return {'error': f'未找到与"{query}"相关的文档'}
        
        # 获取第一个结果
        first_result = results[0]
        doc_id = first_result.get('id')
        
        if not doc_id:
            return {'error': '搜索结果中没有找到有效的文档ID'}
        
        # 获取文档
        try:
            content = self.get_doc(doc_id, language)
            return {
                'doc_id': doc_id,
                'result': first_result,
                'content': content[:1000] + '...' if len(content) > 1000 else content
            }
        except Exception as e:
            return {
                'doc_id': doc_id,
                'result': first_result,
                'error': f'获取文档失败: {str(e)}'
            }


def main():
    """命令行接口"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Context Hub (chub) 包装器')
    subparsers = parser.add_subparsers(dest='command', help='命令')
    
    # 搜索命令
    search_parser = subparsers.add_parser('search', help='搜索文档')
    search_parser.add_argument('query', help='搜索查询')
    search_parser.add_argument('--json', action='store_true', help='JSON输出')
    
    # 获取命令
    get_parser = subparsers.add_parser('get', help='获取文档')
    get_parser.add_argument('doc_id', help='文档ID')
    get_parser.add_argument('--lang', help='语言 (py, js, ts)')
    get_parser.add_argument('-o', '--output', help='输出文件')
    
    # 注解命令
    annotate_parser = subparsers.add_parser('annotate', help='添加注解')
    annotate_parser.add_argument('doc_id', help='文档ID')
    annotate_parser.add_argument('note', help='注解内容')
    
    # 列出注解命令
    subparsers.add_parser('list-annotations', help='列出所有注解')
    
    # 反馈命令
    feedback_parser = subparsers.add_parser('feedback', help='提供反馈')
    feedback_parser.add_argument('doc_id', help='文档ID')
    feedback_parser.add_argument('rating', choices=['up', 'down'], help='评分')
    feedback_parser.add_argument('--label', help='标签')
    
    # 搜索并获取命令
    search_get_parser = subparsers.add_parser('search-get', help='搜索并获取文档')
    search_get_parser.add_argument('query', help='搜索查询')
    search_get_parser.add_argument('--lang', help='语言')
    
    args = parser.parse_args()
    
    wrapper = ChubWrapper()
    
    try:
        if args.command == 'search':
            results = wrapper.search(args.query, args.json)
            if args.json:
                print(json.dumps(results, indent=2, ensure_ascii=False))
            else:
                for result in results:
                    print(f"ID: {result.get('id', 'N/A')}")
                    print(f"名称: {result.get('name', 'N/A')}")
                    print(f"描述: {result.get('description', 'N/A')[:100]}...")
                    print()
        
        elif args.command == 'get':
            content = wrapper.get_doc(args.doc_id, args.lang, args.output)
            if not args.output:
                print(content)
        
        elif args.command == 'annotate':
            result = wrapper.annotate(args.doc_id, args.note)
            print(result)
        
        elif args.command == 'list-annotations':
            result = wrapper.list_annotations()
            print(result)
        
        elif args.command == 'feedback':
            result = wrapper.feedback(args.doc_id, args.rating, args.label)
            print(result)
        
        elif args.command == 'search-get':
            result = wrapper.search_and_get(args.query, args.lang)
            print(json.dumps(result, indent=2, ensure_ascii=False))
        
        else:
            parser.print_help()
    
    except Exception as e:
        print(f"错误: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()