Screen a local Ollama multi-agent app, end to end

ollama serve &                  # skip if it is already running
ollama pull llama3.2:3b         # ~2 GB

mkdir badge-ollama-agent && cd badge-ollama-agent
python3.11 -m venv .venv

# Exercised on Python 3.11.15 (2026-07-26). The OpenTelemetry pins match the
# versions Badge's own staging evidence run was captured on.
langchain-core==1.5.1
langgraph==1.2.9
opentelemetry-sdk==1.44.0
opentelemetry-exporter-otlp-proto-http==1.44.0
httpx==0.28.1

# Extras for the LIVE path: serving the agent over HTTP (`server.py`) so a
# real Badge deployment can screen it.
#
#     ./.venv/bin/pip install -r requirements.txt -r requirements-live.txt
#
# Add -r requirements-ollama.txt too if you want the real local model
# (BADGE_EXAMPLE_BACKEND=ollama) rather than the scripted fake backend.
#
# Deliberately separate from requirements.txt so the offline/CI slice stays
# minimal — CI installs requirements.txt only, and asserts the fake backend
# never imports langchain_ollama.
fastapi==0.140.0
uvicorn==0.51.0

# Extra dependency for the REAL local-model backend
# (BADGE_EXAMPLE_BACKEND=ollama). Install on top of requirements.txt:
#
#     ./.venv/bin/pip install -r requirements.txt -r requirements-ollama.txt
#
# Deliberately kept out of requirements.txt so the offline/CI path never
# needs it — CI asserts the fake backend does not import langchain_ollama.
# You also need the Ollama server itself (https://ollama.com), plus a pulled
# tool-calling model:  ollama pull llama3.2:3b
langchain-ollama==1.1.0

./.venv/bin/pip install -r requirements.txt -r requirements-live.txt -r requirements-ollama.txt

badge_langchain_supervisor_demo          (LangGraph root)
├── supervisor_plan        → 1 LLM call
├── researcher             → vector_search retrieval, web_search, calculator
├── writer                 → word_count
└── supervisor_final       → 1 LLM call

graph = StateGraph(MessagesState)
graph.add_node("supervisor_plan", supervisor_plan)
graph.add_node("researcher", researcher)
graph.add_node("writer", writer)
graph.add_node("supervisor_final", supervisor_final)
graph.add_edge(START, "supervisor_plan")
graph.add_edge("supervisor_plan", "researcher")
graph.add_edge("researcher", "writer")
graph.add_edge("writer", "supervisor_final")
graph.add_edge("supervisor_final", END)
return graph.compile()

def make(temperature: float, tools: list[Any] | None = None) -> Any:
    model = ChatOllama(
        model=OLLAMA_MODEL,
        base_url=OLLAMA_BASE_URL,
        temperature=temperature,
        # Keep responses short: Badge caps each dispatch at 60 s, and a 3B
        # model that rambles will blow through it.
        num_predict=320,
    )
    if tools:
        model = model.bind_tools(tools)
    return model.with_config(metadata={"badge_model_id": OLLAMA_MODEL})

"""LangChain -> OpenTelemetry bridge for Badge's OTLP ingest.

Maps the LangChain callback run tree (run_id / parent_run_id) onto an OTel
span tree and stamps the GenAI semconv attributes Badge's ingest allowlist
accepts:

    gen_ai.request.model        gen_ai.response.model
    gen_ai.usage.input_tokens   gen_ai.usage.output_tokens
    gen_ai.usage.cost           gen_ai.tool.name
    gen_ai.agent.name           badge.response_sha256
    badge.run_id

Everything else — prompts, completions, LangChain metadata — is dropped by
Badge at ingest. This handler deliberately ALSO sets non-whitelisted
attributes (prompt text, raw answer text, framework tags) so the autodiscovery
verification can prove they never reach the DB.

Span-naming rules that matter to Badge:
  * Names are NEVER stored verbatim. They collapse to one of four labels:
    "llm" (name == "chat"/"llm"/... or prefix "chat.", "llm.", "completion.",
    "invoke_model"), "retrieval" (a small fixed set incl. "vector_search"),
    "tool" (forced whenever gen_ai.tool.name is present), else
    "unknown_operation".
  * Names with spaces (e.g. the semconv-recommended "chat gpt-4o-mini") fail
    Badge's identifier regex and collapse to "unknown_operation" — so this
    bridge uses "chat.<model>" (dot separator) for LLM spans.
  * Tool names are always stored as an opaque hash "tool:<sha256[:16]>".
  * Model IDs survive verbatim ONLY if they are in Badge's vendored LiteLLM
    price snapshot; anything else becomes an opaque "model:<sha256[:16]>".
  * Agent names (gen_ai.agent.name) are stored as a bounded projection
    "agent:<display16>~<sha256[:8]>", where the display must fullmatch
    [A-Za-z][A-Za-z0-9_-]{0,15} — 16 chars, leading letter, NO DOTS.
    Whitespace becomes "_"; URL/credential/email/hex shapes, names over 64
    chars, and names whose 16-char prefix cannot be shaped (e.g.
    "v2.supervisor", "3rd-tier-agent") are dropped entirely. At most 8
    distinct agent identities are retained per run AND per Blueprint.
    Spans that carry ONLY an agent name (LangGraph node wrappers) become
    "agent" Blueprint nodes, so the supervisor -> subagent hierarchy is
    discoverable. Keep node names short, alphanumeric, and dot-free for
    them to render.

This bridge is model-agnostic. It is exercised on two backends (see
``agent_app.py``): scripted fake chat models (deterministic, CI) and a real
local model served by Ollama (real tokens, real latency, real variance). On
the fake backend token counts and USD costs are synthesized; on the Ollama
backend both are measured, and cost is derived from energy (see
``LocalEnergyPricer``) because a locally-hosted model has no per-token price.
"""

from __future__ import annotations

import hashlib
import time
from typing import Any
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from opentelemetry import trace
from opentelemetry.trace import Span, Status, StatusCode, set_span_in_context

# Synthesized per-token USD prices, matching Badge's vendored snapshot
# (backend/app/data/model_prices_litellm.json, effective 2026-07-23) so the
# self-reported costs are plausible against Badge's own estimates.
SYNTH_PRICES: dict[str, tuple[float, float]] = {
    "claude-opus-4-6": (0.000005, 0.000025),
    "gpt-4o-mini": (0.00000015, 0.0000006),
}


class LocalEnergyPricer:
    """Derive an honest USD cost for a locally-hosted model call.

    A model running on your own hardware has no per-token price, so Badge's
    vendored price table can never cost it — every span lands in the
    ``unpriced`` bucket unless you self-report. But "free" is not the same as
    "costs nothing": the call burns electricity.

    This prices a call as ``duration x power draw x tariff``, which is a real,
    defensible number rather than a fabricated one:

        cost_usd = (seconds / 3600) x (watts / 1000) x usd_per_kwh

    Mind the /3600: watts are a rate per HOUR, and dropping the conversion
    overstates every call by 3600x (it cost this example one wrong number
    before the unit check caught it).

    The defaults describe the machine this example was exercised on (an Apple
    M4 Pro drawing roughly 40 W above idle while generating) at an EU domestic
    tariff of 0.15 USD/kWh. Both are overridable, and both are *estimates of a
    real quantity* — unlike the fake backend's invented per-token prices. The
    resulting figures are genuinely tiny (~1e-6 USD/call); that is the correct
    and interesting answer, not a rounding error to hide.
    """

    def __init__(self, watts: float = 40.0, usd_per_kwh: float = 0.15) -> None:
        self.watts = watts
        self.usd_per_kwh = usd_per_kwh

    def cost(self, seconds: float) -> float:
        kwh = (seconds / 3600.0) * (self.watts / 1000.0)
        return kwh * self.usd_per_kwh


class BadgeOtelCallbackHandler(BaseCallbackHandler):
    """Convert LangChain callback events into Badge-compatible OTel spans.

    Parameters
    ----------
    tracer:
        An OTel tracer whose provider exports OTLP/HTTP to Badge.
    badge_run_id:
        The Badge run UUID; stamped on every span (required for agent-scope
        ingest tokens; ignored — but harmless — for per-run tokens).
    report_cost_for:
        Models for which the handler sets an explicit ``gen_ai.usage.cost``
        from the ``SYNTH_PRICES`` table (Badge stores it as
        self-reported/ingested cost). Models NOT listed here but with token
        usage exercise Badge's own price-table estimate.
    energy_pricer:
        If set, any LLM span whose model is NOT in ``report_cost_for`` gets a
        measured energy-derived ``gen_ai.usage.cost`` instead. This is the
        local-model path: Badge cannot price ``llama3.2:3b``, so without this
        every span is ``unpriced``.
    probe_dropped_attrs:
        Whether to ALSO stamp deliberately non-whitelisted attributes —
        prompt text, completion text, tool input/output, the raw answer — so
        a verification run can prove Badge drops them.

        **Set this False whenever spans leave the machine.** Badge discards
        these keys at ingest, but "discarded on arrival" is not "never sent":
        with a real exporter configured they travel over the wire first, and
        on the live screening path that text is the customer's prompt. It
        defaults to True because the offline contract probes are this
        example's original purpose; ``server.py`` turns it off automatically
        the moment a real OTLP endpoint is configured.
    """

    # LangChain fires callbacks with run_id/parent_run_id UUIDs; we keep the
    # open span per run_id and parent every child explicitly via context.
    def __init__(
        self,
        tracer: trace.Tracer,
        badge_run_id: str,
        report_cost_for: set[str] | None = None,
        energy_pricer: LocalEnergyPricer | None = None,
        probe_dropped_attrs: bool = True,
    ) -> None:
        self._tracer = tracer
        self._badge_run_id = badge_run_id
        self._report_cost_for = report_cost_for or set()
        self._energy_pricer = energy_pricer
        self._probe = probe_dropped_attrs
        self._spans: dict[UUID, Span] = {}
        self._llm_started_at: dict[UUID, float] = {}
        self.final_answer_sha256: str | None = None
        # Real per-call latencies, in call order. Empty on a fake-model run
        # (those "calls" are dict lookups); the whole point of the Ollama
        # backend is that this list has a distribution worth looking at.
        self.llm_latencies_ms: list[float] = []
        # Measured usage across the whole graph — what the HTTP wrapper
        # reports back to Badge instead of guessing at 4 chars/token.
        self.total_input_tokens = 0
        self.total_output_tokens = 0
        self.tools_called: list[str] = []
        # When Badge dispatches a screening it sends a W3C `traceparent` for
        # a trace id it has ALREADY registered. Setting this makes the root
        # span join that trace instead of starting a new one — the only way
        # an agent-scope OTLP token's spans are accepted.
        self.parent_context: Any | None = None

    # ── internals ──────────────────────────────────────────────────────────

    def _start(self, name: str, run_id: UUID, parent_run_id: UUID | None) -> Span:
        parent_span = self._spans.get(parent_run_id) if parent_run_id else None
        if parent_span is not None:
            ctx = set_span_in_context(parent_span)
        else:
            # Root of the LangChain run tree: attach it to Badge's dispatch
            # trace when there is one, so the whole graph lands inside the
            # trace id Badge registered.
            ctx = self.parent_context
        span = self._tracer.start_span(name, context=ctx)
        span.set_attribute("badge.run_id", self._badge_run_id)
        self._spans[run_id] = span
        return span

    def _end(self, run_id: UUID, error: BaseException | None = None) -> None:
        span = self._spans.pop(run_id, None)
        if span is None:
            return
        if error is not None:
            # The status DESCRIPTION travels on the OTLP wire independently of
            # the attribute whitelist, and Badge never reads it — so raw
            # exception text here is the same "transmitted then discarded"
            # shape `probe_dropped_attrs` exists to prevent. LangChain and
            # provider exceptions routinely embed request content, so when
            # probes are off we send the exception TYPE only.
            detail = str(error) if self._probe else type(error).__name__
            span.set_status(Status(StatusCode.ERROR, detail))
        span.end()

    @staticmethod
    def _serialized_name(serialized: dict[str, Any] | None, **kwargs: Any) -> str:
        if kwargs.get("name"):
            return str(kwargs["name"])
        if serialized:
            if serialized.get("name"):
                return str(serialized["name"])
            if serialized.get("id"):
                return str(serialized["id"][-1])
        return "chain"

    # ── chains (LangGraph nodes, the graph itself) ─────────────────────────

    def on_chain_start(
        self,
        serialized: dict[str, Any],
        inputs: dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        name = self._serialized_name(serialized, **kwargs)
        span = self._start(name, run_id, parent_run_id)
        # WHITELISTED (G2): LangGraph stamps the graph-node name into callback
        # metadata as "langgraph_node" — that is the agent/subagent role.
        # Badge stores it as "agent:<display16>~<hash8>" and derives an "agent"
        # Blueprint node from wrapper spans, so supervisor -> subagent trees
        # survive the projection. Internal chains inside the same graph node
        # inherit the same value and collapse into one node (intended).
        langgraph_node = (metadata or {}).get("langgraph_node")
        if langgraph_node:
            span.set_attribute("gen_ai.agent.name", str(langgraph_node))
        # NOT whitelisted — must be dropped by Badge at ingest:
        if self._probe:
            span.set_attribute("langchain.framework", "langgraph")
            span.set_attribute("badge.agent_role", name)

    def on_chain_end(
        self,
        outputs: dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        span = self._spans.get(run_id)
        if span is not None and parent_run_id is None:
            # Root span: attach the final-answer hash (whitelisted) AND the
            # raw answer text (NOT whitelisted — must be dropped by Badge).
            answer = _extract_answer(outputs)
            if answer:
                digest = hashlib.sha256(answer.encode("utf-8")).hexdigest()
                self.final_answer_sha256 = digest
                span.set_attribute("badge.response_sha256", digest)
                if self._probe:
                    span.set_attribute("badge.answer_text", answer)
        self._end(run_id)

    def on_chain_error(
        self, error: BaseException, *, run_id: UUID, **kwargs: Any
    ) -> None:
        self._end(run_id, error)

    # ── chat models ────────────────────────────────────────────────────────

    def on_chat_model_start(
        self,
        serialized: dict[str, Any],
        messages: list[list[Any]],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        model = (metadata or {}).get("badge_model_id", "unknown-model")
        # "chat.<model>" (dot, not space): survives Badge's identifier regex
        # and its "chat." prefix rule, so the stored label is "llm".
        span = self._start(f"chat.{model}", run_id, parent_run_id)
        self._llm_started_at[run_id] = time.monotonic()
        span.set_attribute("gen_ai.request.model", model)
        # NOT whitelisted — must be dropped by Badge at ingest:
        if self._probe:
            span.set_attribute("gen_ai.prompt.0.content", _first_text(messages))

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        started = self._llm_started_at.pop(run_id, None)
        elapsed = time.monotonic() - started if started is not None else None
        span = self._spans.get(run_id)
        if span is not None:
            message = response.generations[0][0].message  # type: ignore[attr-defined]
            usage = getattr(message, "usage_metadata", None) or {}
            model = message.response_metadata.get("model_name") or "unknown-model"
            span.set_attribute("gen_ai.response.model", model)
            input_tokens = usage.get("input_tokens")
            output_tokens = usage.get("output_tokens")
            if input_tokens is not None:
                span.set_attribute("gen_ai.usage.input_tokens", int(input_tokens))
                self.total_input_tokens += int(input_tokens)
            if output_tokens is not None:
                span.set_attribute("gen_ai.usage.output_tokens", int(output_tokens))
                self.total_output_tokens += int(output_tokens)
            if model in self._report_cost_for and model in SYNTH_PRICES:
                in_price, out_price = SYNTH_PRICES[model]
                cost = (input_tokens or 0) * in_price + (output_tokens or 0) * out_price
                span.set_attribute("gen_ai.usage.cost", round(cost, 10))
            elif self._energy_pricer is not None and elapsed is not None:
                # Local model: no per-token price exists, so price the energy.
                span.set_attribute(
                    "gen_ai.usage.cost", round(self._energy_pricer.cost(elapsed), 12)
                )
            if elapsed is not None:
                self.llm_latencies_ms.append(round(elapsed * 1000, 2))
            # NOT whitelisted — must be dropped by Badge at ingest:
            if self._probe:
                span.set_attribute("gen_ai.completion.0.content", str(message.content))
        self._end(run_id)

    def on_llm_error(
        self, error: BaseException, *, run_id: UUID, **kwargs: Any
    ) -> None:
        self._end(run_id, error)

    # ── tools ──────────────────────────────────────────────────────────────

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        tool_name = self._serialized_name(serialized, **kwargs)
        self.tools_called.append(tool_name)
        span = self._start(f"tool.{tool_name}", run_id, parent_run_id)
        # Whitelisted, but stored as an opaque hash tool:<sha256[:16]>:
        span.set_attribute("gen_ai.tool.name", tool_name)
        # NOT whitelisted — must be dropped by Badge at ingest:
        if self._probe:
            span.set_attribute("gen_ai.tool.input", input_str)

    def on_tool_end(self, output: Any, *, run_id: UUID, **kwargs: Any) -> None:
        span = self._spans.get(run_id)
        if span is not None:
            # NOT whitelisted — must be dropped by Badge at ingest:
            if self._probe:
                span.set_attribute("gen_ai.tool.output", str(output))
        self._end(run_id)

    def on_tool_error(
        self, error: BaseException, *, run_id: UUID, **kwargs: Any
    ) -> None:
        self._end(run_id, error)

    # ── retrievers ─────────────────────────────────────────────────────────

    def on_retriever_start(
        self,
        serialized: dict[str, Any],
        query: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        # "vector_search" is in Badge's fixed retrieval-name set, so this is
        # the one non-LLM, non-tool span shape Badge classifies (label
        # "retrieval") instead of collapsing to unknown_operation.
        self._start("vector_search", run_id, parent_run_id)

    def on_retriever_end(
        self, documents: Any, *, run_id: UUID, **kwargs: Any
    ) -> None:
        self._end(run_id)

    def on_retriever_error(
        self, error: BaseException, *, run_id: UUID, **kwargs: Any
    ) -> None:
        self._end(run_id, error)


def _extract_answer(outputs: dict[str, Any]) -> str | None:
    messages = outputs.get("messages")
    if isinstance(messages, list) and messages:
        last = messages[-1]
        content = getattr(last, "content", None)
        if isinstance(content, str) and content:
            return content
    return None


def _first_text(messages: list[list[Any]]) -> str:
    for batch in messages:
        for message in batch:
            content = getattr(message, "content", None)
            if isinstance(content, str) and content:
                return content[:500]
    return ""

"""LangGraph supervisor + two subagents, on EITHER a real local model or fakes.

Two interchangeable backends behind one graph, selected by
``BADGE_EXAMPLE_BACKEND`` (or the ``backend=`` argument):

``fake`` (default)
    Every model is a ``GenericFakeChatModel`` with scripted ``AIMessage``
    outputs. Deterministic, offline, zero dependencies beyond langchain-core.
    This is the CI path — it must never need Ollama or a network.

``ollama``
    Every model is a real ``ChatOllama`` talking to a local Ollama server.
    Real inference, real token accounting, real latency, real variance —
    and still $0, because the weights run on your own machine.

The topology is identical on both backends, so a real-model trace is
directly comparable to the fake-model baseline:

    supervisor_plan (LLM)
      └─ researcher
           ├─ vector_search retrieval
           ├─ web_search tool
           └─ calculator tool
      └─ writer
           └─ word_count tool
      └─ supervisor_final (LLM)

Graph: START -> supervisor_plan -> researcher -> writer -> supervisor_final -> END
Each node is a LangGraph node, so the LangChain callback tree is exactly what
any LangGraph user's instrumentation would see. Node names are <=16 chars,
letter-initial and dot-free so Badge's ``gen_ai.agent.name`` projection
renders them instead of dropping them.

Two architectures behind one knob, selected by ``BADGE_EXAMPLE_ARCH`` (or the
``arch=`` argument):

``a`` (default)
    The four-role graph above — unchanged, so existing baselines stay valid.

``b``
    Inserts a ``reviewer`` role between writer and supervisor_final:

    START -> supervisor_plan -> researcher -> writer -> reviewer
          -> supervisor_final -> END

    The reviewer is deliberately a PURE LLM node — no tools, greedy — so the
    A->B blueprint diff is exactly one agent node plus rewired edges.
    Badge's observed-architecture versions are content-addressed over the
    whole span graph, and a real model's optional tool calls already mint
    behavioral variants; a tool-free role keeps the intentional change
    distinguishable from that noise.

WHY THE REAL BACKEND EXISTS: the fake backend cannot exercise Badge's
robustness axis at all (the same input always produces byte-identical output)
and its token counts and latencies are invented. A real 3B model produces
genuine run-to-run variance, genuine prompt/completion token counts, and a
genuine cold-start latency spike on the first call.
"""

from __future__ import annotations

import ast
import math
import operator
import os
import re
from typing import Any

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.retrievers import BaseRetriever
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.graph import END, START, MessagesState, StateGraph

# ── backend selection ──────────────────────────────────────────────────────

DEFAULT_BACKEND = os.environ.get("BADGE_EXAMPLE_BACKEND", "fake").strip().lower()

# The local model. llama3.2:3b is the default because it is small enough to
# answer inside Badge's 60 s dispatch cap on consumer hardware AND supports
# native tool calling (many 3B models do not). qwen2.5:3b works identically.
OLLAMA_MODEL = os.environ.get("BADGE_OLLAMA_MODEL", "llama3.2:3b")
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://127.0.0.1:11434")

# "greedy" (default) or "sampled" — see _ollama_models() for the measured
# effect of each on Badge's robustness axis.
TEMPERATURE_PROFILE = (
    os.environ.get("BADGE_TEMPERATURE_PROFILE", "greedy").strip().lower()
)

# "a" (default) or "b" — see the module docstring. Read at build time, not
# import time, so a long-lived server restarted with a new value switches.
DEFAULT_ARCH = "a"


def resolve_arch(arch: str | None) -> str:
    resolved = (arch or os.environ.get("BADGE_EXAMPLE_ARCH", DEFAULT_ARCH)).strip().lower()
    if resolved not in {"a", "b"}:
        raise ValueError(f"unknown architecture {resolved!r}; expected 'a' or 'b'")
    return resolved

# Fake-backend model ids, chosen to exercise three different Badge price-table
# paths (known+self-reported, known+estimated, unknown+unpriced).
SUPERVISOR_MODEL = "claude-opus-4-6"  # in Badge's vendored price table
RESEARCHER_MODEL = "gpt-4o-mini"  # in Badge's vendored price table
WRITER_MODEL = "fable-fake-chat-v1"  # NOT in Badge's price table (on purpose)
REVIEWER_MODEL = "gpt-5-mini"  # arch "b" only; in Badge's vendored price table

QUESTION = (
    "How many EUR does a 3-person team spend per year on a 40 EUR/month "
    "SaaS seat, and summarize why."
)
FINAL_ANSWER = (
    "A 3-person team on a 40 EUR/month seat spends 1440 EUR per year "
    "(3 seats x 40 EUR x 12 months), before any annual-billing discount."
)


# ── tools ──────────────────────────────────────────────────────────────────


@tool
def web_search(query: str) -> str:
    """Search the reference corpus for pricing and market facts."""
    corpus = {
        "pricing": (
            "SaaS seat pricing commonly ranges 30-50 EUR/month; annual billing "
            "discounts ~10%."
        ),
        "screening": (
            "Procurement teams increasingly require independent screening before "
            "granting agents access to production systems."
        ),
        "evaluation": (
            "Benchmarks for agents measure correctness, latency, cost and "
            "robustness. Repeated runs expose nondeterminism a single run hides."
        ),
    }
    lowered = query.lower()
    for key, text in corpus.items():
        if key in lowered:
            return text
    return corpus["pricing"]


# Arithmetic is evaluated with a whitelisted AST walk. NEVER use eval() here:
# on the live-screening path this tool's argument is attacker-influenced (it
# comes from a model that is itself reading a public, unauthenticated prompt).
_BIN_OPS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Mod: operator.mod,
    ast.Pow: operator.pow,
}
_UNARY_OPS = {ast.UAdd: operator.pos, ast.USub: operator.neg}


def _eval_node(node: ast.AST) -> float:
    if isinstance(node, ast.Constant):
        if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
            raise ValueError("only numeric literals are allowed")
        return node.value
    if isinstance(node, ast.BinOp) and type(node.op) in _BIN_OPS:
        left, right = _eval_node(node.left), _eval_node(node.right)
        if isinstance(node.op, ast.Pow) and (abs(right) > 32 or abs(left) > 1e6):
            raise ValueError("exponent out of range")
        return _BIN_OPS[type(node.op)](left, right)
    if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY_OPS:
        return _UNARY_OPS[type(node.op)](_eval_node(node.operand))
    raise ValueError(f"unsupported expression element: {type(node).__name__}")


@tool
def calculator(expression: str) -> str:
    """Evaluate an arithmetic expression, e.g. "3*40*12"."""
    cleaned = expression.strip().replace("x", "*").replace("X", "*")
    if len(cleaned) > 120:
        return "error: expression too long"
    # The formatting tail MUST be inside the try. `1e400` overflows to inf,
    # `1e400-1e400` is nan, and `(-8)**0.5` returns a complex — all of which
    # raise in int()/round() rather than in the walk, and all reachable from
    # <=11 characters of attacker-influenced input.
    try:
        value = _eval_node(ast.parse(cleaned, mode="eval").body)
        if isinstance(value, complex):
            return "error: complex results are not supported"
        if not math.isfinite(value):
            return "error: result is not a finite number"
        if value == int(value):
            return str(int(value))
        return str(round(value, 6))
    except Exception as exc:  # noqa: BLE001 — tool must never raise to the graph
        return f"error: {exc}"


@tool
def word_count(text: str) -> str:
    """Count the words in a draft."""
    return str(len(text.split()))


class ReferenceRetriever(BaseRetriever):
    """Deterministic retriever so on_retriever_start/end callbacks fire.

    Named ``vector_search`` on the span, which is one of the few names Badge
    classifies as a ``retrieval`` node instead of ``unknown_operation``.
    """

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> list[Document]:
        return [Document(page_content="Internal note: our own seat costs 40 EUR/month.")]


TOOLS = [web_search, calculator, word_count]
TOOLS_BY_NAME = {t.name: t for t in TOOLS}
RESEARCHER_TOOLS = [web_search, calculator]
WRITER_TOOLS = [word_count]


# ── fake backend ───────────────────────────────────────────────────────────


def _ai(
    content: str,
    model: str,
    input_tokens: int,
    output_tokens: int,
    tool_calls: list[dict[str, Any]] | None = None,
) -> AIMessage:
    return AIMessage(
        content=content,
        tool_calls=tool_calls or [],
        usage_metadata={
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
        },
        response_metadata={"model_name": model},
    )


def _fake_models() -> dict[str, Any]:
    supervisor = GenericFakeChatModel(
        messages=iter(
            [
                _ai(
                    "Plan: researcher gathers pricing facts, writer drafts the answer.",
                    SUPERVISOR_MODEL,
                    128,
                    41,
                ),
                _ai(FINAL_ANSWER, SUPERVISOR_MODEL, 402, 76),
            ]
        )
    ).with_config(metadata={"badge_model_id": SUPERVISOR_MODEL})

    researcher = GenericFakeChatModel(
        messages=iter(
            [
                _ai(
                    "",
                    RESEARCHER_MODEL,
                    212,
                    33,
                    tool_calls=[
                        {
                            "name": "web_search",
                            "args": {"query": "average SaaS seat price EUR"},
                            "id": "call_ws_1",
                            "type": "tool_call",
                        },
                        {
                            "name": "calculator",
                            "args": {"expression": "3*40*12"},
                            "id": "call_calc_1",
                            "type": "tool_call",
                        },
                    ],
                ),
                _ai(
                    "Found: 40 EUR/month per seat, 3 seats, 12 months = 1440 EUR/year.",
                    RESEARCHER_MODEL,
                    391,
                    58,
                ),
            ]
        )
    ).with_config(metadata={"badge_model_id": RESEARCHER_MODEL})

    writer = GenericFakeChatModel(
        messages=iter(
            [
                _ai(
                    "",
                    WRITER_MODEL,
                    240,
                    22,
                    tool_calls=[
                        {
                            "name": "word_count",
                            "args": {"text": "draft answer about seat pricing"},
                            "id": "call_wc_1",
                            "type": "tool_call",
                        }
                    ],
                ),
                _ai(
                    "Draft: 3 seats x 40 EUR x 12 months = 1440 EUR per year.",
                    WRITER_MODEL,
                    288,
                    47,
                ),
            ]
        )
    ).with_config(metadata={"badge_model_id": WRITER_MODEL})

    reviewer = GenericFakeChatModel(
        messages=iter(
            [
                _ai(
                    "Review: the draft is accurate (3 x 40 x 12 = 1440) and "
                    "concise. Approved without changes.",
                    REVIEWER_MODEL,
                    312,
                    28,
                )
            ]
        )
    ).with_config(metadata={"badge_model_id": REVIEWER_MODEL})

    return {
        "supervisor": supervisor,
        "researcher": researcher,
        "writer": writer,
        # Built for both architectures; arch "a" simply never invokes it. A
        # scripted iterator that is never consumed is inert.
        "reviewer": reviewer,
        # The fake backend scripts exactly two supervisor turns from one
        # iterator, so plan and final share the handle here.
        "supervisor_final": supervisor,
    }


# ── ollama backend ─────────────────────────────────────────────────────────


def _ollama_models() -> dict[str, Any]:
    """Three ChatOllama handles on the same local model, one per graph role.

    They are separate objects (not one shared handle) purely so each role can
    carry its own temperature and tool binding; they all report the same
    ``badge_model_id``, which is the honest thing to send — and which is
    exactly why every span lands in Badge's ``unpriced`` bucket.
    """
    from langchain_ollama import ChatOllama

    def make(temperature: float, tools: list[Any] | None = None) -> Any:
        model = ChatOllama(
            model=OLLAMA_MODEL,
            base_url=OLLAMA_BASE_URL,
            temperature=temperature,
            # Keep responses short: Badge caps each dispatch at 60 s, and a 3B
            # model that rambles will blow through it.
            num_predict=320,
        )
        if tools:
            model = model.bind_tools(tools)
        return model.with_config(metadata={"badge_model_id": OLLAMA_MODEL})

    # Badge's robustness axis scores agreement across repeated runs of the
    # same task, so temperature is not a style knob here — it is the score.
    # Measured on this graph, same prompt, 6 runs each:
    #
    #   sampled  (0.2/0.3/0.5/0.5) -> 6 distinct answers, 17% modal agreement
    #   sampled  (0.2/0.3/0.5/0.0) -> 5 distinct answers, 33% modal agreement
    #   greedy   (0.0/0.0/0.0/0.0) -> 1 distinct answer, 100% modal agreement
    #
    # The middle row is the useful lesson: making only the LAST hop greedy
    # barely helps, because it still receives a different draft every run.
    # Upstream variance propagates, so robustness is a property of the whole
    # graph. `greedy` is the default for exactly that reason.
    profile = TEMPERATURE_PROFILE
    if profile == "greedy":
        temperatures = {
            "supervisor": 0.0,
            "researcher": 0.0,
            "writer": 0.0,
            "reviewer": 0.0,
            "supervisor_final": 0.0,
        }
    elif profile == "sampled":
        temperatures = {
            "supervisor": 0.2,
            "researcher": 0.3,
            "writer": 0.5,
            # The reviewer stays cool even under "sampled": its job is a
            # judgment call on someone else's text, and heating it up only
            # adds variance to a role that exists to REMOVE variance.
            "reviewer": 0.1,
            "supervisor_final": 0.5,
        }
    else:
        raise ValueError(
            f"unknown temperature profile {profile!r}; expected 'greedy' or 'sampled'"
        )

    return {
        "supervisor": make(temperatures["supervisor"]),
        "researcher": make(temperatures["researcher"], RESEARCHER_TOOLS),
        "writer": make(temperatures["writer"], WRITER_TOOLS),
        "reviewer": make(temperatures["reviewer"]),
        "supervisor_final": make(temperatures["supervisor_final"]),
    }


def build_models(backend: str) -> dict[str, Any]:
    if backend == "fake":
        return _fake_models()
    if backend == "ollama":
        return _ollama_models()
    raise ValueError(f"unknown backend {backend!r}; expected 'fake' or 'ollama'")


# ── graph ──────────────────────────────────────────────────────────────────

_PLAN_SYSTEM = (
    "You are a supervisor coordinating a researcher and a writer. "
    "In at most two sentences, state the plan for answering the user's task. "
    "Do not answer the task yourself."
)
_RESEARCH_SYSTEM = (
    "You are a researcher. Gather the facts needed to answer the task. "
    "Use the calculator tool for any arithmetic and the web_search tool for "
    "pricing or market facts. State findings plainly; do not write the final "
    "answer."
)
_WRITE_SYSTEM = (
    "You are a writer. Turn the researcher's findings into a draft answer. "
    "Be concise and concrete."
)
_REVIEW_SYSTEM = (
    "You are a reviewer. Check the draft answer against the task for factual "
    "or arithmetic errors. If it is correct, restate the draft verbatim after "
    "the word APPROVED. If it has an error, output a corrected draft after "
    "the word CORRECTED. Output nothing else."
)
_FINAL_SYSTEM = (
    "You deliver the final answer. Output ONLY the answer itself.\n"
    "Rules:\n"
    "- No preamble, no sign-off, no explanation of your process.\n"
    "- No markdown code fences.\n"
    "- Never describe a function call or emit {\"name\": ..., "
    '"parameters": ...}. You have no tools; just write the answer.\n'
    "- If the task asked for JSON, emit exactly one JSON object using the "
    "field names the task asked for, in the order it asked for them.\n"
    "- If the task asked for prose, answer in at most three sentences."
)

# A 3B model ignores "no code fences" maybe one run in five. Stripping them is
# plain output normalization — the same thing every production wrapper does —
# and it is NOT answer-tampering: it removes packaging, never content.
_FENCE_RE = re.compile(r"^\s*```[a-zA-Z0-9_-]*\s*\n?(.*?)\n?\s*```\s*$", re.DOTALL)


def _normalize_answer(text: str) -> str:
    """Strip markdown fences the model was told not to emit."""
    match = _FENCE_RE.match(text)
    return match.group(1).strip() if match else text.strip()

# A real 3B model will sometimes loop on tool calls. Bound it hard: this is a
# graph that has to terminate inside a 60 s screening dispatch.
MAX_TOOL_ITERATIONS = 3


def _text(message: Any) -> str:
    content = getattr(message, "content", "")
    if isinstance(content, list):  # some providers return content blocks
        return " ".join(
            part.get("text", "") for part in content if isinstance(part, dict)
        ).strip()
    return str(content).strip()


def build_app(backend: str | None = None, arch: str | None = None) -> Any:
    backend = (backend or DEFAULT_BACKEND).strip().lower()
    arch = resolve_arch(arch)
    models = build_models(backend)
    supervisor_model = models["supervisor"]
    researcher_model = models["researcher"]
    writer_model = models["writer"]
    reviewer_model = models["reviewer"]
    final_model = models["supervisor_final"]
    retriever = ReferenceRetriever()

    def _run_tool_loop(
        model: Any, messages: list[Any], config: RunnableConfig
    ) -> list[Any]:
        """One bounded ReAct turn: model -> execute tool_calls -> model."""
        response = model.invoke(messages, config)
        messages = [*messages, response]
        iterations = 0
        while getattr(response, "tool_calls", None) and iterations < MAX_TOOL_ITERATIONS:
            iterations += 1
            for tool_call in response.tool_calls:
                name = tool_call.get("name")
                target = TOOLS_BY_NAME.get(name)
                if target is None:
                    # A real model will occasionally hallucinate a tool name.
                    # Feed the error back rather than crashing the graph.
                    result: Any = f"error: no such tool {name!r}"
                else:
                    try:
                        result = target.invoke(tool_call.get("args") or {}, config)
                    except Exception as exc:  # noqa: BLE001
                        result = f"error: {exc}"
                messages.append(
                    ToolMessage(
                        content=str(result),
                        tool_call_id=tool_call.get("id") or name or "call",
                    )
                )
            response = model.invoke(messages, config)
            messages.append(response)
        return messages

    def supervisor_plan(state: MessagesState, config: RunnableConfig) -> dict[str, Any]:
        task = _text(state["messages"][0])
        plan = supervisor_model.invoke(
            [SystemMessage(content=_PLAN_SYSTEM), HumanMessage(content=task)], config
        )
        return {"messages": [plan]}

    def researcher(state: MessagesState, config: RunnableConfig) -> dict[str, Any]:
        task = _text(state["messages"][0])
        retriever.invoke("seat cost internal notes", config)
        new_messages = _run_tool_loop(
            researcher_model,
            [
                SystemMessage(content=_RESEARCH_SYSTEM),
                HumanMessage(content="Task: " + task),
            ],
            config,
        )
        return {"messages": [new_messages[-1]]}

    def writer(state: MessagesState, config: RunnableConfig) -> dict[str, Any]:
        task = _text(state["messages"][0])
        findings = _text(state["messages"][-1])
        new_messages = _run_tool_loop(
            writer_model,
            [
                SystemMessage(content=_WRITE_SYSTEM),
                HumanMessage(
                    content=f"Task: {task}\n\nResearcher findings:\n{findings}"
                ),
            ],
            config,
        )
        return {"messages": [new_messages[-1]]}

    def reviewer(state: MessagesState, config: RunnableConfig) -> dict[str, Any]:
        task = _text(state["messages"][0])
        draft = _text(state["messages"][-1])
        review = reviewer_model.invoke(
            [
                SystemMessage(content=_REVIEW_SYSTEM),
                HumanMessage(content=f"Task: {task}\n\nDraft answer:\n{draft}"),
            ],
            config,
        )
        return {"messages": [review]}

    def supervisor_final(state: MessagesState, config: RunnableConfig) -> dict[str, Any]:
        task = _text(state["messages"][0])
        draft = _text(state["messages"][-1])
        final = final_model.invoke(
            [
                SystemMessage(content=_FINAL_SYSTEM),
                HumanMessage(content=f"Task: {task}\n\nDraft answer:\n{draft}"),
            ],
            config,
        )
        return {"messages": [final]}

    graph = StateGraph(MessagesState)
    graph.add_node("supervisor_plan", supervisor_plan)
    graph.add_node("researcher", researcher)
    graph.add_node("writer", writer)
    if arch == "b":
        graph.add_node("reviewer", reviewer)
    graph.add_node("supervisor_final", supervisor_final)
    graph.add_edge(START, "supervisor_plan")
    graph.add_edge("supervisor_plan", "researcher")
    graph.add_edge("researcher", "writer")
    if arch == "b":
        graph.add_edge("writer", "reviewer")
        graph.add_edge("reviewer", "supervisor_final")
    else:
        graph.add_edge("writer", "supervisor_final")
    graph.add_edge("supervisor_final", END)
    return graph.compile()


def run(
    callbacks: list[Any],
    prompt: str | None = None,
    backend: str | None = None,
    app: Any | None = None,
    arch: str | None = None,
) -> str:
    """Run the graph once and return the final answer text.

    ``app`` lets a long-lived server compile the graph once and reuse it; the
    fake backend's scripted message iterators are single-use, so it must
    compile a fresh app per call.
    """
    graph = app if app is not None else build_app(backend, arch)
    result = graph.invoke(
        {"messages": [HumanMessage(content=prompt or QUESTION)]},
        config={"callbacks": callbacks, "run_name": "badge_langchain_supervisor_demo"},
    )
    return _normalize_answer(_text(result["messages"][-1]))


if __name__ == "__main__":
    import sys

    cli_prompt = sys.stdin.read().strip() if not sys.stdin.isatty() else None
    print(run([], prompt=cli_prompt or None))

"""HTTP wrapper: serves the LangGraph agent over Badge's screening contract.

This is the piece that turns the example into something Badge can actually
screen. It implements the `badge_native` contract from
/docs/guides/build-multi-agent, and — unlike the deterministic agent in that
guide — it runs a REAL local model, so the score it earns reflects real
variance, real latency and real token counts.

It also implements the full provenance handshake, which is the part no
customer-facing doc currently spells out end to end. When Badge screens an
agent with BPP enabled, the dispatch carries:

    X-Badge-Run: true                        always
    User-Agent: Badge/<version>              always
    traceparent: 00-<trace_id>-<span_id>-01  when BPP is on
    X-Badge-Run-Id: <run uuid>               when BPP is on
    Accept-Signature: ...                    when BPP is on (RFC 9421 hint)

and Badge records `echoed: true` for the round trip only if the response
carries `X-Badge-Trace-Id: <trace_id>` back. This server therefore:

  1. extracts the W3C trace context from the incoming `traceparent`,
  2. runs the graph with every span parented INTO that registered trace,
  3. echoes `X-Badge-Trace-Id`,
  4. exports the spans to Badge's OTLP ingest with an agent-scope token.

That is the only way an agent-scope token's spans are accepted: Badge
rejects any trace id it did not itself register at dispatch (422 "OTLP trace
does not match the registered run dispatch").

Run it:

    BADGE_EXAMPLE_BACKEND=ollama PORT=8787 ./.venv/bin/python server.py

Optional, to also ship traces (agent-scope token from
POST /api/v1/agents/{id}/otlp-token):

    BADGE_OTLP_ENDPOINT=https://api-staging.badgeia.com/api/v1/provenance/otlp
    BADGE_OTLP_TOKEN=<agent-scope token>
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import math
import os
import time
from typing import Any

from fastapi import FastAPI, Request, Response
from fastapi.responses import JSONResponse
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

from agent_app import DEFAULT_BACKEND, OLLAMA_MODEL, build_app, resolve_arch, run
from badge_otel_callback import BadgeOtelCallbackHandler, LocalEnergyPricer

BACKEND = os.environ.get("BADGE_EXAMPLE_BACKEND", DEFAULT_BACKEND).strip().lower()
# Resolved once at startup and passed explicitly, so /health reports the
# architecture this process actually serves — not whatever the env says later.
ARCH = resolve_arch(None)
OTLP_ENDPOINT = os.environ.get("BADGE_OTLP_ENDPOINT", "").strip()
OTLP_TOKEN = os.environ.get("BADGE_OTLP_TOKEN", "").strip()

# Hostile-input budgets. Badge's real screening prompts are a few KB, so
# these are generous; the point is that they are finite. Anyone with the
# tunnel URL can POST to this route.
MAX_BODY_BYTES = int(os.environ.get("BADGE_MAX_BODY_BYTES", 1 << 20))  # 1 MiB
MAX_PROMPT_CHARS = int(os.environ.get("BADGE_MAX_PROMPT_CHARS", 100_000))
# Badge gives up on a dispatch at 60 s; finishing after that wins nothing and
# keeps a worker busy.
RUN_TIMEOUT_SECONDS = float(os.environ.get("BADGE_RUN_TIMEOUT_SECONDS", 60))
MAX_CONCURRENT_RUNS = int(os.environ.get("BADGE_MAX_CONCURRENT_RUNS", 2))

@contextlib.asynccontextmanager
async def _lifespan(_: FastAPI):
    yield
    # Flush queued spans on the way out. BatchSpanProcessor holds the last
    # batch in memory, so without this the final run's trace is simply lost.
    with contextlib.suppress(Exception):
        _provider.force_flush(timeout_millis=10_000)


# Serve exactly two routes and nothing else. A tunnel publishes EVERY path on
# the port, so FastAPI's interactive docs would otherwise go public too.
app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None, lifespan=_lifespan)

_propagator = TraceContextTextMapPropagator()
_pricer = LocalEnergyPricer()
_slots = asyncio.Semaphore(MAX_CONCURRENT_RUNS)

# The graph is compiled once and reused across requests on the real backend.
# The fake backend's scripted message iterators are single-use, so it must
# recompile per request instead.
_APP = build_app(BACKEND, ARCH) if BACKEND != "fake" else None

_provider = TracerProvider(
    resource=Resource.create({"service.name": "badge-langgraph-ollama-reference"})
)
_exporting = False
if OTLP_ENDPOINT and OTLP_TOKEN:
    from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
    from opentelemetry.sdk.trace.export import BatchSpanProcessor

    # Batch, not Simple: exporting inline would add network latency to a
    # response that has a 60 s budget. Badge accepts a run's spans after the
    # run itself completes, which is what makes async export safe here.
    _provider.add_span_processor(
        BatchSpanProcessor(
            OTLPSpanExporter(
                endpoint=OTLP_ENDPOINT,
                headers={"Authorization": f"Bearer {OTLP_TOKEN}"},
                timeout=20,
            )
        )
    )
    _exporting = True

_tracer = _provider.get_tracer("badge.langchain.reference")


def approx_tokens(text: str) -> int:
    """Fallback only — the real counts come from the model's usage metadata."""
    return max(1, math.ceil(len(text) / 4)) if text else 0


@app.get("/health")
async def health() -> dict[str, Any]:
    return {
        "status": "ok",
        "backend": BACKEND,
        "architecture": ARCH,
        "model": OLLAMA_MODEL if BACKEND == "ollama" else "scripted-fake",
        "otlp_export": _exporting,
        "routes": ["GET /health", "POST /execute"],
    }


@app.post("/execute")
async def execute(request: Request) -> Response:
    started = time.perf_counter()

    # This route is PUBLIC and unauthenticated by design (Badge sends no
    # credentials), so treat the body as hostile. Badge's own ingest does the
    # same thing — declared-length precheck plus a streaming cap, because
    # Content-Length is absent on a chunked upload and trivially forged.
    declared = request.headers.get("content-length")
    if declared and declared.isdigit() and int(declared) > MAX_BODY_BYTES:
        return JSONResponse({"error": "request body too large"}, status_code=413)
    chunks: list[bytes] = []
    received = 0
    async for chunk in request.stream():
        received += len(chunk)
        if received > MAX_BODY_BYTES:
            return JSONResponse({"error": "request body too large"}, status_code=413)
        chunks.append(chunk)

    try:
        body = json.loads(b"".join(chunks))
    except Exception:
        return JSONResponse({"error": "body must be valid JSON"}, status_code=400)
    if not isinstance(body, dict) or not isinstance(body.get("prompt"), str):
        return JSONResponse({"error": "'prompt' is required"}, status_code=400)
    if len(body["prompt"]) > MAX_PROMPT_CHARS:
        return JSONResponse({"error": "prompt too long"}, status_code=413)

    # body also carries task_id, max_tokens and — for non-secret tasks —
    # expected_output. We deliberately never read expected_output: grading
    # yourself against the answer key is not work, and it is absent exactly
    # when the task matters most (secret holdout suites).
    prompt = body["prompt"]

    # Adopt Badge's registered trace, if this is a real screening dispatch.
    carrier = {
        key.lower(): value
        for key, value in request.headers.items()
        if key.lower() in {"traceparent", "tracestate"}
    }
    parent_context = _propagator.extract(carrier) if carrier else None
    incoming_trace_id = None
    if parent_context is not None:
        from opentelemetry.trace import get_current_span

        span_context = get_current_span(parent_context).get_span_context()
        if span_context.is_valid:
            incoming_trace_id = format(span_context.trace_id, "032x")

    handler = BadgeOtelCallbackHandler(
        _tracer,
        badge_run_id=request.headers.get("X-Badge-Run-Id", ""),
        report_cost_for=set(),
        energy_pricer=_pricer,
        # OFF whenever spans actually leave this machine. The probe attributes
        # carry prompt and completion text; Badge drops them at ingest, but on
        # this path that text is the caller's prompt and there is no reason to
        # put it on the wire at all. Keep it that way in anything you copy.
        probe_dropped_attrs=not _exporting,
    )
    handler.parent_context = parent_context  # type: ignore[attr-defined]

    graph = _APP if _APP is not None else build_app(BACKEND, ARCH)
    try:
        # `graph.invoke` is SYNCHRONOUS and takes seconds. Calling it directly
        # from an async handler blocks the event loop for the whole run, which
        # serializes every caller AND stops /health answering — and Badge's
        # Test Connection gives up after 10 s. Off-thread it goes, behind a
        # semaphore so a hostile caller cannot spawn unbounded threads.
        async with _slots:
            output = await asyncio.wait_for(
                asyncio.to_thread(run, [handler], prompt, BACKEND, graph),
                timeout=RUN_TIMEOUT_SECONDS,
            )
    except TimeoutError:
        return JSONResponse({"error": "agent timed out"}, status_code=504)
    except Exception:
        # A 200 with a broken body scores `bad_response_shape`, which is worse
        # than an honest 500: it looks like the agent answered. Detail is
        # logged, not returned — the exception text can carry OLLAMA_BASE_URL
        # and local paths, and this endpoint is public.
        logging.exception("agent run failed")
        return JSONResponse({"error": "agent failed"}, status_code=500)

    elapsed = time.perf_counter() - started
    input_tokens = handler.total_input_tokens or approx_tokens(prompt)
    output_tokens = handler.total_output_tokens or approx_tokens(output)

    payload: dict[str, Any] = {
        "output": output,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        # Report the honest number. Omitting it makes Badge estimate at GPT-4o
        # reference pricing, which would overstate a local model by ~1000x.
        "total_cost_usd": round(_pricer.cost(elapsed), 10),
        # Extra keys are preserved in the run trace — a free debug channel.
        "agent": {
            "framework": "langgraph",
            "backend": BACKEND,
            "architecture": ARCH,
            "model": OLLAMA_MODEL if BACKEND == "ollama" else "scripted-fake",
            "llm_calls": len(handler.llm_latencies_ms),
            "llm_latencies_ms": handler.llm_latencies_ms,
            "tools_called": handler.tools_called,
            "local_compute_ms": int(elapsed * 1000),
            "trace_id": incoming_trace_id,
        },
    }

    headers = {}
    if incoming_trace_id:
        # Badge records `echoed: true` for this round trip only if we send
        # the registered trace id back on the response.
        headers["X-Badge-Trace-Id"] = incoming_trace_id
    return JSONResponse(payload, headers=headers)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        app,
        host=os.environ.get("HOST", "127.0.0.1"),
        port=int(os.environ.get("PORT", "8787")),
        log_level="info",
    )

"""Sanity check: run the graph, dump the OTel span tree to stdout.

No Badge, no ingest. Verifies the callback bridge produces the
agent -> subagent -> tool hierarchy and the GenAI attributes before anything
is exported.

Default backend is ``fake``: no network, no LLM, deterministic — this is the
path CI runs. Set ``BADGE_EXAMPLE_BACKEND=ollama`` to run the same assertions
against a real local model (requires ``ollama serve``), and
``BADGE_EXAMPLE_ARCH=b`` to check the reviewer variant's tree.
"""

from __future__ import annotations

import os
import sys

from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

from agent_app import DEFAULT_BACKEND, OLLAMA_MODEL, resolve_arch, run
from badge_otel_callback import BadgeOtelCallbackHandler, LocalEnergyPricer

BACKEND = os.environ.get("BADGE_EXAMPLE_BACKEND", DEFAULT_BACKEND).strip().lower()
ARCH = resolve_arch(None)

provider = TracerProvider(
    resource=Resource.create({"service.name": "langchain-otel-prototype"})
)
exporter = InMemorySpanExporter()
provider.add_span_processor(SimpleSpanProcessor(exporter))
tracer = provider.get_tracer("badge.langchain.demo")

handler = BadgeOtelCallbackHandler(
    tracer,
    badge_run_id="00000000-0000-0000-0000-000000000000",
    # On the fake backend, price the two "claude" spans from the synthetic
    # table. On the real backend nothing is in Badge's price table, so cost
    # comes from measured energy instead.
    report_cost_for={"claude-opus-4-6"} if BACKEND == "fake" else set(),
    energy_pricer=LocalEnergyPricer() if BACKEND == "ollama" else None,
)
answer = run([handler], backend=BACKEND)
provider.force_flush()

spans = exporter.get_finished_spans()
children: dict[int | None, list] = {}
for s in spans:
    parent = s.parent.span_id if s.parent else None
    children.setdefault(parent, []).append(s)


def show(parent_id, depth):
    for s in sorted(children.get(parent_id, []), key=lambda x: x.start_time):
        attrs = {
            k: v
            for k, v in (s.attributes or {}).items()
            if k.startswith(("gen_ai", "badge.response"))
        }
        print("  " * depth + f"- {s.name}  {attrs}")
        show(s.context.span_id, depth + 1)


model_label = OLLAMA_MODEL if BACKEND == "ollama" else "scripted fakes"
print(f"backend={BACKEND} ({model_label})")
print(f"{len(spans)} spans; answer: {answer[:60]}...")
show(None, 0)

if handler.llm_latencies_ms:
    lat = handler.llm_latencies_ms
    print(
        f"\nreal LLM latencies (ms): {lat}"
        f"\n  first call {lat[0]:.0f} ms (includes model load) vs "
        f"median-of-rest {sorted(lat[1:])[len(lat[1:]) // 2]:.0f} ms"
    )

# ── structural assertions (hold on BOTH backends) ─────────────────────────
agent_names = {
    s.attributes.get("gen_ai.agent.name")
    for s in spans
    if s.attributes and s.attributes.get("gen_ai.agent.name")
}
tool_names = {
    s.attributes.get("gen_ai.tool.name")
    for s in spans
    if s.attributes and s.attributes.get("gen_ai.tool.name")
}
llm_spans = [s for s in spans if s.name.startswith("chat.")]
roots = [s for s in spans if s.parent is None]

failures: list[str] = []
expected_agents = {"supervisor_plan", "researcher", "writer", "supervisor_final"}
if ARCH == "b":
    expected_agents |= {"reviewer"}
if agent_names != expected_agents:
    failures.append(f"agent names wrong for arch {ARCH!r}: {sorted(agent_names)}")
if not tool_names:
    failures.append("no tool spans — the model called no tools this run")
if len(llm_spans) < 4:
    failures.append(f"expected >=4 LLM spans, got {len(llm_spans)}")
if len(roots) != 1:
    failures.append(f"expected exactly 1 root span, got {len(roots)}")
if not any(
    s.attributes and s.attributes.get("badge.response_sha256") for s in spans
):
    failures.append("no badge.response_sha256 on the root span")
missing_tokens = [
    s.name
    for s in llm_spans
    if not (s.attributes or {}).get("gen_ai.usage.input_tokens")
]
if missing_tokens:
    failures.append(f"LLM spans missing token usage: {missing_tokens}")

print()
if failures:
    for failure in failures:
        print(f"FAIL  {failure}")
    sys.exit(1)
print(
    f"PASS  1 root, {len(agent_names)} agent nodes {sorted(agent_names)}, "
    f"{len(llm_spans)} LLM spans all with token usage, "
    f"{len(tool_names)} distinct tools {sorted(tool_names)}"
)

"""Register + screen this agent against a REAL Badge deployment.

This points Badge at a PUBLIC tunnel URL, so the arena performs a genuine
HTTPS round trip to the agent, the run lands ``execution_mode=live_endpoint``
on its own, and the whole autodiscovery pipeline runs for real. A localhost
URL cannot work here however convenient it would be: Badge's SSRF validator
correctly refuses to dispatch to one.

    ./.venv/bin/python screen_live.py \
        --base-url https://api-staging.badgeia.com \
        --endpoint https://<your>.trycloudflare.com/execute \
        --email you@example.com --password '...'

Omit --email/--password to register a throwaway account instead (only
sensible against a disposable environment).

Writes evidence/live_<env>.json.
"""

from __future__ import annotations

import argparse
import json
import os
import time
import uuid
from pathlib import Path

import httpx

HERE = Path(__file__).parent
EVIDENCE = HERE / "evidence"

def write_secret_json(path: Path, payload: dict) -> None:
    """Write 0600. This file holds a live JWT and an OTLP ingest token.

    --state-file can point anywhere, so the permission bits are the real
    control here. Do not commit whatever directory this lands in.
    """
    path.write_text(json.dumps(payload, indent=2))
    os.chmod(path, 0o600)


AGENT_NAME = "LangGraph + Ollama reference agent"
AGENT_DESCRIPTION = (
    "Reference implementation for Badge's OpenTelemetry integration: a "
    "LangGraph supervisor with researcher and writer subagents and three "
    "tools, running llama3.2:3b locally via Ollama. Emits OTel GenAI spans "
    "through a callback bridge so Badge can auto-discover its architecture. "
    "Source: the end-to-end Ollama walkthrough in Badge's public docs."
)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", required=True)
    parser.add_argument("--endpoint", required=True, help="public /execute URL")
    parser.add_argument("--email")
    parser.add_argument("--password")
    parser.add_argument("--agent-name", default=AGENT_NAME)
    parser.add_argument("--task-count", type=int, default=3)
    parser.add_argument("--label", default="live")
    parser.add_argument("--reuse-agent-id")
    # Badge's robustness axis needs REPEATED runs of the SAME task (it scores
    # agreement across them). One run per task leaves robustness null.
    parser.add_argument("--repeats", type=int, default=1)
    # Stop after provisioning, so the agent server can be restarted with the
    # agent-scope OTLP token before any screening traffic flows.
    parser.add_argument("--setup-only", action="store_true")
    parser.add_argument("--state-file", default=None)
    args = parser.parse_args()

    # --email without --password (or vice versa) silently fell through to the
    # throwaway-registration branch below, so a user who fat-fingered one flag
    # got a brand-new anonymous account instead of their own — and only found
    # out when their agent was missing from their dashboard. Fail loudly.
    if bool(args.email) != bool(args.password):
        raise SystemExit(
            "--email and --password must be given together (or both omitted, "
            "which registers a throwaway account)."
        )

    base = f"{args.base_url.rstrip('/')}/api/v1"
    client = httpx.Client(timeout=120, follow_redirects=False)
    EVIDENCE.mkdir(exist_ok=True)
    state_path = Path(args.state_file) if args.state_file else (
        EVIDENCE / f"session_{args.label}.json"
    )

    # Resume a previous --setup-only invocation if one is on disk.
    if state_path.exists() and not args.reuse_agent_id:
        state = json.loads(state_path.read_text())
        print(f"resuming session from {state_path}")
    else:
        state = {"base_url": args.base_url, "endpoint": args.endpoint}

    # 1 — authenticate
    if state.get("auth_token"):
        token = state["auth_token"]
    elif args.email and args.password:
        r = client.post(
            f"{base}/auth/login",
            json={"email": args.email, "password": args.password},
        )
        if r.status_code != 200:
            raise SystemExit(f"login failed {r.status_code}: {r.text[:300]}")
        state["auth_mode"] = "existing account"
        token = r.json()["access_token"]
    else:
        email = f"otel-reference-{uuid.uuid4().hex[:8]}@example.com"
        password = f"Otel-Ref-{uuid.uuid4().hex[:10]}!x"
        r = client.post(
            f"{base}/auth/register",
            json={"email": email, "password": password, "name": "OTel Reference"},
        )
        if r.status_code >= 400:
            raise SystemExit(f"register failed {r.status_code}: {r.text[:300]}")
        state["auth_mode"] = "throwaway account"
        state["email"] = email
        # The password is deliberately NOT persisted. The access token below
        # is enough to resume, and a password on disk is a credential with a
        # much longer life than this script's session. Printed once instead —
        # if you need to log back in later, capture it now.
        print(f"throwaway account: {email}  password: {password}")
        token = r.json()["access_token"]
    state["auth_token"] = token
    auth = {"Authorization": f"Bearer {token}"}
    # Persist immediately: a later failure (e.g. the 409 on a duplicate agent
    # name — names are globally unique across ALL accounts, not per-account)
    # would otherwise strand a freshly registered account with no way back in.
    write_secret_json(state_path, state)

    # 2 — register (or reuse) the agent, pointed at the PUBLIC tunnel URL
    if args.reuse_agent_id:
        agent_id = args.reuse_agent_id
    elif state.get("agent_id"):
        agent_id = state["agent_id"]
    else:
        r = client.post(
            f"{base}/agents",
            headers=auth,
            json={
                "name": args.agent_name,
                "description": AGENT_DESCRIPTION,
                "endpoint": args.endpoint,
                "connection_mode": "http_endpoint",
                "payload_schema": "badge_native",
                "is_public": True,
            },
        )
        if r.status_code >= 400:
            raise SystemExit(f"agent create failed {r.status_code}: {r.text[:400]}")
        agent = r.json()
        agent_id = agent["id"]
        state["agent"] = {
            "id": agent_id,
            "slug": agent.get("slug"),
            "is_public": agent.get("is_public"),
        }
    state["agent_id"] = agent_id

    # 2b — mint the AGENT-SCOPE OTLP token. This is the credential a real
    # agent uses: per-run tokens are handed to whoever CREATES the run, not to
    # the agent being screened, so an agent can never see them.
    if not state.get("otlp_agent_token"):
        r = client.post(f"{base}/agents/{agent_id}/otlp-token", headers=auth)
        if r.status_code == 200:
            state["otlp_agent_token"] = r.json().get("otlp_ingest_token")
            state["otlp_agent_token_action"] = r.json().get("action")
            if state["otlp_agent_token_action"] == "rotated":
                # Minting is single-reveal AND rotating: the previous token is
                # dead as of this response. A server started earlier with that
                # token now 401s on every export — screening still passes, the
                # blueprint just stays silently empty. This exact sequence
                # (setup-only under one --label, screening under another →
                # different state file → re-mint) cost a full six-run session
                # of spans on 2026-08-04.
                print(
                    "WARNING: agent OTLP token ROTATED — any agent server "
                    "already running with the previous token is now exporting "
                    "into 401s. Restart it with the new token before screening."
                )
        else:
            state["otlp_agent_token"] = None
            state["otlp_agent_token_error"] = {
                "status": r.status_code,
                "body": r.text[:300],
            }

    if args.setup_only:
        write_secret_json(state_path, state)
        redacted = {k: v for k, v in state.items()
                    if k not in ("auth_token", "otlp_agent_token")}
        print(json.dumps(redacted, indent=2))
        print(f"\nwrote {state_path}")
        print("\nNow restart the agent server with:")
        print(f"  BADGE_OTLP_ENDPOINT={args.base_url.rstrip('/')}/api/v1/provenance/otlp")
        print("  BADGE_OTLP_TOKEN=<otlp_agent_token from " + str(state_path) + ">")
        return

    # 3 — pick tasks
    r = client.get(f"{base}/tasks", params={"page_size": args.task_count}, headers=auth)
    r.raise_for_status()
    tasks = r.json()["items"]
    state["tasks"] = [{"id": t["id"], "title": t.get("title")} for t in tasks]

    # 4 — screen against each task
    runs: list[dict] = []
    schedule = [t for t in tasks for _ in range(max(1, args.repeats))]
    for task in schedule:
        r = client.post(
            f"{base}/runs",
            headers=auth,
            json={"agent_id": agent_id, "task_id": task["id"]},
        )
        if r.status_code >= 400:
            runs.append({"task": task.get("title"), "error": r.text[:300]})
            continue
        created = r.json()
        run_id = created["id"]
        record = {
            "run_id": run_id,
            "task": task.get("title"),
            # Present ONLY when BPP is enabled server-side.
            "otlp_token_minted": bool(created.get("otlp_ingest_token")),
            "otlp_endpoint": created.get("otlp_endpoint"),
        }
        if created.get("otlp_ingest_token"):
            record["otlp_ingest_token"] = created["otlp_ingest_token"]

        for _ in range(90):
            d = client.get(f"{base}/runs/{run_id}", headers=auth)
            d.raise_for_status()
            detail = d.json()
            if detail["status"] in {"completed", "failed", "timeout"}:
                record.update(
                    status=detail["status"],
                    execution_mode=detail.get("execution_mode"),
                    score=detail.get("score"),
                    latency_ms=detail.get("latency_ms"),
                    cost_usd=detail.get("cost_usd"),
                    success=detail.get("success"),
                    error_class=detail.get("error_class"),
                    provenance_level=detail.get("provenance_level"),
                )
                break
            time.sleep(2)
        else:
            record["status"] = "TIMED OUT waiting for terminal status"
        runs.append(record)
        print(json.dumps(record, indent=2))

    state["runs"] = runs

    # 5 — read back the public surfaces
    for name, path in (
        ("agent_detail", f"/agents/{agent_id}"),
        ("agent_stats", f"/agents/{agent_id}/stats"),
        ("agent_fitness", f"/agents/{agent_id}/fitness"),
        ("agent_blueprint", f"/agents/{agent_id}/blueprint"),
    ):
        r = client.get(f"{base}{path}", headers=auth)
        state[name] = (
            r.json() if r.status_code == 200 else {"status": r.status_code, "body": r.text[:300]}
        )

    write_secret_json(state_path, state)
    out = EVIDENCE / f"live_{args.label}.json"
    write_secret_json(out, state)
    print(f"\nwrote {out}")
    print(f"agent_id={agent_id}")


if __name__ == "__main__":
    main()

"""Measure what the fake backend cannot: real run-to-run variance.

Badge's robustness axis scores an agent on whether repeated runs of the same
task agree. A scripted fake model scores a trivially perfect 1.0 — the same
dict lookup returns the same bytes forever — so the fake backend can only
ever produce a baseline that flatters itself.

This runs one prompt N times against the live server and reports the
distributions that actually drive a score: distinct outputs, token counts,
per-call and end-to-end latency (with the cold-start call called out
separately, because the first inference pays the model load).

    ./.venv/bin/python measure_variance.py --runs 5
    ./.venv/bin/python measure_variance.py --runs 5 --prompt "..." --json out.json
"""

from __future__ import annotations

import argparse
import hashlib
import json
import statistics
import time
from pathlib import Path

import httpx

DEFAULT_PROMPT = (
    'Given this JSON: {"user":{"id":42,"name":"Alice"}}. Extract: user id, name.'
)


def percentile(values: list[float], fraction: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(len(ordered) - 1, int(round(fraction * (len(ordered) - 1))))
    return ordered[index]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--runs", type=int, default=5)
    parser.add_argument("--url", default="http://127.0.0.1:8787/execute")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--json", dest="json_path", default=None)
    args = parser.parse_args()

    client = httpx.Client(timeout=120)
    outputs: list[str] = []
    wall_ms: list[float] = []
    input_tokens: list[int] = []
    output_tokens: list[int] = []
    call_latencies: list[float] = []
    first_call_latencies: list[float] = []
    tool_sequences: list[str] = []
    costs: list[float] = []

    for index in range(args.runs):
        started = time.perf_counter()
        response = client.post(
            args.url,
            json={
                "task_id": f"variance-{index}",
                "prompt": args.prompt,
                "max_tokens": 256,
            },
            headers={"X-Badge-Run": "true"},
        )
        elapsed_ms = (time.perf_counter() - started) * 1000
        response.raise_for_status()
        body = response.json()

        outputs.append(body["output"])
        wall_ms.append(elapsed_ms)
        input_tokens.append(body.get("input_tokens", 0))
        output_tokens.append(body.get("output_tokens", 0))
        costs.append(body.get("total_cost_usd", 0.0))
        agent = body.get("agent", {})
        latencies = agent.get("llm_latencies_ms", [])
        if latencies:
            first_call_latencies.append(latencies[0])
            call_latencies.extend(latencies)
        tool_sequences.append(",".join(agent.get("tools_called", [])) or "(none)")
        print(
            f"run {index + 1}/{args.runs}: {elapsed_ms:7.0f} ms  "
            f"{body.get('input_tokens'):>5} in / {body.get('output_tokens'):>4} out  "
            f"tools=[{tool_sequences[-1]}]  {body['output'][:70]!r}"
        )

    distinct = {hashlib.sha256(o.encode()).hexdigest(): o for o in outputs}
    exact_agreement = max(outputs.count(o) for o in set(outputs)) / len(outputs)

    summary = {
        "runs": args.runs,
        "prompt": args.prompt,
        "distinct_outputs": len(distinct),
        "modal_agreement_rate": round(exact_agreement, 3),
        "tool_sequences_distinct": len(set(tool_sequences)),
        "tool_sequences": sorted(set(tool_sequences)),
        "wall_ms": {
            "min": round(min(wall_ms)),
            "median": round(statistics.median(wall_ms)),
            "p95": round(percentile(wall_ms, 0.95)),
            "max": round(max(wall_ms)),
        },
        "llm_call_ms": {
            "count": len(call_latencies),
            "min": round(min(call_latencies)) if call_latencies else 0,
            "median": round(statistics.median(call_latencies)) if call_latencies else 0,
            "p95": round(percentile(call_latencies, 0.95)) if call_latencies else 0,
            "max": round(max(call_latencies)) if call_latencies else 0,
        },
        "first_call_ms": {
            "cold_start_run1": round(first_call_latencies[0])
            if first_call_latencies
            else 0,
            "median_of_later_runs": round(statistics.median(first_call_latencies[1:]))
            if len(first_call_latencies) > 1
            else 0,
        },
        "input_tokens": {
            "min": min(input_tokens),
            "median": round(statistics.median(input_tokens)),
            "max": max(input_tokens),
        },
        "output_tokens": {
            "min": min(output_tokens),
            "median": round(statistics.median(output_tokens)),
            "max": max(output_tokens),
        },
        "cost_usd": {
            "min": min(costs),
            "median": statistics.median(costs),
            "max": max(costs),
            "total": round(sum(costs), 10),
        },
    }

    print("\n" + json.dumps(summary, indent=2))
    print(
        f"\nVERDICT: {len(distinct)} distinct output(s) across {args.runs} runs "
        f"(modal agreement {exact_agreement:.0%}). "
        + (
            "A deterministic agent would print 1."
            if len(distinct) > 1
            else "Byte-identical across runs."
        )
    )
    if len(distinct) > 1:
        print("\ndistinct outputs:")
        for digest, text in distinct.items():
            print(f"  [{digest[:8]}] {text[:160]!r}")

    if args.json_path:
        Path(args.json_path).write_text(json.dumps(summary, indent=2))
        print(f"\nwrote {args.json_path}")


if __name__ == "__main__":
    main()