agent-learning 0.4.1__py3-none-any.whl

agent_learning/scorers/slm/_base.py

@@ -0,0 +1,292 @@
+"""Phi-4-mini-instruct ONNX runner + shared scorer base for Tier 3 SLM.
+
+The runner lazily loads ``onnxruntime-genai`` and the Phi-4-mini-instruct
+INT4 model from a local directory the first time :meth:`SlmRunner.generate`
+is called. Subsequent calls reuse the cached model and tokenizer instances.
+The runner abstraction also makes the per-scorer prompt building and JSON
+parsing trivially mockable in unit tests via :meth:`set_runner`.
+
+Prompt format
+~~~~~~~~~~~~~
+
+The runner applies the Phi-3.5 / Phi-4 chat template:
+
+.. code-block:: text
+
+    <|system|>
+    {system}<|end|>
+    <|user|>
+    {user}<|end|>
+    <|assistant|>
+
+Each subclass of :class:`SlmScorerBase` supplies its own ``system`` and
+``user`` strings; the assistant reply is parsed by
+:func:`parse_verdict` which expects a single JSON object of the form
+``{"verdict": "pass"|"fail", "confidence": <number in [0,1]>}``. The
+parser is permissive about leading or trailing prose because instruction-
+tuned small models occasionally emit a sentence in front of the JSON.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import threading
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Tuple
+
+from ..base import ScoreResult
+
+
+_GENAI_IMPORT_ERROR = (
+    "Tier 3 SLM scorers require the [slm] extra. Install with: "
+    "pip install agent-learning[slm]"
+)
+
+
+def _require_genai() -> Any:
+    """Lazy-import ``onnxruntime_genai``. Raises a friendly error if missing."""
+    try:
+        import onnxruntime_genai as og  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - exercised via tests/mocks
+        raise ImportError(_GENAI_IMPORT_ERROR) from exc
+    return og
+
+
+# ----- model runner --------------------------------------------------------
+
+
+class SlmRunner:
+    """Thin wrapper around onnxruntime-genai's ``Model``+``Tokenizer``+``Generator``.
+
+    The wrapper hides the version-dependent generator loop behind a single
+    :meth:`generate` method. A test harness can replace the runner with any
+    object that implements ``generate(prompt: str, *, max_new_tokens: int,
+    temperature: float) -> str``.
+    """
+
+    _shared_lock = threading.Lock()
+    _shared_runners: Dict[str, "SlmRunner"] = {}
+
+    def __init__(self, model_dir: str) -> None:
+        if not model_dir:
+            raise ValueError("SlmRunner requires a non-empty model_dir.")
+        if not os.path.isdir(model_dir):
+            raise FileNotFoundError(
+                f"SLM model directory not found: {model_dir!r}. "
+                "Download Phi-4-mini-instruct ONNX (INT4) to this path or "
+                "set AGENT_LEARNING_SLM_MODEL_DIR to its location."
+            )
+        og = _require_genai()
+        self._og = og
+        self._model = og.Model(model_dir)
+        self._tokenizer = og.Tokenizer(self._model)
+        self.model_dir = model_dir
+
+    @classmethod
+    def get_shared(cls, model_dir: str) -> "SlmRunner":
+        """Return a process-wide cached runner for ``model_dir``.
+
+        Loading the 3.8 B Phi-4-mini-instruct ONNX model from cold can take
+        several seconds and consume ~2.5 GB of RAM. Sharing one runner across
+        the three Tier 3 scorers keeps inference hot and avoids the
+        re-loading cost.
+        """
+        with cls._shared_lock:
+            runner = cls._shared_runners.get(model_dir)
+            if runner is None:
+                runner = cls(model_dir)
+                cls._shared_runners[model_dir] = runner
+            return runner
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        max_new_tokens: int = 64,
+        temperature: float = 0.0,
+    ) -> str:
+        """Run greedy decoding on ``prompt`` and return the assistant text."""
+        og = self._og
+        tokenizer = self._tokenizer
+        input_tokens = tokenizer.encode(prompt)
+        input_length = len(input_tokens)
+        params = og.GeneratorParams(self._model)
+        params.set_search_options(
+            max_length=input_length + int(max_new_tokens),
+            do_sample=False,
+            temperature=float(max(temperature, 0.0)),
+        )
+        # API surface drift across onnxruntime-genai versions: 0.4 wants
+        # params.input_ids, 0.5+ uses generator.append_tokens. Try both.
+        try:
+            params.input_ids = input_tokens
+        except (AttributeError, TypeError):  # pragma: no cover - version drift
+            pass
+        generator = og.Generator(self._model, params)
+        append = getattr(generator, "append_tokens", None)
+        if callable(append):  # pragma: no cover - version drift
+            try:
+                append(input_tokens)
+            except Exception:
+                pass
+        # Greedy decoding loop.
+        while not generator.is_done():
+            compute_logits = getattr(generator, "compute_logits", None)
+            if callable(compute_logits):  # pragma: no cover - version drift
+                compute_logits()
+            generator.generate_next_token()
+        get_sequence = getattr(generator, "get_sequence", None)
+        if callable(get_sequence):
+            full = get_sequence(0)
+        else:  # pragma: no cover - very old version
+            full = generator.get_next_tokens()
+        # Strip the prompt prefix.
+        output_tokens = list(full)[input_length:]
+        text = tokenizer.decode(output_tokens)
+        return text
+
+
+# ----- prompt rendering ----------------------------------------------------
+
+
+_CHAT_TEMPLATE = (
+    "<|system|>\n{system}<|end|>\n"
+    "<|user|>\n{user}<|end|>\n"
+    "<|assistant|>\n"
+)
+
+
+def render_chat_prompt(*, system: str, user: str) -> str:
+    """Render a Phi-3.5 / Phi-4 chat-style prompt."""
+    return _CHAT_TEMPLATE.format(system=system.strip(), user=user.strip())
+
+
+# ----- verdict parsing -----------------------------------------------------
+
+
+_JSON_RE = re.compile(r"\{[^{}]*\}", re.DOTALL)
+
+
+def parse_verdict(text: str) -> Tuple[float, Dict[str, Any]]:
+    """Extract ``(probability_pass, raw)`` from the model's reply.
+
+    Looks for the first JSON object in ``text``. Falls back to a token
+    heuristic if no JSON is recoverable. The function never raises on a
+    malformed reply; an unrecoverable parse degrades to a neutral 0.5
+    probability so the caller can still emit a score.
+    """
+    raw: Dict[str, Any] = {"text": text}
+    if not text:
+        return 0.5, raw
+    # First, try to parse as raw JSON.
+    candidate: Optional[Dict[str, Any]] = None
+    try:
+        loaded = json.loads(text)
+        if isinstance(loaded, dict):
+            candidate = loaded
+    except (ValueError, TypeError):
+        match = _JSON_RE.search(text)
+        if match is not None:
+            try:
+                loaded = json.loads(match.group(0))
+                if isinstance(loaded, dict):
+                    candidate = loaded
+            except (ValueError, TypeError):
+                candidate = None
+    if candidate is not None:
+        verdict = str(candidate.get("verdict", "")).strip().lower()
+        try:
+            confidence = float(candidate.get("confidence", 0.5))
+        except (TypeError, ValueError):
+            confidence = 0.5
+        confidence = max(0.0, min(1.0, confidence))
+        raw["verdict"] = verdict
+        raw["confidence"] = confidence
+        if verdict.startswith("pass"):
+            return confidence, raw
+        if verdict.startswith("fail"):
+            return 1.0 - confidence, raw
+    # Token-level fallback. Look for an isolated "pass" or "fail" word.
+    lowered = text.lower()
+    pass_hit = re.search(r"\bpass\b", lowered)
+    fail_hit = re.search(r"\bfail\b", lowered)
+    if pass_hit and (not fail_hit or pass_hit.start() < fail_hit.start()):
+        raw["heuristic"] = "pass"
+        return 0.75, raw
+    if fail_hit:
+        raw["heuristic"] = "fail"
+        return 0.25, raw
+    raw["heuristic"] = "neutral"
+    return 0.5, raw
+
+
+# ----- shared base ---------------------------------------------------------
+
+
+@dataclass
+class SlmScorerBase:
+    """Shared scoring scaffolding for Tier 3 small-language-model scorers."""
+
+    name: str = "slm"
+    pass_threshold: float = 0.5
+    max_new_tokens: int = 64
+    temperature: float = 0.0
+    model_dir: str = ""
+    runner: Optional[SlmRunner] = None
+    system_prompt: str = field(default="", repr=False)
+
+    def _get_runner(self) -> SlmRunner:
+        if self.runner is not None:
+            return self.runner
+        if not self.model_dir:
+            raise ValueError(
+                f"{self.name}: model_dir must be set to load the SLM"
+            )
+        self.runner = SlmRunner.get_shared(self.model_dir)
+        return self.runner
+
+    def _build_user_prompt(self, **inputs: Any) -> str:  # pragma: no cover - override
+        raise NotImplementedError
+
+    def _validate(self, **inputs: Any) -> None:  # pragma: no cover - override
+        return
+
+    def _score_from_inputs(self, **inputs: Any) -> ScoreResult:
+        self._validate(**inputs)
+        prompt = render_chat_prompt(
+            system=self.system_prompt,
+            user=self._build_user_prompt(**inputs),
+        )
+        runner = self._get_runner()
+        text = runner.generate(
+            prompt,
+            max_new_tokens=self.max_new_tokens,
+            temperature=self.temperature,
+        )
+        probability, raw = parse_verdict(text)
+        if probability >= self.pass_threshold:
+            label = "pass"
+            confidence = probability
+        else:
+            label = "fail"
+            confidence = 1.0 - probability
+        return ScoreResult(
+            label=label,
+            confidence=confidence,
+            normalized=probability,
+            features={
+                "probability": probability,
+                "model_dir": self.model_dir,
+                "raw": raw,
+            },
+        )
+
+
+__all__ = [
+    "SlmScorerBase",
+    "SlmRunner",
+    "parse_verdict",
+    "render_chat_prompt",
+]

agent_learning/scorers/slm/adherence.py

@@ -0,0 +1,98 @@
+"""Tier 3 SLM adherence scorer.
+
+Asks Phi-4-mini-instruct whether the response adheres to the supplied
+contract (required substrings, forbidden substrings, length bounds, JSON
+shape). The scorer is read-only: it never modifies the contract or the
+response, only scores them against each other.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+from ..base import ScoreResult
+from ._base import SlmRunner, SlmScorerBase
+
+
+_SYSTEM_PROMPT = (
+    "You are a strict evaluator that checks whether an agent's response "
+    "complies with a structural contract. Reply only with a single-line "
+    "JSON object: {\"verdict\": \"pass\" | \"fail\", \"confidence\": "
+    "<number in [0, 1]>}. Output no other text."
+)
+
+
+def _format_contract(contract: Optional[Dict[str, Any]]) -> str:
+    if not contract:
+        return (
+            "(no explicit contract; the response is acceptable as long "
+            "as it is on-topic and free of obvious violations)"
+        )
+    try:
+        return json.dumps(contract, ensure_ascii=False, sort_keys=True)
+    except (TypeError, ValueError):
+        return str(contract)
+
+
+_USER_PROMPT_TEMPLATE = (
+    "CONTRACT:\n{contract}\n\nRESPONSE:\n{response}\n\n"
+    "Does the response satisfy the contract? Reply with the JSON verdict "
+    "only."
+)
+
+
+@dataclass
+class SlmAdherenceScorer(SlmScorerBase):
+    """Phi-4-mini-instruct adherence scorer."""
+
+    name: str = "adherence"
+    system_prompt: str = field(default=_SYSTEM_PROMPT, repr=False)
+
+    @classmethod
+    def load_or_default(
+        cls,
+        model_dir: Optional[str] = None,
+        *,
+        pass_threshold: float = 0.5,
+        max_new_tokens: int = 64,
+        temperature: float = 0.0,
+        runner: Optional[SlmRunner] = None,
+    ) -> "SlmAdherenceScorer":
+        return cls(
+            model_dir=model_dir or "",
+            pass_threshold=pass_threshold,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            runner=runner,
+        )
+
+    def _validate(self, **inputs: object) -> None:
+        if inputs.get("response") is None:
+            raise ValueError(
+                "adherence SLM scorer requires a response"
+            )
+
+    def _build_user_prompt(self, **inputs: object) -> str:
+        return _USER_PROMPT_TEMPLATE.format(
+            contract=_format_contract(inputs.get("contract")),
+            response=str(inputs["response"]).strip(),
+        )
+
+    def score(
+        self,
+        *,
+        response: Optional[str] = None,
+        contract: Optional[Dict[str, Any]] = None,
+        query: Optional[str] = None,
+        **_: object,
+    ) -> ScoreResult:
+        return self._score_from_inputs(
+            response=response,
+            contract=contract,
+            query=query,
+        )
+
+
+__all__ = ["SlmAdherenceScorer"]

agent_learning/scorers/slm/completion.py

@@ -0,0 +1,111 @@
+"""Tier 3 SLM completion scorer.
+
+Asks Phi-4-mini-instruct whether every required token / sub-task is
+present in the response. The expected tokens can be supplied either
+through ``expected_tokens`` directly or through ``contract.completion_tokens``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Sequence
+
+from ..base import ScoreResult
+from ._base import SlmRunner, SlmScorerBase
+
+
+_SYSTEM_PROMPT = (
+    "You are a strict evaluator that checks whether an agent's response "
+    "covers every required item. Reply only with a single-line JSON "
+    "object: {\"verdict\": \"pass\" | \"fail\", \"confidence\": <number "
+    "in [0, 1]>}. Output no other text."
+)
+
+
+_USER_PROMPT_TEMPLATE = (
+    "REQUIRED ITEMS (every one must be covered):\n{items}\n\n"
+    "RESPONSE:\n{response}\n\n"
+    "Does the response cover every required item? Reply with the JSON "
+    "verdict only."
+)
+
+
+def _resolve_targets(
+    expected_tokens: Optional[Sequence[str]],
+    contract: Optional[Dict[str, Any]],
+) -> List[str]:
+    if expected_tokens:
+        return [str(t).strip() for t in expected_tokens if str(t).strip()]
+    if contract:
+        tokens = contract.get("completion_tokens")
+        if isinstance(tokens, Sequence) and not isinstance(tokens, (str, bytes)):
+            return [str(t).strip() for t in tokens if str(t).strip()]
+    return []
+
+
+@dataclass
+class SlmCompletionScorer(SlmScorerBase):
+    """Phi-4-mini-instruct completion scorer."""
+
+    name: str = "completion"
+    system_prompt: str = field(default=_SYSTEM_PROMPT, repr=False)
+
+    @classmethod
+    def load_or_default(
+        cls,
+        model_dir: Optional[str] = None,
+        *,
+        pass_threshold: float = 0.5,
+        max_new_tokens: int = 64,
+        temperature: float = 0.0,
+        runner: Optional[SlmRunner] = None,
+    ) -> "SlmCompletionScorer":
+        return cls(
+            model_dir=model_dir or "",
+            pass_threshold=pass_threshold,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            runner=runner,
+        )
+
+    def _validate(self, **inputs: object) -> None:
+        if inputs.get("response") is None:
+            raise ValueError(
+                "completion SLM scorer requires a response"
+            )
+
+    def _build_user_prompt(self, **inputs: object) -> str:
+        targets = _resolve_targets(
+            inputs.get("expected_tokens"),  # type: ignore[arg-type]
+            inputs.get("contract"),  # type: ignore[arg-type]
+        )
+        if targets:
+            items_str = "\n".join(f"- {t}" for t in targets)
+        else:
+            items_str = (
+                "(no explicit checklist; treat the response as complete "
+                "if it answers the apparent task)"
+            )
+        return _USER_PROMPT_TEMPLATE.format(
+            items=items_str,
+            response=str(inputs["response"]).strip(),
+        )
+
+    def score(
+        self,
+        *,
+        response: Optional[str] = None,
+        expected_tokens: Optional[Sequence[str]] = None,
+        contract: Optional[Dict[str, Any]] = None,
+        query: Optional[str] = None,
+        **_: object,
+    ) -> ScoreResult:
+        return self._score_from_inputs(
+            response=response,
+            expected_tokens=expected_tokens,
+            contract=contract,
+            query=query,
+        )
+
+
+__all__ = ["SlmCompletionScorer"]

agent_learning/scorers/slm/intent.py

@@ -0,0 +1,80 @@
+"""Tier 3 SLM intent scorer.
+
+Asks Phi-4-mini-instruct whether the model's response actually addresses
+the user's question. The scorer requires both ``query`` and ``response``;
+this matches the intent-evaluator surface area used by Tiers 1, 2, and 4.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from ..base import ScoreResult
+from ._base import SlmRunner, SlmScorerBase
+
+
+_SYSTEM_PROMPT = (
+    "You are a strict evaluator that checks whether an agent's response "
+    "addresses the user's question. Reply only with a single-line JSON "
+    "object: {\"verdict\": \"pass\" | \"fail\", \"confidence\": <number "
+    "in [0, 1]>}. Output no other text."
+)
+
+_USER_PROMPT_TEMPLATE = (
+    "QUERY:\n{query}\n\nRESPONSE:\n{response}\n\n"
+    "Does the response actually address the query? Reply with the JSON "
+    "verdict only."
+)
+
+
+@dataclass
+class SlmIntentScorer(SlmScorerBase):
+    """Phi-4-mini-instruct intent scorer."""
+
+    name: str = "intent"
+    system_prompt: str = field(default=_SYSTEM_PROMPT, repr=False)
+
+    @classmethod
+    def load_or_default(
+        cls,
+        model_dir: Optional[str] = None,
+        *,
+        pass_threshold: float = 0.5,
+        max_new_tokens: int = 64,
+        temperature: float = 0.0,
+        runner: Optional[SlmRunner] = None,
+    ) -> "SlmIntentScorer":
+        return cls(
+            model_dir=model_dir or "",
+            pass_threshold=pass_threshold,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            runner=runner,
+        )
+
+    def _validate(self, **inputs: object) -> None:
+        query = inputs.get("query")
+        response = inputs.get("response")
+        if not query or not response:
+            raise ValueError(
+                "intent SLM scorer requires both query and response"
+            )
+
+    def _build_user_prompt(self, **inputs: object) -> str:
+        return _USER_PROMPT_TEMPLATE.format(
+            query=str(inputs["query"]).strip(),
+            response=str(inputs["response"]).strip(),
+        )
+
+    def score(
+        self,
+        *,
+        query: Optional[str] = None,
+        response: Optional[str] = None,
+        **_: object,
+    ) -> ScoreResult:
+        return self._score_from_inputs(query=query, response=response)
+
+
+__all__ = ["SlmIntentScorer"]

agent_learning/scorers/stdlib/__init__.py

@@ -0,0 +1,38 @@
+"""Tier 1 stdlib scorers (zero external dependencies).
+
+These scorers score the rendered response text directly, using only
+the Python standard library:
+
+- :class:`StdlibIntentScorer` is a bag-of-words binary classifier with
+  hashing-trick features and the SDK's pure-Python logistic regression
+  primitive. It returns an unfitted ``pass``-prior policy until fitted.
+- :class:`StdlibAdherenceScorer` is a deterministic rule engine over
+  required substrings, forbidden substrings, length bounds, and an
+  optional JSON-shape requirement. No training is required.
+- :class:`StdlibCompletionScorer` is a token-coverage scorer that
+  measures how many of the expected tokens the response contains.
+  No training is required.
+
+The wire-up contract is the same as every other scorer backend in the
+SDK: each scorer exposes ``score(**kwargs) -> ScoreResult`` and ignores
+extra keyword arguments. The three scorers expect:
+
+- intent: ``query: str``, ``response: str``
+- adherence: ``response: str``, ``contract: dict``
+- completion: ``response: str``, ``expected_tokens: list[str]``
+
+The factory in :mod:`agent_learning.scorers` routes ``tier="stdlib"``
+to this package.
+"""
+
+from __future__ import annotations
+
+from .adherence import StdlibAdherenceScorer
+from .completion import StdlibCompletionScorer
+from .intent import StdlibIntentScorer
+
+__all__ = [
+  "StdlibAdherenceScorer",
+  "StdlibCompletionScorer",
+  "StdlibIntentScorer",
+]