添加即梦

.gitignore

@@ -8,4 +8,9 @@ __pycache__/
 .venv/
 env/
 venv/
-.env
+.env
+
+.idea/
+.vscode/
+.cursor/
+.trae/

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "psutil>=6.0.0",
     "claude-agent-sdk",
     "python-dotenv>=1.0.0",
+    "filelock>=3.25.2",
 ]
 
 [project.optional-dependencies]
@@ -40,4 +41,5 @@ members = [
     "tools/local/run_comfy_workflow",
     "tools/local/task_0cd69d84",
     "tools/local/runcomfy_stop_env",
+    "tools/local/ji_meng",
 ]

tests/test_jimeng.py

@@ -0,0 +1,153 @@
+"""测试 ji_meng 任务工具 — 通过 Router API 调用（范本同 test_liblibai_tool.py）
+
+用法：
+    1. 先启动 Router：uv run python -m tool_agent
+    2. 运行测试：python tests/test_jimeng.py
+
+需已注册 ji_meng_add_task、ji_meng_query_task（或改下方常量）；搜索关键词默认可匹配二者。
+"""
+
+import sys
+import time
+
+if sys.platform == 'win32':
+    sys.stdout.reconfigure(encoding='utf-8')
+
+import httpx
+
+ROUTER_URL = "http://127.0.0.1:8001"
+POLL_INTERVAL_S = 2
+POLL_MAX_WAIT_S = 300
+
+
+def _extract_task_id(data: dict):
+    for key in ("task_id", "taskId", "id", "job_id", "jobId"):
+        v = data.get(key)
+        if v is not None and str(v).strip():
+            return str(v).strip()
+    inner = data.get("data")
+    if isinstance(inner, dict):
+        return _extract_task_id(inner)
+    return None
+
+
+def _terminal_success(data: dict) -> bool:
+    status = (
+        data.get("status")
+        or data.get("task_status")
+        or data.get("taskStatus")
+        or data.get("state")
+    )
+    if status is None and isinstance(data.get("data"), dict):
+        status = data["data"].get("status") or data["data"].get("task_status")
+    if status is None:
+        return False
+    s = str(status).lower()
+    return s in ("completed", "success", "done", "finished", "succeed", "complete")
+
+
+def _terminal_failure(data: dict) -> bool:
+    status = data.get("status") or data.get("task_status") or data.get("state")
+    if status is None and isinstance(data.get("data"), dict):
+        status = data["data"].get("status")
+    if status is None:
+        return False
+    s = str(status).lower()
+    return s in ("failed", "error", "cancelled", "canceled")
+
+
+def main():
+    print("=" * 50)
+    print("测试 ji_meng 任务工具")
+    print("=" * 50)
+
+    # 1. 检查 Router 是否在线
+    try:
+        resp = httpx.get(f"{ROUTER_URL}/health", timeout=3)
+        print(f"Router 状态: {resp.json()}")
+    except httpx.ConnectError:
+        print(f"无法连接 Router ({ROUTER_URL})")
+        print("请先启动: uv run python -m tool_agent")
+        sys.exit(1)
+
+    # 2. 搜索工具，确认已注册
+    print("\n--- 搜索工具 ---")
+    resp = httpx.post(f"{ROUTER_URL}/search_tools", json={"keyword": "ji_meng"})
+    tools = resp.json()
+    print(f"找到 {tools['total']} 个工具")
+    for t in tools["tools"]:
+        print(f"  {t['tool_id']}: {t['name']} (state={t['state']})")
+
+    # 3. 创建 ji_meng_add_task 工具
+    print("\n--- 调用 ji_meng_add_task 工具创建任务 ---")
+    print(f"提示词: simple white line art, cat, black background")
+    print("提交中...")
+
+    resp = httpx.post(
+        f"{ROUTER_URL}/select_tool",
+        json={
+            "tool_id": "ji_meng_add_task",
+            "params": {
+                "task_type": "image",
+                "prompt": "simple white line art, cat, black background",
+            },
+        },
+        timeout=120,
+    )
+
+    result = resp.json()
+    print(f"\n响应状态: {result.get('status')}")
+
+    if result.get("code") != 0:
+        print(f"错误: {result.get('msg')}")
+        sys.exit(1)
+
+    task_id = result.get("data", {}).get("task_id")
+    if not task_id:
+        print("错误: 无法从创建任务响应中解析 task_id")
+        sys.exit(1)
+    print(f"任务 ID: {task_id}")
+
+    # 4. 轮询查询任务
+    print("\n--- 调用 ji_meng_query_task 工具查询任务 ---")
+    deadline = time.monotonic() + POLL_MAX_WAIT_S
+    last = {}
+
+    while time.monotonic() < deadline:
+        resp = httpx.post(
+            f"{ROUTER_URL}/select_tool",
+            json={
+                "tool_id": "ji_meng_query_task",
+                "params": {"task_id": task_id},
+            },
+            timeout=60,
+        )
+        result = resp.json()
+        print(f"\n响应状态: {result.get('status')}")
+
+        if result.get("status") != "success":
+            print(f"错误: {result.get('error')}")
+            sys.exit(1)
+
+        last = result.get("result", {})
+        if not isinstance(last, dict):
+            print(f"非 dict 结果: {last}")
+            time.sleep(POLL_INTERVAL_S)
+            continue
+        if last.get("status") == "error":
+            print(f"错误: {last.get('error')}")
+            sys.exit(1)
+
+        print(f"任务状态: {last.get('status')}")
+        print(f"最终结果: {last}")
+        print("\n测试通过!")
+        return
+
+        time.sleep(POLL_INTERVAL_S)
+
+    print(f"\n等待超时 ({POLL_MAX_WAIT_S}s)，最后一次响应: {last}")
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

tools/local/ji_meng/.env.example

@@ -0,0 +1,12 @@
+# 即梦（或兼容异步任务）上游 HTTP — 与 ji_meng_client.py 一致
+# 复制为 .env 后填写：cp .env.example .env
+
+# 必填：上游服务根 URL（不要末尾 /）
+JI_MENG_API_BASE=
+
+# 可选：Bearer Token；不需要则留空
+JI_MENG_API_KEY=
+
+# 可选：与默认不同时再改（相对 JI_MENG_API_BASE）
+JI_MENG_ADD_TASK_PATH=/add_task
+JI_MENG_QUERY_TASK_PATH=/query_task

tools/local/ji_meng/.gitignore

@@ -0,0 +1,7 @@
+.venv/
+__pycache__/
+*.pyc
+.env
+# 历史 SQLite 占位遗留，勿提交
+ji_meng_tasks.db
+ji_meng_tasks.db-*

tools/local/ji_meng/.python-version

@@ -0,0 +1 @@
+3.12

tools/local/ji_meng/ji_meng_client.py

@@ -0,0 +1,49 @@
+"""即梦（或兼容的异步任务）上游 HTTP 客户端 — 与 liblibai_client 同类，无本地状态。"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import requests
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+class JiMengClient:
+    """通过环境变量配置上游 base URL 与路径，将请求原样转发并返回 JSON。"""
+
+    def __init__(self) -> None:
+        self.base = (os.getenv("JI_MENG_API_BASE") or "https://crawler.aiddit.com/crawler/ji_meng").rstrip("/")
+        if not self.base:
+            raise ValueError("缺少环境变量 JI_MENG_API_BASE（上游服务根 URL，如 https://api.example.com）")
+        self.add_path = os.getenv("JI_MENG_ADD_TASK_PATH", "/add_task")
+        self.query_path = os.getenv("JI_MENG_QUERY_TASK_PATH", "/query_task")
+        self.api_key = os.getenv("JI_MENG_API_KEY", "").strip()
+
+    def _headers(self) -> dict[str, str]:
+        h = {"Content-Type": "application/json"}
+        if self.api_key:
+            h["Authorization"] = f"Bearer {self.api_key}"
+        return h
+
+    def submit_task(self, prompt: str, extra: dict[str, Any] | None = None) -> dict[str, Any]:
+        url: str = f"{self.base}{self.add_path if self.add_path.startswith('/') else '/' + self.add_path}"
+        body: dict[str, Any] = {"task_type": "image", "prompt": prompt}
+        if extra:
+            body.update(extra)
+        resp = requests.post(url, json=body, headers=self._headers(), timeout=120)
+        resp.raise_for_status()
+        return resp.json()
+
+    def query_task(self, task_id: str) -> dict[str, Any]:
+        url = f"{self.base}{self.query_path if self.query_path.startswith('/') else '/' + self.query_path}"
+        resp = requests.post(
+            url,
+            json={"task_id": task_id},
+            headers=self._headers(),
+            timeout=60,
+        )
+        resp.raise_for_status()
+        return resp.json()

tools/local/ji_meng/main.py

@@ -0,0 +1,69 @@
+"""即梦任务工具 — FastAPI 调用层（范本同 liblibai_controlnet）。
+
+无本地缓存：/add_task、/query_task 直接转发到上游 HTTP（见 ji_meng_client.py）。
+
+环境变量：
+  JI_MENG_API_BASE       必填，上游根 URL
+  JI_MENG_ADD_TASK_PATH  可选，默认 /add_task
+  JI_MENG_QUERY_TASK_PATH 可选，默认 /query_task
+  JI_MENG_API_KEY        可选，Bearer Token
+
+注册（由 Agent 写入 registry + sources）：
+  tool_id=ji_meng_add_task / ji_meng_query_task，host_dir=tools/local/ji_meng，
+  endpoint_path 分别为 /add_task、/query_task。
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+
+from ji_meng_client import JiMengClient
+
+app = FastAPI(title="Ji Meng Task API")
+
+
+class AddTaskRequest(BaseModel):
+    prompt: str = Field(..., description="任务描述 / 提示词")
+    extra: dict | None = Field(default=None, description="可选，合并进提交给上游的 JSON body")
+
+
+class QueryTaskRequest(BaseModel):
+    task_id: str = Field(..., description="创建任务接口返回的任务 ID")
+
+
+@app.get("/health")
+def health() -> dict:
+    return {"status": "ok"}
+
+
+@app.post("/add_task")
+def add_task(req: AddTaskRequest) -> dict:
+    try:
+        client = JiMengClient()
+        return client.submit_task(prompt=req.prompt, extra=req.extra)
+    except ValueError as e:
+        raise HTTPException(status_code=503, detail=str(e)) from e
+    except Exception as e:
+        raise HTTPException(status_code=502, detail=str(e)) from e
+
+
+@app.post("/query_task")
+def query_task(req: QueryTaskRequest) -> dict:
+    try:
+        client = JiMengClient()
+        return client.query_task(req.task_id)
+    except ValueError as e:
+        raise HTTPException(status_code=503, detail=str(e)) from e
+    except Exception as e:
+        raise HTTPException(status_code=502, detail=str(e)) from e
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8001)
+    args = parser.parse_args()
+    uvicorn.run(app, host="0.0.0.0", port=args.port)

tools/local/ji_meng/pyproject.toml

@@ -0,0 +1,12 @@
+[project]
+name = "ji-meng"
+version = "0.1.0"
+description = "即梦任务调用层：POST /add_task、/query_task，转发上游 HTTP（无本地状态）"
+requires-python = ">=3.12"
+dependencies = [
+    "fastapi>=0.115.0",
+    "uvicorn>=0.30.0",
+    "pydantic>=2.0.0",
+    "python-dotenv>=1.0.0",
+    "requests>=2.32.0",
+]