diff --git a/app/api/v1/images.py b/app/api/v1/images.py index 1a665fc..43f4091 100644 --- a/app/api/v1/images.py +++ b/app/api/v1/images.py @@ -21,7 +21,7 @@ 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.database import get_db, async_session_factory from app.models.base import OCRImage from app.services.ocr_service import OCRService from app.services.llm_service import LLMService @@ -299,15 +299,11 @@ 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 后结束 + 注意：不使用 Depends(get_db)，因为 StreamingResponse 的生命周期 + 超出 request scope，需要在生成器内部自行管理 session。 """ 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" - # 先获取总记录数 - count_result = await db.execute( - select(func.count(OCRImage.id)).where(OCRImage.status == "completed") - ) - total_records = count_result.scalar() or 0 + # 在生成器内部自行管理 session，避免 get_db 提前关闭 + async with async_session_factory() as db: + # 先获取总记录数 + 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 + 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()") + # 排序 + 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 + # 用于收集结果的队列 + 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) + 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) + checked_count += len(batch_items) - # 发送进度 - progress_data = { - "type": "progress", - "checked": checked_count, - "total": total_records, - "found": found, - } - await result_queue.put(progress_data) + # 发送进度 + progress_data = { + "type": "progress", + "checked": checked_count, + "total": total_records, + "found": found, + } + 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: - 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) + # 找够了，发送结束事件 + 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" - # 启动所有批次的异步任务 - 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" + # 发送结束事件 + yield f"data: {json.dumps({'type': 'done', 'total_found': found, 'total_checked': checked_count}, ensure_ascii=False)}\n\n" return StreamingResponse( event_stream(), diff --git a/app/main.py b/app/main.py index cdf97a5..995c33e 100644 --- a/app/main.py +++ b/app/main.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio import json import logging +import os import threading from contextlib import asynccontextmanager from pathlib import Path @@ -193,7 +194,11 @@ async def health_check(): @app.get("/data/images/{image_name}", tags=["前端"]) async def serve_image(image_name: str): """访问 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(): return JSONResponse(status_code=404, content={"detail": "图片不存在"}) return FileResponse(str(img_path)) diff --git a/app/models/base.py b/app/models/base.py index 1c47e87..5a5205c 100644 --- a/app/models/base.py +++ b/app/models/base.py @@ -81,6 +81,7 @@ class OCRImage(Base, TimestampMixin): String(20), default="pending", comment="状态: pending/processing/completed/failed" ) blocks: Mapped[str | None] = mapped_column(Text, nullable=True, comment="OCR 文本块（JSON）") + error_message: Mapped[str | None] = mapped_column(Text, nullable=True, comment="错误信息") @property def blocks_list(self) -> list[dict]: diff --git a/app/services/llm_service.py b/app/services/llm_service.py index 2b967a2..d248591 100644 --- a/app/services/llm_service.py +++ b/app/services/llm_service.py @@ -286,37 +286,41 @@ class LLMService: 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) + http_client = httpx.AsyncClient(timeout=60.0) + try: + client = AsyncOpenAI( + api_key=settings.MINIMAX_API_KEY, + base_url=settings.MINIMAX_BASE_URL, + http_client=http_client, + ) + 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) + # 解析返回的编号 + numbers = re.findall(r'\d+', content) - if not numbers: - return [] + 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"]) + 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 + return matched_ids + finally: + await http_client.aclose() except Exception as exc: logger.warning("LLM 批量判断失败: %s", exc) return [] diff --git a/app/services/ocr_service.py b/app/services/ocr_service.py index 8db60fa..d3f48b1 100644 --- a/app/services/ocr_service.py +++ b/app/services/ocr_service.py @@ -102,7 +102,7 @@ class PaddleOCRProvider(OCRProviderBase): async def recognize(self, image_path: str) -> OCRResult: """使用 PaddleOCR 识别图片""" self._init_ocr() - loop = asyncio.get_event_loop() + loop = asyncio.get_running_loop() def _run_ocr(): result = self._ocr.ocr(image_path, cls=True) @@ -164,7 +164,7 @@ class AliyunOCRProvider(OCRProviderBase): import asyncio - loop = asyncio.get_event_loop() + loop = asyncio.get_running_loop() def _call_api(): from alibabacloud_tea_openapi.models import Config @@ -249,7 +249,7 @@ class TencentOCRProvider(OCRProviderBase): import base64 import json - loop = asyncio.get_event_loop() + loop = asyncio.get_running_loop() def _call_api(): from tencentcloud.common import credential