feat: 支持 csv、pdf、xlsx等格式
This commit is contained in:
parent
97498e2fd7
commit
c6e020a170
Binary file not shown.
|
@ -0,0 +1,321 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
创建复杂格式的测试文件 - DOCX、PDF、Excel、CSV
|
||||||
|
"""
|
||||||
|
import asyncio
|
||||||
|
from docx import Document
|
||||||
|
from docx.shared import Inches
|
||||||
|
import pandas as pd
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
from reportlab.lib.pagesizes import letter
|
||||||
|
from reportlab.pdfgen import canvas
|
||||||
|
from reportlab.lib.units import inch
|
||||||
|
import io
|
||||||
|
import os
|
||||||
|
|
||||||
|
async def create_complex_docx():
|
||||||
|
"""创建包含表格和图片的复杂 DOCX 文件"""
|
||||||
|
doc = Document()
|
||||||
|
|
||||||
|
# 添加标题
|
||||||
|
doc.add_heading('数据科学技术报告', 0)
|
||||||
|
|
||||||
|
# 添加段落
|
||||||
|
doc.add_paragraph('本报告全面介绍了数据科学领域的核心技术和应用场景。')
|
||||||
|
|
||||||
|
# 添加子标题
|
||||||
|
doc.add_heading('1. 数据科学概述', level=1)
|
||||||
|
|
||||||
|
# 添加详细内容
|
||||||
|
doc.add_paragraph(
|
||||||
|
'数据科学是一个跨学科领域,结合了统计学、计算机科学和领域专业知识。'
|
||||||
|
'它使用科学方法、流程、算法和系统从结构化和非结构化数据中提取知识和洞察。'
|
||||||
|
)
|
||||||
|
|
||||||
|
# 添加表格
|
||||||
|
doc.add_heading('2. 核心技术对比', level=1)
|
||||||
|
table = doc.add_table(rows=1, cols=4)
|
||||||
|
table.style = 'Table Grid'
|
||||||
|
|
||||||
|
# 表头
|
||||||
|
hdr_cells = table.rows[0].cells
|
||||||
|
hdr_cells[0].text = '技术'
|
||||||
|
hdr_cells[1].text = '应用场景'
|
||||||
|
hdr_cells[2].text = '优势'
|
||||||
|
hdr_cells[3].text = '难度等级'
|
||||||
|
|
||||||
|
# 数据行
|
||||||
|
technologies = [
|
||||||
|
('机器学习', '预测分析、分类', '自动化决策', '中等'),
|
||||||
|
('深度学习', '图像识别、NLP', '高精度', '困难'),
|
||||||
|
('数据挖掘', '模式发现', '洞察发现', '简单'),
|
||||||
|
('大数据分析', '海量数据处理', '可扩展性', '中等'),
|
||||||
|
('统计分析', '假设检验', '科学严谨', '简单')
|
||||||
|
]
|
||||||
|
|
||||||
|
for tech, scenario, advantage, difficulty in technologies:
|
||||||
|
row_cells = table.add_row().cells
|
||||||
|
row_cells[0].text = tech
|
||||||
|
row_cells[1].text = scenario
|
||||||
|
row_cells[2].text = advantage
|
||||||
|
row_cells[3].text = difficulty
|
||||||
|
|
||||||
|
# 添加工具和库部分
|
||||||
|
doc.add_heading('3. 常用工具和库', level=1)
|
||||||
|
|
||||||
|
# Python工具
|
||||||
|
doc.add_heading('Python生态系统', level=2)
|
||||||
|
python_tools = [
|
||||||
|
'NumPy - 数值计算基础库',
|
||||||
|
'Pandas - 数据操作和分析',
|
||||||
|
'Matplotlib/Seaborn - 数据可视化',
|
||||||
|
'Scikit-learn - 机器学习',
|
||||||
|
'TensorFlow/PyTorch - 深度学习',
|
||||||
|
'Jupyter Notebook - 交互式开发环境'
|
||||||
|
]
|
||||||
|
|
||||||
|
for tool in python_tools:
|
||||||
|
doc.add_paragraph(tool, style='List Bullet')
|
||||||
|
|
||||||
|
# R工具
|
||||||
|
doc.add_heading('R生态系统', level=2)
|
||||||
|
r_tools = [
|
||||||
|
'dplyr - 数据操作',
|
||||||
|
'ggplot2 - 数据可视化',
|
||||||
|
'caret - 机器学习',
|
||||||
|
'shiny - Web应用开发'
|
||||||
|
]
|
||||||
|
|
||||||
|
for tool in r_tools:
|
||||||
|
doc.add_paragraph(tool, style='List Bullet')
|
||||||
|
|
||||||
|
# 添加流程图说明
|
||||||
|
doc.add_heading('4. 数据科学流程', level=1)
|
||||||
|
|
||||||
|
flow_steps = [
|
||||||
|
'1. 问题定义:明确业务目标和分析需求',
|
||||||
|
'2. 数据收集:获取相关的内部和外部数据源',
|
||||||
|
'3. 数据清洗:处理缺失值、异常值和数据质量问题',
|
||||||
|
'4. 探索性数据分析:理解数据分布和特征关系',
|
||||||
|
'5. 特征工程:创建和选择有价值的特征',
|
||||||
|
'6. 模型建立:选择和训练合适的算法',
|
||||||
|
'7. 模型评估:验证模型性能和泛化能力',
|
||||||
|
'8. 模型部署:将模型集成到生产环境',
|
||||||
|
'9. 监控和维护:持续跟踪模型性能'
|
||||||
|
]
|
||||||
|
|
||||||
|
for step in flow_steps:
|
||||||
|
doc.add_paragraph(step, style='List Number')
|
||||||
|
|
||||||
|
# 添加挑战和趋势
|
||||||
|
doc.add_heading('5. Python在数据科学中的应用', level=1)
|
||||||
|
|
||||||
|
doc.add_paragraph(
|
||||||
|
'Python已成为数据科学领域最受欢迎的编程语言之一。'
|
||||||
|
'下图展示了Python的生态系统:'
|
||||||
|
)
|
||||||
|
|
||||||
|
# 添加图片
|
||||||
|
try:
|
||||||
|
doc.add_picture('/Users/liruwei/Documents/code/project/demo/base_rag/test_files/python.png',
|
||||||
|
width=Inches(4))
|
||||||
|
doc.add_paragraph('图1: Python生态系统', style='Caption')
|
||||||
|
except Exception as e:
|
||||||
|
print(f"警告:无法添加图片到DOCX: {e}")
|
||||||
|
doc.add_paragraph('[此处应显示Python生态系统图片]')
|
||||||
|
|
||||||
|
doc.add_heading('6. 行业挑战与未来趋势', level=1)
|
||||||
|
|
||||||
|
doc.add_paragraph(
|
||||||
|
'数据科学领域面临着数据隐私、算法偏见、可解释性等挑战。'
|
||||||
|
'未来趋势包括自动化机器学习(AutoML)、边缘计算、'
|
||||||
|
'联邦学习和可解释AI等技术的发展。'
|
||||||
|
)
|
||||||
|
|
||||||
|
# 保存文档
|
||||||
|
doc.save('/Users/liruwei/Documents/code/project/demo/base_rag/test_files/complex_data_science.docx')
|
||||||
|
print("已创建复杂的 DOCX 文件: complex_data_science.docx")
|
||||||
|
|
||||||
|
async def create_test_csv():
|
||||||
|
"""创建CSV测试文件"""
|
||||||
|
# 销售数据
|
||||||
|
sales_data = {
|
||||||
|
'日期': ['2024-01-01', '2024-01-02', '2024-01-03', '2024-01-04', '2024-01-05'],
|
||||||
|
'产品': ['笔记本电脑', '台式机', '平板电脑', '智能手机', '耳机'],
|
||||||
|
'销售额': [8500, 6200, 3200, 4500, 280],
|
||||||
|
'数量': [5, 4, 8, 9, 12],
|
||||||
|
'客户类型': ['企业', '个人', '学生', '个人', '学生'],
|
||||||
|
'销售员': ['张三', '李四', '王五', '张三', '李四']
|
||||||
|
}
|
||||||
|
|
||||||
|
df = pd.DataFrame(sales_data)
|
||||||
|
df.to_csv('/Users/liruwei/Documents/code/project/demo/base_rag/test_files/sales_data.csv',
|
||||||
|
index=False, encoding='utf-8')
|
||||||
|
print("已创建 CSV 文件: sales_data.csv")
|
||||||
|
|
||||||
|
async def create_test_excel():
|
||||||
|
"""创建Excel测试文件"""
|
||||||
|
# 创建多个工作表的Excel文件
|
||||||
|
with pd.ExcelWriter('/Users/liruwei/Documents/code/project/demo/base_rag/test_files/company_report.xlsx',
|
||||||
|
engine='openpyxl') as writer:
|
||||||
|
|
||||||
|
# 销售数据表
|
||||||
|
sales_data = {
|
||||||
|
'月份': ['1月', '2月', '3月', '4月', '5月', '6月'],
|
||||||
|
'销售额(万元)': [120, 135, 158, 142, 167, 189],
|
||||||
|
'利润率(%)': [15.2, 16.8, 18.3, 16.9, 19.1, 20.5],
|
||||||
|
'客户数': [856, 923, 1047, 978, 1156, 1289],
|
||||||
|
'新客户': [45, 67, 124, 55, 178, 133]
|
||||||
|
}
|
||||||
|
pd.DataFrame(sales_data).to_excel(writer, sheet_name='销售数据', index=False)
|
||||||
|
|
||||||
|
# 员工信息表
|
||||||
|
employee_data = {
|
||||||
|
'姓名': ['张三', '李四', '王五', '赵六', '钱七'],
|
||||||
|
'部门': ['销售部', '技术部', '市场部', '人事部', '财务部'],
|
||||||
|
'职位': ['销售经理', '高级工程师', '市场专员', 'HR主管', '会计师'],
|
||||||
|
'入职年份': [2020, 2019, 2021, 2018, 2022],
|
||||||
|
'年薪(万元)': [18, 25, 12, 16, 14]
|
||||||
|
}
|
||||||
|
pd.DataFrame(employee_data).to_excel(writer, sheet_name='员工信息', index=False)
|
||||||
|
|
||||||
|
# 产品分析表
|
||||||
|
product_data = {
|
||||||
|
'产品类别': ['电子产品', '服装', '食品', '图书', '家居'],
|
||||||
|
'销售占比(%)': [35.2, 28.6, 15.8, 12.4, 8.0],
|
||||||
|
'平均客单价': [1280, 320, 85, 45, 560],
|
||||||
|
'库存周转率': [4.2, 6.8, 12.5, 8.3, 3.9],
|
||||||
|
'客户满意度': [4.3, 4.1, 4.5, 4.2, 4.0]
|
||||||
|
}
|
||||||
|
pd.DataFrame(product_data).to_excel(writer, sheet_name='产品分析', index=False)
|
||||||
|
|
||||||
|
print("已创建 Excel 文件: company_report.xlsx")
|
||||||
|
|
||||||
|
async def create_test_pdf():
|
||||||
|
"""创建PDF测试文件"""
|
||||||
|
filename = '/Users/liruwei/Documents/code/project/demo/base_rag/test_files/ai_research_report.pdf'
|
||||||
|
|
||||||
|
c = canvas.Canvas(filename, pagesize=letter)
|
||||||
|
width, height = letter
|
||||||
|
|
||||||
|
# 标题
|
||||||
|
c.setFont("Helvetica-Bold", 20)
|
||||||
|
c.drawString(50, height - 50, "Artificial Intelligence Research Report")
|
||||||
|
|
||||||
|
# 副标题
|
||||||
|
c.setFont("Helvetica", 14)
|
||||||
|
c.drawString(50, height - 80, "A Comprehensive Study on Modern AI Technologies")
|
||||||
|
|
||||||
|
# 内容
|
||||||
|
c.setFont("Helvetica", 12)
|
||||||
|
y_position = height - 120
|
||||||
|
|
||||||
|
content = [
|
||||||
|
"1. Introduction",
|
||||||
|
"",
|
||||||
|
"Artificial Intelligence (AI) has become one of the most transformative",
|
||||||
|
"technologies of the 21st century. This report examines the current state",
|
||||||
|
"of AI research and its applications across various industries.",
|
||||||
|
"",
|
||||||
|
"2. Machine Learning Fundamentals",
|
||||||
|
"",
|
||||||
|
"Machine Learning is a subset of AI that enables computers to learn",
|
||||||
|
"without being explicitly programmed. Key approaches include:",
|
||||||
|
"- Supervised Learning: Learning from labeled data",
|
||||||
|
"- Unsupervised Learning: Finding patterns in unlabeled data",
|
||||||
|
"- Reinforcement Learning: Learning through interaction",
|
||||||
|
"",
|
||||||
|
"3. Deep Learning Revolution",
|
||||||
|
"",
|
||||||
|
"Deep Learning has revolutionized AI by enabling:",
|
||||||
|
"- Image Recognition: Achieving human-level accuracy",
|
||||||
|
"- Natural Language Processing: Understanding human language",
|
||||||
|
"- Speech Recognition: Converting speech to text",
|
||||||
|
"- Autonomous Systems: Self-driving cars and robots",
|
||||||
|
"",
|
||||||
|
"4. Applications in Industry",
|
||||||
|
"",
|
||||||
|
"Healthcare: AI assists in medical diagnosis and drug discovery",
|
||||||
|
"Finance: Fraud detection and algorithmic trading",
|
||||||
|
"Transportation: Autonomous vehicles and traffic optimization",
|
||||||
|
"Entertainment: Recommendation systems and content generation",
|
||||||
|
"",
|
||||||
|
"5. Python in AI Development",
|
||||||
|
"",
|
||||||
|
"Python has become the dominant language for AI development due to",
|
||||||
|
"its simplicity and rich ecosystem of libraries.",
|
||||||
|
"",
|
||||||
|
"[Python Ecosystem Image - see below]",
|
||||||
|
"",
|
||||||
|
"6. Challenges and Future Directions",
|
||||||
|
"",
|
||||||
|
"Current challenges include data privacy, algorithmic bias,",
|
||||||
|
"and the need for explainable AI. Future research focuses on",
|
||||||
|
"artificial general intelligence and quantum machine learning.",
|
||||||
|
]
|
||||||
|
|
||||||
|
for line in content:
|
||||||
|
if y_position < 50: # 新页面
|
||||||
|
c.showPage()
|
||||||
|
y_position = height - 50
|
||||||
|
c.setFont("Helvetica", 12)
|
||||||
|
|
||||||
|
if line.startswith(("1.", "2.", "3.", "4.", "5.")):
|
||||||
|
c.setFont("Helvetica-Bold", 12)
|
||||||
|
else:
|
||||||
|
c.setFont("Helvetica", 12)
|
||||||
|
|
||||||
|
c.drawString(50, y_position, line)
|
||||||
|
y_position -= 20
|
||||||
|
|
||||||
|
# 添加图片到PDF
|
||||||
|
try:
|
||||||
|
# 在新页面添加图片
|
||||||
|
if y_position < 200: # 确保有足够空间
|
||||||
|
c.showPage()
|
||||||
|
y_position = height - 50
|
||||||
|
|
||||||
|
# 添加图片标题
|
||||||
|
c.setFont("Helvetica-Bold", 12)
|
||||||
|
c.drawString(50, y_position, "Figure 1: Python Ecosystem")
|
||||||
|
y_position -= 30
|
||||||
|
|
||||||
|
# 添加图片
|
||||||
|
img_path = '/Users/liruwei/Documents/code/project/demo/base_rag/test_files/python.png'
|
||||||
|
if os.path.exists(img_path):
|
||||||
|
c.drawImage(img_path, 50, y_position - 200, width=300, height=150)
|
||||||
|
y_position -= 220
|
||||||
|
else:
|
||||||
|
c.setFont("Helvetica", 10)
|
||||||
|
c.drawString(50, y_position, "[Python ecosystem image would be displayed here]")
|
||||||
|
y_position -= 20
|
||||||
|
except Exception as e:
|
||||||
|
print(f"警告:无法添加图片到PDF: {e}")
|
||||||
|
c.setFont("Helvetica", 10)
|
||||||
|
c.drawString(50, y_position, "[Image placeholder - python.png]")
|
||||||
|
y_position -= 20
|
||||||
|
|
||||||
|
c.save()
|
||||||
|
print("已创建 PDF 文件: ai_research_report.pdf")
|
||||||
|
|
||||||
|
async def main():
|
||||||
|
"""创建所有测试文件"""
|
||||||
|
print("🔨 创建多格式测试文件...")
|
||||||
|
|
||||||
|
# 确保目录存在
|
||||||
|
os.makedirs('/Users/liruwei/Documents/code/project/demo/base_rag/test_files', exist_ok=True)
|
||||||
|
|
||||||
|
await create_complex_docx()
|
||||||
|
await create_test_csv()
|
||||||
|
await create_test_excel()
|
||||||
|
await create_test_pdf()
|
||||||
|
|
||||||
|
print("\n✅ 所有测试文件创建完成!")
|
||||||
|
print("📁 文件列表:")
|
||||||
|
print(" - complex_data_science.docx (带表格的复杂Word文档)")
|
||||||
|
print(" - sales_data.csv (销售数据CSV)")
|
||||||
|
print(" - company_report.xlsx (多工作表Excel)")
|
||||||
|
print(" - ai_research_report.pdf (AI研究报告PDF)")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
asyncio.run(main())
|
|
@ -0,0 +1,207 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
完整的多格式文件测试 - 包含图片的 DOCX、PDF、Excel、CSV
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
import asyncio
|
||||||
|
import warnings
|
||||||
|
from pathlib import Path
|
||||||
|
import shutil
|
||||||
|
|
||||||
|
# 过滤掉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
|
||||||
|
|
||||||
|
|
||||||
|
class AdvancedFormatRAG(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=3)
|
||||||
|
|
||||||
|
if not docs:
|
||||||
|
return "抱歉,没有找到相关信息。"
|
||||||
|
|
||||||
|
# 显示搜索到的文档来源
|
||||||
|
sources = []
|
||||||
|
contexts = []
|
||||||
|
for doc in docs:
|
||||||
|
source = doc.metadata.get("source_file", "未知来源")
|
||||||
|
content = doc.page_content.strip()
|
||||||
|
|
||||||
|
if source not in sources:
|
||||||
|
sources.append(source)
|
||||||
|
contexts.append(content)
|
||||||
|
|
||||||
|
context = "\n\n".join(contexts)
|
||||||
|
sources_str = "、".join(sources)
|
||||||
|
|
||||||
|
return f"基于以下文档({sources_str})的信息:\n\n{context}"
|
||||||
|
|
||||||
|
|
||||||
|
async def test_advanced_formats():
|
||||||
|
"""测试高级文件格式处理"""
|
||||||
|
print("🚀 高级多格式文件处理测试")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
# 清理旧的向量数据库
|
||||||
|
db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/advanced_formats")
|
||||||
|
if db_path.exists():
|
||||||
|
shutil.rmtree(db_path)
|
||||||
|
print("🧹 已清理旧的向量数据库")
|
||||||
|
|
||||||
|
# 创建RAG实例
|
||||||
|
rag = AdvancedFormatRAG(
|
||||||
|
vector_store_name="advanced_formats",
|
||||||
|
retriever_top_k=3,
|
||||||
|
storage_directory="/Users/liruwei/Documents/code/project/demo/base_rag/test_files",
|
||||||
|
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/advanced_status.db",
|
||||||
|
)
|
||||||
|
|
||||||
|
# 测试文件列表 - 包含新创建的文件
|
||||||
|
test_files = [
|
||||||
|
{
|
||||||
|
"file": "complex_data_science.docx",
|
||||||
|
"format": "DOCX",
|
||||||
|
"description": "复杂Word文档(含表格和图片)"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"file": "sales_data.csv",
|
||||||
|
"format": "CSV",
|
||||||
|
"description": "销售数据CSV文件"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"file": "company_report.xlsx",
|
||||||
|
"format": "XLSX",
|
||||||
|
"description": "多工作表Excel文件"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"file": "ai_research_report.pdf",
|
||||||
|
"format": "PDF",
|
||||||
|
"description": "AI研究报告PDF(含图片)"
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
||||||
|
print("📂 处理高级格式文件...")
|
||||||
|
processed_count = 0
|
||||||
|
|
||||||
|
for file_info in test_files:
|
||||||
|
filename = file_info["file"]
|
||||||
|
format_type = file_info["format"]
|
||||||
|
description = file_info["description"]
|
||||||
|
|
||||||
|
file_path = Path("../test_files") / filename
|
||||||
|
|
||||||
|
if not file_path.exists():
|
||||||
|
# 尝试绝对路径
|
||||||
|
file_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/test_files") / filename
|
||||||
|
|
||||||
|
if not file_path.exists():
|
||||||
|
print(f"❌ {format_type}: {filename} - 文件不存在")
|
||||||
|
continue
|
||||||
|
|
||||||
|
print(f"📄 处理 {format_type}: {filename}")
|
||||||
|
print(f" {description}")
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = await rag.ingest(str(file_path))
|
||||||
|
if result and result.get('success'):
|
||||||
|
print(f" ✅ 成功: {result['chunks_count']} 个片段")
|
||||||
|
processed_count += 1
|
||||||
|
else:
|
||||||
|
print(f" ⚠️ 跳过: {result.get('message', '可能已存在')}")
|
||||||
|
if "已经处理完毕" in str(result.get('message', '')):
|
||||||
|
processed_count += 1
|
||||||
|
except Exception as e:
|
||||||
|
print(f" ❌ 失败: {str(e)}")
|
||||||
|
print()
|
||||||
|
|
||||||
|
print(f"📊 处理完成: {processed_count}/{len(test_files)} 个文件")
|
||||||
|
print()
|
||||||
|
|
||||||
|
# 测试针对性查询
|
||||||
|
print("💬 高级格式查询测试...")
|
||||||
|
|
||||||
|
queries = [
|
||||||
|
{
|
||||||
|
"question": "数据科学的核心技术有哪些?",
|
||||||
|
"expected": "complex_data_science.docx"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"question": "销售数据中哪个产品销售额最高?",
|
||||||
|
"expected": "sales_data.csv"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"question": "公司员工信息包含哪些部门?",
|
||||||
|
"expected": "company_report.xlsx"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"question": "人工智能研究面临的挑战是什么?",
|
||||||
|
"expected": "ai_research_report.pdf"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"question": "Python在数据科学中的作用?",
|
||||||
|
"expected": "多个文档"
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
||||||
|
for i, query_info in enumerate(queries, 1):
|
||||||
|
question = query_info["question"]
|
||||||
|
expected = query_info["expected"]
|
||||||
|
|
||||||
|
print(f"\n❓ 查询 {i}: {question}")
|
||||||
|
print(f" 期望来源: {expected}")
|
||||||
|
|
||||||
|
try:
|
||||||
|
answer = await rag.query(question)
|
||||||
|
if "抱歉" not in answer:
|
||||||
|
# 分离来源信息和内容
|
||||||
|
parts = answer.split('\n\n', 1)
|
||||||
|
if len(parts) == 2:
|
||||||
|
source_info = parts[0]
|
||||||
|
content = parts[1]
|
||||||
|
|
||||||
|
print(f" 📚 {source_info}")
|
||||||
|
|
||||||
|
# 显示内容摘要(前150字符)
|
||||||
|
if len(content) > 150:
|
||||||
|
content_preview = content[:150] + "..."
|
||||||
|
else:
|
||||||
|
content_preview = content
|
||||||
|
|
||||||
|
print(f" 💡 {content_preview}")
|
||||||
|
else:
|
||||||
|
print(f" 💡 {answer[:150]}...")
|
||||||
|
else:
|
||||||
|
print(f" 💡 {answer}")
|
||||||
|
except Exception as e:
|
||||||
|
print(f" ❌ 查询失败: {str(e)}")
|
||||||
|
|
||||||
|
print("\n" + "=" * 60)
|
||||||
|
print("🎉 高级多格式文件测试完成!")
|
||||||
|
print("✅ 支持的格式:")
|
||||||
|
print(" 📄 DOCX - Word文档 (含表格、图片)")
|
||||||
|
print(" 📊 CSV - 逗号分隔值文件")
|
||||||
|
print(" 📈 XLSX - Excel工作簿 (多工作表)")
|
||||||
|
print(" 📑 PDF - 便携式文档格式 (含图片)")
|
||||||
|
print()
|
||||||
|
print("🔧 技术特性:")
|
||||||
|
print(" 🔄 异步处理 - 非阻塞I/O操作")
|
||||||
|
print(" 🧠 智能解析 - 自动识别文件格式")
|
||||||
|
print(" 🔍 跨格式查询 - 统一检索接口")
|
||||||
|
print(" 📋 表格数据提取 - 结构化信息处理")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
asyncio.run(test_advanced_formats())
|
|
@ -83,6 +83,26 @@ async def test_multiple_formats():
|
||||||
"file": "deep_learning_guide.docx",
|
"file": "deep_learning_guide.docx",
|
||||||
"format": "DOCX",
|
"format": "DOCX",
|
||||||
"description": "Word文档"
|
"description": "Word文档"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"file": "complex_data_science.docx",
|
||||||
|
"format": "DOCX",
|
||||||
|
"description": "复杂Word文档(含表格)"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"file": "sales_data.csv",
|
||||||
|
"format": "CSV",
|
||||||
|
"description": "CSV数据文件"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"file": "company_report.xlsx",
|
||||||
|
"format": "XLSX",
|
||||||
|
"description": "Excel工作簿"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"file": "ai_research_report.pdf",
|
||||||
|
"format": "PDF",
|
||||||
|
"description": "PDF文档"
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
|
||||||
|
@ -123,7 +143,11 @@ async def test_multiple_formats():
|
||||||
"Python有什么特点?",
|
"Python有什么特点?",
|
||||||
"什么是机器学习?",
|
"什么是机器学习?",
|
||||||
"深度学习的应用领域有哪些?",
|
"深度学习的应用领域有哪些?",
|
||||||
"Web框架有哪些?"
|
"数据科学的核心技术有哪些?",
|
||||||
|
"销售数据中哪个产品销售额最高?",
|
||||||
|
"公司员工的平均年薪是多少?",
|
||||||
|
"人工智能的主要挑战是什么?",
|
||||||
|
"机器学习有哪些类型?"
|
||||||
]
|
]
|
||||||
|
|
||||||
for query in queries:
|
for query in queries:
|
||||||
|
@ -155,7 +179,7 @@ async def test_multiple_formats():
|
||||||
|
|
||||||
print("\n" + "=" * 50)
|
print("\n" + "=" * 50)
|
||||||
print("✅ 多格式文件测试完成!")
|
print("✅ 多格式文件测试完成!")
|
||||||
print("支持的格式: TXT, MD, DOCX")
|
print("支持的格式: TXT, MD, DOCX, CSV, XLSX, PDF")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|
|
@ -567,6 +567,85 @@ class BaseRAG(ABC):
|
||||||
print("请运行: pip install unstructured python-docx")
|
print("请运行: pip install unstructured python-docx")
|
||||||
raise
|
raise
|
||||||
|
|
||||||
|
elif file_extension == '.csv':
|
||||||
|
# CSV文件
|
||||||
|
try:
|
||||||
|
import pandas as pd
|
||||||
|
from langchain_core.documents import Document
|
||||||
|
|
||||||
|
# 读取CSV文件
|
||||||
|
df = pd.read_csv(str(file_path))
|
||||||
|
|
||||||
|
# 将DataFrame转换为文本
|
||||||
|
csv_content = f"CSV文件: {file_path.name}\n\n"
|
||||||
|
csv_content += f"数据概览:\n行数: {len(df)}, 列数: {len(df.columns)}\n\n"
|
||||||
|
csv_content += f"列名: {', '.join(df.columns.tolist())}\n\n"
|
||||||
|
csv_content += "数据内容:\n"
|
||||||
|
csv_content += df.to_string(index=False)
|
||||||
|
|
||||||
|
return [Document(page_content=csv_content, metadata={"source": str(file_path)})]
|
||||||
|
except ImportError:
|
||||||
|
print("警告: 需要安装 pandas 来处理CSV文件")
|
||||||
|
print("请运行: pip install pandas")
|
||||||
|
raise
|
||||||
|
|
||||||
|
elif file_extension in ['.xls', '.xlsx']:
|
||||||
|
# Excel文件
|
||||||
|
try:
|
||||||
|
import pandas as pd
|
||||||
|
from langchain_core.documents import Document
|
||||||
|
|
||||||
|
# 读取Excel文件的所有工作表
|
||||||
|
excel_file = pd.ExcelFile(str(file_path))
|
||||||
|
documents = []
|
||||||
|
|
||||||
|
for sheet_name in excel_file.sheet_names:
|
||||||
|
df = pd.read_excel(str(file_path), sheet_name=sheet_name)
|
||||||
|
|
||||||
|
sheet_content = f"Excel文件: {file_path.name}\n工作表: {sheet_name}\n\n"
|
||||||
|
sheet_content += f"数据概览:\n行数: {len(df)}, 列数: {len(df.columns)}\n\n"
|
||||||
|
sheet_content += f"列名: {', '.join(df.columns.tolist())}\n\n"
|
||||||
|
sheet_content += "数据内容:\n"
|
||||||
|
sheet_content += df.to_string(index=False)
|
||||||
|
|
||||||
|
documents.append(Document(
|
||||||
|
page_content=sheet_content,
|
||||||
|
metadata={"source": str(file_path), "sheet": sheet_name}
|
||||||
|
))
|
||||||
|
|
||||||
|
return documents
|
||||||
|
except ImportError:
|
||||||
|
print("警告: 需要安装 pandas 和 openpyxl 来处理Excel文件")
|
||||||
|
print("请运行: pip install pandas openpyxl")
|
||||||
|
raise
|
||||||
|
|
||||||
|
elif file_extension == '.pdf':
|
||||||
|
# PDF文件
|
||||||
|
try:
|
||||||
|
from langchain_community.document_loaders import PyPDFLoader
|
||||||
|
loader = PyPDFLoader(str(file_path))
|
||||||
|
return loader.load()
|
||||||
|
except ImportError:
|
||||||
|
try:
|
||||||
|
# 备用方案:使用pdfplumber
|
||||||
|
import pdfplumber
|
||||||
|
from langchain_core.documents import Document
|
||||||
|
|
||||||
|
documents = []
|
||||||
|
with pdfplumber.open(str(file_path)) as pdf:
|
||||||
|
for i, page in enumerate(pdf.pages):
|
||||||
|
text = page.extract_text()
|
||||||
|
if text:
|
||||||
|
documents.append(Document(
|
||||||
|
page_content=text,
|
||||||
|
metadata={"source": str(file_path), "page": i + 1}
|
||||||
|
))
|
||||||
|
return documents
|
||||||
|
except ImportError:
|
||||||
|
print("警告: 需要安装 PyPDF2 或 pdfplumber 来处理PDF文件")
|
||||||
|
print("请运行: pip install PyPDF2 pdfplumber")
|
||||||
|
raise
|
||||||
|
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"不支持的文件类型: {file_extension}")
|
raise ValueError(f"不支持的文件类型: {file_extension}")
|
||||||
|
|
||||||
|
|
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
After Width: | Height: | Size: 98 KiB |
|
@ -0,0 +1,6 @@
|
||||||
|
日期,产品,销售额,数量,客户类型,销售员
|
||||||
|
2024-01-01,笔记本电脑,8500,5,企业,张三
|
||||||
|
2024-01-02,台式机,6200,4,个人,李四
|
||||||
|
2024-01-03,平板电脑,3200,8,学生,王五
|
||||||
|
2024-01-04,智能手机,4500,9,个人,张三
|
||||||
|
2024-01-05,耳机,280,12,学生,李四
|
|
|
@ -0,0 +1,6 @@
|
||||||
|
日期,产品,销售额,数量,客户类型,销售员
|
||||||
|
2024-01-01,笔记本电脑,8500,5,企业,张三
|
||||||
|
2024-01-02,台式机,6200,4,个人,李四
|
||||||
|
2024-01-03,平板电脑,3200,8,学生,王五
|
||||||
|
2024-01-04,智能手机,4500,9,个人,张三
|
||||||
|
2024-01-05,耳机,280,12,学生,李四
|
|
Loading…
Reference in New Issue