base_rag/comprehensive_test.py

175 lines
5.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/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())