caudate-cli 0.1.0__py3-none-any.whl

execution/tools/longcat_avatar_tool.py

@@ -0,0 +1,384 @@
+"""LongCat Avatar tool — animate a portrait with speech audio.
+
+The tool drives a local ComfyUI server running the WAN-S2V (Speech-to-Video)
+inference graph at ``cognos_mcp/workflows/wan_s2v_longcat_avatar.json``.
+The underlying model is ``LongCat-Avatar-Single`` (WAN 2.2 S2V fine-tune
+by vantagewithai), loaded from a Q5_1 GGUF.
+
+Flow:
+  1. Resolve ref_image + audio (FileStore id or absolute path)
+  2. Stage them into ComfyUI's input/ dir
+  3. Substitute placeholders + tunables in the workflow JSON
+  4. POST to ComfyUI /prompt, poll /history/<id> until done
+  5. Pull the output video, persist via FileStore, return a markdown link
+
+ComfyUI must be running. The tool does NOT auto-start it — if /system_stats
+is unreachable, it returns an actionable error pointing at the start command.
+Subprocess auto-start can be layered on later (see CONTEXT.md /
+"Caudate" → tool gap log).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import re
+import shutil
+import time
+import uuid
+from pathlib import Path
+from typing import Any
+
+import httpx
+
+from core.schemas import ToolResult
+from execution.tools.base import BaseTool
+
+logger = logging.getLogger(__name__)
+
+
+COMFYUI_ROOT = Path("/home/raveuk/comfy/ComfyUI")
+WORKFLOW_PATH = (
+    Path(__file__).resolve().parent.parent.parent
+    / "cognos_mcp" / "workflows" / "wan_s2v_longcat_avatar.json"
+)
+
+
+def _public_base_url() -> str:
+    return os.environ.get("COGNOS_PUBLIC_URL", "http://127.0.0.1:8000").rstrip("/")
+
+
+class LongCatAvatarTool(BaseTool):
+    mutates = True
+    name = "LongCatAvatar"
+    description = (
+        "Animate a portrait image with a speech audio clip to produce a "
+        "talking-head video. Provide a `ref_image` (FileStore id from a "
+        "prior [file:...] marker, or an absolute path to a portrait), an "
+        "`audio` clip (wav/mp3/flac of the speech to lip-sync), and a "
+        "`prompt` describing the avatar's appearance / setting. Returns "
+        "a `markdown` field — paste it VERBATIM into your reply so the "
+        "UI renders the video inline. Generation takes ~30-90 seconds "
+        "on a 3090 at default settings. Requires a local ComfyUI server "
+        "running on 127.0.0.1:8188 with the LongCat-Avatar GGUF + WAN "
+        "companions installed (see project memory `longcat-to-integrate`)."
+    )
+
+    def __init__(
+        self,
+        file_store: Any | None = None,
+        comfyui_url: str = "http://127.0.0.1:8188",
+        timeout_seconds: float = 600.0,
+    ):
+        self._file_store = file_store
+        self._url = comfyui_url.rstrip("/")
+        self._timeout = timeout_seconds
+        # When set, the tool can synthesize speech from a `text` arg
+        # instead of requiring a pre-recorded `audio` file. Wired by
+        # CognosAgent at startup (agent.py). Keeping it injected
+        # rather than imported avoids a hard dep on the Speak tool here.
+        self._speak_tool: Any | None = None
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "ref_image": {
+                    "type": "string",
+                    "description": "Avatar portrait. Either a FileStore id (uuid from a [file:...] marker) or an absolute path.",
+                },
+                "audio": {
+                    "type": "string",
+                    "description": "Speech clip to lip-sync. FileStore id or absolute path. WAV/MP3/FLAC. Either `audio` OR `text` must be provided.",
+                },
+                "text": {
+                    "type": "string",
+                    "description": "Text to synthesize into speech (Kokoro TTS) and use as the audio driver. Use this instead of `audio` when you want the avatar to say something the user typed. Either `audio` OR `text` must be provided; if both, `audio` wins.",
+                },
+                "voice": {
+                    "type": "string",
+                    "description": "TTS voice id (Kokoro). Only honoured when `text` is supplied. Optional.",
+                },
+                "prompt": {
+                    "type": "string",
+                    "description": "Positive prompt describing the avatar and scene (e.g. 'a young woman speaking expressively, soft studio lighting').",
+                },
+                "negative_prompt": {
+                    "type": "string",
+                    "description": "Negative prompt. Default: 'blurry, distorted, low quality, watermark, deformed face'.",
+                },
+                "width": {"type": "integer", "description": "Output width, multiple of 16. Default 832 (horizontal). Use 480 for vertical/portrait."},
+                "height": {"type": "integer", "description": "Output height, multiple of 16. Default 480."},
+                "length": {"type": "integer", "description": "Frame count. MUST satisfy (length-1) %% 4 == 0. Default 77 (~5s at 16fps); 121 ≈ 7.5s; 161 ≈ 10s."},
+                "seed": {"type": "integer", "description": "Random seed. Default 42; randomize for variety."},
+                "steps": {"type": "integer", "description": "Sampler steps. 20-30 typical. Default 20."},
+                "cfg": {"type": "number", "description": "Classifier-free guidance scale. Default 6.0."},
+                "sampler": {"type": "string", "description": "KSampler sampler_name. Default 'euler'. Try 'dpmpp_2m' or 'uni_pc' for sharper."},
+                "scheduler": {"type": "string", "description": "Sampler scheduler. Default 'simple'. Try 'beta' for cleaner motion."},
+                "fps": {"type": "number", "description": "Output framerate. Default 16 (matches WAN-S2V training)."},
+            },
+            "required": ["ref_image", "prompt"],
+        }
+
+    def set_file_store(self, file_store: Any) -> None:
+        self._file_store = file_store
+
+    def set_speak_tool(self, speak_tool: Any) -> None:
+        """Inject a SpeakTool reference so the `text` parameter can
+        drive Kokoro TTS in-process. Optional — without it, callers
+        must pass `audio` (FileStore id or path)."""
+        self._speak_tool = speak_tool
+
+    # ------------------------------------------------------------------
+
+    def _resolve_input(self, ref: str, *, expected_kind: str | None = None) -> Path:
+        """Resolve a FileStore id or filesystem path to an absolute Path.
+
+        Raises FileNotFoundError on bad ids / missing files.
+        """
+        ref = (ref or "").strip()
+        if not ref:
+            raise ValueError("empty file reference")
+        # FileStore ids are UUIDs (no '/' or '.' typical for paths)
+        if "/" not in ref and "." not in ref and self._file_store is not None:
+            rec = self._file_store.get(ref)
+            if rec is not None:
+                return Path(rec.path)
+        # Otherwise treat as a filesystem path
+        p = Path(ref).expanduser().resolve()
+        if not p.exists():
+            raise FileNotFoundError(f"Not a FileStore id and not a file on disk: {ref}")
+        return p
+
+    def _build_workflow(self, kwargs: dict[str, Any], ref_image_name: str, audio_name: str) -> dict[str, Any]:
+        """Load the template workflow + substitute placeholders + apply overrides."""
+        wf = json.loads(WORKFLOW_PATH.read_text())
+
+        # Placeholder substitutions (the 4 __*__ tokens from the template)
+        wf["5"]["inputs"]["image"] = ref_image_name
+        wf["6"]["inputs"]["audio"] = audio_name
+        wf["7"]["inputs"]["text"] = kwargs["prompt"]
+        wf["8"]["inputs"]["text"] = (
+            kwargs.get("negative_prompt")
+            or "blurry, distorted, low quality, watermark, deformed face"
+        )
+
+        # Tunable overrides — only apply when the caller passed a value.
+        def _maybe(node_id: str, field: str, key: str, cast):
+            if kwargs.get(key) is not None:
+                wf[node_id]["inputs"][field] = cast(kwargs[key])
+
+        _maybe("10", "width",       "width",     int)
+        _maybe("10", "height",      "height",    int)
+        _maybe("10", "length",      "length",    int)
+        _maybe("11", "seed",        "seed",      int)
+        _maybe("11", "steps",       "steps",     int)
+        _maybe("11", "cfg",         "cfg",       float)
+        _maybe("11", "sampler_name","sampler",   str)
+        _maybe("11", "scheduler",   "scheduler", str)
+        _maybe("13", "fps",         "fps",       float)
+
+        return wf
+
+    # ------------------------------------------------------------------
+
+    async def _ensure_comfyui_up(self, client: httpx.AsyncClient) -> str | None:
+        """Probe ComfyUI's /system_stats. Returns None if up, or an error string."""
+        try:
+            r = await client.get(f"{self._url}/system_stats", timeout=5.0)
+            r.raise_for_status()
+            return None
+        except Exception as e:
+            return (
+                f"ComfyUI not reachable at {self._url} ({type(e).__name__}: {e}). "
+                f"Start it with:\n"
+                f"  cd {COMFYUI_ROOT} && python3 main.py --listen 127.0.0.1 --port 8188\n"
+                f"Then re-run this tool."
+            )
+
+    async def _submit_and_wait(self, client: httpx.AsyncClient, workflow: dict) -> tuple[str | None, Path | None, str | None]:
+        """POST the workflow and poll /history until done. Returns (prompt_id, output_path, error_str)."""
+        client_id = uuid.uuid4().hex
+        try:
+            r = await client.post(
+                f"{self._url}/prompt",
+                json={"prompt": workflow, "client_id": client_id},
+                timeout=30.0,
+            )
+            r.raise_for_status()
+            data = r.json()
+        except httpx.HTTPStatusError as e:
+            # ComfyUI returns 400 with a JSON error body when the workflow is invalid
+            body = e.response.text if e.response is not None else "<no body>"
+            return None, None, f"ComfyUI rejected the prompt ({e.response.status_code}): {body[:500]}"
+        except Exception as e:
+            return None, None, f"ComfyUI /prompt POST failed: {type(e).__name__}: {e}"
+
+        prompt_id = data.get("prompt_id")
+        if not prompt_id:
+            return None, None, f"ComfyUI did not return a prompt_id: {data}"
+
+        # Poll /history/<id>
+        deadline = time.time() + self._timeout
+        last_status = None
+        while time.time() < deadline:
+            await asyncio.sleep(2.0)
+            try:
+                h = await client.get(f"{self._url}/history/{prompt_id}", timeout=10.0)
+                h.raise_for_status()
+                hist = h.json().get(prompt_id)
+            except Exception as e:
+                logger.debug(f"history poll failed (will retry): {e}")
+                continue
+            if not hist:
+                continue
+            status = hist.get("status", {})
+            last_status = status
+            status_str = status.get("status_str")
+            if status_str == "error":
+                msgs = status.get("messages", [])
+                return prompt_id, None, f"ComfyUI execution error: {msgs}"
+            if not status.get("completed"):
+                continue
+            # Done — find the SaveVideo output (node 14)
+            outputs = hist.get("outputs", {}) or {}
+            node_out = outputs.get("14", {}) or {}
+            for key in ("videos", "images", "gifs"):
+                items = node_out.get(key) or []
+                if items:
+                    item = items[0]
+                    out_path = (
+                        COMFYUI_ROOT / "output"
+                        / (item.get("subfolder") or "")
+                        / item["filename"]
+                    )
+                    return prompt_id, out_path, None
+            return prompt_id, None, f"ComfyUI completed but no video output found (node 14 outputs: {node_out})"
+
+        return prompt_id, None, f"timed out after {self._timeout}s (last status: {last_status})"
+
+    # ------------------------------------------------------------------
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        prompt = (kwargs.get("prompt") or "").strip()
+        if not prompt:
+            return self._error("No prompt provided")
+        if self._file_store is None:
+            return self._error("LongCatAvatar is not wired to a FileStore. Configure one on the tool.")
+
+        # 1. Resolve inputs. `audio` (FileStore id or path) wins over
+        # `text` if both supplied; if neither, error out.
+        audio_ref = (kwargs.get("audio") or "").strip()
+        text_to_speak = (kwargs.get("text") or "").strip()
+        try:
+            ref_image_src = self._resolve_input(kwargs["ref_image"])
+        except (FileNotFoundError, ValueError) as e:
+            return self._error(str(e))
+
+        if audio_ref:
+            try:
+                audio_src = self._resolve_input(audio_ref)
+            except (FileNotFoundError, ValueError) as e:
+                return self._error(str(e))
+        elif text_to_speak:
+            if self._speak_tool is None:
+                return self._error(
+                    "`text` was supplied but the Speak tool is not wired "
+                    "to this LongCatAvatar. Pass `audio` instead, or "
+                    "ensure CognosAgent wired Speak into the tool."
+                )
+            speak_kwargs: dict[str, Any] = {"text": text_to_speak}
+            if kwargs.get("voice"):
+                speak_kwargs["voice"] = kwargs["voice"]
+            speak_result = await self._speak_tool.execute(**speak_kwargs)
+            status_val = getattr(speak_result, "status", None)
+            # ToolResult.status is an enum; compare by value
+            ok = status_val == "success" or (
+                hasattr(status_val, "value") and status_val.value == "success"
+            )
+            if not ok:
+                return self._error(
+                    f"TTS step (Speak) failed: "
+                    f"{getattr(speak_result, 'error', 'unknown error')}"
+                )
+            # Speak puts the URL into the output string; extract the
+            # file_id from the `.../files/<uuid>/content` segment.
+            # (ToolResult drops metadata kwargs silently in pydantic, so
+            # we can't rely on speak_result.metadata.)
+            output_text = str(getattr(speak_result, "output", "") or "")
+            m = re.search(r"/files/([0-9a-f-]{32,36})/content", output_text)
+            if not m:
+                return self._error(
+                    "TTS step succeeded but could not find the audio "
+                    "FileStore id in Speak's output; cannot proceed."
+                )
+            audio_file_id = m.group(1)
+            try:
+                audio_src = self._resolve_input(audio_file_id)
+            except (FileNotFoundError, ValueError) as e:
+                return self._error(f"Could not resolve TTS output: {e}")
+        else:
+            return self._error(
+                "Provide either `audio` (FileStore id or path) or "
+                "`text` (to synthesize speech). Neither was supplied."
+            )
+
+        # 2. Validate length
+        length = int(kwargs.get("length") or 77)
+        if (length - 1) % 4 != 0:
+            return self._error(
+                f"length must satisfy (length - 1) %% 4 == 0 (the temporal VAE requires it). "
+                f"Got {length}. Try 77, 81, 85, …, 121, …, 161."
+            )
+
+        # 3. Stage inputs into ComfyUI/input/
+        input_dir = COMFYUI_ROOT / "input"
+        try:
+            input_dir.mkdir(parents=True, exist_ok=True)
+        except Exception as e:
+            return self._error(f"Cannot create ComfyUI input dir {input_dir}: {e}")
+        tag = uuid.uuid4().hex[:8]
+        ref_image_name = f"cognos_longcat_ref_{tag}{ref_image_src.suffix or '.png'}"
+        audio_name = f"cognos_longcat_audio_{tag}{audio_src.suffix or '.wav'}"
+        try:
+            shutil.copy2(ref_image_src, input_dir / ref_image_name)
+            shutil.copy2(audio_src, input_dir / audio_name)
+        except Exception as e:
+            return self._error(f"Failed to stage inputs into ComfyUI: {e}")
+
+        # 4. Build workflow
+        try:
+            workflow = self._build_workflow(kwargs, ref_image_name, audio_name)
+        except Exception as e:
+            return self._error(f"Failed to build workflow from template: {e}")
+
+        # 5. Run
+        async with httpx.AsyncClient() as client:
+            up_err = await self._ensure_comfyui_up(client)
+            if up_err is not None:
+                return self._error(up_err)
+            prompt_id, out_path, err = await self._submit_and_wait(client, workflow)
+
+        if err is not None:
+            return self._error(err)
+        if out_path is None or not out_path.exists():
+            return self._error(f"output file missing on disk: {out_path}")
+
+        # 6. Persist + return markdown link
+        try:
+            record = self._file_store.upload(out_path)
+        except Exception as e:
+            return self._error(f"FileStore upload of {out_path} failed: {e}")
+        url = f"{_public_base_url()}/files/{record.id}/content"
+        markdown = f"[{record.filename}]({url})"  # video, not image — no leading !
+        return self._success({
+            "file_id": record.id,
+            "filename": record.filename,
+            "size_bytes": record.size_bytes,
+            "path": str(out_path),
+            "markdown": markdown,
+            "prompt_id": prompt_id,
+        })

execution/tools/mcp_tool.py

@@ -0,0 +1,100 @@
+"""MCP visibility tool — list connected Model Context Protocol servers.
+
+Cognos already supports MCP at runtime: each `CognosAgent` keeps a
+`_mcp_clients` list of connected `MCPClient`s, and the tools they
+expose are merged into the executor's tool palette. This tool just
+makes that surface visible from chat — the agent can ask "which MCP
+servers are connected and what do they expose?" and get a clear
+answer instead of having to dig through config files.
+
+Configuration lives at `data/mcp_servers.json` (see the example in
+the same dir for shape).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from core.schemas import ToolResult
+from execution.tools.base import BaseTool
+
+logger = logging.getLogger(__name__)
+
+
+class MCPListServersTool(BaseTool):
+    name = "MCPListServers"
+    description = (
+        "List all Model Context Protocol (MCP) servers currently "
+        "connected to this Cognos agent, plus the tools each one "
+        "exposes. Useful when the user asks 'what MCP servers do I "
+        "have' or 'why isn't tool X working' — this surfaces what's "
+        "actually wired vs. only configured. No arguments. Configure "
+        "MCP servers in `data/mcp_servers.json`."
+    )
+
+    def __init__(self, agent: Any | None = None):
+        self._agent = agent
+
+    def set_agent(self, agent: Any) -> None:
+        self._agent = agent
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        return {"type": "object", "properties": {}}
+
+    async def execute(self, **_: Any) -> ToolResult:
+        if self._agent is None:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=(
+                    "MCPListServers is not wired to a CognosAgent. "
+                    "The agent must call set_agent() at startup."
+                ),
+            )
+        clients = getattr(self._agent, "_mcp_clients", None)
+        if clients is None:
+            return ToolResult(
+                tool_name=self.name, status="success",
+                output="MCP support is not active in this agent.",
+                metadata={"clients": []},
+            )
+        if not clients:
+            return ToolResult(
+                tool_name=self.name, status="success",
+                output=(
+                    "No MCP servers connected.\n"
+                    "Configure one at `data/mcp_servers.json` (see "
+                    "`data/mcp_servers.json.example` for the shape) "
+                    "and restart Cognos."
+                ),
+                metadata={"clients": []},
+            )
+
+        lines = [f"{len(clients)} MCP server(s) connected:", ""]
+        client_meta: list[dict[str, Any]] = []
+        for c in clients:
+            name = getattr(c, "name", None) or repr(c)
+            tools = []
+            try:
+                tools = c.tools() or []
+            except Exception as e:
+                logger.debug(f"MCP {name}: tools() failed: {e}")
+            tool_names = [
+                getattr(t, "name", None) or t.get("name", "?")
+                if isinstance(t, dict) else str(t)
+                for t in tools
+            ]
+            lines.append(f"  {name}  ({len(tool_names)} tool(s))")
+            for tn in tool_names[:12]:
+                lines.append(f"    - {tn}")
+            if len(tool_names) > 12:
+                lines.append(f"    ... and {len(tool_names) - 12} more")
+            client_meta.append({"name": name, "tool_count": len(tool_names),
+                                "tools": tool_names})
+
+        return ToolResult(
+            tool_name=self.name, status="success",
+            output="\n".join(lines),
+            metadata={"clients": client_meta},
+        )