Upload setup.py

database/setup.py

@@ -1,150 +1,174 @@
-import os
-import sqlite3
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from typing import List, Any
 import json
-from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
-import google.generativeai as genai
+import os
+import re  # 导入正则表达式库
+
+from .tool_recommender import DirectToolRecommender
+from tools.tool_registry import get_tool_by_name
+
+# Agent的思考模板 (保持不变)
+AGENT_PROMPT_TEMPLATE = """
+你是一个强大的AI助理。你的任务是理解用户的问题，并决定是否需要使用工具来回答。
 
-from tools.tool_registry import get_all_tools
+你有以下工具可用：
+{tools}
 
-# --- 配置持久化路径 ---
-DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "data"))
-SQLITE_DB_PATH = os.path.join(DATA_DIR, "tools.metadata.db")
-MILVUS_DATA_PATH = os.path.join(DATA_DIR, "milvus_lite.db")
+如果需要使用工具，请严格按照以下JSON格式进行响应，不要包含任何其他文本或解释：
+{{
+  "tool": "要调用的工具名称",
+  "tool_input": {{ "参数1": "值1", "参数2": "值2" }}
+}}
 
-# --- 模型配置 ---
-EMBEDDING_DIM = 3072
-EMBEDDING_MODEL_NAME = "gemini-embedding-exp-03-07"
-MILVUS_COLLECTION_NAME = "tool_embeddings"
+如果不需要使用任何工具，请直接回答用户的问题。
 
+这是对话历史：
+{chat_history}
 
-def initialize_system():
-    print("--- 开始系统初始化 (最终通关版) ---")
-    os.makedirs(DATA_DIR, exist_ok=True)
+用户问题：{input}
 
-    # --- 正确的初始化顺序 ---
+现在，请你思考并作出回应（JSON或直接回答）：
+"""
+
+
+class SmartAIAgent:
+    def __init__(
+        self,
+        tool_recommender: DirectToolRecommender,
+        registered_tools: List[Any],
+        api_key: str,
+    ):
+        self.tool_recommender = tool_recommender
+        self.registered_tools = registered_tools
+        self.model_name = "gemini-2.5-flash"
+        self.llm = ChatGoogleGenerativeAI(
+            model=self.model_name,
+            google_api_key=api_key,
+            convert_system_message_to_human=True,
+        )
+        self.chat_history = []
+        print(f"LangChain Agent已初始化，使用模型: {self.model_name}。")
+
+    # ------------------- 核心修复在这里！ -------------------
+    # 我们添加一个更健壮的JSON提取函数
+    def _extract_json_from_string(self, text: str) -> dict | None:
+        """从可能包含其他文本的字符串中提取出JSON块。"""
+        # 匹配被 markdown 包裹的JSON
+        match = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL)
+        if match:
+            json_str = match.group(1)
+        else:
+            # 匹配裸露的JSON
+            match = re.search(r"\{.*\}", text, re.DOTALL)
+            if match:
+                json_str = match.group(0)
+            else:
+                return None
+
+        try:
+            return json.loads(json_str)
+        except json.JSONDecodeError:
+            return None
+
+    # ----------------------------------------------------
+
+    def _format_tools_for_prompt(self, tools: List[dict]) -> str:
+        # ... (此函数保持不变) ...
+        if not tools:
+            return "没有可用的工具。"
+        tool_strings = []
+        for tool in tools:
+            try:
+                params = json.loads(tool["parameters"])
+                param_str = ", ".join(
+                    [f"{p_name}: {p_type}" for p_name, p_type in params.items()]
+                )
+                tool_strings.append(
+                    f"- 工具名称: {tool['name']}\n  - 描述: {tool['description']}\n  - 参数: {param_str}"
+                )
+            except (json.JSONDecodeError, TypeError):
+                tool_strings.append(
+                    f"- 工具名称: {tool['name']}\n  - 描述: {tool['description']}\n  - 参数: 无法解析"
+                )
+        return "\n".join(tool_strings)
+
+    def _format_chat_history(self) -> str:
+        # ... (此函数保持不变) ...
+        formatted_history = []
+        for msg in self.chat_history:
+            if isinstance(msg, HumanMessage):
+                formatted_history.append(f"用户: {msg.content}")
+            elif isinstance(msg, AIMessage):
+                formatted_history.append(f"助理: {msg.content}")
+            elif isinstance(msg, ToolMessage):
+                formatted_history.append(f"工具结果: {msg.content}")
+        return "\n".join(formatted_history)
+
+    def stream_run(self, user_input: str):
+        self.chat_history.append(HumanMessage(content=user_input))
+        yield "🤔 正在分析您的问题...\n"
+
+        yield "🔍 正在从工具库中推荐相关工具...\n"
+        recommended_tools_meta = self.tool_recommender.recommend_tools(user_input)
+
+        if not recommended_tools_meta:
+            yield "ℹ️ 未找到相关工具，将直接回答。\n"
+            recommended_tools_prompt = "没有推荐的工具。"
+        else:
+            tool_names = [t["name"] for t in recommended_tools_meta]
+            yield f"✅ 推荐工具: `{', '.join(tool_names)}`\n"
+            recommended_tools_prompt = self._format_tools_for_prompt(
+                recommended_tools_meta
+            )
 
-    # 1. 初始化SQLite并同步工具元数据
-    # 确保SQLite里总是有最新的工具信息
-    _init_sqlite_db()
-    all_tools_definitions = get_all_tools()
-    _sync_tools_to_sqlite(all_tools_definitions)
+        yield f"🧠 正在让AI大脑（{self.model_name}）决定如何行动...\n"
+        prompt = AGENT_PROMPT_TEMPLATE.format(
+            tools=recommended_tools_prompt,
+            chat_history=self._format_chat_history(),
+            input=user_input,
+        )
 
-    # 2. 初始化Milvus并同步向量
-    # 它会从已经填充好的SQLite中读取数据
-    milvus_client = _init_milvus_and_sync_embeddings()
+        llm_response = self.llm.invoke(prompt)
+        llm_decision_content = llm_response.content.strip()
 
-    # 3. 创建工具推荐器
-    from core.tool_recommender import DirectToolRecommender
+        # ------------------- 核心修复在这里！ -------------------
+        # 使用我们新的、更健壮的JSON提取逻辑
+        decision = self._extract_json_from_string(llm_decision_content)
 
-    tool_recommender = DirectToolRecommender(
-        milvus_client=milvus_client, sqlite_db_path=SQLITE_DB_PATH
-    )
+        if decision and "tool" in decision and "tool_input" in decision:
+            # 如果成功提取出有效的工具调用JSON
+            tool_name = decision.get("tool")
+            tool_input = decision.get("tool_input")
 
-    print("--- 系统初始化完成 ---")
-    return all_tools_definitions, tool_recommender
+            yield f"💡 AI决策：调用工具 `{tool_name}`，参数为 `{tool_input}`\n"
 
+            tool_to_execute = get_tool_by_name(tool_name)
+            if tool_to_execute:
+                yield f"⚙️ 正在执行工具 `{tool_name}`...\n"
+                tool_output = tool_to_execute.invoke(tool_input)
+                yield f"📊 工具返回结果：\n---\n{str(tool_output)[:500]}...\n---\n"
 
-def _init_sqlite_db():
-    print(f"SQLite DB 路径: {SQLITE_DB_PATH}")
-    with sqlite3.connect(SQLITE_DB_PATH) as conn:
-        cursor = conn.cursor()
-        cursor.execute(
-            """
-            CREATE TABLE IF NOT EXISTS tools (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                name TEXT UNIQUE NOT NULL,
-                description TEXT NOT NULL,
-                parameters TEXT NOT NULL
-            )
-        """
-        )
-        conn.commit()
-    print("SQLite DB 表已确认存在。")
-
-
-def _sync_tools_to_sqlite(tools_definitions):
-    print("正在同步工具元数据到SQLite...")
-    with sqlite3.connect(SQLITE_DB_PATH) as conn:
-        cursor = conn.cursor()
-        for tool in tools_definitions:
-            cursor.execute("SELECT id FROM tools WHERE name = ?", (tool.name,))
-            if cursor.fetchone() is None:
-                cursor.execute(
-                    "INSERT INTO tools (name, description, parameters) VALUES (?, ?, ?)",
-                    (tool.name, tool.description, json.dumps(tool.args)),
+                self.chat_history.append(
+                    AIMessage(content=json.dumps(decision, ensure_ascii=False))
                 )
-                print(f"  - 新增工具到SQLite: {tool.name}")
-        conn.commit()
-    print("SQLite同步完成。")
-
-
-def _init_milvus_and_sync_embeddings():
-    print(f"Milvus Lite 数据路径: {MILVUS_DATA_PATH}")
-    client = MilvusClient(uri=MILVUS_DATA_PATH)
-
-    # 每次启动都重新创建集合，确保维度正确且数据最新
-    if client.has_collection(collection_name=MILVUS_COLLECTION_NAME):
-        client.drop_collection(collection_name=MILVUS_COLLECTION_NAME)
-        print("发现旧的Milvus集合，已删除以重建。")
-
-    print(f"Milvus集合 '{MILVUS_COLLECTION_NAME}' 正在创建，维度为 {EMBEDDING_DIM}...")
-    fields = [
-        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
-        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=EMBEDDING_DIM),
-    ]
-    schema = CollectionSchema(fields)
-    client.create_collection(collection_name=MILVUS_COLLECTION_NAME, schema=schema)
-
-    index_params = client.prepare_index_params()
-    index_params.add_index(
-        field_name="embedding", index_type="AUTOINDEX", metric_type="L2"
-    )
-    client.create_index(
-        collection_name=MILVUS_COLLECTION_NAME, index_params=index_params
-    )
-    print("Milvus集合和索引创建完成。")
-
-    # --- 关键：现在我们才同步嵌入 ---
-    _sync_tool_embeddings_to_milvus(client)
-
-    client.load_collection(collection_name=MILVUS_COLLECTION_NAME)
-    return client
-
-
-def _sync_tool_embeddings_to_milvus(milvus_client):
-    print("正在同步工具嵌入到Milvus...")
-    api_key = os.environ.get("GEMINI_API_KEY")
-    if not api_key:
-        print("错误：无法找到GEMINI_API_KEY。")
-        return
-    genai.configure(api_key=api_key)
-
-    with sqlite3.connect(SQLITE_DB_PATH) as conn:
-        cursor = conn.cursor()
-        cursor.execute("SELECT id, description FROM tools")
-        all_tools_in_db = cursor.fetchall()
-
-    if not all_tools_in_db:
-        print("SQLite中没有工具可同步，这是一个错误！")
-        return
-
-    print(f"从SQLite发现 {len(all_tools_in_db)} 个工具，正在生成嵌入...")
-    docs_to_embed = [tool[1] for tool in all_tools_in_db]
-
-    print(f"使用嵌入模型: {EMBEDDING_MODEL_NAME}")
-    result = genai.embed_content(
-        model=EMBEDDING_MODEL_NAME,
-        content=docs_to_embed,
-        task_type="retrieval_document",
-    )
-
-    embeddings = result["embedding"]
-    tool_ids_to_insert = [tool[0] for tool in all_tools_in_db]
-
-    data_to_insert = [
-        {"id": tool_id, "embedding": embedding}
-        for tool_id, embedding in zip(tool_ids_to_insert, embeddings)
-    ]
-
-    milvus_client.insert(collection_name=MILVUS_COLLECTION_NAME, data=data_to_insert)
-    print(f"成功将 {len(data_to_insert)} 个新嵌入插入到Milvus。")
+                self.chat_history.append(
+                    ToolMessage(content=str(tool_output), tool_call_id="N/A")
+                )
+
+                yield "✍️ 正在根据工具结果生成最终回答...\n\n"
+                final_answer_prompt = f"基于以下对话历史和最新的工具结果，请为用户生成一个最终的、完整的、自然的回答。\n\n对话历史：\n{self._format_chat_history()}\n\n请直接回答，不要提及你的思考过程。"
+                final_answer_stream = self.llm.stream(final_answer_prompt)
+                full_final_answer = ""
+                for chunk in final_answer_stream:
+                    yield chunk.content
+                    full_final_answer += chunk.content
+                self.chat_history.append(AIMessage(content=full_final_answer))
+            else:
+                yield f"❌ 错误：AI决策调用的工具 `{tool_name}` 不存在。\n"
+        else:
+            # 如果没有提取出JSON，或者JSON格式不正确，则认为是直接回答
+            yield "✅ AI决策：直接回答。\n\n"
+            yield llm_decision_content
+            self.chat_history.append(AIMessage(content=llm_decision_content))
+        # ----------------------------------------------------