|
||
---|---|---|
examples | ||
scripts | ||
src/base_rag | ||
.gitignore | ||
LICENSE | ||
README.md | ||
RERANK_GUIDE.md | ||
pyproject.toml | ||
requirements.txt |
README.md
Base RAG
简洁的RAG基础库,支持多种embedding模型和Chroma向量数据库。
安装
pip install base-rag
使用
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("问题")