搭建单机版的知识库问题系统

ollama run deepseek-r1:1.5b
ollama run nomic-embed-text:latest

from pathlib import Path
import re
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
from chromadb.config import Settings
import sys
from typing import Optional


class RAGSystem:
    def __init__(self, ollama_base_url: str = "http://localhost:11434"):
        # 初始化配置
        self.data_dir = Path("./data")
        self.db_dir = Path("./chroma_db")
        self.ollama_base_url = ollama_base_url
        self._init_components()

    def _init_components(self):
        # 初始化LLM和Embedding模型
        self.llm = Ollama(
            model="deepseek-r1:1.5b",
            base_url=self.ollama_base_url,
            temperature=0.1
        )

        self.embed_model = OllamaEmbedding(
            model_name="nomic-embed-text:latest",
            base_url=self.ollama_base_url
        )

        # 初始化ChromaDB
        self.chroma_client = chromadb.PersistentClient(
            path=str(self.db_dir),
            settings=Settings(allow_reset=True)
        )
        self.vector_store = ChromaVectorStore(
            chroma_collection=self.chroma_client.get_or_create_collection(
                "rag_demo")
        )

        # 配置文本分割器
        self.node_parser = SentenceSplitter(
            chunk_size=512,
            chunk_overlap=50
        )

    def build_index(self):
        # 加载文档
        documents = SimpleDirectoryReader(
            input_dir=self.data_dir,
            recursive=True,
            required_exts=[".txt", ".docx", ".jsonl"]
        ).load_data()

        print(f"成功加载 {len(documents)} 个文档")
        for doc in documents:
            print(f"- 文档ID: {doc.doc_id}, 长度: {len(doc.text)} 字符")

        # 处理文档并构建索引
        index = VectorStoreIndex.from_documents(
            documents=documents,
            embed_model=self.embed_model,
            transformations=[self.node_parser],
            vector_store=self.vector_store,
            show_progress=True
        )

        # 验证索引
        print("\n索引验证:")
        print(f"总节点数: {len(index.docstore.docs)}")
        sample_node_id = next(iter(index.docstore.docs.keys()))
        sample_node = index.docstore.get_node(sample_node_id)
        print(
            f"示例节点内容 ({len(sample_node.text)} 字符):\n{sample_node.text[:200]}...")

        return index

    def query_engine(self, index):
        return index.as_query_engine(
            llm=self.llm,
            similarity_top_k=3,
            response_mode="compact",
            timeout=60,  # 设置60秒超时
            # 优化检索参数
            # vector_store_query_mode="hybrid",  # 混合检索模式
            vector_store_kwargs={
                "search_kwargs": {"k": 50}  # 扩大候选集提升召回
            }
        )


if __name__ == "__main__":

    rag = RAGSystem()

    if "ingest" in sys.argv:
        print("Building vector index...")
        index = rag.build_index()
        index.storage_context.persist(persist_dir=rag.db_dir)
        print(f"Index persisted to {rag.db_dir}")

    if "query" in sys.argv:
        from llama_index.core import StorageContext
        from llama_index.core import load_index_from_storage  # 添加正确的导入

        # 加载已有索引
        # 修复后代码（从持久化目录加载）
        storage_context = StorageContext.from_defaults(
            persist_dir=rag.db_dir
        )
        # 修复后代码（使用load_index_from_storage）
        index = load_index_from_storage(
            storage_context=storage_context,
            embed_model=rag.embed_model
        )

        ##
        # 验证索引
        print("\n索引验证:")
        print(f"总节点数: {len(index.docstore.docs)}")
        sample_node_id = next(iter(index.docstore.docs.keys()))
        sample_node = index.docstore.get_node(sample_node_id)
        print(
            f"示例节点内容 ({len(sample_node.text)} 字符):\n{sample_node.text[:200]}...")

        # 获取查询问题
        if len(sys.argv) < 3:
            print("请提供查询问题,格式: python src/rag_pipeline.py query '你的问题'")
            sys.exit(1)

        query = sys.argv[2]
        qe = rag.query_engine(index)
        # 添加详细调试信息
        print("\n=== 查询调试信息 ===")
        print(f"查询问题: {query}")

        # 显示检索到的节点
        print("\n检索到的节点:")
        nodes = qe.retrieve(query)
        for i, node in enumerate(nodes):
            print(f"节点 {i+1}:")
            print(f"分数: {node.score:.4f}")
            print(f"内容: {node.text[:200]}...\n")

        # 执行查询
        response = qe.query(query)

        # 输出结果
        print("\n=== 最终答案 ===")
        # 过滤掉包含"<think>"的响应内容
        filtered_response = re.sub(
            r'<think>.*?</think>', '', response.response, flags=re.DOTALL).strip()
        final_answer = filtered_response if filtered_response else "未找到相关答案"
        print(final_answer)

        # 显示参考文档详情
        print("\n=== 参考文档 ===")
        if response.source_nodes:
            for node in response.source_nodes:
                print(f"- 相似度: {node.score:.4f}")
                print(f"  元数据: {node.metadata}")
                print(f"  内容: {node.text[:200]}...\n")
        else:
            print("未找到相关参考文档")

poetry run python src/rag_pipeline.py ingest

  poetry run python src/rag_pipeline.py query "伽利略支持哪些SDK?"