open-interpretability 0.0.20__py3-none-any.whl

openinterp/agents/base.py

@@ -0,0 +1,229 @@
+"""Abstract LLM agent with structured outputs.
+
+Two backends: a LiteLLM model string (routed through instructor), or ``external/<name>``,
+where any LLM, agent or person answers a prompt file written to disk.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Any, ClassVar, TypeVar
+
+from pydantic import BaseModel, ValidationError
+
+from openinterp.core import costs
+from openinterp.core.agent_names import EXTERNAL_PREFIX, spelling_problem
+from openinterp.core.exceptions import AgentError
+from openinterp.core.paths import DEFAULT_EXCHANGE_DIR
+
+T = TypeVar("T", bound=BaseModel)
+
+DEFAULT_MODEL = os.environ.get("OPENINTERP_AGENT_MODEL", "anthropic/claude-opus-5-5")
+MODEL_ALIASES_ENV = "OPENINTERP_MODEL_ALIASES"
+# LiteLLM providers that cache a prompt prefix only where a ``cache_control`` mark ends it.
+# Others (OpenAI, DeepSeek, Gemini) cache a repeated prefix without a mark.
+CACHE_CONTROL_PROVIDERS = frozenset({"anthropic", "bedrock"})
+
+
+def recorded_model(model: str | None) -> str | None:
+    """The name to record for ``model``, after ``$OPENINTERP_MODEL_ALIASES``.
+
+    This is for external answerers behind a proxy. An example is Claude Code routed
+    through a LiteLLM proxy to Grok, whose subagents report the Claude model they think
+    they are. The variable holds ``reported=actual`` pairs separated by commas
+    (``claude-haiku-5-5=grok-4-fast,claude-opus-5-5=grok-4``). Or it holds one bare name
+    that replaces every external model.
+
+    Only ``external/<name>`` is rewritten. A LiteLLM string is what gets called, so it
+    already names the model. Raises ``ValueError`` for an external name spelled other
+    than the site would show it (:func:`openinterp.core.agent_names.spelling_problem`).
+    So one model keeps one name.
+    """
+    spec = os.environ.get(MODEL_ALIASES_ENV, "").strip()
+    if not model or not model.startswith(EXTERNAL_PREFIX):
+        return model
+    if not spec:
+        return _checked(model)
+    name = model.removeprefix(EXTERNAL_PREFIX)
+    aliases, default = {}, None
+    for entry in filter(None, (e.strip() for e in spec.split(","))):
+        reported, sep, actual = entry.partition("=")
+        if not sep:
+            default = entry.removeprefix(EXTERNAL_PREFIX)
+        elif reported.strip() and actual.strip():
+            aliases[reported.strip().removeprefix(EXTERNAL_PREFIX)] = actual.strip()
+        else:
+            raise ValueError(f"${MODEL_ALIASES_ENV}: malformed entry {entry!r}")
+    actual = aliases.get(name, default or name)
+    return _checked(EXTERNAL_PREFIX + actual.removeprefix(EXTERNAL_PREFIX))
+
+
+def _checked(model: str) -> str:
+    if problem := spelling_problem(model):
+        raise ValueError(f"{model!r}: {problem}")
+    return model
+
+
+class LLMAgent:
+    """Base class for agents that return validated Pydantic objects.
+
+    ``model`` is either a LiteLLM model string (``"anthropic/..."``, ``"openai/..."``,
+    ``"ollama/..."``) or ``"external/<name>"``. An external model is answered through
+    files in ``exchange_dir``. The agent writes ``<Agent>.prompt.md``, with the response's
+    JSON schema. Then it waits for ``<Agent>.response.json``. The full model string is what
+    gets recorded, so name the model that actually answered, e.g. ``external/gpt-5``.
+    Pass ``client`` to inject a pre-built instructor client (useful in tests).
+    """
+
+    system_prompt: ClassVar[str] = ""
+    role: ClassVar[str] = "agent"  # the stage its calls are timed as
+
+    def __init__(
+        self,
+        model: str = DEFAULT_MODEL,
+        *,
+        temperature: float | None = None,
+        max_retries: int = 2,
+        client: Any | None = None,
+        exchange_dir: Path | None = None,
+        timeout: float | None = None,
+        poll_interval: float = 2.0,
+    ) -> None:
+        self.model = model
+        self.temperature = temperature
+        self.max_retries = max_retries
+        self._client = client
+        self.exchange_dir = exchange_dir or Path(
+            os.environ.get("OPENINTERP_EXCHANGE_DIR", DEFAULT_EXCHANGE_DIR)
+        )
+        self.timeout = timeout or float(os.environ.get("OPENINTERP_EXCHANGE_TIMEOUT", 3600))
+        self.poll_interval = poll_interval
+
+    @property
+    def client(self) -> Any:
+        """The instructor client over LiteLLM, made on first use."""
+        if self._client is None:
+            import instructor
+            import litellm
+
+            self._client = instructor.from_litellm(litellm.completion)
+        return self._client
+
+    @classmethod
+    def render_prompt(
+        cls, response_model: type[BaseModel], prompt: str, response_path: Path | str
+    ) -> str:
+        """The full prompt for an external answerer: system prompt, JSON schema and task.
+
+        The parts that are the same for every unit come first, so a provider can cache
+        them as a prefix. The task and the response path come last.
+        """
+        schema = json.dumps(response_model.model_json_schema(), indent=2)
+        return (
+            f"# System\n\n{cls.system_prompt}\n\n"
+            f"# Response schema\n\n```json\n{schema}\n```\n\n# Prompt\n\n{prompt}\n\n"
+            f"# Response\n\nWrite a JSON object matching the schema above to "
+            f"`{response_path}`.\n"
+        )
+
+    def system_message(self) -> dict[str, Any]:
+        """The system message, marked as the end of a cacheable prefix where needed.
+
+        The response schema goes to the provider as a tool, which comes before the system
+        message, so the cached prefix holds both.
+        """
+        if self.model.split("/", 1)[0] not in CACHE_CONTROL_PROVIDERS:
+            return {"role": "system", "content": self.system_prompt}
+        block = {"type": "text", "text": self.system_prompt, "cache_control": {"type": "ephemeral"}}
+        return {"role": "system", "content": [block]}
+
+    def structured(self, response_model: type[T], prompt: str) -> T:
+        """Ask the model for a ``response_model`` and record the call's cost under this role."""
+        with costs.stage(self.role):
+            if self.model.startswith(EXTERNAL_PREFIX):
+                result = self._external(response_model, prompt)
+                chars = len(self.system_prompt) + len(prompt)
+                call = costs.LLMCall(
+                    agent=self.role, model=self.model, external=True, prompt_chars=chars
+                )
+                costs.add_call(call)
+                return result
+            return self._completion(response_model, prompt)
+
+    def _completion(self, response_model: type[T], prompt: str) -> T:
+        kwargs: dict[str, Any] = {}
+        if self.temperature is not None:
+            kwargs["temperature"] = self.temperature
+        responses: list[Any] = []  # one raw response per attempt, for usage and cost
+        client, hook = self.client, responses.append
+        if hasattr(client, "on"):  # instructor's hooks; a stub client may lack them
+            client.on("completion:response", hook)
+        try:
+            return client.chat.completions.create(
+                model=self.model,
+                response_model=response_model,
+                max_retries=self.max_retries,
+                messages=[self.system_message(), {"role": "user", "content": prompt}],
+                **kwargs,
+            )
+        except Exception as exc:  # instructor/litellm raise many provider-specific types
+            raise AgentError(
+                f"{type(self).__name__} ({self.model}) failed to produce "
+                f"{response_model.__name__}: {exc}"
+            ) from exc
+        finally:
+            if hasattr(client, "off"):
+                client.off("completion:response", hook)
+            chars = len(self.system_prompt) + len(prompt)
+            costs.add_call(
+                costs.LLMCall(
+                    agent=self.role,
+                    model=self.model,
+                    prompt_chars=chars,
+                    **costs.usage_of(responses),
+                )
+            )
+
+    def _external(self, response_model: type[T], prompt: str) -> T:
+        """Write the prompt to disk and wait for a valid JSON response from anyone."""
+        name = type(self).__name__
+        self.exchange_dir.mkdir(parents=True, exist_ok=True)
+        prompt_path = self.exchange_dir / f"{name}.prompt.md"
+        response_path = self.exchange_dir / f"{name}.response.json"
+        error_path = self.exchange_dir / f"{name}.error.txt"
+        for stale in (response_path, error_path):
+            stale.unlink(missing_ok=True)
+        _write_whole(prompt_path, self.render_prompt(response_model, prompt, response_path))
+        print(f"Waiting for {response_path} (prompt: {prompt_path})…", file=sys.stderr)
+        deadline = time.monotonic() + self.timeout
+        seen = None  # the last unparseable response, rejected only once it stops changing
+        while time.monotonic() < deadline:
+            text = response_path.read_text() if response_path.exists() else ""
+            if text.strip():
+                try:
+                    result = response_model.model_validate_json(text)
+                except ValidationError as exc:
+                    if text != seen:  # possibly caught mid-write: look again next poll
+                        seen = text
+                        time.sleep(self.poll_interval)
+                        continue
+                    seen = None
+                    _write_whole(error_path, str(exc))
+                    response_path.unlink()
+                    print(f"Invalid response, see {error_path}", file=sys.stderr)
+                else:
+                    error_path.unlink(missing_ok=True)
+                    return result
+            time.sleep(self.poll_interval)
+        raise AgentError(f"{name} ({self.model}) got no response at {response_path}.")
+
+
+def _write_whole(path: Path, text: str) -> None:
+    """Write ``text`` so a reader polling for ``path`` never sees it half-written."""
+    partial = path.with_name(f".{path.name}.partial")
+    partial.write_text(text)
+    partial.replace(path)

openinterp/agents/judge.py

@@ -0,0 +1,149 @@
+"""Judge agent: scores an explanation on held-out windows from the explanation alone.
+
+Detection asks which windows activate the pole. Fuzzing asks whether the marked tokens
+in each window are the ones that activate it (Paulo et al., 2025). The judge never sees
+activations. A :class:`ScoringTask` holds only the windows' text and markings. The
+labels are measured again when the answers are scored.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from pydantic import BaseModel, Field
+
+from openinterp.agents.base import LLMAgent
+from openinterp.core.costs import ExperimentCosts
+from openinterp.core.fence import FENCE_RULE, fence
+from openinterp.core.schemas import (
+    ComponentIdentifier,
+    DictionaryProfile,
+    HeadProfile,
+    SuperpositionProfile,
+)
+from openinterp.score_controls import ControlTask
+
+MARK_OPEN, MARK_CLOSE = "<<", ">>"
+
+
+class TaskWindow(BaseModel):
+    """One held-out window as the judge sees it: text, tokens and the marked tokens."""
+
+    text: str
+    tokens: list[str]
+    highlighted: list[int] = Field(default_factory=list)
+
+
+class ScoringTask(BaseModel):
+    """What ``open-interpretability score --emit-task`` writes: everything but the labels."""
+
+    identifier: ComponentIdentifier
+    head_profile: HeadProfile | None = None  # a head's strength is defined by its mean output
+    dictionary_profile: DictionaryProfile | None = None  # and a feature's by its dictionary
+    superposition_profile: SuperpositionProfile | None = (
+        None  # and a superposition's by its members (#3767)
+    )
+    # v0.8: the unit isn't measured at a window's first token (NeuronAuditSubmission)
+    first_token_floor: bool | None = None
+    explanation: str
+    threshold: float
+    window_tokens: int
+    quantiles: int
+    quantile_edges: list[float]
+    seed: int
+    corpus: str
+    corpus_sha256: str | None = None
+    candidate_windows: int
+    active_windows: int
+    detection: list[TaskWindow]
+    fuzzing: list[TaskWindow]
+    controls: list[ControlTask] = Field(default_factory=list)  # judged on the same windows
+    fuzzing_decoys: str | None = None  # "matched-shape" since issue #3505 (ExplanationScores)
+    costs: ExperimentCosts | None = None  # of drawing the windows
+
+
+class DetectionAnswer(BaseModel):  # noqa: D101 (a docstring would join the JSON schema that the agent reads)
+    active: list[int] = Field(
+        description="Numbers of the texts in which, going by the explanation, the neuron "
+        "fires on at least one token. Leave out every other text."
+    )
+
+
+class FuzzingAnswer(BaseModel):  # noqa: D101 (a docstring would join the JSON schema that the agent reads)
+    correct: list[int] = Field(
+        description="Numbers of the texts whose marked tokens are exactly where, going by the "
+        "explanation, the neuron fires. Leave out texts marked wrongly."
+    )
+
+
+def _show(tokens: Sequence[str], marked: Sequence[int] = ()) -> str:
+    hot = set(marked)
+    text = "".join(f"{MARK_OPEN}{t}{MARK_CLOSE}" if i in hot else t for i, t in enumerate(tokens))
+    return text.replace("\n", "⏎")
+
+
+class ExplanationJudge(LLMAgent):
+    """The judge: applies an explanation to windows, never seeing activations."""
+
+    role = "judge"
+    system_prompt = (
+        "You test how well a short explanation of a neuron (or attention head, or dictionary "
+        "latent) in a language "
+        "model predicts where it fires. You see only the explanation and some short texts, "
+        "never its activations. Apply the explanation literally, as written: don't guess "
+        "what the neuron really does, and don't reward texts for being on a related topic "
+        "if the explanation's condition isn't met. Texts are excerpts and may start or end "
+        "mid-word; line breaks are shown as ⏎. " + FENCE_RULE
+    )
+
+    @staticmethod
+    def detection_prompt(
+        explanation: str, windows: Sequence[TaskWindow], noun: str = "neuron"
+    ) -> str:
+        """The detection prompt: which of ``windows`` make the unit fire."""
+        texts = "\n".join(f"{n}. {_show(w.tokens)}" for n, w in enumerate(windows, 1))
+        return (
+            f"Explanation of the {noun}:\n\n{fence(explanation)}\n\n"
+            f"Any number of the texts below, possibly none or all, make the {noun} fire on at "
+            "least one token. List the numbers of the texts in which the explanation says it "
+            f"fires.\n\n{fence(texts)}"
+        )
+
+    @staticmethod
+    def fuzzing_prompt(
+        explanation: str, windows: Sequence[TaskWindow], noun: str = "neuron"
+    ) -> str:
+        """The fuzzing prompt: in which ``windows`` the marked tokens are where the unit fires."""
+        texts = "\n".join(
+            f"{n}. {_show(w.tokens, w.highlighted)}" for n, w in enumerate(windows, 1)
+        )
+        return (
+            f"Explanation of the {noun}:\n\n{fence(explanation)}\n\n"
+            f"In each text below, some tokens are marked {MARK_OPEN}like this{MARK_CLOSE}. In "
+            f"some texts the marks are exactly the tokens where the {noun} fires; in the rest "
+            f"they are on other tokens, or the {noun} doesn't fire in that text at all. List "
+            f"the numbers of the texts whose marks are where the explanation says the {noun} "
+            f"fires.\n\n{fence(texts)}"
+        )
+
+    def detect(
+        self, explanation: str, windows: Sequence[TaskWindow], noun: str = "neuron"
+    ) -> DetectionAnswer:
+        """Ask the judge the detection question."""
+        prompt = self.detection_prompt(explanation, windows, noun)
+        return self.structured(DetectionAnswer, prompt)
+
+    def fuzz(
+        self, explanation: str, windows: Sequence[TaskWindow], noun: str = "neuron"
+    ) -> FuzzingAnswer:
+        """Ask the judge the fuzzing question."""
+        return self.structured(FuzzingAnswer, self.fuzzing_prompt(explanation, windows, noun))
+
+
+def answered(numbers: Sequence[int], count: int, what: str) -> list[bool]:
+    """Per-window yes/no from a judge's list of 1-based text numbers."""
+    bad = sorted({n for n in numbers if not 1 <= n <= count})
+    if bad:
+        raise ValueError(f"{what} answer names texts {bad}, but there are only {count}")
+    chosen = set(numbers)
+    return [n in chosen for n in range(1, count + 1)]