本地RAG尝试
This commit is contained in:
232
app/rag/example.py
Normal file
232
app/rag/example.py
Normal file
@@ -0,0 +1,232 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
RAG 系统使用示例
|
||||
|
||||
演示如何使用 app/rag 模块进行知识检索。
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# 添加项目根目录到路径
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
from langchain_community.llms import VLLMOpenAI
|
||||
|
||||
|
||||
def setup_environment():
|
||||
"""设置环境变量"""
|
||||
# 设置 Qdrant 连接信息（根据实际情况修改）
|
||||
os.environ.setdefault("QDRANT_URL", "http://115.190.121.151:6333")
|
||||
# 如果需要 API 密钥，请设置 QDRANT_API_KEY
|
||||
|
||||
print("环境变量已设置")
|
||||
print(f"QDRANT_URL: {os.environ.get('QDRANT_URL')}")
|
||||
|
||||
|
||||
def demonstrate_basic_rag():
|
||||
"""演示基础 RAG 功能"""
|
||||
print("\n" + "="*60)
|
||||
print("演示: 基础 RAG 检索 (Level 1)")
|
||||
print("="*60)
|
||||
|
||||
# 创建嵌入模型（使用 OpenAI 兼容的本地模型）
|
||||
embeddings = OpenAIEmbeddings(
|
||||
openai_api_base="http://localhost:8000/v1", # 本地 VLLM 服务
|
||||
openai_api_key="no-key-needed",
|
||||
model="text-embedding-ada-002", # 假设的模型名称
|
||||
)
|
||||
|
||||
# 创建 RAG 流水线
|
||||
from app.rag import RAGPipeline, RAGConfig, RAGLevel
|
||||
|
||||
config = RAGConfig(
|
||||
collection_name="documents", # 你的集合名称
|
||||
rag_level=RAGLevel.BASIC,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
pipeline = RAGPipeline(
|
||||
embeddings=embeddings,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# 示例查询
|
||||
query = "公司报销流程是什么？"
|
||||
print(f"\n查询: {query}")
|
||||
|
||||
try:
|
||||
result = pipeline.retrieve(query)
|
||||
print(f"找到 {len(result.documents)} 个相关文档")
|
||||
|
||||
# 格式化上下文
|
||||
context = pipeline.format_context(result.documents)
|
||||
print(f"\n上下文预览:\n{context[:500]}...")
|
||||
|
||||
except Exception as e:
|
||||
print(f"检索失败: {e}")
|
||||
print("请确保 Qdrant 服务正常运行且集合存在")
|
||||
|
||||
|
||||
def demonstrate_hybrid_rag():
|
||||
"""演示混合 RAG 功能"""
|
||||
print("\n" + "="*60)
|
||||
print("演示: 混合 RAG 检索 (Level 2)")
|
||||
print("="*60)
|
||||
|
||||
embeddings = OpenAIEmbeddings(
|
||||
openai_api_base="http://localhost:8000/v1",
|
||||
openai_api_key="no-key-needed",
|
||||
model="text-embedding-ada-002",
|
||||
)
|
||||
|
||||
from app.rag import RAGPipeline, RAGConfig, RAGLevel
|
||||
|
||||
config = RAGConfig(
|
||||
collection_name="documents",
|
||||
rag_level=RAGLevel.HYBRID,
|
||||
dense_k=10,
|
||||
sparse_k=10,
|
||||
rerank_top_n=5,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
pipeline = RAGPipeline(
|
||||
embeddings=embeddings,
|
||||
config=config,
|
||||
)
|
||||
|
||||
query = "如何申请年假？"
|
||||
print(f"\n查询: {query}")
|
||||
|
||||
try:
|
||||
result = pipeline.retrieve(query)
|
||||
print(f"找到 {len(result.documents)} 个重排序后的文档")
|
||||
|
||||
except Exception as e:
|
||||
print(f"检索失败: {e}")
|
||||
|
||||
|
||||
def demonstrate_rag_fusion():
|
||||
"""演示 RAG-Fusion 功能"""
|
||||
print("\n" + "="*60)
|
||||
print("演示: RAG-Fusion (Level 3)")
|
||||
print("="*60)
|
||||
|
||||
embeddings = OpenAIEmbeddings(
|
||||
openai_api_base="http://localhost:8000/v1",
|
||||
openai_api_key="no-key-needed",
|
||||
model="text-embedding-ada-002",
|
||||
)
|
||||
|
||||
# 创建语言模型用于查询改写
|
||||
llm = VLLMOpenAI(
|
||||
openai_api_base="http://localhost:8000/v1",
|
||||
openai_api_key="no-key-needed",
|
||||
model_name="Qwen2.5-7B-Instruct", # 你的本地模型
|
||||
temperature=0.3,
|
||||
max_tokens=512,
|
||||
)
|
||||
|
||||
from app.rag import RAGPipeline, RAGConfig, RAGLevel
|
||||
|
||||
config = RAGConfig(
|
||||
collection_name="documents",
|
||||
rag_level=RAGLevel.FUSION,
|
||||
num_queries=3,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
pipeline = RAGPipeline(
|
||||
embeddings=embeddings,
|
||||
llm=llm,
|
||||
config=config,
|
||||
)
|
||||
|
||||
query = "项目上线需要哪些审批？"
|
||||
print(f"\n查询: {query}")
|
||||
|
||||
try:
|
||||
result = pipeline.retrieve(query)
|
||||
print(f"找到 {len(result.documents)} 个文档 (经过多路查询改写和重排序)")
|
||||
|
||||
except Exception as e:
|
||||
print(f"检索失败: {e}")
|
||||
|
||||
|
||||
def demonstrate_agentic_rag():
|
||||
"""演示 Agentic RAG 功能"""
|
||||
print("\n" + "="*60)
|
||||
print("演示: Agentic RAG (Level 4)")
|
||||
print("="*60)
|
||||
|
||||
embeddings = OpenAIEmbeddings(
|
||||
openai_api_base="http://localhost:8000/v1",
|
||||
openai_api_key="no-key-needed",
|
||||
model="text-embedding-ada-002",
|
||||
)
|
||||
|
||||
llm = VLLMOpenAI(
|
||||
openai_api_base="http://localhost:8000/v1",
|
||||
openai_api_key="no-key-needed",
|
||||
model_name="Qwen2.5-7B-Instruct",
|
||||
temperature=0.3,
|
||||
max_tokens=512,
|
||||
)
|
||||
|
||||
from app.rag import create_agentic_rag_pipeline
|
||||
|
||||
try:
|
||||
# 创建 Agentic RAG 流水线
|
||||
agentic_rag = create_agentic_rag_pipeline(
|
||||
embeddings=embeddings,
|
||||
agent_llm=llm,
|
||||
config={
|
||||
"collection_name": "documents",
|
||||
"verbose": True,
|
||||
},
|
||||
)
|
||||
|
||||
print("Agentic RAG 流水线创建成功！")
|
||||
print(f"- 绑定的模型: {agentic_rag['llm']}")
|
||||
print(f"- RAG 工具: {agentic_rag['tool'].name}")
|
||||
|
||||
# 演示工具调用
|
||||
print("\n工具调用示例:")
|
||||
response = agentic_rag["tool"].invoke({"query": "员工福利有哪些？"})
|
||||
print(f"工具响应预览: {response[:200]}...")
|
||||
|
||||
except Exception as e:
|
||||
print(f"创建 Agentic RAG 失败: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
print("RAG 系统演示")
|
||||
print("="*60)
|
||||
|
||||
# 设置环境
|
||||
setup_environment()
|
||||
|
||||
# 演示各级功能
|
||||
demonstrate_basic_rag()
|
||||
demonstrate_hybrid_rag()
|
||||
demonstrate_rag_fusion()
|
||||
demonstrate_agentic_rag()
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("演示完成!")
|
||||
print("="*60)
|
||||
|
||||
print("\n使用说明:")
|
||||
print("1. 确保 Qdrant 服务运行且集合已创建")
|
||||
print("2. 根据需要修改 embeddings 和 llm 配置")
|
||||
print("3. 在 Agent 系统中导入并使用 app.rag.tools.search_knowledge_base_tool")
|
||||
print("4. 将工具绑定到你的 Agent 模型")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user