zyk-rag 0.1.0__py3-none-any.whl

zyk_rag-0.1.0.dist-info/METADATA

@@ -0,0 +1,83 @@
+Metadata-Version: 2.4
+Name: zyk_rag
+Version: 0.1.0
+Summary: A modular RAG retrieval class based on LangChain and Chroma
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: chromadb
+Requires-Dist: langchain-chroma
+Requires-Dist: langchain-community
+Requires-Dist: langchain-core
+Requires-Dist: langchain-openai
+Requires-Dist: langchain-classic
+Requires-Dist: jieba
+Requires-Dist: rank-bm25
+Requires-Dist: python-dotenv
+Requires-Dist: dashscope
+Provides-Extra: publish
+Requires-Dist: build; extra == "publish"
+Requires-Dist: twine; extra == "publish"
+
+# langchain-modular-rag
+
+提供一个 `ZykRAG` 类，用于连接已有的 Chroma 知识库并执行在线检索。
+
+Python 3.11 及以上。
+
+## 安装
+
+```bash
+pip install zyk_rag
+```
+
+## 使用
+
+```python
+from zyk_rag import ZykRAG
+
+rag = ModularRAG({
+    "persist_directory": "./chroma_db",
+    "collection_name": "study",
+    "embedding_model": "text-embedding-v4",
+    "top_k": 5,
+    "recall_top_k": 20,
+    "pre_retrieval": {
+        "enabled": False,
+        "rewrite_query": False,
+        "hyde": False,
+        "model": "qwen-plus",
+    },
+    "hybrid_retrieval": {
+        "enabled": True,
+        "weights": [0.5, 0.5],
+    },
+    "post_retrieval": {
+        "enabled": True,
+        "deduplicate": True,
+        "rerank": False,
+        "rerank_model": "qwen3-rerank",
+        "similarity_filter": False,
+        "similarity_threshold": 0.4,
+        "compress_context": False,
+        "llm_model": "qwen-plus",
+        "reorder": False,
+    },
+})
+
+results = rag.search("公司打卡时间？")
+for document in results:
+    print(document.metadata)
+    print(document.page_content)
+```
+
+## 环境变量
+
+需要配置通义 API Key：
+
+```bash
+export DASHSCOPE_API_KEY="你的 API Key"
+或
+setx DASHSCOPE_API_KEY="你的 API Key"
+```
+
+`persist_directory` 指向已经完成入库的 Chroma 数据库。这个包当前只提供在线检索，不包含你的个人知识库数据，也不会自动创建或导入知识库。

zyk_rag-0.1.0.dist-info/RECORD

@@ -0,0 +1,5 @@
+zyk_rag.py,sha256=6lOI84O1sezHMZWeGAhjB4ovYQCq6ZE_f5_aLQxrTt0,12757
+zyk_rag-0.1.0.dist-info/METADATA,sha256=h1-2iXfyqj8lPIxtRikl3hgRpiy1CzMdwSInAegJz-0,2096
+zyk_rag-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+zyk_rag-0.1.0.dist-info/top_level.txt,sha256=gxW03LggUjSL0D65cpDypMuIN1nQe2uNwFOds0hFDsk,8
+zyk_rag-0.1.0.dist-info/RECORD,,

zyk_rag-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

zyk_rag-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+zyk_rag

zyk_rag.py

@@ -0,0 +1,318 @@
+
+class ZykRAG:
+    def __init__(self, config: dict):
+        import os
+        from pathlib import Path
+        from dotenv import load_dotenv
+        from langchain_chroma import Chroma
+        from langchain_community.embeddings import DashScopeEmbeddings
+
+        load_dotenv()
+        self.config = config
+        directory = config.get("persist_directory", "../chroma_db")
+        if not (Path(directory).expanduser() / "chroma.sqlite3").is_file():
+            raise FileNotFoundError(f"知识库不存在：{directory}")
+        self.embeddings = DashScopeEmbeddings(
+            model=config.get("embedding_model", "text-embedding-v4"),
+            dashscope_api_key=os.getenv("DASHSCOPE_API_KEY"),
+        )
+        self.db = Chroma(
+            collection_name=config.get("collection_name", "study"),
+            persist_directory=str(Path(directory).expanduser()),
+            embedding_function=self.embeddings,
+            create_collection_if_not_exists=False,
+        )
+
+    def search(self, question: str, top_k=None, filters=None):
+        from langchain_core.documents import Document
+
+        if not question.strip():
+            raise ValueError("问题不能为空")
+        limit = self.config.get("top_k", 5) if top_k is None else top_k
+        if not isinstance(limit, int) or limit <= 0:
+            raise ValueError("top_k 必须是正整数")
+        count = max(self.config.get("recall_top_k", 20), limit)
+        pre = self.config.get("pre_retrieval", {})
+        hybrid = self.config.get("hybrid_retrieval", {})
+        post = self.config.get("post_retrieval", {})
+
+        query = question
+        if pre.get("enabled", False):
+            query = self.normalize_query(query)
+            if pre.get("rewrite_query", False):
+                query = self.rewrite_query(query, model=pre.get("model", "qwen-plus"))
+        keyword_query = query
+        dense_query = query
+        if pre.get("enabled", False) and pre.get("hyde", False):
+            dense_query = self.generate_hyde(query, model=pre.get("model", "qwen-plus"))
+
+        conditions = []
+        for condition in (self.config.get("filters"), filters):
+            if condition:
+                conditions.extend({key: value} for key, value in condition.items())
+        where = None
+        if len(conditions) == 1:
+            where = conditions[0]
+        elif len(conditions) > 1:
+            where = {"$and": conditions}
+
+        if hybrid.get("enabled", False):
+            import jieba
+            from langchain_community.retrievers import BM25Retriever
+            from langchain_core.runnables import RunnableLambda
+
+            stored = self.db.get(where=where, include=["documents", "metadatas"])
+            corpus = []
+            for index, text in enumerate(stored["documents"] or []):
+                corpus.append(Document(
+                    id=stored["ids"][index],
+                    page_content=text,
+                    metadata=stored["metadatas"][index] or {},
+                ))
+            if not corpus:
+                return []
+            vector_retriever = self.db.as_retriever(
+                search_type="similarity",
+                search_kwargs={"k": count, "filter": where},
+            )
+            if dense_query != keyword_query:
+                vector_retriever = (
+                    RunnableLambda(lambda incoming_query: dense_query) | vector_retriever
+                )
+            bm25_retriever = BM25Retriever.from_documents(
+                corpus,
+                preprocess_func=lambda text: list(jieba.cut(text)),
+                k=count,
+            )
+            documents = self.rrf_fusion(
+                keyword_query, vector_retriever, bm25_retriever,
+                weights=hybrid.get("weights", [0.5, 0.5]),
+            )
+        else:
+            documents = self.vector_search(dense_query, self.db, count, where)
+
+        if post.get("enabled", False):
+            if post.get("deduplicate", True):
+                documents = self.deduplicate(documents)
+            if post.get("rerank", False):
+                documents = self.rerank_documents(
+                    question, documents, top_k=len(documents),
+                    model=post.get("rerank_model", "qwen3-rerank"),
+                )
+            if post.get("similarity_filter", False):
+                documents = self.similarity_filter(
+                    question, documents, self.embeddings,
+                    threshold=post.get("similarity_threshold", 0.4),
+                )
+
+        documents = documents[:limit]
+        for rank, document in enumerate(documents, start=1):
+            document.metadata["retrieval_rank"] = rank
+
+        if post.get("enabled", False):
+            if post.get("compress_context", False):
+                documents = self.compress_context(
+                    question, documents, model=post.get("llm_model", "qwen-plus")
+                )
+            if post.get("reorder", False):
+                documents = self.reorder_documents(documents)
+        for position, document in enumerate(documents, start=1):
+            document.metadata["context_position"] = position
+        return documents
+    @staticmethod
+    # ==================== 1. 归一化查询词：移除所有空格并首尾去空格 ====================
+    def normalize_query(question: str):
+        import re
+        return re.sub(r"\s+", " ", question).strip()  # 移除所有空格并首尾去空格
+
+    # ==================== 2. 改写查询词：根据模型提示改写为适合知识库检索的独立查询 ====================
+    @staticmethod
+    def rewrite_query(question: str, model="qwen-plus"):
+        import os
+        from dotenv import load_dotenv
+        from langchain_openai import ChatOpenAI
+
+        load_dotenv()
+        llm = ChatOpenAI(
+            model=model,
+            api_key=os.getenv("DASHSCOPE_API_KEY"),
+            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
+            temperature=0,
+        )
+        response = llm.invoke([
+            ("system", "把问题改写为适合知识库检索的独立查询。保留原意，不编造缺失上下文，只输出查询。"),
+            ("human", question),
+        ])
+        return response.content.strip()
+
+
+    # ==================== 3. 生成假设性答案：根据模型提示生成约100字的假设性答案 ====================
+    @staticmethod
+    def generate_hyde(question: str, model="qwen-plus"):
+        import os
+        from dotenv import load_dotenv
+        from langchain_openai import ChatOpenAI
+
+        load_dotenv()
+        llm = ChatOpenAI(
+            model=model,
+            api_key=os.getenv("DASHSCOPE_API_KEY"),
+            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
+            temperature=0,
+        )
+        response = llm.invoke([
+            ("system", "为问题生成约100字的假设性答案，用于语义检索。只输出假设答案，不将其视为事实证据。"),
+            ("human", question),
+        ])
+        return response.content.strip()
+
+
+    # ==================== 4. 稠密：向量检索召回文档 ====================
+    @staticmethod
+    def vector_search(question: str, db, top_k=5, filters=None):
+        retriever = db.as_retriever(
+            search_type="similarity",
+            search_kwargs={"k": top_k, "filter": filters},
+        )
+        documents = retriever.invoke(question)
+        for document in documents:
+            document.metadata["chunk_id"] = document.id
+        return documents
+
+
+    # ==================== 5. 稀疏：BM25检索召回文档 ====================
+    @staticmethod
+    def bm25_search(question: str, documents, top_k=5):
+        import jieba
+        from langchain_community.retrievers import BM25Retriever
+
+        if not documents:
+            return []
+        bm25_retriever = BM25Retriever.from_documents(
+            documents,
+            preprocess_func=lambda text: list(jieba.cut(text)),
+            k=top_k,
+        )
+        return bm25_retriever.invoke(question)
+
+
+    # ==================== 6. RRF融合：向量检索和BM25检索融合 ====================
+    @staticmethod
+    def rrf_fusion(question: str, vector_retriever, bm25_retriever, weights=None, rrf_k=60):
+        from langchain_classic.retrievers import EnsembleRetriever
+
+        retriever = EnsembleRetriever(
+            retrievers=[vector_retriever, bm25_retriever],
+            weights=weights or [0.5, 0.5],
+            c=rrf_k,
+        )
+        return retriever.invoke(question)
+
+
+    # ==================== 7. 去重：去重 ====================
+    @staticmethod
+    def deduplicate(documents):
+        seen = set()
+        results = []
+        for document in documents:
+            if document.page_content not in seen:
+                seen.add(document.page_content)
+                results.append(document)
+        return results
+
+    # ==================== 8. 重排序：根据模型提示对文档进行重排序 ====================
+
+    @staticmethod
+    def rerank_documents(question: str, documents, top_k=5, model="qwen3-rerank"):
+        import os
+        from dotenv import load_dotenv
+        from langchain_community.document_compressors import DashScopeRerank
+
+        if not documents:
+            return []
+        load_dotenv()
+        reranker = DashScopeRerank(
+            model=model,
+            dashscope_api_key=os.getenv("DASHSCOPE_API_KEY"),
+            top_n=top_k,
+        )
+        return list(reranker.compress_documents(documents, question))
+
+
+    # ==================== 9. 相似度过滤：根据相似度阈值过滤文档 ====================
+    @staticmethod
+    def similarity_filter(question: str, documents, embeddings, threshold=0.4):
+        import math
+
+        if not documents:
+            return []
+        query_vector = embeddings.embed_query(question)
+        vectors = embeddings.embed_documents([
+            document.page_content for document in documents
+        ])
+        query_norm = math.sqrt(sum(value * value for value in query_vector))
+
+        results = []
+        for document, vector in zip(documents, vectors):
+            vector_norm = math.sqrt(sum(value * value for value in vector))
+            if query_norm == 0 or vector_norm == 0:
+                continue
+            similarity = sum(
+                query_value * value for query_value, value in zip(query_vector, vector)
+            ) / (query_norm * vector_norm)
+            if similarity >= threshold:
+                result = document.model_copy(deep=True)
+                result.metadata["cosine_similarity"] = similarity
+                results.append(result)
+        return results
+
+
+
+    # ==================== 10. 上下文压缩：根据模型提示对文档进行压缩 ====================
+    @staticmethod
+    def compress_context(question: str, documents, model="qwen-plus"):
+        import os
+        from dotenv import load_dotenv
+        from langchain_openai import ChatOpenAI
+        from langchain_classic.retrievers.document_compressors import LLMChainExtractor
+
+        if not documents:
+            return []
+        load_dotenv()
+        llm = ChatOpenAI(
+            model=model,
+            api_key=os.getenv("DASHSCOPE_API_KEY"),
+            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
+            temperature=0,
+        )
+        compressor = LLMChainExtractor.from_llm(llm)
+        results = []
+        for document in documents:
+            is_table = (
+                    document.metadata.get("content_type") == "table"
+                    or "<table" in document.page_content.lower()
+            )
+            if is_table:
+                results.append(document)
+            else:
+                results.extend(compressor.compress_documents([document], question))
+        return results
+
+    # ==================== 11. 上下文重排：根据模型提示对文档进行重排 ====================
+    @staticmethod
+    def reorder_documents(documents):
+        from langchain_community.document_transformers import LongContextReorder
+
+        return list(LongContextReorder().transform_documents(documents))
+
+    from langchain_core.documents import Document
+
+    documents = [
+        Document(id=str(index), page_content=f"相关性排名第{index}的内容")
+        for index in range(1, 6)
+    ]
+    for document in reorder_documents(documents):
+        print(document.page_content)
+
+
+