init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
200
core/agentic_media.py
Normal file
200
core/agentic_media.py
Normal file
@@ -0,0 +1,200 @@
|
||||
"""
|
||||
Agentic RAG - 富媒体处理 Mixin
|
||||
|
||||
包含图表查找、图片提取、富媒体附加等方法
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
import logging
|
||||
|
||||
from .agentic_base import logger
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RichMediaMixin:
|
||||
"""富媒体处理方法"""
|
||||
|
||||
def _find_figure(self, query: str, contexts: list, source: str = None) -> dict:
|
||||
"""精确查找图表，带 fallback"""
|
||||
patterns = [
|
||||
r'图\s*(\d+[\.\-]\d+)',
|
||||
r'Fig\.?\s*(\d+[\.\-]\d+)',
|
||||
r'Figure\s*(\d+[\.\-]\d+)',
|
||||
]
|
||||
|
||||
target_figure = None
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, query, re.IGNORECASE)
|
||||
if match:
|
||||
target_figure = match.group(1).replace('-', '.')
|
||||
break
|
||||
|
||||
if not target_figure:
|
||||
return {"found": False}
|
||||
|
||||
# 从 contexts 中查找
|
||||
for ctx in contexts:
|
||||
meta = ctx.get('meta', {})
|
||||
fig_num = meta.get('figure_number', '')
|
||||
if fig_num == target_figure:
|
||||
if not source or meta.get('source') == source:
|
||||
return {
|
||||
"found": True,
|
||||
"chunk_id": meta.get('chunk_id'),
|
||||
"source": meta.get('source'),
|
||||
"page": meta.get('page'),
|
||||
"caption": meta.get('caption'),
|
||||
"image_path": meta.get('image_path'),
|
||||
}
|
||||
|
||||
# Fallback: 直接查向量库
|
||||
try:
|
||||
from knowledge.manager import get_kb_manager
|
||||
kb_mgr = get_kb_manager()
|
||||
coll = kb_mgr.get_collection('public_kb')
|
||||
|
||||
if coll:
|
||||
where_conditions = [{'chunk_type': {'$in': ['image', 'chart']}}]
|
||||
if source:
|
||||
where_conditions.append({'source': source})
|
||||
|
||||
result = coll.get(
|
||||
where={'$and': where_conditions} if len(where_conditions) > 1 else where_conditions[0],
|
||||
include=['metadatas', 'documents']
|
||||
)
|
||||
|
||||
for meta, doc in zip(result.get('metadatas', []), result.get('documents', [])):
|
||||
if meta.get('figure_number') == target_figure:
|
||||
return {
|
||||
"found": True,
|
||||
"chunk_id": meta.get('chunk_id'),
|
||||
"source": meta.get('source'),
|
||||
"page": meta.get('page'),
|
||||
"caption": meta.get('caption'),
|
||||
"image_path": meta.get('image_path'),
|
||||
}
|
||||
caption = meta.get('caption', '') or (doc if doc else '')
|
||||
if f"图{target_figure}" in caption or f"图 {target_figure}" in caption:
|
||||
return {
|
||||
"found": True,
|
||||
"chunk_id": meta.get('chunk_id'),
|
||||
"source": meta.get('source'),
|
||||
"page": meta.get('page'),
|
||||
"caption": meta.get('caption'),
|
||||
"image_path": meta.get('image_path'),
|
||||
}
|
||||
except Exception as e:
|
||||
logger.warning(f"_find_figure fallback 查询失败: {e}")
|
||||
|
||||
return {"found": False}
|
||||
|
||||
def _get_images_for_source(self, source: str, collections: list = None) -> list:
|
||||
"""直接从向量库获取指定文件的所有图片"""
|
||||
try:
|
||||
from knowledge.manager import get_kb_manager
|
||||
kb_mgr = get_kb_manager()
|
||||
except ImportError:
|
||||
return []
|
||||
|
||||
images = []
|
||||
seen_ids = set()
|
||||
|
||||
target_collections = collections or ['public_kb']
|
||||
|
||||
for kb_name in target_collections:
|
||||
try:
|
||||
coll = kb_mgr.get_collection(kb_name)
|
||||
if not coll:
|
||||
continue
|
||||
|
||||
result = coll.get(
|
||||
where={'source': source},
|
||||
include=['metadatas']
|
||||
)
|
||||
|
||||
for meta in result.get('metadatas', []):
|
||||
images_json = meta.get('images_json')
|
||||
if images_json:
|
||||
try:
|
||||
imgs = json.loads(images_json)
|
||||
for img in imgs:
|
||||
img_id = img.get('id')
|
||||
if img_id and img_id not in seen_ids:
|
||||
seen_ids.add(img_id)
|
||||
images.append({
|
||||
"id": img_id,
|
||||
"caption": img.get("caption", ""),
|
||||
"url": f"/images/{img_id}",
|
||||
"page": img.get("page") or meta.get("page"),
|
||||
"source": source,
|
||||
"width": img.get("width"),
|
||||
"height": img.get("height")
|
||||
})
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning(f"从 {kb_name} 获取图片失败: {e}")
|
||||
continue
|
||||
|
||||
return images
|
||||
|
||||
def _extract_rich_media(self, contexts: list, sources_filter: list = None, max_images: int = 10,
|
||||
max_tables: int = 5) -> dict:
|
||||
"""从检索结果中提取富媒体（图片、表格）"""
|
||||
images = []
|
||||
tables = []
|
||||
seen_image_ids = set()
|
||||
seen_table_ids = set()
|
||||
|
||||
for ctx in contexts:
|
||||
meta = ctx.get('meta', {})
|
||||
source = meta.get('source', '')
|
||||
|
||||
# 过滤来源
|
||||
if sources_filter and source not in sources_filter:
|
||||
continue
|
||||
|
||||
# 提取图片
|
||||
images_json = meta.get('images_json')
|
||||
if images_json:
|
||||
try:
|
||||
imgs = json.loads(images_json)
|
||||
for img in imgs:
|
||||
img_id = img.get('id')
|
||||
if img_id and img_id not in seen_image_ids:
|
||||
seen_image_ids.add(img_id)
|
||||
images.append({
|
||||
"id": img_id,
|
||||
"caption": img.get("caption", ""),
|
||||
"url": f"/images/{img_id}",
|
||||
"page": img.get("page") or meta.get("page"),
|
||||
"source": source,
|
||||
"type": img.get("type", "image")
|
||||
})
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
# 提取表格
|
||||
table_json = meta.get('table_json')
|
||||
if table_json:
|
||||
try:
|
||||
tbl = json.loads(table_json)
|
||||
tbl_id = tbl.get('id') or meta.get('chunk_id')
|
||||
if tbl_id and tbl_id not in seen_table_ids:
|
||||
seen_table_ids.add(tbl_id)
|
||||
tables.append({
|
||||
"id": tbl_id,
|
||||
"caption": tbl.get("caption", ""),
|
||||
"markdown": tbl.get("markdown", ""),
|
||||
"page": meta.get("page"),
|
||||
"source": source
|
||||
})
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
return {
|
||||
"images": images[:max_images],
|
||||
"tables": tables[:max_tables]
|
||||
}
|
||||
Reference in New Issue
Block a user