Build a multi-agent app Badge can screen

      YOUR MACHINE                                BADGE
 ┌─────────────────────┐                 ┌──────────────────────────┐
 │  agent.py           │                 │  register the agent      │
 │  (LangGraph graph)  │                 │  POST /api/v1/agents     │
 │        │            │                 │           │              │
 │  server.py          │   tunnel        │  screen it               │
 │  POST /execute ◄────┼─────────────────┼── POST <your endpoint>   │
 │        │            │                 │           │              │
 │  improve one thing  │                 │  score + radar + trace   │
 │        ▲            │                 │  /stats /fitness /trace  │
 └────────┼────────────┘                 └───────────┼──────────────┘
          └──────────── read, interpret, decide ◄────┘

POST /execute HTTP/1.1
Content-Type: application/json
User-Agent: Badge/0.1.108
X-Badge-Run: true

{"task_id": "0789d1c1-...",
 "prompt": "Given this JSON: {\"user\": {\"id\": 42, ...}}. Extract: user id, ...",
 "max_tokens": 256,
 "expected_output": "{\"id\": 42, ...}"}

{"output": "{\"id\": 42, \"name\": \"Alice\", \"email\": \"alice@example.com\", \"city\": \"Madrid\"}",
 "input_tokens": 52,
 "output_tokens": 19}

traceparent: 00-<trace_id>-<span_id>-01
X-Badge-Run-Id: <run uuid>
Accept-Signature: ...

START -> supervisor -> researcher -> supervisor -> writer -> supervisor -> END

mkdir badge-agent && cd badge-agent
python3 -m venv .venv
./.venv/bin/pip install langgraph langchain-core fastapi uvicorn

"""A LangGraph supervisor with two subagents and two tools.

    START -> supervisor -> researcher -> supervisor -> writer -> supervisor -> END

The supervisor classifies the task and routes; the researcher gathers evidence
with the tools; the writer renders the answer in the shape the task asked for.
Every decision is deterministic Python, so screening costs $0 in provider
tokens and the same input always produces the same output. Swap the helpers
for model calls and the graph shape stays the same.
"""

from __future__ import annotations

import json
import math
import re
from typing import Any, TypedDict

from langchain_core.tools import tool
from langgraph.graph import END, START, StateGraph

# --- Tool 1: pull structure out of raw text ---------------------------------

_TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")


def _find_json_objects(text: str) -> list[dict[str, Any]]:
    """Scan for balanced {...} spans and return those that parse as objects."""
    found: list[dict[str, Any]] = []
    depth, start = 0, -1
    for index, char in enumerate(text):
        if char == "{":
            if depth == 0:
                start = index
            depth += 1
        elif char == "}" and depth > 0:
            depth -= 1
            if depth == 0 and start >= 0:
                try:
                    parsed = json.loads(text[start : index + 1])
                except ValueError:
                    parsed = None
                if isinstance(parsed, dict):
                    found.append(parsed)
                start = -1
    return found


def _flatten(obj: Any, prefix: str = "") -> dict[str, Any]:
    """Flatten nested dicts/lists into dotted paths -> scalar values."""
    flat: dict[str, Any] = {}
    if isinstance(obj, dict):
        for key, value in obj.items():
            flat.update(_flatten(value, f"{prefix}.{key}" if prefix else str(key)))
    elif isinstance(obj, list):
        for position, value in enumerate(obj):
            flat.update(_flatten(value, f"{prefix}[{position}]"))
    elif prefix:
        flat[prefix] = obj
    return flat


@tool("extract_structured")
def extract_structured(text: str) -> str:
    """Extract embedded JSON objects from text and flatten them to dotted paths.

    Use this whenever a task asks you to extract or reshape data that is
    already present in the prompt.
    """
    records = _find_json_objects(text)
    flat: dict[str, Any] = {}
    for record in records:
        flat.update(_flatten(record))
    return json.dumps({"json_records": records, "flat": flat})

# --- Tool 2: retrieve from an offline corpus --------------------------------

CORPUS = [
    {
        "id": "doc-emea-market",
        "title": "EMEA agent market overview",
        "text": (
            "Enterprise adoption of AI agents in EMEA accelerated through 2026. "
            "Procurement teams increasingly require independent screening before "
            "granting agents access to production systems. Regulated industries "
            "lead the shift because audit trails are already mandatory."
        ),
    },
    {
        "id": "doc-eval-methods",
        "title": "Evaluation methods for agents",
        "text": (
            "Benchmarks for agents measure correctness, latency, cost, and "
            "robustness. Repeated runs of the same task expose nondeterminism "
            "that a single run hides. Deterministic agents trade flexibility "
            "for repeatable results."
        ),
    },
    {
        "id": "doc-hiring",
        "title": "Hiring workflows for agents",
        "text": (
            "Teams that treat agents like hires screen them, compare candidates "
            "on shared tasks, and re-screen after every change. A public track "
            "record shortens the trust-building phase of adoption."
        ),
    },
]

FOCUS_WEIGHT = 3.0
_STOPWORDS = {"a", "an", "and", "for", "from", "in", "is", "of", "on", "that", "the", "to"}


def _tokens(text: str) -> list[str]:
    return [t.lower() for t in _TOKEN_RE.findall(text) if t.lower() not in _STOPWORDS]


def _idf_weights() -> dict[str, float]:
    seen: dict[str, int] = {}
    for document in CORPUS:
        for token in set(_tokens(f"{document['title']} {document['text']}")):
            seen[token] = seen.get(token, 0) + 1
    return {token: math.log(1 + len(CORPUS) / count) for token, count in seen.items()}


@tool("search_corpus")
def search_corpus(query: str, focus: str = "") -> str:
    """Retrieve passages from the agent's offline knowledge corpus.

    `query` carries the broad topic; `focus` carries the specific angle and is
    weighted more heavily, so one topic can be searched from several
    directions. Returns ranked hits as JSON.
    """
    weights = _idf_weights()
    query_weights: dict[str, float] = {}
    for token in _tokens(query):
        query_weights[token] = query_weights.get(token, 0.0) + 1.0
    for token in _tokens(focus):
        query_weights[token] = query_weights.get(token, 0.0) + FOCUS_WEIGHT

    scored = []
    for document in CORPUS:
        frequencies: dict[str, int] = {}
        for token in _tokens(f"{document['title']} {document['text']}"):
            frequencies[token] = frequencies.get(token, 0) + 1
        score = sum(
            weight * weights.get(token, 0.0) * frequencies.get(token, 0)
            for token, weight in query_weights.items()
        )
        if score > 0:
            scored.append(
                {
                    "id": document["id"],
                    "title": document["title"],
                    "score": round(score, 3),
                    "text": document["text"],
                }
            )
    scored.sort(key=lambda hit: hit["score"], reverse=True)
    return json.dumps({"query": query, "focus": focus, "hits": scored[:2]})


TOOLS_BY_NAME = {t.name: t for t in (extract_structured, search_corpus)}

# --- Graph state ------------------------------------------------------------


class AgentState(TypedDict, total=False):
    prompt: str
    intent: str
    aspects: list[str]
    completed: list[str]
    next_worker: str
    findings: dict[str, Any]
    tool_calls: list[str]
    output: str


# --- Supervisor: classify, then route ---------------------------------------


def _classify(prompt: str) -> str:
    """Deterministic stand-in for a routing LLM call."""
    lowered = prompt.lower()
    if "{" in prompt and "}" in prompt and "extract" in lowered:
        return "json_extraction"
    if re.search(r"\bresearch\b|\bsummar(y|ise|ize)\b|\bcovering\b", lowered):
        return "research_summary"
    return "generic"


def _split_list(blob: str) -> list[str]:
    parts = re.split(r",\s*(?:and\s+)?|\s+and\s+", blob.strip())
    return [part.strip(" .;:\n") for part in parts if part.strip(" .;:\n")]


def _requested_aspects(prompt: str) -> list[str]:
    """Read the 'covering: a, b, and c' clause a task uses to define sections."""
    match = re.search(r"covering:?\s*(.+?)(?:\.\s|\.$|$)", prompt, re.IGNORECASE | re.DOTALL)
    return _split_list(match.group(1)) if match else []


def supervisor(state: AgentState) -> AgentState:
    """Classify on the first pass, then decide which subagent runs next."""
    if "intent" not in state:
        return {
            "intent": _classify(state.get("prompt", "")),
            "aspects": _requested_aspects(state.get("prompt", "")),
            "completed": [],
            "findings": {},
            "tool_calls": [],
            "next_worker": "researcher",
        }
    completed = state.get("completed", [])
    if "researcher" not in completed:
        return {"next_worker": "researcher"}
    if "writer" not in completed:
        return {"next_worker": "writer"}
    return {"next_worker": "done"}


def route(state: AgentState) -> str:
    worker = state.get("next_worker", "done")
    return worker if worker in ("researcher", "writer") else "done"


# --- Subagent 1: researcher — gathers evidence, never writes the answer -----


def _call_tool(tool_calls: list[str], name: str, arguments: dict[str, Any]) -> Any:
    tool_calls.append(name)
    return json.loads(TOOLS_BY_NAME[name].invoke(arguments))


def researcher(state: AgentState) -> AgentState:
    prompt = state.get("prompt", "")
    intent = state.get("intent", "generic")
    tool_calls = list(state.get("tool_calls", []))
    findings: dict[str, Any] = dict(state.get("findings", {}))

    if intent == "json_extraction":
        findings["structured"] = _call_tool(tool_calls, "extract_structured", {"text": prompt})
    elif intent == "research_summary":
        topic_match = re.search(r"research\s+(?:the\s+)?(.+?)(?:\.|\bUse\b|$)", prompt, re.IGNORECASE)
        topic = topic_match.group(1).strip() if topic_match else prompt[:80]
        hits: dict[str, Any] = {}
        for aspect in state.get("aspects") or ["overview"]:
            # One retrieval per requested section, searched from that angle.
            hits[aspect] = _call_tool(tool_calls, "search_corpus", {"query": topic, "focus": aspect})
        findings["topic"], findings["hits"] = topic, hits
    else:
        findings["structured"] = _call_tool(tool_calls, "extract_structured", {"text": prompt})

    return {
        "findings": findings,
        "tool_calls": tool_calls,
        "completed": [*state.get("completed", []), "researcher"],
    }

# --- Subagent 2: writer — renders findings in the requested shape -----------


def _resolve_field(phrase: str, flat_keys: list[str]) -> str | None:
    """Map a requested field ('user id') onto a flattened key path ('user.id').

    Scores candidates by token overlap, weighting the final path segment
    double so 'user id' resolves to user.id rather than user.name.
    """
    wanted = set(_TOKEN_RE.findall(phrase.lower()))
    best: tuple[float, str | None] = (0.0, None)
    for key in flat_keys:
        segments = [s.lower() for s in re.split(r"[.\[\]]+", key) if s]
        score = sum(2.0 for s in segments[-1:] if s in wanted)
        score += sum(1.0 for s in segments[:-1] if s in wanted)
        if score > best[0]:
            best = (score, key)
    return best[1]


def _write_json_extraction(state: AgentState) -> str:
    flat = state.get("findings", {}).get("structured", {}).get("flat", {})
    match = re.search(r"extract:?\s*(.+?)(?:\.\s|\.$|$)", state.get("prompt", ""), re.IGNORECASE | re.DOTALL)
    result: dict[str, Any] = {}
    for phrase in _split_list(match.group(1)) if match else []:
        key = _resolve_field(phrase, list(flat))
        if key is not None:
            label = [s for s in re.split(r"[.\[\]]+", key) if s][-1]
            result[label] = flat[key]
    return json.dumps(result or flat)


def _write_research_summary(state: AgentState) -> str:
    findings = state.get("findings", {})
    lines = [f"# Research summary: {findings.get('topic', 'the requested topic')}", ""]
    sources: list[str] = []
    for aspect, payload in findings.get("hits", {}).items():
        lines.append(f"## {aspect[:1].upper()}{aspect[1:]}")
        for hit in payload.get("hits", [])[:1]:
            sources.append(hit["id"])
            lines.append(f"- {hit['text']}")
        lines.append("")
    lines.append(f"Sources consulted: {', '.join(sorted(set(sources))) or 'none'}")
    lines.append(f"Tool calls made: {len(state.get('tool_calls', []))}")
    return "\n".join(lines).strip()


def _write_generic(state: AgentState) -> str:
    structured = state.get("findings", {}).get("structured", {})
    return json.dumps(
        {
            "task": state.get("prompt", "")[:200],
            "extracted": structured.get("flat", {}),
            "notes": "No specialised handler matched; returned best-effort structure.",
        },
        indent=2,
    )


_WRITERS = {
    "json_extraction": _write_json_extraction,
    "research_summary": _write_research_summary,
}


def writer(state: AgentState) -> AgentState:
    render = _WRITERS.get(state.get("intent", "generic"), _write_generic)
    return {
        "output": render(state),
        "completed": [*state.get("completed", []), "writer"],
    }


# --- Wire the graph ---------------------------------------------------------


def build_graph() -> Any:
    graph = StateGraph(AgentState)
    graph.add_node("supervisor", supervisor)
    graph.add_node("researcher", researcher)
    graph.add_node("writer", writer)
    graph.add_edge(START, "supervisor")
    graph.add_conditional_edges(
        "supervisor",
        route,
        {"researcher": "researcher", "writer": "writer", "done": END},
    )
    graph.add_edge("researcher", "supervisor")
    graph.add_edge("writer", "supervisor")
    return graph.compile()


GRAPH = build_graph()


def run_agent(prompt: str) -> dict[str, Any]:
    """Execute the graph once and return the answer plus execution metadata."""
    final: AgentState = GRAPH.invoke({"prompt": prompt})
    return {
        "output": final.get("output", ""),
        "intent": final.get("intent", "unknown"),
        "tool_calls": final.get("tool_calls", []),
    }


if __name__ == "__main__":
    import sys

    print(json.dumps(run_agent(sys.stdin.read()), indent=2))

"""HTTP wrapper implementing Badge's badge_native screening contract."""

from __future__ import annotations

import math
import os
import time
from typing import Any

from agent import run_agent
from fastapi import FastAPI, Request, Response
from fastapi.responses import JSONResponse

# Report the honest cost of a run. This agent makes no provider calls, so its
# true cost is $0.00. If you omit total_cost_usd, Badge estimates your cost
# from the token counts at GPT-4o reference pricing instead.
REPORT_COST = os.environ.get("BADGE_REPORT_COST", "").lower() in {"1", "true", "yes"}

# Serve exactly two routes and nothing else. A tunnel publishes every path on
# the port, so the interactive docs FastAPI mounts by default would go public.
app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)


def approx_tokens(text: str) -> int:
    """~4 characters per token. Replace with your provider's real usage numbers."""
    return max(1, math.ceil(len(text) / 4)) if text else 0


@app.get("/health")
async def health() -> dict[str, Any]:
    return {"status": "ok", "routes": ["GET /health", "POST /execute"]}


@app.post("/execute")
async def execute(request: Request) -> Response:
    started = time.perf_counter()
    try:
        body = await request.json()
    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)

    # body also carries task_id, max_tokens and — for non-secret tasks —
    # expected_output. Never read expected_output: grading yourself against
    # the answer key is not work, and it is absent on secret holdout tasks.
    result = run_agent(body["prompt"])
    output = result["output"]

    payload: dict[str, Any] = {
        "output": output,
        "input_tokens": approx_tokens(body["prompt"]),
        "output_tokens": approx_tokens(output),
        # Anything extra you return is preserved in the run trace on Badge —
        # a free debugging channel for your graph's internals.
        "agent": {
            "framework": "langgraph",
            "route": result["intent"],
            "tool_calls": result["tool_calls"],
            "local_compute_ms": int((time.perf_counter() - started) * 1000),
        },
    }
    if REPORT_COST:
        payload["total_cost_usd"] = 0.0
    return JSONResponse(payload)


if __name__ == "__main__":
    import uvicorn

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

BADGE_REPORT_COST=1 PORT=8787 ./.venv/bin/python server.py

curl -s http://127.0.0.1:8787/health

curl -s -X POST http://127.0.0.1:8787/execute \
  -H 'Content-Type: application/json' \
  -d '{"task_id":"local","prompt":"Given this JSON: {\"user\":{\"id\":42,\"name\":\"Alice\"}}. Extract: user id, full name.","max_tokens":256}'

curl -s -o /dev/null -w '%{http_code}\n' http://127.0.0.1:8787/openapi.json   # expect 404

for i in 1 2 3; do
  curl -s -X POST http://127.0.0.1:8787/execute \
    -H 'Content-Type: application/json' \
    -d '{"task_id":"local","prompt":"Given this JSON: {\"user\":{\"id\":42}}. Extract: user id.","max_tokens":256}'
  echo
done | sort -u | wc -l   # expect 1