127 lines
4.0 KiB
Python
127 lines
4.0 KiB
Python
#!/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()
|