jarvis-ai-assistant 0.1.74__tar.gz → 0.1.76__tar.gz

{jarvis_ai_assistant-0.1.74/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.1.76}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.74
+Version: 0.1.76
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire

{jarvis_ai_assistant-0.1.74 → jarvis_ai_assistant-0.1.76}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jarvis-ai-assistant"
-version = "0.1.74"
+version = "0.1.76"
 description = "Jarvis: An AI assistant that uses tools to interact with the system"
 readme = "README.md"
 authors = [{ name = "Your Name", email = "your.email@example.com" }]

{jarvis_ai_assistant-0.1.74 → jarvis_ai_assistant-0.1.76}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="jarvis-ai-assistant",
-    version="0.1.74",
+    version="0.1.76",
     author="skyfire",
     author_email="skyfireitdiy@hotmail.com",
     description="An AI assistant that uses various tools to interact with the system",

{jarvis_ai_assistant-0.1.74 → jarvis_ai_assistant-0.1.76}/src/jarvis/__init__.py

@@ -1,3 +1,3 @@
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.74" 
+__version__ = "0.1.76" 

{jarvis_ai_assistant-0.1.74 → jarvis_ai_assistant-0.1.76}/src/jarvis/agent.py

@@ -2,6 +2,9 @@ import time
 from typing import Dict, List, Optional
 
 import yaml
+import numpy as np
+import faiss
+import json
 
 from .models.registry import PlatformRegistry
 from .tools import ToolRegistry
@@ -9,6 +12,7 @@ from .utils import PrettyOutput, OutputType, get_multiline_input, while_success
 import os
 from datetime import datetime
 from prompt_toolkit import prompt
+from sentence_transformers import SentenceTransformer
 
 class Agent:
     def __init__(self, name: str = "Jarvis", is_sub_agent: bool = False):
@@ -26,7 +30,37 @@ class Agent:
         self.is_sub_agent = is_sub_agent
         self.prompt = ""
         self.conversation_turns = 0  
-
+        
+        # 从环境变量加载嵌入模型配置
+        self.embedding_model_name = os.environ.get("JARVIS_EMBEDDING_MODEL", "BAAI/bge-large-zh-v1.5")
+        self.embedding_dimension = 1536  # Default for many embedding models
+        
+        # 初始化嵌入模型
+        try:
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"
+            PrettyOutput.print(f"正在加载嵌入模型: {self.embedding_model_name}...", OutputType.INFO)
+            self.embedding_model = SentenceTransformer(self.embedding_model_name)
+            
+            # 预热模型并获取正确的维度
+            test_text = "这是一段测试文本，用于确保模型完全加载。"
+            test_embedding = self.embedding_model.encode(test_text, 
+                                                      convert_to_tensor=True,
+                                                      normalize_embeddings=True)
+            self.embedding_dimension = len(test_embedding)
+            PrettyOutput.print("嵌入模型加载完成", OutputType.SUCCESS)
+            
+            # 初始化HNSW索引（使用正确的维度）
+            hnsw_index = faiss.IndexHNSWFlat(self.embedding_dimension, 16)
+            hnsw_index.hnsw.efConstruction = 40
+            hnsw_index.hnsw.efSearch = 16
+            self.methodology_index = faiss.IndexIDMap(hnsw_index)
+            
+        except Exception as e:
+            PrettyOutput.print(f"加载嵌入模型失败: {str(e)}", OutputType.ERROR)
+            raise
+            
+        # 初始化方法论相关属性
+        self.methodology_data = []
 
     @staticmethod
     def extract_tool_calls(content: str) -> List[Dict]:
@@ -87,16 +121,78 @@ class Agent:
                     sleep_time = 30
                 continue
 
+    def _create_methodology_embedding(self, methodology_text: str) -> np.ndarray:
+        """为方法论文本创建嵌入向量"""
+        try:
+            # 对长文本进行截断
+            max_length = 512
+            text = ' '.join(methodology_text.split()[:max_length])
+            
+            # 使用sentence_transformers模型获取嵌入向量
+            embedding = self.embedding_model.encode([text], 
+                                                 convert_to_tensor=True,
+                                                 normalize_embeddings=True)
+            vector = np.array(embedding, dtype=np.float32)
+            return vector[0]  # 返回第一个向量，因为我们只编码了一个文本
+        except Exception as e:
+            PrettyOutput.print(f"创建方法论嵌入向量失败: {str(e)}", OutputType.ERROR)
+            return np.zeros(self.embedding_dimension, dtype=np.float32)
 
-    def _load_methodology(self) -> Dict[str, str]:
-        """加载方法论"""
+    def _load_methodology(self, user_input: str) -> Dict[str, str]:
+        """加载方法论并构建向量索引"""
+        PrettyOutput.print("加载方法论...", OutputType.PLANNING)
         user_jarvis_methodology = os.path.expanduser("~/.jarvis_methodology")
-        if os.path.exists(user_jarvis_methodology):
+        if not os.path.exists(user_jarvis_methodology):
+            return {}
+
+        try:
             with open(user_jarvis_methodology, "r", encoding="utf-8") as f:
                 data = yaml.safe_load(f)
-            PrettyOutput.print(f"从 {user_jarvis_methodology} 加载方法论: {', '.join(data.keys())}", OutputType.INFO)
-            return data
-        return {}
+
+            # 重置数据结构
+            self.methodology_data = []
+            vectors = []
+            ids = []
+
+            # 为每个方法论创建嵌入向量
+            for i, (key, value) in enumerate(data.items()):
+                PrettyOutput.print(f"向量化方法论: {key} ...", OutputType.INFO)
+                methodology_text = f"{key}\n{value}"
+                embedding = self._create_methodology_embedding(methodology_text)
+                vectors.append(embedding)
+                ids.append(i)
+                self.methodology_data.append({"key": key, "value": value})
+
+            if vectors:
+                vectors_array = np.vstack(vectors)
+                self.methodology_index.add_with_ids(vectors_array, np.array(ids))
+                query_embedding = self._create_methodology_embedding(user_input)
+                k = min(5, len(self.methodology_data))
+                PrettyOutput.print(f"检索方法论...", OutputType.INFO)
+                distances, indices = self.methodology_index.search(
+                    query_embedding.reshape(1, -1), k
+                )
+
+                relevant_methodologies = {}
+                for dist, idx in zip(distances[0], indices[0]):
+                    if idx >= 0:
+                        similarity = 1.0 / (1.0 + float(dist))
+                        methodology = self.methodology_data[idx]
+                        PrettyOutput.print(
+                            f"方法论 '{methodology['key']}' 相似度: {similarity:.3f}",
+                            OutputType.INFO
+                        )
+                        if similarity >= 0.5:
+                            relevant_methodologies[methodology["key"]] = methodology["value"]
+                        
+                if relevant_methodologies:
+                    return relevant_methodologies
+
+            return {}
+
+        except Exception as e:
+            PrettyOutput.print(f"加载方法论时发生错误: {str(e)}", OutputType.ERROR)
+            return {}
 
     def _summarize_and_clear_history(self) -> None:
         """总结当前对话历史并清空历史记录，只保留系统消息和总结
@@ -195,6 +291,38 @@ class Agent:
         self.prompt = summary_prompt
         return self._call_model(self.prompt)
 
+    def choose_tools(self, user_input: str) -> List[Dict]:
+        """根据用户输入选择工具"""
+        PrettyOutput.print("选择工具...", OutputType.PLANNING)
+        tools = self.tool_registry.get_all_tools()
+        prompt = f"""你是一个工具选择专家，请根据用户输入选择合适的工具，返回可能使用到的工具的名称。以下是可用工具：
+"""
+        for tool in tools:
+            prompt += f"- {tool['name']}: {tool['description']}\n"
+        prompt += f"用户输入: {user_input}\n"
+        prompt += f"请返回可能使用到的工具的名称，如果无法确定，请返回空列表。"
+        prompt += f"返回的格式为：\n"
+        prompt += f"<TOOL_CHOICE_START>\n"
+        prompt += f"tool_name1\n"
+        prompt += f"tool_name2\n"
+        prompt += f"<TOOL_CHOICE_END>\n"
+        model = PlatformRegistry.get_global_platform()
+        model.set_suppress_output(True)
+        try:
+            response = model.chat(prompt)
+            response = response.replace("<TOOL_CHOICE_START>", "").replace("<TOOL_CHOICE_END>", "")
+            tools_name = response.split("\n")
+            choosed_tools = []
+            for tool_name in tools_name:
+                for tool in tools:
+                    if tool['name'] == tool_name:
+                        choosed_tools.append(tool)
+                        break
+            return choosed_tools
+        except Exception as e:
+            PrettyOutput.print(f"工具选择失败: {str(e)}", OutputType.ERROR)
+            return []
+
     def run(self, user_input: str, file_list: Optional[List[str]] = None, keep_history: bool = False) -> str:
         """处理用户输入并返回响应，返回任务总结报告
         
@@ -207,29 +335,35 @@ class Agent:
             str: 任务总结报告
         """
         try:
+            PrettyOutput.section("准备环境", OutputType.PLANNING)
             if file_list:
                 self.model.upload_files(file_list)
 
             # 加载方法论
-            methodology = self._load_methodology()
+            methodology = self._load_methodology(user_input)
             methodology_prompt = ""
             if methodology:
                 methodology_prompt = f"""这是以往处理问题的标准方法论，如果当前任务与此类似，可参考：
 {methodology}
 
 """
+            tools_prompt = ""
+
+            # 选择工具
+            tools = self.choose_tools(user_input)
+            if tools:
+                tools_prompt += "可用工具:\n"
+                for tool in tools:
+                    PrettyOutput.print(f"选择工具: {tool['name']}", OutputType.INFO)
+                    tools_prompt += f"- 名称: {tool['name']}\n"
+                    tools_prompt += f"  描述: {tool['description']}\n"
+                    tools_prompt += f"  参数: {tool['parameters']}\n"
 
-            self.clear_history()
-            self.conversation_turns = 0  
-            
             # 显示任务开始
             PrettyOutput.section(f"开始新任务: {self.name}", OutputType.PLANNING)
 
-            tools_prompt = "可用工具:\n"
-            for tool in self.tool_registry.get_all_tools():
-                tools_prompt += f"- 名称: {tool['name']}\n"
-                tools_prompt += f"  描述: {tool['description']}\n"
-                tools_prompt += f"  参数: {tool['parameters']}\n"
+            self.clear_history()
+            self.conversation_turns = 0  
 
             self.model.set_system_message(f"""你是 {self.name}，一个问题处理能力强大的 AI 助手。
                                           

{jarvis_ai_assistant-0.1.74 → jarvis_ai_assistant-0.1.76/src/jarvis_ai_assistant.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.74
+Version: 0.1.76
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire