base_rag/examples/simple_test.py

105 lines
3.0 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
"""
简单的文件处理测试
"""
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())