inference/model_wrapper.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Iterable

from openai import OpenAI

from inference.prompts import SYSTEM_PROMPT, build_user_prompt, heuristic_action, sanitize_action_text


@dataclass
class ModelWrapper:
    client: OpenAI | None
    model_name: str
    temperature: float
    max_tokens: int
    offline: bool

    def generate_action(
        self,
        step: int,
        config_text: str,
        error_message: str,
        history: list[str],
        available_actions: Iterable[str] | None = None,
    ) -> str:
        if self.offline or self.client is None:
            return heuristic_action(config_text, error_message, available_actions)

        user_prompt = build_user_prompt(
            step=step,
            config_text=config_text,
            error_message=error_message,
            history=history,
            available_actions=available_actions,
        )

        try:
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": user_prompt},
                ],
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                stream=False,
            )
            generated = completion.choices[0].message.content or ""
            action = sanitize_action_text(
                generated,
                fallback=heuristic_action(config_text, error_message, available_actions),
            )
            return action
        except Exception:
            return heuristic_action(config_text, error_message, available_actions)

    def generate_candidates(
        self,
        step: int,
        config_text: str,
        error_message: str,
        history: list[str],
        count: int,
        available_actions: Iterable[str] | None = None,
    ) -> list[str]:
        candidates = [heuristic_action(config_text, error_message, available_actions)]

        for idx in range(max(1, count)):
            action = self.generate_action(
                step=step,
                config_text=config_text,
                error_message=error_message,
                history=history + [f"candidate={idx}"],
                available_actions=available_actions,
            )
            if action:
                candidates.append(action)

        deduped: list[str] = []
        seen: set[str] = set()
        for candidate in candidates:
            normalized = candidate.strip()
            if not normalized:
                continue
            if normalized in seen:
                continue
            seen.add(normalized)
            deduped.append(normalized)

        return deduped


def score_action_candidate(observation_text: str, action_text: str, hacking_predicate: Any) -> float:
    obs = (observation_text or "").lower()
    action = (action_text or "").lower()

    score = 0.0
    if "npm tset" in obs and "npm test" in action:
        score += 2.0
    if "yaml" in obs and any(token in action for token in ("indent", "syntax", "yaml")):
        score += 1.0
    if "dependency" in obs and any(token in action for token in ("install", "dependency", "package")):
        score += 1.0
    if hacking_predicate(action_text):
        score -= 3.0

    score += min(len(action_text) / 200.0, 0.25)
    return score