base_rag/test_file_formats.py

97 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
"""
测试 md、docx 文件格式的处理
"""
import asyncio
import sys
import os
# 添加项目路径
sys.path.append('/Users/liruwei/Documents/code/project/demo/base_rag/src')
from base_rag.core import BaseRAG
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_formats():
"""测试不同文件格式的处理"""
print("=== 测试文件格式处理 ===\n")
# 初始化RAG系统
rag = SimpleRAG(
vector_store_name="test_formats_kb",
retriever_top_k=2,
persist_directory="/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db",
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/test_status.db"
)
# 测试文件列表
test_files = [
"/Users/liruwei/Documents/code/project/demo/base_rag/test_files/python_guide.md",
"/Users/liruwei/Documents/code/project/demo/base_rag/test_files/machine_learning.md",
"/Users/liruwei/Documents/code/project/demo/base_rag/test_files/deep_learning_guide.docx"
]
# 处理每个文件
for file_path in test_files:
print(f"处理文件: {os.path.basename(file_path)}")
try:
result = await rag.process_file_to_vector_store(file_path)
if result:
print(f"✅ 成功处理: {result}")
else:
print(f"❌ 处理失败或文件已存在")
except Exception as e:
print(f"❌ 处理出错: {str(e)}")
print()
# 测试查询
print("=== 测试查询功能 ===\n")
test_queries = [
"Python的基本语法有哪些",
"什么是机器学习?",
"深度学习的核心概念是什么?",
"神经网络的主要架构有哪些?"
]
for query in test_queries:
print(f"查询: {query}")
try:
response = await rag.query(query)
print(f"回答: {response}\n")
except Exception as e:
print(f"❌ 查询出错: {str(e)}\n")
if __name__ == "__main__":
asyncio.run(test_file_formats())