odysseus/src/editor_agent_model.py

"""Editor agent/planner model resolution.

The image-editor Quick Edit planner (`/api/editor/plan`, `/api/editor/agent`)
runs on an explicitly configured local vision model when one is reachable,
with the configured chat model as a fallback. This module owns:

  * the local model base URLs / model id (settings-overridable),
  * an ordered candidate list (per-request ``agent_model`` override, then the
    local Qwen3.5-9B, then the hosted default/utility chat model),
  * a vision fallback chain: local model WITH the canvas screenshot, else a
    text-only call whose image is described by a separate VL caption model,
  * a lightweight reachability probe for the configured local bases.

Everything is best-effort: when no local server is reachable the planner
transparently falls back to the previously-used hosted chat model, so the
editor agent never regresses to "no plan".
"""

from __future__ import annotations

import contextvars
import json
import logging
import os
import re
import threading
from typing import Any, Dict, List, Optional, Tuple

import httpx

logger = logging.getLogger(__name__)

# Local endpoints are opt-in through settings or ODYSSEUS_EDITOR_AGENT_BASES.
DEFAULT_LOCAL_BASES = []
DEFAULT_LOCAL_MODEL = "Qwen/Qwen3.5-9B"

# Vision caption model used when the editor-agent model cannot accept images.
DEFAULT_CAPTION_BASE = ""
DEFAULT_CAPTION_MODEL = "Qwen/Qwen3.5-9B"

# Served-model-id cache per base (populated on first use). The local servers 400
# with "Failed to load model" when asked for an id they do not serve (e.g. the
# client's default "Qwen3.5-9B-MLX-8bit", or a stale caption model), so we
# resolve the requested id against GET {base}/models and fall back to one that
# is actually served.
_MODELS_CACHE: Dict[str, set] = {}

_CAPTION_INSTRUCTION = (
    "Describe this image in one or two sentences for a photo editor: the main "
    "subject, the background/scene, lighting, dominant colours, and any obvious "
    "flaws (blown highlights, colour cast, clutter, blur)."
)


def _setting(key: str, default: Any = "") -> Any:
    try:
        from src.settings import load_settings

        value = load_settings().get(key)
        return value if value not in (None, "") else default
    except Exception:
        return default


def local_bases() -> List[str]:
    override = _setting("editor_agent_base_url")
    if override:
        return [str(override)]
    env = os.getenv("ODYSSEUS_EDITOR_AGENT_BASES")
    if env:
        bases = [b.strip() for b in env.split(",") if b.strip()]
        if bases:
            return bases
    return list(DEFAULT_LOCAL_BASES)


def local_model() -> str:
    return str(_setting("editor_agent_model", DEFAULT_LOCAL_MODEL))


def caption_base() -> str:
    return str(_setting("editor_caption_base_url", DEFAULT_CAPTION_BASE))


def caption_model() -> str:
    return str(_setting("editor_caption_model", DEFAULT_CAPTION_MODEL))


# Per-request / per-plan context (contextvars keep this correct under the
# async event loop without sharing state between concurrent editor turns).
_override: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    "editor_agent_model_override", default=None
)
_last_model: contextvars.ContextVar[Optional[Tuple[str, str]]] = contextvars.ContextVar(
    "editor_agent_last_model", default=None
)


def set_override(value: Optional[str]) -> None:
    _override.set(value or None)


def get_override() -> Optional[str]:
    return _override.get()


def set_last_model(model: str, source: str) -> None:
    _last_model.set((model, source))


def get_last_model() -> Optional[Tuple[str, str]]:
    return _last_model.get()


def _chat_url(base: str) -> str:
    b = str(base or "").rstrip("/")
    if b.endswith("/chat/completions"):
        return b
    return b + "/chat/completions"


def _split_override(override: str) -> Tuple[str, str]:
    """Parse ``base::model`` (or a bare model id)."""
    raw = str(override or "").strip()
    if "::" in raw:
        base, model = raw.split("::", 1)
        return base.strip(), model.strip()
    return "", raw


def _has_image(messages: List[dict]) -> bool:
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and part.get("type") == "image_url":
                    return True
    return False


def _first_image(messages: List[dict]) -> Optional[str]:
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and part.get("type") == "image_url":
                    url = (part.get("image_url") or {}).get("url")
                    if isinstance(url, str):
                        return url
    return None


def _strip_images(messages: List[dict]) -> List[dict]:
    out: List[dict] = []
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            text = "\n".join(
                str(part.get("text") or "")
                for part in content
                if isinstance(part, dict) and part.get("type") == "text"
            )
            out.append({"role": msg.get("role", "user"), "content": text})
        else:
            out.append({"role": msg.get("role", "user"), "content": content})
    return out


def _with_caption(messages: List[dict], caption: str) -> List[dict]:
    note = f"\n\n[Canvas screenshot caption: {caption}]"
    out = [dict(m) for m in messages]
    for msg in reversed(out):
        if msg.get("role") == "user":
            msg["content"] = str(msg.get("content") or "") + note
            break
    return out


async def _post_chat(
    base: str,
    model: str,
    messages: List[dict],
    headers: Optional[dict],
    timeout: float,
    max_tokens: int = 1024,
    response_format: Optional[dict] = None,
) -> Optional[str]:
    if not base or not model:
        return None
    payload: Dict[str, Any] = {
        "model": model,
        "messages": messages,
        "temperature": 0.2,
        "max_tokens": max_tokens,
        "stream": False,
    }
    # Constrain the planner to a strict JSON ops schema (structured outputs /
    # grammar) so the model can only emit valid editor ops. Unknown servers
    # ignore the field; the server-side schema validation is the real enforcer.
    if isinstance(response_format, dict):
        payload["response_format"] = response_format
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(_chat_url(base), json=payload, headers=headers or {})
            if resp.status_code != 200:
                logger.debug(
                    "editor agent chat %s -> %s (%s)", base, resp.status_code, resp.text[:200]
                )
                return None
            data = resp.json()
    except Exception as exc:
        logger.debug("editor agent chat %s failed: %s", base, exc)
        return None
    try:
        return data["choices"][0]["message"]["content"]
    except Exception:
        return None


async def _served_models(base: str, timeout: float = 3.0) -> set:
    """Return the set of model ids a base serves (cached; empty on failure)."""
    if not base:
        return set()
    cached = _MODELS_CACHE.get(base)
    if cached is not None:
        return cached
    ids: set = set()
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.get(base.rstrip("/") + "/models")
            if resp.status_code == 200:
                ids = {
                    str(m.get("id"))
                    for m in (resp.json().get("data") or [])
                    if isinstance(m, dict) and m.get("id")
                }
    except Exception:
        ids = set()
    _MODELS_CACHE[base] = ids
    return ids


async def _resolve_model(base: str, requested: str, fallback: str) -> str:
    """Return a model id ``base`` actually serves.

    If the requested id is served, keep it; otherwise use ``fallback`` (when
    served) or any served id. When the served list is unknown (server down),
    keep the requested id (or ``fallback``).
    """
    requested = str(requested or "").strip()
    fallback = str(fallback or "").strip()
    ids = await _served_models(base)
    if ids:
        if requested and requested in ids:
            return requested
        if fallback in ids:
            return fallback
        return sorted(ids)[0]
    return requested or fallback


async def caption_image(data_url: str, owner: Optional[str] = None, timeout: float = 90.0) -> Optional[str]:
    if not isinstance(data_url, str) or not data_url.startswith("data:image/"):
        return None
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": _CAPTION_INSTRUCTION},
                {"type": "image_url", "image_url": {"url": data_url}},
            ],
        }
    ]
    base = caption_base()
    if not base:
        return None
    model = await _resolve_model(base, caption_model(), local_model())
    return await _post_chat(base, model, messages, None, timeout, max_tokens=200)


_DONE_RE = re.compile(r"^\s*DONE\b", re.I)
_RETRY_RE = re.compile(r"RETRY:?\s*(.+)", re.I | re.S)


def _parse_self_check(text: str) -> Dict[str, Any]:
    s = str(text or "").strip()
    if _DONE_RE.match(s):
        return {"done": True, "issues": [], "retry_hint": "", "text": s}
    m = _RETRY_RE.search(s)
    if m or re.search(r"\bretry\b", s, re.I):
        hint = (m.group(1).strip() if m else s) or s
        return {"done": False, "issues": [hint], "retry_hint": hint, "text": s}
    return {"done": True, "issues": [], "retry_hint": "", "text": s}


async def self_check(
    before: Any,
    after: Any,
    prompt: str,
    model: Optional[str] = None,
    owner: Optional[str] = None,
    timeout: float = 30.0,
) -> Dict[str, Any]:
    """Same-origin VL self-check: compare BEFORE vs AFTER against the intent.

    The browser cannot call the local VL server directly (it sends no CORS
    headers), so ``/api/editor/self-check`` proxies here. Returns
    ``{done, issues, retry_hint, text}``; on any failure (bad input, the model
    is unreachable, or a timeout) it degrades to ``done=True`` with
    ``reason="unavailable"`` so the client never gets stuck retrying.
    """
    if not (
        isinstance(before, str) and before.startswith("data:image/")
        and isinstance(after, str) and after.startswith("data:image/")
    ):
        return {"done": True, "issues": [], "retry_hint": "",
                "reason": "unavailable", "text": "DONE"}
    bases = local_bases() or list(DEFAULT_LOCAL_BASES)
    if not bases:
        return {"done": True, "issues": [], "retry_hint": "",
                "reason": "unavailable", "text": "DONE"}
    base = bases[0]
    mdl = await _resolve_model(base, str(model or local_model()), local_model())
    messages = [{
        "role": "user",
        "content": [
            {"type": "text", "text": str(prompt or "")},
            {"type": "image_url", "image_url": {"url": before}},
            {"type": "image_url", "image_url": {"url": after}},
        ],
    }]
    text = await _post_chat(base, mdl, messages, None, timeout, max_tokens=140)
    if not text:
        return {"done": True, "issues": [], "retry_hint": "",
                "reason": "unavailable", "text": "DONE"}
    out = _parse_self_check(text)
    out["model"] = mdl
    out["source"] = "local"
    return out


def _resolve_override_candidate(override: str, owner: Optional[str]) -> Optional[Dict[str, Any]]:
    """Resolve an explicit agent_model override under server authority.

    Guarantees:
      - Client-provided raw URLs never determine arbitrary outbound network destinations.
      - Registered endpoint IDs resolve only to enabled endpoints visible to the owner.
      - Shared endpoints (owner IS NULL) remain visible to authorized users.
      - Disabled, foreign-owned, forged, or unconfigured destinations fail closed (return None).
      - Admin privileges do not allow arbitrary unconfigured destinations.
    """
    raw = str(override or "").strip()
    if not raw:
        return None

    if "::" in raw:
        target, requested_model = raw.split("::", 1)
        target = target.strip()
        requested_model = requested_model.strip()
    else:
        target = ""
        requested_model = raw

    # 1. Explicit local selection or matching a configured local base
    if target == "local":
        bases = local_bases()
        if not bases:
            return None
        return {
            "base": bases[0],
            "model": requested_model or local_model(),
            "headers": None,
            "source": "override",
        }
    for lb in local_bases():
        if target and target.rstrip("/") == lb.rstrip("/"):
            return {
                "base": lb,
                "model": requested_model or local_model(),
                "headers": None,
                "source": "override",
            }

    # 2. Explicit hosted fallback selection
    try:
        from src.endpoint_resolver import resolve_endpoint

        for prefix in ("utility", "default"):
            try:
                url, model, headers = resolve_endpoint(prefix, owner=owner)
            except Exception:
                continue
            if url and model:
                if target in ("default", "hosted"):
                    return {
                        "base": url,
                        "model": requested_model or model,
                        "headers": headers,
                        "source": "hosted",
                    }
                if target and target.rstrip("/") == url.rstrip("/"):
                    if not requested_model or requested_model == model:
                        return {
                            "base": url,
                            "model": requested_model or model,
                            "headers": headers,
                            "source": "hosted",
                        }
    except Exception:
        pass
    if target in ("default", "hosted"):
        return None

    # 3. Bare model id without endpoint prefix
    if not target:
        if requested_model == "local":
            bases = local_bases()
            return {"base": bases[0], "model": local_model(), "headers": None, "source": "local"} if bases else None
        if requested_model in ("default", "hosted"):
            return _resolve_override_candidate("default::", owner)

        # Exact match with local model
        if requested_model == local_model():
            bases = local_bases()
            if bases:
                return {"base": bases[0], "model": local_model(), "headers": None, "source": "local"}

        # Exact match with hosted default/utility
        try:
            from src.endpoint_resolver import resolve_endpoint

            for prefix in ("utility", "default"):
                try:
                    url, model, headers = resolve_endpoint(prefix, owner=owner)
                except Exception:
                    continue
                if url and model and model == requested_model:
                    return {"base": url, "model": model, "headers": headers, "source": "hosted"}
        except Exception:
            pass

        # Check registered database endpoints for this bare model
        try:
            from core.database import SessionLocal, ModelEndpoint
            from src.auth_helpers import owner_filter, _auth_disabled
            from src.endpoint_resolver import (
                resolve_endpoint_runtime,
                build_chat_url,
                build_headers,
                _endpoint_enabled_models,
                _endpoint_hidden_models,
            )

            db = SessionLocal()
            try:
                if not owner and not _auth_disabled():
                    q = db.query(ModelEndpoint).filter(
                        ModelEndpoint.is_enabled == True,
                        ModelEndpoint.owner == None,
                    )
                else:
                    q = owner_filter(
                        db.query(ModelEndpoint).filter(ModelEndpoint.is_enabled == True),
                        ModelEndpoint,
                        owner,
                    )
                matching_eps = []
                for ep in q.all():
                    enabled_models = _endpoint_enabled_models(ep)
                    hidden_models = _endpoint_hidden_models(ep)
                    if requested_model in enabled_models and requested_model not in hidden_models:
                        matching_eps.append(ep)

                # Avoid model ID collision across providers: only accept if unique
                if len(matching_eps) == 1:
                    ep = matching_eps[0]
                    base, api_key = resolve_endpoint_runtime(ep, owner=owner)
                    return {
                        "base": build_chat_url(base),
                        "model": requested_model,
                        "headers": build_headers(api_key, base),
                        "source": "registered",
                        "endpoint_id": ep.id,
                    }
            finally:
                db.close()
        except Exception:
            pass
        return None

    # 4. Target is an endpoint_id or a registered legacy endpoint URL
    try:
        from core.database import SessionLocal, ModelEndpoint
        from src.auth_helpers import owner_filter, _auth_disabled
        from src.endpoint_resolver import (
            resolve_endpoint_runtime,
            build_chat_url,
            build_headers,
            _endpoint_enabled_models,
            _endpoint_hidden_models,
            _first_chat_model,
        )

        db = SessionLocal()
        try:
            if not owner and not _auth_disabled():
                q = db.query(ModelEndpoint).filter(
                    ModelEndpoint.is_enabled == True,
                    ModelEndpoint.owner == None,
                )
            else:
                q = owner_filter(
                    db.query(ModelEndpoint).filter(ModelEndpoint.is_enabled == True),
                    ModelEndpoint,
                    owner,
                )
            # Lookup by ID
            ep = q.filter(ModelEndpoint.id == target).first()

            # If not found by ID, match legacy registered base URL (strictly registered rows only)
            if ep is None and ("://" in target or target.startswith("http")):
                target_norm = target.rstrip("/")
                for cand_ep in q.all():
                    cand_base = (cand_ep.base_url or "").rstrip("/")
                    if (
                        cand_base == target_norm
                        or cand_base + "/v1" == target_norm
                        or cand_base == target_norm + "/v1"
                    ):
                        ep = cand_ep
                        break

            if ep is None:
                # Target is disabled, foreign-owned, nonexistent, or arbitrary unconfigured URL -> fail closed!
                return None

            enabled_models = _endpoint_enabled_models(ep)
            hidden_models = _endpoint_hidden_models(ep)
            if requested_model:
                if requested_model in hidden_models:
                    return None
                if enabled_models and requested_model not in enabled_models:
                    # Model not served by this endpoint -> fail closed against forged model IDs
                    return None
                model = requested_model
            else:
                model = _first_chat_model(enabled_models)
                if not model:
                    return None

            base, api_key = resolve_endpoint_runtime(ep, owner=owner)
            return {
                "base": build_chat_url(base),
                "model": model,
                "headers": build_headers(api_key, base),
                "source": "registered",
                "endpoint_id": ep.id,
            }
        finally:
            db.close()
    except Exception:
        return None


def _candidates(owner: Optional[str], override: Optional[str]) -> List[Dict[str, Any]]:
    cands: List[Dict[str, Any]] = []
    for base in local_bases():
        cands.append({"base": base, "model": local_model(), "headers": None, "source": "local"})
    try:
        from src.endpoint_resolver import resolve_endpoint

        for prefix in ("utility", "default"):
            try:
                url, model, headers = resolve_endpoint(prefix, owner=owner)
            except Exception:
                continue
            if url and model:
                cands.append({"base": url, "model": model, "headers": headers, "source": "hosted"})
                break
    except Exception:
        pass
    if override:
        cand = _resolve_override_candidate(override, owner)
        if cand is not None:
            cands.insert(0, cand)
            return cands
        # Unauthorized, invalid, foreign, disabled, or forged override:
        # Fail closed! Never silently dispatch sensitive user images/state to an unexpected fallback.
        return []
    return cands


async def plan_with_editor_agent(
    messages: List[dict],
    owner: Optional[str] = None,
    override: Optional[str] = None,
    timeout: float = 90.0,
    response_format: Optional[dict] = None,
) -> Optional[Tuple[str, str, str]]:
    """Return (raw_plan_text, model_used, source) or None.

    Tries the local model WITH the screenshot, then the same model text-only
    with a VL caption of the screenshot, then the hosted chat model.

    ``response_format`` (a JSON-schema structured-output constraint) is sent
    ONLY to the local editor-agent bases, so a hosted fallback model that does
    not support the field never rejects the request.
    """
    has_img = _has_image(messages)
    text_messages = _strip_images(messages) if has_img else messages
    image = _first_image(messages) if has_img else None
    caption: Optional[str] = None
    _local = {str(b).rstrip("/") for b in local_bases()}
    for cand in _candidates(owner, override):
        base, model = cand["base"], cand["model"]
        headers = cand.get("headers")
        is_local = cand["source"] == "local" or str(base).rstrip("/") in _local
        rf = response_format if (response_format and is_local) else None
        extra = {"response_format": rf} if rf else {}
        if has_img:
            raw = await _post_chat(base, model, messages, headers, timeout, **extra)
            if raw:
                return raw, model, cand["source"]
            if caption is None and image:
                caption = await caption_image(image, owner)
            if caption:
                raw = await _post_chat(
                    base, model, _with_caption(text_messages, caption), headers, timeout, **extra,
                )
                if raw:
                    return raw, model, cand["source"] + "+caption"
        else:
            raw = await _post_chat(
                base, model, text_messages, headers, timeout, **extra
            )
            if raw:
                return raw, model, cand["source"]
    return None


async def list_editor_agent_models(owner: Optional[str] = None) -> List[Dict[str, Any]]:
    """Models available to the editor agent, for the model picker.

    Lists the local Qwen3.5-9B (one entry per reachable base) plus the hosted
    default/utility chat model that the planner falls back to.
    """
    models: List[Dict[str, Any]] = []
    seen: set = set()
    for base in local_bases():
        model = local_model()
        key = ("local", base, model)
        if key in seen:
            continue
        seen.add(key)
        models.append({
            "id": model,
            "label": f"{model} · local",
            "source": "local",
            "vision": True,
            "base_url": base,
            "endpoint_id": "local",
        })
    try:
        from src.endpoint_resolver import resolve_endpoint
        from src.settings import get_user_setting, load_settings

        for prefix in ("utility", "default"):
            try:
                url, model, _headers = resolve_endpoint(prefix, owner=owner)
            except Exception:
                continue
            if url and model:
                key = ("hosted", model)
                if key not in seen:
                    seen.add(key)
                    ep_id = None
                    try:
                        settings = load_settings()
                        ep_id = (get_user_setting(f"{prefix}_endpoint_id", owner or "", settings.get(f"{prefix}_endpoint_id", "")) or "").strip() or None
                    except Exception:
                        pass
                    models.append({
                        "id": model,
                        "label": model,
                        "source": "hosted",
                        "vision": None,
                        "base_url": url,
                        "endpoint_id": ep_id or prefix,
                    })
                break
    except Exception:
        pass
    return models


# ---------------------------------------------------------------------------
# Reachability probe (best effort, non-blocking).
# ---------------------------------------------------------------------------

_PROBE: Dict[str, bool] = {}


def probe_models(base: str, timeout: float = 3.0) -> bool:
    try:
        with httpx.Client(timeout=timeout) as client:
            resp = client.get(str(base).rstrip("/") + "/models")
            return resp.status_code == 200
    except Exception:
        return False


def _probe_worker() -> None:
    for base in local_bases():
        ok = probe_models(base)
        _PROBE[base] = ok
        logger.info("editor agent local base %s reachable=%s", base, ok)
    cap = caption_base()
    if cap:
        _PROBE[cap] = probe_models(cap)
        logger.info("editor agent caption base %s reachable=%s", cap, _PROBE[cap])


def startup_probe() -> None:
    """Kick off a background reachability probe (never blocks startup)."""
    try:
        threading.Thread(target=_probe_worker, name="editor-agent-probe", daemon=True).start()
    except Exception:
        logger.debug("editor agent probe thread failed", exc_info=True)


def probe_status() -> Dict[str, bool]:
    return dict(_PROBE)