Phase 3: Provider abstraction

import os
import time
from typing import Protocol

from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

class Completion:
    """One reply from a provider, with what it cost to get it."""

    def __init__(self, message, model: str, prompt_tokens: int, completion_tokens: int, latency_ms: int):
        self.message = message
        self.model = model
        self.prompt_tokens = prompt_tokens
        self.completion_tokens = completion_tokens
        self.latency_ms = latency_ms

class Provider(Protocol):
    def complete(self, messages: list, tools: list | None = None) -> Completion: ...

class GroqProvider:
    def __init__(self, model: str | None = None):
        self.model = model or os.environ.get("MODEL", "openai/gpt-oss-120b")
        self.client = OpenAI(
            api_key=os.environ["GROQ_API_KEY"],
            base_url="https://api.groq.com/openai/v1",
        )

    def complete(self, messages: list, tools: list | None = None) -> Completion:
        kwargs = {"model": self.model, "messages": messages}
        if tools:
            kwargs["tools"] = tools

        start = time.perf_counter()
        response = self.client.chat.completions.create(**kwargs)
        latency_ms = round((time.perf_counter() - start) * 1000)

        return Completion(
            message=response.choices[0].message,
            model=self.model,
            prompt_tokens=response.usage.prompt_tokens,
            completion_tokens=response.usage.completion_tokens,
            latency_ms=latency_ms,
        )

PROVIDERS = {"groq": GroqProvider}

def get_provider(name: str | None = None) -> Provider:
    name = name or os.environ.get("BACKEND", "groq")
    if name not in PROVIDERS:
        raise ValueError(f"unknown backend: {name}")
    return PROVIDERS[name]()

if __name__ == "__main__":
    provider = get_provider()
    result = provider.complete([{"role": "user", "content": "Say hello in one sentence."}])
    print(result.message.content)

import json
import sys
import time

from backend import get_provider
from guardrails import NEEDS_APPROVAL, check, log
from tools import REGISTRY, TOOLS

MAX_STEPS = 10
SYSTEM_PROMPT = "You are a coding assistant. When a task needs a tool, call the tool right away. Do not announce what you are about to do, and never ask the user for permission in your reply. The harness asks for approval by itself."

provider = get_provider()

def approve(name: str, args: dict) -> bool:
    print(f"\nThe model wants to use {name}:")
    for key, value in args.items():
        print(f"  {key}: {value}")
    return input("Allow? [y/N] ").strip().lower() == "y"

def execute(name: str, args: dict) -> tuple[str, str]:
    problem = check(name, args)
    if problem:
        return "blocked", f"blocked: {problem}"
    if name in NEEDS_APPROVAL and not approve(name, args):
        return "denied", "denied: the user said no"
    try:
        return "ok", REGISTRY[name](**args)
    except Exception as error:
        return "error", f"error: {error}"

def run(task: str) -> str:
    run_id = time.strftime("%Y%m%d-%H%M%S")
    log({"run_id": run_id, "event": "task", "task": task})
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": task},
    ]
    total_tokens = 0
    total_latency_ms = 0

    for step in range(MAX_STEPS):
        result = provider.complete(messages, tools=TOOLS)
        message = result.message
        total_tokens += result.prompt_tokens + result.completion_tokens
        total_latency_ms += result.latency_ms
        log(
            {
                "run_id": run_id,
                "event": "usage",
                "step": step + 1,
                "prompt_tokens": result.prompt_tokens,
                "completion_tokens": result.completion_tokens,
                "latency_ms": result.latency_ms,
            }
        )

        if not message.tool_calls:
            log({"run_id": run_id, "event": "answer", "content": message.content})
            print(f"\n{result.model}: {total_tokens} tokens, {total_latency_ms}ms total")
            return message.content

        messages.append(
            {
                "role": "assistant",
                "content": message.content,
                "tool_calls": [
                    {
                        "id": call.id,
                        "type": "function",
                        "function": {
                            "name": call.function.name,
                            "arguments": call.function.arguments,
                        },
                    }
                    for call in message.tool_calls
                ],
            }
        )

        for call in message.tool_calls:
            name = call.function.name
            try:
                args = json.loads(call.function.arguments)
                outcome, result_text = execute(name, args)
            except Exception as error:
                args, outcome, result_text = call.function.arguments, "error", f"error: {error}"
            print(f"[{step + 1}] {name} {call.function.arguments} -> {outcome}")
            log(
                {
                    "run_id": run_id,
                    "event": "tool_call",
                    "step": step + 1,
                    "tool": name,
                    "args": args,
                    "outcome": outcome,
                    "result": result_text[:500],
                }
            )
            messages.append(
                {"role": "tool", "tool_call_id": call.id, "content": result_text}
            )

    log({"run_id": run_id, "event": "stopped", "reason": "step limit"})
    return "stopped: hit the step limit"

if __name__ == "__main__":
    print(run(sys.argv[1]))