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