caudate-cli 0.1.0__py3-none-any.whl

execution/tools/edit_image_tool.py

@@ -0,0 +1,262 @@
+"""EditImage — modify an existing image with FLUX img2img or Kontext.
+
+Two modes selected by the LLM (or auto-routed):
+
+  - **`reimagine`** (FLUX-schnell img2img): preserves composition,
+    restyles content. Use for "watercolor version", "anime style",
+    "more vibrant colors", "blurry photo of the same scene".
+    Strength controls how much to deviate (0=identity, 1=fresh).
+
+  - **`edit`** (FLUX.1-Kontext-dev): instruction-driven editing.
+    Use for "make the wings blue", "remove the flower", "add a
+    forest behind", "change the time to sunset". The model edits
+    only what the instruction names; the rest of the image stays.
+
+  - **`auto`** (default): keyword heuristic picks one.
+    Imperative-edit verbs ("change", "make", "remove", "add",
+    "replace", "turn") → `edit`. Otherwise → `reimagine`.
+
+Source can be a local path, a `file://` URI, or a `files/<id>` /
+`/files/<id>/content` reference resolved via FileStore. The output
+is persisted to FileStore exactly like Draw — returns a markdown
+image link the chat UI renders inline.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import urllib.parse
+import urllib.request
+from pathlib import Path
+from typing import Any
+
+from core.image import ImageEdit, edit_to_file_store, make_image_edit
+from core.schemas import ToolResult
+from execution.tools.base import BaseTool
+
+logger = logging.getLogger(__name__)
+
+
+# Keyword sets for the auto-routing heuristic. Brittle by design —
+# the LLM can override by passing `mode` explicitly. We bias toward
+# `edit` (Kontext) because in practice it preserves source content
+# far better than FLUX-schnell img2img, even for stylistic prompts.
+_EDIT_VERBS = {
+    "change", "make", "remove", "add", "replace", "turn",
+    "delete", "erase", "swap", "convert", "transform", "recolor",
+    "set", "put", "insert", "give", "fix", "edit",
+}
+_STYLE_KEYWORDS = {
+    "cartoon", "comic", "anime", "manga", "watercolor", "oil painting",
+    "pencil sketch", "pixel art", "cel shaded", "ghibli", "pop art",
+    "vintage", "noir",
+}
+# Only reach for img2img when the prompt is a *very mild* tweak —
+# Kontext doesn't lose source identity at high stylization strength.
+_REIMAGINE_HINTS_STRONG = {"slightly", "subtle", "soft tweak", "minor"}
+
+
+def _detect_mode(prompt: str) -> str:
+    p = prompt.lower().strip()
+    if any(h in p for h in _REIMAGINE_HINTS_STRONG):
+        return "reimagine"
+    # Style-transfer keywords or imperative-edit verbs both go to
+    # Kontext — it handles photos-with-faces correctly.
+    if any(k in p for k in _STYLE_KEYWORDS):
+        return "edit"
+    first = re.split(r"\s+", p, maxsplit=1)[0]
+    if first in _EDIT_VERBS:
+        return "edit"
+    return "edit"  # default — Kontext is safer than FLUX-schnell img2img
+
+
+class EditImageTool(BaseTool):
+    mutates = True
+    name = "EditImage"
+    description = (
+        "Modify an existing image. Two modes: `reimagine` (FLUX-schnell "
+        "img2img — restyle/transform while preserving composition; pass "
+        "`strength` 0-1) or `edit` (FLUX.1-Kontext — instruction edit "
+        "like 'make the wings blue' or 'remove the flower'). Use `mode` "
+        "to pick explicitly; default `auto` routes by prompt verbs. "
+        "Source: local path or 'files/<id>' from a previous Draw / "
+        "DescribeImage / upload. Output: markdown image link the chat "
+        "UI renders inline.\n\n"
+        "STRONGLY PREFER `mode=edit` (Kontext) when:\n"
+        "  - The source contains people/faces you want to preserve\n"
+        "  - You want a STRONG style change (cartoon/anime/watercolor/oil-painting/comic) on a photo\n"
+        "  - You want a targeted change ('make X blue', 'remove Y')\n"
+        "Use `mode=reimagine` only for mild stylistic tweaks where some "
+        "drift is OK. FLUX-schnell img2img has a narrow strength sweet "
+        "spot; on photos with faces it usually fabricates rather than "
+        "transforming.\n\n"
+        "NOTE: Kontext is non-commercial license — for commercial use, "
+        "fall back to `mode=reimagine`."
+    )
+
+    def __init__(
+        self,
+        file_store: Any | None = None,
+        backend: str = "diffusers",
+        backend_kwargs: dict | None = None,
+    ):
+        self._file_store = file_store
+        self._backend_name = backend
+        self._backend_kwargs = backend_kwargs or {}
+        self._backend: ImageEdit | None = None
+
+    @property
+    def input_schema(self) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "image": {
+                    "type": "string",
+                    "description": (
+                        "Source image: absolute path, 'files/<file_id>', "
+                        "or a /files/<id>/content URL from a prior tool."
+                    ),
+                },
+                "prompt": {
+                    "type": "string",
+                    "description": (
+                        "Instruction or restyling prompt. Examples: "
+                        "'make the butterfly wings blue' (edit), "
+                        "'watercolor version of this scene' (reimagine)."
+                    ),
+                },
+                "mode": {
+                    "type": "string",
+                    "enum": ["auto", "reimagine", "edit"],
+                    "description": (
+                        "auto = keyword heuristic; reimagine = "
+                        "FLUX-schnell img2img; edit = FLUX-Kontext "
+                        "instruction edit. Default auto."
+                    ),
+                    "default": "auto",
+                },
+                "strength": {
+                    "type": "number",
+                    "description": (
+                        "img2img only (mode=reimagine): how much to "
+                        "deviate from source. 0=identity, 1=fresh. "
+                        "Default 0.75."
+                    ),
+                    "default": 0.75,
+                },
+                "size": {
+                    "type": "string",
+                    "description": "Output dims, e.g. '1024x1024'. Default uses source size.",
+                },
+                "seed": {
+                    "type": "integer",
+                    "description": "Seed for reproducibility. Omit for random.",
+                },
+                "steps": {
+                    "type": "integer",
+                    "description": "Override inference steps. Defaults: 4 for reimagine, 28 for edit.",
+                },
+            },
+            "required": ["image", "prompt"],
+        }
+
+    def set_file_store(self, file_store: Any) -> None:
+        self._file_store = file_store
+
+    def _get_backend(self) -> ImageEdit:
+        if self._backend is None:
+            self._backend = make_image_edit(self._backend_name, **self._backend_kwargs)
+        return self._backend
+
+    def _resolve_image(self, ref: str) -> bytes:
+        """Resolve an image reference to raw bytes."""
+        ref = ref.strip()
+        # /files/<id>/content URL or files/<id> shorthand
+        m = re.search(r"/?files/([0-9a-f-]{8,})(?:/content)?(?:[?#].*)?$", ref)
+        if m and self._file_store is not None:
+            rec = self._file_store.get(m.group(1))
+            if rec is None:
+                raise FileNotFoundError(f"file not found in store: {m.group(1)}")
+            return Path(rec.path).read_bytes()
+        # http(s) URL — fetch
+        if ref.startswith(("http://", "https://")):
+            with urllib.request.urlopen(ref, timeout=30) as r:
+                return r.read()
+        # file:// URI
+        if ref.startswith("file://"):
+            ref = urllib.parse.urlparse(ref).path
+        # Plain path
+        path = Path(ref).expanduser().resolve()
+        if not path.exists() or not path.is_file():
+            raise FileNotFoundError(f"image file not found: {path}")
+        return path.read_bytes()
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        image_ref = (kwargs.get("image") or "").strip()
+        prompt = (kwargs.get("prompt") or "").strip()
+        if not image_ref or not prompt:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error="`image` and `prompt` are both required",
+            )
+        if self._file_store is None:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error="EditImage is not wired to a FileStore.",
+            )
+
+        mode = (kwargs.get("mode") or "auto").lower()
+        if mode == "auto":
+            mode = _detect_mode(prompt)
+            logger.info(f"EditImage auto-routed to mode={mode!r}")
+        if mode not in ("reimagine", "edit"):
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"unknown mode {mode!r}; use 'auto' / 'reimagine' / 'edit'",
+            )
+
+        try:
+            image_bytes = self._resolve_image(image_ref)
+        except (FileNotFoundError, OSError, urllib.error.URLError) as e:
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"could not read source image: {e}",
+            )
+
+        edit_kwargs: dict[str, Any] = {}
+        if "strength" in kwargs and kwargs["strength"] is not None:
+            edit_kwargs["strength"] = float(kwargs["strength"])
+        if kwargs.get("size"):
+            edit_kwargs["size"] = kwargs["size"]
+        if kwargs.get("seed") is not None:
+            edit_kwargs["seed"] = int(kwargs["seed"])
+        if kwargs.get("steps") is not None:
+            edit_kwargs["steps"] = int(kwargs["steps"])
+
+        try:
+            backend = self._get_backend()
+            record = await edit_to_file_store(
+                backend, self._file_store, image_bytes, prompt,
+                mode=mode, **edit_kwargs,
+            )
+        except Exception as e:
+            logger.exception("EditImage failed")
+            return ToolResult(
+                tool_name=self.name, status="error",
+                error=f"edit failed: {e}",
+            )
+
+        url = f"http://127.0.0.1:8000/files/{record.id}/content"
+        viewer = f"http://127.0.0.1:8000/artifact/{record.id}"
+        md = f"![{prompt[:120]}]({url})\n\n[📄 Open in viewer]({viewer})"
+        return ToolResult(
+            tool_name=self.name, status="success",
+            output=md,
+            metadata={
+                "file_id": record.id,
+                "filename": record.filename,
+                "mode": mode,
+                "prompt": prompt,
+            },
+        )

execution/tools/edit_tool.py

@@ -0,0 +1,245 @@
+"""Edit tool — surgical string replacement in a file.
+
+Two modes:
+
+  - **Normal:** `old_string` → `new_string`, single or all occurrences.
+
+  - **FIM gap-fill:** when `new_string` contains the literal marker
+    `<|FIM|>`, the tool splits `new_string` on the marker, calls the
+    project's FIM model with the surrounding file context, and fills
+    the gap before applying the replacement. Lets the agent write
+    half a function and let a code-model fill the body without
+    re-emitting the whole edit.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from core.schemas import ToolResult
+from execution.tools.base import BaseTool
+
+
+FIM_MARKER = "<|FIM|>"
+
+
+class EditTool(BaseTool):
+    mutates = True
+    name = "Edit"
+    description = (
+        "Replace an exact string in a file with a new string. The old_string "
+        "must appear EXACTLY once in the file unless replace_all is true. "
+        "Preserves all surrounding content.\n\n"
+        f"FIM mode: include the literal marker `{FIM_MARKER}` once in "
+        "new_string to have a code-model fill that gap using the "
+        "surrounding file context. Pass `fim_model` to override the "
+        "default (qwen2.5-coder:1.5b)."
+    )
+
+    @property
+    def input_schema(self) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "Path to the file to edit",
+                },
+                "old_string": {
+                    "type": "string",
+                    "description": "Exact text to replace",
+                },
+                "new_string": {
+                    "type": "string",
+                    "description": (
+                        "Replacement text (must differ from old_string). "
+                        f"Include `{FIM_MARKER}` once to trigger FIM gap-fill."
+                    ),
+                },
+                "replace_all": {
+                    "type": "boolean",
+                    "description": "If true, replace every occurrence. Default false.",
+                    "default": False,
+                },
+                "fim_model": {
+                    "type": "string",
+                    "description": (
+                        "Override FIM model id (e.g. ollama/qwen3-coder-next:q4_K_M). "
+                        "Only used when new_string contains the FIM marker."
+                    ),
+                },
+                "fim_max_tokens": {
+                    "type": "integer",
+                    "description": "Max tokens for the FIM completion (default 128).",
+                },
+            },
+            "required": ["path", "old_string", "new_string"],
+        }
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        path = kwargs.get("path", "")
+        old = kwargs.get("old_string", "")
+        new = kwargs.get("new_string", "")
+        replace_all = kwargs.get("replace_all", False)
+        fim_model = kwargs.get("fim_model")
+        fim_max_tokens = kwargs.get("fim_max_tokens")
+
+        if not path:
+            return self._error("No path provided")
+        if not old:
+            return self._error("old_string cannot be empty")
+
+        try:
+            p = Path(path)
+            content = p.read_text()
+        except FileNotFoundError:
+            return self._error(f"File not found: {path}")
+        except Exception as e:
+            return self._error(str(e))
+
+        count = content.count(old)
+        if count == 0:
+            return self._error(f"old_string not found in {path}")
+        if count > 1 and not replace_all:
+            return self._error(
+                f"old_string matches {count} places in {path}; "
+                "set replace_all=true or provide a more specific string"
+            )
+
+        # FIM gap-fill mode: marker in new_string → call code-model
+        # with the surrounding file context, splice the result in.
+        # Only one marker supported per call; replace_all+FIM is
+        # rejected because filling the same gap N times would loop the
+        # same completion into every spot.
+        fim_used = False
+        if FIM_MARKER in new:
+            if new.count(FIM_MARKER) > 1:
+                return self._error(
+                    f"new_string contains {new.count(FIM_MARKER)} FIM "
+                    "markers; only one is supported per call"
+                )
+            if replace_all:
+                return self._error(
+                    "replace_all=true is incompatible with FIM gap-fill"
+                )
+            local_prefix, local_suffix = new.split(FIM_MARKER, 1)
+            # Locate the first occurrence in the file (we've verified count >= 1).
+            idx = content.find(old)
+            file_before = content[:idx]
+            file_after = content[idx + len(old):]
+            fim_prefix = file_before + local_prefix
+            fim_suffix = local_suffix + file_after
+            try:
+                from llm.provider import fim_complete, DEFAULT_FIM_MODEL
+                gap = await fim_complete(
+                    prefix=fim_prefix,
+                    suffix=fim_suffix,
+                    model=fim_model or DEFAULT_FIM_MODEL,
+                    max_tokens=fim_max_tokens,
+                )
+            except ValueError as e:
+                return self._error(f"FIM model rejected: {e}")
+            except Exception as e:
+                return self._error(f"FIM call failed: {e}")
+            if old == local_prefix + gap + local_suffix:
+                return self._error(
+                    "FIM produced text identical to old_string — nothing to change"
+                )
+            new = local_prefix + gap + local_suffix
+            fim_used = True
+
+        if old == new:
+            return self._error("old_string and new_string are identical")
+
+        new_content = content.replace(old, new) if replace_all else content.replace(old, new, 1)
+
+        try:
+            p.write_text(new_content)
+        except Exception as e:
+            return self._error(str(e))
+
+        n = count if replace_all else 1
+        msg = f"Replaced {n} occurrence(s) in {path}"
+        if fim_used:
+            msg += " (FIM gap-fill)"
+        return self._success(msg)
+
+
+class FillInMiddleTool(BaseTool):
+    """Standalone FIM tool — routable target for Caudate's tool head.
+
+    Edit already supports FIM via the marker mode, but exposing FIM as
+    its own named tool lets Caudate's contrastive head learn to route
+    to it directly (e.g. for cases where the agent just wants the
+    completion string without applying an edit). Read-only: returns
+    the generated text, doesn't touch the filesystem.
+    """
+
+    mutates = False
+    name = "FillInMiddle"
+    description = (
+        "Generate the text that should fill the gap between `prefix` "
+        "and `suffix` using a FIM-trained code model. Returns the "
+        "completion string only — does not write to any file. Use for "
+        "code-completion lookups, structural gap-fills, or when you "
+        "want to inspect the model's suggestion before applying it. "
+        "Default model: qwen2.5-coder:1.5b (low-latency); override "
+        "with `model` for heavier completions."
+    )
+
+    @property
+    def input_schema(self) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "prefix": {
+                    "type": "string",
+                    "description": "Text before the gap.",
+                },
+                "suffix": {
+                    "type": "string",
+                    "description": "Text after the gap.",
+                    "default": "",
+                },
+                "model": {
+                    "type": "string",
+                    "description": (
+                        "Override model id (e.g. ollama/qwen3-coder-next:q4_K_M). "
+                        "Must be a FIM-trained code model on Ollama."
+                    ),
+                },
+                "max_tokens": {
+                    "type": "integer",
+                    "description": "Max tokens for the completion (default 128).",
+                },
+                "temperature": {
+                    "type": "number",
+                    "description": "Sampling temperature (default 0.1).",
+                },
+            },
+            "required": ["prefix"],
+        }
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        prefix = kwargs.get("prefix", "")
+        if not prefix:
+            return self._error("prefix is required")
+        suffix = kwargs.get("suffix", "")
+        model = kwargs.get("model")
+        max_tokens = kwargs.get("max_tokens")
+        temperature = kwargs.get("temperature")
+        try:
+            from llm.provider import fim_complete, DEFAULT_FIM_MODEL
+            completion = await fim_complete(
+                prefix=prefix,
+                suffix=suffix,
+                model=model or DEFAULT_FIM_MODEL,
+                max_tokens=max_tokens,
+                temperature=temperature,
+            )
+        except ValueError as e:
+            return self._error(str(e))
+        except Exception as e:
+            return self._error(f"FIM call failed: {e}")
+        return self._success(completion)

execution/tools/file_tool.py

@@ -0,0 +1,90 @@
+"""File tools — read and write files."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from core.schemas import ToolResult
+from execution.tools.base import BaseTool
+
+
+class FileReadTool(BaseTool):
+    name = "Read"
+    description = "Read the contents of a file. Returns the file text. Supports offset and limit for large files."
+
+    @property
+    def input_schema(self) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "Absolute or relative path to the file to read",
+                },
+                "offset": {
+                    "type": "integer",
+                    "description": "Line number to start reading from (1-based)",
+                },
+                "limit": {
+                    "type": "integer",
+                    "description": "Maximum number of lines to read",
+                },
+            },
+            "required": ["path"],
+        }
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        path = kwargs.get("path", "")
+        if not path:
+            return self._error("No path provided")
+        offset = kwargs.get("offset")
+        limit = kwargs.get("limit")
+        try:
+            content = Path(path).read_text()
+            if offset or limit:
+                lines = content.splitlines(keepends=True)
+                start = (offset - 1) if offset and offset > 0 else 0
+                end = (start + limit) if limit else len(lines)
+                content = "".join(lines[start:end])
+            return self._success(content)
+        except FileNotFoundError:
+            return self._error(f"File not found: {path}")
+        except Exception as e:
+            return self._error(str(e))
+
+
+class FileWriteTool(BaseTool):
+    mutates = True
+    name = "Write"
+    description = "Write content to a file. Creates parent directories if needed. Overwrites existing content."
+
+    @property
+    def input_schema(self) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "Path to the file to write",
+                },
+                "content": {
+                    "type": "string",
+                    "description": "Content to write to the file",
+                },
+            },
+            "required": ["path", "content"],
+        }
+
+    async def execute(self, **kwargs: Any) -> ToolResult:
+        path = kwargs.get("path", "")
+        content = kwargs.get("content", "")
+        if not path:
+            return self._error("No path provided")
+        try:
+            p = Path(path)
+            p.parent.mkdir(parents=True, exist_ok=True)
+            p.write_text(content)
+            return self._success(f"Written {len(content)} chars to {path}")
+        except Exception as e:
+            return self._error(str(e))