fix issue
This commit is contained in:
@@ -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(),
|
||||||
|
|||||||
@@ -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))
|
||||||
|
|||||||
@@ -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]:
|
||||||
|
|||||||
@@ -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 []
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user