feat: v1.2.0 - 图片去重与管理、微信机器人优化、搜索设置可配置
主要功能: - 图片上传时 OCR 内容去重（3个上传端点统一使用公共函数 _check_ocr_duplicate） - 图片管理 Tab：展示所有图片、手动删除、一键去重 - 搜索结果详情弹窗增加删除按钮（带确认弹窗） - 图片管理卡片点击查看详情（复用 showOcrDetailModal） - 搜索限制和 LLM 批量判断数量可通过网站设置 - MiniMax API 调用添加 reasoning_split=True - 企业微信机器人：WebSocket 长连接、图片搜索、配置化搜索数量 - 版本号升级至 1.2.0
This commit is contained in:
EduBrain Dev
5
app/__init__.py
Normal file
5
app/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""
|
||||
中文直播教育知识库系统 - edu-brain
|
||||
"""
|
||||
|
||||
__version__ = "1.2.0"
|
||||
3
app/api/__init__.py
Normal file
3
app/api/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
API 路由包
|
||||
"""
|
||||
3
app/api/v1/__init__.py
Normal file
3
app/api/v1/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
API v1 路由包
|
||||
"""
|
||||
565
app/api/v1/images.py
Normal file
565
app/api/v1/images.py
Normal file
@@ -0,0 +1,565 @@
|
||||
"""
|
||||
图片 OCR API 路由
|
||||
OCR 识别 + BM25 倒排索引搜索
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, Form, Query
|
||||
from fastapi.responses import StreamingResponse, Response
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import select, func, text, delete
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.config import settings
|
||||
from app.database import get_db
|
||||
from app.models.base import OCRImage
|
||||
from app.services.ocr_service import OCRService
|
||||
from app.services.search_engine import BM25Index, tokenize
|
||||
from app.services.llm_service import LLMService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter()
|
||||
|
||||
ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
|
||||
|
||||
# 全局 BM25 索引实例（在 main.py 中初始化）
|
||||
bm25_index: Optional[BM25Index] = None
|
||||
|
||||
|
||||
def _check_image_ext(filename: str) -> str:
|
||||
suffix = os.path.splitext(filename)[1].lower()
|
||||
if suffix not in ALLOWED_EXTENSIONS:
|
||||
raise HTTPException(status_code=400, detail=f"不支持的图片格式: {suffix}")
|
||||
return suffix
|
||||
|
||||
|
||||
def _save_ocr_result(record: OCRImage, ocr_result, db: AsyncSession):
|
||||
"""保存 OCR 识别结果到数据库"""
|
||||
record.ocr_text = ocr_result.text
|
||||
record.confidence = ocr_result.confidence
|
||||
record.provider = ocr_result.provider
|
||||
record.status = "completed"
|
||||
# 保留 keywords（如果有）
|
||||
if hasattr(ocr_result, 'keywords') and ocr_result.keywords:
|
||||
record.tags = json.dumps(ocr_result.keywords, ensure_ascii=False)
|
||||
|
||||
|
||||
def _add_to_index(record: OCRImage):
|
||||
"""将 OCR 结果添加到 BM25 索引"""
|
||||
if bm25_index is None or not record.ocr_text:
|
||||
return
|
||||
# 索引内容 = OCR 文本（权重最高）
|
||||
index_text = record.ocr_text
|
||||
# 如果有 tags/story_summary 也加入索引
|
||||
if record.tags:
|
||||
try:
|
||||
tags = json.loads(record.tags)
|
||||
if tags:
|
||||
index_text += "\n" + " ".join(tags)
|
||||
except Exception:
|
||||
pass
|
||||
if record.story_summary:
|
||||
index_text += "\n" + record.story_summary
|
||||
|
||||
bm25_index.add_document(
|
||||
doc_id=record.id,
|
||||
text=index_text,
|
||||
metadata={
|
||||
"file_path": record.file_path,
|
||||
"tags": json.loads(record.tags) if record.tags else [],
|
||||
"story_summary": record.story_summary or "",
|
||||
"confidence": record.confidence,
|
||||
"provider": record.provider,
|
||||
"created_at": str(record.created_at),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _check_ocr_duplicate(ocr_text: str, db: AsyncSession) -> Optional[int]:
|
||||
"""检查 OCR 文本是否已存在，返回已存在记录的 ID 或 None"""
|
||||
if not ocr_text:
|
||||
return None
|
||||
dup_query = select(OCRImage.id).where(
|
||||
OCRImage.ocr_text == ocr_text,
|
||||
OCRImage.status == "completed",
|
||||
).limit(1)
|
||||
dup_result = await db.execute(dup_query)
|
||||
return dup_result.scalar_one_or_none()
|
||||
|
||||
|
||||
# ─── 上传 ───
|
||||
|
||||
@router.post("/upload", summary="上传图片")
|
||||
async def upload_image(file: UploadFile, db: AsyncSession = Depends(get_db)):
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="文件名不能为空")
|
||||
_check_image_ext(file.filename)
|
||||
image_record = OCRImage(file_path=file.filename, status="pending")
|
||||
db.add(image_record)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
return {"id": image_record.id, "file_path": image_record.file_path, "original_filename": file.filename, "status": image_record.status, "created_at": str(image_record.created_at)}
|
||||
|
||||
|
||||
# ─── 单文件识别 ───
|
||||
|
||||
@router.post("/{image_id}/recognize", summary="识别图片文字")
|
||||
async def recognize_image(image_id: int, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(OCRImage).where(OCRImage.id == image_id))
|
||||
image_record = result.scalar_one_or_none()
|
||||
if image_record is None:
|
||||
raise HTTPException(status_code=404, detail="图片记录不存在")
|
||||
if not os.path.exists(image_record.file_path):
|
||||
raise HTTPException(status_code=404, detail=f"图片文件不存在: {image_record.file_path}")
|
||||
|
||||
image_record.status = "processing"
|
||||
await db.flush()
|
||||
|
||||
try:
|
||||
ocr_result = await OCRService.recognize(image_record.file_path)
|
||||
_save_ocr_result(image_record, ocr_result, db)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
_add_to_index(image_record)
|
||||
|
||||
return {
|
||||
"id": image_record.id,
|
||||
"file_path": image_record.file_path,
|
||||
"ocr_text": image_record.ocr_text,
|
||||
"tags": json.loads(image_record.tags) if image_record.tags else [],
|
||||
"confidence": image_record.confidence,
|
||||
"provider": image_record.provider,
|
||||
"status": image_record.status,
|
||||
"block_count": len(ocr_result.blocks),
|
||||
}
|
||||
except Exception as e:
|
||||
image_record.status = "failed"
|
||||
image_record.error_message = str(e)
|
||||
await db.flush()
|
||||
raise HTTPException(status_code=500, detail=f"OCR 识别失败: {str(e)}")
|
||||
|
||||
|
||||
# ─── 批量识别 ───
|
||||
|
||||
@router.post("/batch-recognize", summary="批量上传并识别图片")
|
||||
async def batch_recognize(files: List[UploadFile], db: AsyncSession = Depends(get_db)):
|
||||
results = []
|
||||
for file in files:
|
||||
if not file.filename:
|
||||
results.append({"filename": "unknown", "status": "error", "message": "文件名为空"})
|
||||
continue
|
||||
suffix = _check_image_ext(file.filename)
|
||||
dest_path = settings.images_dir / file.filename
|
||||
counter = 1
|
||||
stem = os.path.splitext(file.filename)[0]
|
||||
while dest_path.exists():
|
||||
dest_path = settings.images_dir / f"{stem}_{counter}{suffix}"
|
||||
counter += 1
|
||||
content = await file.read()
|
||||
dest_path.write_bytes(content)
|
||||
image_record = OCRImage(file_path=str(dest_path), status="processing")
|
||||
db.add(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
ocr_result = await OCRService.recognize(str(dest_path))
|
||||
# OCR 内容去重
|
||||
existing_id = await _check_ocr_duplicate(ocr_result.text, db)
|
||||
if existing_id is not None:
|
||||
# 已存在相同内容，回滚当前记录，删除图片文件
|
||||
await db.delete(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
os.unlink(dest_path)
|
||||
except OSError:
|
||||
pass
|
||||
results.append({
|
||||
"filename": file.filename,
|
||||
"file_path": str(dest_path),
|
||||
"status": "duplicate",
|
||||
"duplicate_of": existing_id,
|
||||
"ocr_text": ocr_result.text,
|
||||
})
|
||||
continue
|
||||
_save_ocr_result(image_record, ocr_result, db)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
_add_to_index(image_record)
|
||||
results.append({
|
||||
"id": image_record.id, "filename": file.filename,
|
||||
"file_path": str(dest_path), "status": "completed",
|
||||
"ocr_text": ocr_result.text,
|
||||
"tags": json.loads(image_record.tags) if image_record.tags else [],
|
||||
"confidence": ocr_result.confidence, "provider": ocr_result.provider,
|
||||
"block_count": len(ocr_result.blocks),
|
||||
})
|
||||
except Exception as e:
|
||||
image_record.status = "failed"
|
||||
image_record.error_message = str(e)
|
||||
await db.flush()
|
||||
results.append({"id": image_record.id, "filename": file.filename, "file_path": str(dest_path), "status": "failed", "message": str(e)})
|
||||
return {"total": len(results), "success": sum(1 for r in results if r["status"] == "completed"), "failed": sum(1 for r in results if r["status"] == "failed"), "duplicates": sum(1 for r in results if r["status"] == "duplicate"), "results": results}
|
||||
|
||||
|
||||
# ─── 服务器路径导入 ───
|
||||
|
||||
class PathImportRequest(BaseModel):
|
||||
paths: List[str]
|
||||
recursive: bool = False
|
||||
|
||||
|
||||
@router.post("/import-paths", summary="从服务器路径批量导入图片")
|
||||
async def import_from_paths(data: PathImportRequest, db: AsyncSession = Depends(get_db)):
|
||||
all_paths = []
|
||||
for p in data.paths:
|
||||
if os.path.isfile(p):
|
||||
if os.path.splitext(p)[1].lower() in ALLOWED_EXTENSIONS:
|
||||
all_paths.append(p)
|
||||
elif os.path.isdir(p):
|
||||
walk = os.walk(p) if data.recursive else [(p, [], os.listdir(p))]
|
||||
for root, dirs, files in walk:
|
||||
for f in sorted(files):
|
||||
fp = os.path.join(root, f)
|
||||
if os.path.isfile(fp) and os.path.splitext(f)[1].lower() in ALLOWED_EXTENSIONS:
|
||||
all_paths.append(fp)
|
||||
if not all_paths:
|
||||
return {"total": 0, "success": 0, "failed": 0, "message": "未找到图片文件", "results": []}
|
||||
results = []
|
||||
for file_path in all_paths:
|
||||
image_record = OCRImage(file_path=file_path, status="processing")
|
||||
db.add(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
ocr_result = await OCRService.recognize(file_path)
|
||||
# OCR 内容去重
|
||||
existing_id = await _check_ocr_duplicate(ocr_result.text, db)
|
||||
if existing_id is not None:
|
||||
await db.delete(image_record)
|
||||
await db.flush()
|
||||
results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "duplicate", "duplicate_of": existing_id, "ocr_text": ocr_result.text})
|
||||
continue
|
||||
_save_ocr_result(image_record, ocr_result, db)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
_add_to_index(image_record)
|
||||
results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "completed", "ocr_text": ocr_result.text, "tags": json.loads(image_record.tags) if image_record.tags else [], "confidence": ocr_result.confidence, "provider": ocr_result.provider})
|
||||
except Exception as e:
|
||||
image_record.status = "failed"
|
||||
image_record.error_message = str(e)
|
||||
await db.flush()
|
||||
results.append({"id": image_record.id, "filename": os.path.basename(file_path), "file_path": file_path, "status": "failed", "message": str(e)})
|
||||
return {"total": len(results), "success": sum(1 for r in results if r["status"] == "completed"), "failed": sum(1 for r in results if r["status"] == "failed"), "results": results}
|
||||
|
||||
|
||||
# ─── 一步上传+识别 ───
|
||||
|
||||
@router.post("/recognize-direct", summary="直接上传并识别图片")
|
||||
async def recognize_direct(file: UploadFile, db: AsyncSession = Depends(get_db)):
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="文件名不能为空")
|
||||
suffix = _check_image_ext(file.filename)
|
||||
dest_path = settings.images_dir / file.filename
|
||||
counter = 1
|
||||
stem = os.path.splitext(file.filename)[0]
|
||||
while dest_path.exists():
|
||||
dest_path = settings.images_dir / f"{stem}_{counter}{suffix}"
|
||||
counter += 1
|
||||
content = await file.read()
|
||||
dest_path.write_bytes(content)
|
||||
image_record = OCRImage(file_path=str(dest_path), status="processing")
|
||||
db.add(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
ocr_result = await OCRService.recognize(str(dest_path))
|
||||
# OCR 内容去重
|
||||
existing_id = await _check_ocr_duplicate(ocr_result.text, db)
|
||||
if existing_id is not None:
|
||||
await db.delete(image_record)
|
||||
await db.flush()
|
||||
try:
|
||||
os.unlink(dest_path)
|
||||
except OSError:
|
||||
pass
|
||||
return {
|
||||
"id": image_record.id, "file_path": str(dest_path),
|
||||
"original_filename": file.filename,
|
||||
"status": "duplicate", "duplicate_of": existing_id,
|
||||
"ocr_text": ocr_result.text,
|
||||
}
|
||||
_save_ocr_result(image_record, ocr_result, db)
|
||||
await db.flush()
|
||||
await db.refresh(image_record)
|
||||
_add_to_index(image_record)
|
||||
return {
|
||||
"id": image_record.id, "file_path": image_record.file_path,
|
||||
"original_filename": file.filename,
|
||||
"ocr_text": image_record.ocr_text,
|
||||
"tags": json.loads(image_record.tags) if image_record.tags else [],
|
||||
"confidence": image_record.confidence, "provider": image_record.provider,
|
||||
"status": image_record.status,
|
||||
"blocks": [{"text": b.text, "bbox": b.bbox, "confidence": b.confidence} for b in ocr_result.blocks],
|
||||
}
|
||||
except Exception as e:
|
||||
image_record.status = "failed"
|
||||
image_record.error_message = str(e)
|
||||
await db.flush()
|
||||
raise HTTPException(status_code=500, detail=f"OCR 识别失败: {str(e)}")
|
||||
|
||||
|
||||
# ─── LLM 精排搜索（SSE 流式） ───
|
||||
|
||||
@router.get("/search", summary="AI 搜索图片（流式）")
|
||||
async def search_images(
|
||||
keyword: str,
|
||||
limit: int = Query(5, ge=1, le=100),
|
||||
sort: str = Query("time_desc"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
AI 搜索图片（SSE 流式返回）。
|
||||
1. 从 DB 取所有图片，按指定排序
|
||||
2. 每批 5 条发给 Qwen3-8B 判断
|
||||
3. 并发池最多 10 个请求
|
||||
4. 匹配结果实时 SSE 推送
|
||||
5. 找够 limit 个或遍历完 DB 后结束
|
||||
"""
|
||||
|
||||
async def event_stream():
|
||||
found = 0
|
||||
offset = 0
|
||||
batch_size = settings.JUDGE_BATCH_SIZE
|
||||
max_concurrent = 10
|
||||
semaphore = asyncio.Semaphore(max_concurrent)
|
||||
|
||||
# 发送开始事件
|
||||
yield f"data: {json.dumps({'type': 'start', 'keyword': keyword, 'limit': limit}, ensure_ascii=False)}\n\n"
|
||||
|
||||
# 先获取总记录数
|
||||
count_result = await db.execute(
|
||||
select(func.count(OCRImage.id)).where(OCRImage.status == "completed")
|
||||
)
|
||||
total_records = count_result.scalar() or 0
|
||||
|
||||
if total_records == 0:
|
||||
yield f"data: {json.dumps({'type': 'done', 'total_found': 0, 'total_checked': 0}, ensure_ascii=False)}\n\n"
|
||||
return
|
||||
|
||||
# 排序
|
||||
order_clause = OCRImage.created_at.desc()
|
||||
if sort == "time_asc":
|
||||
order_clause = OCRImage.created_at.asc()
|
||||
elif sort == "random":
|
||||
order_clause = text("RANDOM()")
|
||||
|
||||
# 用于收集结果的队列
|
||||
result_queue = asyncio.Queue()
|
||||
checked_count = 0
|
||||
|
||||
async def process_batch(batch_items: list, batch_offset: int):
|
||||
"""处理一批图片：发给 LLM 判断"""
|
||||
nonlocal checked_count
|
||||
async with semaphore:
|
||||
articles = [{"id": item.id, "text": item.ocr_text or ""} for item in batch_items]
|
||||
matched_ids = await LLMService.judge_batch(keyword, articles)
|
||||
|
||||
checked_count += len(batch_items)
|
||||
|
||||
# 发送进度
|
||||
progress_data = {
|
||||
"type": "progress",
|
||||
"checked": checked_count,
|
||||
"total": total_records,
|
||||
"found": found,
|
||||
}
|
||||
await result_queue.put(progress_data)
|
||||
|
||||
# 如果有匹配且未达上限，发送结果
|
||||
for match_id in matched_ids:
|
||||
if found >= limit:
|
||||
break
|
||||
item = next((i for i in batch_items if i.id == match_id), None)
|
||||
if item:
|
||||
file_name = item.file_path.split('/')[-1] if item.file_path else ''
|
||||
# 读取图片并转为 base64（用于发送到企业微信等）
|
||||
image_base64 = ""
|
||||
if file_name:
|
||||
img_path = settings.images_dir / file_name
|
||||
if img_path.exists():
|
||||
image_base64 = base64.b64encode(img_path.read_bytes()).decode()
|
||||
result_data = {
|
||||
"type": "result",
|
||||
"id": item.id,
|
||||
"file_path": item.file_path,
|
||||
"image_url": file_name,
|
||||
"image_base64": image_base64,
|
||||
"ocr_text": item.ocr_text,
|
||||
"ocr_text_preview": (item.ocr_text[:150] + "...") if item.ocr_text and len(item.ocr_text) > 150 else item.ocr_text,
|
||||
"tags": json.loads(item.tags) if item.tags else [],
|
||||
"confidence": item.confidence,
|
||||
"provider": item.provider,
|
||||
"created_at": str(item.created_at),
|
||||
}
|
||||
await result_queue.put(result_data)
|
||||
|
||||
# 启动所有批次的异步任务
|
||||
tasks = []
|
||||
while offset < total_records:
|
||||
query = select(OCRImage).where(OCRImage.status == "completed").order_by(order_clause).offset(offset).limit(batch_size)
|
||||
result = await db.execute(query)
|
||||
batch_items = result.scalars().all()
|
||||
|
||||
if not batch_items:
|
||||
break
|
||||
|
||||
tasks.append(asyncio.create_task(process_batch(batch_items, offset)))
|
||||
offset += batch_size
|
||||
|
||||
# 等待所有任务完成
|
||||
async def wait_all():
|
||||
await asyncio.gather(*tasks)
|
||||
await result_queue.put({"type": "all_done"})
|
||||
|
||||
asyncio.create_task(wait_all())
|
||||
|
||||
# 从队列读取结果并 yield
|
||||
while True:
|
||||
try:
|
||||
data = await asyncio.wait_for(result_queue.get(), timeout=300)
|
||||
except asyncio.TimeoutError:
|
||||
break
|
||||
|
||||
if data.get("type") == "all_done":
|
||||
break
|
||||
|
||||
if data.get("type") == "result":
|
||||
found += 1
|
||||
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||||
if found >= limit:
|
||||
# 找够了，发送结束事件
|
||||
yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n"
|
||||
return
|
||||
elif data.get("type") == "progress":
|
||||
yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||||
|
||||
# 发送结束事件
|
||||
yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n"
|
||||
|
||||
return StreamingResponse(
|
||||
event_stream(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ─── 删除单张图片 ───
|
||||
|
||||
@router.delete("/{image_id}", summary="删除单张图片", status_code=204)
|
||||
async def delete_image(image_id: int, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(OCRImage).where(OCRImage.id == image_id))
|
||||
image_record = result.scalar_one_or_none()
|
||||
if image_record is None:
|
||||
raise HTTPException(status_code=404, detail="图片记录不存在")
|
||||
# 从 BM25 索引移除
|
||||
if bm25_index is not None:
|
||||
bm25_index.remove_document(image_record.id)
|
||||
# 删除磁盘上的图片文件
|
||||
try:
|
||||
os.unlink(image_record.file_path)
|
||||
except OSError:
|
||||
pass
|
||||
# 删除数据库记录
|
||||
await db.delete(image_record)
|
||||
await db.flush()
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
# ─── 其他接口 ───
|
||||
|
||||
@router.get("/{image_id}", summary="获取图片识别结果")
|
||||
async def get_image_result(image_id: int, db: AsyncSession = Depends(get_db)):
|
||||
result = await db.execute(select(OCRImage).where(OCRImage.id == image_id))
|
||||
image_record = result.scalar_one_or_none()
|
||||
if image_record is None:
|
||||
raise HTTPException(status_code=404, detail="图片记录不存在")
|
||||
return image_record.to_dict()
|
||||
|
||||
|
||||
@router.get("", summary="获取图片列表")
|
||||
async def list_images(page: int = 1, page_size: int = 20, status: Optional[str] = None, db: AsyncSession = Depends(get_db)):
|
||||
query = select(OCRImage)
|
||||
count_query = select(func.count(OCRImage.id))
|
||||
if status:
|
||||
query = query.where(OCRImage.status == status)
|
||||
count_query = count_query.where(OCRImage.status == status)
|
||||
total_result = await db.execute(count_query)
|
||||
total = total_result.scalar() or 0
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(OCRImage.created_at.desc()).offset(offset).limit(page_size)
|
||||
result = await db.execute(query)
|
||||
items = result.scalars().all()
|
||||
return {"total": total, "page": page, "page_size": page_size, "items": [item.to_dict() for item in items]}
|
||||
|
||||
|
||||
# ─── 一键去重 ───
|
||||
|
||||
@router.post("/dedup", summary="一键去重：删除 OCR 内容重复的图片")
|
||||
async def dedup_images(db: AsyncSession = Depends(get_db)):
|
||||
# 查询所有 status='completed' 且 ocr_text 不为空的记录
|
||||
query = select(OCRImage).where(
|
||||
OCRImage.status == "completed",
|
||||
OCRImage.ocr_text.isnot(None),
|
||||
OCRImage.ocr_text != "",
|
||||
).order_by(OCRImage.created_at.asc())
|
||||
result = await db.execute(query)
|
||||
all_records = result.scalars().all()
|
||||
|
||||
total_checked = len(all_records)
|
||||
# 按 ocr_text 分组
|
||||
groups: dict[str, list[OCRImage]] = defaultdict(list)
|
||||
for record in all_records:
|
||||
groups[record.ocr_text].append(record)
|
||||
|
||||
duplicates_found = 0
|
||||
deleted = 0
|
||||
kept = 0
|
||||
|
||||
for ocr_text, group in groups.items():
|
||||
if len(group) <= 1:
|
||||
kept += 1
|
||||
continue
|
||||
# 保留 created_at 最早的一条（已按 asc 排序，第一条即最早）
|
||||
duplicates_found += len(group) - 1
|
||||
kept += 1
|
||||
for record in group[1:]:
|
||||
# 从 BM25 索引移除
|
||||
if bm25_index is not None:
|
||||
bm25_index.remove_document(record.id)
|
||||
# 删除磁盘上的图片文件
|
||||
try:
|
||||
os.unlink(record.file_path)
|
||||
except OSError:
|
||||
pass
|
||||
# 删除数据库记录
|
||||
await db.delete(record)
|
||||
deleted += 1
|
||||
|
||||
await db.flush()
|
||||
return {
|
||||
"total_checked": total_checked,
|
||||
"duplicates_found": duplicates_found,
|
||||
"deleted": deleted,
|
||||
"kept": kept,
|
||||
}
|
||||
215
app/api/v1/import_export.py
Normal file
215
app/api/v1/import_export.py
Normal file
@@ -0,0 +1,215 @@
|
||||
"""
|
||||
导入导出 API 路由
|
||||
提供文件导入和知识库导出接口
|
||||
支持 .md / .txt / .docx 格式
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.config import settings
|
||||
from app.database import get_db
|
||||
from app.services.import_service import ImportService
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# 支持的导入文件格式
|
||||
SUPPORTED_EXTENSIONS = (".md", ".txt", ".docx")
|
||||
|
||||
|
||||
@router.post("/file", summary="从文件导入知识")
|
||||
async def import_from_file(
|
||||
file: UploadFile = File(..., description="上传文件（.md / .txt / .docx）"),
|
||||
course_name: Optional[str] = Form(None, description="课程名称"),
|
||||
teacher_name: Optional[str] = Form(None, description="讲师名称"),
|
||||
live_date: Optional[str] = Form(None, description="直播日期 (YYYY-MM-DD)"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
上传文件并导入到知识库。
|
||||
|
||||
- Markdown (.md)：解析 frontmatter 元数据，按 ## 标题分页
|
||||
- 纯文本 (.txt)：整体作为一个页面
|
||||
- Word 文档 (.docx)：提取文本，按标题分页
|
||||
- 导入后自动分块、生成嵌入向量、建立全文索引
|
||||
"""
|
||||
# 验证文件类型
|
||||
if not file.filename:
|
||||
raise HTTPException(status_code=400, detail="文件名不能为空")
|
||||
|
||||
suffix = os.path.splitext(file.filename)[1].lower()
|
||||
if suffix not in SUPPORTED_EXTENSIONS:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"仅支持 {', '.join(SUPPORTED_EXTENSIONS)} 文件",
|
||||
)
|
||||
|
||||
# 保存上传文件到 transcripts 目录
|
||||
dest_dir = settings.transcripts_dir
|
||||
dest_path = dest_dir / file.filename
|
||||
|
||||
# 如果文件已存在，添加序号
|
||||
counter = 1
|
||||
original_stem = os.path.splitext(file.filename)[0]
|
||||
while dest_path.exists():
|
||||
dest_path = dest_dir / f"{original_stem}_{counter}{suffix}"
|
||||
counter += 1
|
||||
|
||||
# 写入文件
|
||||
content = await file.read()
|
||||
dest_path.write_bytes(content)
|
||||
|
||||
# 执行导入
|
||||
service = ImportService(db)
|
||||
try:
|
||||
result = await service.import_file(
|
||||
file_path=str(dest_path),
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
live_date=live_date,
|
||||
)
|
||||
return {
|
||||
"message": "导入成功",
|
||||
"file": str(dest_path),
|
||||
**result,
|
||||
}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"导入失败: {str(e)}")
|
||||
|
||||
|
||||
@router.post("/directory", summary="从目录批量导入")
|
||||
async def import_from_directory(
|
||||
directory: str = Form(..., description="数据目录中的子目录名（相对于 transcripts 目录）"),
|
||||
course_name: Optional[str] = Form(None, description="课程名称"),
|
||||
teacher_name: Optional[str] = Form(None, description="讲师名称"),
|
||||
live_date: Optional[str] = Form(None, description="直播日期"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
批量导入指定目录下的所有文件（.md / .txt / .docx）。
|
||||
"""
|
||||
target_dir = settings.transcripts_dir / directory
|
||||
if not target_dir.exists():
|
||||
raise HTTPException(status_code=404, detail=f"目录不存在: {directory}")
|
||||
|
||||
service = ImportService(db)
|
||||
total_pages = 0
|
||||
total_chunks = 0
|
||||
files_processed = 0
|
||||
errors: list[str] = []
|
||||
|
||||
for filepath in sorted(target_dir.iterdir()):
|
||||
if filepath.suffix.lower() not in SUPPORTED_EXTENSIONS:
|
||||
continue
|
||||
|
||||
try:
|
||||
result = await service.import_file(
|
||||
file_path=str(filepath),
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
live_date=live_date,
|
||||
)
|
||||
total_pages += result["pages"]
|
||||
total_chunks += result["chunks"]
|
||||
files_processed += 1
|
||||
except Exception as e:
|
||||
errors.append(f"{filepath.name}: {str(e)}")
|
||||
|
||||
return {
|
||||
"message": "批量导入完成",
|
||||
"files_processed": files_processed,
|
||||
"total_pages": total_pages,
|
||||
"total_chunks": total_chunks,
|
||||
"errors": errors if errors else None,
|
||||
}
|
||||
|
||||
|
||||
@router.get("/export", summary="导出知识库")
|
||||
async def export_knowledge_base(
|
||||
format: str = "json",
|
||||
course_name: Optional[str] = None,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
导出知识库内容。
|
||||
|
||||
Args:
|
||||
format: 导出格式，支持 json / markdown
|
||||
course_name: 按课程名称过滤导出
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
query = text("""
|
||||
SELECT id, title, content, source_file, course_name, teacher_name,
|
||||
live_date, page_number, metadata_json, created_at
|
||||
FROM knowledge_pages
|
||||
WHERE (:course_name IS NULL OR course_name = :course_name)
|
||||
ORDER BY created_at DESC
|
||||
""")
|
||||
result = await db.execute(query, {"course_name": course_name})
|
||||
pages = result.fetchall()
|
||||
|
||||
if format == "markdown":
|
||||
# 导出为 Markdown 文件
|
||||
lines = []
|
||||
for row in pages:
|
||||
lines.append(f"# {row.title}")
|
||||
lines.append("")
|
||||
meta_parts = []
|
||||
if row.course_name:
|
||||
meta_parts.append(f"课程: {row.course_name}")
|
||||
if row.teacher_name:
|
||||
meta_parts.append(f"讲师: {row.teacher_name}")
|
||||
if row.live_date:
|
||||
meta_parts.append(f"日期: {row.live_date}")
|
||||
if meta_parts:
|
||||
lines.append(" | ".join(meta_parts))
|
||||
lines.append("")
|
||||
lines.append(row.content)
|
||||
lines.append("")
|
||||
lines.append("---")
|
||||
lines.append("")
|
||||
|
||||
export_content = "\n".join(lines)
|
||||
ext = "md"
|
||||
else:
|
||||
# 导出为 JSON
|
||||
pages_data = []
|
||||
for row in pages:
|
||||
pages_data.append({
|
||||
"id": row.id,
|
||||
"title": row.title,
|
||||
"content": row.content,
|
||||
"source_file": row.source_file,
|
||||
"course_name": row.course_name,
|
||||
"teacher_name": row.teacher_name,
|
||||
"live_date": row.live_date,
|
||||
"page_number": row.page_number,
|
||||
"metadata_json": row.metadata_json,
|
||||
"created_at": str(row.created_at),
|
||||
})
|
||||
export_content = json.dumps(pages_data, ensure_ascii=False, indent=2)
|
||||
ext = "json"
|
||||
|
||||
# 保存到 exports 目录
|
||||
export_dir = settings.exports_dir
|
||||
export_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"export_{timestamp}.{ext}"
|
||||
export_path = export_dir / filename
|
||||
export_path.write_text(export_content, encoding="utf-8")
|
||||
|
||||
return {
|
||||
"message": "导出成功",
|
||||
"file": filename,
|
||||
"pages": len(pages),
|
||||
"format": format,
|
||||
}
|
||||
112
app/api/v1/pages.py
Normal file
112
app/api/v1/pages.py
Normal file
@@ -0,0 +1,112 @@
|
||||
"""
|
||||
知识页面 API 路由
|
||||
提供知识页面的 CRUD 接口
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
from app.schemas.page import PageCreate, PageDetailResponse, PageListResponse, PageResponse, PageUpdate
|
||||
from app.services.page_service import PageService
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("", response_model=PageListResponse, summary="获取知识页面列表")
|
||||
async def list_pages(
|
||||
page: int = Query(1, ge=1, description="页码"),
|
||||
page_size: int = Query(20, ge=1, le=100, description="每页数量"),
|
||||
course_name: Optional[str] = Query(None, description="按课程名称过滤"),
|
||||
teacher_name: Optional[str] = Query(None, description="按讲师名称过滤"),
|
||||
keyword: Optional[str] = Query(None, description="关键词搜索"),
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""分页查询知识页面列表，支持按课程/讲师/关键词过滤"""
|
||||
service = PageService(db)
|
||||
result = await service.list_pages(
|
||||
page=page,
|
||||
page_size=page_size,
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
keyword=keyword,
|
||||
)
|
||||
return PageListResponse(
|
||||
total=result["total"],
|
||||
page=result["page"],
|
||||
page_size=result["page_size"],
|
||||
items=[PageResponse.model_validate(p) for p in result["items"]],
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{page_id}", response_model=PageDetailResponse, summary="获取知识页面详情")
|
||||
async def get_page(
|
||||
page_id: int,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""根据 ID 获取知识页面详情，包含关联的分块信息"""
|
||||
service = PageService(db)
|
||||
page = await service.get_page(page_id)
|
||||
if page is None:
|
||||
raise HTTPException(status_code=404, detail="知识页面不存在")
|
||||
|
||||
chunks = await service.get_page_chunks(page_id)
|
||||
return PageDetailResponse(
|
||||
**page.to_dict(),
|
||||
chunks=[c.to_dict() for c in chunks],
|
||||
)
|
||||
|
||||
|
||||
@router.post("", response_model=PageResponse, status_code=201, summary="创建知识页面")
|
||||
async def create_page(
|
||||
data: PageCreate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""创建新的知识页面"""
|
||||
service = PageService(db)
|
||||
page = await service.create_page(data)
|
||||
return PageResponse.model_validate(page)
|
||||
|
||||
|
||||
@router.put("/{page_id}", response_model=PageResponse, summary="更新知识页面")
|
||||
async def update_page(
|
||||
page_id: int,
|
||||
data: PageUpdate,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""更新知识页面信息"""
|
||||
service = PageService(db)
|
||||
page = await service.update_page(page_id, data)
|
||||
if page is None:
|
||||
raise HTTPException(status_code=404, detail="知识页面不存在")
|
||||
return PageResponse.model_validate(page)
|
||||
|
||||
|
||||
@router.delete("/{page_id}", status_code=204, summary="删除知识页面")
|
||||
async def delete_page(
|
||||
page_id: int,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""删除知识页面及其关联的所有分块"""
|
||||
service = PageService(db)
|
||||
success = await service.delete_page(page_id)
|
||||
if not success:
|
||||
raise HTTPException(status_code=404, detail="知识页面不存在")
|
||||
|
||||
|
||||
@router.post("/{page_id}/reindex", summary="重新索引知识页面")
|
||||
async def reindex_page(
|
||||
page_id: int,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""重新对知识页面进行分块和向量化（页面内容更新后使用）"""
|
||||
service = PageService(db)
|
||||
try:
|
||||
result = await service.reindex_page(page_id)
|
||||
return result
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
30
app/api/v1/search.py
Normal file
30
app/api/v1/search.py
Normal file
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
语义搜索 API 路由
|
||||
提供向量搜索 + 全文搜索混合查询接口
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.database import get_db
|
||||
from app.schemas.search import SearchRequest, SearchResponse
|
||||
from app.services.search_service import SearchService
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("", response_model=SearchResponse, summary="语义搜索")
|
||||
async def semantic_search(
|
||||
request: SearchRequest,
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
对知识库进行语义搜索。
|
||||
|
||||
支持向量相似度搜索和中文全文搜索的混合排序。
|
||||
可按课程名称、讲师名称、直播日期范围过滤结果。
|
||||
"""
|
||||
service = SearchService(db)
|
||||
return await service.search(request)
|
||||
223
app/api/v1/settings.py
Normal file
223
app/api/v1/settings.py
Normal file
@@ -0,0 +1,223 @@
|
||||
"""
|
||||
系统设置 API 路由
|
||||
提供运行时配置的读取和修改接口
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class SettingsResponse(BaseModel):
|
||||
"""设置响应（隐藏敏感信息的中间位）"""
|
||||
embedding_provider: str
|
||||
embedding_model: str
|
||||
embedding_dimensions: int
|
||||
ocr_provider: str
|
||||
llm_provider: str
|
||||
minimax_base_url: str
|
||||
minimax_embedding_model: str
|
||||
minimax_chat_model: str
|
||||
deepseek_base_url: str
|
||||
deepseek_ocr_model: str
|
||||
openai_base_url: str | None
|
||||
chunk_size: int
|
||||
chunk_overlap: int
|
||||
search_limit: int
|
||||
judge_batch_size: int
|
||||
# API Key 只返回是否已配置（不返回实际值）
|
||||
has_minimax_key: bool
|
||||
has_deepseek_key: bool
|
||||
has_openai_key: bool
|
||||
has_zhipu_key: bool
|
||||
has_dashscope_key: bool
|
||||
has_aliyun_ocr: bool
|
||||
has_tencent_ocr: bool
|
||||
|
||||
|
||||
@router.get("", response_model=SettingsResponse, summary="获取当前设置")
|
||||
async def get_settings():
|
||||
"""获取当前系统设置（API Key 脱敏）"""
|
||||
return SettingsResponse(
|
||||
embedding_provider=settings.EMBEDDING_PROVIDER.value,
|
||||
embedding_model=settings.EMBEDDING_MODEL,
|
||||
embedding_dimensions=settings.EMBEDDING_DIMENSIONS,
|
||||
ocr_provider=settings.OCR_PROVIDER.value,
|
||||
llm_provider=settings.LLM_PROVIDER,
|
||||
minimax_base_url=settings.MINIMAX_BASE_URL,
|
||||
minimax_embedding_model=settings.MINIMAX_EMBEDDING_MODEL,
|
||||
minimax_chat_model=settings.MINIMAX_CHAT_MODEL,
|
||||
deepseek_base_url=settings.DEEPSEEK_BASE_URL,
|
||||
deepseek_ocr_model=settings.DEEPSEEK_OCR_MODEL,
|
||||
openai_base_url=settings.OPENAI_BASE_URL,
|
||||
chunk_size=settings.CHUNK_SIZE,
|
||||
chunk_overlap=settings.CHUNK_OVERLAP,
|
||||
search_limit=settings.SEARCH_LIMIT,
|
||||
judge_batch_size=settings.JUDGE_BATCH_SIZE,
|
||||
has_minimax_key=bool(settings.MINIMAX_API_KEY),
|
||||
has_deepseek_key=bool(settings.DEEPSEEK_API_KEY),
|
||||
has_openai_key=bool(settings.OPENAI_API_KEY),
|
||||
has_zhipu_key=bool(settings.ZHIPU_API_KEY),
|
||||
has_dashscope_key=bool(settings.DASHSCOPE_API_KEY),
|
||||
has_aliyun_ocr=bool(settings.ALIYUN_OCR_ACCESS_KEY),
|
||||
has_tencent_ocr=bool(settings.TENCENT_OCR_SECRET_ID),
|
||||
)
|
||||
|
||||
|
||||
class SettingsUpdate(BaseModel):
|
||||
"""设置更新请求（所有字段可选）"""
|
||||
embedding_provider: str | None = None
|
||||
embedding_model: str | None = None
|
||||
embedding_dimensions: int | None = None
|
||||
ocr_provider: str | None = None
|
||||
llm_provider: str | None = None
|
||||
minimax_api_key: str | None = None
|
||||
minimax_base_url: str | None = None
|
||||
minimax_embedding_model: str | None = None
|
||||
minimax_chat_model: str | None = None
|
||||
deepseek_api_key: str | None = None
|
||||
deepseek_base_url: str | None = None
|
||||
deepseek_ocr_model: str | None = None
|
||||
openai_api_key: str | None = None
|
||||
openai_base_url: str | None = None
|
||||
zhipu_api_key: str | None = None
|
||||
dashscope_api_key: str | None = None
|
||||
aliyun_ocr_access_key: str | None = None
|
||||
aliyun_ocr_secret: str | None = None
|
||||
tencent_ocr_secret_id: str | None = None
|
||||
tencent_ocr_secret_key: str | None = None
|
||||
chunk_size: int | None = None
|
||||
chunk_overlap: int | None = None
|
||||
search_limit: int | None = None
|
||||
judge_batch_size: int | None = None
|
||||
|
||||
|
||||
@router.put("", summary="更新设置")
|
||||
async def update_settings(data: SettingsUpdate):
|
||||
"""
|
||||
更新系统设置（运行时生效）。
|
||||
修改嵌入模型或 OCR 提供商后，对应的单例服务会自动重置。
|
||||
"""
|
||||
from app.config import EmbeddingProvider, OCRProvider
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
from app.services.ocr_service import OCRService
|
||||
from app.services.llm_service import LLMService
|
||||
|
||||
update_data = data.model_dump(exclude_none=True)
|
||||
if not update_data:
|
||||
raise HTTPException(status_code=400, detail="没有需要更新的字段")
|
||||
|
||||
provider_changed = False
|
||||
|
||||
# 映射字段名（API 用 snake_case，Settings 用 UPPER_CASE）
|
||||
field_map = {
|
||||
"embedding_provider": "EMBEDDING_PROVIDER",
|
||||
"embedding_model": "EMBEDDING_MODEL",
|
||||
"embedding_dimensions": "EMBEDDING_DIMENSIONS",
|
||||
"ocr_provider": "OCR_PROVIDER",
|
||||
"llm_provider": "LLM_PROVIDER",
|
||||
"minimax_api_key": "MINIMAX_API_KEY",
|
||||
"minimax_base_url": "MINIMAX_BASE_URL",
|
||||
"minimax_embedding_model": "MINIMAX_EMBEDDING_MODEL",
|
||||
"minimax_chat_model": "MINIMAX_CHAT_MODEL",
|
||||
"deepseek_api_key": "DEEPSEEK_API_KEY",
|
||||
"deepseek_base_url": "DEEPSEEK_BASE_URL",
|
||||
"deepseek_ocr_model": "DEEPSEEK_OCR_MODEL",
|
||||
"openai_api_key": "OPENAI_API_KEY",
|
||||
"openai_base_url": "OPENAI_BASE_URL",
|
||||
"zhipu_api_key": "ZHIPU_API_KEY",
|
||||
"dashscope_api_key": "DASHSCOPE_API_KEY",
|
||||
"aliyun_ocr_access_key": "ALIYUN_OCR_ACCESS_KEY",
|
||||
"aliyun_ocr_secret": "ALIYUN_OCR_SECRET",
|
||||
"tencent_ocr_secret_id": "TENCENT_OCR_SECRET_ID",
|
||||
"tencent_ocr_secret_key": "TENCENT_OCR_SECRET_KEY",
|
||||
"chunk_size": "CHUNK_SIZE",
|
||||
"chunk_overlap": "CHUNK_OVERLAP",
|
||||
"search_limit": "SEARCH_LIMIT",
|
||||
"judge_batch_size": "JUDGE_BATCH_SIZE",
|
||||
}
|
||||
|
||||
for api_field, config_field in field_map.items():
|
||||
if api_field in update_data:
|
||||
value = update_data[api_field]
|
||||
# 枚举类型需要转换
|
||||
if config_field == "EMBEDDING_PROVIDER":
|
||||
value = EmbeddingProvider(value)
|
||||
provider_changed = True
|
||||
elif config_field == "OCR_PROVIDER":
|
||||
value = OCRProvider(value)
|
||||
provider_changed = True
|
||||
setattr(settings, config_field, value)
|
||||
|
||||
# 如果提供商变了，重置单例
|
||||
if provider_changed:
|
||||
EmbeddingService.reset()
|
||||
OCRService.reset()
|
||||
LLMService.reset()
|
||||
|
||||
return {"message": "设置已更新", "provider_changed": provider_changed}
|
||||
|
||||
|
||||
@router.get("/stats", summary="获取知识库统计")
|
||||
async def get_stats():
|
||||
"""获取知识库统计信息"""
|
||||
from sqlalchemy import text
|
||||
from app.database import async_session_factory
|
||||
|
||||
async with async_session_factory() as session:
|
||||
result = await session.execute(text("""
|
||||
SELECT
|
||||
(SELECT COUNT(*) FROM knowledge_pages) AS page_count,
|
||||
(SELECT COUNT(*) FROM knowledge_chunks) AS chunk_count,
|
||||
(SELECT COUNT(*) FROM knowledge_chunks WHERE embedding IS NOT NULL) AS embedded_count,
|
||||
(SELECT COUNT(*) FROM ocr_images) AS image_count,
|
||||
(SELECT COUNT(*) FROM ocr_images WHERE status = 'completed') AS ocr_completed_count
|
||||
"""))
|
||||
row = result.fetchone()
|
||||
|
||||
return {
|
||||
"page_count": row.page_count,
|
||||
"chunk_count": row.chunk_count,
|
||||
"embedded_count": row.embedded_count,
|
||||
"image_count": row.image_count,
|
||||
"ocr_completed_count": row.ocr_completed_count,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/test/embedding", summary="测试嵌入模型连接")
|
||||
async def test_embedding():
|
||||
"""测试当前嵌入模型是否可用"""
|
||||
try:
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
result = await EmbeddingService.embed_single("测试")
|
||||
return {"success": True, "dimension": len(result), "message": f"嵌入模型正常，维度: {len(result)}"}
|
||||
except Exception as e:
|
||||
return {"success": False, "message": f"嵌入模型测试失败: {str(e)}"}
|
||||
|
||||
|
||||
@router.post("/test/llm", summary="测试 LLM 连接")
|
||||
async def test_llm():
|
||||
"""测试当前 LLM 是否可用"""
|
||||
try:
|
||||
from app.services.llm_service import LLMService
|
||||
result = await LLMService.chat([
|
||||
{"role": "user", "content": "请回复\"连接正常\""}
|
||||
], max_tokens=20)
|
||||
return {"success": True, "message": f"LLM 正常，回复: {result[:100]}"}
|
||||
except Exception as e:
|
||||
return {"success": False, "message": f"LLM 测试失败: {str(e)}"}
|
||||
|
||||
|
||||
@router.post("/test/ocr", summary="测试 OCR（使用示例文本）")
|
||||
async def test_ocr():
|
||||
"""测试当前 OCR 提供商是否可用（不实际上传图片，只检查配置）"""
|
||||
try:
|
||||
from app.services.ocr_service import OCRService
|
||||
provider = OCRService.get_instance()
|
||||
return {"success": True, "message": f"OCR 提供商 {provider.name} 已就绪"}
|
||||
except Exception as e:
|
||||
return {"success": False, "message": f"OCR 测试失败: {str(e)}"}
|
||||
216
app/config.py
Normal file
216
app/config.py
Normal file
@@ -0,0 +1,216 @@
|
||||
"""
|
||||
配置管理模块
|
||||
使用 pydantic-settings 从环境变量 / .env 文件加载配置
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class EmbeddingProvider(str, Enum):
|
||||
"""嵌入模型提供商"""
|
||||
OPENAI = "openai"
|
||||
ZHIPU = "zhipu"
|
||||
LOCAL_BGE = "local_bge"
|
||||
DASHSCOPE = "dashscope"
|
||||
MINIMAX = "minimax"
|
||||
|
||||
|
||||
class OCRProvider(str, Enum):
|
||||
"""OCR 识别提供商"""
|
||||
PADDLEOCR = "paddleocr"
|
||||
ALIYUN = "aliyun"
|
||||
TENCENT = "tencent"
|
||||
AUTO = "auto"
|
||||
DEEPSEEK = "deepseek"
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""
|
||||
全局配置类，所有配置项均可通过环境变量覆盖。
|
||||
环境变量优先级：系统环境变量 > .env 文件 > 默认值
|
||||
"""
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env",
|
||||
env_file_encoding="utf-8",
|
||||
case_sensitive=False,
|
||||
extra="ignore",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 数据库 ────────────────────────────
|
||||
DATABASE_URL: str = Field(
|
||||
default="postgresql+asyncpg://postgres:postgres@db:5432/edu_brain",
|
||||
description="PostgreSQL 异步连接字符串",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 嵌入模型 ────────────────────────────
|
||||
EMBEDDING_PROVIDER: EmbeddingProvider = Field(
|
||||
default=EmbeddingProvider.LOCAL_BGE,
|
||||
description="嵌入模型提供商: openai / zhipu / local_bge / dashscope",
|
||||
)
|
||||
EMBEDDING_MODEL: str = Field(
|
||||
default="BAAI/bge-large-zh-v1.5",
|
||||
description="嵌入模型名称",
|
||||
)
|
||||
EMBEDDING_DIMENSIONS: int = Field(
|
||||
default=1024,
|
||||
description="嵌入向量维度",
|
||||
)
|
||||
|
||||
# ── OpenAI ──
|
||||
OPENAI_API_KEY: Optional[str] = Field(default=None, description="OpenAI API Key")
|
||||
OPENAI_BASE_URL: Optional[str] = Field(
|
||||
default=None,
|
||||
description="OpenAI 兼容接口地址（可用于代理或第三方兼容服务）",
|
||||
)
|
||||
|
||||
# ── 智谱 AI ──
|
||||
ZHIPU_API_KEY: Optional[str] = Field(default=None, description="智谱 AI API Key")
|
||||
|
||||
# ── 阿里云 DashScope ──
|
||||
DASHSCOPE_API_KEY: Optional[str] = Field(
|
||||
default=None, description="阿里云 DashScope API Key"
|
||||
)
|
||||
|
||||
# ── 本地 BGE 模型 ──
|
||||
LOCAL_BGE_MODEL_PATH: Optional[str] = Field(
|
||||
default=None,
|
||||
description="本地 BGE 模型路径，为空则自动从 HuggingFace 下载",
|
||||
)
|
||||
|
||||
# ──────────────────────────── OCR ────────────────────────────
|
||||
OCR_PROVIDER: OCRProvider = Field(
|
||||
default=OCRProvider.AUTO,
|
||||
description="OCR 提供商: paddleocr / aliyun / tencent / auto",
|
||||
)
|
||||
|
||||
# ── 阿里云 OCR ──
|
||||
ALIYUN_OCR_ACCESS_KEY: Optional[str] = Field(
|
||||
default=None, description="阿里云 OCR AccessKey ID"
|
||||
)
|
||||
ALIYUN_OCR_SECRET: Optional[str] = Field(
|
||||
default=None, description="阿里云 OCR AccessKey Secret"
|
||||
)
|
||||
|
||||
# ── 腾讯云 OCR ──
|
||||
TENCENT_OCR_SECRET_ID: Optional[str] = Field(
|
||||
default=None, description="腾讯云 OCR SecretId"
|
||||
)
|
||||
TENCENT_OCR_SECRET_KEY: Optional[str] = Field(
|
||||
default=None, description="腾讯云 OCR SecretKey"
|
||||
)
|
||||
|
||||
# ── MiniMax ──
|
||||
MINIMAX_API_KEY: Optional[str] = Field(default=None, description="MiniMax API Key")
|
||||
MINIMAX_BASE_URL: str = Field(
|
||||
default="https://api.minimaxi.com/v1",
|
||||
description="MiniMax API 地址",
|
||||
)
|
||||
MINIMAX_EMBEDDING_MODEL: str = Field(
|
||||
default="MiniMax-Text-Embedding",
|
||||
description="MiniMax 嵌入模型名称",
|
||||
)
|
||||
MINIMAX_CHAT_MODEL: str = Field(
|
||||
default="MiniMax-M2",
|
||||
description="MiniMax Chat 模型名称",
|
||||
)
|
||||
|
||||
# ── DeepSeek ──
|
||||
DEEPSEEK_API_KEY: Optional[str] = Field(default=None, description="DeepSeek API Key")
|
||||
DEEPSEEK_BASE_URL: str = Field(
|
||||
default="https://api.deepseek.com",
|
||||
description="DeepSeek API 地址（自部署时改为你的服务器地址）",
|
||||
)
|
||||
DEEPSEEK_OCR_MODEL: str = Field(
|
||||
default="deepseek-chat",
|
||||
description="DeepSeek OCR 使用的模型名称",
|
||||
)
|
||||
|
||||
# ── LLM（用于查询扩展、实体检测等） ──
|
||||
LLM_PROVIDER: str = Field(
|
||||
default="minimax",
|
||||
description="LLM 提供商: minimax / deepseek / openai",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 数据目录 ────────────────────────────
|
||||
DATA_DIR: str = Field(
|
||||
default="/app/data",
|
||||
description="数据根目录（ transcripts / images / exports 的父目录）",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 文本分块 ────────────────────────────
|
||||
CHUNK_SIZE: int = Field(default=300, description="文本分块大小（字符数）")
|
||||
CHUNK_OVERLAP: int = Field(default=50, description="分块重叠字符数")
|
||||
|
||||
# ──────────────────────────── 图片搜索 ────────────────────────────
|
||||
SEARCH_LIMIT: int = Field(default=3, description="图片搜索默认返回数量(1-10)")
|
||||
JUDGE_BATCH_SIZE: int = Field(default=10, description="每次LLM判断的文章数量(1-50)")
|
||||
|
||||
# ──────────────────────────── 服务端口 ────────────────────────────
|
||||
HOST: str = Field(default="0.0.0.0", description="服务监听地址")
|
||||
PORT: int = Field(default=8000, description="服务监听端口")
|
||||
DEBUG: bool = Field(default=False, description="调试模式")
|
||||
|
||||
# ──────────────────────────── CORS ────────────────────────────
|
||||
CORS_ORIGINS: list[str] = Field(
|
||||
default=["*"],
|
||||
description="允许的跨域来源列表",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 企业微信智能机器人（WebSocket 长连接） ────────────────────────────
|
||||
WEWORK_BOT_ENABLED: bool = Field(
|
||||
default=False,
|
||||
description="是否启用企业微信智能机器人（独立进程运行）",
|
||||
)
|
||||
WEWORK_BOT_ID: Optional[str] = Field(
|
||||
default=None,
|
||||
description="企业微信智能机器人 ID",
|
||||
)
|
||||
WEWORK_BOT_SECRET: Optional[str] = Field(
|
||||
default=None,
|
||||
description="企业微信智能机器人 Secret",
|
||||
)
|
||||
WEWORK_BOT_SEARCH_LIMIT: int = Field(
|
||||
default=3,
|
||||
ge=1,
|
||||
le=10,
|
||||
description="企业微信机器人搜索返回图片数量上限",
|
||||
)
|
||||
EDUBRAIN_BASE_URL: str = Field(
|
||||
default="http://localhost:8765",
|
||||
description="EduBrain 服务地址（供企业微信机器人调用搜索接口）",
|
||||
)
|
||||
|
||||
# ──────────────────────────── 属性方法 ────────────────────────────
|
||||
|
||||
@property
|
||||
def transcripts_dir(self) -> Path:
|
||||
"""直播文字稿目录"""
|
||||
return Path(self.DATA_DIR) / "transcripts"
|
||||
|
||||
@property
|
||||
def images_dir(self) -> Path:
|
||||
"""聊天截图目录"""
|
||||
return Path(self.DATA_DIR) / "images"
|
||||
|
||||
@property
|
||||
def exports_dir(self) -> Path:
|
||||
"""导出文件目录"""
|
||||
return Path(self.DATA_DIR) / "exports"
|
||||
|
||||
def ensure_dirs(self) -> None:
|
||||
"""确保所有数据目录存在"""
|
||||
for d in (self.transcripts_dir, self.images_dir, self.exports_dir):
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
# ──────────────────────────── 全局单例 ────────────────────────────
|
||||
|
||||
settings = Settings()
|
||||
107
app/database.py
Normal file
107
app/database.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
数据库连接模块
|
||||
支持 PostgreSQL（生产）和 SQLite（本地开发）两种后端
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from pathlib import Path
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import (
|
||||
AsyncSession,
|
||||
async_sessionmaker,
|
||||
create_async_engine,
|
||||
)
|
||||
from sqlalchemy.pool import NullPool
|
||||
|
||||
from app.config import settings
|
||||
|
||||
# ──────────────────────────── 判断数据库类型 ────────────────────────────
|
||||
|
||||
DB_URL = settings.DATABASE_URL
|
||||
IS_SQLITE = DB_URL.startswith("sqlite")
|
||||
|
||||
# ──────────────────────────── 异步引擎 ────────────────────────────
|
||||
|
||||
if IS_SQLITE:
|
||||
# SQLite 配置
|
||||
sqlite_path = DB_URL.replace("sqlite+aiosqlite:///", "")
|
||||
if not sqlite_path.startswith("/"):
|
||||
sqlite_path = str(Path(__file__).parent.parent / sqlite_path)
|
||||
|
||||
engine = create_async_engine(
|
||||
f"sqlite+aiosqlite:///{sqlite_path}",
|
||||
echo=settings.DEBUG,
|
||||
)
|
||||
else:
|
||||
# PostgreSQL 配置
|
||||
engine = create_async_engine(
|
||||
DB_URL,
|
||||
echo=settings.DEBUG,
|
||||
poolclass=NullPool,
|
||||
pool_pre_ping=True,
|
||||
)
|
||||
|
||||
# ──────────────────────────── 会话工厂 ────────────────────────────
|
||||
|
||||
async_session_factory = async_sessionmaker(
|
||||
bind=engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False,
|
||||
autoflush=False,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 依赖注入 ────────────────────────────
|
||||
|
||||
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
||||
"""
|
||||
FastAPI 依赖项：获取异步数据库会话。
|
||||
"""
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
yield session
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
|
||||
# ──────────────────────────── 生命周期辅助 ────────────────────────────
|
||||
|
||||
async def init_db() -> None:
|
||||
"""初始化数据库：建表"""
|
||||
from app.models.base import Base
|
||||
|
||||
async with engine.begin() as conn:
|
||||
if IS_SQLITE:
|
||||
# SQLite 启用 WAL 模式和外键约束
|
||||
await conn.execute(text("PRAGMA journal_mode=WAL"))
|
||||
await conn.execute(text("PRAGMA foreign_keys=ON"))
|
||||
# 创建所有表
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
if not IS_SQLITE:
|
||||
# PostgreSQL 需要创建扩展
|
||||
async with engine.begin() as conn:
|
||||
try:
|
||||
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS zhparser"))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
await conn.execute(text(
|
||||
"CREATE TEXT SEARCH CONFIGURATION IF NOT EXISTS chinese_zh (PARSER = zhparser)"
|
||||
))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def close_db() -> None:
|
||||
"""关闭数据库引擎"""
|
||||
await engine.dispose()
|
||||
171
app/main.py
Normal file
171
app/main.py
Normal file
@@ -0,0 +1,171 @@
|
||||
"""
|
||||
FastAPI 应用入口
|
||||
挂载 v1 API 路由、CORS 中间件、生命周期管理
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
|
||||
|
||||
from app.config import settings
|
||||
from app.database import close_db, init_db
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
"""
|
||||
应用生命周期管理：
|
||||
- 启动时：初始化数据目录、验证数据库连接、预热嵌入服务、初始化 BM25 索引
|
||||
- 关闭时：释放数据库连接
|
||||
"""
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
# ── 启动 ──
|
||||
settings.ensure_dirs()
|
||||
await init_db()
|
||||
|
||||
# 预热嵌入服务（懒加载，首次调用时初始化）
|
||||
try:
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
EmbeddingService.get_instance()
|
||||
except Exception as exc:
|
||||
_logger.warning("嵌入服务预热失败（将在首次使用时重试）: %s", exc)
|
||||
|
||||
# ── 初始化 BM25 索引 ──
|
||||
try:
|
||||
from app.services.search_engine import BM25Index
|
||||
import app.api.v1.images as images_module
|
||||
|
||||
_images_bm25 = BM25Index()
|
||||
|
||||
# 加载数据库中已有的 OCR 记录到索引
|
||||
from app.database import async_session_factory
|
||||
from app.models.base import OCRImage
|
||||
from sqlalchemy import select
|
||||
|
||||
async with async_session_factory() as session:
|
||||
result = await session.execute(
|
||||
select(OCRImage).where(OCRImage.status == "completed")
|
||||
)
|
||||
records = result.scalars().all()
|
||||
for record in records:
|
||||
if record.ocr_text:
|
||||
index_text = record.ocr_text
|
||||
if record.tags:
|
||||
try:
|
||||
tags = json.loads(record.tags)
|
||||
if tags:
|
||||
index_text += "\n" + " ".join(tags)
|
||||
except Exception:
|
||||
pass
|
||||
if record.story_summary:
|
||||
index_text += "\n" + record.story_summary
|
||||
_images_bm25.add_document(
|
||||
doc_id=record.id,
|
||||
text=index_text,
|
||||
metadata={
|
||||
"file_path": record.file_path,
|
||||
"tags": json.loads(record.tags) if record.tags else [],
|
||||
"story_summary": record.story_summary or "",
|
||||
"confidence": record.confidence,
|
||||
"provider": record.provider,
|
||||
"created_at": str(record.created_at),
|
||||
},
|
||||
)
|
||||
|
||||
images_module.bm25_index = _images_bm25
|
||||
_logger.info("BM25 索引加载完成: %d 条记录", _images_bm25.size)
|
||||
except Exception as exc:
|
||||
_logger.warning("BM25 索引初始化失败: %s", exc)
|
||||
|
||||
yield
|
||||
|
||||
# ── 关闭 ──
|
||||
await close_db()
|
||||
|
||||
|
||||
# ──────────────────────────── 创建应用 ────────────────────────────
|
||||
|
||||
app = FastAPI(
|
||||
title="EduBrain - 中文直播教育知识库",
|
||||
description="基于向量检索的中文直播教育知识库系统，支持直播文字稿导入、OCR 截图识别、语义搜索等功能。",
|
||||
version="1.2.0",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
# ──────────────────────────── CORS 中间件 ────────────────────────────
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.CORS_ORIGINS,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 全局异常处理 ────────────────────────────
|
||||
|
||||
@app.exception_handler(Exception)
|
||||
async def global_exception_handler(request: Request, exc: Exception) -> JSONResponse:
|
||||
"""全局未捕获异常处理"""
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error("未处理异常: %s", exc, exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={"detail": "服务器内部错误", "message": str(exc) if settings.DEBUG else "请稍后重试"},
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 注册路由 ────────────────────────────
|
||||
|
||||
from app.api.v1.pages import router as pages_router
|
||||
from app.api.v1.search import router as search_router
|
||||
from app.api.v1.import_export import router as import_export_router
|
||||
from app.api.v1.images import router as images_router
|
||||
from app.api.v1.settings import router as settings_router
|
||||
|
||||
app.include_router(pages_router, prefix="/api/v1/pages", tags=["知识页面"])
|
||||
app.include_router(search_router, prefix="/api/v1/search", tags=["语义搜索"])
|
||||
app.include_router(import_export_router, prefix="/api/v1/import", tags=["导入导出"])
|
||||
app.include_router(images_router, prefix="/api/v1/images", tags=["图片OCR"])
|
||||
app.include_router(settings_router, prefix="/api/v1/settings", tags=["系统设置"])
|
||||
|
||||
|
||||
# ──────────────────────────── 健康检查 ────────────────────────────
|
||||
|
||||
@app.get("/health", tags=["系统"])
|
||||
async def health_check():
|
||||
"""健康检查接口"""
|
||||
return {"status": "ok", "version": "1.2.0"}
|
||||
|
||||
|
||||
# ──────────────────────────── 图片文件访问 ────────────────────────────
|
||||
|
||||
@app.get("/data/images/{image_name}", tags=["前端"])
|
||||
async def serve_image(image_name: str):
|
||||
"""访问 images 目录下的图片文件（用于前端预览）"""
|
||||
img_path = settings.images_dir / image_name
|
||||
if not img_path.exists():
|
||||
return JSONResponse(status_code=404, content={"detail": "图片不存在"})
|
||||
return FileResponse(str(img_path))
|
||||
|
||||
|
||||
# ──────────────────────────── 前端页面 ────────────────────────────
|
||||
|
||||
@app.get("/", tags=["前端"])
|
||||
async def index():
|
||||
"""返回前端管理页面"""
|
||||
static_dir = Path(__file__).parent.parent / "static"
|
||||
index_path = static_dir / "index.html"
|
||||
if index_path.exists():
|
||||
return FileResponse(str(index_path))
|
||||
return RedirectResponse(url="/docs")
|
||||
4
app/mcp/__init__.py
Normal file
4
app/mcp/__init__.py
Normal file
@@ -0,0 +1,4 @@
|
||||
"""
|
||||
MCP (Model Context Protocol) Server 包
|
||||
提供 MCP 协议接口，允许 AI Agent 直接查询知识库
|
||||
"""
|
||||
297
app/mcp/server.py
Normal file
297
app/mcp/server.py
Normal file
@@ -0,0 +1,297 @@
|
||||
"""
|
||||
MCP Server 实现
|
||||
基于 MCP SDK 提供知识库查询工具，供 AI Agent 调用
|
||||
|
||||
工具列表:
|
||||
1. search_knowledge - 语义搜索知识库
|
||||
2. get_page - 获取知识页面详情
|
||||
3. list_courses - 列出所有课程
|
||||
4. get_page_chunks - 获取页面的分块内容
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ──────────────────────────── 创建 MCP Server ────────────────────────────
|
||||
|
||||
mcp = FastMCP(
|
||||
name="EduBrain-Knowledge",
|
||||
instructions=(
|
||||
"中文直播教育知识库助手。你可以使用以下工具来搜索和查询知识库中的内容：\n"
|
||||
"- search_knowledge: 语义搜索知识库，支持按课程、讲师、日期过滤\n"
|
||||
"- get_page: 获取指定知识页面的完整内容\n"
|
||||
"- list_courses: 列出知识库中所有课程\n"
|
||||
"- get_page_chunks: 获取页面的分块内容（用于精确定位）\n"
|
||||
"- search_images: 根据关键字搜索已 OCR 识别的聊天截图，返回图片路径"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 工具定义 ────────────────────────────
|
||||
|
||||
@mcp.tool()
|
||||
async def search_knowledge(
|
||||
query: str,
|
||||
top_k: int = 5,
|
||||
course_name: str = "",
|
||||
teacher_name: str = "",
|
||||
threshold: float = 0.5,
|
||||
) -> str:
|
||||
"""
|
||||
语义搜索知识库。
|
||||
|
||||
Args:
|
||||
query: 搜索查询文本（自然语言描述你要查找的知识点）
|
||||
top_k: 返回结果数量，默认 5
|
||||
course_name: 按课程名称过滤（可选）
|
||||
teacher_name: 按讲师名称过滤（可选）
|
||||
threshold: 相似度阈值（0-1），默认 0.5
|
||||
|
||||
Returns:
|
||||
搜索结果列表，包含标题、内容摘要和相似度分数
|
||||
"""
|
||||
from app.database import async_session_factory
|
||||
from app.schemas.search import SearchRequest
|
||||
from app.services.search_service import SearchService
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
request = SearchRequest(
|
||||
query=query,
|
||||
top_k=top_k,
|
||||
course_name=course_name or None,
|
||||
teacher_name=teacher_name or None,
|
||||
threshold=threshold,
|
||||
use_fulltext=True,
|
||||
)
|
||||
service = SearchService(session)
|
||||
response = await service.search(request)
|
||||
|
||||
if not response.results:
|
||||
return f"未找到与「{query}」相关的知识内容。"
|
||||
|
||||
# 格式化搜索结果
|
||||
output_parts = [
|
||||
f"找到 {response.total} 条相关结果（耗时 {response.elapsed_ms:.0f}ms）：\n"
|
||||
]
|
||||
for i, result in enumerate(response.results, 1):
|
||||
output_parts.append(
|
||||
f"【{i}】{result.page_title} "
|
||||
f"(相似度: {result.score:.2%})"
|
||||
)
|
||||
if result.course_name:
|
||||
output_parts.append(f" 课程: {result.course_name}")
|
||||
if result.teacher_name:
|
||||
output_parts.append(f" 讲师: {result.teacher_name}")
|
||||
# 截断内容避免过长
|
||||
content = result.content[:200]
|
||||
if len(result.content) > 200:
|
||||
content += "..."
|
||||
output_parts.append(f" 内容: {content}")
|
||||
output_parts.append("")
|
||||
|
||||
return "\n".join(output_parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 搜索失败: %s", e, exc_info=True)
|
||||
return f"搜索出错: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_page(page_id: int) -> str:
|
||||
"""
|
||||
获取指定知识页面的完整内容。
|
||||
|
||||
Args:
|
||||
page_id: 知识页面 ID
|
||||
|
||||
Returns:
|
||||
页面的完整内容，包括标题、正文、元数据
|
||||
"""
|
||||
from app.database import async_session_factory
|
||||
from app.services.page_service import PageService
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
service = PageService(session)
|
||||
page = await service.get_page(page_id)
|
||||
|
||||
if page is None:
|
||||
return f"未找到 ID 为 {page_id} 的知识页面。"
|
||||
|
||||
parts = [
|
||||
f"标题: {page.title}",
|
||||
]
|
||||
if page.course_name:
|
||||
parts.append(f"课程: {page.course_name}")
|
||||
if page.teacher_name:
|
||||
parts.append(f"讲师: {page.teacher_name}")
|
||||
if page.live_date:
|
||||
parts.append(f"直播日期: {page.live_date}")
|
||||
parts.append(f"\n{page.content}")
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 获取页面失败: %s", e, exc_info=True)
|
||||
return f"获取页面出错: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def list_courses() -> str:
|
||||
"""
|
||||
列出知识库中所有课程。
|
||||
|
||||
Returns:
|
||||
课程列表，包含课程名称、讲师、页面数量等信息
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
from app.database import async_session_factory
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
result = await session.execute(text("""
|
||||
SELECT
|
||||
course_name,
|
||||
teacher_name,
|
||||
COUNT(*) as page_count,
|
||||
MIN(live_date) as earliest_date,
|
||||
MAX(live_date) as latest_date
|
||||
FROM knowledge_pages
|
||||
WHERE course_name IS NOT NULL
|
||||
GROUP BY course_name, teacher_name
|
||||
ORDER BY page_count DESC
|
||||
"""))
|
||||
rows = result.fetchall()
|
||||
|
||||
if not rows:
|
||||
return "知识库中暂无课程数据。"
|
||||
|
||||
parts = [f"共 {len(rows)} 个课程：\n"]
|
||||
for i, row in enumerate(rows, 1):
|
||||
parts.append(
|
||||
f"{i}. {row.course_name}"
|
||||
f"（讲师: {row.teacher_name or '未知'}，"
|
||||
f"页面: {row.page_count}，"
|
||||
f"日期: {row.earliest_date or '未知'} ~ {row.latest_date or '未知'}）"
|
||||
)
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 列出课程失败: %s", e, exc_info=True)
|
||||
return f"列出课程出错: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def search_images(
|
||||
keyword: str,
|
||||
top_k: int = 10,
|
||||
) -> str:
|
||||
"""
|
||||
根据关键字搜索已 OCR 识别的聊天截图。
|
||||
|
||||
Args:
|
||||
keyword: 搜索关键字（在 OCR 识别文本中查找）
|
||||
top_k: 返回结果数量，默认 10
|
||||
|
||||
Returns:
|
||||
匹配的图片列表，包含图片路径和 OCR 文本预览
|
||||
"""
|
||||
from sqlalchemy import select
|
||||
from app.database import async_session_factory
|
||||
from app.models.base import OCRImage
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
like_pattern = f"%{keyword}%"
|
||||
result = await session.execute(
|
||||
select(OCRImage)
|
||||
.where(
|
||||
OCRImage.status == "completed",
|
||||
OCRImage.ocr_text.ilike(like_pattern),
|
||||
)
|
||||
.order_by(OCRImage.created_at.desc())
|
||||
.limit(top_k)
|
||||
)
|
||||
images = result.scalars().all()
|
||||
|
||||
if not images:
|
||||
return f"未找到包含「{keyword}」的图片。"
|
||||
|
||||
parts = [f"找到 {len(images)} 张包含「{keyword}」的图片：\n"]
|
||||
for i, img in enumerate(images, 1):
|
||||
preview = (img.ocr_text[:100] + "...") if img.ocr_text and len(img.ocr_text) > 100 else (img.ocr_text or "")
|
||||
parts.append(
|
||||
f"【{i}】图片路径: {img.file_path}\n"
|
||||
f" OCR 预览: {preview}\n"
|
||||
f" 置信度: {img.confidence:.0%} | 识别方式: {img.provider}\n"
|
||||
)
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 图片搜索失败: %s", e, exc_info=True)
|
||||
return f"图片搜索出错: {str(e)}"
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_page_chunks(page_id: int) -> str:
|
||||
"""
|
||||
获取知识页面的分块内容（用于精确定位知识点）。
|
||||
|
||||
Args:
|
||||
page_id: 知识页面 ID
|
||||
|
||||
Returns:
|
||||
该页面的所有分块内容列表
|
||||
"""
|
||||
from app.database import async_session_factory
|
||||
from app.services.page_service import PageService
|
||||
|
||||
async with async_session_factory() as session:
|
||||
try:
|
||||
service = PageService(session)
|
||||
chunks = await service.get_page_chunks(page_id)
|
||||
|
||||
if not chunks:
|
||||
return f"页面 {page_id} 没有分块数据。"
|
||||
|
||||
parts = [f"页面 {page_id} 共 {len(chunks)} 个分块：\n"]
|
||||
for i, chunk in enumerate(chunks, 1):
|
||||
content = chunk.content[:300]
|
||||
if len(chunk.content) > 300:
|
||||
content += "..."
|
||||
parts.append(f"【分块 {i}】(ID: {chunk.id})\n{content}\n")
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("MCP 获取分块失败: %s", e, exc_info=True)
|
||||
return f"获取分块出错: {str(e)}"
|
||||
|
||||
|
||||
# ──────────────────────────── 启动入口 ────────────────────────────
|
||||
|
||||
def run_mcp_server(transport: str = "stdio", port: int = 8001):
|
||||
"""
|
||||
启动 MCP Server。
|
||||
|
||||
Args:
|
||||
transport: 传输协议，支持 "stdio"（默认）或 "sse"
|
||||
port: SSE 模式下的端口号
|
||||
"""
|
||||
if transport == "sse":
|
||||
mcp.run(transport="sse", port=port)
|
||||
else:
|
||||
mcp.run(transport="stdio")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_mcp_server()
|
||||
7
app/models/__init__.py
Normal file
7
app/models/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
SQLAlchemy ORM 模型包
|
||||
"""
|
||||
|
||||
from app.models.base import Base, TimestampMixin
|
||||
|
||||
__all__ = ["Base", "TimestampMixin"]
|
||||
123
app/models/base.py
Normal file
123
app/models/base.py
Normal file
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
SQLAlchemy 基础模型 + 通用 Mixin
|
||||
提供 DeclarativeBase、时间戳 Mixin 和公共序列化方法
|
||||
兼容 PostgreSQL（pgvector）和 SQLite（本地开发）
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict
|
||||
|
||||
from sqlalchemy import DateTime, Float, Integer, String, Text, func
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||
|
||||
from app.database import IS_SQLITE
|
||||
|
||||
# SQLite 下不使用 pgvector，用 Text 存储向量（JSON 格式）
|
||||
if IS_SQLITE:
|
||||
VectorType = Text
|
||||
else:
|
||||
from pgvector.sqlalchemy import Vector
|
||||
VectorType = Vector
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
"""所有 ORM 模型的基类"""
|
||||
pass
|
||||
|
||||
|
||||
class TimestampMixin:
|
||||
"""
|
||||
时间戳 Mixin，为模型自动添加 created_at / updated_at 字段。
|
||||
"""
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
comment="创建时间",
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
server_default=func.now(),
|
||||
onupdate=func.now(),
|
||||
comment="更新时间",
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""将模型实例转换为字典"""
|
||||
result: Dict[str, Any] = {}
|
||||
for column in self.__table__.columns: # type: ignore[attr-defined]
|
||||
value = getattr(self, column.name, None)
|
||||
if isinstance(value, datetime):
|
||||
value = value.isoformat()
|
||||
elif hasattr(value, "tolist"):
|
||||
value = value.tolist()
|
||||
result[column.name] = value
|
||||
return result
|
||||
|
||||
|
||||
# ──────────────────────────── 业务模型 ────────────────────────────
|
||||
|
||||
class KnowledgePage(Base, TimestampMixin):
|
||||
"""知识页面模型"""
|
||||
__tablename__ = "knowledge_pages"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
title: Mapped[str] = mapped_column(String(500), nullable=False, comment="页面标题")
|
||||
content: Mapped[str] = mapped_column(Text, nullable=False, comment="页面正文内容")
|
||||
source_file: Mapped[str | None] = mapped_column(String(500), nullable=True, comment="来源文件名")
|
||||
course_name: Mapped[str | None] = mapped_column(String(200), nullable=True, comment="课程名称")
|
||||
teacher_name: Mapped[str | None] = mapped_column(String(100), nullable=True, comment="讲师名称")
|
||||
live_date: Mapped[str | None] = mapped_column(String(20), nullable=True, comment="直播日期")
|
||||
page_number: Mapped[int | None] = mapped_column(Integer, nullable=True, comment="原始页码")
|
||||
metadata_json: Mapped[str | None] = mapped_column(Text, nullable=True, comment="额外元数据")
|
||||
|
||||
|
||||
class KnowledgeChunk(Base, TimestampMixin):
|
||||
"""知识分块模型"""
|
||||
__tablename__ = "knowledge_chunks"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
page_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True, comment="关联页面 ID")
|
||||
chunk_index: Mapped[int] = mapped_column(Integer, nullable=False, comment="分块序号")
|
||||
content: Mapped[str] = mapped_column(Text, nullable=False, comment="分块文本内容")
|
||||
embedding: Mapped[list | None] = mapped_column(
|
||||
VectorType(768) if not IS_SQLITE else Text,
|
||||
nullable=True,
|
||||
comment="嵌入向量",
|
||||
)
|
||||
search_vector: Mapped[Any | None] = mapped_column(
|
||||
Text, nullable=True, comment="全文搜索向量（仅 PG）",
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result = super().to_dict()
|
||||
result.pop("embedding", None)
|
||||
result.pop("search_vector", None)
|
||||
return result
|
||||
|
||||
|
||||
class OCRImage(Base, TimestampMixin):
|
||||
"""OCR 图片记录模型"""
|
||||
__tablename__ = "ocr_images"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
file_path: Mapped[str] = mapped_column(String(500), nullable=False, comment="图片文件路径")
|
||||
ocr_text: Mapped[str | None] = mapped_column(Text, nullable=True, comment="OCR 识别全文")
|
||||
confidence: Mapped[float | None] = mapped_column(Float, nullable=True, comment="平均置信度")
|
||||
provider: Mapped[str | None] = mapped_column(String(50), nullable=True, comment="OCR 提供商")
|
||||
status: Mapped[str] = mapped_column(String(20), default="pending", comment="处理状态")
|
||||
error_message: Mapped[str | None] = mapped_column(Text, nullable=True, comment="错误信息")
|
||||
tags: Mapped[str | None] = mapped_column(Text, nullable=True, comment="LLM 提取的标签，JSON 数组格式")
|
||||
story_summary: Mapped[str | None] = mapped_column(Text, nullable=True, comment="Qwen3-8B 提炼的故事摘要")
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result = super().to_dict()
|
||||
# 将 tags 从 JSON 字符串解析为列表返回
|
||||
if result.get("tags") and isinstance(result["tags"], str):
|
||||
try:
|
||||
import json
|
||||
result["tags"] = json.loads(result["tags"])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
return result
|
||||
32
app/schemas/__init__.py
Normal file
32
app/schemas/__init__.py
Normal file
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
Pydantic Schemas 包
|
||||
用于 API 请求/响应的数据校验与序列化
|
||||
"""
|
||||
|
||||
from app.schemas.page import (
|
||||
PageCreate,
|
||||
PageResponse,
|
||||
PageListResponse,
|
||||
)
|
||||
from app.schemas.search import (
|
||||
SearchRequest,
|
||||
SearchResponse,
|
||||
SearchResult,
|
||||
)
|
||||
from app.schemas.ocr import (
|
||||
OCRResultSchema,
|
||||
OCRBlockSchema,
|
||||
ImageUploadResponse,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"PageCreate",
|
||||
"PageResponse",
|
||||
"PageListResponse",
|
||||
"SearchRequest",
|
||||
"SearchResponse",
|
||||
"SearchResult",
|
||||
"OCRResultSchema",
|
||||
"OCRBlockSchema",
|
||||
"ImageUploadResponse",
|
||||
]
|
||||
47
app/schemas/ocr.py
Normal file
47
app/schemas/ocr.py
Normal file
@@ -0,0 +1,47 @@
|
||||
"""
|
||||
OCR 相关的 Pydantic Schema
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class OCRBlockSchema(BaseModel):
|
||||
"""OCR 识别的文本块"""
|
||||
text: str = Field(..., description="文本块内容")
|
||||
bbox: Optional[list[float]] = Field(None, description="边界框 [x1, y1, x2, y2]")
|
||||
confidence: float = Field(..., description="置信度 (0-1)")
|
||||
|
||||
|
||||
class OCRResultSchema(BaseModel):
|
||||
"""OCR 识别结果"""
|
||||
text: str = Field(..., description="识别全文")
|
||||
blocks: list[OCRBlockSchema] = Field(default_factory=list, description="文本块列表")
|
||||
confidence: float = Field(..., description="平均置信度 (0-1)")
|
||||
provider: str = Field(..., description="使用的 OCR 提供商")
|
||||
|
||||
|
||||
class ImageUploadResponse(BaseModel):
|
||||
"""图片上传响应"""
|
||||
id: int
|
||||
file_path: str
|
||||
status: str
|
||||
created_at: datetime
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class ImageOCRResponse(BaseModel):
|
||||
"""图片 OCR 结果响应"""
|
||||
id: int
|
||||
file_path: str
|
||||
ocr_text: Optional[str] = None
|
||||
confidence: Optional[float] = None
|
||||
provider: Optional[str] = None
|
||||
status: str
|
||||
error_message: Optional[str] = None
|
||||
created_at: datetime
|
||||
model_config = {"from_attributes": True}
|
||||
63
app/schemas/page.py
Normal file
63
app/schemas/page.py
Normal file
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
知识页面相关的 Pydantic Schema
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class PageCreate(BaseModel):
|
||||
"""创建知识页面请求"""
|
||||
title: str = Field(..., max_length=500, description="页面标题")
|
||||
content: str = Field(..., description="页面正文内容")
|
||||
source_file: Optional[str] = Field(None, max_length=500, description="来源文件名")
|
||||
course_name: Optional[str] = Field(None, max_length=200, description="课程名称")
|
||||
teacher_name: Optional[str] = Field(None, max_length=100, description="讲师名称")
|
||||
live_date: Optional[str] = Field(None, max_length=20, description="直播日期 (YYYY-MM-DD)")
|
||||
page_number: Optional[int] = Field(None, description="原始页码")
|
||||
metadata_json: Optional[str] = Field(None, description="额外元数据 (JSON)")
|
||||
|
||||
|
||||
class PageUpdate(BaseModel):
|
||||
"""更新知识页面请求（所有字段可选）"""
|
||||
title: Optional[str] = Field(None, max_length=500, description="页面标题")
|
||||
content: Optional[str] = Field(None, description="页面正文内容")
|
||||
course_name: Optional[str] = Field(None, max_length=200, description="课程名称")
|
||||
teacher_name: Optional[str] = Field(None, max_length=100, description="讲师名称")
|
||||
live_date: Optional[str] = Field(None, max_length=20, description="直播日期")
|
||||
page_number: Optional[int] = Field(None, description="原始页码")
|
||||
metadata_json: Optional[str] = Field(None, description="额外元数据")
|
||||
|
||||
|
||||
class PageResponse(BaseModel):
|
||||
"""知识页面响应"""
|
||||
id: int
|
||||
title: str
|
||||
content: str
|
||||
source_file: Optional[str] = None
|
||||
course_name: Optional[str] = None
|
||||
teacher_name: Optional[str] = None
|
||||
live_date: Optional[str] = None
|
||||
page_number: Optional[int] = None
|
||||
metadata_json: Optional[str] = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
model_config = {"from_attributes": True}
|
||||
|
||||
|
||||
class PageListResponse(BaseModel):
|
||||
"""知识页面列表响应（带分页）"""
|
||||
total: int = Field(..., description="总记录数")
|
||||
page: int = Field(..., description="当前页码")
|
||||
page_size: int = Field(..., description="每页数量")
|
||||
items: list[PageResponse] = Field(default_factory=list, description="页面列表")
|
||||
|
||||
|
||||
class PageDetailResponse(PageResponse):
|
||||
"""知识页面详情响应（包含分块信息）"""
|
||||
chunks: list[dict[str, Any]] = Field(default_factory=list, description="关联的分块列表")
|
||||
42
app/schemas/search.py
Normal file
42
app/schemas/search.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
语义搜索相关的 Pydantic Schema
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
"""语义搜索请求"""
|
||||
query: str = Field(..., min_length=1, max_length=1000, description="搜索查询文本")
|
||||
top_k: int = Field(default=10, ge=1, le=100, description="返回结果数量")
|
||||
course_name: Optional[str] = Field(None, description="按课程名称过滤")
|
||||
teacher_name: Optional[str] = Field(None, description="按讲师名称过滤")
|
||||
live_date_from: Optional[str] = Field(None, description="直播日期起始 (YYYY-MM-DD)")
|
||||
live_date_to: Optional[str] = Field(None, description="直播日期结束 (YYYY-MM-DD)")
|
||||
threshold: float = Field(default=0.5, ge=0.0, le=1.0, description="相似度阈值")
|
||||
use_fulltext: bool = Field(default=True, description="是否同时使用全文搜索混合排序")
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""单条搜索结果"""
|
||||
chunk_id: int = Field(..., description="分块 ID")
|
||||
page_id: int = Field(..., description="知识页面 ID")
|
||||
page_title: str = Field(..., description="页面标题")
|
||||
content: str = Field(..., description="分块文本内容")
|
||||
score: float = Field(..., description="相似度分数")
|
||||
course_name: Optional[str] = None
|
||||
teacher_name: Optional[str] = None
|
||||
live_date: Optional[str] = None
|
||||
highlight: Optional[str] = Field(None, description="高亮摘要")
|
||||
|
||||
|
||||
class SearchResponse(BaseModel):
|
||||
"""搜索结果响应"""
|
||||
query: str = Field(..., description="原始查询文本")
|
||||
total: int = Field(..., description="结果总数")
|
||||
results: list[SearchResult] = Field(default_factory=list, description="搜索结果列表")
|
||||
elapsed_ms: float = Field(..., description="搜索耗时（毫秒）")
|
||||
17
app/services/__init__.py
Normal file
17
app/services/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
业务逻辑服务层包
|
||||
"""
|
||||
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
from app.services.ocr_service import OCRService
|
||||
from app.services.search_service import SearchService
|
||||
from app.services.import_service import ImportService
|
||||
from app.services.page_service import PageService
|
||||
|
||||
__all__ = [
|
||||
"EmbeddingService",
|
||||
"OCRService",
|
||||
"SearchService",
|
||||
"ImportService",
|
||||
"PageService",
|
||||
]
|
||||
328
app/services/embedding_service.py
Normal file
328
app/services/embedding_service.py
Normal file
@@ -0,0 +1,328 @@
|
||||
"""
|
||||
嵌入服务模块
|
||||
支持多种嵌入模型提供商：OpenAI / 智谱 / 阿里云 DashScope / 本地 BGE
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from app.config import EmbeddingProvider, settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ──────────────────────────── 抽象基类 ────────────────────────────
|
||||
|
||||
class EmbeddingProviderBase(abc.ABC):
|
||||
"""嵌入模型提供商抽象基类"""
|
||||
|
||||
@abc.abstractmethod
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
批量生成文本嵌入向量。
|
||||
|
||||
Args:
|
||||
texts: 待嵌入的文本列表
|
||||
|
||||
Returns:
|
||||
嵌入向量列表，每个向量是 float 的列表
|
||||
|
||||
Raises:
|
||||
Exception: 嵌入生成失败时抛出异常
|
||||
"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
"""生成单条文本的嵌入向量"""
|
||||
...
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def dimension(self) -> int:
|
||||
"""嵌入向量维度"""
|
||||
...
|
||||
|
||||
|
||||
# ──────────────────────────── OpenAI 实现 ────────────────────────────
|
||||
|
||||
class OpenAIEmbedding(EmbeddingProviderBase):
|
||||
"""OpenAI 嵌入模型（兼容 OpenAI API 格式的服务）"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=settings.OPENAI_API_KEY or "EMPTY",
|
||||
base_url=settings.OPENAI_BASE_URL,
|
||||
)
|
||||
self._model = settings.EMBEDDING_MODEL or "text-embedding-3-small"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 OpenAI 兼容接口批量生成嵌入"""
|
||||
# OpenAI 接口单次最多 2048 条，自动分批
|
||||
batch_size = 2048
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
response = await self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
dimensions=self.dimension if self.dimension <= 3072 else None,
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 智谱 AI 实现 ────────────────────────────
|
||||
|
||||
class ZhipuEmbedding(EmbeddingProviderBase):
|
||||
"""智谱 AI 嵌入模型"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from zhipuai import ZhipuAI
|
||||
|
||||
self._client = ZhipuAI(api_key=settings.ZHIPU_API_KEY)
|
||||
self._model = settings.EMBEDDING_MODEL or "embedding-3"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用智谱 AI 接口批量生成嵌入"""
|
||||
import asyncio
|
||||
|
||||
# 智谱 API 单次最多 16 条
|
||||
batch_size = 16
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
# 智谱 SDK 是同步的，用线程池包装
|
||||
loop = asyncio.get_event_loop()
|
||||
response = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
),
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 阿里云 DashScope 实现 ────────────────────────────
|
||||
|
||||
class DashscopeEmbedding(EmbeddingProviderBase):
|
||||
"""阿里云 DashScope 嵌入模型"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
import dashscope
|
||||
dashscope.api_key = settings.DASHSCOPE_API_KEY
|
||||
self._model = settings.EMBEDDING_MODEL or "text-embedding-v3"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 DashScope 接口批量生成嵌入"""
|
||||
import asyncio
|
||||
from dashscope import TextEmbedding
|
||||
|
||||
batch_size = 25 # DashScope 单次最多 25 条
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _call() -> list:
|
||||
resp = TextEmbedding.call(
|
||||
model=self._model,
|
||||
input=batch,
|
||||
dimension=self.dimension,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(f"DashScope 嵌入失败: {resp.code} - {resp.message}")
|
||||
return [item["embedding"] for item in resp.output["embeddings"]]
|
||||
|
||||
batch_embeddings = await loop.run_in_executor(None, _call)
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 本地 BGE 实现 ────────────────────────────
|
||||
|
||||
class LocalBGEEmbedding(EmbeddingProviderBase):
|
||||
"""
|
||||
本地 BGE 嵌入模型
|
||||
使用 sentence-transformers 加载模型，在 CPU/GPU 上本地推理
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._model = None
|
||||
self._dim = settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
def _load_model(self):
|
||||
"""延迟加载模型（首次调用时初始化）"""
|
||||
if self._model is not None:
|
||||
return
|
||||
|
||||
logger.info("正在加载本地 BGE 嵌入模型: %s", settings.EMBEDDING_MODEL)
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
model_path = settings.LOCAL_BGE_MODEL_PATH or settings.EMBEDDING_MODEL
|
||||
self._model = SentenceTransformer(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
# 获取实际维度
|
||||
test_embedding = self._model.encode(["测试"])
|
||||
self._dim = len(test_embedding[0])
|
||||
logger.info("BGE 模型加载完成，维度: %d", self._dim)
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return self._dim
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""使用本地模型批量生成嵌入"""
|
||||
import asyncio
|
||||
|
||||
self._load_model()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _encode() -> List[List[float]]:
|
||||
embeddings = self._model.encode(texts, normalize_embeddings=True)
|
||||
return embeddings.tolist()
|
||||
|
||||
return await loop.run_in_executor(None, _encode)
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── MiniMax 实现 ────────────────────────────
|
||||
|
||||
class MiniMaxEmbedding(EmbeddingProviderBase):
|
||||
"""
|
||||
MiniMax 嵌入模型
|
||||
通过 MiniMax 的 OpenAI 兼容接口调用嵌入 API
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=settings.MINIMAX_API_KEY or "EMPTY",
|
||||
base_url=settings.MINIMAX_BASE_URL,
|
||||
)
|
||||
self._model = settings.MINIMAX_EMBEDDING_MODEL
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return settings.EMBEDDING_DIMENSIONS
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""调用 MiniMax 兼容接口批量生成嵌入"""
|
||||
batch_size = 64
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
response = await self._client.embeddings.create(
|
||||
input=batch,
|
||||
model=self._model,
|
||||
)
|
||||
batch_embeddings = [item.embedding for item in response.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
async def embed_single(self, text: str) -> List[float]:
|
||||
results = await self.embed_batch([text])
|
||||
return results[0]
|
||||
|
||||
|
||||
# ──────────────────────────── 工厂类 ────────────────────────────
|
||||
|
||||
class EmbeddingService:
|
||||
"""
|
||||
嵌入服务工厂类
|
||||
根据配置自动选择嵌入模型提供商，提供全局单例访问
|
||||
"""
|
||||
|
||||
_instance: Optional[EmbeddingProviderBase] = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> EmbeddingProviderBase:
|
||||
"""获取嵌入服务单例实例"""
|
||||
if cls._instance is not None:
|
||||
return cls._instance
|
||||
|
||||
provider_map = {
|
||||
EmbeddingProvider.OPENAI: OpenAIEmbedding,
|
||||
EmbeddingProvider.ZHIPU: ZhipuEmbedding,
|
||||
EmbeddingProvider.DASHSCOPE: DashscopeEmbedding,
|
||||
EmbeddingProvider.LOCAL_BGE: LocalBGEEmbedding,
|
||||
EmbeddingProvider.MINIMAX: MiniMaxEmbedding,
|
||||
}
|
||||
|
||||
provider_cls = provider_map.get(settings.EMBEDDING_PROVIDER)
|
||||
if provider_cls is None:
|
||||
raise ValueError(f"不支持的嵌入模型提供商: {settings.EMBEDDING_PROVIDER}")
|
||||
|
||||
cls._instance = provider_cls()
|
||||
logger.info(
|
||||
"嵌入服务初始化完成: provider=%s, model=%s",
|
||||
settings.EMBEDDING_PROVIDER.value,
|
||||
settings.EMBEDDING_MODEL,
|
||||
)
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
async def embed_batch(cls, texts: List[str]) -> List[List[float]]:
|
||||
"""便捷方法：批量生成嵌入"""
|
||||
instance = cls.get_instance()
|
||||
return await instance.embed_batch(texts)
|
||||
|
||||
@classmethod
|
||||
async def embed_single(cls, text: str) -> List[float]:
|
||||
"""便捷方法：生成单条嵌入"""
|
||||
instance = cls.get_instance()
|
||||
return await instance.embed_single(text)
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""重置单例（用于测试或切换配置）"""
|
||||
cls._instance = None
|
||||
452
app/services/import_service.py
Normal file
452
app/services/import_service.py
Normal file
@@ -0,0 +1,452 @@
|
||||
"""
|
||||
导入服务
|
||||
负责从 Markdown / 纯文本 / Word 文件导入知识页面，并进行文本分块和向量化
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.config import settings
|
||||
from app.database import IS_SQLITE
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _parse_frontmatter(content: str) -> tuple[dict, str]:
|
||||
"""
|
||||
手动解析 YAML frontmatter。
|
||||
格式：--- 开头和结尾包围的 YAML 块。
|
||||
"""
|
||||
import yaml
|
||||
|
||||
if not content.startswith("---"):
|
||||
return {}, content
|
||||
|
||||
parts = content.split("---", 2)
|
||||
if len(parts) < 3:
|
||||
return {}, content
|
||||
|
||||
try:
|
||||
metadata = yaml.safe_load(parts[1]) or {}
|
||||
except Exception:
|
||||
metadata = {}
|
||||
|
||||
body = parts[2].strip()
|
||||
return metadata, body
|
||||
|
||||
|
||||
class ImportService:
|
||||
"""知识导入服务"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self.db = db
|
||||
|
||||
async def import_file(
|
||||
self,
|
||||
file_path: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
从文件导入知识内容。
|
||||
|
||||
支持的格式:
|
||||
- Markdown (.md): 解析 frontmatter 元数据，按标题分页
|
||||
- 纯文本 (.txt): 整体作为一个页面
|
||||
- Word 文档 (.docx): 提取文本，按标题分页
|
||||
|
||||
Args:
|
||||
file_path: 文件路径
|
||||
course_name: 课程名称（覆盖 frontmatter）
|
||||
teacher_name: 讲师名称（覆盖 frontmatter）
|
||||
live_date: 直播日期（覆盖 frontmatter）
|
||||
|
||||
Returns:
|
||||
导入结果统计
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"文件不存在: {file_path}")
|
||||
|
||||
suffix = path.suffix.lower()
|
||||
|
||||
if suffix == ".md":
|
||||
content = path.read_text(encoding="utf-8")
|
||||
return await self._import_markdown(
|
||||
content, path.name, course_name, teacher_name, live_date
|
||||
)
|
||||
elif suffix == ".txt":
|
||||
content = path.read_text(encoding="utf-8")
|
||||
return await self._import_text(
|
||||
content, path.name, course_name, teacher_name, live_date
|
||||
)
|
||||
elif suffix == ".docx":
|
||||
return await self._import_docx(
|
||||
str(path), path.name, course_name, teacher_name, live_date
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"不支持的文件格式: {suffix}，仅支持 .md、.txt 和 .docx")
|
||||
|
||||
async def _import_markdown(
|
||||
self,
|
||||
content: str,
|
||||
filename: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""导入 Markdown 文件，解析 frontmatter 并按标题分页"""
|
||||
# 解析 frontmatter
|
||||
fm_data, body = _parse_frontmatter(content)
|
||||
|
||||
# frontmatter 中的元数据可作为默认值
|
||||
fm_course = fm_data.get("course", fm_data.get("course_name"))
|
||||
fm_teacher = fm_data.get("teacher", fm_data.get("teacher_name"))
|
||||
fm_date = fm_data.get("date", fm_data.get("live_date"))
|
||||
|
||||
final_course = course_name or fm_course
|
||||
final_teacher = teacher_name or fm_teacher
|
||||
final_date = live_date or (str(fm_date) if fm_date else None)
|
||||
|
||||
# 按二级标题（##）拆分页面
|
||||
sections = self._split_markdown_sections(body)
|
||||
|
||||
if not sections:
|
||||
# 没有二级标题，整体作为一个页面
|
||||
sections = [{"title": Path(filename).stem, "content": body}]
|
||||
|
||||
imported_pages = 0
|
||||
imported_chunks = 0
|
||||
|
||||
for idx, section in enumerate(sections):
|
||||
page_number = idx + 1
|
||||
|
||||
# 插入知识页面
|
||||
page_id = await self._insert_page(
|
||||
title=section["title"],
|
||||
content=section["content"],
|
||||
source_file=filename,
|
||||
course_name=final_course,
|
||||
teacher_name=final_teacher,
|
||||
live_date=final_date,
|
||||
page_number=page_number,
|
||||
)
|
||||
imported_pages += 1
|
||||
|
||||
# 分块并向量化
|
||||
chunks = self._chunk_text(section["content"])
|
||||
chunk_count = await self._insert_chunks(page_id, chunks)
|
||||
imported_chunks += chunk_count
|
||||
|
||||
logger.info(
|
||||
"Markdown 导入完成: file=%s, pages=%d, chunks=%d",
|
||||
filename, imported_pages, imported_chunks,
|
||||
)
|
||||
|
||||
return {
|
||||
"file": filename,
|
||||
"pages": imported_pages,
|
||||
"chunks": imported_chunks,
|
||||
}
|
||||
|
||||
async def _import_text(
|
||||
self,
|
||||
content: str,
|
||||
filename: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""导入纯文本文件，整体作为一个页面"""
|
||||
title = Path(filename).stem
|
||||
|
||||
page_id = await self._insert_page(
|
||||
title=title,
|
||||
content=content,
|
||||
source_file=filename,
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
live_date=live_date,
|
||||
page_number=1,
|
||||
)
|
||||
|
||||
chunks = self._chunk_text(content)
|
||||
chunk_count = await self._insert_chunks(page_id, chunks)
|
||||
|
||||
logger.info(
|
||||
"文本导入完成: file=%s, chunks=%d",
|
||||
filename, chunk_count,
|
||||
)
|
||||
|
||||
return {
|
||||
"file": filename,
|
||||
"pages": 1,
|
||||
"chunks": chunk_count,
|
||||
}
|
||||
|
||||
async def _import_docx(
|
||||
self,
|
||||
file_path: str,
|
||||
filename: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
导入 Word 文档（.docx）。
|
||||
提取所有段落文本，按标题（Heading 1/2）拆分为多个知识页面。
|
||||
"""
|
||||
from docx import Document
|
||||
|
||||
doc = Document(file_path)
|
||||
|
||||
# 提取所有段落，保留标题层级信息
|
||||
sections: List[dict] = [] # [{"title": "...", "content": "...", "level": int}]
|
||||
current_title = Path(filename).stem
|
||||
current_content: List[str] = []
|
||||
current_level = 0
|
||||
|
||||
for para in doc.paragraphs:
|
||||
style_name = (para.style.name or "").lower()
|
||||
|
||||
# 判断是否为标题段落
|
||||
if style_name.startswith("heading"):
|
||||
try:
|
||||
level = int(style_name.replace("heading", "").strip())
|
||||
except ValueError:
|
||||
level = 1
|
||||
|
||||
# 保存上一个段落
|
||||
if current_content:
|
||||
text = "\n\n".join(current_content).strip()
|
||||
if text:
|
||||
sections.append({
|
||||
"title": current_title,
|
||||
"content": text,
|
||||
"level": current_level,
|
||||
})
|
||||
|
||||
# 开始新段落
|
||||
current_title = para.text.strip() or f"第 {len(sections) + 1} 节"
|
||||
current_content = []
|
||||
current_level = level
|
||||
else:
|
||||
text = para.text.strip()
|
||||
if text:
|
||||
current_content.append(text)
|
||||
|
||||
# 保存最后一个段落
|
||||
if current_content:
|
||||
text = "\n\n".join(current_content).strip()
|
||||
if text:
|
||||
sections.append({
|
||||
"title": current_title,
|
||||
"content": text,
|
||||
"level": current_level,
|
||||
})
|
||||
|
||||
if not sections:
|
||||
raise ValueError(f"Word 文档内容为空: {filename}")
|
||||
|
||||
# 如果只有一个段落且没有标题，整体作为一个页面
|
||||
if len(sections) == 1 and sections[0]["level"] == 0:
|
||||
sections[0]["title"] = Path(filename).stem
|
||||
|
||||
imported_pages = 0
|
||||
imported_chunks = 0
|
||||
|
||||
for idx, section in enumerate(sections):
|
||||
page_id = await self._insert_page(
|
||||
title=section["title"],
|
||||
content=section["content"],
|
||||
source_file=filename,
|
||||
course_name=course_name,
|
||||
teacher_name=teacher_name,
|
||||
live_date=live_date,
|
||||
page_number=idx + 1,
|
||||
)
|
||||
imported_pages += 1
|
||||
|
||||
chunks = self._chunk_text(section["content"])
|
||||
chunk_count = await self._insert_chunks(page_id, chunks)
|
||||
imported_chunks += chunk_count
|
||||
|
||||
logger.info(
|
||||
"Word 文档导入完成: file=%s, pages=%d, chunks=%d",
|
||||
filename, imported_pages, imported_chunks,
|
||||
)
|
||||
|
||||
return {
|
||||
"file": filename,
|
||||
"pages": imported_pages,
|
||||
"chunks": imported_chunks,
|
||||
}
|
||||
|
||||
async def _insert_page(
|
||||
self,
|
||||
title: str,
|
||||
content: str,
|
||||
source_file: str,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
live_date: Optional[str] = None,
|
||||
page_number: Optional[int] = None,
|
||||
) -> int:
|
||||
"""插入知识页面记录，返回页面 ID"""
|
||||
sql = text("""
|
||||
INSERT INTO knowledge_pages (title, content, source_file, course_name, teacher_name, live_date, page_number)
|
||||
VALUES (:title, :content, :source_file, :course_name, :teacher_name, :live_date, :page_number)
|
||||
RETURNING id
|
||||
""")
|
||||
result = await self.db.execute(sql, {
|
||||
"title": title,
|
||||
"content": content,
|
||||
"source_file": source_file,
|
||||
"course_name": course_name,
|
||||
"teacher_name": teacher_name,
|
||||
"live_date": live_date,
|
||||
"page_number": page_number,
|
||||
})
|
||||
row = result.fetchone()
|
||||
await self.db.flush()
|
||||
return row.id
|
||||
|
||||
async def _insert_chunks(self, page_id: int, chunks: List[str]) -> int:
|
||||
"""
|
||||
批量插入分块记录并生成嵌入向量。
|
||||
|
||||
Args:
|
||||
page_id: 关联的知识页面 ID
|
||||
chunks: 分块文本列表
|
||||
|
||||
Returns:
|
||||
插入的分块数量
|
||||
"""
|
||||
if not chunks:
|
||||
return 0
|
||||
|
||||
# 批量生成嵌入向量
|
||||
try:
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
embeddings = await EmbeddingService.embed_batch(chunks)
|
||||
except Exception as exc:
|
||||
logger.error("嵌入向量生成失败，分块将不包含向量: %s", exc)
|
||||
embeddings = [None] * len(chunks)
|
||||
|
||||
# 批量插入
|
||||
for idx, (chunk_text, embedding) in enumerate(zip(chunks, embeddings)):
|
||||
if IS_SQLITE:
|
||||
# SQLite: 向量存为 JSON 文本
|
||||
import json as _json
|
||||
embedding_str = _json.dumps(embedding) if embedding else None
|
||||
sql = text("""
|
||||
INSERT INTO knowledge_chunks (page_id, chunk_index, content, embedding)
|
||||
VALUES (:page_id, :chunk_index, :content, :embedding)
|
||||
""")
|
||||
else:
|
||||
# PostgreSQL: 使用 pgvector 类型
|
||||
embedding_str = "[" + ",".join(str(x) for x in embedding) + "]" if embedding else None
|
||||
sql = text("""
|
||||
INSERT INTO knowledge_chunks (page_id, chunk_index, content, embedding)
|
||||
VALUES (:page_id, :chunk_index, :content, :embedding::vector)
|
||||
""")
|
||||
await self.db.execute(sql, {
|
||||
"page_id": page_id,
|
||||
"chunk_index": idx,
|
||||
"content": chunk_text,
|
||||
"embedding": embedding_str,
|
||||
})
|
||||
|
||||
# 更新全文搜索向量（仅 PostgreSQL）
|
||||
if not IS_SQLITE:
|
||||
await self.db.execute(text("""
|
||||
UPDATE knowledge_chunks
|
||||
SET search_vector = to_tsvector('chinese_zh', content)
|
||||
WHERE page_id = :page_id
|
||||
"""), {"page_id": page_id})
|
||||
|
||||
await self.db.flush()
|
||||
return len(chunks)
|
||||
|
||||
@staticmethod
|
||||
def _split_markdown_sections(body: str) -> List[dict]:
|
||||
"""
|
||||
按二级标题（##）拆分 Markdown 内容为多个段落。
|
||||
|
||||
Returns:
|
||||
[{"title": "...", "content": "..."}, ...]
|
||||
"""
|
||||
sections: List[dict] = []
|
||||
# 匹配 ## 开头的标题
|
||||
pattern = re.compile(r"^##\s+(.+)$", re.MULTILINE)
|
||||
matches = list(pattern.finditer(body))
|
||||
|
||||
if not matches:
|
||||
return []
|
||||
|
||||
for i, match in enumerate(matches):
|
||||
title = match.group(1).strip()
|
||||
start = match.end()
|
||||
end = matches[i + 1].start() if i + 1 < len(matches) else len(body)
|
||||
content = body[start:end].strip()
|
||||
if content:
|
||||
sections.append({"title": title, "content": content})
|
||||
|
||||
return sections
|
||||
|
||||
@staticmethod
|
||||
def _chunk_text(text: str) -> List[str]:
|
||||
"""
|
||||
将文本按固定大小分块，保留重叠部分。
|
||||
|
||||
使用简单的字符数分块策略，按段落边界切分以保持语义完整性。
|
||||
"""
|
||||
chunk_size = settings.CHUNK_SIZE
|
||||
overlap = settings.CHUNK_OVERLAP
|
||||
|
||||
if not text or len(text) <= chunk_size:
|
||||
return [text] if text.strip() else []
|
||||
|
||||
# 按段落分割
|
||||
paragraphs = re.split(r"\n{2,}", text)
|
||||
chunks: List[str] = []
|
||||
current_chunk = ""
|
||||
|
||||
for para in paragraphs:
|
||||
para = para.strip()
|
||||
if not para:
|
||||
continue
|
||||
|
||||
if len(current_chunk) + len(para) + 2 <= chunk_size:
|
||||
# 当前块还能容纳这个段落
|
||||
if current_chunk:
|
||||
current_chunk += "\n\n" + para
|
||||
else:
|
||||
current_chunk = para
|
||||
else:
|
||||
# 当前块已满，保存并开始新块
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk)
|
||||
|
||||
if len(para) > chunk_size:
|
||||
# 单个段落超过 chunk_size，强制切分
|
||||
for i in range(0, len(para), chunk_size - overlap):
|
||||
piece = para[i : i + chunk_size]
|
||||
if piece.strip():
|
||||
chunks.append(piece)
|
||||
current_chunk = ""
|
||||
else:
|
||||
current_chunk = para
|
||||
|
||||
if current_chunk.strip():
|
||||
chunks.append(current_chunk)
|
||||
|
||||
return chunks
|
||||
479
app/services/llm_service.py
Normal file
479
app/services/llm_service.py
Normal file
@@ -0,0 +1,479 @@
|
||||
"""
|
||||
LLM 服务模块
|
||||
提供统一的 LLM 调用接口，用于查询扩展、标签提取等任务
|
||||
支持 MiniMax / DeepSeek / OpenAI（均为 OpenAI 兼容格式）
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Optional
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMService:
|
||||
"""
|
||||
LLM 服务类
|
||||
根据配置的 LLM_PROVIDER 自动选择对应的 API 端点。
|
||||
所有支持的提供商均兼容 OpenAI Chat Completions 格式。
|
||||
"""
|
||||
|
||||
_client = None
|
||||
|
||||
@classmethod
|
||||
def _get_client(cls):
|
||||
"""延迟初始化 OpenAI 兼容客户端"""
|
||||
if cls._client is not None:
|
||||
return cls._client
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
provider = settings.LLM_PROVIDER.lower()
|
||||
|
||||
if provider == "minimax":
|
||||
cls._client = AsyncOpenAI(
|
||||
api_key=settings.MINIMAX_API_KEY or "EMPTY",
|
||||
base_url=settings.MINIMAX_BASE_URL,
|
||||
)
|
||||
cls._model = settings.MINIMAX_CHAT_MODEL
|
||||
elif provider == "deepseek":
|
||||
cls._client = AsyncOpenAI(
|
||||
api_key=settings.DEEPSEEK_API_KEY or "EMPTY",
|
||||
base_url=settings.DEEPSEEK_BASE_URL,
|
||||
)
|
||||
cls._model = settings.DEEPSEEK_OCR_MODEL
|
||||
elif provider == "openai":
|
||||
cls._client = AsyncOpenAI(
|
||||
api_key=settings.OPENAI_API_KEY or "EMPTY",
|
||||
base_url=settings.OPENAI_BASE_URL,
|
||||
)
|
||||
cls._model = settings.EMBEDDING_MODEL
|
||||
else:
|
||||
raise ValueError(f"不支持的 LLM 提供商: {provider}")
|
||||
|
||||
logger.info("LLM 服务初始化完成: provider=%s, model=%s", provider, cls._model)
|
||||
return cls._client
|
||||
|
||||
@staticmethod
|
||||
def _clean_reasoning_content(content: str) -> str:
|
||||
"""清理推理模型的思考过程，只保留最终回答"""
|
||||
import re
|
||||
|
||||
# 方式1: 清理 🧠...💠 标签（MiniMax M2.7 reasoning 模式）
|
||||
content = re.sub(r'🧠[\s\S]*?💠', '', content)
|
||||
|
||||
# 方式2: 清理 💭(U+1F4AD) 包裹的思考块
|
||||
if '💭' in content:
|
||||
parts = content.split('💭')
|
||||
answer_parts = [parts[i] for i in range(len(parts)) if i % 2 == 1]
|
||||
if answer_parts:
|
||||
content = ''.join(answer_parts)
|
||||
|
||||
# 方式3: 清理 </think 标签（M2.7 有时使用）
|
||||
think_end = content.find('</think')
|
||||
if think_end >= 0:
|
||||
after = content[think_end + len('</think'):]
|
||||
after = re.sub(r'^[>\n\s]*', '', after)
|
||||
content = after
|
||||
|
||||
# 方式4: 清理 \x0f 分隔的思考块
|
||||
if '\x0f' in content:
|
||||
parts = content.split('\x0f')
|
||||
content = ''.join(parts[1::2]) if len(parts) > 1 else parts[-1] if parts else content
|
||||
|
||||
# 方式5: 清理 0x00 等特殊字符
|
||||
content = content.replace('\x00', '').strip()
|
||||
|
||||
# 方式6: 如果内容仍然很长（说明还有思考残留），取最后一个短段落作为答案
|
||||
if len(content) > 100 and '\n\n' in content:
|
||||
segments = [s.strip() for s in content.split('\n\n')]
|
||||
for seg in reversed(segments):
|
||||
if seg and len(seg) <= 100:
|
||||
content = seg
|
||||
break
|
||||
|
||||
return content.strip()
|
||||
|
||||
@classmethod
|
||||
async def chat(cls, messages: List[dict], temperature: float = 0.3, max_tokens: int = 1024) -> str:
|
||||
"""
|
||||
发送聊天请求。
|
||||
|
||||
Args:
|
||||
messages: 消息列表，格式 [{"role": "system/user/assistant", "content": "..."}]
|
||||
temperature: 采样温度
|
||||
max_tokens: 最大生成 token 数
|
||||
|
||||
Returns:
|
||||
模型回复文本（已清理思考过程）
|
||||
"""
|
||||
client = cls._get_client()
|
||||
|
||||
# MiniMax 模型默认开启思考模式，使用 reasoning_split 将思考内容
|
||||
# 分离到 reasoning_details 字段，content 只保留最终回答
|
||||
extra_body = {}
|
||||
if settings.LLM_PROVIDER.lower() == "minimax":
|
||||
extra_body["reasoning_split"] = True
|
||||
|
||||
response = await client.chat.completions.create(
|
||||
model=cls._model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
**({"extra_body": extra_body} if extra_body else {}),
|
||||
)
|
||||
content = response.choices[0].message.content or ""
|
||||
|
||||
# 兜底清理：以防 reasoning_split 未生效或使用了其他模型
|
||||
content = cls._clean_reasoning_content(content)
|
||||
|
||||
return content
|
||||
|
||||
@classmethod
|
||||
async def expand_query(cls, query: str) -> List[str]:
|
||||
"""
|
||||
查询扩展：将原始查询改写为多个语义相近的表述。
|
||||
用于混合搜索时提高召回率。
|
||||
|
||||
Args:
|
||||
query: 原始查询文本
|
||||
|
||||
Returns:
|
||||
扩展后的查询列表（包含原始查询）
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个搜索查询优化助手。用户会给你一个搜索查询，"
|
||||
"你需要生成 2 个语义相近但表述不同的替代查询。"
|
||||
"只输出 JSON 数组格式，不要添加任何其他文字。"
|
||||
'例如：["原始查询", "替代查询1", "替代查询2"]'
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"请为以下查询生成 2 个替代表述：{query}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
result = await cls.chat(messages, temperature=0.3, max_tokens=256)
|
||||
# 尝试解析 JSON
|
||||
result = result.strip()
|
||||
if result.startswith("```"):
|
||||
result = result.split("\n", 1)[-1]
|
||||
if result.endswith("```"):
|
||||
result = result[:-3]
|
||||
result = result.strip()
|
||||
queries = json.loads(result)
|
||||
if isinstance(queries, list):
|
||||
# 去重并限制数量
|
||||
seen = set()
|
||||
unique = []
|
||||
for q in queries:
|
||||
q = str(q).strip()
|
||||
if q and q not in seen:
|
||||
seen.add(q)
|
||||
unique.append(q)
|
||||
return unique[:3]
|
||||
except Exception as exc:
|
||||
logger.warning("查询扩展失败: %s，使用原始查询", exc)
|
||||
|
||||
return [query]
|
||||
|
||||
@classmethod
|
||||
async def detect_entities(cls, text: str) -> List[dict]:
|
||||
"""
|
||||
实体检测：从文本中提取知识点、人物、概念等实体。
|
||||
|
||||
Args:
|
||||
text: 待分析的文本
|
||||
|
||||
Returns:
|
||||
实体列表，格式 [{"name": "...", "type": "knowledge_point|person|concept|technique", "description": "..."}]
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个教育内容分析助手。用户会给你一段直播课的文字稿，"
|
||||
"你需要从中提取关键实体（知识点、技巧、人物、概念等）。"
|
||||
"只输出 JSON 数组格式，不要添加任何其他文字。\n"
|
||||
'格式：[{"name": "实体名称", "type": "knowledge_point|technique|person|concept", "description": "简要描述"}]\n'
|
||||
"type 说明：\n"
|
||||
"- knowledge_point: 知识点（如公式、定理、定义）\n"
|
||||
"- technique: 技巧/方法（如解题技巧、记忆方法）\n"
|
||||
"- person: 人物\n"
|
||||
"- concept: 概念/术语"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"请从以下文字中提取关键实体：\n\n{text[:3000]}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
result = await cls.chat(messages, temperature=0.1, max_tokens=2048)
|
||||
result = result.strip()
|
||||
if result.startswith("```"):
|
||||
result = result.split("\n", 1)[-1]
|
||||
if result.endswith("```"):
|
||||
result = result[:-3]
|
||||
result = result.strip()
|
||||
entities = json.loads(result)
|
||||
if isinstance(entities, list):
|
||||
return entities
|
||||
except Exception as exc:
|
||||
logger.warning("实体检测失败: %s", exc)
|
||||
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
async def extract_tags(cls, text: str) -> List[str]:
|
||||
"""
|
||||
从文本中提取语义标签。
|
||||
用于图片 OCR 内容的标签化和搜索查询的标签化。
|
||||
|
||||
Args:
|
||||
text: 待提取标签的文本
|
||||
|
||||
Returns:
|
||||
标签列表，如 ["教育", "高考数学", "二次方程", "解题技巧"]
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个内容标签提取助手。用户会给你一段文字内容，"
|
||||
"你需要从中提取 3-10 个能概括内容主题的标签。\n"
|
||||
"标签要求：\n"
|
||||
"1. 用简短的词语（2-6个字）\n"
|
||||
"2. 涵盖主题、场景、情感、关键实体等维度\n"
|
||||
"3. 考虑同义词和近义词（如'不想工作'也打上'躺平'标签）\n"
|
||||
"4. 只输出 JSON 数组，不要其他文字\n"
|
||||
'例如：["高考数学", "二次方程", "韦达定理", "解题技巧"]'
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"请从以下内容中提取标签：\n\n{text[:2000]}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
result = await cls.chat(messages, temperature=0.3, max_tokens=256)
|
||||
result = result.strip()
|
||||
if result.startswith("```"):
|
||||
result = result.split("\n", 1)[-1]
|
||||
if result.endswith("```"):
|
||||
result = result[:-3]
|
||||
result = result.strip()
|
||||
tags = json.loads(result)
|
||||
if isinstance(tags, list):
|
||||
# 去重、清洗
|
||||
seen = set()
|
||||
clean = []
|
||||
for t in tags:
|
||||
t = str(t).strip()
|
||||
if 1 <= len(t) <= 20 and t not in seen:
|
||||
seen.add(t)
|
||||
clean.append(t)
|
||||
return clean[:15]
|
||||
except Exception as exc:
|
||||
logger.warning("标签提取失败: %s", exc)
|
||||
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
async def extract_search_tags(cls, query: str) -> List[str]:
|
||||
"""
|
||||
从用户搜索查询中提取标签，用于匹配数据库中存储的标签。
|
||||
和 extract_tags 类似，但更侧重于提取搜索意图相关的标签。
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个搜索意图分析助手。用户会给你一个搜索查询，"
|
||||
"你需要提取 3-8 个可能匹配的标签词。\n"
|
||||
"标签要求：\n"
|
||||
"1. 用简短的词语（2-6个字）\n"
|
||||
"2. 包含同义词和近义词\n"
|
||||
"3. 包含上义词和下义词\n"
|
||||
"4. 只输出 JSON 数组\n"
|
||||
'例如：用户搜"孩子躺平"，输出 ["躺平", "不上班", "年轻人", "就业", "摆烂", "啃老"]'
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"请从以下搜索查询中提取匹配标签：{query}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
result = await cls.chat(messages, temperature=0.3, max_tokens=256)
|
||||
result = result.strip()
|
||||
if result.startswith("```"):
|
||||
result = result.split("\n", 1)[-1]
|
||||
if result.endswith("```"):
|
||||
result = result[:-3]
|
||||
result = result.strip()
|
||||
tags = json.loads(result)
|
||||
if isinstance(tags, list):
|
||||
seen = set()
|
||||
clean = []
|
||||
for t in tags:
|
||||
t = str(t).strip()
|
||||
if 1 <= len(t) <= 20 and t not in seen:
|
||||
seen.add(t)
|
||||
clean.append(t)
|
||||
return clean[:10]
|
||||
except Exception as exc:
|
||||
logger.warning("搜索标签提取失败: %s", exc)
|
||||
|
||||
# 降级：直接用原始查询词作为标签
|
||||
return [query.strip()] if query.strip() else []
|
||||
|
||||
@classmethod
|
||||
async def summarize_story(cls, text: str) -> str:
|
||||
"""
|
||||
用 Qwen3-8B 提炼图片 OCR 文本的故事内容。
|
||||
非思考模式，直接输出。
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个内容提炼助手。用户会给你一段从图片中识别出的文字内容，"
|
||||
"你需要提炼出其中讲述的故事或事件。\n"
|
||||
"要求：\n"
|
||||
"1. 用简洁的自然语言描述图片内容讲述的故事或事件\n"
|
||||
"2. 保留关键人物、事件、情感、场景等核心信息\n"
|
||||
"3. 50-200字\n"
|
||||
"4. 直接输出提炼内容，不要加前缀"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"/no_think\n请提炼以下内容的故事：\n\n{text[:3000]}",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
from openai import AsyncOpenAI
|
||||
client = AsyncOpenAI(
|
||||
api_key=settings.OPENAI_API_KEY,
|
||||
base_url=settings.OPENAI_BASE_URL,
|
||||
)
|
||||
response = await client.chat.completions.create(
|
||||
model="Qwen/Qwen3-8B",
|
||||
messages=messages,
|
||||
temperature=0.3,
|
||||
max_tokens=512,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||
)
|
||||
content = response.choices[0].message.content or ""
|
||||
# 清理可能的 /no_think 回显
|
||||
content = content.replace("/no_think", "").strip()
|
||||
return content
|
||||
except Exception as exc:
|
||||
logger.warning("故事提炼失败: %s", exc)
|
||||
return ""
|
||||
|
||||
@classmethod
|
||||
async def judge_batch(cls, query: str, articles: List[dict]) -> List[int]:
|
||||
"""
|
||||
用 LLM 批量判断哪些文章符合查询。
|
||||
|
||||
Args:
|
||||
query: 用户搜索查询
|
||||
articles: [{"id": int, "text": str}, ...]，最多10个
|
||||
|
||||
Returns:
|
||||
匹配的文章 id 列表（空列表表示都不匹配）
|
||||
"""
|
||||
if not articles:
|
||||
return []
|
||||
|
||||
# 构建文章列表
|
||||
articles_text = ""
|
||||
for i, article in enumerate(articles, 1):
|
||||
text = (article.get("text") or "")[:800]
|
||||
articles_text += f"{i}. {text}\n"
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"你是一个严格的语义搜索匹配引擎。用户会给你一个搜索查询和多篇文章片段，"
|
||||
"你需要判断哪些文章与查询在**核心语义上高度相关**。\n\n"
|
||||
"判断原则：\n"
|
||||
"- 文章的**核心主题**必须与查询直接相关，而不是仅仅提及了查询中的某个词\n"
|
||||
"- 例如：查询\"爸爸出轨\"，文章只提到\"爸爸带孩子玩\"是不相关的，"
|
||||
"只有文章讨论出轨、婚姻危机、夫妻矛盾等才算相关\n"
|
||||
"- 例如：查询\"孩子不想上学\"，文章提到\"孩子打游戏不去学校\"是相关的，"
|
||||
"但只提到\"孩子\"或\"学校\"其他方面则不相关\n\n"
|
||||
"不相关的典型情况（必须排除）：\n"
|
||||
"- 仅包含查询中的某个词，但讨论的是完全不同的话题\n"
|
||||
"- 主题相近但具体内容无关（如查询出轨，文章讲的是正常的夫妻相处）\n"
|
||||
"- 只有一个宽泛的关联词，没有实质性的语义联系\n\n"
|
||||
"只输出匹配的文章编号，用逗号分隔。\n"
|
||||
"如果没有文章与查询相关，只输出 0。\n"
|
||||
"不要输出任何解释。宁可漏判，不可误判。"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"搜索查询: {query}\n\n{articles_text}\n请输出匹配的文章编号:",
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
from openai import AsyncOpenAI
|
||||
client = AsyncOpenAI(
|
||||
api_key=settings.MINIMAX_API_KEY,
|
||||
base_url=settings.MINIMAX_BASE_URL,
|
||||
http_client=httpx.AsyncClient(timeout=60.0),
|
||||
)
|
||||
response = await client.chat.completions.create(
|
||||
model=settings.MINIMAX_CHAT_MODEL,
|
||||
messages=messages,
|
||||
temperature=0.0,
|
||||
max_tokens=512,
|
||||
stream=False,
|
||||
extra_body={"reasoning_split": True},
|
||||
)
|
||||
content = response.choices[0].message.content or ""
|
||||
content = cls._clean_reasoning_content(content)
|
||||
|
||||
# 解析返回的编号
|
||||
numbers = re.findall(r'\d+', content)
|
||||
|
||||
if not numbers:
|
||||
return []
|
||||
|
||||
matched_ids = []
|
||||
for num_str in numbers:
|
||||
num = int(num_str)
|
||||
if num == 0:
|
||||
continue # 0 表示都不匹配
|
||||
if 1 <= num <= len(articles):
|
||||
matched_ids.append(articles[num - 1]["id"])
|
||||
|
||||
return matched_ids
|
||||
except Exception as exc:
|
||||
logger.warning("LLM 批量判断失败: %s", exc)
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""重置客户端（用于测试或切换配置）"""
|
||||
cls._client = None
|
||||
540
app/services/ocr_service.py
Normal file
540
app/services/ocr_service.py
Normal file
@@ -0,0 +1,540 @@
|
||||
"""
|
||||
OCR 识别服务模块
|
||||
支持多种 OCR 提供商：PaddleOCR（本地）/ 阿里云 / 腾讯云
|
||||
auto 模式下本地优先 + 云端 fallback
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from app.config import OCRProvider, settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ──────────────────────────── 数据结构 ────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class OCRBlock:
|
||||
"""OCR 识别的文本块"""
|
||||
text: str
|
||||
bbox: Optional[List[float]] = None # [x1, y1, x2, y2]
|
||||
confidence: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class OCRResult:
|
||||
"""OCR 识别结果"""
|
||||
text: str
|
||||
blocks: List[OCRBlock] = field(default_factory=list)
|
||||
confidence: float = 0.0
|
||||
provider: str = ""
|
||||
keywords: List[str] = field(default_factory=list) # OCR 时提取的关键词标签
|
||||
|
||||
|
||||
# ──────────────────────────── 抽象基类 ────────────────────────────
|
||||
|
||||
class OCRProviderBase(abc.ABC):
|
||||
"""OCR 提供商抽象基类"""
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def name(self) -> str:
|
||||
"""提供商名称"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""
|
||||
识别图片中的文字。
|
||||
|
||||
Args:
|
||||
image_path: 图片文件路径
|
||||
|
||||
Returns:
|
||||
OCRResult 包含识别文本、文本块和置信度
|
||||
|
||||
Raises:
|
||||
Exception: 识别失败时抛出异常
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# ──────────────────────────── PaddleOCR 实现（本地） ────────────────────────────
|
||||
|
||||
class PaddleOCRProvider(OCRProviderBase):
|
||||
"""
|
||||
PaddleOCR 本地识别
|
||||
优点：免费、离线可用、中文识别效果好
|
||||
缺点：需要较多内存、首次加载模型较慢
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._ocr = None
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "paddleocr"
|
||||
|
||||
def _init_ocr(self):
|
||||
"""延迟初始化 PaddleOCR 引擎"""
|
||||
if self._ocr is not None:
|
||||
return
|
||||
|
||||
logger.info("正在初始化 PaddleOCR 引擎...")
|
||||
from paddleocr import PaddleOCR
|
||||
|
||||
self._ocr = PaddleOCR(
|
||||
use_angle_cls=True,
|
||||
lang="ch",
|
||||
use_gpu=False, # 默认 CPU，生产环境可通过环境变量控制
|
||||
show_log=False,
|
||||
enable_mkldnn=True, # 启用 MKLDNN 加速
|
||||
)
|
||||
logger.info("PaddleOCR 引擎初始化完成")
|
||||
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""使用 PaddleOCR 识别图片"""
|
||||
self._init_ocr()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _run_ocr():
|
||||
result = self._ocr.ocr(image_path, cls=True)
|
||||
return result
|
||||
|
||||
raw_result = await loop.run_in_executor(None, _run_ocr)
|
||||
|
||||
# 解析 PaddleOCR 返回结果
|
||||
blocks: List[OCRBlock] = []
|
||||
all_text_parts: List[str] = []
|
||||
total_confidence = 0.0
|
||||
|
||||
if raw_result and raw_result[0]:
|
||||
for line in raw_result[0]:
|
||||
bbox_points = line[0]
|
||||
text = line[1][0]
|
||||
confidence = float(line[1][1])
|
||||
|
||||
# 将四个角点转为 [x1, y1, x2, y2]
|
||||
xs = [p[0] for p in bbox_points]
|
||||
ys = [p[1] for p in bbox_points]
|
||||
bbox = [min(xs), min(ys), max(xs), max(ys)]
|
||||
|
||||
blocks.append(OCRBlock(text=text, bbox=bbox, confidence=confidence))
|
||||
all_text_parts.append(text)
|
||||
total_confidence += confidence
|
||||
|
||||
full_text = "\n".join(all_text_parts)
|
||||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||||
|
||||
return OCRResult(
|
||||
text=full_text,
|
||||
blocks=blocks,
|
||||
confidence=avg_confidence,
|
||||
provider=self.name,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 阿里云 OCR 实现 ────────────────────────────
|
||||
|
||||
class AliyunOCRProvider(OCRProviderBase):
|
||||
"""
|
||||
阿里云 OCR 识别
|
||||
使用阿里云文字识别 API（通用文字识别）
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
if not settings.ALIYUN_OCR_ACCESS_KEY or not settings.ALIYUN_OCR_SECRET:
|
||||
raise ValueError("阿里云 OCR 需要配置 ALIYUN_OCR_ACCESS_KEY 和 ALIYUN_OCR_SECRET")
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "aliyun"
|
||||
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""调用阿里云 OCR API 识别图片"""
|
||||
import base64
|
||||
import json
|
||||
|
||||
import asyncio
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _call_api():
|
||||
from alibabacloud_tea_openapi.models import Config
|
||||
from alibabacloud_ocr_api20210707.client import Client
|
||||
from alibabacloud_ocr_api20210707.models import RecognizeGeneralRequest
|
||||
from alibabacloud_tea_util.models import RuntimeOptions
|
||||
|
||||
config = Config(
|
||||
access_key_id=settings.ALIYUN_OCR_ACCESS_KEY,
|
||||
access_key_secret=settings.ALIYUN_OCR_SECRET,
|
||||
endpoint="ocr-api.cn-hangzhou.aliyuncs.com",
|
||||
)
|
||||
client = Client(config)
|
||||
|
||||
# 读取图片并 Base64 编码
|
||||
with open(image_path, "rb") as f:
|
||||
image_bytes = f.read()
|
||||
image_b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
request = RecognizeGeneralRequest(
|
||||
body=json.dumps({"url": f"data:image/png;base64,{image_b64}"}),
|
||||
)
|
||||
runtime = RuntimeOptions()
|
||||
|
||||
response = client.recognize_general_with_options(request, runtime)
|
||||
return response
|
||||
|
||||
response = await loop.run_in_executor(None, _call_api)
|
||||
|
||||
if response.body and response.body.data:
|
||||
data = json.loads(response.body.data)
|
||||
blocks: List[OCRBlock] = []
|
||||
all_text_parts: List[str] = []
|
||||
total_confidence = 0.0
|
||||
|
||||
# 阿里云返回格式可能包含多个文本块
|
||||
if isinstance(data, dict) and "content" in data:
|
||||
text = data["content"]
|
||||
all_text_parts.append(text)
|
||||
blocks.append(OCRBlock(text=text, confidence=1.0))
|
||||
total_confidence = 1.0
|
||||
elif isinstance(data, list):
|
||||
for item in data:
|
||||
text = item.get("text", item.get("word", ""))
|
||||
score = float(item.get("score", item.get("confidence", 1.0)))
|
||||
blocks.append(OCRBlock(text=text, confidence=score))
|
||||
all_text_parts.append(text)
|
||||
total_confidence += score
|
||||
|
||||
full_text = "\n".join(all_text_parts)
|
||||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||||
|
||||
return OCRResult(
|
||||
text=full_text,
|
||||
blocks=blocks,
|
||||
confidence=avg_confidence,
|
||||
provider=self.name,
|
||||
)
|
||||
|
||||
return OCRResult(text="", blocks=[], confidence=0.0, provider=self.name)
|
||||
|
||||
|
||||
# ──────────────────────────── 腾讯云 OCR 实现 ────────────────────────────
|
||||
|
||||
class TencentOCRProvider(OCRProviderBase):
|
||||
"""
|
||||
腾讯云 OCR 识别
|
||||
使用腾讯云文字识别 API（通用印刷体识别）
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
if not settings.TENCENT_OCR_SECRET_ID or not settings.TENCENT_OCR_SECRET_KEY:
|
||||
raise ValueError("腾讯云 OCR 需要配置 TENCENT_OCR_SECRET_ID 和 TENCENT_OCR_SECRET_KEY")
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "tencent"
|
||||
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""调用腾讯云 OCR API 识别图片"""
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _call_api():
|
||||
from tencentcloud.common import credential
|
||||
from tencentcloud.common.profile.client_profile import ClientProfile
|
||||
from tencentcloud.common.profile.http_profile import HttpProfile
|
||||
from tencentcloud.ocr.v20181119 import ocr_client, models
|
||||
|
||||
cred = credential.Credential(
|
||||
settings.TENCENT_OCR_SECRET_ID,
|
||||
settings.TENCENT_OCR_SECRET_KEY,
|
||||
)
|
||||
http_profile = HttpProfile()
|
||||
http_profile.endpoint = "ocr.tencentcloudapi.com"
|
||||
|
||||
client_profile = ClientProfile()
|
||||
client_profile.httpProfile = http_profile
|
||||
|
||||
client = ocr_client.OcrClient(cred, "", client_profile)
|
||||
|
||||
# 读取图片并 Base64 编码
|
||||
with open(image_path, "rb") as f:
|
||||
image_bytes = f.read()
|
||||
image_b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
req = models.GeneralBasicOCRRequest()
|
||||
req.ImageBase64 = image_b64
|
||||
|
||||
resp = client.GeneralBasicOCR(req)
|
||||
return resp
|
||||
|
||||
resp = await loop.run_in_executor(None, _call_api)
|
||||
|
||||
blocks: List[OCRBlock] = []
|
||||
all_text_parts: List[str] = []
|
||||
total_confidence = 0.0
|
||||
|
||||
if resp.TextDetections:
|
||||
for item in resp.TextDetections:
|
||||
text = item.DetectedText or ""
|
||||
confidence = item.Confidence / 100.0 if item.Confidence else 0.0
|
||||
|
||||
# 解析多边形顶点
|
||||
bbox = None
|
||||
if item.Polygon:
|
||||
xs = [p.X for p in item.Polygon]
|
||||
ys = [p.Y for p in item.Polygon]
|
||||
bbox = [min(xs), min(ys), max(xs), max(ys)]
|
||||
|
||||
blocks.append(OCRBlock(text=text, bbox=bbox, confidence=confidence))
|
||||
all_text_parts.append(text)
|
||||
total_confidence += confidence
|
||||
|
||||
full_text = "\n".join(all_text_parts)
|
||||
avg_confidence = total_confidence / len(blocks) if blocks else 0.0
|
||||
|
||||
return OCRResult(
|
||||
text=full_text,
|
||||
blocks=blocks,
|
||||
confidence=avg_confidence,
|
||||
provider=self.name,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── DeepSeek OCR 实现 ────────────────────────────
|
||||
|
||||
class DeepSeekOCRProvider(OCRProviderBase):
|
||||
"""
|
||||
DeepSeek Vision OCR
|
||||
通过 OpenAI 兼容接口发送图片 base64，让视觉模型识别文字。
|
||||
支持 DeepSeek 官方 API 和 Ollama 自部署（如 deepseek-ocr 模型）。
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
if not settings.DEEPSEEK_API_KEY:
|
||||
raise ValueError("DeepSeek OCR 需要配置 DEEPSEEK_API_KEY")
|
||||
import httpx
|
||||
from openai import AsyncOpenAI
|
||||
self._client = AsyncOpenAI(
|
||||
api_key=settings.DEEPSEEK_API_KEY,
|
||||
base_url=settings.DEEPSEEK_BASE_URL,
|
||||
http_client=httpx.AsyncClient(timeout=120.0), # OCR 较慢，超时 120 秒
|
||||
)
|
||||
self._model = settings.DEEPSEEK_OCR_MODEL
|
||||
# 判断是否为 Ollama 部署
|
||||
self._is_ollama = "ollama" in settings.DEEPSEEK_API_KEY.lower() or \
|
||||
"11434" in settings.DEEPSEEK_BASE_URL
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "deepseek"
|
||||
|
||||
async def recognize(self, image_path: str) -> OCRResult:
|
||||
"""使用 DeepSeek Vision 模型识别图片中的文字"""
|
||||
import base64
|
||||
|
||||
# 读取图片并编码为 base64
|
||||
with open(image_path, "rb") as f:
|
||||
image_bytes = f.read()
|
||||
base64_image = base64.b64encode(image_bytes).decode("utf-8")
|
||||
|
||||
# 根据图片格式确定 MIME 类型
|
||||
suffix = Path(image_path).suffix.lower()
|
||||
mime_map = {
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".bmp": "image/bmp",
|
||||
".webp": "image/webp",
|
||||
}
|
||||
mime_type = mime_map.get(suffix, "image/png")
|
||||
|
||||
system_prompt = (
|
||||
"你是一个 OCR 文字识别工具。"
|
||||
"请识别图片中的所有文字，逐行输出，不要添加任何解释。"
|
||||
)
|
||||
|
||||
response = await self._client.chat.completions.create(
|
||||
model=self._model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": system_prompt,
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:{mime_type};base64,{base64_image}"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "请识别这张图片中的所有文字内容。",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
max_tokens=4096,
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
text = response.choices[0].message.content or ""
|
||||
|
||||
# 清理特殊 token（deepseek-ocr 模型可能输出 <｜begin of sentence｜> 等）
|
||||
import re
|
||||
text = re.sub(r'<｜[^>]*｜>', '', text)
|
||||
text = re.sub(r'<\|[^>]*\|>', '', text)
|
||||
|
||||
# 清理可能的 markdown 包裹
|
||||
text = text.strip()
|
||||
if text.startswith("```"):
|
||||
text = text.split("\n", 1)[-1]
|
||||
if text.endswith("```"):
|
||||
text = text[:-3]
|
||||
text = text.strip()
|
||||
|
||||
lines = [line for line in text.split("\n") if line.strip()]
|
||||
blocks = [OCRBlock(text=line, confidence=1.0) for line in lines]
|
||||
|
||||
return OCRResult(
|
||||
text=text,
|
||||
blocks=blocks,
|
||||
confidence=1.0,
|
||||
provider=self.name,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────── 工厂类 ────────────────────────────
|
||||
|
||||
class OCRService:
|
||||
"""
|
||||
OCR 服务工厂类
|
||||
根据配置自动选择 OCR 提供商。
|
||||
auto 模式下：本地 PaddleOCR 优先，失败后 fallback 到云端 OCR。
|
||||
"""
|
||||
|
||||
_instance: Optional[OCRProviderBase] = None
|
||||
|
||||
@classmethod
|
||||
def _create_provider(cls, provider: OCRProvider) -> OCRProviderBase:
|
||||
"""根据提供商枚举创建对应的 OCR 实例"""
|
||||
provider_map = {
|
||||
OCRProvider.PADDLEOCR: PaddleOCRProvider,
|
||||
OCRProvider.ALIYUN: AliyunOCRProvider,
|
||||
OCRProvider.TENCENT: TencentOCRProvider,
|
||||
OCRProvider.DEEPSEEK: DeepSeekOCRProvider,
|
||||
}
|
||||
provider_cls = provider_map.get(provider)
|
||||
if provider_cls is None:
|
||||
raise ValueError(f"不支持的 OCR 提供商: {provider}")
|
||||
return provider_cls()
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> OCRProviderBase:
|
||||
"""获取 OCR 服务单例实例"""
|
||||
if cls._instance is not None:
|
||||
return cls._instance
|
||||
|
||||
if settings.OCR_PROVIDER == OCRProvider.AUTO:
|
||||
# auto 模式返回 PaddleOCR，fallback 逻辑在 recognize_auto 中处理
|
||||
cls._instance = PaddleOCRProvider()
|
||||
else:
|
||||
cls._instance = cls._create_provider(settings.OCR_PROVIDER)
|
||||
|
||||
logger.info("OCR 服务初始化完成: provider=%s", cls._instance.name)
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
async def recognize(cls, image_path: str) -> OCRResult:
|
||||
"""
|
||||
识别图片文字。
|
||||
auto 模式下：PaddleOCR 优先，失败后依次尝试阿里云、腾讯云。
|
||||
"""
|
||||
if settings.OCR_PROVIDER == OCRProvider.AUTO:
|
||||
return await cls._recognize_auto(image_path)
|
||||
|
||||
instance = cls.get_instance()
|
||||
return await instance.recognize(image_path)
|
||||
|
||||
@classmethod
|
||||
async def _recognize_auto(cls, image_path: str) -> OCRResult:
|
||||
"""
|
||||
auto 模式识别逻辑：
|
||||
1. 先尝试本地 PaddleOCR
|
||||
2. 如果失败且配置了阿里云，尝试阿里云 OCR
|
||||
3. 如果仍失败且配置了腾讯云，尝试腾讯云 OCR
|
||||
"""
|
||||
# 第一步：本地 PaddleOCR
|
||||
try:
|
||||
provider = PaddleOCRProvider()
|
||||
result = await provider.recognize(image_path)
|
||||
if result.text.strip():
|
||||
logger.info("PaddleOCR 识别成功，文本长度: %d", len(result.text))
|
||||
return result
|
||||
logger.warning("PaddleOCR 返回空文本，尝试 fallback")
|
||||
except Exception as exc:
|
||||
logger.warning("PaddleOCR 识别失败: %s，尝试 fallback", exc)
|
||||
|
||||
# 第二步：DeepSeek OCR（如果配置了）
|
||||
if settings.DEEPSEEK_API_KEY:
|
||||
try:
|
||||
provider = DeepSeekOCRProvider()
|
||||
result = await provider.recognize(image_path)
|
||||
if result.text.strip():
|
||||
logger.info("DeepSeek OCR 识别成功，文本长度: %d", len(result.text))
|
||||
return result
|
||||
logger.warning("DeepSeek OCR 返回空文本，尝试 fallback")
|
||||
except Exception as exc:
|
||||
logger.warning("DeepSeek OCR 识别失败: %s，尝试 fallback", exc)
|
||||
|
||||
# 第三步：阿里云 OCR
|
||||
if settings.ALIYUN_OCR_ACCESS_KEY and settings.ALIYUN_OCR_SECRET:
|
||||
try:
|
||||
provider = AliyunOCRProvider()
|
||||
result = await provider.recognize(image_path)
|
||||
if result.text.strip():
|
||||
logger.info("阿里云 OCR 识别成功，文本长度: %d", len(result.text))
|
||||
return result
|
||||
logger.warning("阿里云 OCR 返回空文本，继续 fallback")
|
||||
except Exception as exc:
|
||||
logger.warning("阿里云 OCR 识别失败: %s，继续 fallback", exc)
|
||||
|
||||
# 第四步：腾讯云 OCR
|
||||
if settings.TENCENT_OCR_SECRET_ID and settings.TENCENT_OCR_SECRET_KEY:
|
||||
try:
|
||||
provider = TencentOCRProvider()
|
||||
result = await provider.recognize(image_path)
|
||||
if result.text.strip():
|
||||
logger.info("腾讯云 OCR 识别成功，文本长度: %d", len(result.text))
|
||||
return result
|
||||
logger.warning("腾讯云 OCR 返回空文本")
|
||||
except Exception as exc:
|
||||
logger.warning("腾讯云 OCR 识别失败: %s", exc)
|
||||
|
||||
# 所有 OCR 都失败
|
||||
return OCRResult(
|
||||
text="",
|
||||
blocks=[],
|
||||
confidence=0.0,
|
||||
provider="none",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""重置单例（用于测试或切换配置）"""
|
||||
cls._instance = None
|
||||
202
app/services/page_service.py
Normal file
202
app/services/page_service.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""
|
||||
知识页面服务
|
||||
提供知识页面的 CRUD 操作
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import func, select, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.base import KnowledgeChunk, KnowledgePage
|
||||
from app.schemas.page import PageCreate, PageResponse, PageUpdate
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PageService:
|
||||
"""知识页面服务"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self.db = db
|
||||
|
||||
async def create_page(self, data: PageCreate) -> KnowledgePage:
|
||||
"""创建知识页面"""
|
||||
page = KnowledgePage(
|
||||
title=data.title,
|
||||
content=data.content,
|
||||
source_file=data.source_file,
|
||||
course_name=data.course_name,
|
||||
teacher_name=data.teacher_name,
|
||||
live_date=data.live_date,
|
||||
page_number=data.page_number,
|
||||
metadata_json=data.metadata_json,
|
||||
)
|
||||
self.db.add(page)
|
||||
await self.db.flush()
|
||||
await self.db.refresh(page)
|
||||
return page
|
||||
|
||||
async def get_page(self, page_id: int) -> Optional[KnowledgePage]:
|
||||
"""根据 ID 获取知识页面"""
|
||||
result = await self.db.execute(
|
||||
select(KnowledgePage).where(KnowledgePage.id == page_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def list_pages(
|
||||
self,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
course_name: Optional[str] = None,
|
||||
teacher_name: Optional[str] = None,
|
||||
keyword: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
分页查询知识页面列表。
|
||||
|
||||
Args:
|
||||
page: 页码（从 1 开始）
|
||||
page_size: 每页数量
|
||||
course_name: 按课程名称过滤
|
||||
teacher_name: 按讲师名称过滤
|
||||
keyword: 关键词搜索（标题或内容）
|
||||
|
||||
Returns:
|
||||
{"total": int, "page": int, "page_size": int, "items": list}
|
||||
"""
|
||||
query = select(KnowledgePage)
|
||||
count_query = select(func.count(KnowledgePage.id))
|
||||
|
||||
# 过滤条件
|
||||
if course_name:
|
||||
query = query.where(KnowledgePage.course_name == course_name)
|
||||
count_query = count_query.where(KnowledgePage.course_name == course_name)
|
||||
if teacher_name:
|
||||
query = query.where(KnowledgePage.teacher_name == teacher_name)
|
||||
count_query = count_query.where(KnowledgePage.teacher_name == teacher_name)
|
||||
if keyword:
|
||||
like_pattern = f"%{keyword}%"
|
||||
query = query.where(
|
||||
KnowledgePage.title.ilike(like_pattern)
|
||||
| KnowledgePage.content.ilike(like_pattern)
|
||||
)
|
||||
count_query = count_query.where(
|
||||
KnowledgePage.title.ilike(like_pattern)
|
||||
| KnowledgePage.content.ilike(like_pattern)
|
||||
)
|
||||
|
||||
# 总数
|
||||
total_result = await self.db.execute(count_query)
|
||||
total = total_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.order_by(KnowledgePage.created_at.desc()).offset(offset).limit(page_size)
|
||||
result = await self.db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"items": items,
|
||||
}
|
||||
|
||||
async def update_page(self, page_id: int, data: PageUpdate) -> Optional[KnowledgePage]:
|
||||
"""更新知识页面"""
|
||||
page = await self.get_page(page_id)
|
||||
if page is None:
|
||||
return None
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
for field, value in update_data.items():
|
||||
setattr(page, field, value)
|
||||
|
||||
await self.db.flush()
|
||||
await self.db.refresh(page)
|
||||
return page
|
||||
|
||||
async def delete_page(self, page_id: int) -> bool:
|
||||
"""删除知识页面及其关联分块"""
|
||||
page = await self.get_page(page_id)
|
||||
if page is None:
|
||||
return False
|
||||
|
||||
# 先删除关联的分块
|
||||
await self.db.execute(
|
||||
text("DELETE FROM knowledge_chunks WHERE page_id = :page_id"),
|
||||
{"page_id": page_id},
|
||||
)
|
||||
# 再删除页面
|
||||
await self.db.delete(page)
|
||||
await self.db.flush()
|
||||
return True
|
||||
|
||||
async def get_page_chunks(self, page_id: int) -> List[KnowledgeChunk]:
|
||||
"""获取知识页面的所有分块"""
|
||||
result = await self.db.execute(
|
||||
select(KnowledgeChunk)
|
||||
.where(KnowledgeChunk.page_id == page_id)
|
||||
.order_by(KnowledgeChunk.chunk_index)
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def reindex_page(self, page_id: int) -> dict:
|
||||
"""
|
||||
重新索引知识页面：重新分块并生成嵌入向量。
|
||||
用于页面内容更新后刷新向量索引。
|
||||
"""
|
||||
page = await self.get_page(page_id)
|
||||
if page is None:
|
||||
raise ValueError(f"知识页面不存在: {page_id}")
|
||||
|
||||
# 删除旧分块
|
||||
await self.db.execute(
|
||||
text("DELETE FROM knowledge_chunks WHERE page_id = :page_id"),
|
||||
{"page_id": page_id},
|
||||
)
|
||||
|
||||
# 重新分块
|
||||
from app.services.import_service import ImportService
|
||||
chunks = ImportService._chunk_text(page.content)
|
||||
|
||||
if not chunks:
|
||||
await self.db.flush()
|
||||
return {"page_id": page_id, "chunks": 0}
|
||||
|
||||
# 生成嵌入并插入
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
try:
|
||||
embeddings = await EmbeddingService.embed_batch(chunks)
|
||||
except Exception as exc:
|
||||
logger.error("重新索引嵌入失败: %s", exc)
|
||||
embeddings = [None] * len(chunks)
|
||||
|
||||
for idx, (chunk_text, embedding) in enumerate(zip(chunks, embeddings)):
|
||||
embedding_str = None
|
||||
if embedding:
|
||||
embedding_str = "[" + ",".join(str(x) for x in embedding) + "]"
|
||||
|
||||
await self.db.execute(text("""
|
||||
INSERT INTO knowledge_chunks (page_id, chunk_index, content, embedding)
|
||||
VALUES (:page_id, :chunk_index, :content, :embedding::vector)
|
||||
"""), {
|
||||
"page_id": page_id,
|
||||
"chunk_index": idx,
|
||||
"content": chunk_text,
|
||||
"embedding": embedding_str,
|
||||
})
|
||||
|
||||
# 更新全文搜索向量
|
||||
await self.db.execute(text("""
|
||||
UPDATE knowledge_chunks
|
||||
SET search_vector = to_tsvector('chinese_zh', content)
|
||||
WHERE page_id = :page_id
|
||||
"""), {"page_id": page_id})
|
||||
|
||||
await self.db.flush()
|
||||
return {"page_id": page_id, "chunks": len(chunks)}
|
||||
185
app/services/search_engine.py
Normal file
185
app/services/search_engine.py
Normal file
@@ -0,0 +1,185 @@
|
||||
"""
|
||||
BM25 倒排索引搜索引擎
|
||||
用于图片 OCR 内容的全文搜索
|
||||
参考 Meilisearch / Whoosh 设计
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import math
|
||||
import logging
|
||||
from collections import Counter, defaultdict
|
||||
from typing import List, Dict, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 尝试导入 jieba，失败则用简单的字符分割
|
||||
try:
|
||||
import jieba
|
||||
jieba.setLogLevel(logging.WARNING)
|
||||
def tokenize(text: str) -> List[str]:
|
||||
"""中文分词"""
|
||||
return [w for w in jieba.cut_for_search(text) if len(w.strip()) > 1]
|
||||
except ImportError:
|
||||
logger.warning("jieba 未安装，使用简单分词（pip install jieba 获得更好效果）")
|
||||
def tokenize(text: str) -> List[str]:
|
||||
"""简单分词：按空格和标点分割"""
|
||||
return [w for w in re.split(r'[\s,.\-;:!?"\'/\\|()\[\]{}<>，。；：！？""''、（）【】]+', text) if len(w.strip()) > 1]
|
||||
|
||||
|
||||
class BM25Index:
|
||||
"""
|
||||
BM25 倒排索引
|
||||
- 支持增量添加文档
|
||||
- 支持删除文档
|
||||
- 支持中文分词（jieba）
|
||||
- BM25 评分排序
|
||||
"""
|
||||
|
||||
def __init__(self, k1: float = 1.5, b: float = 0.75):
|
||||
"""
|
||||
Args:
|
||||
k1: 词频饱和参数（默认 1.5）
|
||||
b: 文档长度归一化参数（默认 0.75）
|
||||
"""
|
||||
self.k1 = k1
|
||||
self.b = b
|
||||
|
||||
# 倒排索引：term -> {doc_id: tf}
|
||||
self.inverted_index: Dict[str, Dict[int, int]] = defaultdict(lambda: defaultdict(int))
|
||||
|
||||
# 文档信息
|
||||
self.doc_lengths: Dict[int, int] = {} # doc_id -> 文档长度（词数）
|
||||
self.doc_count: int = 0
|
||||
self.avg_doc_length: float = 0.0
|
||||
|
||||
# 文档存储（存完整数据）
|
||||
self.documents: Dict[int, dict] = {}
|
||||
|
||||
# IDF 缓存
|
||||
self._idf_cache: Dict[str, float] = {}
|
||||
|
||||
def add_document(self, doc_id: int, text: str, metadata: Optional[dict] = None):
|
||||
"""
|
||||
添加文档到索引
|
||||
|
||||
Args:
|
||||
doc_id: 文档 ID
|
||||
text: 待索引的文本内容
|
||||
metadata: 附加元数据（如 file_path, tags 等）
|
||||
"""
|
||||
# 如果已存在，先删除
|
||||
if doc_id in self.documents:
|
||||
self.remove_document(doc_id)
|
||||
|
||||
tokens = tokenize(text)
|
||||
self.doc_lengths[doc_id] = len(tokens)
|
||||
self.documents[doc_id] = {
|
||||
"text": text,
|
||||
"tokens": tokens,
|
||||
"metadata": metadata or {},
|
||||
}
|
||||
|
||||
# 构建倒排索引
|
||||
tf = Counter(tokens)
|
||||
for term, count in tf.items():
|
||||
self.inverted_index[term][doc_id] = count
|
||||
|
||||
self.doc_count += 1
|
||||
self._update_avg_doc_length()
|
||||
self._idf_cache.clear() # 文档数变了，IDF 需要重算
|
||||
|
||||
def remove_document(self, doc_id: int):
|
||||
"""从索引中删除文档"""
|
||||
if doc_id not in self.documents:
|
||||
return
|
||||
|
||||
# 从倒排索引中移除
|
||||
tokens = self.documents[doc_id].get("tokens", [])
|
||||
for term in set(tokens):
|
||||
if term in self.inverted_index and doc_id in self.inverted_index[term]:
|
||||
del self.inverted_index[term][doc_id]
|
||||
if not self.inverted_index[term]:
|
||||
del self.inverted_index[term]
|
||||
|
||||
del self.documents[doc_id]
|
||||
del self.doc_lengths[doc_id]
|
||||
self.doc_count -= 1
|
||||
self._update_avg_doc_length()
|
||||
self._idf_cache.clear()
|
||||
|
||||
def _update_avg_doc_length(self):
|
||||
"""更新平均文档长度"""
|
||||
if self.doc_count > 0:
|
||||
self.avg_doc_length = sum(self.doc_lengths.values()) / self.doc_count
|
||||
else:
|
||||
self.avg_doc_length = 0.0
|
||||
|
||||
def _idf(self, term: str) -> float:
|
||||
"""
|
||||
计算逆文档频率 IDF
|
||||
IDF(t) = ln((N - df(t) + 0.5) / (df(t) + 0.5) + 1)
|
||||
其中 N 是文档总数，df(t) 是包含 term 的文档数
|
||||
"""
|
||||
if term in self._idf_cache:
|
||||
return self._idf_cache[term]
|
||||
|
||||
df = len(self.inverted_index.get(term, {}))
|
||||
idf = math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1)
|
||||
|
||||
self._idf_cache[term] = idf
|
||||
return idf
|
||||
|
||||
def search(self, query: str, top_k: int = 10) -> List[Tuple[int, float]]:
|
||||
"""
|
||||
BM25 搜索
|
||||
|
||||
Args:
|
||||
query: 搜索查询（会自动分词）
|
||||
top_k: 返回前 K 个结果
|
||||
|
||||
Returns:
|
||||
[(doc_id, score), ...] 按 BM25 分数降序排列
|
||||
"""
|
||||
if self.doc_count == 0 or self.avg_doc_length == 0:
|
||||
return []
|
||||
|
||||
query_tokens = tokenize(query)
|
||||
if not query_tokens:
|
||||
return []
|
||||
|
||||
scores: Dict[int, float] = defaultdict(float)
|
||||
|
||||
for term in query_tokens:
|
||||
idf = self._idf(term)
|
||||
postings = self.inverted_index.get(term, {})
|
||||
|
||||
for doc_id, tf in postings.items():
|
||||
dl = self.doc_lengths[doc_id]
|
||||
# BM25 公式
|
||||
numerator = tf * (self.k1 + 1)
|
||||
denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_doc_length)
|
||||
scores[doc_id] += idf * numerator / denominator
|
||||
|
||||
# 按分数降序排序
|
||||
ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
|
||||
return ranked[:top_k]
|
||||
|
||||
def get_document(self, doc_id: int) -> Optional[dict]:
|
||||
"""获取文档完整数据"""
|
||||
return self.documents.get(doc_id)
|
||||
|
||||
def clear(self):
|
||||
"""清空索引"""
|
||||
self.inverted_index.clear()
|
||||
self.doc_lengths.clear()
|
||||
self.documents.clear()
|
||||
self._idf_cache.clear()
|
||||
self.doc_count = 0
|
||||
self.avg_doc_length = 0.0
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
"""索引中的文档数量"""
|
||||
return self.doc_count
|
||||
286
app/services/search_service.py
Normal file
286
app/services/search_service.py
Normal file
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
语义搜索服务
|
||||
支持 PostgreSQL（向量+全文混合）和 SQLite（LIKE 关键词搜索）
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.config import settings
|
||||
from app.database import IS_SQLITE
|
||||
from app.schemas.search import SearchRequest, SearchResult, SearchResponse
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SearchService:
|
||||
"""语义搜索服务"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self.db = db
|
||||
|
||||
async def search(self, request: SearchRequest) -> SearchResponse:
|
||||
"""
|
||||
执行搜索。
|
||||
- PostgreSQL: 向量搜索 + 中文全文搜索混合排序
|
||||
- SQLite: LIKE 关键词搜索（无向量能力）
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
if IS_SQLITE:
|
||||
return await self._search_sqlite(request, start_time)
|
||||
else:
|
||||
return await self._search_postgres(request, start_time)
|
||||
|
||||
# ──────────────────────────── SQLite 搜索 ────────────────────────────
|
||||
|
||||
async def _search_sqlite(self, request: SearchRequest, start_time: float) -> SearchResponse:
|
||||
"""SQLite 模式：使用 LIKE 关键词搜索"""
|
||||
# LLM 查询扩展（可选）
|
||||
expanded_queries = [request.query]
|
||||
try:
|
||||
from app.services.llm_service import LLMService
|
||||
expanded_queries = await LLMService.expand_query(request.query)
|
||||
except Exception as exc:
|
||||
logger.debug("查询扩展跳过: %s", exc)
|
||||
|
||||
results: list[SearchResult] = []
|
||||
seen_chunk_ids = set()
|
||||
|
||||
for q in expanded_queries:
|
||||
like_pattern = f"%{q}%"
|
||||
sql = text("""
|
||||
SELECT
|
||||
kc.id AS chunk_id,
|
||||
kp.id AS page_id,
|
||||
kp.title AS page_title,
|
||||
kc.content,
|
||||
kp.course_name,
|
||||
kp.teacher_name,
|
||||
kp.live_date
|
||||
FROM knowledge_chunks kc
|
||||
JOIN knowledge_pages kp ON kc.page_id = kp.id
|
||||
WHERE kc.content LIKE :pattern
|
||||
AND (:course IS NULL OR kp.course_name = :course)
|
||||
AND (:teacher IS NULL OR kp.teacher_name = :teacher)
|
||||
ORDER BY kc.id
|
||||
LIMIT :limit
|
||||
""")
|
||||
params = {
|
||||
"pattern": like_pattern,
|
||||
"course": request.course_name,
|
||||
"teacher": request.teacher_name,
|
||||
"limit": request.top_k,
|
||||
}
|
||||
result = await self.db.execute(sql, params)
|
||||
rows = result.fetchall()
|
||||
|
||||
for row in rows:
|
||||
if row.chunk_id not in seen_chunk_ids:
|
||||
seen_chunk_ids.add(row.chunk_id)
|
||||
# 简单的相关性评分：关键词出现次数
|
||||
score = min(1.0, row.content.lower().count(q.lower()) * 0.2 + 0.5)
|
||||
results.append(SearchResult(
|
||||
chunk_id=row.chunk_id,
|
||||
page_id=row.page_id,
|
||||
page_title=row.page_title,
|
||||
content=row.content,
|
||||
score=round(score, 4),
|
||||
course_name=row.course_name,
|
||||
teacher_name=row.teacher_name,
|
||||
live_date=row.live_date,
|
||||
highlight=None,
|
||||
))
|
||||
|
||||
results.sort(key=lambda x: x.score, reverse=True)
|
||||
results = results[:request.top_k]
|
||||
|
||||
elapsed_ms = (time.time() - start_time) * 1000
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
total=len(results),
|
||||
results=results,
|
||||
elapsed_ms=round(elapsed_ms, 2),
|
||||
)
|
||||
|
||||
# ──────────────────────────── PostgreSQL 搜索 ────────────────────────────
|
||||
|
||||
async def _search_postgres(self, request: SearchRequest, start_time: float) -> SearchResponse:
|
||||
"""PostgreSQL 模式：向量搜索 + 中文全文搜索混合"""
|
||||
# LLM 查询扩展
|
||||
expanded_queries = [request.query]
|
||||
if request.use_fulltext:
|
||||
try:
|
||||
from app.services.llm_service import LLMService
|
||||
expanded_queries = await LLMService.expand_query(request.query)
|
||||
except Exception as exc:
|
||||
logger.debug("查询扩展跳过: %s", exc)
|
||||
|
||||
# 生成查询向量
|
||||
from app.services.embedding_service import EmbeddingService
|
||||
all_query_embeddings = await EmbeddingService.embed_batch(expanded_queries)
|
||||
|
||||
# 构建过滤条件
|
||||
where_clauses = []
|
||||
params: dict = {}
|
||||
|
||||
if request.course_name:
|
||||
where_clauses.append("kp.course_name = :course_name")
|
||||
params["course_name"] = request.course_name
|
||||
if request.teacher_name:
|
||||
where_clauses.append("kp.teacher_name = :teacher_name")
|
||||
params["teacher_name"] = request.teacher_name
|
||||
if request.live_date_from:
|
||||
where_clauses.append("kp.live_date >= :live_date_from")
|
||||
params["live_date_from"] = request.live_date_from
|
||||
if request.live_date_to:
|
||||
where_clauses.append("kp.live_date <= :live_date_to")
|
||||
params["live_date_to"] = request.live_date_to
|
||||
|
||||
where_sql = ""
|
||||
if where_clauses:
|
||||
where_sql = "AND " + " AND ".join(where_clauses)
|
||||
|
||||
# 向量搜索
|
||||
vector_results = []
|
||||
params["limit"] = request.top_k * 2
|
||||
for q_embedding in all_query_embeddings:
|
||||
embedding_str = "[" + ",".join(str(x) for x in q_embedding) + "]"
|
||||
vector_sql = f"""
|
||||
SELECT
|
||||
kc.id AS chunk_id,
|
||||
kp.id AS page_id,
|
||||
kp.title AS page_title,
|
||||
kc.content,
|
||||
1 - (kc.embedding <=> :embedding::vector) AS vector_score,
|
||||
0.0 AS text_score,
|
||||
kp.course_name,
|
||||
kp.teacher_name,
|
||||
kp.live_date
|
||||
FROM knowledge_chunks kc
|
||||
JOIN knowledge_pages kp ON kc.page_id = kp.id
|
||||
WHERE kc.embedding IS NOT NULL
|
||||
{where_sql}
|
||||
ORDER BY kc.embedding <=> :embedding::vector
|
||||
LIMIT :limit
|
||||
"""
|
||||
v_params = {**params, "embedding": embedding_str}
|
||||
result = await self.db.execute(text(vector_sql), v_params)
|
||||
vector_results.extend(result.fetchall())
|
||||
|
||||
# 全文搜索
|
||||
text_results = []
|
||||
if request.use_fulltext:
|
||||
for q in expanded_queries:
|
||||
clean_query = q.replace("'", "''")
|
||||
fulltext_sql = f"""
|
||||
SELECT
|
||||
kc.id AS chunk_id,
|
||||
kp.id AS page_id,
|
||||
kp.title AS page_title,
|
||||
kc.content,
|
||||
0.0 AS vector_score,
|
||||
ts_rank_cd(kc.search_vector, to_tsquery('chinese_zh', :query)) AS text_score,
|
||||
kp.course_name,
|
||||
kp.teacher_name,
|
||||
kp.live_date,
|
||||
ts_headline('chinese_zh', kc.content, to_tsquery('chinese_zh', :query),
|
||||
'MaxWords=50, MinWords=20, ShortWord=2') AS highlight
|
||||
FROM knowledge_chunks kc
|
||||
JOIN knowledge_pages kp ON kc.page_id = kp.id
|
||||
WHERE kc.search_vector @@ to_tsquery('chinese_zh', :query)
|
||||
{where_sql}
|
||||
ORDER BY text_score DESC
|
||||
LIMIT :limit
|
||||
"""
|
||||
ft_params = {**params, "query": clean_query}
|
||||
ft_result = await self.db.execute(text(fulltext_sql), ft_params)
|
||||
text_results.extend(ft_result.fetchall())
|
||||
|
||||
# 合并结果
|
||||
merged: dict[int, dict] = {}
|
||||
|
||||
for row in vector_results:
|
||||
chunk_id = row.chunk_id
|
||||
if chunk_id not in merged:
|
||||
merged[chunk_id] = {
|
||||
"chunk_id": chunk_id,
|
||||
"page_id": row.page_id,
|
||||
"page_title": row.page_title,
|
||||
"content": row.content,
|
||||
"vector_score": row.vector_score,
|
||||
"text_score": 0.0,
|
||||
"course_name": row.course_name,
|
||||
"teacher_name": row.teacher_name,
|
||||
"live_date": row.live_date,
|
||||
"highlight": None,
|
||||
}
|
||||
else:
|
||||
merged[chunk_id]["vector_score"] = max(
|
||||
merged[chunk_id]["vector_score"], row.vector_score
|
||||
)
|
||||
|
||||
for row in text_results:
|
||||
chunk_id = row.chunk_id
|
||||
if chunk_id not in merged:
|
||||
merged[chunk_id] = {
|
||||
"chunk_id": chunk_id,
|
||||
"page_id": row.page_id,
|
||||
"page_title": row.page_title,
|
||||
"content": row.content,
|
||||
"vector_score": 0.0,
|
||||
"text_score": row.text_score or 0.0,
|
||||
"course_name": row.course_name,
|
||||
"teacher_name": row.teacher_name,
|
||||
"live_date": row.live_date,
|
||||
"highlight": row.highlight,
|
||||
}
|
||||
else:
|
||||
merged[chunk_id]["text_score"] = max(
|
||||
merged[chunk_id]["text_score"], row.text_score or 0.0
|
||||
)
|
||||
if row.highlight:
|
||||
merged[chunk_id]["highlight"] = row.highlight
|
||||
|
||||
# 计算综合得分
|
||||
max_text_score = max(
|
||||
(r["text_score"] for r in merged.values()), default=1.0
|
||||
) or 1.0
|
||||
|
||||
scored_results = []
|
||||
for item in merged.values():
|
||||
normalized_text = item["text_score"] / max_text_score if max_text_score > 0 else 0
|
||||
combined_score = 0.7 * item["vector_score"] + 0.3 * normalized_text
|
||||
|
||||
if combined_score >= request.threshold:
|
||||
scored_results.append(
|
||||
SearchResult(
|
||||
chunk_id=item["chunk_id"],
|
||||
page_id=item["page_id"],
|
||||
page_title=item["page_title"],
|
||||
content=item["content"],
|
||||
score=round(combined_score, 4),
|
||||
course_name=item["course_name"],
|
||||
teacher_name=item["teacher_name"],
|
||||
live_date=item["live_date"],
|
||||
highlight=item["highlight"],
|
||||
)
|
||||
)
|
||||
|
||||
scored_results.sort(key=lambda x: x.score, reverse=True)
|
||||
scored_results = scored_results[:request.top_k]
|
||||
|
||||
elapsed_ms = (time.time() - start_time) * 1000
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
total=len(scored_results),
|
||||
results=scored_results,
|
||||
elapsed_ms=round(elapsed_ms, 2),
|
||||
)
|
||||
362
app/wework_bot.py
Normal file
362
app/wework_bot.py
Normal file
@@ -0,0 +1,362 @@
|
||||
"""
|
||||
企业微信智能机器人服务（WebSocket 长连接模式）
|
||||
|
||||
使用官方 SDK (wecom-aibot-python-sdk) 的 WSClient 类，
|
||||
封装了连接管理、心跳、断线重连等能力。
|
||||
|
||||
作为独立进程运行：python -m app.wework_bot
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
|
||||
import httpx
|
||||
from aibot import WSClient, WSClientOptions, generate_req_id
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# 加载 .env 文件（独立运行时需要）
|
||||
load_dotenv()
|
||||
|
||||
logger = logging.getLogger("wework_bot")
|
||||
|
||||
# 分片上传参数（参考官方文档）
|
||||
_CHUNK_SIZE = 512 * 1024 # 单个分片最大 512KB（Base64 编码前）
|
||||
|
||||
|
||||
class WeComBotService:
|
||||
"""企业微信智能机器人服务（WebSocket 长连接模式）"""
|
||||
|
||||
def __init__(self):
|
||||
self.bot_id = os.getenv("WEWORK_BOT_ID", "")
|
||||
self.bot_secret = os.getenv("WEWORK_BOT_SECRET", "")
|
||||
self.edubrain_base_url = os.getenv(
|
||||
"EDUBRAIN_BASE_URL", "http://localhost:8765"
|
||||
)
|
||||
self.search_limit = int(os.getenv("WEWORK_BOT_SEARCH_LIMIT", "3"))
|
||||
self.enabled = os.getenv("WEWORK_BOT_ENABLED", "false").lower() == "true"
|
||||
self.ws_client: WSClient | None = None
|
||||
|
||||
def start(self):
|
||||
"""启动机器人服务"""
|
||||
if not self.enabled:
|
||||
logger.info("企业微信机器人未启用 (WEWORK_BOT_ENABLED=false)")
|
||||
return
|
||||
|
||||
if not self.bot_id or not self.bot_secret:
|
||||
logger.error("WEWORK_BOT_ID 或 WEWORK_BOT_SECRET 未配置")
|
||||
return
|
||||
|
||||
self.ws_client = WSClient(
|
||||
WSClientOptions(
|
||||
bot_id=self.bot_id,
|
||||
secret=self.bot_secret,
|
||||
)
|
||||
)
|
||||
|
||||
# 注册事件
|
||||
self.ws_client.on("authenticated", self._on_authenticated)
|
||||
self.ws_client.on("message.text", self._on_text_message)
|
||||
self.ws_client.on("event.enter_chat", self._on_enter_chat)
|
||||
self.ws_client.on("disconnected", self._on_disconnected)
|
||||
self.ws_client.on("error", self._on_error)
|
||||
|
||||
logger.info("启动企业微信智能机器人...")
|
||||
self.ws_client.run()
|
||||
|
||||
# ──────────────────────────── 事件处理 ────────────────────────────
|
||||
|
||||
def _on_authenticated(self):
|
||||
"""认证成功回调"""
|
||||
logger.info("企业微信机器人认证成功")
|
||||
|
||||
async def _on_enter_chat(self, frame: dict):
|
||||
"""进入会话事件 -> 回复欢迎语"""
|
||||
await self.ws_client.reply_welcome(
|
||||
frame,
|
||||
{
|
||||
"msgtype": "text",
|
||||
"text": {
|
||||
"content": "你好！我是知识库图片搜索助手，发送关键词即可搜索相关图片。"
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
async def _on_text_message(self, frame: dict):
|
||||
"""收到文本消息 -> 搜索图片 -> 流式回复文本 + 逐张发送图片"""
|
||||
body = frame.get("body", {})
|
||||
text_content = body.get("text", {}).get("content", "")
|
||||
|
||||
# 去掉 @机器人 的前缀
|
||||
text_content = re.sub(r"@.*?\s*", "", text_content).strip()
|
||||
|
||||
if not text_content:
|
||||
stream_id = generate_req_id("stream")
|
||||
await self.ws_client.reply_stream(
|
||||
frame, stream_id, "请输入搜索关键词", True
|
||||
)
|
||||
return
|
||||
|
||||
# 清理口语化前缀，提取关键词
|
||||
keyword = re.sub(
|
||||
r"(帮我|请|找|搜|搜索|查|查一下|来|给我|要|想要|的?图|图片|相关|关于|一下)",
|
||||
"", text_content,
|
||||
).strip()
|
||||
if not keyword:
|
||||
keyword = text_content # 兜底
|
||||
|
||||
stream_id = generate_req_id("stream")
|
||||
|
||||
# 立即回复：开始搜索
|
||||
await self.ws_client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
f"\U0001f50d 正在搜索「{keyword}」...",
|
||||
False,
|
||||
)
|
||||
|
||||
# 调用 EduBrain 搜索
|
||||
try:
|
||||
results = await self._search_images(keyword)
|
||||
|
||||
if not results:
|
||||
await self.ws_client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
f"未找到与「{text_content}」相关的图片",
|
||||
True,
|
||||
)
|
||||
return
|
||||
|
||||
# 构建文本回复内容
|
||||
content_parts = [f"找到 {len(results)} 张相关图片："]
|
||||
for i, r in enumerate(results, 1):
|
||||
preview = r.get("ocr_text_preview", "")[:80]
|
||||
content_parts.append(f"\n{i}. {preview}...")
|
||||
|
||||
content = "\n".join(content_parts)
|
||||
|
||||
# 完成流式消息（仅文本，不使用 msg_item）
|
||||
await self.ws_client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
content,
|
||||
True,
|
||||
)
|
||||
|
||||
# 逐张上传临时素材并发送图片消息
|
||||
for i, r in enumerate(results, 1):
|
||||
b64 = r.get("image_base64", "")
|
||||
if not b64:
|
||||
continue
|
||||
|
||||
try:
|
||||
# 解码 base64 获取原始图片数据
|
||||
image_data = base64.b64decode(b64)
|
||||
file_name = r.get("image_url", f"image_{i}.jpg")
|
||||
|
||||
# 上传临时素材（分片上传）
|
||||
media_id = await self._upload_temp_media(
|
||||
image_data, file_name
|
||||
)
|
||||
if not media_id:
|
||||
logger.warning("第 %d 张图片上传失败，跳过", i)
|
||||
continue
|
||||
|
||||
# 发送图片消息（通过 aibot_respond_msg）
|
||||
await self.ws_client.reply(
|
||||
frame,
|
||||
{
|
||||
"msgtype": "image",
|
||||
"image": {"media_id": media_id},
|
||||
},
|
||||
)
|
||||
logger.info("已发送第 %d 张图片 (media_id: %s)", i, media_id)
|
||||
except Exception as e:
|
||||
logger.error("发送第 %d 张图片失败: %s", i, e, exc_info=True)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("搜索失败: %s", e, exc_info=True)
|
||||
await self.ws_client.reply_stream(
|
||||
frame,
|
||||
stream_id,
|
||||
f"搜索出错：{e}",
|
||||
True,
|
||||
)
|
||||
|
||||
def _on_disconnected(self, reason: str):
|
||||
"""连接断开回调"""
|
||||
logger.warning("企业微信连接断开: %s", reason)
|
||||
|
||||
def _on_error(self, error):
|
||||
"""错误回调"""
|
||||
logger.error("企业微信机器人错误: %s", error)
|
||||
|
||||
# ──────────────────────── 上传临时素材 ────────────────────────
|
||||
|
||||
async def _upload_temp_media(
|
||||
self, file_data: bytes, filename: str
|
||||
) -> str | None:
|
||||
"""
|
||||
通过 WebSocket 通道上传临时素材（三步分片上传）。
|
||||
|
||||
流程：aibot_upload_media_init → aibot_upload_media_chunk × N → aibot_upload_media_finish
|
||||
|
||||
Args:
|
||||
file_data: 文件原始二进制数据
|
||||
filename: 文件名
|
||||
|
||||
Returns:
|
||||
成功返回 media_id，失败返回 None
|
||||
"""
|
||||
total_size = len(file_data)
|
||||
total_chunks = math.ceil(total_size / _CHUNK_SIZE)
|
||||
md5 = hashlib.md5(file_data).hexdigest()
|
||||
ws = self.ws_client._ws_manager
|
||||
|
||||
# 第一步：上传初始化
|
||||
init_req_id = generate_req_id("upload")
|
||||
try:
|
||||
init_result = await ws.send_reply(
|
||||
init_req_id,
|
||||
{
|
||||
"type": "image",
|
||||
"filename": filename,
|
||||
"total_size": total_size,
|
||||
"total_chunks": total_chunks,
|
||||
"md5": md5,
|
||||
},
|
||||
cmd="aibot_upload_media_init",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("上传初始化失败: %s", e)
|
||||
return None
|
||||
|
||||
upload_id = init_result.get("body", {}).get("upload_id", "")
|
||||
if not upload_id:
|
||||
logger.error("上传初始化未返回 upload_id: %s", init_result)
|
||||
return None
|
||||
|
||||
logger.info(
|
||||
"上传初始化成功: upload_id=%s, total_size=%d, chunks=%d",
|
||||
upload_id, total_size, total_chunks,
|
||||
)
|
||||
|
||||
# 第二步：分片上传
|
||||
for chunk_idx in range(total_chunks):
|
||||
start = chunk_idx * _CHUNK_SIZE
|
||||
end = min(start + _CHUNK_SIZE, total_size)
|
||||
chunk_data = file_data[start:end]
|
||||
chunk_b64 = base64.b64encode(chunk_data).decode("utf-8")
|
||||
|
||||
chunk_req_id = generate_req_id("upload_chunk")
|
||||
try:
|
||||
await ws.send_reply(
|
||||
chunk_req_id,
|
||||
{
|
||||
"upload_id": upload_id,
|
||||
"chunk_index": chunk_idx, # 分片索引从 0 开始
|
||||
"base64_data": chunk_b64,
|
||||
},
|
||||
cmd="aibot_upload_media_chunk",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"分片上传失败 (chunk %d/%d): %s",
|
||||
chunk_idx + 1, total_chunks, e,
|
||||
)
|
||||
return None
|
||||
|
||||
logger.info("所有分片上传完成: upload_id=%s", upload_id)
|
||||
|
||||
# 第三步：完成上传
|
||||
finish_req_id = generate_req_id("upload_finish")
|
||||
try:
|
||||
finish_result = await ws.send_reply(
|
||||
finish_req_id,
|
||||
{"upload_id": upload_id},
|
||||
cmd="aibot_upload_media_finish",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("完成上传失败: %s", e)
|
||||
return None
|
||||
|
||||
media_id = finish_result.get("body", {}).get("media_id", "")
|
||||
if not media_id:
|
||||
logger.error("完成上传未返回 media_id: %s", finish_result)
|
||||
return None
|
||||
|
||||
logger.info("上传临时素材成功: media_id=%s", media_id)
|
||||
return media_id
|
||||
|
||||
# ──────────────────────────── 搜索调用 ────────────────────────────
|
||||
|
||||
async def _get_search_limit(self) -> int:
|
||||
"""从后端 API 实时获取当前 search_limit 配置"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
resp = await client.get(f"{self.edubrain_base_url}/api/v1/settings")
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
return data.get("search_limit", self.search_limit)
|
||||
except Exception as e:
|
||||
logger.warning("获取 search_limit 失败，使用默认值: %s", e)
|
||||
return self.search_limit
|
||||
|
||||
async def _search_images(self, keyword: str) -> list:
|
||||
"""
|
||||
调用 EduBrain 搜索接口（SSE 流式），收集所有匹配结果。
|
||||
|
||||
Args:
|
||||
keyword: 搜索关键词
|
||||
|
||||
Returns:
|
||||
匹配结果列表
|
||||
"""
|
||||
results = []
|
||||
url = f"{self.edubrain_base_url}/api/v1/images/search"
|
||||
limit = await self._get_search_limit()
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
async with client.stream(
|
||||
"GET",
|
||||
url,
|
||||
params={
|
||||
"keyword": keyword,
|
||||
"limit": limit,
|
||||
"sort": "time_desc",
|
||||
},
|
||||
) as resp:
|
||||
async for line in resp.aiter_lines():
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
try:
|
||||
data = json.loads(line[6:])
|
||||
if data.get("type") == "result":
|
||||
results.append(data)
|
||||
elif data.get("type") == "done":
|
||||
break
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def run_wework_bot():
|
||||
"""入口函数"""
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
|
||||
)
|
||||
service = WeComBotService()
|
||||
service.start()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_wework_bot()
|
||||
Reference in New Issue
Block a user