agent.py

from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Protocol


@dataclass(frozen=True)
class Say:
    text: str


@dataclass(frozen=True)
class Call:
    tool: str
    arg: str


Reply = Say | Call


class Model(Protocol):
    def respond(self, system: str, history: list[str]) -> Reply: ...


Tool = Callable[[str], str]


@dataclass
class Agent:
    model: Model                                          # 1. Model
    system: str                                           # 2. Instructions
    history: list[str] = field(default_factory=list)      # 3. Memory
    tools: dict[str, Tool] = field(default_factory=dict)  # 4. Tools

    def run(self, user_input: str) -> str:
        self.history.append(f"user: {user_input}")
        while True:  # real agents cap the iterations
            match self.model.respond(self.system, self.history):
                case Say(text):
                    self.history.append(f"agent: {text}")
                    return text
                case Call(tool, arg):
                    fn = self.tools.get(tool)
                    result = fn(arg) if fn else f"no such tool: {tool}"
                    self.history.append(f"tool[{tool}]: {result}")
                    # loop again: the model sees the result and decides what's next


def read_file(path: str) -> str:
    try:
        return f"{len(Path(path).read_text())} bytes"
    except OSError as e:
        return f"error: {e}"


class FakeModel:
    def respond(self, system: str, history: list[str]) -> Reply:
        last = history[-1] if history else ""
        if last.startswith("tool["):
            return Say(f"Done: {last}")
        if last.startswith("user: read "):
            return Call("read_file", last.removeprefix("user: read ").strip())
        return Say("I can read files. Try: read <path>")


def main() -> None:
    agent = Agent(
        model=FakeModel(),
        system="You can read files.",
        tools={"read_file": read_file},
    )
    while True:
        try:
            line = input("> ")
        except EOFError:
            break
        print(agent.run(line.strip()))


if __name__ == "__main__":
    main()