base_rag/README.md

837 B
Raw Blame History

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("问题")