Designing Tools

{
  "type": "function",
  "name": "search_issues",
  "strict": true,
  "parameters": {
    "properties": {
      "query": {"description": "Keywords or an exact term, e.g. 'ORB-1402' or 'ran twice'.",
                "title": "Query", "type": "string"},
      "state": {"anyOf": [{"enum": ["open", "closed"], "type": "string"}, {"type": "null"}],
                "description": "Only open or closed issues; null for both.", "title": "State"},
      "limit": {"default": 5, "description": "Maximum issues to return.",
                "maximum": 20, "minimum": 1, "title": "Limit", "type": "integer"}
    },
    "required": ["query", "state", "limit"],
    "title": "IssuesSearchArgs",
    "type": "object",
    "additionalProperties": false
  },
  "description": "Search Orbit's issue tracker for bug reports and user questions."
}

# script_id: day_036_tool_design/schema_preview
# Print the strict Responses tool the SDK generates from a Pydantic model. Runs offline (no key).
# Install: uv add "openai>=3.46,<4" pydantic
import json
from typing import Literal

from openai import pydantic_responses_function_tool
from pydantic import BaseModel, Field


class IssuesSearchArgs(BaseModel):
    query: str = Field(description="Keywords or an exact term, e.g. 'ORB-1402' or 'ran twice'.")
    state: Literal["open", "closed"] | None = Field(None, description="Only open or closed issues; null for both.")
    limit: int = Field(5, ge=1, le=20, description="Maximum issues to return.")


tool = pydantic_responses_function_tool(
    IssuesSearchArgs,
    name="search_issues",
    description="Search Orbit's issue tracker for bug reports and user questions.",
)
print(json.dumps(tool, indent=2))

{
  "type": "function",
  "name": "search_issues",
  "strict": true,
  "parameters": {
    "properties": {
... (the rest matches the JSON in "Reading a Strict Schema by Hand")

# script_id: day_036_tool_design/first_pydantic_tool
# One question -> one typed tool call. Needs OPENAI_API_KEY; ATLAS_MODEL is optional.
# Install: uv add "openai>=3.46,<4" pydantic
import os
from typing import Literal

from openai import OpenAI, pydantic_responses_function_tool
from pydantic import BaseModel, Field

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
client = OpenAI()


class IssuesSearchArgs(BaseModel):
    query: str = Field(description="Keywords or an exact term, e.g. 'ORB-1402' or 'ran twice'.")
    state: Literal["open", "closed"] | None = Field(None, description="Only open or closed issues; null for both.")
    limit: int = Field(5, ge=1, le=20, description="Maximum issues to return.")


search_issues = pydantic_responses_function_tool(
    IssuesSearchArgs, name="search_issues", description="Search Orbit's issue tracker for bug reports.")
resp = client.responses.parse(model=MODEL, input="Find open issues that mention ORB-1101.", tools=[search_issues])
for item in resp.output:
    if item.type == "function_call":
        print(item.name, repr(item.parsed_arguments))

search_issues IssuesSearchArgs(query='ORB-1101', state='open', limit=5)

# script_id: day_036_tool_design/strict_limits
# Strict mode's limits, offline: a local rule check, a schema strict can't express, a valid-but-wrong call.
# Install: uv add "openai>=3.46,<4" pydantic    Env: ORBIT_DATA (default datasets/orbit)
import json
import os
import sqlite3
from contextlib import closing
from pathlib import Path

from openai import pydantic_responses_function_tool
from pydantic import BaseModel, Field

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))


def strict_problems(schema: dict, path: str = "$") -> list[str]:
    """Check the two documented strict rules before the API sees the schema."""
    problems = []
    if schema.get("type") == "object":
        props = schema.get("properties", {})
        if schema.get("additionalProperties") is not False:
            problems.append(f"{path}: additionalProperties must be false")
        if missing := sorted(set(props) - set(schema.get("required", []))):
            problems.append(f"{path}: not in required: {missing}")
        for key, sub in props.items():
            problems += strict_problems(sub, f"{path}.{key}")
    for i, sub in enumerate(schema.get("anyOf", [])):
        problems += strict_problems(sub, f"{path}.anyOf[{i}]")
    if isinstance(schema.get("items"), dict):
        problems += strict_problems(schema["items"], f"{path}[]")
    for name, sub in schema.get("$defs", {}).items():
        problems += strict_problems(sub, f"$defs.{name}")
    return problems

# script_id: day_036_tool_design/strict_limits
# (a) Hand-written: `state` is optional the plain JSON Schema way, which strict rejects.
hand_written = {"type": "object", "additionalProperties": False,
                "properties": {"query": {"type": "string"}, "state": {"type": "string"}},
                "required": ["query"]}
print("(a)", strict_problems(hand_written))


# (b) A free-form map: Pydantic accepts it and the SDK emits it, but strict can't express it.
class LabelArgs(BaseModel):
    issue_id: int
    labels: dict[str, str] = Field(description="label name -> value")


print("(b)", strict_problems(pydantic_responses_function_tool(LabelArgs, name="set_labels")["parameters"]))


# (c) Valid shape, wrong fact: strict guarantees an int, not an issue that exists.
class IssueArgs(BaseModel):
    issue_id: int = Field(description="Issue number, e.g. 143.")


args = IssueArgs.model_validate_json('{"issue_id": 99999}')  # passes: 99999 IS an int
with closing(sqlite3.connect(f"file:{DATA_DIR / 'issues' / 'orbit.db'}?mode=ro", uri=True)) as db:
    row = db.execute("SELECT title FROM issues WHERE id = ?", (args.issue_id,)).fetchone()
    top = db.execute("SELECT max(id) FROM issues").fetchone()[0]
if row is None:
    print("(c)", json.dumps({"error": f"No issue #{args.issue_id}.",
                             "hint": f"Issue ids run 1-{top}; use search_issues to find one by topic."}))

(a) ["$: not in required: ['state']"]
(b) ['$.labels: additionalProperties must be false']
(c) {"error": "No issue #99999.", "hint": "Issue ids run 1-240; use search_issues to find one by topic."}

# script_id: day_036_tool_design/atlas_tools
# atlas/tools.py v2: @tool registry, validating dispatch(), the v2 Orbit toolset and the v1 set for the eval.
# Install: uv add "openai>=3.46,<4" pydantic tiktoken    Env: ORBIT_DATA (default datasets/orbit)
# Services: none; reads docs/*.md and a read-only orbit.db. Runs offline. Save as atlas_tools.py.
import ast
import difflib
import inspect
import json
import operator
import os
import re
import sqlite3
from contextlib import closing
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable, Literal

import tiktoken
from openai import pydantic_responses_function_tool
from pydantic import BaseModel, Field, ValidationError

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))


@dataclass
class ToolSpec:
    name: str
    args_model: type[BaseModel]
    fn: Callable[[Any], Any]
    description: str
    side_effect: bool = False  # metadata only today; Day 37 puts these tools behind approval

    def as_tool(self) -> dict:
        return pydantic_responses_function_tool(self.args_model, name=self.name, description=self.description)

# script_id: day_036_tool_design/atlas_tools
REGISTRY: dict[str, ToolSpec] = {}  # v2: the canonical Atlas tools
LEGACY: dict[str, ToolSpec] = {}    # v1: kept only so the eval can compare


class ToolInputError(Exception):
    """A well-formed call that can't be served. `hint` tells the model what to do instead."""

    def __init__(self, error: str, hint: str) -> None:
        super().__init__(error)
        self.hint = hint


def tool(args_model: type[BaseModel], *, name: str | None = None, side_effect: bool = False,
         registry: dict[str, ToolSpec] = REGISTRY) -> Callable:
    """Register fn(args) as a tool. The docstring becomes the description the model reads."""
    def register(fn: Callable) -> Callable:
        spec = ToolSpec(name or fn.__name__, args_model, fn, inspect.cleandoc(fn.__doc__ or ""), side_effect)
        if spec.name in registry:
            raise ValueError(f"duplicate tool name {spec.name!r}")  # names are a contract
        registry[spec.name] = spec
        return fn
    return register


def tool_defs(registry: dict[str, ToolSpec] = REGISTRY) -> list[dict]:
    """The definitions you send as `tools=`; regenerated, never hand-edited."""
    return [spec.as_tool() for spec in registry.values()]

# script_id: day_036_tool_design/atlas_tools
def _err(error: str, hint: str) -> str:
    return json.dumps({"error": error, "hint": hint})


def dispatch(name: str, arguments: str, registry: dict[str, ToolSpec] = REGISTRY) -> str:
    """Validate, run, and return a string the model can read. Bad input never raises."""
    spec = registry.get(name)
    if spec is None:
        return _err(f"Unknown tool {name!r}.", f"Available tools: {', '.join(registry)}.")
    try:
        args = spec.args_model.model_validate_json(arguments)
    except ValidationError as exc:
        where = "; ".join(f"{'.'.join(map(str, e['loc'])) or 'arguments'}: {e['msg']}" for e in exc.errors())
        return _err(f"Invalid arguments for {name}: {where}.", "Fix those fields and call again.")
    try:
        result = spec.fn(args)
    except ToolInputError as exc:
        return _err(str(exc), exc.hint)
    return result if isinstance(result, str) else json.dumps(result, ensure_ascii=False)

# script_id: day_036_tool_design/atlas_tools
@lru_cache(maxsize=1)
def load_docs() -> dict[str, dict]:
    """doc_id -> {title, applies_to, heads, body}, parsed from each page's front matter."""
    docs = {}
    for path in sorted((DATA_DIR / "docs").glob("*.md")):
        _, front, body = path.read_text(encoding="utf-8").split("---", 2)
        meta = {k.strip(): v.strip().strip('"') for k, v in (ln.split(":", 1) for ln in front.strip().splitlines())}
        heads = " ".join(ln for ln in body.splitlines() if ln.startswith("#")).lower()
        docs[path.stem] = {"title": meta.get("title", path.stem), "applies_to": meta.get("applies_to", "all"),
                           "heads": heads, "body": body}
    return docs


STOP = {"the", "and", "for", "what", "does", "how", "with", "mean", "about", "are", "any", "has", "anyone"}


def keywords(text: str) -> list[str]:
    return [w for w in re.findall(r"[a-z0-9_]+(?:-[a-z0-9_]+)*", text.lower()) if len(w) > 2 and w not in STOP]


def rank_docs(query: str, k: int) -> list[tuple[int, str]]:
    words = keywords(query)
    # capped body hits + bonuses for title and headings, so one page defining "ORB-1202" beats ten mentions
    scored = [(sum(min(d["body"].lower().count(w), 3) + 5 * (w in d["title"].lower()) + 5 * (w in d["heads"])
                   for w in words), doc_id)
              for doc_id, d in load_docs().items()]
    return sorted((s for s in scored if s[0] > 0), reverse=True)[:k]


class DocsSearchArgs(BaseModel):
    query: str = Field(description="Keywords or an exact term, e.g. 'ORB-1202' or 'lease_seconds'.")
    k: int = Field(3, ge=1, le=10, description="How many pages to return.")


@tool(DocsSearchArgs)
def search_docs(args: DocsSearchArgs) -> dict:
    """Search Orbit's documentation: how-to guides, configuration keys, CLI flags and ORB-xxxx error codes.
    Use for "how do I...", "what does X mean" and "what is the default for Y".
    Do NOT use for bug reports or what users experienced: use search_issues.
    Returns doc_id, title and a one-line snippet; call get_doc_section(doc_id) for the text."""
    words, hits = keywords(args.query), []
    for _, doc_id in rank_docs(args.query, args.k):
        doc = load_docs()[doc_id]
        line = next((ln for ln in doc["body"].splitlines() if any(w in ln.lower() for w in words)), "")
        hits.append({"doc_id": doc_id, "title": doc["title"], "snippet": line.strip()[:160]})
    return {"results": hits}

# script_id: day_036_tool_design/atlas_tools
class DocSectionArgs(BaseModel):
    doc_id: str = Field(description="A doc_id from search_docs, e.g. 'retries-and-backoff'.")
    heading: str | None = Field(None, description="A '## ' heading to return only that section; null for the page start.")


@tool(DocSectionArgs)
def get_doc_section(args: DocSectionArgs) -> dict:
    """Read one Orbit doc page, or one '## ' section of it, by doc_id. Use after search_docs,
    or when the user names a page. Returns at most 1,500 characters."""
    docs = load_docs()
    if args.doc_id not in docs:
        close = [d for d in docs if args.doc_id in d] or difflib.get_close_matches(args.doc_id, docs, n=1, cutoff=0.4)
        hint = f"Closest: {close[0]!r}. " if close else ""
        raise ToolInputError(f"Unknown doc_id {args.doc_id!r}.", hint + "Call search_docs first to get a valid doc_id.")
    sections = re.split(r"(?m)^## ", docs[args.doc_id]["body"])
    text = sections[0]
    if args.heading:
        match = [s for s in sections[1:] if s.lower().startswith(args.heading.lower())]
        if not match:
            heads = [s.splitlines()[0] for s in sections[1:]]
            raise ToolInputError(f"No section {args.heading!r} in {args.doc_id}.", f"Real headings: {heads}.")
        text = "## " + match[0]
    return {"doc_id": args.doc_id, "text": text[:1500], "truncated": len(text) > 1500}

# script_id: day_036_tool_design/atlas_tools
def connect_ro() -> sqlite3.Connection:  # built on Day 33: the file refuses writes
    return sqlite3.connect(f"file:{DATA_DIR / 'issues' / 'orbit.db'}?mode=ro", uri=True)


class IssuesSearchArgs(BaseModel):
    query: str = Field(description="Keywords or an exact term, e.g. 'DST ran twice' or 'ORB-1402'.")
    state: Literal["open", "closed"] | None = Field(None, description="Only open or closed issues; null for both.")
    limit: int = Field(5, ge=1, le=20, description="Maximum issues to return.")
    detail: Literal["concise", "detailed"] = Field(
        "concise", description="'concise': id, title, state. 'detailed' adds the first 600 characters of the body.")


@tool(IssuesSearchArgs)
def search_issues(args: IssuesSearchArgs) -> dict:
    """Search Orbit's issue tracker: bug reports, feature requests and what users experienced.
    Use for "has anyone reported...", "are there open issues about..." or a symptom seen in production.
    Do NOT use for how-to questions or error-code meanings: use search_docs. To read one issue, use get_issue."""
    words = keywords(args.query)[:6] or [args.query.lower()]
    score = " + ".join(["(lower(title || ' ' || body) LIKE ?)"] * len(words))  # placeholders only
    where = f"({score}) > 0" + (" AND state = ?" if args.state else "")
    sql = f"SELECT id, title, state, body FROM issues WHERE {where} ORDER BY ({score}) DESC, id DESC LIMIT ?"
    likes = [f"%{w}%" for w in words]
    with closing(connect_ro()) as db:
        rows = db.execute(sql, likes + ([args.state] if args.state else []) + likes + [args.limit]).fetchall()
    hits = [{"id": i, "title": t, "state": s} | ({"body": b[:600]} if args.detail == "detailed" else {})
            for i, t, s, b in rows]
    return {"results": hits, "hint": "Call get_issue(issue_id) for the body, comments and labels."}

# script_id: day_036_tool_design/atlas_tools
class IssueArgs(BaseModel):
    issue_id: int = Field(description="Issue number, e.g. 143. Not an ORB-xxxx error code.")
    include: list[Literal["body", "comments", "labels"]] = Field(
        default_factory=lambda: ["body"], description="Parts to add to id, title and state.")


@tool(IssueArgs)
def get_issue(args: IssueArgs) -> dict:
    """Fetch one Orbit issue by number, optionally with its body, first comments and labels.
    Use when the user gives an issue number, or after search_issues. Long text is truncated."""
    with closing(connect_ro()) as db:
        row = db.execute("SELECT id, title, state, body FROM issues WHERE id = ?", (args.issue_id,)).fetchone()
        if row is None:
            top = db.execute("SELECT max(id) FROM issues").fetchone()[0]
            raise ToolInputError(f"No issue #{args.issue_id}.", f"Issue ids run 1-{top}. For an ORB-xxxx code "
                                 "use search_docs; to find issues by topic use search_issues.")
        out: dict[str, Any] = {"id": row[0], "title": row[1], "state": row[2]}
        if "body" in args.include:
            out["body"] = row[3][:800]
        if "comments" in args.include:
            out["comments"] = [{"author": a, "text": b[:300]} for a, b in db.execute(
                "SELECT author, body FROM comments WHERE issue_id = ? ORDER BY position LIMIT 3", (row[0],))]
        if "labels" in args.include:
            out["labels"] = [n for (n,) in db.execute(
                "SELECT l.name FROM labels l JOIN issue_labels il ON il.label_id = l.id WHERE il.issue_id = ?",
                (row[0],))]
    return out

# script_id: day_036_tool_design/atlas_tools
class CalcArgs(BaseModel):
    expression: str = Field(description="Arithmetic only: numbers, + - * / // % ** and parentheses, e.g. '(2+4+8)*1250'.")


_OPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv,
        ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod, ast.Pow: operator.pow}


def _eval(node: ast.AST) -> float:
    if isinstance(node, ast.Constant) and type(node.value) in (int, float):
        return node.value
    if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
        return -_eval(node.operand)
    if isinstance(node, ast.BinOp) and type(node.op) in _OPS:
        left, right = _eval(node.left), _eval(node.right)
        if isinstance(node.op, ast.Pow) and abs(right) > 100:
            raise ToolInputError("Exponent too large.", "Keep exponents at or below 100.")
        return _OPS[type(node.op)](left, right)
    raise ToolInputError(f"Unsupported element: {type(node).__name__}.", "Use only numbers, + - * / // % ** and ( ).")


@tool(CalcArgs)
def calculate(args: CalcArgs) -> dict:
    """Evaluate an arithmetic expression exactly. Use it for ANY arithmetic (lease and heartbeat maths,
    retry delays, throughput) instead of computing in your head. Not for dates or unit conversion."""
    try:
        return {"expression": args.expression, "result": _eval(ast.parse(args.expression, mode="eval").body)}
    except (SyntaxError, ZeroDivisionError) as exc:
        raise ToolInputError(f"Cannot evaluate {args.expression!r}: {exc}.", "Send a plain arithmetic expression.") from exc

# script_id: day_036_tool_design/atlas_tools
class Q(BaseModel):
    q: str


class DocId(BaseModel):
    id: str


class IssueId(BaseModel):
    issue_id: int


@tool(Q, name="search", registry=LEGACY)
def v1_search(args: Q) -> dict:
    """Search."""
    return {"docs": [load_docs()[d]["body"][:3000] for _, d in rank_docs(args.q, 3)]}


@tool(DocId, name="get", registry=LEGACY)
def v1_get(args: DocId) -> str:
    """Get a document."""
    return load_docs().get(args.id, {}).get("body", "not found")


@tool(Q, name="lookup_issue", registry=LEGACY)
def v1_lookup_issue(args: Q) -> dict:
    """Look up an issue."""
    ids = [int(args.q.strip("#"))] if args.q.strip("#").isdigit() else \
        [h["id"] for h in search_issues(IssuesSearchArgs(query=args.q))["results"]]
    with closing(connect_ro()) as db:  # whole bodies: the v1 token bill
        return {"issues": [db.execute("SELECT id, title, state, body FROM issues WHERE id = ?", (i,)).fetchone()
                           for i in ids]}


@tool(IssueId, name="get_comments", registry=LEGACY)
def v1_get_comments(args: IssueId) -> dict:
    """Get comments."""
    return get_issue(IssueArgs(issue_id=args.issue_id, include=["comments"]))


@tool(CalcArgs, name="calc", registry=LEGACY)
def v1_calc(args: CalcArgs) -> dict:
    """Calculator."""
    return calculate(args)

# script_id: day_036_tool_design/atlas_tools
TOOLS: list[dict] = tool_defs(REGISTRY)  # the v2 definitions you send as tools=
V1_TOOLS: list[dict] = tool_defs(LEGACY)


@lru_cache(maxsize=1)
def _encoding() -> tiktoken.Encoding:
    return tiktoken.get_encoding("o200k_base")


def count_tokens(text: str) -> int:
    """Built on Day 4: approximate for GPT-6. The input-token count endpoint gives exact numbers."""
    return len(_encoding().encode(text, disallowed_special=()))


if __name__ == "__main__":
    print([t["name"] for t in TOOLS])
    print(dispatch("search_docs", '{"query": "ORB-1202", "k": 2}'))
    print(dispatch("calculate", '{"expression": "300 / 30"}'))

['search_docs', 'get_doc_section', 'search_issues', 'get_issue', 'calculate']
{"results": [{"doc_id": "error-codes", ...
{"expression": "300 / 30", "result": 10.0}

# script_id: day_036_tool_design/actionable_errors
# Five bad calls through dispatch(); each returns an error the model can act on. Runs offline.
# Install: uv add "openai>=3.46,<4" pydantic tiktoken    Needs atlas_tools.py next to this file.
from atlas_tools import dispatch

BAD_CALLS = [
    ("get_doc_section", '{"doc_id": "retries", "heading": null}'),
    ("get_doc_section", '{"doc_id": "retries-and-backoff", "heading": "Jitter"}'),
    # strict mode blocks limit=50, but non-strict paths (MCP in Day 39, tests) don't
    ("search_issues", '{"query": "DST", "state": "open", "limit": 50, "detail": "concise"}'),
    ("get_issue", '{"issue_id": 99999, "include": []}'),
    ("lookup_issue", '{"q": "ORB-1202"}'),
]
for name, arguments in BAD_CALLS:
    print(f"{name}: {dispatch(name, arguments)}")

get_doc_section: {"error": "Unknown doc_id 'retries'.", "hint": "Closest: 'retries-and-backoff'. Call search_docs first to get a valid doc_id."}
get_doc_section: {"error": "No section 'Jitter' in retries-and-backoff.", "hint": "Real headings: ['Global policy: `[retries]`', 'Per-job policy: `@app.job(retries=...)`', 'Backoff strategies', 'What counts as a failed attempt', 'Permanent errors', 'What it looks like in the logs', 'Coming from 2.x', 'See also']."}
search_issues: {"error": "Invalid arguments for search_issues: limit: Input should be less than or equal to 20.", "hint": "Fix those fields and call again."}
get_issue: {"error": "No issue #99999.", "hint": "Issue ids run 1-240. For an ORB-xxxx code use search_docs; to find issues by topic use search_issues."}
lookup_issue: {"error": "Unknown tool 'lookup_issue'.", "hint": "Available tools: search_docs, get_doc_section, search_issues, get_issue, calculate."}

# script_id: day_036_tool_design/concise_vs_detailed
# Tool-output tokens: v1 lookup_issue vs v2 search_issues (concise / detailed). Runs offline.
# Install: uv add "openai>=3.46,<4" pydantic tiktoken    Needs atlas_tools.py next to this file.
import json

from atlas_tools import LEGACY, count_tokens, dispatch

QUERIES = ["DST ran twice", "reconnect storm redis", "dashboard oidc", "ORB-1101 lease expired", "memory leak"]
print(f"{'query':24} {'v1':>6} {'concise':>8} {'detailed':>9}")
for q in QUERIES:
    v1 = dispatch("lookup_issue", json.dumps({"q": q}), LEGACY)
    v2 = {d: dispatch("search_issues", json.dumps({"query": q, "state": None, "limit": 5, "detail": d}))
          for d in ("concise", "detailed")}
    print(f"{q:24} {count_tokens(v1):>6} {count_tokens(v2['concise']):>8} {count_tokens(v2['detailed']):>9}")

query                        v1  concise  detailed
DST ran twice              1691      196      1268
reconnect storm redis     34816      234      1128
dashboard oidc             1218      190       954
ORB-1101 lease expired     1632      173      1163
memory leak                1197      205      1128

# script_id: day_036_tool_design/tool_eval
# 20-case tool-choice eval, v1 vs v2 toolsets, x3 repeats (~120 calls). Needs OPENAI_API_KEY.
# Install: uv add "openai>=3.46,<4" pydantic tiktoken    Env: ATLAS_MODEL, ORBIT_DATA, REPEATS (default 3)
# Needs atlas_tools.py next to this file. Save as tool_eval.py.
import os
from statistics import mean

from openai import OpenAI

from atlas_tools import LEGACY, REGISTRY, TOOLS, V1_TOOLS, count_tokens, dispatch

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
client = OpenAI()
INSTRUCTIONS = "You are Atlas, an assistant for the maintainers of Orbit, a Python background-jobs library."
PRICES = {"gpt-6-luna": {"input": 0.10, "cached": 0.01, "output": 0.50}}  # $/1M tokens, as of 2026-11; verify


class CostTracker:  # built on Day 12: compact
    def __init__(self) -> None:
        self.total_usd = 0.0

    def record(self, model: str, input_tokens: int, output_tokens: int, cached: int = 0) -> float:
        p = PRICES.get(model, PRICES["gpt-6-luna"])
        usd = ((input_tokens - cached) * p["input"] + cached * p["cached"] + output_tokens * p["output"]) / 1e6
        self.total_usd += usd
        return usd

# script_id: day_036_tool_design/tool_eval
TOOL_CHOICE_CASES = [  # expected_tool: a v2 name, or None = answer without a tool
    {"question": "Thanks, that's all for today!", "expected_tool": None},
    {"question": "Rewrite this to sound friendlier: 'Fix your config, it is wrong.'", "expected_tool": None},
    {"question": "In one sentence, what is a background job queue in general?", "expected_tool": None},
    {"question": "What does ORB-1202 mean?", "expected_tool": "search_docs"},
    {"question": "How do I fix ORB-1101 LeaseExpired?", "expected_tool": "search_docs"},
    {"question": "What is the default value of max_payload_kb?", "expected_tool": "search_docs"},
    {"question": "What does the heartbeat_seconds setting control?", "expected_tool": "search_docs"},
    {"question": "What does the orbit doctor command check?", "expected_tool": "search_docs"},
    {"question": "Has anyone reported jobs running twice on the DST change?", "expected_tool": "search_issues"},
    {"question": "How does Orbit handle cron schedules across a DST change?", "expected_tool": "search_docs"},
    {"question": "Have users reported a reconnect storm on 3.1.0?", "expected_tool": "search_issues"},
    {"question": "How do I configure exponential backoff for retries?", "expected_tool": "search_docs"},
    {"question": "Compute (2+4+8+16+32+64) * 1250 for our retry budget.", "expected_tool": "calculate"},
    {"question": "1840 jobs/min at 2.7 s each: what is 1840 * 2.7 / 60 worker slots?", "expected_tool": "calculate"},
    {"question": "Show me issue #143.", "expected_tool": "get_issue"},
    {"question": "What did people say in the comments on issue 161?", "expected_tool": "get_issue"},
    {"question": "Read the 'Backoff strategies' section of the retries-and-backoff page.", "expected_tool": "get_doc_section"},
    {"question": "Open the 'Heartbeats' section of the leases-and-heartbeats doc.", "expected_tool": "get_doc_section"},
    {"question": "Find open issues about dashboard OIDC login.", "expected_tool": "search_issues"},
    {"question": "How do I deploy Orbit workers on Kubernetes?", "expected_tool": "search_docs"},
]

# script_id: day_036_tool_design/tool_eval
V1_EQUIV = {"search": {"search_docs"}, "get": {"get_doc_section"}, "lookup_issue": {"search_issues", "get_issue"},
            "get_comments": {"get_issue"}, "calc": {"calculate"}}
TOOLSETS = {"v1": (V1_TOOLS, LEGACY, V1_EQUIV), "v2": (TOOLS, REGISTRY, {n: {n} for n in REGISTRY})}
tracker = CostTracker()


def choose(question: str, tools: list[dict]):
    """One request. Return the first function_call item, or None if the model answered directly."""
    resp = client.responses.create(model=MODEL, instructions=INSTRUCTIONS, input=question,
                                   tools=tools, parallel_tool_calls=False)
    u = resp.usage
    tracker.record(MODEL, u.input_tokens, u.output_tokens, u.input_tokens_details.cached_tokens)
    return next((i for i in resp.output if i.type == "function_call"), None)


def run_eval(cases: list[dict], task, check, repeats: int = 1) -> list[dict]:  # built on Day 20: compact
    rows = []
    for run in range(repeats):
        for case in cases:
            out = task(case)
            rows.append({"run": run, "case": case, "out": out, "ok": check(case, out)})
    return rows

# script_id: day_036_tool_design/tool_eval
def score_toolset(label: str, repeats: int) -> dict:
    tools, registry, equiv = TOOLSETS[label]
    spent_before = tracker.total_usd

    def task(case: dict) -> dict:
        call = choose(case["question"], tools)
        if call is None:
            return {"tool": None, "tokens": 0}
        return {"tool": call.name, "tokens": count_tokens(dispatch(call.name, call.arguments, registry))}

    def check(case: dict, out: dict) -> bool:
        if case["expected_tool"] is None:
            return out["tool"] is None
        return case["expected_tool"] in equiv.get(out["tool"], set())

    rows = run_eval(TOOL_CHOICE_CASES, task, check, repeats)
    called = [r["out"]["tokens"] for r in rows if r["out"]["tool"]]
    return {"acc": [sum(r["ok"] for r in rows if r["run"] == k) for k in range(repeats)],
            "tokens": mean(called) if called else 0, "usd": tracker.total_usd - spent_before,
            "misses": sorted({(r["case"]["question"][:45], str(r["out"]["tool"])) for r in rows if not r["ok"]})}


if __name__ == "__main__":
    repeats = int(os.environ.get("REPEATS", "3"))
    print(f"{'toolset':8} {'correct per run':18} {'tool-output tokens':>19} {'$':>8}")
    for label in ("v1", "v2"):
        r = score_toolset(label, repeats)
        print(f"{label:8} {str(r['acc']):18} {r['tokens']:>19.0f} {r['usd']:>8.4f}")
        for question, picked in r["misses"]:
            print(f"   miss: {question!r} -> {picked}")

toolset  correct per run    tool-output tokens        $
v1       [12, 13, 12]                     3900   0.0120
   miss: 'What does ORB-1202 mean?' -> lookup_issue
   miss: 'How do I fix ORB-1101 LeaseExpired?' -> lookup_issue
   miss: 'Compute (2+4+8+16+32+64) * 1250 for our ret' -> None
   ...
v2       [18, 19, 18]                      230   0.0150
   miss: 'How does Orbit handle cron schedules across' -> search_issues