Bubbles/commands/ai_router.py

import re
import json
import logging
from typing import Dict, Callable, Optional, Any, Tuple
from dataclasses import dataclass, field
from .context import MessageContext

logger = logging.getLogger(__name__)

@dataclass
class AIFunction:
    """AI可调用的功能定义"""
    name: str                          # 功能唯一标识名
    handler: Callable                  # 处理函数
    description: str                   # 功能描述（给AI看的）
    examples: list[str] = field(default_factory=list)  # 示例用法
    params_description: str = ""       # 参数说明
    
class AIRouter:
    """AI智能路由器"""
    
    def __init__(self):
        self.functions: Dict[str, AIFunction] = {}
        self.logger = logger
        
    def register(self, name: str, description: str, examples: list[str] = None, params_description: str = ""):
        """
        装饰器：注册一个功能到AI路由器
        
        @ai_router.register(
            name="weather_query",
            description="查询指定城市的天气预报",
            examples=["北京天气怎么样", "查一下上海的天气", "明天深圳会下雨吗"],
            params_description="城市名称"
        )
        def handle_weather(ctx: MessageContext, params: str) -> bool:
            # 实现天气查询逻辑
            pass
        """
        def decorator(func: Callable) -> Callable:
            ai_func = AIFunction(
                name=name,
                handler=func,
                description=description,
                examples=examples or [],
                params_description=params_description
            )
            self.functions[name] = ai_func
            self.logger.info(f"AI路由器注册功能: {name} - {description}")
            return func
        
        return decorator
    
    def _build_ai_prompt(self) -> str:
        """构建给AI的系统提示词，包含所有可用功能的信息"""
        prompt = """你是一个智能路由助手。根据用户的输入，判断用户的意图并返回JSON格式的响应。

可用的功能列表：
"""
        for name, func in self.functions.items():
            prompt += f"\n- {name}: {func.description}"
            if func.params_description:
                prompt += f"\n  参数: {func.params_description}"
            if func.examples:
                prompt += f"\n  示例: {', '.join(func.examples[:3])}"
            prompt += "\n"
        
        prompt += """
分析用户输入，严格按照以下格式返回JSON：

如果用户需要使用上述功能之一，返回：
{
    "action_type": "function",
    "function_name": "上述功能列表中的功能名",
    "params": "从用户输入中提取的参数"
}

如果用户只是聊天或者不匹配任何功能，返回：
{
    "action_type": "chat"
}

示例：
- 用户输入"北京天气怎么样" -> {"action_type": "function", "function_name": "weather_query", "params": "北京"}
- 用户输入"看看新闻" -> {"action_type": "function", "function_name": "news_query", "params": ""}
- 用户输入"你好" -> {"action_type": "chat"}
- 用户输入"查一下Python教程" -> {"action_type": "function", "function_name": "perplexity_search", "params": "Python教程"}

重要：
1. action_type 只能是 "function" 或 "chat"
2. 只返回JSON，无需其他解释
3. function_name 必须完全匹配上述功能列表中的名称
"""
        return prompt
    
    def route(self, ctx: MessageContext) -> Tuple[bool, Optional[Dict[str, Any]]]:
        """
        AI路由决策
        
        返回: (是否处理成功, AI决策结果)
        """
        print(f"[AI路由器] route方法被调用")
        
        if not ctx.text:
            print("[AI路由器] ctx.text为空，返回False")
            return False, None
            
        # 获取AI模型
        chat_model = getattr(ctx, 'chat', None)
        if not chat_model:
            chat_model = getattr(ctx.robot, 'chat', None) if ctx.robot else None
            
        if not chat_model:
            print("[AI路由器] 无可用的AI模型")
            self.logger.error("AI路由器：无可用的AI模型")
            return False, None
        
        print(f"[AI路由器] 找到AI模型: {type(chat_model)}")
        
        try:
            # 构建系统提示词
            system_prompt = self._build_ai_prompt()
            print(f"[AI路由器] 已构建系统提示词，长度: {len(system_prompt)}")
            
            # 让AI分析用户意图
            user_input = f"用户输入：{ctx.text}"
            print(f"[AI路由器] 准备调用AI分析意图: {user_input}")
            
            ai_response = chat_model.get_answer(
                user_input, 
                wxid=ctx.get_receiver(),
                system_prompt_override=system_prompt
            )
            
            print(f"[AI路由器] AI响应: {ai_response}")
            
            # 解析AI返回的JSON
            json_match = re.search(r'\{.*\}', ai_response, re.DOTALL)
            if not json_match:
                self.logger.warning(f"AI路由器：无法从AI响应中提取JSON - {ai_response}")
                return False, None
                
            decision = json.loads(json_match.group(0))
            
            # 验证决策格式
            action_type = decision.get("action_type")
            if action_type not in ["chat", "function"]:
                self.logger.warning(f"AI路由器：未知的action_type - {action_type}")
                return False, None
            
            # 如果是功能调用，验证功能名
            if action_type == "function":
                function_name = decision.get("function_name")
                if function_name not in self.functions:
                    self.logger.warning(f"AI路由器：未知的功能名 - {function_name}")
                    return False, None
            
            self.logger.info(f"AI路由决策: {decision}")
            return True, decision
            
        except json.JSONDecodeError as e:
            self.logger.error(f"AI路由器：解析JSON失败 - {e}")
            return False, None
        except Exception as e:
            self.logger.error(f"AI路由器：处理异常 - {e}")
            return False, None
    
    def dispatch(self, ctx: MessageContext) -> bool:
        """
        执行AI路由分发
        
        返回: 是否成功处理
        """
        print(f"[AI路由器] dispatch被调用，消息内容: {ctx.text}")
        
        # 获取AI路由决策
        success, decision = self.route(ctx)
        print(f"[AI路由器] route返回 - success: {success}, decision: {decision}")
        
        if not success or not decision:
            print("[AI路由器] route失败或无决策，返回False")
            return False
        
        action_type = decision.get("action_type")
        
        # 如果是聊天，返回False让后续处理器处理
        if action_type == "chat":
            self.logger.info("AI路由器：识别为聊天意图，交给聊天处理器")
            return False
        
        # 如果是功能调用
        if action_type == "function":
            function_name = decision.get("function_name")
            params = decision.get("params", "")
            
            func = self.functions.get(function_name)
            if not func:
                self.logger.error(f"AI路由器：功能 {function_name} 未找到")
                return False
            
            try:
                self.logger.info(f"AI路由器：调用功能 {function_name}，参数: {params}")
                result = func.handler(ctx, params)
                return result
            except Exception as e:
                self.logger.error(f"AI路由器：执行功能 {function_name} 出错 - {e}")
                return False
        
        return False

# 创建全局AI路由器实例
ai_router = AIRouter()