base_rag/test_docx.py

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())