feat: 测试内容
This commit is contained in:
parent
d7ad7a8342
commit
af8b7f65fb
|
@ -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?
|
||||
RAG(Retrieval-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()
|
|
@ -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()
|
|
@ -26,19 +26,29 @@ class SimpleRAG(BaseRAG):
|
|||
|
||||
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:
|
||||
return "抱歉,没有找到相关信息。"
|
||||
|
||||
# 简单的回答拼接
|
||||
context = "\n".join([doc.page_content for doc in docs])
|
||||
return f"基于以下信息回答:\n{context}"
|
||||
# 显示搜索到的文档来源
|
||||
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",
|
||||
|
@ -46,45 +56,69 @@ def test_file_processing():
|
|||
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框架。
|
||||
# Python基础知识
|
||||
python_file = test_dir / "python_basics.txt"
|
||||
python_file.write_text("""
|
||||
Python是一种高级编程语言,由Guido van Rossum于1991年创建。
|
||||
Python具有简洁易读的语法,适合初学者学习编程。
|
||||
Python是解释型语言,支持面向对象、函数式等多种编程范式。
|
||||
Python的设计哲学强调代码的可读性和简洁性。
|
||||
""", encoding="utf-8")
|
||||
|
||||
print("1. 处理知识文件...")
|
||||
result = rag.ingest(str(knowledge_file))
|
||||
print(f"处理结果: {result['message']}")
|
||||
print(f"文档片段数: {result.get('chunks_count', 0)}")
|
||||
print()
|
||||
# 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是什么?",
|
||||
"Python有哪些应用领域?",
|
||||
"有哪些Python Web框架?"
|
||||
"Python是谁创建的?",
|
||||
"Flask和Django有什么区别?",
|
||||
"Pandas是做什么的?",
|
||||
"什么是NumPy?",
|
||||
"FastAPI有什么特点?"
|
||||
]
|
||||
|
||||
|
||||
for question in questions:
|
||||
print(f"问题: {question}")
|
||||
answer = rag.query(question)
|
||||
print(f"回答: {answer[:100]}...")
|
||||
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=== 测试完成 ===")
|
||||
|
||||
|
||||
|
|
|
@ -0,0 +1,5 @@
|
|||
|
||||
NumPy是Python中用于科学计算的基础库,提供多维数组对象。
|
||||
Pandas是强大的数据分析和处理库,提供DataFrame数据结构。
|
||||
Matplotlib是Python的绘图库,用于创建静态、动态和交互式图表。
|
||||
Scikit-learn是机器学习库,提供各种算法和工具。
|
|
@ -0,0 +1,5 @@
|
|||
|
||||
Python是一种高级编程语言,由Guido van Rossum于1991年创建。
|
||||
Python具有简洁易读的语法,适合初学者学习编程。
|
||||
Python是解释型语言,支持面向对象、函数式等多种编程范式。
|
||||
Python的设计哲学强调代码的可读性和简洁性。
|
|
@ -0,0 +1,5 @@
|
|||
|
||||
Flask是一个轻量级的Python Web框架,易于学习和使用。
|
||||
Django是一个功能丰富的Python Web框架,适合大型项目开发。
|
||||
FastAPI是现代的Python Web框架,专为构建API而设计。
|
||||
Tornado是一个可扩展的非阻塞Web服务器和Web应用框架。
|
|
@ -0,0 +1,5 @@
|
|||
|
||||
NumPy是Python中用于科学计算的基础库,提供多维数组对象。
|
||||
Pandas是强大的数据分析和处理库,提供DataFrame数据结构。
|
||||
Matplotlib是Python的绘图库,用于创建静态、动态和交互式图表。
|
||||
Scikit-learn是机器学习库,提供各种算法和工具。
|
|
@ -0,0 +1,5 @@
|
|||
|
||||
Python是一种高级编程语言,由Guido van Rossum于1991年创建。
|
||||
Python具有简洁易读的语法,适合初学者学习编程。
|
||||
Python是解释型语言,支持面向对象、函数式等多种编程范式。
|
||||
Python的设计哲学强调代码的可读性和简洁性。
|
|
@ -0,0 +1,5 @@
|
|||
|
||||
Flask是一个轻量级的Python Web框架,易于学习和使用。
|
||||
Django是一个功能丰富的Python Web框架,适合大型项目开发。
|
||||
FastAPI是现代的Python Web框架,专为构建API而设计。
|
||||
Tornado是一个可扩展的非阻塞Web服务器和Web应用框架。
|
BIN
test_status.db
BIN
test_status.db
Binary file not shown.
Loading…
Reference in New Issue