97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
测试 md、docx 文件格式的处理
|
||
"""
|
||
import asyncio
|
||
import sys
|
||
import os
|
||
|
||
# 添加项目路径
|
||
sys.path.append('/Users/liruwei/Documents/code/project/demo/base_rag/src')
|
||
|
||
from base_rag.core import BaseRAG
|
||
|
||
|
||
class SimpleRAG(BaseRAG):
|
||
"""简单的RAG实现示例"""
|
||
|
||
async def ingest(self, file_path: str, **kwargs):
|
||
"""实现文档导入逻辑"""
|
||
return await self.process_file_to_vector_store(file_path, **kwargs)
|
||
|
||
async def query(self, question: str) -> str:
|
||
"""实现简单的查询逻辑"""
|
||
docs = await self.similarity_search_with_rerank(question, k=2)
|
||
|
||
if not docs:
|
||
return "抱歉,没有找到相关信息。"
|
||
|
||
# 显示搜索到的文档来源
|
||
sources = []
|
||
contexts = []
|
||
for doc in docs:
|
||
source = doc.metadata.get("source_file", "未知来源")
|
||
if source not in sources:
|
||
sources.append(source)
|
||
contexts.append(doc.page_content.strip())
|
||
|
||
context = "\n\n".join(contexts)
|
||
sources_str = "、".join(sources)
|
||
|
||
return f"基于以下文档({sources_str})的信息:\n\n{context}"
|
||
|
||
async def test_file_formats():
|
||
"""测试不同文件格式的处理"""
|
||
print("=== 测试文件格式处理 ===\n")
|
||
|
||
# 初始化RAG系统
|
||
rag = SimpleRAG(
|
||
vector_store_name="test_formats_kb",
|
||
retriever_top_k=2,
|
||
persist_directory="/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db",
|
||
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/test_status.db"
|
||
)
|
||
|
||
# 测试文件列表
|
||
test_files = [
|
||
"/Users/liruwei/Documents/code/project/demo/base_rag/test_files/python_guide.md",
|
||
"/Users/liruwei/Documents/code/project/demo/base_rag/test_files/machine_learning.md",
|
||
"/Users/liruwei/Documents/code/project/demo/base_rag/test_files/deep_learning_guide.docx"
|
||
]
|
||
|
||
# 处理每个文件
|
||
for file_path in test_files:
|
||
print(f"处理文件: {os.path.basename(file_path)}")
|
||
|
||
try:
|
||
result = await rag.process_file_to_vector_store(file_path)
|
||
if result:
|
||
print(f"✅ 成功处理: {result}")
|
||
else:
|
||
print(f"❌ 处理失败或文件已存在")
|
||
except Exception as e:
|
||
print(f"❌ 处理出错: {str(e)}")
|
||
|
||
print()
|
||
|
||
# 测试查询
|
||
print("=== 测试查询功能 ===\n")
|
||
|
||
test_queries = [
|
||
"Python的基本语法有哪些?",
|
||
"什么是机器学习?",
|
||
"深度学习的核心概念是什么?",
|
||
"神经网络的主要架构有哪些?"
|
||
]
|
||
|
||
for query in test_queries:
|
||
print(f"查询: {query}")
|
||
try:
|
||
response = await rag.query(query)
|
||
print(f"回答: {response}\n")
|
||
except Exception as e:
|
||
print(f"❌ 查询出错: {str(e)}\n")
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(test_file_formats())
|