128 lines
4.9 KiB
Python
128 lines
4.9 KiB
Python
"""
|
||
图片查询工具 — 封装 EduBrain 图片 API
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import httpx
|
||
|
||
BASE_URL = os.getenv("IMAGE_API_BASE", "http://localhost:8765")
|
||
|
||
|
||
def _build_image_url(image_url: str) -> str:
|
||
"""拼接完整的图片访问地址"""
|
||
if image_url.startswith("http"):
|
||
return image_url
|
||
return f"{BASE_URL}{image_url}"
|
||
|
||
|
||
# ============================================================
|
||
# Tool: AI 搜索图片（SSE 流式）
|
||
# ============================================================
|
||
def search_images(keyword: str, limit: int = 5) -> str:
|
||
"""根据自然语言描述搜索图片库，返回语义匹配的图片列表。
|
||
|
||
使用 AI 语义搜索，支持模糊查询。例如搜"孩子不想上学"会匹配
|
||
包含亲子关系、厌学等相关内容的图片。
|
||
|
||
Args:
|
||
keyword: 搜索关键词，支持自然语言描述
|
||
limit: 最多返回多少条匹配结果，默认 5
|
||
"""
|
||
results = []
|
||
try:
|
||
with httpx.Client(timeout=300.0) as client:
|
||
with client.stream(
|
||
"GET",
|
||
f"{BASE_URL}/api/v1/images/search",
|
||
params={"keyword": keyword, "limit": limit},
|
||
) as resp:
|
||
for line in resp.iter_lines():
|
||
if not line.startswith("data: "):
|
||
continue
|
||
event = json.loads(line[6:])
|
||
event_type = event.get("type")
|
||
|
||
if event_type == "result":
|
||
# 透传接口返回的所有字段（含 base64、image_url 等）
|
||
results.append(event)
|
||
except httpx.ConnectError:
|
||
return json.dumps({"error": f"无法连接到图片服务（{BASE_URL}），请确认服务已启动。"}, ensure_ascii=False)
|
||
except Exception as e:
|
||
return json.dumps({"error": f"搜索图片时出错: {e}"}, ensure_ascii=False)
|
||
|
||
if not results:
|
||
return json.dumps({"keyword": keyword, "total": 0, "images": []}, ensure_ascii=False)
|
||
|
||
return json.dumps({"keyword": keyword, "total": len(results), "images": results}, ensure_ascii=False)
|
||
|
||
|
||
# ============================================================
|
||
# Tool: 获取图片详情
|
||
# ============================================================
|
||
def get_image_detail(image_id: int) -> str:
|
||
"""根据图片 ID 获取图片的详细信息，包括完整 OCR 文本。
|
||
|
||
Args:
|
||
image_id: 图片记录 ID
|
||
"""
|
||
try:
|
||
resp = httpx.get(f"{BASE_URL}/api/v1/images/{image_id}", timeout=30.0)
|
||
if resp.status_code == 404:
|
||
return json.dumps({"error": f"未找到 ID 为 {image_id} 的图片。"}, ensure_ascii=False)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
except httpx.ConnectError:
|
||
return json.dumps({"error": f"无法连接到图片服务（{BASE_URL}），请确认服务已启动。"}, ensure_ascii=False)
|
||
except Exception as e:
|
||
return json.dumps({"error": f"获取图片详情时出错: {e}"}, ensure_ascii=False)
|
||
|
||
result = {
|
||
"id": data["id"],
|
||
"image_url": _build_image_url(data.get("image_url", data.get("file_path", ""))),
|
||
"ocr_text": data.get("ocr_text", ""),
|
||
"tags": data.get("tags", []),
|
||
"confidence": data.get("confidence"),
|
||
"provider": data.get("provider"),
|
||
"created_at": data.get("created_at", ""),
|
||
}
|
||
return json.dumps(result, ensure_ascii=False)
|
||
|
||
|
||
# ============================================================
|
||
# Tool: 获取图片列表
|
||
# ============================================================
|
||
def list_images(page: int = 1, page_size: int = 20, status: str = "completed") -> str:
|
||
"""获取图片库列表，支持分页和状态筛选。
|
||
|
||
Args:
|
||
page: 页码，默认第 1 页
|
||
page_size: 每页数量，默认 20
|
||
status: 筛选状态，可选 completed/pending/processing/failed，默认 completed
|
||
"""
|
||
try:
|
||
resp = httpx.get(
|
||
f"{BASE_URL}/api/v1/images",
|
||
params={"page": page, "page_size": page_size, "status": status},
|
||
timeout=30.0,
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
except httpx.ConnectError:
|
||
return json.dumps({"error": f"无法连接到图片服务（{BASE_URL}），请确认服务已启动。"}, ensure_ascii=False)
|
||
except Exception as e:
|
||
return json.dumps({"error": f"获取图片列表时出错: {e}"}, ensure_ascii=False)
|
||
|
||
total = data.get("total", 0)
|
||
items = []
|
||
for item in data.get("items", []):
|
||
items.append({
|
||
"id": item["id"],
|
||
"image_url": _build_image_url(item.get("image_url", item.get("file_path", ""))),
|
||
"ocr_text_preview": (item.get("ocr_text") or "")[:150],
|
||
"tags": item.get("tags", []),
|
||
"created_at": item.get("created_at", ""),
|
||
})
|
||
|
||
return json.dumps({"total": total, "page": page, "images": items}, ensure_ascii=False)
|