judgetap 0.2.1.dev40001__tar.gz → 0.2.1.dev42001__tar.gz

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: judgetap
-Version: 0.2.1.dev40001
+Version: 0.2.1.dev42001
 Summary: Fast typed decisions (choice, score, yes/no) across Jev-style engines, plus a guard for coding agents. Early development.
 Project-URL: Homepage, https://github.com/mergesafe-ai/judgetap
 Author-email: Omer Bar-Ness <omer@zsquared.io>
@@ -101,6 +101,31 @@ Every PR here is reviewed by [MergeSafe](https://mergesafe.ai) before merge. A f
 
 Early development. Design in [docs/SPEC.md](docs/SPEC.md); roadmap in the issues.
 
+## Open judge models
+
+Open models now match Jev on most decision tasks. LangWatch's comparison reports, over the tasks each model could run: Jev 81.9%, Eikos-27B 80.0%, Shisa DE-1 79.7%, AutoJev-27B 78.9% ([langwatch.ai/compare/jev-vs-all](https://langwatch.ai/compare/jev-vs-all); their numbers, not ours). Run one behind any OpenAI-compatible server and use it through the `llm` engine in **logprobs** mode, which reads each option's probability from the model's token logprobs instead of asking for JSON:
+
+```bash
+# vLLM (GPU; a 27B model needs roughly an H100):
+vllm serve <hf-repo-of-the-model> --served-model-name judge
+# or Ollama:
+ollama serve   # after `ollama pull <model>`
+```
+
+```python
+import os
+
+import judgetap as jt
+
+os.environ["OPENAI_API_BASE"] = (
+    "http://127.0.0.1:8000/v1"  # vLLM; Ollama: http://127.0.0.1:11434/v1
+)
+os.environ["OPENAI_API_KEY"] = "local"  # any value for a local server
+jt.configure(jt.engines.load("llm:openai/judge?logprobs"))
+```
+
+Logprobs mode makes one short call per question (one generated token), run concurrently (up to 8 at a time), and each call is counted in the metrics. The probabilities come from the model's token distribution; they are not calibrated (answers stay `calibrated=False`), so use `judgetap eval` to check how well they track accuracy on your cases. If the server doesn't support logprobs, the engine falls back to JSON mode on its own.
+
 ## Contributing
 
 See [CONTRIBUTING.md](CONTRIBUTING.md).

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/README.md

@@ -84,6 +84,31 @@ Every PR here is reviewed by [MergeSafe](https://mergesafe.ai) before merge. A f
 
 Early development. Design in [docs/SPEC.md](docs/SPEC.md); roadmap in the issues.
 
+## Open judge models
+
+Open models now match Jev on most decision tasks. LangWatch's comparison reports, over the tasks each model could run: Jev 81.9%, Eikos-27B 80.0%, Shisa DE-1 79.7%, AutoJev-27B 78.9% ([langwatch.ai/compare/jev-vs-all](https://langwatch.ai/compare/jev-vs-all); their numbers, not ours). Run one behind any OpenAI-compatible server and use it through the `llm` engine in **logprobs** mode, which reads each option's probability from the model's token logprobs instead of asking for JSON:
+
+```bash
+# vLLM (GPU; a 27B model needs roughly an H100):
+vllm serve <hf-repo-of-the-model> --served-model-name judge
+# or Ollama:
+ollama serve   # after `ollama pull <model>`
+```
+
+```python
+import os
+
+import judgetap as jt
+
+os.environ["OPENAI_API_BASE"] = (
+    "http://127.0.0.1:8000/v1"  # vLLM; Ollama: http://127.0.0.1:11434/v1
+)
+os.environ["OPENAI_API_KEY"] = "local"  # any value for a local server
+jt.configure(jt.engines.load("llm:openai/judge?logprobs"))
+```
+
+Logprobs mode makes one short call per question (one generated token), run concurrently (up to 8 at a time), and each call is counted in the metrics. The probabilities come from the model's token distribution; they are not calibrated (answers stay `calibrated=False`), so use `judgetap eval` to check how well they track accuracy on your cases. If the server doesn't support logprobs, the engine falls back to JSON mode on its own.
+
 ## Contributing
 
 See [CONTRIBUTING.md](CONTRIBUTING.md).

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/docs/SPEC.md

@@ -38,7 +38,8 @@ sj.batch([...questions], context=...)                -> list[Decision]  # one pa
 | `laya` | local, open weights (Apache-2.0) | via transformers; GPU optional |
 | `julia` | local, open weights (Apache-2.0), 144M | Julia-1 runtime from its model repo; CPU by default; Jev-shaped API; 2-20 options per question |
 | `agentjev` | local, open weights | ~50 ms per pass |
-| `llm` | any structured-output LLM via LiteLLM | OpenAI, Gemini, Anthropic, Ollama; probabilities are the model's own JSON estimate, flagged `calibrated=False` (reading logprobs is a later enhancement) |
+| `llm` | any structured-output LLM via LiteLLM | OpenAI, Gemini, Anthropic, Ollama; probabilities are the model's own JSON estimate, flagged `calibrated=False` (see `?logprobs` below) |
+| `llm:<model>?logprobs` | same, OpenAI-compatible servers exposing `logprobs` | options listed as letters, one single-token call per question; the distribution is the letters' `top_logprobs` renormalised (missing letters get 0, none present is an error); token probabilities, not calibrated (`calibrated=False`; check with `judgetap eval`); per-question calls run concurrently (max 8) and each is reported as a Call; more than 26 options is a local error raised before any request; falls back to JSON mode for good if the provider rejects logprobs |
 
 Config: `judgetap.toml` or env vars; `sj.configure(engines=[...])` in code. Keys never logged.
 

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "judgetap"
-version = "0.2.1.dev40001"
+version = "0.2.1.dev42001"
 description = "Fast typed decisions (choice, score, yes/no) across Jev-style engines, plus a guard for coding agents. Early development."
 readme = "README.md"
 license = "Apache-2.0"

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/src/judgetap/__init__.py

@@ -23,7 +23,7 @@ from judgetap.errors import (
 )
 from judgetap.types import Decision, Question
 
-__version__ = "0.2.1.dev40001"  # x-release-please-version
+__version__ = "0.2.1.dev42001"  # x-release-please-version
 
 __all__ = [
     "Cascade",

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/src/judgetap/engines/__init__.py

@@ -3,6 +3,7 @@
 load("jev")                       -> TypeSafe Jev (TYPESAFE_API_KEY)
 load("jev:jev-1.13.0")            -> Jev pinned to a model version
 load("llm:openai/gpt-4o-mini")    -> any LiteLLM model (pip install judgetap[llm])
+    load("llm:openai/m?logprobs")     -> same, probabilities from token logprobs
 """
 
 from __future__ import annotations
@@ -41,7 +42,10 @@ def load(spec: str | None = None) -> Engine:
             raise JudgetapError("llm engine needs a model, e.g. llm:openai/gpt-4o-mini")
         from judgetap.engines.llm import LLMEngine
 
-        return LLMEngine(model=arg)
+        model, _, flags = arg.partition("?")
+        if flags not in ("", "logprobs"):
+            raise JudgetapError(f"unknown llm option {flags!r}; known: logprobs")
+        return LLMEngine(model=model, logprobs=flags == "logprobs")
     if name == "laya":
         from judgetap.engines.laya import DEFAULT_MODEL, LayaEngine
 

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/src/judgetap/engines/julia.py

@@ -6,16 +6,18 @@ the request is built with the Jev adapter's helpers. Its runtime ships in the
 model repo rather than on PyPI:
 
     python -c "from huggingface_hub import snapshot_download; \\
-        snapshot_download('SupersonicLabs/Julia-1', local_dir='Julia-1')"
-    pip install -e ./Julia-1
+        import os; snapshot_download('SupersonicLabs/Julia-1', \\
+        local_dir=os.path.expanduser('~/.judgetap/models/Julia-1'))"
+    pip install -e ~/.judgetap/models/Julia-1
 
 Spec `julia` loads the checkpoint at $JUDGETAP_JULIA_PATH (default
 `~/.judgetap/models/Julia-1`); `julia:<path>` names it. A relative path is
-taken from ~/.judgetap/models, never the working directory: the guard runs
-inside untrusted repos, which must not be able to supply the model that
-judges them. $JUDGETAP_JULIA_DEVICE picks the device (default `cpu`).
+taken from ~/.judgetap/models, never the working directory, and may not
+escape it (`../x` is rejected): the guard runs inside untrusted repos, which
+must not be able to supply the model that judges them. $JUDGETAP_JULIA_DEVICE picks the device (default `cpu`).
 
-Loaded models are cached per process (path, device), so a long-running
+Loaded models are cached per process (path, device), at most
+MAX_CACHED_MODELS of them (least recently used evicted), so a long-running
 program loads each once. The guard hook is a new process per action, so an
 in-process model still loads on every guarded action there: for the guard,
 prefer a server engine (AgentJev, or a TypeSafe-compatible URL).
@@ -36,7 +38,8 @@ from judgetap.errors import JudgetapError
 from judgetap.types import NO, YES, Question
 
 DEFAULT_PATH = "Julia-1"
-_models: dict[tuple[str, str], Any] = {}
+MAX_CACHED_MODELS = 2
+_models: dict[tuple[str, str], Any] = {}  # insertion order = LRU order
 _models_lock = threading.Lock()
 
 
@@ -47,9 +50,21 @@ def _models_dir() -> Path:
 
 
 def resolve_path(path: str) -> str:
-    """Absolute paths as given; anything else under ~/.judgetap/models."""
+    """Absolute paths as given; anything else under ~/.judgetap/models.
+
+    A relative path that resolves outside the models directory (``../x``,
+    or a symlink out of it) is rejected.
+    """
     p = Path(path).expanduser()
-    return str(p if p.is_absolute() else _models_dir() / p)
+    if p.is_absolute():
+        return str(p)
+    base = _models_dir()
+    if not (base / p).resolve().is_relative_to(base.resolve()):
+        raise ValueError(
+            f"julia model path {path!r} escapes {base}; "
+            "use an absolute path for a model outside it"
+        )
+    return str(base / p)
 
 
 # Julia's native limits: 2-20 options per question.
@@ -90,7 +105,9 @@ class JuliaEngine:
         key = (self.path, self.device)
         # One load per process and model, even with concurrent first calls.
         with _models_lock:
-            if key not in _models:
+            if key in _models:
+                _models[key] = _models.pop(key)  # mark most recently used
+            else:
                 try:
                     from julia import load_model
                 except ImportError as err:
@@ -106,6 +123,8 @@ class JuliaEngine:
                     max_length=8192,
                     head_length=512,
                 )
+                while len(_models) > MAX_CACHED_MODELS:
+                    del _models[next(iter(_models))]
             self._runtime = _models[key]
         return self._runtime
 

judgetap-0.2.1.dev42001/src/judgetap/engines/llm.py

@@ -0,0 +1,274 @@
+"""Any chat model through LiteLLM.
+
+Two modes:
+- JSON (default): the model is asked for a JSON distribution. Those numbers
+  are its own estimate, so answers are marked `calibrated=False`.
+- logprobs (`llm:<model>?logprobs`): each question is one call that lists the
+  options as letters and reads the letter's token probabilities from the
+  server's `top_logprobs`. These are the model's token probabilities, not a
+  validated confidence, so answers are still `calibrated=False`; use
+  `judgetap eval` to check how well they track accuracy. One call per
+  question (run concurrently, at most MAX_CONCURRENT at once), each generating
+  a single token; every call is reported so metrics count them. If the
+  provider rejects logprobs, the engine falls back to JSON mode for good.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import math
+import string
+import time
+from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+from judgetap.engine import Call, Context, RawAnswer, plain_context
+from judgetap.errors import JudgetapError
+from judgetap.types import Question
+
+SYSTEM = (
+    "You answer typed questions about the given state. For each question, "
+    "return a probability for every allowed option; probabilities for one "
+    "question sum to 1. Reply with JSON only: "
+    '{"answers": {"<question id>": {"<option>": <probability>, ...}, ...}}'
+)
+
+
+class LLMError(JudgetapError):
+    """The model call failed or its reply could not be parsed."""
+
+
+def _prompt(questions: Sequence[Question], context: Context) -> str:
+    state = plain_context(context)
+    state = state if isinstance(state, str) else json.dumps(state)
+    lines = [f"State:\n{state}\n", "Questions:"]
+    for i, q in enumerate(questions):
+        order = " (ordered lowest to highest)" if q.kind == "score" else ""
+        lines.append(
+            f"- q{i}: {q.text}\n  options{order}: {json.dumps(list(q.options))}"
+        )
+    return "\n".join(lines)
+
+
+LETTERS = string.ascii_uppercase
+MAX_CONCURRENT = 8  # per-question logprobs calls in flight at once
+LOGPROBS_SYSTEM = (
+    "You answer one question about the given state by choosing one option. "
+    "Reply with the option's letter only, nothing else."
+)
+
+
+def _letter_prompt(question: Question, context: Context) -> str:
+    state = plain_context(context)
+    state = state if isinstance(state, str) else json.dumps(state)
+    order = " (ordered lowest to highest)" if question.kind == "score" else ""
+    lines = [f"State:\n{state}\n", f"Question: {question.text}", f"Options{order}:"]
+    lines += [f"{LETTERS[i]}. {option}" for i, option in enumerate(question.options)]
+    lines.append("Answer with one letter.")
+    return "\n".join(lines)
+
+
+def _field(obj: Any, name: str) -> Any:
+    """Read a field from an OpenAI-shaped object or a plain dict."""
+    return obj.get(name) if isinstance(obj, dict) else getattr(obj, name, None)
+
+
+def _letter_distribution(
+    question: Question, response: Any, model: str
+) -> dict[str, float]:
+    """Option -> probability from the first token's top_logprobs, matching
+    letters case-insensitively and ignoring surrounding whitespace."""
+    try:
+        first = _field(_field(_field(response, "choices")[0], "logprobs"), "content")[0]
+        tops = _field(first, "top_logprobs") or []
+    except (TypeError, IndexError, AttributeError) as err:
+        raise LLMError(f"{model} returned no logprobs: {err!r}") from err
+    letters = {LETTERS[i]: option for i, option in enumerate(question.options)}
+    mass = dict.fromkeys(question.options, 0.0)
+    for top in tops:
+        token = str(_field(top, "token") or "").strip().upper()
+        logprob = _field(top, "logprob")
+        if token in letters and isinstance(logprob, int | float):
+            mass[letters[token]] += math.exp(logprob)
+    total = sum(mass.values())
+    if total <= 0:
+        raise LLMError(f"{model} put no probability on any option letter")
+    return {option: p / total for option, p in mass.items()}
+
+
+def _rejects_logprobs(litellm: Any, err: Exception) -> bool:
+    """True when the provider refused the logprobs parameters themselves.
+    Only provider errors count: judgetap's own errors (an option-count limit,
+    an unparseable reply) never switch the engine to JSON mode."""
+    if isinstance(err, JudgetapError):
+        return False
+    kinds = tuple(
+        k
+        for k in (
+            getattr(litellm, "UnsupportedParamsError", None),
+            getattr(litellm, "BadRequestError", None),
+        )
+        if isinstance(k, type)
+    )
+    if kinds and isinstance(err, getattr(litellm, "UnsupportedParamsError", ())):
+        return True
+    return bool(kinds) and isinstance(err, kinds) and "logprob" in str(err).lower()
+
+
+class LLMEngine:
+    def __init__(
+        self, model: str, *, logprobs: bool = False, **completion_kwargs: Any
+    ) -> None:
+        self.name = f"llm:{model}"
+        self.model = model
+        self.logprobs = logprobs
+        self._kwargs = completion_kwargs
+
+    def _messages(self, questions: Sequence[Question], context: Context) -> list[dict]:
+        return [
+            {"role": "system", "content": SYSTEM},
+            {"role": "user", "content": _prompt(questions, context)},
+        ]
+
+    def _parse(self, questions: Sequence[Question], response: Any) -> list[RawAnswer]:
+        try:
+            content = response.choices[0].message.content
+            answers = json.loads(content)["answers"]
+            try:
+                from litellm import completion_cost
+
+                cost = completion_cost(completion_response=response) / len(questions)
+            except Exception:  # noqa: BLE001 -- LiteLLM raises assorted errors for unpriced models
+                cost = None
+            return [
+                RawAnswer(dict(answers[f"q{i}"]), cost_usd=cost, calibrated=False)
+                for i in range(len(questions))
+            ]
+        except (KeyError, TypeError, ValueError, IndexError, AttributeError) as err:
+            raise LLMError(f"could not parse {self.model} reply: {err!r}") from err
+
+    def _check_options(self, questions: Sequence[Question]) -> None:
+        """Before any provider call: a local limit, never a provider refusal."""
+        for q in questions:
+            if len(q.options) > len(LETTERS):
+                raise LLMError(
+                    f"logprobs mode handles up to {len(LETTERS)} options; "
+                    f"{q.text!r} has {len(q.options)}"
+                )
+
+    def _letter_call(self, question: Question, context: Context) -> dict[str, Any]:
+        return {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": LOGPROBS_SYSTEM},
+                {"role": "user", "content": _letter_prompt(question, context)},
+            ],
+            "max_tokens": 1,
+            "logprobs": True,
+            "top_logprobs": 20,
+            **self._kwargs,
+        }
+
+    def _letter_answer(self, question: Question, response: Any) -> RawAnswer:
+        try:
+            from litellm import completion_cost
+
+            cost = completion_cost(completion_response=response)
+        except Exception:  # noqa: BLE001 -- LiteLLM raises assorted errors for unpriced models
+            cost = None
+        return RawAnswer(
+            _letter_distribution(question, response, self.model),
+            cost_usd=cost,
+            calibrated=False,  # token probabilities, not validated confidence
+        )
+
+    def _with_calls(
+        self, answers: list[RawAnswer], timings: list[float]
+    ) -> list[RawAnswer]:
+        """Every answer carries the batch's calls: one per provider request."""
+        calls = tuple(Call(self.name, ms, True, 1) for ms in timings)
+        return [
+            RawAnswer(
+                a.distribution,
+                cost_usd=a.cost_usd,
+                calibrated=a.calibrated,
+                calls=calls,
+            )
+            for a in answers
+        ]
+
+    def _one(
+        self, litellm: Any, q: Question, context: Context
+    ) -> tuple[RawAnswer, float]:
+        start = time.perf_counter()
+        response = litellm.completion(**self._letter_call(q, context))
+        return self._letter_answer(q, response), (time.perf_counter() - start) * 1000
+
+    def decide(
+        self, questions: Sequence[Question], context: Context
+    ) -> Sequence[RawAnswer]:
+        litellm = _import_litellm()
+        if self.logprobs:
+            self._check_options(questions)
+            try:
+                workers = min(MAX_CONCURRENT, max(1, len(questions)))
+                with ThreadPoolExecutor(max_workers=workers) as pool:
+                    done = list(
+                        pool.map(lambda q: self._one(litellm, q, context), questions)
+                    )
+                return self._with_calls([a for a, _ in done], [ms for _, ms in done])
+            except Exception as err:
+                if not _rejects_logprobs(litellm, err):
+                    raise
+                self.logprobs = False  # this provider can't: JSON mode from now on
+        response = litellm.completion(
+            model=self.model,
+            messages=self._messages(questions, context),
+            response_format={"type": "json_object"},
+            **self._kwargs,
+        )
+        return self._parse(questions, response)
+
+    async def adecide(
+        self, questions: Sequence[Question], context: Context
+    ) -> Sequence[RawAnswer]:
+        litellm = _import_litellm()
+        if self.logprobs:
+            self._check_options(questions)
+            gate = asyncio.Semaphore(MAX_CONCURRENT)
+
+            async def one(q: Question) -> tuple[RawAnswer, float]:
+                async with gate:
+                    start = time.perf_counter()
+                    response = await litellm.acompletion(
+                        **self._letter_call(q, context)
+                    )
+                    ms = (time.perf_counter() - start) * 1000
+                return self._letter_answer(q, response), ms
+
+            try:
+                done = await asyncio.gather(*(one(q) for q in questions))
+                return self._with_calls([a for a, _ in done], [ms for _, ms in done])
+            except Exception as err:
+                if not _rejects_logprobs(litellm, err):
+                    raise
+                self.logprobs = False
+        response = await litellm.acompletion(
+            model=self.model,
+            messages=self._messages(questions, context),
+            response_format={"type": "json_object"},
+            **self._kwargs,
+        )
+        return self._parse(questions, response)
+
+
+def _import_litellm():
+    try:
+        import litellm
+    except ImportError as err:
+        raise LLMError(
+            "the llm engine needs LiteLLM: pip install 'judgetap[llm]'"
+        ) from err
+    return litellm

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/src/judgetap/guard/install.py

@@ -200,7 +200,9 @@ def agentjev_up(
     return True
 
 
-SPEC_PATTERN = re.compile(r"[A-Za-z0-9:_./@+\[\]-]+")  # [ ] for IPv6 hosts
+SPEC_PATTERN = re.compile(
+    r"[A-Za-z0-9:_./@+\[\]?-]+"
+)  # [ ] for IPv6 hosts, ? for llm options
 
 
 KNOWN_ENGINES = frozenset({"jev", "llm", "laya", "julia", "agentjev", "typesafe"})

{judgetap-0.2.1.dev40001 → judgetap-0.2.1.dev42001}/tests/test_engine_julia.py

@@ -167,3 +167,33 @@ def test_model_loads_once_per_process_even_concurrently(tmp_path, monkeypatch):
         t.join()
     assert len(loads) == 1
     assert len({id(e._runtime) for e in engines}) == 1
+
+
+@pytest.mark.parametrize("spec", ["../Julia-1", "a/../../x", "../../etc"])
+def test_relative_paths_escaping_models_dir_are_rejected(spec, tmp_path, monkeypatch):
+    monkeypatch.setenv("JUDGETAP_HOME", str(tmp_path / "home"))
+    with pytest.raises(ValueError, match="escapes"):
+        JuliaEngine(spec)
+
+
+def test_relative_path_inside_models_dir_with_dotdot_is_allowed(tmp_path, monkeypatch):
+    monkeypatch.setenv("JUDGETAP_HOME", str(tmp_path / "home"))
+    path = JuliaEngine("a/../Julia-1").path
+    assert path == str(tmp_path / "home" / "models" / "a" / ".." / "Julia-1")
+
+
+def test_model_cache_keeps_two_most_recently_used(fake_julia, tmp_path, monkeypatch):
+    from judgetap.engines import julia as jmod
+
+    monkeypatch.setenv("JUDGETAP_HOME", str(tmp_path))
+    monkeypatch.setattr(jmod, "_models", {})
+    jmod.JuliaEngine("A")._get_runtime()
+    jmod.JuliaEngine("B")._get_runtime()
+    jmod.JuliaEngine("A")._get_runtime()  # A is now most recent
+    jmod.JuliaEngine("C")._get_runtime()  # evicts B
+    models = tmp_path / "models"
+    assert [k[0] for k in jmod._models] == [str(models / "A"), str(models / "C")]
+    assert len(fake_julia.loads) == 3
+    jmod.JuliaEngine("B")._get_runtime()  # reloads B, evicts A
+    assert len(fake_julia.loads) == 4
+    assert [k[0] for k in jmod._models] == [str(models / "C"), str(models / "B")]

judgetap-0.2.1.dev42001/tests/test_llm_logprobs.py

@@ -0,0 +1,266 @@
+import asyncio
+import math
+import sys
+import types
+
+import pytest
+
+import judgetap as jt
+from judgetap.engines import load
+from judgetap.engines.llm import LLMEngine, LLMError
+
+
+def lp_response(tops):
+    """An OpenAI-shaped chat completion whose first token has these top_logprobs."""
+    content = [
+        types.SimpleNamespace(
+            top_logprobs=[
+                types.SimpleNamespace(token=t, logprob=math.log(p)) for t, p in tops
+            ]
+        )
+    ]
+    choice = types.SimpleNamespace(
+        logprobs=types.SimpleNamespace(content=content),
+        message=types.SimpleNamespace(content="A"),
+    )
+    return types.SimpleNamespace(choices=[choice])
+
+
+def json_response(text):
+    msg = types.SimpleNamespace(content=text)
+    return types.SimpleNamespace(
+        choices=[types.SimpleNamespace(message=msg, logprobs=None)]
+    )
+
+
+class FakeLiteLLM(types.ModuleType):
+    def __init__(self, handler):
+        super().__init__("litellm")
+        self.handler = handler
+        self.calls = []
+
+        class UnsupportedParamsError(Exception):
+            pass
+
+        self.UnsupportedParamsError = UnsupportedParamsError
+
+    def completion(self, **kw):
+        self.calls.append(kw)
+        return self.handler(self, kw)
+
+    async def acompletion(self, **kw):
+        return self.completion(**kw)
+
+    def completion_cost(self, completion_response):
+        raise ValueError("unpriced")
+
+
+@pytest.fixture
+def fake(monkeypatch):
+    def install(handler):
+        mod = FakeLiteLLM(handler)
+        monkeypatch.setitem(sys.modules, "litellm", mod)
+        return mod
+
+    return install
+
+
+def test_distribution_from_letter_logprobs(fake):
+    lit = fake(
+        lambda m, kw: lp_response([("A", 0.7), (" b", 0.2), ("C", 0.05), ("Z", 0.05)])
+    )
+    d = jt.choice(
+        "route?",
+        ["billing", "tech", "sales"],
+        engine=LLMEngine("openai/x", logprobs=True),
+    )
+    assert (
+        d.value == "billing" and not d.calibrated
+    )  # token probabilities, not calibrated
+    assert d.distribution["billing"] == pytest.approx(0.7 / 0.95)
+    assert d.distribution["tech"] == pytest.approx(0.2 / 0.95)
+    call = lit.calls[0]
+    assert (
+        call["max_tokens"] == 1
+        and call["logprobs"] is True
+        and call["top_logprobs"] == 20
+    )
+    assert "A. billing" in call["messages"][1]["content"]
+
+
+def test_missing_letters_get_zero(fake):
+    fake(lambda m, kw: lp_response([("B", 0.9), ("x", 0.1)]))
+    d = jt.yesno("q", engine=LLMEngine("openai/x", logprobs=True))
+    assert d.distribution == {"yes": 0.0, "no": 1.0}
+
+
+def test_no_option_letter_is_an_error_so_a_cascade_falls_through(fake):
+    fake(lambda m, kw: lp_response([("x", 0.9), ("y", 0.1)]))
+    with pytest.raises(LLMError, match="no probability"):
+        jt.yesno("q", engine=LLMEngine("openai/x", logprobs=True))
+
+
+def test_one_call_per_question(fake):
+    lit = fake(lambda m, kw: lp_response([("A", 1.0)]))
+    jt.batch(
+        [jt.Question.yesno("a"), jt.Question.yesno("b")],
+        engine=LLMEngine("openai/x", logprobs=True),
+    )
+    assert len(lit.calls) == 2
+
+
+def test_provider_rejecting_logprobs_falls_back_to_json_for_good(fake):
+    def handler(m, kw):
+        if kw.get("logprobs"):
+            raise m.UnsupportedParamsError("logprobs is not supported")
+        return json_response('{"answers": {"q0": {"yes": 0.8, "no": 0.2}}}')
+
+    lit = fake(handler)
+    engine = LLMEngine("openai/x", logprobs=True)
+    d = jt.yesno("q", engine=engine)
+    assert d.value == "yes" and not d.calibrated and engine.logprobs is False
+    jt.yesno("q", engine=engine)
+    assert sum(1 for c in lit.calls if c.get("logprobs")) == 1  # not retried
+
+
+def test_other_errors_are_not_swallowed(fake):
+    def handler(m, kw):
+        raise RuntimeError("rate limited")
+
+    fake(handler)
+    with pytest.raises(RuntimeError):
+        jt.yesno("q", engine=LLMEngine("openai/x", logprobs=True))
+
+
+def test_async_logprobs(fake):
+    fake(lambda m, kw: lp_response([("A", 0.6), ("B", 0.4)]))
+    d = asyncio.run(jt.ayesno("q", engine=LLMEngine("openai/x", logprobs=True)))
+    assert d.value == "yes" and d.p == pytest.approx(0.6)
+
+
+def test_dict_shaped_logprobs(fake):
+    choice = types.SimpleNamespace(
+        logprobs={"content": [{"top_logprobs": [{"token": "b", "logprob": 0.0}]}]}
+    )
+    resp = types.SimpleNamespace(choices=[choice])
+    fake(lambda m, kw: resp)
+    assert jt.yesno("q", engine=LLMEngine("openai/x", logprobs=True)).value == "no"
+
+
+def test_too_many_options_for_letters():
+    q = jt.Question.choice("q", [str(i) for i in range(27)])
+    with pytest.raises(LLMError, match="up to 26"):
+        LLMEngine("m", logprobs=True)._check_options([q])
+
+
+def test_spec_parsing():
+    assert load("llm:openai/qwen?logprobs").logprobs is True
+    e = load("llm:openai/qwen")
+    assert e.logprobs is False and e.model == "openai/qwen"
+    with pytest.raises(jt.JudgetapError):
+        load("llm:openai/qwen?bogus")
+
+
+def test_spec_can_be_saved_to_guard_toml(tmp_path):
+    from judgetap.guard.install import write_engine
+
+    write_engine(tmp_path, "llm:openai/qwen?logprobs")
+    assert (
+        'engine = "llm:openai/qwen?logprobs"' in (tmp_path / "guard.toml").read_text()
+    )
+
+
+def test_option_limit_is_local_and_keeps_logprobs_enabled(fake):
+    lit = fake(lambda m, kw: lp_response([("A", 0.9), ("B", 0.1)]))
+    engine = LLMEngine("openai/x", logprobs=True)
+    too_many = [str(i) for i in range(27)]
+    with pytest.raises(LLMError, match="up to 26 options"):
+        jt.choice("q", too_many, engine=engine)
+    assert engine.logprobs is True and lit.calls == []  # no request, no fallback
+    d = jt.choice("q", ["a", "b"], engine=engine)
+    assert d.value == "a" and lit.calls[-1].get("logprobs")
+
+
+def test_bad_request_mentioning_logprobs_falls_back(fake):
+    def handler(m, kw):
+        if kw.get("logprobs"):
+            raise m.BadRequestError("this model does not support logprobs")
+        return json_response('{"answers": {"q0": {"yes": 0.8, "no": 0.2}}}')
+
+    lit = fake(handler)
+
+    class BadRequestError(Exception):
+        pass
+
+    lit.BadRequestError = BadRequestError
+    engine = LLMEngine("openai/x", logprobs=True)
+    assert jt.yesno("q", engine=engine).value == "yes" and engine.logprobs is False
+
+
+def test_batch_questions_run_concurrently_in_order(fake):
+    import threading
+
+    lock, state = threading.Lock(), {"now": 0, "peak": 0}
+    gate = threading.Barrier(3, timeout=2)
+
+    def handler(m, kw):
+        with lock:
+            state["now"] += 1
+            state["peak"] = max(state["peak"], state["now"])
+        gate.wait()  # only passes if all three requests are in flight together
+        with lock:
+            state["now"] -= 1
+        letter = "A" if "first" in kw["messages"][1]["content"] else "B"
+        return lp_response([(letter, 1.0)])
+
+    fake(handler)
+    qs = [
+        jt.Question.yesno("first"),
+        jt.Question.yesno("second"),
+        jt.Question.yesno("third"),
+    ]
+    out = jt.batch(qs, engine=LLMEngine("openai/x", logprobs=True))
+    assert state["peak"] == 3
+    assert [d.value for d in out] == ["yes", "no", "no"]
+
+
+def test_async_batch_runs_concurrently(fake):
+    lit = fake(lambda m, kw: lp_response([("A", 1.0)]))
+    active = {"now": 0, "peak": 0}
+
+    async def acompletion(**kw):
+        active["now"] += 1
+        active["peak"] = max(active["peak"], active["now"])
+        await asyncio.sleep(0.01)
+        active["now"] -= 1
+        return lp_response([("A", 1.0)])
+
+    lit.acompletion = acompletion
+    qs = [jt.Question.yesno(str(i)) for i in range(5)]
+    asyncio.run(jt.abatch(qs, engine=LLMEngine("openai/x", logprobs=True)))
+    assert active["peak"] == 5
+
+
+def test_plain_dict_response_is_read(fake):
+    def handler(m, kw):
+        return {
+            "choices": [
+                {
+                    "logprobs": {
+                        "content": [{"top_logprobs": [{"token": "B", "logprob": 0.0}]}]
+                    }
+                }
+            ]
+        }
+
+    fake(handler)
+    d = jt.choice("q", ["a", "b"], engine=LLMEngine("openai/x", logprobs=True))
+    assert d.value == "b"
+
+
+def test_each_request_is_reported_as_a_call(fake):
+    fake(lambda m, kw: lp_response([("A", 1.0)]))
+    qs = [jt.Question.yesno(str(i)) for i in range(3)]
+    out = jt.batch(qs, engine=LLMEngine("openai/x", logprobs=True))
+    assert len(out[0].calls) == 3
+    assert all(c.engine == "llm:openai/x" and c.questions == 1 for c in out[0].calls)

judgetap-0.2.1.dev40001/src/judgetap/engines/llm.py

@@ -1,103 +0,0 @@
-"""Any chat model through LiteLLM, asked for a JSON distribution.
-
-LLMs rarely expose calibrated probabilities for a label set, so answers
-from this engine are marked `calibrated=False`: the numbers are the model's
-own estimate, useful for ranking and escalation, not as true frequencies.
-"""
-
-from __future__ import annotations
-
-import json
-from collections.abc import Sequence
-from typing import Any
-
-from judgetap.engine import Context, RawAnswer, plain_context
-from judgetap.errors import JudgetapError
-from judgetap.types import Question
-
-SYSTEM = (
-    "You answer typed questions about the given state. For each question, "
-    "return a probability for every allowed option; probabilities for one "
-    "question sum to 1. Reply with JSON only: "
-    '{"answers": {"<question id>": {"<option>": <probability>, ...}, ...}}'
-)
-
-
-class LLMError(JudgetapError):
-    """The model call failed or its reply could not be parsed."""
-
-
-def _prompt(questions: Sequence[Question], context: Context) -> str:
-    state = plain_context(context)
-    state = state if isinstance(state, str) else json.dumps(state)
-    lines = [f"State:\n{state}\n", "Questions:"]
-    for i, q in enumerate(questions):
-        order = " (ordered lowest to highest)" if q.kind == "score" else ""
-        lines.append(
-            f"- q{i}: {q.text}\n  options{order}: {json.dumps(list(q.options))}"
-        )
-    return "\n".join(lines)
-
-
-class LLMEngine:
-    def __init__(self, model: str, **completion_kwargs: Any) -> None:
-        self.name = f"llm:{model}"
-        self.model = model
-        self._kwargs = completion_kwargs
-
-    def _messages(self, questions: Sequence[Question], context: Context) -> list[dict]:
-        return [
-            {"role": "system", "content": SYSTEM},
-            {"role": "user", "content": _prompt(questions, context)},
-        ]
-
-    def _parse(self, questions: Sequence[Question], response: Any) -> list[RawAnswer]:
-        try:
-            content = response.choices[0].message.content
-            answers = json.loads(content)["answers"]
-            try:
-                from litellm import completion_cost
-
-                cost = completion_cost(completion_response=response) / len(questions)
-            except Exception:  # noqa: BLE001 -- LiteLLM raises assorted errors for unpriced models
-                cost = None
-            return [
-                RawAnswer(dict(answers[f"q{i}"]), cost_usd=cost, calibrated=False)
-                for i in range(len(questions))
-            ]
-        except (KeyError, TypeError, ValueError, IndexError, AttributeError) as err:
-            raise LLMError(f"could not parse {self.model} reply: {err!r}") from err
-
-    def decide(
-        self, questions: Sequence[Question], context: Context
-    ) -> Sequence[RawAnswer]:
-        litellm = _import_litellm()
-        response = litellm.completion(
-            model=self.model,
-            messages=self._messages(questions, context),
-            response_format={"type": "json_object"},
-            **self._kwargs,
-        )
-        return self._parse(questions, response)
-
-    async def adecide(
-        self, questions: Sequence[Question], context: Context
-    ) -> Sequence[RawAnswer]:
-        litellm = _import_litellm()
-        response = await litellm.acompletion(
-            model=self.model,
-            messages=self._messages(questions, context),
-            response_format={"type": "json_object"},
-            **self._kwargs,
-        )
-        return self._parse(questions, response)
-
-
-def _import_litellm():
-    try:
-        import litellm
-    except ImportError as err:
-        raise LLMError(
-            "the llm engine needs LiteLLM: pip install 'judgetap[llm]'"
-        ) from err
-    return litellm