jarvis-ai-assistant 0.1.218__py3-none-any.whl → 0.1.220__py3-none-any.whl

jarvis/jarvis_rag/rag_pipeline.py

@@ -0,0 +1,177 @@
+import os
+from typing import List, Literal, Optional, cast
+
+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_mode,
+    get_rag_vector_db_path,
+    get_rag_embedding_cache_path,
+    get_rag_embedding_models,
+)
+
+
+class JarvisRAGPipeline:
+    """
+    The main orchestrator for the RAG pipeline.
+
+    This class integrates the embedding manager, retriever, and LLM to provide
+    a complete pipeline for adding documents and querying them.
+    """
+
+    def __init__(
+        self,
+        llm: Optional[LLMInterface] = None,
+        embedding_mode: Optional[Literal["performance", "accuracy"]] = None,
+        db_path: Optional[str] = None,
+        collection_name: str = "jarvis_rag_collection",
+    ):
+        """
+        Initializes the RAG pipeline.
+
+        Args:
+            llm: An instance of a class implementing LLMInterface.
+                 If None, defaults to the ToolAgent_LLM.
+            embedding_mode: The mode for the local embedding model. If None, uses config value.
+            db_path: Path to the persistent vector database. If None, uses config value.
+            collection_name: Name of the collection in the vector database.
+        """
+        # Determine the embedding model to isolate data paths
+        _embedding_mode = embedding_mode or get_rag_embedding_mode()
+        embedding_models = get_rag_embedding_models()
+        model_name = embedding_models[_embedding_mode]["model_name"]
+        sanitized_model_name = model_name.replace("/", "_").replace("\\", "_")
+
+        # If a specific db_path is given, use it. Otherwise, create a model-specific path.
+        _final_db_path = (
+            str(db_path)
+            if db_path
+            else os.path.join(get_rag_vector_db_path(), sanitized_model_name)
+        )
+        # Always create a model-specific cache path.
+        _final_cache_path = os.path.join(
+            get_rag_embedding_cache_path(), sanitized_model_name
+        )
+
+        self.embedding_manager = EmbeddingManager(
+            mode=cast(Literal["performance", "accuracy"], _embedding_mode),
+            cache_dir=_final_cache_path,
+        )
+        self.retriever = ChromaRetriever(
+            embedding_manager=self.embedding_manager,
+            db_path=_final_db_path,
+            collection_name=collection_name,
+        )
+        # Default to the ToolAgent_LLM unless a specific LLM is provided
+        self.llm = llm if llm is not None else ToolAgent_LLM()
+        self.reranker = Reranker()
+        # Use a standard LLM for the query rewriting task, not the agent
+        self.query_rewriter = QueryRewriter(JarvisPlatform_LLM())
+
+        print("✅ JarvisRAGPipeline 初始化成功。")
+
+    def add_documents(self, documents: List[Document]):
+        """
+        Adds documents to the vector knowledge base.
+
+        Args:
+            documents: A list of LangChain Document objects to add.
+        """
+        self.retriever.add_documents(documents)
+
+    def _create_prompt(
+        self, query: str, context_docs: List[Document], source_files: List[str]
+    ) -> str:
+        """Creates the final prompt for the LLM or Agent."""
+        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:
+        """
+        Performs a query against the knowledge base using a multi-query
+        retrieval and reranking pipeline.
+
+        Args:
+            query_text: The user's original question.
+            n_results: The number of final relevant chunks to retrieve.
+
+        Returns:
+            The answer generated by the LLM.
+        """
+        # 1. Rewrite the original query into multiple queries
+        rewritten_queries = self.query_rewriter.rewrite(query_text)
+
+        # 2. Retrieve initial candidates for each rewritten query
+        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)
+
+        # De-duplicate the candidate documents
+        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. Rerank the unified candidate pool against the *original* query
+        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 "我在提供的文档中找不到任何相关信息来回答您的问题。"
+
+        # Print the sources of the final retrieved documents
+        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. Create the final prompt and generate the answer
+        # We use the original query_text for the final prompt to the 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:
+    """
+    A reranker class that uses a Cross-Encoder model to re-score and sort
+    documents based on their relevance to a given query.
+    """
+
+    def __init__(self, model_name: str = "BAAI/bge-reranker-base"):
+        """
+        Initializes the Reranker.
+
+        Args:
+            model_name (str): The name of the Cross-Encoder model to use.
+        """
+        print(f"🔍 正在初始化重排模型: {model_name}...")
+        self.model = CrossEncoder(model_name)
+        print("✅ 重排模型初始化成功。")
+
+    def rerank(
+        self, query: str, documents: List[Document], top_n: int = 5
+    ) -> List[Document]:
+        """
+        Reranks a list of documents based on their relevance to the query.
+
+        Args:
+            query (str): The user's query.
+            documents (List[Document]): The list of documents retrieved from the initial search.
+            top_n (int): The number of top documents to return after reranking.
+
+        Returns:
+            List[Document]: A sorted list of the most relevant documents.
+        """
+        if not documents:
+            return []
+
+        # Create pairs of [query, document_content] for scoring
+        pairs = [[query, doc.page_content] for doc in documents]
+
+        # Get scores from the Cross-Encoder model
+        scores = self.model.predict(pairs)
+
+        # Combine documents with their scores and sort
+        doc_with_scores = list(zip(documents, scores))
+        doc_with_scores.sort(key=lambda x: x[1], reverse=True)
+
+        # Return the top N documents
+        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:
+    """
+    A retriever class that combines dense vector search (ChromaDB) and
+    sparse keyword search (BM25) for hybrid retrieval.
+    """
+
+    def __init__(
+        self,
+        embedding_manager: EmbeddingManager,
+        db_path: str,
+        collection_name: str = "jarvis_rag_collection",
+    ):
+        """
+        Initializes the ChromaRetriever.
+
+        Args:
+            embedding_manager: An instance of EmbeddingManager.
+            db_path: The file path for ChromaDB's persistent storage.
+            collection_name: The name of the collection within ChromaDB.
+        """
+        self.embedding_manager = embedding_manager
+        self.db_path = db_path
+        self.collection_name = collection_name
+
+        # Initialize ChromaDB client
+        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 Index setup
+        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):
+        """Loads the BM25 index from disk or initializes a new one."""
+        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):
+        """Saves the BM25 index to disk."""
+        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
+    ):
+        """
+        Splits, embeds, and adds documents to both ChromaDB and the BM25 index.
+        """
+        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
+
+        # Extract content, metadata, and generate IDs
+        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))]
+
+        # Add to 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 集合中。")
+
+        # Update and save BM25 index
+        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]:
+        """
+        Performs hybrid retrieval using both vector search and BM25,
+        then fuses the results using Reciprocal Rank Fusion (RRF).
+        """
+        # 1. Vector Search (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,  # Retrieve more results for fusion
+        )
+
+        # 2. Keyword Search (BM25)
+        bm25_docs = []
+        if self.bm25_index:
+            tokenized_query = query.split()
+            doc_scores = self.bm25_index.get_scores(tokenized_query)
+
+            # Get all documents from Chroma to match with BM25 scores
+            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:
+                # Create a mapping from index to document
+                bm25_results_with_docs = [
+                    (
+                        all_documents[i],
+                        all_metadatas[i],
+                        score,
+                    )
+                    for i, score in enumerate(doc_scores)
+                    if score > 0
+                ]
+
+            # Sort by score and take top results
+            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. Reciprocal Rank Fusion (RRF)
+        fused_scores: Dict[str, float] = {}
+        k = 60  # RRF ranking constant
+
+        # Process vector results
+        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)
+
+            # Create a map from document text to its ID for BM25 fusion
+            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
+                    )
+
+        # Sort fused results
+        sorted_fused_results = sorted(
+            fused_scores.items(), key=lambda x: x[1], reverse=True
+        )
+
+        # Get the final documents from ChromaDB based on fused ranking
+        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

jarvis/jarvis_tools/search_web.py

@@ -1,10 +1,20 @@
 # -*- coding: utf-8 -*-
+"""A tool for searching the web."""
 from typing import Any, Dict
 
+import httpx
+from bs4 import BeautifulSoup
+from ddgs import DDGS
+
+from jarvis.jarvis_agent import Agent
 from jarvis.jarvis_platform.registry import PlatformRegistry
+from jarvis.jarvis_utils.http import get as http_get
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 
 
 class SearchWebTool:
+    """A class to handle web searches."""
+
     name = "search_web"
     description = "搜索互联网上的信息"
     parameters = {
@@ -12,18 +22,124 @@ class SearchWebTool:
         "properties": {"query": {"type": "string", "description": "具体的问题"}},
     }
 
-    def execute(self, args: Dict[str, Any]) -> Dict[str, Any]:  # type: ignore
+    def _search_with_ddgs(self, query: str, agent: Agent) -> Dict[str, Any]:
+        # pylint: disable=too-many-locals, broad-except
+        """Performs a web search, scrapes content, and summarizes the results."""
+        try:
+            PrettyOutput.print("▶️ 使用 DuckDuckGo 开始网页搜索...", OutputType.INFO)
+            results = list(DDGS().text(query, max_results=5))
+
+            if not results:
+                return {
+                    "stdout": "未找到搜索结果。",
+                    "stderr": "未找到搜索结果。",
+                    "success": False,
+                }
+
+            urls = [r["href"] for r in results]
+            full_content = ""
+            visited_urls = []
+
+            for url in urls:
+                try:
+                    PrettyOutput.print(f"📄 正在抓取内容: {url}", OutputType.INFO)
+                    response = http_get(url, timeout=10.0, follow_redirects=True)
+                    soup = BeautifulSoup(response.text, "lxml")
+                    body = soup.find("body")
+                    if body:
+                        full_content += body.get_text(" ", strip=True) + "\n\n"
+                        visited_urls.append(url)
+                except httpx.HTTPStatusError as e:
+                    PrettyOutput.print(
+                        f"⚠️ HTTP错误 {e.response.status_code} 访问 {url}",
+                        OutputType.WARNING,
+                    )
+                except httpx.RequestError as e:
+                    PrettyOutput.print(f"⚠️ 请求错误: {e}", OutputType.WARNING)
+
+            if not full_content.strip():
+                return {
+                    "stdout": "无法从任何URL抓取有效内容。",
+                    "stderr": "抓取内容失败。",
+                    "success": False,
+                }
+
+            url_list_str = "\n".join(f"  - {u}" for u in visited_urls)
+            PrettyOutput.print(
+                f"🔍 已成功访问并处理以下URL:\n{url_list_str}", OutputType.INFO
+            )
+
+            PrettyOutput.print("🧠 正在总结内容...", OutputType.INFO)
+            summary_prompt = f"请为查询“{query}”总结以下内容：\n\n{full_content}"
+
+            if not agent.model:
+                return {
+                    "stdout": "",
+                    "stderr": "用于总结的Agent模型未找到。",
+                    "success": False,
+                }
+
+            platform_name = agent.model.platform_name()
+            model_name = agent.model.name()
+
+            model = PlatformRegistry().create_platform(platform_name)
+            if not model:
+                return {
+                    "stdout": "",
+                    "stderr": "无法创建用于总结的模型。",
+                    "success": False,
+                }
+
+            model.set_model_name(model_name)
+            model.set_suppress_output(False)
+            summary = model.chat_until_success(summary_prompt)
+
+            return {"stdout": summary, "stderr": "", "success": True}
+
+        except Exception as e:
+            PrettyOutput.print(f"❌ 网页搜索过程中发生错误: {e}", OutputType.ERROR)
+            return {
+                "stdout": "",
+                "stderr": f"网页搜索过程中发生错误: {e}",
+                "success": False,
+            }
+
+    def execute(self, args: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Executes the web search.
+
+        If the agent's model supports a native web search, it uses it.
+        Otherwise, it falls back to using DuckDuckGo Search and scraping pages.
+        """
         query = args.get("query")
-        model = PlatformRegistry().get_normal_platform()
-        model.set_web(True)
-        model.set_suppress_output(False)  # type: ignore
-        return {
-            "stdout": model.chat_until_success(query),  # type: ignore
-            "stderr": "",
-            "success": True,
-        }
+        agent = args.get("agent")
+
+        if not query:
+            return {"stdout": "", "stderr": "缺少查询参数。", "success": False}
+
+        if not isinstance(agent, Agent) or not agent.model:
+            return {
+                "stdout": "",
+                "stderr": "Agent或Agent模型未找到。",
+                "success": False,
+            }
+
+        if agent.model.support_web():
+            model = PlatformRegistry().create_platform(agent.model.platform_name())
+            if not model:
+                return {"stdout": "", "stderr": "无法创建模型。", "success": False}
+            model.set_model_name(agent.model.name())
+            model.set_web(True)
+            model.set_suppress_output(False)
+            return {
+                "stdout": model.chat_until_success(query),
+                "stderr": "",
+                "success": True,
+            }
+
+        return self._search_with_ddgs(query, agent)
 
     @staticmethod
     def check() -> bool:
-        """检查当前平台是否支持web功能"""
-        return PlatformRegistry().get_normal_platform().support_web()
+        """Check if the tool is available."""
+        return True

jarvis/jarvis_utils/config.py

@@ -3,6 +3,7 @@ import os
 from functools import lru_cache
 from typing import Any, Dict, List
 
+import torch
 import yaml  # type: ignore
 
 from jarvis.jarvis_utils.builtin_replace_map import BUILTIN_REPLACE_MAP
@@ -248,3 +249,73 @@ def get_mcp_config() -> List[Dict[str, Any]]:
         List[Dict[str, Any]]: MCP配置项列表，如果未配置则返回空列表
     """
     return GLOBAL_CONFIG_DATA.get("JARVIS_MCP", [])
+
+
+# ==============================================================================
+# RAG Framework Configuration
+# ==============================================================================
+
+EMBEDDING_MODELS = {
+    "performance": {
+        "model_name": "BAAI/bge-base-zh-v1.5",
+        "model_kwargs": {"device": "cuda" if torch.cuda.is_available() else "cpu"},
+        "encode_kwargs": {"normalize_embeddings": True},
+        "show_progress": True,
+    },
+    "accuracy": {
+        "model_name": "BAAI/bge-large-zh-v1.5",
+        "model_kwargs": {"device": "cuda" if torch.cuda.is_available() else "cpu"},
+        "encode_kwargs": {"normalize_embeddings": True},
+        "show_progress": True,
+    },
+}
+
+
+def get_rag_config() -> Dict[str, Any]:
+    """
+    获取RAG框架的配置。
+
+    返回:
+        Dict[str, Any]: RAG配置字典
+    """
+    return GLOBAL_CONFIG_DATA.get("JARVIS_RAG", {})
+
+
+def get_rag_embedding_models() -> Dict[str, Any]:
+    """
+    获取RAG嵌入模型的定义。
+
+    返回:
+        Dict[str, Any]: 嵌入模型配置字典
+    """
+    return EMBEDDING_MODELS
+
+
+def get_rag_embedding_mode() -> str:
+    """
+    获取RAG嵌入模型的模式。
+
+    返回:
+        str: 'performance' 或 'accuracy'
+    """
+    return get_rag_config().get("embedding_mode", "performance")
+
+
+def get_rag_embedding_cache_path() -> str:
+    """
+    获取RAG嵌入缓存的路径。
+
+    返回:
+        str: 缓存路径
+    """
+    return get_rag_config().get("embedding_cache_path", ".jarvis/rag/embeddings")
+
+
+def get_rag_vector_db_path() -> str:
+    """
+    获取RAG向量数据库的路径。
+
+    返回:
+        str: 数据库路径
+    """
+    return get_rag_config().get("vector_db_path", ".jarvis/rag/vectordb")