diff --git a/examples/quick_start.py b/examples/quick_start.py index 3274c09..afceb0a 100644 --- a/examples/quick_start.py +++ b/examples/quick_start.py @@ -1,26 +1,45 @@ -"""Base RAG 使用示例""" +import sys +import os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) 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 ingest(self, documents): + for doc in documents: + self.vector_store.add_texts([doc]) + + def query(self, question, k=3): + docs = self.vector_store.similarity_search(question, k=k) + return docs - def query(self, question: str) -> str: - docs = self.similarity_search(question) - return f"找到 {len(docs)} 个相关文档" + +def main(): + config = { + "model_name": "sentence-transformers/all-MiniLM-L6-v2", + "embedding_type": "local" + } + + rag = SimpleRAG(embedding_config=config) + print("RAG初始化完成!") + + # 添加一些文档 + documents = [ + "苹果是一种水果,味道甜美,营养丰富。", + "苹果公司是一家科技公司,生产iPhone和Mac等产品。", + "Python是一种编程语言,简单易学,功能强大。" + ] + + print("正在添加文档...") + rag.ingest(documents) + print("文档添加完成!") + + # 测试查询 + print("\n正在查询: '什么是苹果?'") + result = rag.query("什么是苹果?") + print(f"查询结果: {result}") 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) + main() diff --git a/pyproject.toml b/pyproject.toml index f378939..4f8e5e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "langchain-community>=0.3.0", "langchain-openai>=0.2.0", "langchain-chroma>=0.1.0", + "langchain-huggingface>=0.1.0", "chromadb>=0.4.0", "openai>=1.0.0", "tiktoken>=0.5.0", diff --git a/requirements.txt b/requirements.txt index b9a5b64..ab43e84 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ langchain>=0.3.0 langchain-community>=0.3.0 langchain-openai>=0.2.0 langchain-chroma>=0.1.0 +langchain-huggingface>=0.1.0 chromadb>=0.4.0 openai>=1.0.0 tiktoken>=0.5.0 diff --git a/src/base_rag/core.py b/src/base_rag/core.py index 4c93477..86de225 100644 --- a/src/base_rag/core.py +++ b/src/base_rag/core.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from typing import List, Optional, Dict, ClassVar, Union import threading -from langchain_community.embeddings import HuggingFaceEmbeddings +from langchain_huggingface import HuggingFaceEmbeddings from langchain_openai import OpenAIEmbeddings from langchain.embeddings.base import Embeddings from langchain_chroma import Chroma