重构七天训练营老师

image_agent/image_tools.py

@@ -0,0 +1,127 @@
+"""
+图片查询工具 — 封装 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)