tiny_trigger/llm.py

from __future__ import annotations

import json
import re
from typing import Any

from pydantic import BaseModel, ValidationError

from .automation import AutomationDocument, automation_schema


SYSTEM_PROMPT = """You compile home automation requests into Tiny Trigger rules.
Return JSON only. Never return code, markdown, explanations, or tool calls.
The root object MUST include a non-empty "rules" array.
Use video conditions in when: present, count, near, far, moving.
Use state gates in gate: enabled, cooldown.
Use trigger.on for edge behavior: while, enter, exit, change.
Use only these action types: simulate, webhook.
Use trigger.on="enter" for state assertions like "must be on", "should be on", "keep on", "turn on when", or "notify when".
Use trigger.on="exit" with a plain then action list for requests like "when the person leaves", "when it disappears", or "when it stops meeting the condition".
Use trigger.on="while" only when the user explicitly wants repeated actions while a condition remains true, usually with a cooldown.
When the request says one object is near, next to, beside, at, by, close to, or in front of another object, you MUST emit a near condition.
Do not replace a near relation with two present conditions.
Use max_gap_percent for near/far box-edge distance. It is the largest horizontal/vertical edge gap between boxes in normalized frame percent; touching or overlapping boxes have gap 0.
Use moving for simple same-object displacement across sampled frames, such as "car moving" or "person walks". Use min_displacement_ratio, default 0.15 (displacement as a fraction of the object's own box size), window_frames, minimum/default 3, and max_missing_frames, default 1.
Do not generate speed, direction, long-gap re-identification, or trajectory path rules.
If the user mentions elapsed time since an action or limiting repeat fires, encode it as gate.cooldown.
If the user asks for one action when a condition starts and another action when it stops, use trigger.on="change" and then.enter / then.exit.
"""

DEFAULT_REPLICATE_MODEL = "openai/gpt-5.2"
DEFAULT_OPENAI_MODEL = "gpt-5.5"
DEFAULT_ANTHROPIC_MODEL = "claude-sonnet-4-6"


class LLMCompileResult(BaseModel):
    raw_text: str
    document: AutomationDocument


def _invalid_json_error(provider: str, raw_text: str, error: Exception) -> ValueError:
    preview = raw_text.strip()
    if len(preview) > 1200:
        preview = preview[:1200] + "..."
    return ValueError(
        f"{provider} returned text that was not valid Tiny Trigger JSON. "
        f"Validation error: {error}. Raw response: {preview}"
    )


def compile_automation_with_replicate(
    *,
    instruction: str,
    class_names: list[str],
    api_token: str,
    model: str = DEFAULT_REPLICATE_MODEL,
    reasoning_effort: str = "medium",
    timeout: float = 600.0,
) -> LLMCompileResult:
    """Compile natural language into validated rules through Replicate."""
    api_token = _clean_api_key(api_token, "Replicate")
    try:
        import replicate
    except ImportError as exc:  # pragma: no cover - dependency guard
        raise RuntimeError("Install replicate to use the Replicate compiler.") from exc

    user_prompt = _build_user_prompt(instruction=instruction, class_names=class_names)
    raw_text = _stream_replicate_completion(
        model=model,
        prompt=_provider_prompt(user_prompt),
        api_token=api_token,
        reasoning_effort=reasoning_effort,
        timeout=timeout,
        replicate_module=replicate,
    )
    try:
        return _validate_compile_result(raw_text)
    except (json.JSONDecodeError, ValidationError, ValueError) as exc:
        raise _invalid_json_error("Replicate", raw_text, exc) from exc


def compile_automation_with_openai(
    *,
    instruction: str,
    class_names: list[str],
    api_key: str,
    model: str = DEFAULT_OPENAI_MODEL,
    timeout: float = 120.0,
) -> LLMCompileResult:
    """Compile natural language into validated rules through OpenAI."""
    api_key = _clean_api_key(api_key, "OpenAI")
    user_prompt = _build_user_prompt(instruction=instruction, class_names=class_names)
    try:
        import openai
    except ImportError:
        try:
            import requests
        except ImportError as exc:  # pragma: no cover - dependency guard
            raise RuntimeError("Install openai or requests to use the OpenAI compiler.") from exc
        raw_text = _post_openai_chat_completion(
            endpoint="https://api.openai.com/v1/chat/completions",
            api_key=api_key,
            model=model,
            user_prompt=user_prompt,
            timeout=timeout,
            requests_module=requests,
        )
    else:
        raw_text = _openai_chat_completion(
            api_key=api_key,
            model=model,
            user_prompt=user_prompt,
            timeout=timeout,
            openai_module=openai,
        )
    try:
        return _validate_compile_result(raw_text)
    except (json.JSONDecodeError, ValidationError, ValueError) as exc:
        raise _invalid_json_error("OpenAI", raw_text, exc) from exc


def compile_automation_with_anthropic(
    *,
    instruction: str,
    class_names: list[str],
    api_key: str,
    model: str = DEFAULT_ANTHROPIC_MODEL,
    timeout: float = 120.0,
) -> LLMCompileResult:
    """Compile natural language into validated rules through Anthropic Claude."""
    api_key = _clean_api_key(api_key, "Anthropic")
    user_prompt = _build_user_prompt(instruction=instruction, class_names=class_names)
    try:
        import anthropic
    except ImportError:
        try:
            import requests
        except ImportError as exc:  # pragma: no cover - dependency guard
            raise RuntimeError("Install anthropic or requests to use the Anthropic compiler.") from exc
        raw_text = _post_anthropic_message(
            endpoint="https://api.anthropic.com/v1/messages",
            api_key=api_key,
            model=model,
            user_prompt=user_prompt,
            timeout=timeout,
            requests_module=requests,
        )
    else:
        raw_text = _anthropic_message(
            api_key=api_key,
            model=model,
            user_prompt=user_prompt,
            timeout=timeout,
            anthropic_module=anthropic,
        )
    try:
        return _validate_compile_result(raw_text)
    except (json.JSONDecodeError, ValidationError, ValueError) as exc:
        raise _invalid_json_error("Anthropic", raw_text, exc) from exc


def _provider_prompt(user_prompt: str) -> str:
    return f"{SYSTEM_PROMPT}\n\n{user_prompt}\n\nReturn only the JSON object."


def _stream_replicate_completion(
    *,
    model: str,
    prompt: str,
    api_token: str,
    reasoning_effort: str,
    timeout: float,
    replicate_module: Any,
) -> str:
    _split_replicate_model(model)
    payload = {
        "prompt": prompt,
        "messages": [],
        "verbosity": "medium",
        "reasoning_effort": reasoning_effort,
    }
    client = replicate_module.Client(api_token=api_token)
    chunks: list[str] = []
    try:
        for event in client.stream(model, input=payload):
            chunks.append(str(event))
    except Exception as exc:
        raise RuntimeError(_provider_exception_message("Replicate", exc)) from exc
    text = "".join(chunks).strip()
    if not text:
        raise ValueError("Replicate stream returned no output.")
    return text


def _split_replicate_model(model: str) -> tuple[str, str]:
    parts = model.strip().split("/", 1)
    if len(parts) != 2 or not all(parts):
        raise ValueError("Replicate model must be in owner/model format, for example openai/gpt-5.2.")
    return parts[0], parts[1]


def _post_openai_chat_completion(
    *,
    endpoint: str,
    api_key: str,
    model: str,
    user_prompt: str,
    timeout: float,
    requests_module: Any,
) -> str:
    payload = _chat_payload(model=model, user_prompt=user_prompt, response_format="json_object")
    try:
        response = requests_module.post(
            endpoint,
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            json=payload,
            timeout=timeout,
        )
        response.raise_for_status()
    except Exception as exc:
        raise RuntimeError(_provider_exception_message("OpenAI", exc)) from exc
    body = response.json()
    return body["choices"][0]["message"]["content"]


def _openai_chat_completion(
    *,
    api_key: str,
    model: str,
    user_prompt: str,
    timeout: float,
    openai_module: Any,
) -> str:
    client = openai_module.OpenAI(api_key=api_key, timeout=timeout)
    try:
        response = client.chat.completions.create(
            **_chat_payload(model=model, user_prompt=user_prompt, response_format="json_object")
        )
    except Exception as exc:
        raise RuntimeError(_provider_exception_message("OpenAI", exc)) from exc
    content = response.choices[0].message.content
    if isinstance(content, list):
        return "".join(str(part.get("text", "")) for part in content if isinstance(part, dict))
    return str(content or "")


def _post_anthropic_message(
    *,
    endpoint: str,
    api_key: str,
    model: str,
    user_prompt: str,
    timeout: float,
    requests_module: Any,
) -> str:
    try:
        response = requests_module.post(
            endpoint,
            headers={
                "x-api-key": api_key,
                "anthropic-version": "2023-06-01",
                "Content-Type": "application/json",
            },
            json={
                "model": model,
                "system": SYSTEM_PROMPT,
                "messages": [{"role": "user", "content": user_prompt}],
                "max_tokens": 512,
                "temperature": 0,
            },
            timeout=timeout,
        )
        response.raise_for_status()
    except Exception as exc:
        raise RuntimeError(_provider_exception_message("Anthropic", exc)) from exc
    body = response.json()
    chunks = body.get("content") or []
    return "".join(str(chunk.get("text", "")) for chunk in chunks if isinstance(chunk, dict))


def _anthropic_message(
    *,
    api_key: str,
    model: str,
    user_prompt: str,
    timeout: float,
    anthropic_module: Any,
) -> str:
    client = anthropic_module.Anthropic(api_key=api_key, timeout=timeout)
    try:
        response = client.messages.create(
            model=model,
            system=SYSTEM_PROMPT,
            messages=[{"role": "user", "content": user_prompt}],
            max_tokens=512,
            temperature=0,
        )
    except Exception as exc:
        raise RuntimeError(_provider_exception_message("Anthropic", exc)) from exc
    chunks = getattr(response, "content", []) or []
    texts: list[str] = []
    for chunk in chunks:
        text = getattr(chunk, "text", None)
        if text is None and isinstance(chunk, dict):
            text = chunk.get("text")
        if text:
            texts.append(str(text))
    return "".join(texts)


def _validate_compile_result(raw_text: str) -> LLMCompileResult:
    data = json.loads(extract_json_object(raw_text))
    document = AutomationDocument.model_validate(data)
    if not document.rules:
        raise ValueError("LLM response must include a non-empty rules array.")
    return LLMCompileResult(raw_text=raw_text, document=document)


def _clean_api_key(api_key: str | None, provider: str) -> str:
    cleaned = (api_key or "").strip()
    if not cleaned:
        raise ValueError(f"Paste a {provider} API key.")
    if any(char in cleaned for char in ("\n", "\r", "\t")):
        raise ValueError(f"{provider} API key must be a single-line token with no whitespace.")
    if "Traceback" in cleaned or 'File "' in cleaned:
        raise ValueError(
            f"{provider} API key field looks like it contains a pasted error log, not an API key."
        )
    return cleaned


def _provider_exception_message(provider: str, exc: Exception) -> str:
    parts = [f"{provider} API request failed"]
    status = getattr(exc, "status", None) or getattr(exc, "status_code", None)
    if status:
        parts.append(f"status {status}")
    detail = getattr(exc, "detail", None)
    if detail:
        parts.append(str(detail))
    response = getattr(exc, "response", None)
    if response is not None:
        response_status = getattr(response, "status_code", None)
        if response_status and not status:
            parts.append(f"status {response_status}")
        try:
            body = response.json()
        except Exception:
            body = getattr(response, "text", "")
        if body:
            parts.append(str(body))
    if len(parts) == 1:
        parts.append(str(exc))
    message = ". ".join(part for part in parts if part)
    status_text = str(status or "")
    if response is not None:
        response_status = getattr(response, "status_code", None)
        if response_status:
            status_text = str(response_status)
    lower_message = message.lower()
    if status_text == "429" or "throttled" in lower_message or "rate limit" in lower_message:
        message += (
            ". This is a provider rate limit, separate from account credits. "
            "Wait a bit or switch provider."
        )
    if status_text == "404" or "not found" in lower_message:
        message += (
            ". This usually means the configured model is unavailable, deprecated, "
            "or misspelled for this provider."
        )
    return message


def extract_json_object(text: str) -> str:
    stripped = text.strip()
    if stripped.startswith("```"):
        stripped = re.sub(r"^```(?:json)?", "", stripped, flags=re.IGNORECASE).strip()
        stripped = re.sub(r"```$", "", stripped).strip()
    if stripped.startswith("{") and stripped.endswith("}"):
        return stripped

    match = re.search(r"\{.*\}", stripped, flags=re.DOTALL)
    if not match:
        raise ValueError("LLM response did not contain a JSON object.")
    return match.group(0)


def _build_user_prompt(*, instruction: str, class_names: list[str]) -> str:
    class_hint = ", ".join(class_names) if class_names else "Use labels from the request."
    return f"""Available detection labels: {class_hint}

Automation request:
{instruction}

Return a JSON object matching this high-level shape:
{{
  "rules": [
    {{
      "name": "short-kebab-case-name",
      "when": {{"all": [{{"present": {{"label": "cat", "min_count": 1}}}}]}},
      "trigger": {{"on": "enter"}},
      "gate": {{"enabled": true}},
      "then": [{{"type": "simulate", "name": "action name"}}]
    }}
  ]
}}

Examples:

User: If person near steering wheel then you have to turn on pc.
JSON:
{{
  "rules": [
    {{
      "name": "person-near-steering-wheel",
      "when": {{
        "all": [
          {{"present": {{"label": "person", "min_count": 1}}}},
          {{"near": {{"a": "person", "b": "steering wheel", "max_gap_percent": 16}}}}
        ]
      }},
      "trigger": {{"on": "enter"}},
      "gate": {{"enabled": true}},
      "then": [{{"type": "simulate", "name": "turn on pc"}}]
    }}
  ]
}}

User: If package is at door and 15 minutes since last notification, notify me.
JSON:
{{
  "rules": [
    {{
      "name": "package-at-door",
      "when": {{
        "all": [
          {{"present": {{"label": "package", "min_count": 1}}}},
          {{"near": {{"a": "package", "b": "door", "max_gap_percent": 16}}}}
        ]
      }},
      "trigger": {{"on": "while"}},
      "gate": {{"enabled": true, "cooldown": {{"key": "package-at-door", "minutes": 15}}}},
      "then": [{{"type": "simulate", "name": "notify me"}}]
    }}
  ]
}}

User: While there is a guitar in the scene, amplifier must be on.
JSON:
{{
  "rules": [
    {{
      "name": "guitar-amplifier-on",
      "when": {{
        "all": [
          {{"present": {{"label": "guitar", "min_count": 1}}}}
        ]
      }},
      "trigger": {{"on": "enter"}},
      "gate": {{"enabled": true}},
      "then": [{{"type": "simulate", "name": "turn on amplifier"}}]
    }}
  ]
}}

User: If a car is moving, notify me.
JSON:
{{
  "rules": [
    {{
      "name": "car-moving",
      "when": {{
        "all": [
          {{"moving": {{"label": "car", "min_displacement_ratio": 0.15, "window_frames": 3, "max_missing_frames": 1}}}}
        ]
      }},
      "trigger": {{"on": "enter"}},
      "gate": {{"enabled": true}},
      "then": [{{"type": "simulate", "name": "notify me"}}]
    }}
  ]
}}

User: If person is near monitor turn on lights. When they leave, turn off lights.
JSON:
{{
  "rules": [
    {{
      "name": "monitor-presence-lights",
      "when": {{
        "all": [
          {{"near": {{"a": "person", "b": "monitor", "max_gap_percent": 16}}}}
        ]
      }},
      "trigger": {{"on": "change"}},
      "gate": {{"enabled": true}},
      "then": {{
        "enter": [{{"type": "simulate", "name": "turn on lights"}}],
        "exit": [{{"type": "simulate", "name": "turn off lights"}}]
      }}
    }}
  ]
}}

Full validation schema:
{json.dumps(automation_schema(), indent=2)}
"""


def _build_repair_prompt(*, original_prompt: str, bad_response: str, error: str) -> str:
    return f"""{original_prompt}

Your previous response failed validation.

Validation error:
{error}

Previous response:
{bad_response}

Return corrected JSON only. The root object MUST include a non-empty "rules" array.
"""


def _chat_payload(*, model: str, user_prompt: str, response_format: str) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "model": model,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ],
        "max_tokens": 512,
        "temperature": 0,
        "stream": False,
    }
    if response_format == "json_schema":
        payload["response_format"] = {
            "type": "json_schema",
            "json_schema": {
                "name": "tiny_trigger_automation",
                "strict": True,
                "schema": automation_schema(),
            },
        }
    else:
        payload["response_format"] = {"type": "json_object"}
    return payload