base_rag/examples/simple_test.py

89 lines
2.5 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
from pathlib import Path
# 添加源码路径
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()