commit c10119b2dc5dbc96f5ed38ed82e497edd0f91bf1 Author: 李如威 Date: Mon Jul 28 10:44:56 2025 +0800 feat: init project diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..46fdecc --- /dev/null +++ b/.gitignore @@ -0,0 +1,17 @@ +# Python +__pycache__/ +*.py[cod] +*.so + +# Distribution +build/ +dist/ +*.egg-info/ + +# Environments +.venv +venv/ + +# Project specific +chroma_db/ +.DS_Store diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..7747c98 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Ruwei Li + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..d8b2f7b --- /dev/null +++ b/README.md @@ -0,0 +1,43 @@ +# Base RAG + +简洁的RAG基础库,支持多种embedding模型和Chroma向量数据库。 + +## 安装 + +```bash +pip install base-rag +``` + +## 使用 + +```python +from base_rag import BaseRAG + +class MyRAG(BaseRAG): + def ingest(self, file_path: str): + documents = self.load_and_split_documents(file_path) + self.add_documents_to_vector_store(documents) + + def query(self, question: str) -> str: + docs = self.similarity_search(question) + return f"找到 {len(docs)} 个相关文档" + +# OpenAI API +config = { + "type": "openai", + "model": "text-embedding-3-small", + "api_key": "your-api-key" +} + +# 本地模型 +config = { + "type": "local", + "model_name": "sentence-transformers/all-MiniLM-L6-v2" +} + +rag = MyRAG(embedding_config=config) +rag.ingest("document.txt") +result = rag.query("问题") +``` + +你只需要继承这个基类,实现 `ingest()` 和 `query()` 两个方法即可定制不同的 RAG 流程。如果你需要,我可以帮你写一个继承类样例。是否继续? diff --git a/examples/quick_start.py b/examples/quick_start.py new file mode 100644 index 0000000..2d704dc --- /dev/null +++ b/examples/quick_start.py @@ -0,0 +1,27 @@ +"""Base RAG 使用示例""" + +from base_rag import BaseRAG + +class SimpleRAG(BaseRAG): + def ingest(self, file_path: str): + documents = self.load_and_split_documents(file_path) + self.add_documents_to_vector_store(documents) + print(f"导入 {len(documents)} 个文档") + + def query(self, question: str) -> str: + docs = self.similarity_search(question) + return f"找到 {len(docs)} 个相关文档" + +if __name__ == "__main__": + # 本地模型配置 + config = { + "type": "local", + "model_name": "sentence-transformers/all-MiniLM-L6-v2" + } + + rag = SimpleRAG(embedding_config=config) + print("RAG初始化完成!") + + # rag.ingest("your_document.txt") + # result = rag.query("你的问题") + # print(result) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f378939 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,27 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "base-rag" +version = "0.1.0" +description = "简洁的RAG基础库" +readme = "README.md" +license = {text = "MIT"} +requires-python = ">=3.8" +dependencies = [ + "langchain>=0.3.0", + "langchain-community>=0.3.0", + "langchain-openai>=0.2.0", + "langchain-chroma>=0.1.0", + "chromadb>=0.4.0", + "openai>=1.0.0", + "tiktoken>=0.5.0", + "sentence-transformers>=2.2.0", +] + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-dir] +"" = "src" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b9a5b64 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +langchain>=0.3.0 +langchain-community>=0.3.0 +langchain-openai>=0.2.0 +langchain-chroma>=0.1.0 +chromadb>=0.4.0 +openai>=1.0.0 +tiktoken>=0.5.0 +sentence-transformers>=2.2.0 \ No newline at end of file diff --git a/scripts/build.sh b/scripts/build.sh new file mode 100755 index 0000000..cdbe652 --- /dev/null +++ b/scripts/build.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# 构建脚本 + +set -e + +echo "清理..." +rm -rf build/ dist/ *.egg-info/ + +echo "构建..." +python -m build + +echo "完成! 输出在 dist/ 目录" diff --git a/src/base_rag/__init__.py b/src/base_rag/__init__.py new file mode 100644 index 0000000..b1f1647 --- /dev/null +++ b/src/base_rag/__init__.py @@ -0,0 +1,6 @@ +"""简洁的RAG基础库""" + +from .core import BaseRAG + +__version__ = "0.1.0" +__all__ = ["BaseRAG"] diff --git a/src/base_rag/core.py b/src/base_rag/core.py new file mode 100644 index 0000000..4c93477 --- /dev/null +++ b/src/base_rag/core.py @@ -0,0 +1,206 @@ +from abc import ABC, abstractmethod +from typing import List, Optional, Dict, ClassVar, Union +import threading + +from langchain_community.embeddings import HuggingFaceEmbeddings +from langchain_openai import OpenAIEmbeddings +from langchain.embeddings.base import Embeddings +from langchain_chroma import Chroma +from langchain_community.document_loaders import TextLoader +from langchain.text_splitter import RecursiveCharacterTextSplitter +from langchain.chains import RetrievalQA +from langchain.llms.base import BaseLLM +from langchain.schema import Document + + +class BaseRAG(ABC): + # 类级别的模型缓存 + _embedding_models: ClassVar[Dict[str, Embeddings]] = {} + # 线程锁,保护模型缓存的并发访问 + _lock: ClassVar[threading.Lock] = threading.Lock() + + def __init__( + self, + vector_store_name: str = "default", + embedding_config: Optional[Dict] = None, + retriever_top_k: int = 3, + llm: Optional[BaseLLM] = None, + persist_directory: str = "./chroma_db", + ): + """ + 初始化基础RAG类。 + :param vector_store_name: 向量库名字(用于区分不同知识库) + :param embedding_config: 嵌入模型配置,支持本地和API模式 + :param retriever_top_k: 检索返回的文档数量 + :param llm: 可选的对话模型 + :param persist_directory: Chroma持久化目录 + + embedding_config 示例: + 本地模型名称: {"type": "local", "model_name": "sentence-transformers/all-MiniLM-L6-v2"} + 本地模型路径: {"type": "local", "model_path": "/path/to/your/model"} + OpenAI API: {"type": "openai", "model": "text-embedding-ada-002", "api_key": "sk-..."} + """ + self.vector_store_name = vector_store_name + self.embedding_config = embedding_config or { + "type": "local", + "model_name": "sentence-transformers/all-MiniLM-L6-v2", + } + self.retriever_top_k = retriever_top_k + self.llm = llm + self.persist_directory = persist_directory + + # 使用缓存的嵌入模型 + config_key = self._get_config_key(self.embedding_config) + self.embedding_model = self._get_or_create_embedding_model( + config_key, self.embedding_config + ) + + # 初始化 Chroma 向量库 + self.vector_store = Chroma( + collection_name=vector_store_name, + embedding_function=self.embedding_model, + persist_directory=persist_directory, + ) + + @staticmethod + def _get_config_key(config: Dict) -> str: + """ + 根据配置生成唯一的缓存键 + """ + config_type = config.get("type", "local") + if config_type == "local": + # 支持本地路径和模型名称两种方式 + if "model_path" in config: + return f"local_path_{config['model_path'].replace('/', '_').replace('\\', '_')}" + else: + return f"local_name_{config.get('model_name', 'default')}" + elif config_type == "openai": + return f"openai_{config.get('model', 'text-embedding-ada-002')}" + else: + return f"{config_type}_{config.get('model', 'default')}" + + @classmethod + def _get_or_create_embedding_model( + cls, config_key: str, config: Dict + ) -> Embeddings: + """ + 获取或创建嵌入模型(带缓存,线程安全) + """ + # 双重检查锁定模式,先检查是否已存在(避免不必要的锁开销) + if config_key in cls._embedding_models: + print(f"使用缓存的嵌入模型: {config_key}") + return cls._embedding_models[config_key] + + # 获取锁,进行安全的创建操作 + with cls._lock: + # 再次检查,防止在等待锁期间其他线程已经创建了模型 + if config_key not in cls._embedding_models: + print(f"正在创建嵌入模型: {config_key}") + cls._embedding_models[config_key] = cls._create_embedding_model(config) + else: + print(f"使用缓存的嵌入模型: {config_key}") + + return cls._embedding_models[config_key] + + @staticmethod + def _create_embedding_model(config: Dict) -> Embeddings: + """ + 根据配置创建嵌入模型 + """ + config_type = config.get("type", "local") + + if config_type == "local": + # 支持本地路径和模型名称两种方式 + if "model_path" in config: + model_path = config["model_path"] + print(f"从本地路径加载模型: {model_path}") + return HuggingFaceEmbeddings( + model_name=model_path, + model_kwargs=config.get("model_kwargs", {"device": "cpu"}), + encode_kwargs=config.get( + "encode_kwargs", {"normalize_embeddings": True} + ), + ) + else: + model_name = config.get( + "model_name", "sentence-transformers/all-MiniLM-L6-v2" + ) + print(f"从HuggingFace Hub加载模型: {model_name}") + return HuggingFaceEmbeddings( + model_name=model_name, + model_kwargs=config.get("model_kwargs", {"device": "cpu"}), + encode_kwargs=config.get( + "encode_kwargs", {"normalize_embeddings": True} + ), + ) + + elif config_type == "openai": + from langchain_openai import OpenAIEmbeddings + + return OpenAIEmbeddings( + model=config.get("model", "text-embedding-3-small"), + api_key=config.get("api_key"), + base_url=config.get("api_base"), + max_retries=config.get("max_retries", 3), + ) + + else: + raise ValueError( + f"不支持的嵌入模型类型: {config_type},支持的类型: 'local', 'openai'" + ) + + def load_and_split_documents(self, file_path: str) -> List[Document]: + """ + 加载并切分文档,可被子类重写实现不同的切分方式。 + """ + loader = TextLoader(file_path, encoding="utf-8") + documents = loader.load() + + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + return splitter.split_documents(documents) + + def add_documents_to_vector_store(self, documents: List[Document]): + """ + 将文档添加到 Chroma 向量库。 + """ + if documents: + self.vector_store.add_documents(documents) + self.vector_store.persist() # 持久化数据 + + def build_retriever(self): + """ + 构建检索器,可被子类或外部替换。 + """ + return self.vector_store.as_retriever(search_kwargs={"k": self.retriever_top_k}) + + def build_qa_chain(self): + """ + 构建 QA 链。 + """ + if not self.llm: + raise ValueError("LLM模型未设置") + retriever = self.build_retriever() + return RetrievalQA.from_chain_type( + llm=self.llm, retriever=retriever, return_source_documents=True + ) + + def similarity_search(self, query: str, k: int = None) -> List[Document]: + """ + 相似性搜索。 + """ + k = k or self.retriever_top_k + return self.vector_store.similarity_search(query, k=k) + + @abstractmethod + def ingest(self, *args, **kwargs): + """ + 子类需实现的文档导入逻辑。 + """ + pass + + @abstractmethod + def query(self, question: str) -> str: + """ + 子类需实现的问答逻辑。 + """ + pass