175 lines
5.6 KiB
Python
175 lines
5.6 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
文件格式支持综合测试 - 测试 txt、md、docx 文件
|
||
"""
|
||
import asyncio
|
||
import sys
|
||
import os
|
||
from pathlib import Path
|
||
|
||
# 添加项目路径
|
||
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=3)
|
||
|
||
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 comprehensive_file_format_test():
|
||
"""全面的文件格式测试"""
|
||
print("=" * 60)
|
||
print("🚀 Base RAG 异步文件格式支持测试")
|
||
print("=" * 60)
|
||
print()
|
||
|
||
# 初始化RAG系统
|
||
rag = SimpleRAG(
|
||
vector_store_name="comprehensive_test_kb",
|
||
retriever_top_k=3,
|
||
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 = [
|
||
{
|
||
"format": "TXT",
|
||
"path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/knowledge.txt",
|
||
"description": "纯文本文件 - 基础知识"
|
||
},
|
||
{
|
||
"format": "MD",
|
||
"path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/python_guide.md",
|
||
"description": "Markdown文件 - Python编程指南"
|
||
},
|
||
{
|
||
"format": "MD",
|
||
"path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/machine_learning.md",
|
||
"description": "Markdown文件 - 机器学习入门"
|
||
},
|
||
{
|
||
"format": "DOCX",
|
||
"path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/deep_learning_guide.docx",
|
||
"description": "Word文档 - 深度学习技术指南"
|
||
}
|
||
]
|
||
|
||
print("📂 文件处理测试\n")
|
||
|
||
# 处理每个文件
|
||
processed_files = []
|
||
for file_info in test_files:
|
||
file_path = file_info["path"]
|
||
format_name = file_info["format"]
|
||
description = file_info["description"]
|
||
|
||
if not Path(file_path).exists():
|
||
print(f"❌ {format_name}: {description} - 文件不存在")
|
||
continue
|
||
|
||
print(f"📄 处理 {format_name} 格式: {description}")
|
||
|
||
try:
|
||
result = await rag.process_file_to_vector_store(file_path)
|
||
if result and result.get('success'):
|
||
print(f" ✅ 成功: {result['chunks_count']} 个文档片段")
|
||
processed_files.append(file_info)
|
||
else:
|
||
print(f" ⚠️ 跳过: {result.get('message', '文件可能已存在')}")
|
||
processed_files.append(file_info)
|
||
except Exception as e:
|
||
print(f" ❌ 失败: {str(e)}")
|
||
|
||
print()
|
||
|
||
print("=" * 60)
|
||
print("💬 跨格式查询测试")
|
||
print("=" * 60)
|
||
print()
|
||
|
||
# 测试跨格式查询
|
||
test_queries = [
|
||
{
|
||
"question": "Python的基本语法有哪些?",
|
||
"expected_format": "MD (Python指南)"
|
||
},
|
||
{
|
||
"question": "什么是机器学习?有哪些类型?",
|
||
"expected_format": "MD (机器学习)"
|
||
},
|
||
{
|
||
"question": "深度学习有哪些核心组件?",
|
||
"expected_format": "DOCX (深度学习)"
|
||
},
|
||
{
|
||
"question": "深度学习的应用领域都有哪些?",
|
||
"expected_format": "DOCX (深度学习)"
|
||
}
|
||
]
|
||
|
||
for i, query_info in enumerate(test_queries, 1):
|
||
question = query_info["question"]
|
||
expected = query_info["expected_format"]
|
||
|
||
print(f"❓ 查询 {i}: {question}")
|
||
print(f" 期望来源: {expected}")
|
||
|
||
try:
|
||
response = await rag.query(question)
|
||
print(f" 💡 回答: {response[:200]}...")
|
||
print()
|
||
except Exception as e:
|
||
print(f" ❌ 查询失败: {str(e)}")
|
||
print()
|
||
|
||
print("=" * 60)
|
||
print("📊 测试总结")
|
||
print("=" * 60)
|
||
|
||
format_stats = {}
|
||
for file_info in processed_files:
|
||
format_name = file_info["format"]
|
||
if format_name not in format_stats:
|
||
format_stats[format_name] = 0
|
||
format_stats[format_name] += 1
|
||
|
||
print(f"✅ 成功处理的文件格式:")
|
||
for format_name, count in format_stats.items():
|
||
print(f" • {format_name}: {count} 个文件")
|
||
|
||
print(f"\n🎉 异步RAG系统文件格式支持测试完成!")
|
||
print(f" 支持格式: TXT, MD, DOCX")
|
||
print(f" 异步处理: ✅ 完全支持")
|
||
print(f" 跨格式查询: ✅ 完全支持")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(comprehensive_file_format_test())
|