fix issue

This commit is contained in:
EduBrain Dev
2026-04-14 21:25:41 +08:00
parent 23eabca58d
commit 4d5665d369
5 changed files with 145 additions and 137 deletions

View File

@@ -21,7 +21,7 @@ from sqlalchemy import select, func, text, delete
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from app.config import settings from app.config import settings
from app.database import get_db from app.database import get_db, async_session_factory
from app.models.base import OCRImage from app.models.base import OCRImage
from app.services.ocr_service import OCRService from app.services.ocr_service import OCRService
from app.services.llm_service import LLMService from app.services.llm_service import LLMService
@@ -299,15 +299,11 @@ async def search_images(
keyword: str, keyword: str,
limit: int = Query(5, ge=1, le=100), limit: int = Query(5, ge=1, le=100),
sort: str = Query("time_desc"), sort: str = Query("time_desc"),
db: AsyncSession = Depends(get_db),
): ):
""" """
AI 搜索图片SSE 流式返回)。 AI 搜索图片SSE 流式返回)。
1. 从 DB 取所有图片,按指定排序 注意:不使用 Depends(get_db),因为 StreamingResponse 的生命周期
2. 每批 5 条发给 Qwen3-8B 判断 超出 request scope需要在生成器内部自行管理 session。
3. 并发池最多 10 个请求
4. 匹配结果实时 SSE 推送
5. 找够 limit 个或遍历完 DB 后结束
""" """
async def event_stream(): async def event_stream():
@@ -320,114 +316,116 @@ async def search_images(
# 发送开始事件 # 发送开始事件
yield f"data: {json.dumps({'type': 'start', 'keyword': keyword, 'limit': limit}, ensure_ascii=False)}\n\n" yield f"data: {json.dumps({'type': 'start', 'keyword': keyword, 'limit': limit}, ensure_ascii=False)}\n\n"
# 先获取总记录数 # 在生成器内部自行管理 session避免 get_db 提前关闭
count_result = await db.execute( async with async_session_factory() as db:
select(func.count(OCRImage.id)).where(OCRImage.status == "completed") # 先获取总记录数
) count_result = await db.execute(
total_records = count_result.scalar() or 0 select(func.count(OCRImage.id)).where(OCRImage.status == "completed")
)
total_records = count_result.scalar() or 0
if total_records == 0: if total_records == 0:
yield f"data: {json.dumps({'type': 'done', 'total_found': 0, 'total_checked': 0}, ensure_ascii=False)}\n\n" yield f"data: {json.dumps({'type': 'done', 'total_found': 0, 'total_checked': 0}, ensure_ascii=False)}\n\n"
return return
# 排序 # 排序
order_clause = OCRImage.created_at.desc() order_clause = OCRImage.created_at.desc()
if sort == "time_asc": if sort == "time_asc":
order_clause = OCRImage.created_at.asc() order_clause = OCRImage.created_at.asc()
elif sort == "random": elif sort == "random":
order_clause = text("RANDOM()") order_clause = text("RANDOM()")
# 用于收集结果的队列 # 用于收集结果的队列
result_queue = asyncio.Queue() result_queue = asyncio.Queue()
checked_count = 0 checked_count = 0
async def process_batch(batch_items: list, batch_offset: int): async def process_batch(batch_items: list, batch_offset: int):
"""处理一批图片:发给 LLM 判断""" """处理一批图片:发给 LLM 判断"""
nonlocal checked_count nonlocal checked_count
async with semaphore: async with semaphore:
articles = [{"id": item.id, "text": item.ocr_text or ""} for item in batch_items] articles = [{"id": item.id, "text": item.ocr_text or ""} for item in batch_items]
matched_ids = await LLMService.judge_batch(keyword, articles) matched_ids = await LLMService.judge_batch(keyword, articles)
checked_count += len(batch_items) checked_count += len(batch_items)
# 发送进度 # 发送进度
progress_data = { progress_data = {
"type": "progress", "type": "progress",
"checked": checked_count, "checked": checked_count,
"total": total_records, "total": total_records,
"found": found, "found": found,
} }
await result_queue.put(progress_data) await result_queue.put(progress_data)
# 如果有匹配且未达上限,发送结果 # 如果有匹配且未达上限,发送结果
for match_id in matched_ids: 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,
"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: if found >= limit:
break # 找够了,发送结束事件
item = next((i for i in batch_items if i.id == match_id), None) yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n"
if item: return
file_name = item.file_path.split('/')[-1] if item.file_path else '' elif data.get("type") == "progress":
# 读取图片并转为 base64用于发送到企业微信等 yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
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,
"confidence": item.confidence,
"provider": item.provider,
"created_at": str(item.created_at),
}
await result_queue.put(result_data)
# 启动所有批次的异步任务 # 发送结束事件
tasks = [] yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n"
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( return StreamingResponse(
event_stream(), event_stream(),

View File

@@ -8,6 +8,7 @@ from __future__ import annotations
import asyncio import asyncio
import json import json
import logging import logging
import os
import threading import threading
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from pathlib import Path from pathlib import Path
@@ -193,7 +194,11 @@ async def health_check():
@app.get("/data/images/{image_name}", tags=["前端"]) @app.get("/data/images/{image_name}", tags=["前端"])
async def serve_image(image_name: str): async def serve_image(image_name: str):
"""访问 images 目录下的图片文件(用于前端预览)""" """访问 images 目录下的图片文件(用于前端预览)"""
img_path = settings.images_dir / image_name # 防止路径遍历攻击
safe_name = os.path.basename(image_name)
if safe_name != image_name or ".." in image_name:
return JSONResponse(status_code=400, content={"detail": "非法文件名"})
img_path = settings.images_dir / safe_name
if not img_path.exists(): if not img_path.exists():
return JSONResponse(status_code=404, content={"detail": "图片不存在"}) return JSONResponse(status_code=404, content={"detail": "图片不存在"})
return FileResponse(str(img_path)) return FileResponse(str(img_path))

View File

@@ -81,6 +81,7 @@ class OCRImage(Base, TimestampMixin):
String(20), default="pending", comment="状态: pending/processing/completed/failed" String(20), default="pending", comment="状态: pending/processing/completed/failed"
) )
blocks: Mapped[str | None] = mapped_column(Text, nullable=True, comment="OCR 文本块JSON") blocks: Mapped[str | None] = mapped_column(Text, nullable=True, comment="OCR 文本块JSON")
error_message: Mapped[str | None] = mapped_column(Text, nullable=True, comment="错误信息")
@property @property
def blocks_list(self) -> list[dict]: def blocks_list(self) -> list[dict]:

View File

@@ -286,37 +286,41 @@ class LLMService:
try: try:
from openai import AsyncOpenAI from openai import AsyncOpenAI
client = AsyncOpenAI( http_client = httpx.AsyncClient(timeout=60.0)
api_key=settings.MINIMAX_API_KEY, try:
base_url=settings.MINIMAX_BASE_URL, client = AsyncOpenAI(
http_client=httpx.AsyncClient(timeout=60.0), api_key=settings.MINIMAX_API_KEY,
) base_url=settings.MINIMAX_BASE_URL,
response = await client.chat.completions.create( http_client=http_client,
model=settings.MINIMAX_CHAT_MODEL, )
messages=messages, response = await client.chat.completions.create(
temperature=0.0, model=settings.MINIMAX_CHAT_MODEL,
max_tokens=512, messages=messages,
stream=False, temperature=0.0,
extra_body={"reasoning_split": True}, max_tokens=512,
) stream=False,
content = response.choices[0].message.content or "" extra_body={"reasoning_split": True},
content = cls._clean_reasoning_content(content) )
content = response.choices[0].message.content or ""
content = cls._clean_reasoning_content(content)
# 解析返回的编号 # 解析返回的编号
numbers = re.findall(r'\d+', content) numbers = re.findall(r'\d+', content)
if not numbers: if not numbers:
return [] return []
matched_ids = [] matched_ids = []
for num_str in numbers: for num_str in numbers:
num = int(num_str) num = int(num_str)
if num == 0: if num == 0:
continue # 0 表示都不匹配 continue # 0 表示都不匹配
if 1 <= num <= len(articles): if 1 <= num <= len(articles):
matched_ids.append(articles[num - 1]["id"]) matched_ids.append(articles[num - 1]["id"])
return matched_ids return matched_ids
finally:
await http_client.aclose()
except Exception as exc: except Exception as exc:
logger.warning("LLM 批量判断失败: %s", exc) logger.warning("LLM 批量判断失败: %s", exc)
return [] return []

View File

@@ -102,7 +102,7 @@ class PaddleOCRProvider(OCRProviderBase):
async def recognize(self, image_path: str) -> OCRResult: async def recognize(self, image_path: str) -> OCRResult:
"""使用 PaddleOCR 识别图片""" """使用 PaddleOCR 识别图片"""
self._init_ocr() self._init_ocr()
loop = asyncio.get_event_loop() loop = asyncio.get_running_loop()
def _run_ocr(): def _run_ocr():
result = self._ocr.ocr(image_path, cls=True) result = self._ocr.ocr(image_path, cls=True)
@@ -164,7 +164,7 @@ class AliyunOCRProvider(OCRProviderBase):
import asyncio import asyncio
loop = asyncio.get_event_loop() loop = asyncio.get_running_loop()
def _call_api(): def _call_api():
from alibabacloud_tea_openapi.models import Config from alibabacloud_tea_openapi.models import Config
@@ -249,7 +249,7 @@ class TencentOCRProvider(OCRProviderBase):
import base64 import base64
import json import json
loop = asyncio.get_event_loop() loop = asyncio.get_running_loop()
def _call_api(): def _call_api():
from tencentcloud.common import credential from tencentcloud.common import credential