blank/backend/conversation_manager.py

"""
多轮对话管理器 - 处理包含工具调用的多轮对话
"""
import asyncio
from typing import Dict, List, Any, Optional
import httpx
from anthropic import Anthropic
from mcp_client import MCPClient
from tool_converter import ToolConverter
from tool_handler import ToolCallHandler


# ========== 默认组件提示（降级方案） ==========
# 当前端未发送组件列表时使用的默认提示
DEFAULT_COMPONENTS_PROMPT = """### 可用的 json-render 组件

你可以使用基础组件来展示结构化数据：
- `card`: 卡片容器
- `stack`: 布局容器
- `heading`: 标题 (h1-h6)
- `text`: 文本内容
- `button`: 按钮
- `badge`: 徽章标签
- `code-block`: 代码块
- `data-table`: 数据表格

将组件 JSON 包含在 markdown 代码块中返回。
"""


# ========== 基础 System Prompt ==========
# 不包含组件列表，组件列表由前端动态提供
BASE_SYSTEM_PROMPT = """你是一个 AI 助手，可以通过调用 MCP 工具来帮助用户完成任务。

## 重要：你可以返回 UI 组件

除了普通文本，你还可以返回 **json-render 组件 spec** 来展示更丰富的 UI。组件 spec 是一个 JSON 对象，前端会自动渲染成 UI 组件。

### 组件格式

返回组件时，使用以下 JSON 格式（在你的回复中包裹在 ```json 中）：

```json
{
  "type": "组件类型",
  "其他属性": "..."
}
```

### 使用指南

1. **调用工具后，优先使用对应组件展示结果**
   - 根据工具返回的数据类型选择合适的组件
   - 组件必须放在代码块中返回

2. **你可以混合文本和组件**
   - 先用文本解释结果
   - 然后用组件展示数据

3. **给出操作建议时使用按钮组件**
   - 用户查看列表后，建议"下一页"、"筛选"等操作
   - 使用 `button` 或 `suggestion-buttons` 组件

### 小说相关操作（重要）

当用户说"查看小说：xxx"、"小说详情：xxx"或点击小说卡片时：
1. 从消息中提取小说标题或 ID
2. 使用 `get_novel_detail` 工具获取小说数据（需要提供 novel_id 参数）
3. 使用 `novel-detail` 组件展示结果，格式如下：
```json
{
  "type": "novel-detail",
  "novel": {
    "id": "小说ID",
    "title": "小说标题",
    "author": "作者",
    "category": "分类",
    "description": "简介",
    "status": "状态",
    "chapterCount": 章节数,
    "wordCount": 字数,
    "viewCount": 阅读量,
    "isVip": 是否VIP
  }
}
```
"""


def create_anthropic_client(api_key: str, base_url: str) -> Anthropic:
    """
    创建 Anthropic 客户端，支持自定义认证格式

    自定义 API 代理需要 'Authorization: Bearer <token>' 格式，
    而不是 Anthropic SDK 默认的 'x-api-key' header。
    """
    # 创建自定义 httpx client，设置正确的 Authorization header
    http_client = httpx.Client(
        headers={"Authorization": f"Bearer {api_key}"},
        timeout=120.0
    )
    return Anthropic(base_url=base_url, http_client=http_client)


class ConversationManager:
    """管理包含工具调用的多轮对话"""

    def __init__(
        self,
        api_key: str,
        base_url: str,
        model: str,
        session_id: str = None,
        mcp_tokens: dict = None,
        components_prompt: str = None,
        enabled_mcp_list: list = None  # 新增：前端传递的已启用 MCP 列表
    ):
        self.api_key = api_key
        self.base_url = base_url
        self.model = model
        self.session_id = session_id
        self.mcp_tokens = mcp_tokens or {}  # MCP 服务器 token 映射
        self.enabled_mcp_list = enabled_mcp_list  # 前端传递的已启用 MCP 列表

        # 组件提示词（由前端动态提供）
        self.components_prompt = components_prompt or DEFAULT_COMPONENTS_PROMPT

        # 构建完整的系统提示词
        self.system_prompt = self._build_system_prompt()

        # DEBUG: 打印接收到的 token 和启用列表
        print(f"[DEBUG ConversationManager.__init__] mcp_tokens keys: {list(self.mcp_tokens.keys())}")
        print(f"[DEBUG ConversationManager.__init__] enabled_mcp_list: {self.enabled_mcp_list}")
        print(f"[DEBUG ConversationManager.__init__] Using dynamic components: {components_prompt is not None}")
        for k, v in self.mcp_tokens.items():
            print(f"[DEBUG ConversationManager.__init__]   {k}: {v[:30] if v else 'None'}...")

        self.tool_handler = ToolCallHandler(session_id=session_id, mcp_tokens=mcp_tokens)
        self._cached_tools = None
        self._tool_to_server_map = {}  # 工具名到服务器 ID 的映射
        # 使用自定义 client，支持 Bearer token 认证
        self.client = create_anthropic_client(api_key, base_url)

    def _build_system_prompt(self) -> str:
        """构建完整的系统提示词（基础提示 + 组件列表）"""
        return BASE_SYSTEM_PROMPT + "\n\n" + self.components_prompt

    async def get_available_tools(self) -> List[Dict[str, Any]]:
        """获取可用的 Claude 格式工具列表（带缓存）"""
        if self._cached_tools is not None:
            return self._cached_tools

        # 从 MCP 服务器发现工具（带 token 和启用列表）
        mcp_tools = await MCPClient.get_all_tools_with_tokens_async(
            self.session_id,
            self.mcp_tokens,
            self.enabled_mcp_list  # 传递前端传递的已启用 MCP 列表
        )

        # 转换为 Claude 格式
        claude_tools = []
        for tool in mcp_tools:
            claude_tool = ToolConverter.mcp_to_claude_tool(tool)
            claude_tools.append(claude_tool)

            # 构建工具名到服务器 ID 的映射
            server_id = tool.get("_server_id", "")
            if server_id:
                self._tool_to_server_map[claude_tool["name"]] = server_id

        self._cached_tools = claude_tools

        return claude_tools

    @classmethod
    async def get_tools_async(cls, session_id: str = None) -> List[Dict[str, Any]]:
        """
        类方法：获取可用的工具列表（异步）

        用于 API 端点直接调用，无需创建完整实例
        """
        mcp_tools = await MCPClient.get_all_tools_async(session_id)
        return ToolConverter.convert_mcp_tools(mcp_tools)

    @staticmethod
    def get_tools(session_id: str = None) -> List[Dict[str, Any]]:
        """
        静态方法：获取可用的工具列表（同步）

        用于 API 端点直接调用
        """
        return asyncio.run(ConversationManager.get_tools_async(session_id))

    async def chat(
        self,
        user_message: str,
        conversation_history: List[Dict[str, Any]] = None,
        max_turns: int = 5
    ) -> Dict[str, Any]:
        """
        执行多轮对话（自动处理工具调用）

        Args:
            user_message: 用户消息
            conversation_history: 对话历史
            max_turns: 最大对话轮数（防止无限循环）

        Returns:
            最终响应和对话历史
        """
        if conversation_history is None:
            conversation_history = []

        messages = conversation_history.copy()
        messages.append({
            "role": "user",
            "content": user_message
        })

        current_messages = messages
        response_text = ""
        tool_calls_made = []

        for turn in range(max_turns):
            # 获取可用工具
            tools = await self.get_available_tools()

            # 调用 Claude API
            if tools:
                response = self.client.messages.create(
                    model=self.model,
                    max_tokens=4096,
                    system=self.system_prompt,  # 使用动态系统提示
                    messages=current_messages,
                    tools=tools
                )
            else:
                response = self.client.messages.create(
                    model=self.model,
                    max_tokens=4096,
                    system=self.system_prompt,  # 使用动态系统提示
                    messages=current_messages
                )

            # 检查响应中是否有 tool_use
            content_blocks = []
            tool_use_blocks = []
            text_blocks = []

            for block in response.content:
                block_type = getattr(block, "type", None)

                if block_type == "tool_use":
                    # 工具调用块
                    block_dict = {
                        "type": "tool_use",
                        "id": getattr(block, "id", ""),
                        "name": getattr(block, "name", ""),
                        "input": getattr(block, "input", {})
                    }
                    content_blocks.append(block_dict)
                    tool_use_blocks.append(block_dict)
                else:
                    # 文本块
                    text_content = getattr(block, "text", "")
                    if text_content:
                        text_blocks.append({
                            "type": "text",
                            "text": text_content
                        })
                        content_blocks.append({
                            "type": "text",
                            "text": text_content
                        })
                        response_text += text_content

            # 如果没有工具调用，返回结果
            if not tool_use_blocks:
                return {
                    "response": response_text,
                    "messages": current_messages,
                    "tool_calls": tool_calls_made
                }

            # 处理工具调用
            tool_results = await self.tool_handler.process_tool_use_blocks(
                tool_use_blocks,
                self._tool_to_server_map
            )

            # 记录工具调用
            for tr in tool_results:
                tool_calls_made.append({
                    "tool": tr.get("tool_name"),
                    "result": tr.get("result", {})
                })

            # 构建工具结果消息
            tool_result_message = ToolCallHandler.create_tool_result_message(
                tool_results
            )

            # 添加到消息历史
            current_messages.append({
                "role": "assistant",
                "content": content_blocks
            })
            current_messages.append(tool_result_message)

        # 达到最大轮数
        return {
            "response": response_text,
            "messages": current_messages,
            "tool_calls": tool_calls_made,
            "warning": "达到最大对话轮数"
        }

    @staticmethod
    def format_history_for_claude(history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        格式化对话历史为 Claude API 格式

        Args:
            history: 原始对话历史

        Returns:
            Claude API 格式的消息列表
        """
        formatted = []

        for msg in history:
            role = msg.get("role")
            content = msg.get("content")

            if role == "user":
                if isinstance(content, str):
                    formatted.append({"role": "user", "content": content})
                elif isinstance(content, list):
                    formatted.append({"role": "user", "content": content})
            elif role == "assistant":
                if isinstance(content, str):
                    formatted.append({"role": "assistant", "content": content})
                elif isinstance(content, list):
                    formatted.append({"role": "assistant", "content": content})

        return formatted