93 lines
2.6 KiB
Python
93 lines
2.6 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)
|
||
|
||
if not docs:
|
||
return "抱歉,没有找到相关信息。"
|
||
|
||
# 简单的回答拼接
|
||
context = "\n".join([doc.page_content for doc in docs])
|
||
return f"基于以下信息回答:\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)
|
||
|
||
# 创建一个知识文件
|
||
knowledge_file = test_dir / "knowledge.txt"
|
||
knowledge_file.write_text("""
|
||
Python是一种高级编程语言。
|
||
它具有简洁的语法和强大的功能。
|
||
Python广泛应用于Web开发、数据科学、人工智能等领域。
|
||
机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。
|
||
Flask和Django是流行的Python Web框架。
|
||
""", encoding="utf-8")
|
||
|
||
print("1. 处理知识文件...")
|
||
result = rag.ingest(str(knowledge_file))
|
||
print(f"处理结果: {result['message']}")
|
||
print(f"文档片段数: {result.get('chunks_count', 0)}")
|
||
print()
|
||
|
||
print("2. 查询测试...")
|
||
questions = [
|
||
"Python是什么?",
|
||
"Python有哪些应用领域?",
|
||
"有哪些Python Web框架?"
|
||
]
|
||
|
||
for question in questions:
|
||
print(f"问题: {question}")
|
||
answer = rag.query(question)
|
||
print(f"回答: {answer[:100]}...")
|
||
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()
|