jarvis-ai-assistant 0.1.219__py3-none-any.whl → 0.1.221__py3-none-any.whl

jarvis/jarvis_rag/llm_interface.py

@@ -0,0 +1,128 @@
+from abc import ABC, abstractmethod
+import os
+import os
+from abc import ABC, abstractmethod
+
+from jarvis.jarvis_agent import Agent as JarvisAgent
+from jarvis.jarvis_platform.base import BasePlatform
+from jarvis.jarvis_platform.registry import PlatformRegistry
+
+
+class LLMInterface(ABC):
+    """
+    大型语言模型接口的抽象基类。
+
+    该类定义了与远程LLM交互的标准接口。
+    任何LLM提供商（如OpenAI、Anthropic等）都应作为该接口的子类来实现。
+    """
+
+    @abstractmethod
+    def generate(self, prompt: str, **kwargs) -> str:
+        """
+        根据给定的提示从LLM生成响应。
+
+        参数:
+            prompt: 发送给LLM的输入提示。
+            **kwargs: LLM API调用的其他关键字参数
+                      （例如，temperature, max_tokens）。
+
+        返回:
+            由LLM生成的文本响应。
+        """
+        pass
+
+
+class ToolAgent_LLM(LLMInterface):
+    """
+    LLMInterface的一个实现，它使用一个能操作工具的JarvisAgent来生成最终响应。
+    """
+
+    def __init__(self):
+        """
+        初始化工具-代理 LLM 包装器。
+        """
+        print("🤖 已初始化工具 Agent 作为最终应答者。")
+        self.allowed_tools = ["read_code", "execute_script"]
+        # 为代理提供一个通用的系统提示
+        self.system_prompt = "You are a helpful assistant. Please answer the user's question based on the provided context. You can use tools to find more information if needed."
+        self.summary_prompt = """
+<report>
+请为本次问答任务生成一个总结报告，包含以下内容：
+
+1. **原始问题**: 重述用户最开始提出的问题。
+2. **关键信息来源**: 总结你是基于哪些关键信息或文件得出的结论。
+3. **最终答案**: 给出最终的、精炼的回答。
+</report>
+"""
+
+    def generate(self, prompt: str, **kwargs) -> str:
+        """
+        使用受限的工具集运行JarvisAgent以生成答案。
+
+        参数:
+            prompt: 要发送给代理的完整提示，包括上下文。
+            **kwargs: 已忽略，为保持接口兼容性而保留。
+
+        返回:
+            由代理生成的最终答案。
+        """
+        try:
+            # 使用RAG上下文的特定设置初始化代理
+            agent = JarvisAgent(
+                system_prompt=self.system_prompt,
+                use_tools=self.allowed_tools,
+                auto_complete=True,
+                use_methodology=False,
+                use_analysis=False,
+                need_summary=True,
+                summary_prompt=self.summary_prompt,
+            )
+
+            # 代理的run方法需要'user_input'参数
+            final_answer = agent.run(user_input=prompt)
+            return str(final_answer)
+
+        except Exception as e:
+            print(f"❌ Agent 在执行过程中发生错误: {e}")
+            return "错误: Agent 未能成功生成回答。"
+
+
+class JarvisPlatform_LLM(LLMInterface):
+    """
+    项目内部平台的LLMInterface实现。
+
+    该类使用PlatformRegistry来获取配置的“普通”模型。
+    """
+
+    def __init__(self):
+        """
+        初始化Jarvis平台LLM客户端。
+        """
+        try:
+            self.registry = PlatformRegistry.get_global_platform_registry()
+            self.platform: BasePlatform = self.registry.get_normal_platform()
+            self.platform.set_suppress_output(
+                False
+            )  # 确保模型没有控制台输出
+            print(f"🚀 已初始化 Jarvis 平台 LLM，模型: {self.platform.name()}")
+        except Exception as e:
+            print(f"❌ 初始化 Jarvis 平台 LLM 失败: {e}")
+            raise
+
+    def generate(self, prompt: str, **kwargs) -> str:
+        """
+        向本地平台模型发送提示并返回响应。
+
+        参数:
+            prompt: 用户的提示。
+            **kwargs: 已忽略，为保持接口兼容性而保留。
+
+        返回:
+            由平台模型生成的响应。
+        """
+        try:
+            # 使用健壮的chat_until_success方法
+            return self.platform.chat_until_success(prompt)
+        except Exception as e:
+            print(f"❌ 调用 Jarvis 平台模型时发生错误: {e}")
+            return "错误: 无法从本地LLM获取响应。"

jarvis/jarvis_rag/query_rewriter.py

@@ -0,0 +1,62 @@
+from typing import List
+from .llm_interface import LLMInterface
+
+
+class QueryRewriter:
+    """
+    使用LLM将用户的查询重写为多个不同的搜索查询，以提高检索召回率。
+    """
+
+    def __init__(self, llm: LLMInterface):
+        """
+        初始化QueryRewriter。
+
+        参数:
+            llm: 实现LLMInterface接口的类的实例。
+        """
+        self.llm = llm
+        self.rewrite_prompt_template = self._create_prompt_template()
+
+    def _create_prompt_template(self) -> str:
+        """为多查询重写任务创建提示模板。"""
+        return """
+你是一个精通检索的AI助手。你的任务是将以下这个单一的用户问题，从不同角度改写成 3 个不同的、但语义上相关的搜索查询。这有助于在知识库中进行更全面的搜索。
+
+请遵循以下原则：
+1.  **多样性**：生成的查询应尝试使用不同的关键词和表述方式。
+2.  **保留核心意图**：所有查询都必须围绕原始问题的核心意图。
+3.  **简洁性**：每个查询都应该是独立的、可以直接用于搜索的短语或问题。
+4.  **格式要求**：请直接输出 3 个查询，每个查询占一行，用换行符分隔。不要添加任何编号、前缀或解释。
+
+原始问题:
+---
+{query}
+---
+
+3个改写后的查询 (每行一个):
+"""
+
+    def rewrite(self, query: str) -> List[str]:
+        """
+        使用LLM将用户查询重写为多个查询。
+
+        参数:
+            query: 原始用户查询。
+
+        返回:
+            一个经过重写、搜索优化的查询列表。
+        """
+        prompt = self.rewrite_prompt_template.format(query=query)
+        print(f"✍️  正在将原始查询重写为多个搜索查询...")
+
+        response_text = self.llm.generate(prompt)
+        rewritten_queries = [
+            line.strip() for line in response_text.strip().split("\n") if line.strip()
+        ]
+
+        # 同时包含原始查询以保证鲁棒性
+        if query not in rewritten_queries:
+            rewritten_queries.insert(0, query)
+
+        print(f"✅ 生成了 {len(rewritten_queries)} 个查询变体。")
+        return rewritten_queries

jarvis/jarvis_rag/rag_pipeline.py

@@ -0,0 +1,174 @@
+import os
+from typing import List, Optional
+
+from langchain.docstore.document import Document
+
+from .embedding_manager import EmbeddingManager
+from .llm_interface import JarvisPlatform_LLM, LLMInterface, ToolAgent_LLM
+from .query_rewriter import QueryRewriter
+from .reranker import Reranker
+from .retriever import ChromaRetriever
+from jarvis.jarvis_utils.config import (
+    get_rag_embedding_model,
+    get_rag_rerank_model,
+    get_rag_vector_db_path,
+    get_rag_embedding_cache_path,
+)
+
+
+class JarvisRAGPipeline:
+    """
+    RAG管道的主要协调器。
+
+    该类集成了嵌入管理器、检索器和LLM，为添加文档和查询
+    提供了一个完整的管道。
+    """
+
+    def __init__(
+        self,
+        llm: Optional[LLMInterface] = None,
+        embedding_model: Optional[str] = None,
+        db_path: Optional[str] = None,
+        collection_name: str = "jarvis_rag_collection",
+    ):
+        """
+        初始化RAG管道。
+
+        参数:
+            llm: 实现LLMInterface接口的类的实例。
+                 如果为None，则默认为ToolAgent_LLM。
+            embedding_model: 嵌入模型的名称。如果为None，则使用配置值。
+            db_path: 持久化向量数据库的路径。如果为None，则使用配置值。
+            collection_name: 向量数据库中集合的名称。
+        """
+        # 确定嵌入模型以隔离数据路径
+        model_name = embedding_model or get_rag_embedding_model()
+        sanitized_model_name = model_name.replace("/", "_").replace("\\", "_")
+
+        # 如果给定了特定的db_path，则使用它。否则，创建一个特定于模型的路径。
+        _final_db_path = (
+            str(db_path)
+            if db_path
+            else os.path.join(get_rag_vector_db_path(), sanitized_model_name)
+        )
+        # 始终创建一个特定于模型的缓存路径。
+        _final_cache_path = os.path.join(
+            get_rag_embedding_cache_path(), sanitized_model_name
+        )
+
+        self.embedding_manager = EmbeddingManager(
+            model_name=model_name,
+            cache_dir=_final_cache_path,
+        )
+        self.retriever = ChromaRetriever(
+            embedding_manager=self.embedding_manager,
+            db_path=_final_db_path,
+            collection_name=collection_name,
+        )
+        # 除非提供了特定的LLM，否则默认为ToolAgent_LLM
+        self.llm = llm if llm is not None else ToolAgent_LLM()
+        self.reranker = Reranker(model_name=get_rag_rerank_model())
+        # 使用标准LLM执行查询重写任务，而不是代理
+        self.query_rewriter = QueryRewriter(JarvisPlatform_LLM())
+
+        print("✅ JarvisRAGPipeline 初始化成功。")
+
+    def add_documents(self, documents: List[Document]):
+        """
+        将文档添加到向量知识库。
+
+        参数:
+            documents: 要添加的LangChain文档对象列表。
+        """
+        self.retriever.add_documents(documents)
+
+    def _create_prompt(
+        self, query: str, context_docs: List[Document], source_files: List[str]
+    ) -> str:
+        """为LLM或代理创建最终的提示。"""
+        context = "\n\n".join([doc.page_content for doc in context_docs])
+        sources_text = "\n".join([f"- {source}" for source in source_files])
+
+        prompt_template = f"""
+        你是一个专家助手。请根据用户的问题，结合下面提供的参考信息来回答。
+
+        **重要**: 提供的上下文和文件列表**仅供参考**，可能不完整或已过时。在回答前，你应该**优先使用工具（如 read_code）来获取最新、最准确的信息**。
+
+        参考文件列表:
+        ---
+        {sources_text}
+        ---
+
+        参考上下文:
+        ---
+        {context}
+        ---
+
+        问题: {query}
+
+        回答:
+        """
+        return prompt_template.strip()
+
+    def query(self, query_text: str, n_results: int = 5) -> str:
+        """
+        使用多查询检索和重排管道对知识库执行查询。
+
+        参数:
+            query_text: 用户的原始问题。
+            n_results: 要检索的最终相关块的数量。
+
+        返回:
+            由LLM生成的答案。
+        """
+        # 1. 将原始查询重写为多个查询
+        rewritten_queries = self.query_rewriter.rewrite(query_text)
+
+        # 2. 为每个重写的查询检索初始候选文档
+        all_candidate_docs = []
+        for q in rewritten_queries:
+            print(f"🔍 正在为查询变体 '{q}' 进行混合检索...")
+            candidates = self.retriever.retrieve(q, n_results=n_results * 2)
+            all_candidate_docs.extend(candidates)
+
+        # 对候选文档进行去重
+        unique_docs_dict = {doc.page_content: doc for doc in all_candidate_docs}
+        unique_candidate_docs = list(unique_docs_dict.values())
+
+        if not unique_candidate_docs:
+            return "我在提供的文档中找不到任何相关信息来回答您的问题。"
+
+        # 3. 根据*原始*查询对统一的候选池进行重排
+        print(
+            f"🔍 正在对 {len(unique_candidate_docs)} 个候选文档进行重排（基于原始问题）..."
+        )
+        retrieved_docs = self.reranker.rerank(
+            query_text, unique_candidate_docs, top_n=n_results
+        )
+
+        if not retrieved_docs:
+            return "我在提供的文档中找不到任何相关信息来回答您的问题。"
+
+        # 打印最终检索到的文档的来源
+        sources = sorted(
+            list(
+                {
+                    doc.metadata["source"]
+                    for doc in retrieved_docs
+                    if "source" in doc.metadata
+                }
+            )
+        )
+        if sources:
+            print(f"📚 根据以下文档回答:")
+            for source in sources:
+                print(f"  - {source}")
+
+        # 4. 创建最终提示并生成答案
+        # 我们使用原始的query_text作为给LLM的最终提示
+        prompt = self._create_prompt(query_text, retrieved_docs, sources)
+
+        print("🤖 正在从LLM生成答案...")
+        answer = self.llm.generate(prompt)
+
+        return answer

jarvis/jarvis_rag/reranker.py

@@ -0,0 +1,56 @@
+from typing import List
+
+from langchain.docstore.document import Document
+from sentence_transformers.cross_encoder import (  # type: ignore
+    CrossEncoder,
+)
+
+
+class Reranker:
+    """
+    一个重排器类，使用Cross-Encoder模型根据文档与给定查询的相关性
+    对文档进行重新评分和排序。
+    """
+
+    def __init__(self, model_name: str):
+        """
+        初始化重排器。
+
+        参数:
+            model_name (str): 要使用的Cross-Encoder模型的名称。
+        """
+        print(f"🔍 正在初始化重排模型: {model_name}...")
+        self.model = CrossEncoder(model_name)
+        print("✅ 重排模型初始化成功。")
+
+    def rerank(
+        self, query: str, documents: List[Document], top_n: int = 5
+    ) -> List[Document]:
+        """
+        根据文档与查询的相关性对文档列表进行重排。
+
+        参数:
+            query (str): 用户的查询。
+            documents (List[Document]): 从初始搜索中检索到的文档列表。
+            top_n (int): 重排后要返回的顶部文档数。
+
+        返回:
+            List[Document]: 一个已排序的最相关文档列表。
+        """
+        if not documents:
+            return []
+
+        # 创建 [查询, 文档内容] 对用于评分
+        pairs = [[query, doc.page_content] for doc in documents]
+
+        # 从Cross-Encoder模型获取分数
+        scores = self.model.predict(pairs)
+
+        # 将文档与它们的分数结合并排序
+        doc_with_scores = list(zip(documents, scores))
+        doc_with_scores.sort(key=lambda x: x[1], reverse=True)  # type: ignore
+
+        # 返回前N个文档
+        reranked_docs = [doc for doc, score in doc_with_scores[:top_n]]
+
+        return reranked_docs

jarvis/jarvis_rag/retriever.py

@@ -0,0 +1,201 @@
+import os
+import pickle
+from typing import Any, Dict, List, cast
+
+import chromadb
+from langchain.docstore.document import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from rank_bm25 import BM25Okapi  # type: ignore
+
+from .embedding_manager import EmbeddingManager
+
+
+class ChromaRetriever:
+    """
+    一个检索器类，它结合了密集向量搜索（ChromaDB）和稀疏关键字搜索（BM25）
+    以实现混合检索。
+    """
+
+    def __init__(
+        self,
+        embedding_manager: EmbeddingManager,
+        db_path: str,
+        collection_name: str = "jarvis_rag_collection",
+    ):
+        """
+        初始化ChromaRetriever。
+
+        参数:
+            embedding_manager: EmbeddingManager的实例。
+            db_path: ChromaDB持久化存储的文件路径。
+            collection_name: ChromaDB中集合的名称。
+        """
+        self.embedding_manager = embedding_manager
+        self.db_path = db_path
+        self.collection_name = collection_name
+
+        # 初始化ChromaDB客户端
+        self.client = chromadb.PersistentClient(path=self.db_path)
+        self.collection = self.client.get_or_create_collection(
+            name=self.collection_name
+        )
+        print(
+            f"✅ ChromaDB 客户端已在 '{db_path}' 初始化，集合为 '{collection_name}'。"
+        )
+
+        # BM25索引设置
+        self.bm25_index_path = os.path.join(self.db_path, f"{collection_name}_bm25.pkl")
+        self._load_or_initialize_bm25()
+
+    def _load_or_initialize_bm25(self):
+        """从磁盘加载BM25索引或初始化一个新索引。"""
+        if os.path.exists(self.bm25_index_path):
+            print("🔍 正在加载现有的 BM25 索引...")
+            with open(self.bm25_index_path, "rb") as f:
+                data = pickle.load(f)
+                self.bm25_corpus = data["corpus"]
+                self.bm25_index = BM25Okapi(self.bm25_corpus)
+            print("✅ BM25 索引加载成功。")
+        else:
+            print("⚠️ 未找到 BM25 索引，将初始化一个新的。")
+            self.bm25_corpus = []
+            self.bm25_index = None
+
+    def _save_bm25_index(self):
+        """将BM25索引保存到磁盘。"""
+        if self.bm25_index:
+            print("💾 正在保存 BM25 索引...")
+            with open(self.bm25_index_path, "wb") as f:
+                pickle.dump({"corpus": self.bm25_corpus, "index": self.bm25_index}, f)
+            print("✅ BM25 索引保存成功。")
+
+    def add_documents(
+        self, documents: List[Document], chunk_size=1000, chunk_overlap=100
+    ):
+        """
+        将文档拆分、嵌入，并添加到ChromaDB和BM25索引中。
+        """
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size, chunk_overlap=chunk_overlap
+        )
+        chunks = text_splitter.split_documents(documents)
+
+        print(f"📄 已将 {len(documents)} 个文档拆分为 {len(chunks)} 个块。")
+
+        if not chunks:
+            return
+
+        # 提取内容、元数据并生成ID
+        chunk_texts = [chunk.page_content for chunk in chunks]
+        metadatas = [chunk.metadata for chunk in chunks]
+        start_id = self.collection.count()
+        ids = [f"doc_{i}" for i in range(start_id, start_id + len(chunks))]
+
+        # 添加到ChromaDB
+        embeddings = self.embedding_manager.embed_documents(chunk_texts)
+        self.collection.add(
+            ids=ids,
+            embeddings=cast(Any, embeddings),
+            documents=chunk_texts,
+            metadatas=cast(Any, metadatas),
+        )
+        print(f"✅ 成功将 {len(chunks)} 个块添加到 ChromaDB 集合中。")
+
+        # 更新并保存BM25索引
+        tokenized_chunks = [doc.split() for doc in chunk_texts]
+        self.bm25_corpus.extend(tokenized_chunks)
+        self.bm25_index = BM25Okapi(self.bm25_corpus)
+        self._save_bm25_index()
+
+    def retrieve(self, query: str, n_results: int = 5) -> List[Document]:
+        """
+        使用向量搜索和BM25执行混合检索，然后使用倒数排序融合（RRF）
+        对结果进行融合。
+        """
+        # 1. 向量搜索 (ChromaDB)
+        query_embedding = self.embedding_manager.embed_query(query)
+        vector_results = self.collection.query(
+            query_embeddings=cast(Any, [query_embedding]),
+            n_results=n_results * 2,  # 检索更多结果用于融合
+        )
+
+        # 2. 关键字搜索 (BM25)
+        bm25_docs = []
+        if self.bm25_index:
+            tokenized_query = query.split()
+            doc_scores = self.bm25_index.get_scores(tokenized_query)
+
+            # 从Chroma获取所有文档以匹配BM25分数
+            all_docs_in_collection = self.collection.get()
+            all_documents = all_docs_in_collection.get("documents")
+            all_metadatas = all_docs_in_collection.get("metadatas")
+
+            bm25_results_with_docs = []
+            if all_documents and all_metadatas:
+                # 创建从索引到文档的映射
+                bm25_results_with_docs = [
+                    (
+                        all_documents[i],
+                        all_metadatas[i],
+                        score,
+                    )
+                    for i, score in enumerate(doc_scores)
+                    if score > 0
+                ]
+
+            # 按分数排序并取最高结果
+            bm25_results_with_docs.sort(key=lambda x: x[2], reverse=True)
+
+            for doc_text, metadata, _ in bm25_results_with_docs[: n_results * 2]:
+                bm25_docs.append(Document(page_content=doc_text, metadata=metadata))
+
+        # 3. 倒数排序融合 (RRF)
+        fused_scores: Dict[str, float] = {}
+        k = 60  # RRF排名常数
+
+        # 处理向量结果
+        if vector_results and vector_results["ids"] and vector_results["documents"]:
+            vec_ids = vector_results["ids"][0]
+            vec_texts = vector_results["documents"][0]
+
+            for rank, doc_id in enumerate(vec_ids):
+                fused_scores[doc_id] = fused_scores.get(doc_id, 0) + 1 / (k + rank)
+
+            # 为BM25融合创建从文档文本到其ID的映射
+            doc_text_to_id = {text: doc_id for text, doc_id in zip(vec_texts, vec_ids)}
+
+            for rank, doc in enumerate(bm25_docs):
+                bm25_doc_id = doc_text_to_id.get(doc.page_content)
+                if bm25_doc_id:
+                    fused_scores[bm25_doc_id] = fused_scores.get(bm25_doc_id, 0) + 1 / (
+                        k + rank
+                    )
+
+        # 对融合结果进行排序
+        sorted_fused_results = sorted(
+            fused_scores.items(), key=lambda x: x[1], reverse=True
+        )
+
+        # 根据融合排名从ChromaDB获取最终文档
+        final_doc_ids = [item[0] for item in sorted_fused_results[:n_results]]
+
+        if not final_doc_ids:
+            return []
+
+        final_docs_data = self.collection.get(ids=final_doc_ids)
+
+        retrieved_docs = []
+        if final_docs_data:
+            final_documents = final_docs_data.get("documents")
+            final_metadatas = final_docs_data.get("metadatas")
+
+            if final_documents and final_metadatas:
+                for doc_text, metadata in zip(final_documents, final_metadatas):
+                    if doc_text is not None and metadata is not None:
+                        retrieved_docs.append(
+                            Document(
+                                page_content=cast(str, doc_text), metadata=metadata
+                            )
+                        )
+
+        return retrieved_docs