92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
测试 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_docx_format():
|
|
"""测试docx文件格式处理"""
|
|
print("=== 测试DOCX文件格式处理 ===\n")
|
|
|
|
# 初始化RAG系统
|
|
rag = SimpleRAG(
|
|
vector_store_name="test_docx_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"
|
|
)
|
|
|
|
# 测试docx文件
|
|
docx_file = "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/deep_learning_guide.docx"
|
|
|
|
print(f"处理文件: {os.path.basename(docx_file)}")
|
|
|
|
try:
|
|
result = await rag.process_file_to_vector_store(docx_file)
|
|
if result:
|
|
print(f"✅ 成功处理: {result}")
|
|
else:
|
|
print(f"❌ 处理失败或文件已存在")
|
|
except Exception as e:
|
|
print(f"❌ 处理出错: {str(e)}")
|
|
|
|
print()
|
|
|
|
# 测试查询
|
|
print("=== 测试针对DOCX文档的查询 ===\n")
|
|
|
|
test_queries = [
|
|
"深度学习是什么?",
|
|
"卷积神经网络有什么特点?",
|
|
"深度学习有哪些应用领域?",
|
|
"深度学习面临哪些技术挑战?"
|
|
]
|
|
|
|
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_docx_format())
|