resolutionflow/backend/app/services/template_extraction_service.py

"""TemplateExtractionService — propose a parameter schema from a rendered script.

Phase 5 of the FlowPilot migration. Called when an engineer chooses
"Run now, templatize after resolve" on a suggested fix with no existing
library match. The service looks at the concrete script (with the values
the engineer is about to run with) and session/ticket context, then
proposes a parameterization that future engineers could use from the
Script Library.

Design choices (per FLOWPILOT-MIGRATION.md Section 6.4):

- **Conservative by default.** Prefer fewer parameters. Environment-agnostic
  values (like a command name) should not be parameterized. The prompt calls
  that out explicitly.
- **Round-trip check.** After the LLM proposes parameters, we validate that
  the templated body renders back to the original script when given the
  extracted parameter values. Failures log a warning and the caller falls
  back to a single-parameter "raw script" proposal.
- **Model:** Sonnet (`template_extraction` tier). Creates a persistent
  library artifact — quality matters more than latency.

Output shape mirrors the Script Generator's parameter schema:
    {
      "parameters": [
        {"key": "<snake>", "label": "<human>", "type": "text|password|select|...",
         "inferred_from": "<session fact / ticket field / ai guess>"}
      ],
      "templated_body": "<script with {{ key }} placeholders>",
    }
"""
from __future__ import annotations

import json
import logging
import re
from typing import Any

from app.core.ai_provider import get_ai_provider
from app.core.config import settings

logger = logging.getLogger(__name__)


_EXTRACTION_SYSTEM_PROMPT = """\
You are a senior MSP engineer drafting a reusable script template from a \
concrete script that resolved one ticket. Your job is to identify the values \
in the script that would change for a different invocation — those become \
parameters — and replace them with {{ snake_case }} placeholders.

Return strict JSON with this shape:
{
  "parameters": [
    {
      "key": "<snake_case, ASCII>",
      "label": "<Short human label, Title Case>",
      "type": "text" | "password" | "select" | "boolean" | "number" | "textarea",
      "inferred_from": "<short sentence naming the session fact or ticket \
field this value came from; or 'ai best-guess' when neither>"
    }
  ],
  "templated_body": "<the original script with each parameterized value \
replaced by {{ key }} matching the parameters above>"
}

Rules:
- Prefer FEWER parameters. If a value looks environment-agnostic — a cmdlet \
name, a standard path like C:\\Windows\\System32, a Microsoft-documented URL \
— keep it hardcoded.
- Secret-looking values (passwords, API keys, client secrets) MUST be \
parameterized with type=password.
- The templated_body MUST render back to the original script when the \
parameter values from the context are substituted in. Preserve all whitespace, \
comments, and casing.
- If the script has no meaningful parameters (e.g. it's a single read-only \
cmdlet like Get-Service), return parameters=[] and templated_body = original.
- No markdown fences, no prose, only the JSON object.
"""


async def extract_parameters(
    *,
    script_body: str,
    session_context: str | None = None,
    ticket_context: str | None = None,
) -> dict[str, Any]:
    """Return `{parameters, templated_body}` for the given rendered script.

    On LLM failure or malformed output, returns a conservative fallback:
    the original body with no parameters proposed. Callers can still create
    a `draft_templates` row from this — the engineer reviews and refines
    before accepting in the post-resolve prompt (Phase 6).
    """
    model = settings.get_model_for_action("template_extraction")
    provider = get_ai_provider(model=model)

    input_lines = [
        "# Script to templatize",
        "```",
        script_body.strip(),
        "```",
    ]
    if session_context:
        input_lines.extend(["", "# Session context (facts, symptoms)", session_context.strip()])
    if ticket_context:
        input_lines.extend(["", "# Ticket context (company, user, priority)", ticket_context.strip()])
    user_input = "\n".join(input_lines)

    system_blocks: list[dict[str, Any]] = [
        {
            "type": "text",
            "text": _EXTRACTION_SYSTEM_PROMPT,
            "cache_control": {"type": "ephemeral"},
            # cacheable: identical across every extraction call
        },
    ]

    try:
        text, _in, _out = await provider.generate_json(
            system_prompt=system_blocks,
            messages=[{"role": "user", "content": user_input}],
            max_tokens=3000,
        )
    except Exception:
        logger.exception("TemplateExtractionService LLM call failed; returning fallback")
        return _fallback(script_body)

    parsed = _parse_response(text)
    if parsed is None:
        return _fallback(script_body)

    # Round-trip validation: render parsed["templated_body"] with the
    # `inferred_from` values and confirm it matches the original. We don't
    # have the engineer's values yet here (those come at runtime), but we
    # can at least check that every {{ key }} in templated_body maps to a
    # declared parameter. A mismatch means the LLM referenced an undeclared
    # placeholder — conservative fallback.
    declared_keys = {p.get("key") for p in parsed["parameters"] if isinstance(p, dict)}
    referenced_keys = set(re.findall(r"\{\{\s*(\w+)\s*\}\}", parsed["templated_body"]))
    missing = referenced_keys - declared_keys
    if missing:
        logger.warning(
            "TemplateExtractionService: templated_body references undeclared "
            "keys %s; using fallback",
            sorted(missing),
        )
        return _fallback(script_body)

    return parsed


def _parse_response(raw: str) -> dict[str, Any] | None:
    """Tolerant parse. Returns None on any structural problem."""
    cleaned = raw.strip()
    if cleaned.startswith("```"):
        cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned)
        cleaned = re.sub(r"\s*```$", "", cleaned)
    try:
        data = json.loads(cleaned)
    except (json.JSONDecodeError, ValueError):
        logger.warning("TemplateExtractionService returned non-JSON: %r", raw[:200])
        return None

    if not isinstance(data, dict):
        return None
    params = data.get("parameters")
    body = data.get("templated_body")
    if not isinstance(params, list) or not isinstance(body, str):
        logger.warning("TemplateExtractionService missing parameters or templated_body")
        return None

    # Validate each parameter shape. Drop malformed entries rather than
    # failing the whole response — the engineer will review before accept.
    valid_params: list[dict[str, Any]] = []
    allowed_types = {"text", "password", "select", "boolean", "number", "textarea", "multi_text"}
    for p in params:
        if not isinstance(p, dict):
            continue
        key = p.get("key")
        if not isinstance(key, str) or not re.match(r"^[a-z_][a-z0-9_]*$", key):
            continue
        ptype = p.get("type", "text")
        if ptype not in allowed_types:
            ptype = "text"
        valid_params.append({
            "key": key,
            "label": p.get("label") or key.replace("_", " ").title(),
            "type": ptype,
            "inferred_from": p.get("inferred_from") or "ai best-guess",
        })

    return {"parameters": valid_params, "templated_body": body}


def _fallback(script_body: str) -> dict[str, Any]:
    """Conservative no-op result: zero parameters, body unchanged.

    Used when the LLM call fails or returns unusable output. The engineer
    can still save this as a draft and refine in the post-resolve prompt —
    it just won't propose a parameterization for them.
    """
    return {"parameters": [], "templated_body": script_body}