Termination, Budgets & Graceful Degradation

# Colab cell 1 — pure Python: runs with no key and no client.
import time

class Budget:
    # Pull current per-MTok prices from your provider's pricing page.
    # Never hardcode from memory; keep them in one place so tests can pin them.
    PRICE_IN_PER_MTOK = 0.0   # TODO: fill from pricing page
    PRICE_OUT_PER_MTOK = 0.0  # TODO: fill from pricing page

    def __init__(self, max_iterations: int = 15,
                 max_usd: float = 0.50, max_seconds: float = 60.0):
        self.max_iterations = max_iterations
        self.max_usd = max_usd
        self.deadline = time.monotonic() + max_seconds
        self.iterations = 0
        self.usd = 0.0

    def add_call(self, usage) -> None:
        self.iterations += 1
        self.usd += (usage.input_tokens * self.PRICE_IN_PER_MTOK +
                     usage.output_tokens * self.PRICE_OUT_PER_MTOK) / 1_000_000

    def exhausted(self) -> str | None:
        """Return a human-readable reason, or None if we may continue."""
        if self.iterations >= self.max_iterations:
            return f"iteration cap ({self.max_iterations}) reached"
        if self.usd >= self.max_usd:
            return f"cost budget exceeded ({self.usd:.3f} USD)"
        if time.monotonic() >= self.deadline:
            return "wall-clock deadline passed"
        return None

budget = Budget(max_iterations=3, max_usd=0.25, max_seconds=30)
print(budget.exhausted())   # None -> all guards green, the loop may proceed

# Colab cell 2 — run cell 1 first (it defines Budget). This cell adds the
# key, the client, and a fake repo, then runs the fully guarded loop.
!pip install -q anthropic

import os
try:
    from google.colab import userdata
    os.environ["ANTHROPIC_API_KEY"] = userdata.get("ANTHROPIC_API_KEY")
except Exception:
    from getpass import getpass
    os.environ.setdefault("ANTHROPIC_API_KEY", getpass("Anthropic API key: "))

import anthropic

client = anthropic.Anthropic()
MODEL = "claude-sonnet-5"

# The same tiny in-memory "repo" as lesson 2, so the tools do real work.
REPO = {
    "README.md": "Sample service. Retry policy is configured in config/app.yaml.",
    "config/app.yaml": "retries: 5\nbackoff: exponential\ntimeout_s: 30\n",
    "src/retry.py": "def backoff(attempt):\n    return min(2 ** attempt, 60)\n",
}

def list_dir(path: str = "") -> str:
    hits = [p for p in REPO if p.startswith(path)]
    return "\n".join(sorted(hits)) if hits else f"Nothing under {path!r}."

def grep(pattern: str) -> str:
    hits = [f"{p}: {line}" for p, body in REPO.items()
            for line in body.splitlines() if pattern.lower() in line.lower()]
    return "\n".join(hits) if hits else f"No lines match {pattern!r}."

def read_file(path: str) -> str:
    return REPO.get(path, f"No file at {path!r}.")

IMPL = {"list_dir": list_dir, "grep": grep, "read_file": read_file}
TOOLS = [
    {"name": "list_dir", "description": "List repo paths under a prefix.",
     "input_schema": {"type": "object",
                      "properties": {"path": {"type": "string"}}, "required": []}},
    {"name": "grep", "description": "Find lines matching a substring.",
     "input_schema": {"type": "object",
                      "properties": {"pattern": {"type": "string"}},
                      "required": ["pattern"]}},
    {"name": "read_file", "description": "Read one file in full, by path.",
     "input_schema": {"type": "object",
                      "properties": {"path": {"type": "string"}},
                      "required": ["path"]}},
]

def execute_all(content) -> list:
    # naive executor: run every requested tool, pair every result.
    # Lesson 5 hardens this into SafeExecutor (errors, budgets, repeats).
    return [{"type": "tool_result", "tool_use_id": block.id,
             "content": IMPL[block.name](**block.input)}
            for block in content if block.type == "tool_use"]

FINISH_TOOL = {
    "name": "finish",
    "description": (
        "Submit your final answer. Call exactly once, when you have enough "
        "evidence. Every claim must cite a file path you actually read."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "answer": {"type": "string"},
            "citations": {"type": "array", "items": {"type": "string"},
                          "description": "file paths supporting the answer"},
        },
        "required": ["answer", "citations"],
    },
}

def run(question: str, budget: Budget) -> dict:
    messages = [{"role": "user", "content": question}]
    while True:
        reason = budget.exhausted()
        if reason is not None:                    # check BEFORE paying
            return best_effort(messages, reason)

        resp = client.messages.create(
            model=MODEL, max_tokens=2048,
            tools=TOOLS + [FINISH_TOOL], messages=messages,
        )
        budget.add_call(resp.usage)

        finish = next((b for b in resp.content
                       if b.type == "tool_use" and b.name == "finish"), None)
        if finish is not None:
            return {"answer": finish.input["answer"],
                    "citations": finish.input["citations"],
                    "complete": True}

        if resp.stop_reason != "tool_use":
            # model stopped talking without calling finish — nudge it.
            # Not a one-shot: if it keeps stalling this fires every turn;
            # the budget guard above is what actually bounds the retries.
            messages.append({"role": "assistant", "content": resp.content})
            messages.append({"role": "user", "content":
                "Call the finish tool with your answer and citations."})
            continue

        messages.append({"role": "assistant", "content": resp.content})
        messages.append({"role": "user",
                         "content": execute_all(resp.content)})

def best_effort(messages, reason: str) -> dict:
    """Budget is gone. One last cheap call, NO tools, to salvage an answer."""
    wrap_up = messages + [{"role": "user", "content":
        "Budget exhausted (" + reason + "). Using only what you have "
        "already found, give your best answer and state explicitly what "
        "you could not verify."}]
    resp = client.messages.create(model=MODEL, max_tokens=1024,
                                  messages=wrap_up)
    answer = next(b.text for b in resp.content if b.type == "text")
    return {"answer": answer, "citations": [],
            "complete": False, "stop_reason": reason}

print(run("How does this service configure retries?", Budget()))