src/iac_secfix_agents/llm.py

from __future__ import annotations

import inspect
import json
from dataclasses import dataclass
from typing import Any

from pydantic import ValidationError

from .config import RuntimeConfig, vllm_extra_body
from .schemas import FixCandidate


def extract_json_object(text: str) -> str | None:
    start = text.find("{")
    if start == -1:
        return None

    depth = 0
    in_string = False
    escaped = False
    for index in range(start, len(text)):
        char = text[index]
        if in_string:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
            continue
        if char == '"':
            in_string = True
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return text[start : index + 1]
    return None


def coerce_fix_candidate(raw: dict[str, Any]) -> dict[str, Any]:
    patch = raw.get("patch") or raw.get("diff") or ""
    fixed_file = raw.get("fixed_file") or raw.get("content") or raw.get("file_content")
    resolved = raw.get("resolved_policy_ids") or raw.get("resolved") or []
    if isinstance(resolved, str):
        resolved = [resolved]
    return {
        "patch": patch,
        "fixed_file": fixed_file,
        "resolved_policy_ids": resolved,
        "explanation": raw.get("explanation") or raw.get("summary") or "",
        "verification_commands": raw.get("verification_commands") or [],
        "risk_notes": raw.get("risk_notes") or [],
        "requires_human_approval": bool(raw.get("requires_human_approval", True)),
    }


def parse_fix_candidate(text: str) -> tuple[FixCandidate | None, bool]:
    json_text = extract_json_object(text)
    if json_text is None:
        return None, False
    try:
        return FixCandidate.model_validate_json(json_text), True
    except ValidationError:
        try:
            raw = json.loads(json_text)
            if not isinstance(raw, dict):
                return None, False
            return FixCandidate.model_validate(coerce_fix_candidate(raw)), True
        except Exception:
            return None, False


@dataclass(slots=True)
class ChatResult:
    text: str
    used_json_mode: bool


class OpenAICompatibleLLM:
    """Thin OpenAI-compatible client used by the PatchAgent."""

    def __init__(self, config: RuntimeConfig):
        from openai import OpenAI

        self.config = config
        self.client = OpenAI(base_url=config.base_url, api_key=config.api_key)

    def chat(self, messages: list[dict[str, str]], json_mode: bool = True) -> ChatResult:
        kwargs: dict[str, Any] = {
            "model": self.config.model,
            "messages": messages,
            "temperature": self.config.temperature,
            "top_p": self.config.top_p,
            "max_tokens": self.config.max_tokens,
            "extra_body": vllm_extra_body(),
        }
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}

        try:
            completion = self.client.chat.completions.create(**kwargs)
            return ChatResult(completion.choices[0].message.content or "{}", used_json_mode=json_mode)
        except Exception:
            kwargs.pop("response_format", None)
            kwargs.pop("extra_body", None)
            completion = self.client.chat.completions.create(**kwargs)
            return ChatResult(completion.choices[0].message.content or "{}", used_json_mode=False)


def make_pydantic_openai_model(config: RuntimeConfig) -> Any | None:
    """Create a PydanticAI OpenAI-compatible model when the library is present."""

    try:
        from pydantic_ai.models.openai import OpenAIChatModel
        from pydantic_ai.providers.openai import OpenAIProvider

        provider = OpenAIProvider(base_url=config.base_url, api_key=config.api_key)
        return OpenAIChatModel(config.model, provider=provider)
    except Exception:
        pass

    try:
        from pydantic_ai.models.openai import OpenAIModel
        from pydantic_ai.providers.openai import OpenAIProvider

        provider = OpenAIProvider(base_url=config.base_url, api_key=config.api_key)
        return OpenAIModel(config.model, provider=provider)
    except Exception:
        return None


def make_pydantic_agent(model: Any, output_model: type[Any], system_prompt: str) -> Any | None:
    """Compatibility wrapper for PydanticAI versions.

    Some environments use `result_type`, newer ones may use `output_type`. The
    notebook hit this exact API mismatch, so the submission code adapts at runtime.
    """

    if model is None:
        return None
    try:
        from pydantic_ai import Agent

        signature = inspect.signature(Agent)
        kwargs: dict[str, Any] = {}
        if "output_type" in signature.parameters:
            kwargs["output_type"] = output_model
        elif "result_type" in signature.parameters:
            kwargs["result_type"] = output_model
        if "instructions" in signature.parameters:
            kwargs["instructions"] = system_prompt
        elif "system_prompt" in signature.parameters:
            kwargs["system_prompt"] = system_prompt
        if "model" in signature.parameters:
            kwargs["model"] = model
            return Agent(**kwargs)
        return Agent(model, **kwargs)
    except Exception:
        return None