logitgate 0.2.0__py3-none-any.whl

logitgate/__init__.py

@@ -0,0 +1,103 @@
+from logitgate.backends.base import (
+    GenerativeBackend,
+    LogitBackend,
+    StubBackend,
+    infer_backend_spec,
+    resolve_generative_backend,
+    resolve_logit_backend,
+)
+from logitgate.backends.llamacpp import LlamaCppBackend
+from logitgate.backends.ollama import OllamaBackend
+from logitgate.backends.openai_compat import OpenAICompatBackend
+from logitgate.backends.recurrent import RecurrentPyTorchBackend
+from logitgate.cascade import Cascade
+from logitgate.config import (
+    LogitGateDefaults,
+    configure,
+    get_default_backend,
+    get_default_save_failures,
+    get_default_snapshot_dir,
+    reset_config,
+)
+from logitgate.debiasing.inversion import (
+    compile_batch_binary_prompt,
+    compile_binary_prompt,
+    compile_choice_prompt,
+    compile_score_prompt,
+)
+from logitgate.debiasing.permutation import (
+    KNOWN_MODEL_TEMPERATURES,
+    CalibrationResult,
+    auto_calibrate_backend,
+    combine_symmetric_log_odds,
+    fit_calibration,
+    get_model_temperature,
+)
+from logitgate.debiasing.softmax import (
+    compute_binary_log_odds,
+    compute_normalized_probabilities,
+    extract_token_logprob,
+)
+from logitgate.gates.ast_gate import ASTGate, ContractGate
+from logitgate.gates.base import BaseGate
+from logitgate.gates.judge_gate import JudgeGate
+from logitgate.gates.logit_gate import ChoiceGate, LogitGate, ScoreGate
+from logitgate.integrations.deepeval import LogitGateMetric, LogitGateSystemOneModel
+from logitgate.integrations.langgraph import as_conditional_edge
+from logitgate.integrations.pytest_plugin import check
+from logitgate.models import (
+    DecisionStatus,
+    LayerTrace,
+    Verdict,
+)
+from logitgate.snapshot import VerdictSnapshot
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "KNOWN_MODEL_TEMPERATURES",
+    "ASTGate",
+    "BaseGate",
+    "CalibrationResult",
+    "Cascade",
+    "ChoiceGate",
+    "ContractGate",
+    "DecisionStatus",
+    "GenerativeBackend",
+    "JudgeGate",
+    "LayerTrace",
+    "LlamaCppBackend",
+    "LogitBackend",
+    "LogitGate",
+    "LogitGateDefaults",
+    "LogitGateMetric",
+    "LogitGateSystemOneModel",
+    "OllamaBackend",
+    "OpenAICompatBackend",
+    "RecurrentPyTorchBackend",
+    "ScoreGate",
+    "StubBackend",
+    "Verdict",
+    "VerdictSnapshot",
+    "as_conditional_edge",
+    "auto_calibrate_backend",
+    "check",
+    "combine_symmetric_log_odds",
+    "compile_batch_binary_prompt",
+    "compile_binary_prompt",
+    "compile_choice_prompt",
+    "compile_score_prompt",
+    "compute_binary_log_odds",
+    "compute_normalized_probabilities",
+    "configure",
+    "extract_token_logprob",
+    "fit_calibration",
+    "get_default_backend",
+    "get_default_save_failures",
+    "get_default_snapshot_dir",
+    "get_model_temperature",
+    "infer_backend_spec",
+    "reset_config",
+    "resolve_generative_backend",
+    "resolve_logit_backend",
+]

logitgate/backends/__init__.py

@@ -0,0 +1,25 @@
+from logitgate.backends.base import (
+    GenerativeBackend,
+    LogitBackend,
+    StubBackend,
+    infer_backend_spec,
+    resolve_generative_backend,
+    resolve_logit_backend,
+)
+from logitgate.backends.llamacpp import LlamaCppBackend
+from logitgate.backends.ollama import OllamaBackend
+from logitgate.backends.openai_compat import OpenAICompatBackend
+from logitgate.backends.recurrent import RecurrentPyTorchBackend
+
+__all__ = [
+    "GenerativeBackend",
+    "LlamaCppBackend",
+    "LogitBackend",
+    "OllamaBackend",
+    "OpenAICompatBackend",
+    "RecurrentPyTorchBackend",
+    "StubBackend",
+    "infer_backend_spec",
+    "resolve_generative_backend",
+    "resolve_logit_backend",
+]

logitgate/backends/base.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import abc
+import asyncio
+import math
+import re
+from typing import Any
+from urllib.parse import urlparse
+
+
+def strip_reasoning_blocks(text: str) -> str:
+    """Strips <think>...</think> reasoning traces from model responses."""
+    cleaned = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
+    if "<think>" in cleaned:
+        cleaned = re.sub(r"<think>.*$", "", cleaned, flags=re.DOTALL)
+    return cleaned.strip()
+
+
+class LogitBackend(abc.ABC):
+    """Abstract base class for extracting first-token log-probabilities."""
+
+    @abc.abstractmethod
+    def query_top_logprobs(
+        self,
+        prompt: str,
+        top_k: int = 10,
+        allowed_tokens: list[str] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Returns top_k token logprobs: [{'token': str, 'logprob': float}, ...]"""
+
+    async def aquery_top_logprobs(
+        self,
+        prompt: str,
+        top_k: int = 10,
+        allowed_tokens: list[str] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Async variant of query_top_logprobs."""
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            None, self.query_top_logprobs, prompt, top_k, allowed_tokens
+        )
+
+
+class GenerativeBackend(abc.ABC):
+    """Abstract base class for generative reasoning & diagnostics."""
+
+    @abc.abstractmethod
+    def generate(self, prompt: str, max_tokens: int | None = None, **kwargs: Any) -> tuple[str, int]:
+        """Returns (generated_text, tokens_generated)."""
+
+    async def agenerate(self, prompt: str, max_tokens: int | None = None, **kwargs: Any) -> tuple[str, int]:
+        """Async variant of generate."""
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, lambda: self.generate(prompt, max_tokens=max_tokens, **kwargs))
+
+
+class StubBackend(LogitBackend, GenerativeBackend):
+    """
+    Offline deterministic stub backend for sub-millisecond unit tests and CI.
+    Records all prompted inputs in self.prompt_log.
+    """
+
+    def __init__(
+        self,
+        default_prob: float = 0.95,
+        forward_top_logprobs: list[dict[str, Any]] | None = None,
+        reverse_top_logprobs: list[dict[str, Any]] | None = None,
+        logprob_responses: list[list[dict[str, Any]]] | None = None,
+        diagnostic_text: str = "[VERDICT: PASS] All criteria met.",
+        tokens_generated: int = 0,
+    ) -> None:
+        self.default_prob = default_prob
+        self.forward_top_logprobs = forward_top_logprobs
+        self.reverse_top_logprobs = reverse_top_logprobs
+        self.logprob_responses = list(logprob_responses) if logprob_responses else None
+        self.diagnostic_text = diagnostic_text
+        self.tokens_generated = tokens_generated
+        self.prompt_log: list[str] = []
+        self._response_idx = 0
+
+    def __repr__(self) -> str:
+        return f"StubBackend(default_prob={self.default_prob})"
+
+    def query_top_logprobs(
+        self,
+        prompt: str,
+        top_k: int = 10,
+        allowed_tokens: list[str] | None = None,
+    ) -> list[dict[str, Any]]:
+        self.prompt_log.append(prompt)
+
+        # 1. Explicit queue response
+        if self.logprob_responses:
+            resp = self.logprob_responses[self._response_idx % len(self.logprob_responses)]
+            self._response_idx += 1
+            return resp
+
+        # 2. Check if this is Pass 2 (reversed mapping where A is Fail / violates)
+        is_reverse_pass = (
+            "A: The candidate artifact violates" in prompt
+            or "A: The content violates" in prompt
+            or "A: The candidate does not" in prompt
+        )
+
+        if is_reverse_pass:
+            if self.reverse_top_logprobs is not None:
+                return self.reverse_top_logprobs
+            p = max(min(self.default_prob, 0.9999), 0.0001)
+            return [
+                {"token": "B", "logprob": math.log(p)},
+                {"token": "A", "logprob": math.log(1.0 - p)},
+            ]
+        else:
+            if self.forward_top_logprobs is not None:
+                return self.forward_top_logprobs
+            p = max(min(self.default_prob, 0.9999), 0.0001)
+            return [
+                {"token": "A", "logprob": math.log(p)},
+                {"token": "B", "logprob": math.log(1.0 - p)},
+            ]
+
+    async def aquery_top_logprobs(
+        self,
+        prompt: str,
+        top_k: int = 10,
+        allowed_tokens: list[str] | None = None,
+    ) -> list[dict[str, Any]]:
+        return self.query_top_logprobs(prompt, top_k, allowed_tokens)
+
+    def generate(self, prompt: str, max_tokens: int | None = None, **kwargs: Any) -> tuple[str, int]:
+        self.prompt_log.append(prompt)
+        return self.diagnostic_text, self.tokens_generated
+
+    async def agenerate(self, prompt: str, max_tokens: int | None = None, **kwargs: Any) -> tuple[str, int]:
+        return self.generate(prompt, max_tokens=max_tokens, **kwargs)
+
+
+def infer_backend_spec(backend: Any) -> str:
+    """Infers a canonical URI specifier string from a resolved backend instance."""
+    if isinstance(backend, str):
+        return backend
+    if isinstance(backend, StubBackend):
+        return f"stub://{backend.default_prob}" if backend.default_prob != 0.95 else "stub://default"
+    cls_name = backend.__class__.__name__
+    base_url = getattr(backend, "base_url", None) or ""
+    model = getattr(backend, "model", None) or ""
+    if cls_name == "OllamaBackend":
+        host = base_url.replace("http://", "").replace("https://", "").rstrip("/") or "localhost:11434"
+        return f"ollama://{host}/{model}" if model else f"ollama://{host}"
+    if cls_name == "OpenAICompatBackend":
+        host = base_url.replace("https://", "").replace("http://", "").rstrip("/")
+        scheme = "openai" if base_url.startswith("https://") else "vllm"
+        return f"{scheme}://{host}"
+    if cls_name == "LlamaCppBackend":
+        if base_url:
+            host = base_url.replace("http://", "").replace("https://", "").rstrip("/")
+            return f"llamacpp://{host}/{model}" if model else f"llamacpp://{host}"
+        return f"llamacpp:///{model}" if model else "llamacpp"
+    return "stub://default"
+
+
+def resolve_logit_backend(
+    backend_spec: str | LogitBackend | None = None,
+    model: str | None = None,
+    **kwargs: Any,
+) -> LogitBackend:
+    if isinstance(backend_spec, LogitBackend):
+        return backend_spec
+
+    from logitgate.config import get_default_backend, get_default_model
+
+    explicit_model = model
+    if backend_spec is None or backend_spec == "default":
+        default_spec = get_default_backend()
+        if model is None:
+            model = get_default_model()
+        if (
+            default_spec == "stub://default"
+            and explicit_model is not None
+            and explicit_model not in ("default", "stub")
+            and not explicit_model.startswith("stub")
+        ):
+            backend_spec = f"ollama://localhost:11434/{explicit_model}"
+        else:
+            backend_spec = default_spec
+    elif model is None:
+        model = get_default_model()
+
+    spec = str(backend_spec).strip()
+    if spec.startswith("stub://") or spec == "stub" or spec == "stub://default":
+        prob = 0.95
+        if "/" in spec and spec not in ("stub://", "stub://default"):
+            part = spec.split("/")[-1]
+            try:
+                prob = float(part)
+            except ValueError:
+                pass
+        return StubBackend(default_prob=prob, **kwargs)
+
+    if spec.startswith("ollama://"):
+        from logitgate.backends.ollama import OllamaBackend
+        parsed = urlparse(spec)
+        netloc = parsed.netloc or "localhost:11434"
+        base_url = f"http://{netloc}"
+        parsed_model = parsed.path.lstrip("/") if parsed.path else None
+        effective_model = model or parsed_model or "qwen2.5-coder:14b"
+        return OllamaBackend(base_url=base_url, model=effective_model, **kwargs)
+
+    if spec.startswith(("openai://", "vllm://")):
+        from logitgate.backends.openai_compat import OpenAICompatBackend
+        parsed = urlparse(spec)
+        scheme = "https" if spec.startswith("openai://") else "http"
+        base_url = f"{scheme}://{parsed.netloc}{parsed.path}"
+        effective_model = model or "gpt-4o-mini"
+        return OpenAICompatBackend(base_url=base_url, model=effective_model, **kwargs)
+
+    if spec.startswith("llamacpp://") or spec == "llamacpp":
+        from logitgate.backends.llamacpp import LlamaCppBackend
+        if spec == "llamacpp":
+            return LlamaCppBackend(model=model, **kwargs)
+        parsed = urlparse(spec)
+        if parsed.netloc:
+            base_url = f"http://{parsed.netloc}"
+            parsed_model = parsed.path.lstrip("/") if parsed.path else None
+            effective_model_opt: str | None = model or parsed_model or None
+            return LlamaCppBackend(base_url=base_url, model=effective_model_opt, **kwargs)
+        else:
+            effective_model_opt = model or (parsed.path.lstrip("/") if parsed.path else None)
+            return LlamaCppBackend(model=effective_model_opt, **kwargs)
+
+    if spec.startswith("recurrent://") or spec == "recurrent":
+        from logitgate.backends.recurrent import RecurrentPyTorchBackend
+        if spec.startswith("recurrent://"):
+            parsed_rec = urlparse(spec)
+            target_model = model or (f"{parsed_rec.netloc}{parsed_rec.path}" if parsed_rec.netloc else None) or "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16"
+        else:
+            target_model = model or "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16"
+        return RecurrentPyTorchBackend(model_name=target_model, **kwargs)
+
+    # Fallback to Ollama or Stub
+    if "stub" in spec:
+        return StubBackend(**kwargs)
+
+    from logitgate.backends.ollama import OllamaBackend
+    return OllamaBackend(model=model or spec, **kwargs)
+
+
+def resolve_generative_backend(
+    backend_spec: str | GenerativeBackend | None = None,
+    model: str | None = None,
+    **kwargs: Any,
+) -> GenerativeBackend:
+    if isinstance(backend_spec, GenerativeBackend):
+        return backend_spec
+
+    backend = resolve_logit_backend(backend_spec, model=model, **kwargs)
+    if isinstance(backend, GenerativeBackend):
+        return backend
+    raise TypeError(f"Backend {backend} does not implement GenerativeBackend")