feat: 测试内容

This commit is contained in:
李如威 2025-08-07 17:10:49 +08:00
parent d7ad7a8342
commit af8b7f65fb
10 changed files with 90 additions and 234 deletions

View File

@ -1,126 +0,0 @@
#!/usr/bin/env python3
"""
文件处理示例
演示如何使用BaseRAG的文件处理功能
"""
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 "抱歉,没有找到相关信息。"
# 简单的回答拼接实际应用中应该使用LLM
context = "\n".join([doc.page_content for doc in docs])
return f"基于以下信息回答:\n{context}"
def main():
# 创建RAG实例
rag = SimpleRAG(
vector_store_name="file_demo",
retriever_top_k=3,
storage_directory="./demo_documents",
status_db_path="./demo_file_status.db"
)
# 创建测试文件
test_dir = Path("./test_files")
test_dir.mkdir(exist_ok=True)
# 创建测试文本文件
txt_file = test_dir / "test_document.txt"
txt_file.write_text("""
这是一个测试文档
它包含了关于人工智能的信息
人工智能是计算机科学的一个分支致力于创建智能机器
机器学习是人工智能的一个重要组成部分
深度学习是机器学习的一个子领域
""", encoding="utf-8")
# 创建测试Markdown文件
md_file = test_dir / "test_markdown.md"
md_file.write_text("""
# RAG系统介绍
## 什么是RAG
RAGRetrieval-Augmented Generation是一种结合了检索和生成的AI技术
## RAG的优势
- 能够利用外部知识库
- 提高回答的准确性
- 支持实时更新知识
## 应用场景
RAG系统广泛应用于问答系统知识管理等领域
""", encoding="utf-8")
print("=== 文件处理示例 ===\n")
# 1. 处理文本文件
print("1. 处理文本文件...")
result1 = rag.ingest(str(txt_file))
print(f"处理结果: {result1}\n")
# 2. 处理Markdown文件
print("2. 处理Markdown文件...")
result2 = rag.ingest(str(md_file))
print(f"处理结果: {result2}\n")
# 3. 再次处理同一个文件(应该跳过)
print("3. 再次处理文本文件(测试重复处理)...")
result3 = rag.ingest(str(txt_file))
print(f"处理结果: {result3}\n")
# 4. 查看所有文件状态
print("4. 查看所有文件处理状态...")
all_files = rag.get_file_processing_status()
for file_info in all_files:
print(f"文件: {file_info['filename']}")
print(f"类型: {file_info['file_type']}")
print(f"状态: {file_info['status']}")
print(f"处理时间: {file_info['updated_at']}")
print("---")
# 5. 查看已完成的文件
print("\n5. 查看已完成处理的文件...")
completed_files = rag.list_files_by_status(FileStatus.COMPLETED)
print(f"已完成处理的文件数量: {len(completed_files)}")
# 6. 测试搜索功能
print("\n6. 测试搜索功能...")
questions = [
"什么是人工智能?",
"RAG有什么优势",
"机器学习是什么?"
]
for question in questions:
print(f"\n问题: {question}")
answer = rag.query(question)
print(f"回答: {answer[:200]}...") # 只显示前200个字符
print("\n=== 示例完成 ===")
if __name__ == "__main__":
main()

View File

@ -1,82 +0,0 @@
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from base_rag import BaseRAG
class SimpleRAG(BaseRAG):
def ingest(self, documents):
for doc in documents:
self.vector_store.add_texts([doc])
def query(self, question, k=3):
# 使用带重排功能的搜索
print(f"查询参数 k={k}")
docs = self.similarity_search_with_rerank(question, k=k)
print(f"重排后返回文档数量: {len(docs)}")
return docs
def query_without_rerank(self, question, k=3):
# 不使用重排的普通搜索
docs = self.similarity_search(question, k=k)
return docs
def main():
# 重排配置 - 使用BGE重排模型
rerank_config = {
"enabled": True,
"type": "local",
"model": "BAAI/bge-reranker-base",
# 注意这里不设置top_k让它在查询时动态决定
}
rag = SimpleRAG(rerank_config=rerank_config)
print("RAG系统含重排功能初始化完成!")
# 添加更多测试文档
# documents = [
# "苹果是一种红色或绿色的水果味道甜美营养丰富含有丰富的维生素C。",
# "苹果公司是一家总部位于美国的科技公司以生产iPhone、iPad、Mac等产品而闻名。",
# "Python是一种高级编程语言简单易学功能强大广泛用于数据科学和机器学习。",
# "苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。",
# "苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。",
# "苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。",
# "iPhone是苹果公司生产的智能手机具有先进的技术和优秀的用户体验。",
# "机器学习是人工智能的一个分支Python是机器学习领域最流行的编程语言之一。",
# "苹果公司主要产品是iPhone、iPad、Mac等",
# "使用多线程可以提高程序在多核处理器上的运行效率。",
# "数据库设计应该遵循第三范式,以减少数据冗余。",
# "使用更高效的数据结构(如哈希表)可以提升算法效率。",
# "程序界面应该简洁明了,方便用户操作。",
# "尽量避免使用全局变量,减少潜在的并发问题。",
# "在Python中使用NumPy等库可以加快数值计算速度。",
# "使用SSD替代传统HDD可提升文件读写速度。",
# "将代码中的重复逻辑封装为函数,有助于维护但对性能影响小。",
# "内存泄漏会导致程序运行越来越慢,应及时释放资源。",
# "使用代码缓存如Memoization可以避免重复计算。",
# ]
# print("正在添加文档...")
# rag.ingest(documents)
# print(f"文档添加完成! 共添加了 {len(documents)} 个文档")
# 测试查询并比较重排效果
test_query = "Python中的list操作较慢尽量使用数组库"
print(f"\n测试查询: '{test_query}'")
print("\n=== 不使用重排的结果 ===")
result_no_rerank = rag.query_without_rerank(test_query, k=3)
for i, doc in enumerate(result_no_rerank, 1):
print(f"{i}. {doc.page_content}")
print("\n=== 使用重排的结果 ===")
result_with_rerank = rag.query(test_query, k=3)
for i, doc in enumerate(result_with_rerank, 1):
print(f"{i}. {doc.page_content}")
if __name__ == "__main__":
main()

View File

@ -26,19 +26,29 @@ class SimpleRAG(BaseRAG):
def query(self, question: str) -> str: def query(self, question: str) -> str:
"""实现简单的查询逻辑""" """实现简单的查询逻辑"""
docs = self.similarity_search_with_rerank(question) docs = self.similarity_search_with_rerank(question, k=2)
if not docs: if not docs:
return "抱歉,没有找到相关信息。" return "抱歉,没有找到相关信息。"
# 简单的回答拼接 # 显示搜索到的文档来源
context = "\n".join([doc.page_content for doc in docs]) sources = []
return f"基于以下信息回答:\n{context}" 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(): def test_file_processing():
print("=== 文件处理功能测试 ===\n") print("=== 文件处理功能测试 ===\n")
# 创建RAG实例 # 创建RAG实例
rag = SimpleRAG( rag = SimpleRAG(
vector_store_name="test_kb", vector_store_name="test_kb",
@ -46,45 +56,69 @@ def test_file_processing():
storage_directory="./test_docs", storage_directory="./test_docs",
status_db_path="./test_status.db" status_db_path="./test_status.db"
) )
# 创建测试文件 # 创建测试文件
test_dir = Path("./test_files") test_dir = Path("./test_files")
test_dir.mkdir(exist_ok=True) test_dir.mkdir(exist_ok=True)
# 创建多个不同主题的知识文件
# 创建一个知识文件 # Python基础知识
knowledge_file = test_dir / "knowledge.txt" python_file = test_dir / "python_basics.txt"
knowledge_file.write_text(""" python_file.write_text("""
Python是一种高级编程语言 Python是一种高级编程语言由Guido van Rossum于1991年创建
它具有简洁的语法和强大的功能 Python具有简洁易读的语法适合初学者学习编程
Python广泛应用于Web开发数据科学人工智能等领域 Python是解释型语言支持面向对象函数式等多种编程范式
机器学习库如scikit-learnTensorFlow和PyTorch都支持Python Python的设计哲学强调代码的可读性和简洁性
Flask和Django是流行的Python Web框架
""", encoding="utf-8") """, encoding="utf-8")
print("1. 处理知识文件...") # Web框架知识
result = rag.ingest(str(knowledge_file)) web_file = test_dir / "web_frameworks.txt"
print(f"处理结果: {result['message']}") web_file.write_text("""
print(f"文档片段数: {result.get('chunks_count', 0)}") Flask是一个轻量级的Python Web框架易于学习和使用
print() 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. 查询测试...") print("2. 查询测试...")
questions = [ questions = [
"Python是什么", "Python是谁创建的",
"Python有哪些应用领域", "Flask和Django有什么区别",
"有哪些Python Web框架" "Pandas是做什么的",
"什么是NumPy",
"FastAPI有什么特点"
] ]
for question in questions: for question in questions:
print(f"问题: {question}") print(f"问题: {question}")
answer = rag.query(question) answer = rag.query(question)
print(f"回答: {answer[:100]}...") print(f"回答: {answer[:150]}...")
print("-" * 50)
print() print()
print("3. 查看文件状态...") print("3. 查看文件状态...")
files = rag.get_file_processing_status() files = rag.get_file_processing_status()
for file_info in files: for file_info in files:
print(f"文件: {file_info['filename']} | 状态: {file_info['status']}") print(f"文件: {file_info['filename']} | 状态: {file_info['status']}")
print("\n=== 测试完成 ===") print("\n=== 测试完成 ===")

View File

@ -0,0 +1,5 @@
NumPy是Python中用于科学计算的基础库提供多维数组对象。
Pandas是强大的数据分析和处理库提供DataFrame数据结构。
Matplotlib是Python的绘图库用于创建静态、动态和交互式图表。
Scikit-learn是机器学习库提供各种算法和工具。

View File

@ -0,0 +1,5 @@
Python是一种高级编程语言由Guido van Rossum于1991年创建。
Python具有简洁易读的语法适合初学者学习编程。
Python是解释型语言支持面向对象、函数式等多种编程范式。
Python的设计哲学强调代码的可读性和简洁性。

View File

@ -0,0 +1,5 @@
Flask是一个轻量级的Python Web框架易于学习和使用。
Django是一个功能丰富的Python Web框架适合大型项目开发。
FastAPI是现代的Python Web框架专为构建API而设计。
Tornado是一个可扩展的非阻塞Web服务器和Web应用框架。

View File

@ -0,0 +1,5 @@
NumPy是Python中用于科学计算的基础库提供多维数组对象。
Pandas是强大的数据分析和处理库提供DataFrame数据结构。
Matplotlib是Python的绘图库用于创建静态、动态和交互式图表。
Scikit-learn是机器学习库提供各种算法和工具。

View File

@ -0,0 +1,5 @@
Python是一种高级编程语言由Guido van Rossum于1991年创建。
Python具有简洁易读的语法适合初学者学习编程。
Python是解释型语言支持面向对象、函数式等多种编程范式。
Python的设计哲学强调代码的可读性和简洁性。

View File

@ -0,0 +1,5 @@
Flask是一个轻量级的Python Web框架易于学习和使用。
Django是一个功能丰富的Python Web框架适合大型项目开发。
FastAPI是现代的Python Web框架专为构建API而设计。
Tornado是一个可扩展的非阻塞Web服务器和Web应用框架。

Binary file not shown.