44 lines
1.0 KiB
Markdown
44 lines
1.0 KiB
Markdown
# 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 流程。如果你需要,我可以帮你写一个继承类样例。是否继续?
|