Upload 8 files

database/setup.py

@@ -1,151 +1,150 @@
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
-from typing import List, Any
-import json
 import os
-import re # 导入正则表达式库
+import sqlite3
+import json
+from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
+import google.generativeai as genai
 
-from .tool_recommender import DirectToolRecommender
-from tools.tool_registry import get_tool_by_name
+from tools.tool_registry import get_all_tools
 
-# Agent的思考模板 (保持不变)
-AGENT_PROMPT_TEMPLATE = """
-你是一个强大的AI助理。你的任务是理解用户的问题，并决定是否需要使用工具来回答。
+# --- 配置持久化路径 ---
+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")
 
-你有以下工具可用：
-{tools}
+# --- 模型配置 ---
+EMBEDDING_DIM = 3072
+EMBEDDING_MODEL_NAME = "gemini-embedding-exp-03-07"
+MILVUS_COLLECTION_NAME = "tool_embeddings"
 
-如果需要使用工具，请严格按照以下JSON格式进行响应，不要包含任何其他文本或解释：
-{{
-  "tool": "要调用的工具名称",
-  "tool_input": {{ "参数1": "值1", "参数2": "值2" }}
-}}
 
-如果不需要使用任何工具，请直接回答用户的问题。
+def initialize_system():
+    print("--- 开始系统初始化 (最终通关版) ---")
+    os.makedirs(DATA_DIR, exist_ok=True)
 
-这是对话历史：
-{chat_history}
+    # --- 正确的初始化顺序 ---
 
-用户问题：{input}
+    # 1. 初始化SQLite并同步工具元数据
+    # 确保SQLite里总是有最新的工具信息
+    _init_sqlite_db()
+    all_tools_definitions = get_all_tools()
+    _sync_tools_to_sqlite(all_tools_definitions)
 
-现在，请你思考并作出回应（JSON或直接回答）：
-"""
+    # 2. 初始化Milvus并同步向量
+    # 它会从已经填充好的SQLite中读取数据
+    milvus_client = _init_milvus_and_sync_embeddings()
 
-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)
-
-        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
+    # 3. 创建工具推荐器
+    from core.tool_recommender import DirectToolRecommender
+
+    tool_recommender = DirectToolRecommender(
+        milvus_client=milvus_client, sqlite_db_path=SQLITE_DB_PATH
+    )
+
+    print("--- 系统初始化完成 ---")
+    return all_tools_definitions, tool_recommender
+
+
+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
+            )
+        """
         )
-        
-        llm_response = self.llm.invoke(prompt)
-        llm_decision_content = llm_response.content.strip()
-        
-        # ------------------- 核心修复在这里！ -------------------
-        # 使用我们新的、更健壮的JSON提取逻辑
-        decision = self._extract_json_from_string(llm_decision_content)
-        
-        if decision and "tool" in decision and "tool_input" in decision:
-            # 如果成功提取出有效的工具调用JSON
-            tool_name = decision.get("tool")
-            tool_input = decision.get("tool_input")
-            
-            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"
-                
-                self.chat_history.append(AIMessage(content=json.dumps(decision, ensure_ascii=False)))
-                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))
-        # ----------------------------------------------------
+        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)),
+                )
+                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。")