base_rag/examples/simple_test.py

127 lines
4.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 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实现示例"""
def ingest(self, file_path: str, **kwargs):
"""实现文档导入逻辑"""
return self.process_file_to_vector_store(file_path, **kwargs)
def query(self, question: str) -> str:
"""实现简单的查询逻辑"""
docs = 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}"
def test_file_processing():
print("=== 文件处理功能测试 ===\n")
# 创建RAG实例
rag = SimpleRAG(
vector_store_name="test_kb",
retriever_top_k=2,
storage_directory="./test_docs",
status_db_path="./test_status.db"
)
# 创建测试文件
test_dir = Path("./test_files")
test_dir.mkdir(exist_ok=True)
# 创建多个不同主题的知识文件
# Python基础知识
python_file = test_dir / "python_basics.txt"
python_file.write_text("""
Python是一种高级编程语言由Guido van Rossum于1991年创建。
Python具有简洁易读的语法适合初学者学习编程。
Python是解释型语言支持面向对象、函数式等多种编程范式。
Python的设计哲学强调代码的可读性和简洁性。
""", encoding="utf-8")
# Web框架知识
web_file = test_dir / "web_frameworks.txt"
web_file.write_text("""
Flask是一个轻量级的Python Web框架易于学习和使用。
Django是一个功能丰富的Python Web框架适合大型项目开发。
FastAPI是现代的Python Web框架专为构建API而设计。
Tornado是一个可扩展的非阻塞Web服务器和Web应用框架。
""", encoding="utf-8")
# 数据科学知识
datascience_file = test_dir / "data_science.txt"
datascience_file.write_text("""
NumPy是Python中用于科学计算的基础库提供多维数组对象。
Pandas是强大的数据分析和处理库提供DataFrame数据结构。
Matplotlib是Python的绘图库用于创建静态、动态和交互式图表。
Scikit-learn是机器学习库提供各种算法和工具。
""", encoding="utf-8")
print("1. 处理多个知识文件...")
files_to_process = [python_file, web_file, datascience_file]
for file_path in files_to_process:
result = 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 = rag.query(question)
print(f"回答: {answer[:150]}...")
print("-" * 50)
print()
print("3. 查看文件状态...")
files = rag.get_file_processing_status()
for file_info in files:
print(f"文件: {file_info['filename']} | 状态: {file_info['status']}")
print("\n=== 测试完成 ===")
if __name__ == "__main__":
test_file_processing()