105 lines
3.0 KiB
Python
105 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
简单的文件处理测试
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import asyncio
|
||
import warnings
|
||
from pathlib import Path
|
||
|
||
# 过滤掉PyTorch的FutureWarning
|
||
warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
|
||
|
||
# 添加源码路径
|
||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))
|
||
|
||
from base_rag.core import BaseRAG, FileStatus
|
||
|
||
|
||
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_processing():
|
||
print("=== 文件处理功能测试 ===\n")
|
||
|
||
# 创建RAG实例
|
||
rag = SimpleRAG(
|
||
vector_store_name="test_kb",
|
||
retriever_top_k=2,
|
||
storage_directory="./test_files", # 统一使用test_files目录
|
||
status_db_path="./status.db", # 统一数据库名称
|
||
)
|
||
|
||
# 使用现有的测试文件
|
||
test_dir = Path("./test_files")
|
||
|
||
# 使用已有的测试文件
|
||
python_file = test_dir / "python_basics.txt"
|
||
web_file = test_dir / "web_frameworks.txt"
|
||
datascience_file = test_dir / "data_science.txt"
|
||
|
||
print("1. 处理多个知识文件...")
|
||
files_to_process = [python_file, web_file, datascience_file]
|
||
|
||
for file_path in files_to_process:
|
||
result = await rag.ingest(str(file_path), chunk_size=200, chunk_overlap=20)
|
||
print(
|
||
f"处理 {file_path.name}: {result['message']} (片段数: {result.get('chunks_count', 0)})"
|
||
)
|
||
print()
|
||
|
||
print("2. 查询测试...")
|
||
questions = [
|
||
"Python是谁创建的?",
|
||
"Flask和Django有什么区别?",
|
||
"Pandas是做什么的?",
|
||
"什么是NumPy?",
|
||
"FastAPI有什么特点?",
|
||
]
|
||
|
||
for question in questions:
|
||
print(f"问题: {question}")
|
||
answer = await rag.query(question)
|
||
print(f"回答: {answer[:150]}...")
|
||
print("-" * 50)
|
||
print()
|
||
|
||
print("3. 查看文件状态...")
|
||
files = await rag.get_file_processing_status()
|
||
for file_info in files:
|
||
print(f"文件: {file_info['filename']} | 状态: {file_info['status']}")
|
||
|
||
print("\n=== 测试完成 ===")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(test_file_processing())
|