diff --git a/README.md b/README.md index 350da526..ad203766 100644 --- a/README.md +++ b/README.md @@ -474,6 +474,36 @@ def postmortem(incident: str) -> str: A method can also run on a live thread whose calls share history, have its proposals checked by a gate that runs as its post-condition, or read memory that the optimizer can update. See the [methods documentation](docs/methods.md) for details. +## Teams + +The `ai_functions.experimental.teams` module puts a lead AI Function over a few live members. The lead holds every member as a typed tool and decides whom to ask, and each member keeps its own conversation across the lead's asks. Hooks add everything else, such as a review of the answer (see `examples/teams/status_desk.py`): + +```python +from ai_functions.experimental.teams import Accept, FunctionMember, MethodMember, Revise, Team +from ai_functions.runtime import Budget, Cost + + +class HouseStyle: + def on_answer(self, work, answer): + if "Resolved at" in answer: + return Accept() + return Revise(feedback='End with the restore time as "Resolved at HH:MM UTC".', cap=2) + + +team = Team( + status_update, + members=[ + MethodMember(analyst, "timeline", ask_estimate=Cost(dollars=0.01)), + FunctionMember(check_claim, name="skeptic", ask_estimate=Cost(dollars=0.01)), + ], + hooks=[HouseStyle()], +) +run = await team.run("Payment errors on checkout this morning; the issue is fixed.", budget=Budget(Cost(dollars=0.50))) +print(run.answer) +``` + +A `Revise` sends the lead back to work with the feedback, and a revised answer goes back through every reviewer before it ships. The run's transcript records every member call and revision as a typed record, and `budget=` bounds what the run may spend in tokens, dollars or both, with each member reserving its `ask_estimate` before it's asked. See the [teams documentation](docs/teams.md) for details. + ## Security Code execution is off by default. The `"local"` mode validates generated code with AST checks, restricts imports to the allowlist you pass in, and applies timeouts. But it is not a sandbox: it cannot stop resource exhaustion (an infinite loop, runaway memory allocation) and offers no process-level isolation. For production, run AI Functions inside a container or other isolated environment, which adds the process isolation and resource limits that `"local"` mode cannot provide. For untrusted input, use `"disabled"` mode. diff --git a/docs/teams.md b/docs/teams.md new file mode 100644 index 00000000..8b5d8382 --- /dev/null +++ b/docs/teams.md @@ -0,0 +1,211 @@ +# Teams + +An AI Function answers one request on one thread. The `ai_functions.experimental.teams` module puts one lead AI Function over a few live members. The lead holds every member as a typed tool, and it decides whom to ask and what to ask them. Hooks add everything else: extra tools, a rewritten request, a review of the answer, a check before every ask. + +This enables: + +- **A lead that delegates through typed tools.** Each member is a tool named after it. A member keeps its own conversation across the lead's asks, so a follow-up question can refer to an earlier answer. +- **Members from functions and methods.** `FunctionMember` wraps an `AIFunction` and `MethodMember` wraps one `@ai_method` of a `MethodAgent`. Any object with a name and `spawn`, `ask` and `retire` is a member too. +- **Reviewed answers.** A hook's `on_answer` returns `Accept()` or `Revise(feedback=...)`. A revised answer goes back through every reviewer before it ships, and every reviewer's budget of revisions is bounded. +- **A transcript the team writes.** Every member call and every revision is a typed record, written by the code that carried it rather than reported by the model. +- **Budgets in tokens and dollars.** `Team.run(budget=...)` charges the run everything it spends, prices tokens by the model that spent them, and reserves each member ask before it runs. + +Note: the module is experimental, and its API may change in future releases. + +## Contents + +- [A lead, members and a run](#a-lead-members-and-a-run) +- [Members](#members) +- [Hooks](#hooks) +- [Reviewing the answer: `Accept` and `Revise`](#reviewing-the-answer-accept-and-revise) +- [The transcript](#the-transcript) +- [Budgets and deadlines](#budgets-and-deadlines) +- [Cancelling a run](#cancelling-a-run) +- [Errors](#errors) +- [Examples](#examples) + +## A lead, members and a run + +```python +from ai_functions import ai_function +from ai_functions.experimental.teams import FunctionMember, Team + + +@ai_function +def research(topic: str) -> str: + """Report what you know about {topic}, in three sentences.""" + + +@ai_function(system_prompt="Answer the question. Ask the scout before you decide.") +def answer(question: str) -> str: + """{question}""" + + +team = Team(answer, [FunctionMember(research, name="scout")]) +run = await team.run("Why are there two tides a day?") +print(run.answer) +``` + +The lead is an ordinary AI Function. `Team(lead, members, hooks=())` builds the team, and `await team.run(request)` does the rest. It spawns the lead's thread, spawns each member as a child of it, runs the lead with every member as a tool, walks the hooks over the answer, and retires everybody in a `finally`, even when something raised. The request lands in the lead's first positional parameter. + +With no `coordinator`, a run stands up a private `InMemoryCoordinator` and `LocalWorker` and closes them when it ends, which is what a script wants. Pass a coordinator of your own to put the run's threads in a larger runtime's registry and event log, and `parent_id=` to make the lead a child of a thread you already hold. The whole run is then one subtree under that parent. + +`run` returns a `TeamRun`: + +| Field | What it holds | +|--------------------------------|------------------------------------------------------------------------------------------------------| +| `answer` | The lead's final answer, in the lead's output type | +| `transcript` | The typed records of the run, in order (see [The transcript](#the-transcript)) | +| `hooks_data` | What hooks put in `Workspace.data` | +| `blocked`, `blocked_by` | Whether a mandatory reviewer ran out of revisions without approving, and which ones | +| `cancelled`, `cancel_reason` | Whether the run was stopped while live, and the first reason given | +| `cost` | The run's measured spend as a `Cost`: its tokens, and the dollars of the usage the price table priced | +| `run_id`, `lead_id` | The run's identity and the lead's thread id | + +Empty fields drop out of `model_dump()`, so a bare run dumps its answer, its identity and its cost. + +A `Team` instance runs one run at a time. Members hold live handles and hooks keep per-run state, so a second `run` while the first is in flight raises `WiringError` before anything spawns. Build two `Team` instances to run in parallel. Sequential runs on one instance are fine. + +## Members + +A member is something with a name that spawns, answers one request at a time, and retires. Two adapters cover the common shapes. + +`FunctionMember(function, *, name=None, parameter=None, ask_estimate=None, **overrides)` wraps an `AIFunction`. `name` is the member's name on the team and its tool name on the lead's side, and defaults to `function.name`. Pass one when two members share a function. Tool names allow only letters, digits, `_` and `-`, so any other character is mapped to `_` on the tool, while the transcript keeps the real name. `parameter` names the keyword the request arrives as. It defaults to the function's first positional parameter, and the function's other parameters keep their defaults. `overrides` are `ThreadConfig` overrides for the member's thread, such as `model=`. + +`MethodMember(agent, method, *, name=None, parameter=None, ask_estimate=None, **overrides)` puts one `@ai_method` of a `MethodAgent` on a live thread. Its name defaults to the method's tool name, `{owner}_{method}`, so the lead's tool, the coordinator's thread name and every error agree. One agent can contribute several members, one per method. + +```python +from ai_functions.experimental.teams import MethodMember + +team = Team(postmortem, [MethodMember(analyst, "timeline"), MethodMember(analyst, "causes")]) +``` + +The lead sees each member as a tool that takes one `request: str`. The tool's description tells the lead that the member doesn't see its conversation, so the request has to say everything the member needs. A member that exposes a `description` attribute (the `Described` protocol) gets that text appended to its tool's description, so the lead can choose between members by what they do. + +A member's answer reaches the lead the way any Strands tool result does. Text passes through as is, and anything else is formatted as its JSON. The transcript records the same text. + +### Writing your own member + +Any object that satisfies the `Member` protocol joins a team: + +```python +class Member(Protocol): + name: str + + async def spawn(self, coordinator: Coordinator, *, parent_id: ThreadId | None = None) -> Spawned: ... + async def ask(self, request: str) -> Any: ... + async def retire(self) -> None: ... +``` + +`spawn` returns anything with an `id`; the core reads nothing else from it. `retire` must be idempotent, because the core calls it on every path out of a run. A member that also exposes `equip(config_hook)` (the `Equippable` protocol) receives the tools hooks hand to members. A member without it is skipped rather than refused, and the transcript records an `unequipped` note naming the member and the hooks whose tools it goes without. + +`current_call()` returns the `CallContext` of the ask in flight on the current task, with the run id, the member's name and the ask's ordinal. A member that does work of its own during an ask can name that work after the ask that caused it. A `Team.run` started inside a member's ask picks the context up on its own and re-runs the outer team's `before_ask` checks before each of its own asks, so stopping the outer run also stops the inner one. + +### The team is the only connection between its threads + +A thread on a coordinator normally gets the runtime's `list_threads` and `send_message` tools. A team turns both off on every thread it spawns: the lead, and every member that `FunctionMember` or `MethodMember` spawns. Both adapters refuse `coordinator_tools_enabled=True`. So the lead reaches a member only through the member's tool, and every ask it makes is recorded on the transcript and runs under the run's checks, deadline and budget. Members share what they know only through tools a hook hands them. A member you write yourself owns its spawn and should turn those tools off the same way. + +## Hooks + +A hook is any object that implements one or more of these methods. Every method is optional, and the core skips the ones a hook doesn't have. + +| Method | Runs | +|---------------------------------------------|-----------------------------------------------------------------------------------------| +| `on_assemble(work)` | Once the members are live, before the lead runs | +| `on_request(work, request) -> str` | Before the lead's first cycle; returns the request the lead sees, folded in hook order | +| `tools_for_lead(work, ctx) -> tools` | On every lead cycle; returns extra tools for that cycle | +| `tools_for_member(work, member, ctx)` | On every cycle of each equippable member; returns extra tools for that cycle | +| `before_ask(work, member)` | Before every member ask, the lead's and a hook's; raise `StopRun` to stop the run | +| `on_answer(work, answer) -> Accept/Revise` | On each answer the lead gives (see below) | +| `on_teardown(work)` | When the run is over, successfully or not, while the members are still live | +| `on_closed(work)` | When everything is retired; `work.usage` is the final total and `work.fault` is set | + +`tools_for_lead` and `tools_for_member` are synchronous, because the runtime builds a cycle's tools synchronously. The other methods may be sync or async, and the core awaits whichever it finds. `Team(hooks=...)` is typed `PartialTeamHook`, any object with at least one of these methods, so a type checker flags an object that has none of them. + +Every method receives the run's `Workspace`. It carries the request, the live `members`, the lead's handle, the `transcript`, a `data` dictionary that the run publishes as `TeamRun.hooks_data`, and the run's `run_id`. A workspace is one run and compares by identity, so a hook instance that serves many runs keeps per-run state keyed by the workspace. A hook that asks a member a question of its own calls `await work.ask(member, request)`, which runs under the same `before_ask` checks and cancellation as the lead's asks and returns the typed answer. The ask's tokens count toward the run's cost, but it reserves nothing against a budget, has no deadline, and leaves the transcript to the hook. + +A model-fixable failure in a hook's tool should come back as `"error: ..."` text, so the model can recover. A failure in persistence should raise. + +## Reviewing the answer: `Accept` and `Revise` + +```python +import re + +from ai_functions.experimental.teams import Accept, Revise, Workspace + + +class HouseStyle: + """Send the update back until it states the restore time the way the status page prints it.""" + + def on_answer(self, work: Workspace, answer: str) -> Accept | Revise: + if re.search(r"Resolved at \d{2}:\d{2} UTC", answer): + return Accept() + return Revise(feedback='End with the restore time as "Resolved at HH:MM UTC".', cap=2) +``` + +`on_answer` returns exactly one of two verdicts. `Accept()` passes the answer on. `Revise(feedback=..., cap=2, mandatory=False)` re-runs the lead with the feedback. `cap` is how many revisions this hook may ask for in the run, read from its latest verdict. Anything else returned from `on_answer` raises `WiringError` naming the hook, so a hook that forgets to return can't pass for an approval. + +Reviewers run in hook order. After any `Revise`, the walk starts again at the first reviewer, so a revised answer goes back through every reviewer that had already accepted the earlier draft. An answer ships only after one full pass in which every reviewer accepts it or has run out of revisions. Each restart spends one revision from some hook's cap, so the loop is bounded by the sum of the caps. A hook that runs out of revisions lets the answer through, and the transcript records a `revise_cap` entry for it once. + +`mandatory=True` turns that pass-through into a blocked result. When the final pass includes a mandatory reviewer that ran out of revisions without approving, `run` returns `TeamRun(blocked=True, blocked_by=(...))` with the last answer still attached. A mandatory reviewer that ran out on one draft and then accepted the next one approved what ships, so it blocks nothing. Hooks are labeled by their class name, with a `#n` suffix when two hooks share a class. + +A revision runs the lead again on its live thread, with `Your answer was reviewed and needs revision.` and the feedback as its first argument. The lead keeps its history, so it can revise its own draft with what its members already told it. Its docstring is rendered again with the feedback in place of the request, though, so keep standing instructions (whom to ask, what form to answer in) in the lead's `system_prompt` and keep the docstring to the request. Otherwise the lead reads its instructions again and may ask its members the same questions a second time. + +## The transcript + +`TeamRun.transcript` lists what happened in the run as frozen typed records: + +| Record | `kind` | Written when | Fields | +|---------------|---------------|-------------------------------------------------|---------------------------------------------------------| +| `MemberCall` | `member` | The lead asked a member a question | `member`, `request`, `answer`, `call_id` | +| `Revision` | `revise` | A hook sent the answer back | `hook`, `round`, `feedback` | +| `RevisionCap` | `revise_cap` | A hook ran out of revisions | `hook`, `rounds` | +| `Note` | a hook's own | A hook appended something it wants remembered | `kind`, `data` | + +The core writes the first three, and hooks append `Note` entries. A `Note` can't use the core's three kinds (`CORE_TRANSCRIPT_KINDS`), so a reader filtering on `kind` never mistakes a hook's entry for the core's. `MemberCall.call_id` is `run_id:member:n`, where `n` counts that member's asks from one. `MemberCall.answer` is the text the lead read. When the ask raised, it's `error:` followed by the exception, and the lead sees a failed tool call. + +## Budgets and deadlines + +```python +from ai_functions.runtime import Budget, Cost + +team = Team( + status_update, + [ + MethodMember(analyst, "timeline", ask_estimate=Cost(dollars=0.01)), + FunctionMember(check_claim, name="skeptic", ask_estimate=Cost(dollars=0.01)), + ], +) +run = await team.run("Payment errors on checkout this morning.", budget=Budget(Cost(dollars=0.50))) +print(run.cost) +``` + +`ai_functions.runtime.Budget` is a cost bound in tokens, dollars or both, with reservations, and `Budget.child()` nests one budget under another. Under `Team.run(budget=...)` the run holds one charge on the budget for its measured spend: the lead's tokens, every member's, and any hook's asks. The charge is refreshed around every ask and closed at the same figure `TeamRun.cost` reports. Each ask the lead makes also reserves the member's estimate before it runs, so a budget the run has already spent refuses the next ask. A hook's own asks are charged with the run but reserve nothing. + +A member's estimate is its `ask_estimate` until it has measured an ask of its own, then the costliest of its recent measured asks. A budget's `on_unknown` policy decides what happens to a member with no estimate. The default, `"refuse"`, rejects such a member with a `WiringError` before anything spawns. `"estimate"` holds the budget's `default_estimate` instead. A refused reservation asks the member nothing and reaches the lead as `error:` text, so the lead can finish with what it has. + +Dollars come from the coordinator's price table, `ai_functions.runtime.PRICES` by default, which holds list prices for current Claude models and applies the Bedrock regional premium to geographic inference profiles. Each token usage event names the model that spent it, and the table prices it by that id. To price other models or negotiated rates, pass `InMemoryCoordinator(prices=PRICES.with_prices({...}))` and run the team on that coordinator. Usage the table can't price adds no dollars, so `TeamRun.cost.dollars` sums only the priced usage, and it's `None` when nothing was priced. That usage isn't treated as free, though: once a run has spent any, every later ask under a budget that bounds dollars is refused, with an error naming the model. + +`Team(ask_timeout_s=...)` sets a deadline, in seconds, on every member ask. A member that carries an `ask_timeout_s` attribute of its own is held to that instead. Past the deadline the member is retired and the lead reads `error:` text that says so. A retired member isn't asked again in that run: a later call to its tool asks nothing and returns `error:` text. `None`, the default, bounds nothing. + +## Cancelling a run + +`team.cancel(reason)` stops the instance's live run from the same event loop and returns whether a run was live. The lead is interrupted, every member is retired, the teardown hooks run, and `run` returns `TeamRun(cancelled=True, cancel_reason=reason)`. The first reason wins. A `before_ask` hook that raises `StopRun(reason)` stops the run the same way, before the ask it was checking, and the lead's model isn't called again. Cancelling the caller's own task still propagates as usual. + +## Errors + +| Error | Raised when | +|-----------------------------------|-----------------------------------------------------------------------------------------------------------------------------| +| `WiringError` | two members share a tool name, a hook returns something other than a verdict, a run is already in flight on the instance, a member can't take hook tools, a member has no estimate under a refusing budget, or a member adapter can't be wired | +| `StopRun` | a `before_ask` hook raises it to stop the run; `run` returns a cancelled `TeamRun` rather than propagating it | +| `BudgetExhausted`, `UnknownCost` | a reservation is refused; inside a team run they reach the lead as `error:` text | + +`WiringError` subclasses `TeamError`, an `AIFunctionError` whose `function_name` names the lead or member that raised it. `StopRun` subclasses `asyncio.CancelledError`. `BudgetExhausted` and `UnknownCost` subclass `ai_functions.runtime.BudgetError`. + +## Examples + +| Example | Shows | +|-------------------------|-------------------------------------------------------------------------------------------------------------------------| +| `teams/status_desk.py` | A lead over a method member and a function member, a hook that sends the first draft back, and a run under a dollar budget | + +Run it from the `examples/` folder with `uv run teams/status_desk.py`. The [examples guide](../examples/README.md#teams) lists all commands and setup requirements. diff --git a/examples/README.md b/examples/README.md index 1bada499..6fefe4db 100644 --- a/examples/README.md +++ b/examples/README.md @@ -61,6 +61,7 @@ Lean automatically; see [Verified](#verified) for native compilation setup. | Remember information and improve results from feedback | [Memory](#memory) | | Choose models according to cost and expected success | [Economics](#economics) | | Give an object AI methods and hand them to other agents | [Methods](#methods) | +| Lead a team of agents, review its answer, and bound its cost | [Teams](#teams) | | Use Claude Code or Kiro as a thread | [Integrations](#integrations) | | Verify the agent's work or compile a verified implementation | [Verified](#verified) | | Explore a larger learning loop or benchmark | [Projects](#projects) | @@ -167,6 +168,16 @@ with a gate. See [the methods guide](../docs/methods.md). | [Analyst team](methods/analyst_team.py) | Call a method once, hand every method to an orchestrator as a typed tool, and keep one method on a live thread. | `uv run methods/analyst_team.py` | | [Gated proposer](methods/gated_proposer.py) | Check each proposal against a policy cap, retry with the gate's report, and filter one-shot proposals from forked threads. | `uv run methods/gated_proposer.py` | +## Teams + +Put a lead AI function over a few live members it calls as typed tools, review +the lead's answer with a hook, and run the team under a budget. See +[the teams guide](../docs/teams.md). + +| Example | What it teaches | Run | +| --- | --- | --- | +| [Status desk](teams/status_desk.py) | Hold a method member and a function member as the lead's tools, send the first draft back with a hook, and read the transcript and the run's cost under a dollar budget. | `uv run teams/status_desk.py` | + ## Integrations These examples run external agent sessions as threads. Run them from diff --git a/examples/teams/status_desk.py b/examples/teams/status_desk.py new file mode 100644 index 00000000..6c489851 --- /dev/null +++ b/examples/teams/status_desk.py @@ -0,0 +1,151 @@ +"""Teams: a lead over live members, with a hook that reviews the answer. + +An incident desk writes the customer status update for a payment outage. The +lead is an ordinary ``@ai_function``. It holds two members as typed tools: an +analyst's ``timeline`` method on a live thread (``MethodMember``) and a +skeptic function that checks one claim against the notes (``FunctionMember``). +Each member keeps its own conversation across the lead's questions. + +``HouseStyle`` is a hook. Its ``on_answer`` reads the lead's draft and returns +``Accept()`` or ``Revise(feedback=...)``, which sends the lead back to work with +the feedback. The team runs under a budget in dollars, and the result carries +the answer, a transcript of every member call and revision, and what the run +cost. +""" + +import asyncio +import re + +from example_helpers import models +from example_helpers.utils import display, rule +from pydantic import BaseModel, Field + +from ai_functions import ai_function +from ai_functions.experimental.methods import MethodAgent, ai_method +from ai_functions.experimental.teams import ( + Accept, + FunctionMember, + MemberCall, + MethodMember, + Revise, + Revision, + RevisionCap, + Team, + Workspace, +) +from ai_functions.runtime import Budget, Cost + +NOTES = """\ +09:02 UTC deploy 4172 rolls out to checkout: new payment client, retry limit changed from 5 to 0 +09:05 UTC p99 latency on /pay climbs from 300 ms to 2.1 s +09:07 UTC error rate on /pay reaches 4%; the card processor reports no incidents on its side +09:12 UTC on-call rolls back deploy 4172 +09:15 UTC latency and error rate on /pay are back to baseline +""" + + +class Finding(BaseModel): + answer: str = Field(description="The answer, in one or two sentences.") + evidence: list[str] = Field(description="The note lines the answer rests on, quoted exactly.") + + +class Analyst(MethodAgent): + """The analyst for one service. Its notes live on ``self`` and reach every prompt.""" + + model = models.small + + def __init__(self, service: str, notes: str) -> None: + self.service = service + self.notes = notes + + @ai_method(Finding, description="Answer a question about when things happened") + def timeline(self, question: str) -> Finding: + """You're the analyst for the {self.service} service. Answer from these notes only: + + {self.notes} + Question: {question} + """ + + +# The body returns the prompt, so it can put the module's notes in front of the claim. +@ai_function[str](model=models.small) +def check_claim(claim: str) -> str: + return ( + 'Say whether these incident notes support this claim. Reply "supported" or "unsupported", ' + f"then quote the note line that decides it.\n\nNotes:\n{NOTES}\nClaim: {claim}" + ) + + +DESK = """\ +You write customer status updates for incidents, in at most three sentences. Ask the analyst for +the facts before you write, and have the skeptic check your root-cause claim. Customers read the +update, so leave out deploy numbers and internal names.""" + + +# The team's request lands in the lead's first parameter. A reviewer's feedback lands there too, +# on the same live thread, so the standing instructions live in the system prompt and the +# docstring is only the request: a revision then reads as a follow-up to the lead's own draft. +@ai_function(model=models.medium, system_prompt=DESK) +def status_update(incident: str) -> str: + """{incident}""" + + +class HouseStyle: + """Send the update back until it states the restore time the way the status page prints it.""" + + RESOLVED = re.compile(r"Resolved at \d{2}:\d{2} UTC") + + def on_answer(self, work: Workspace, answer: str) -> Accept | Revise: + if self.RESOLVED.search(str(answer)): + return Accept() + return Revise( + feedback='End the update with the restore time in the form "Resolved at HH:MM UTC".', + cap=2, + ) + + +def _short(text: str, limit: int = 110) -> str: + flat = " ".join(text.split()) + return flat if len(flat) <= limit else flat[: limit - 3] + "..." + + +async def main() -> None: + analyst = Analyst("checkout", NOTES) + team = Team( + status_update, + members=[ + # The lead sees this member as the `analyst_timeline` tool. Under a budget, each ask + # reserves the member's estimate until the member has measured an ask of its own. + MethodMember(analyst, "timeline", ask_estimate=Cost(dollars=0.01)), + FunctionMember(check_claim, name="skeptic", ask_estimate=Cost(dollars=0.01)), + ], + hooks=[HouseStyle()], + ) + + rule("One team run under a $0.50 budget") + run = await team.run( + "Payment errors on checkout this morning; the issue is fixed.", + budget=Budget(Cost(dollars=0.50)), + ) + display("Status update", run.answer) + + rule("The transcript: written by the team, not reported by the model") + for entry in run.transcript: + if isinstance(entry, MemberCall): + print(f"member {entry.call_id}") + print(f" asked: {_short(entry.request)}") + print(f" answered: {_short(entry.answer)}") + elif isinstance(entry, Revision): + print(f"revise {entry.hook} round {entry.round}: {entry.feedback}") + elif isinstance(entry, RevisionCap): + print(f"revise_cap {entry.hook} after {entry.rounds} rounds; the last answer passed through") + else: + print(f"{entry.kind} {entry.data}") + + rule("What the run cost") + dollars = "unpriced" if run.cost.dollars is None else f"${run.cost.dollars:.4f}" + print(f"{run.cost.tokens} tokens, {dollars}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/spec/ai_functions/experimental/teams/__init__.pyi b/spec/ai_functions/experimental/teams/__init__.pyi new file mode 100644 index 00000000..fe6b0b4d --- /dev/null +++ b/spec/ai_functions/experimental/teams/__init__.pyi @@ -0,0 +1,135 @@ +"""Teams: a lead over live members, with hooks for everything that is not that. + +The core (:mod:`.core`) owns exactly what every team needs — spawn the +members, run the lead with the members as typed tools, review the answer +through a bounded loop, retire everybody — and every other capability arrives +as a hook. The contracts layer (:mod:`.contracts`, :mod:`.types`) is the +vocabulary those parties exchange and nothing else; :class:`Team` and its +:class:`TeamRun` report are the composition built on it. +:class:`FunctionMember` is the ``AIFunction``-backed member shape and +:class:`MethodMember` the ``MethodAgent``-backed one. + +Four protocols and one state plane. :class:`Member` is what a member +promises: a name, and the three verbs the core performs on it — ``spawn``, +``ask``, ``retire``. :class:`TeamHook` is what a hook may implement; every +method is optional and absence means skip, so a team's ``hooks`` are typed +:data:`PartialTeamHook`, any object with at least one of those methods. +:class:`Equippable` is the optional seam a member exposes so hook-contributed +tools can ride into its cycles, and :class:`Described` the optional seam a +member's ``description`` reaches the lead through, as part of the tool the +lead calls it by. :class:`Workspace` is one run's shared state, handed to +every hook method as one argument so a new field breaks no existing hook. + +The answer loop speaks in two verdicts. :class:`Accept` passes the answer on; +:class:`Revise` re-runs the lead with feedback under a cap the hook itself +chooses. What the loop and the members did is written to the transcript as +frozen records — a :class:`MemberCall` per question the lead asked a member, +a :class:`Revision` per re-run, a :class:`RevisionCap` when a hook's budget ran +out, a :class:`Note` for anything a hook wants remembered — each recorded by +the wire that carried the event, never reported by the model. + +The protocols compose with this repository's runtime: ``spawn`` takes an +:class:`~ai_functions.protocols.Coordinator` and a +:class:`~ai_functions.types.ThreadId` parent, and the lead's handle on the +workspace is a :class:`~ai_functions.handle.ThreadHandle`. A run can be +bounded by a :class:`~ai_functions.runtime.budget.Budget` in tokens, dollars +or both. + +Invariants: + T1 — a member is a name and three verbs. ``spawn`` puts it on a live + thread as a child of the given parent and returns something with an + ``id``; ``ask`` runs one cycle and returns its result; ``retire`` tears the + thread down and is idempotent, because an unwind loop must not crash + mid-unwind. + + T2 — a verdict is exactly one of ``Accept`` or ``Revise``. ``Accept`` + carries nothing and refuses extra fields; ``Revise`` carries feedback and + a cap. Anything else returned from ``on_answer`` is a wiring bug for the + core to raise on, never a silent pass. + + T3 — ``Revise.cap`` is the hook's own budget for this finding, read off + the latest verdict. A negative cap is a construction error because it + would silently behave as zero. Exhaustion is not an error: the answer + passes on and the transcript records that the cap, not the hook, ended + the loop. ``Revise.mandatory`` (default off) makes that exhaustion a + blocked ``TeamRun`` instead of ordinary success (core C8). + + T4 — a ``Workspace`` is one run and has identity, not value equality: two + runs never share one, two workspaces built from the same inputs are not + equal, and per-run state may key by the workspace itself or its + ``run_id``. ``Workspace.calls`` counts each member's asks, so + ``MemberCall.call_id`` is ``run_id:member:n``. + + T5 — transcript entries are frozen records with a ``kind``. The core owns + ``member``, ``revise``, and ``revise_cap``; a hook's ``Note`` may not reuse + them, so a reader filtering on kind never mistakes a hook's entry for the + core's. + + T6 — every ``TeamHook`` method is optional, and presence is checked with + ``getattr``, never ``isinstance``. ``tools_for_lead`` and + ``tools_for_member`` are synchronous; the lifecycle methods may be sync or + async and the core awaits whichever it finds. + + T7 — the contracts and core modules import nothing from any other + experimental package; :mod:`.members` reaches sideways only into + ``experimental/methods``. Methods never imports teams. + + T8 — a member's description reaches the lead through one seam. + ``Described`` is the seam (``description: str | None``); the core reads it + by ``isinstance`` and appends the text to the member tool's generic + description, and a ``None`` or blank description, or a member without the + seam, leaves the generic text byte for byte. + + T9 — the ask in flight is readable. The core sets a :class:`CallContext` + around every member ask and :func:`current_call` reads it back on the + same task, so a member that opens work of its own names that work after + the ask that caused it without a new argument on ``Member``. +""" + +from __future__ import annotations + +from .contracts import Described, Equippable, Member, PartialTeamHook, Spawned, TeamHook, Workspace +from .core import CallContext, StopRun, Team, TeamRun, WiringError, current_call +from .members import FunctionMember, MethodMember +from .types import ( + CORE_TRANSCRIPT_KINDS, + Accept, + MemberCall, + Note, + Revise, + Revision, + RevisionCap, + RunId, + TeamError, + TranscriptEntry, + Verdict, +) + +__all__ = [ + "Accept", + "CORE_TRANSCRIPT_KINDS", + "CallContext", + "Described", + "Equippable", + "FunctionMember", + "Member", + "MemberCall", + "MethodMember", + "Note", + "PartialTeamHook", + "Revise", + "Revision", + "RevisionCap", + "RunId", + "Spawned", + "StopRun", + "Team", + "TeamError", + "TeamHook", + "TeamRun", + "TranscriptEntry", + "Verdict", + "WiringError", + "Workspace", + "current_call", +] diff --git a/spec/ai_functions/experimental/teams/contracts.pyi b/spec/ai_functions/experimental/teams/contracts.pyi new file mode 100644 index 00000000..9488a76f --- /dev/null +++ b/spec/ai_functions/experimental/teams/contracts.pyi @@ -0,0 +1,337 @@ +"""The seams of a team: what a member promises, what a hook may implement, what a run shares. + +Deliberately protocols rather than base classes: the members worth having +already exist and already differ, and a hook implements the two methods it +needs and nothing else. The core (:mod:`.core`) says what a team *does* with a +member and a hook; this module says what shapes may *be* one. Only +:class:`Workspace` is concrete, because it is state, not behavior — the one +argument every hook method receives. + +Everything here binds to this repository's runtime. A :class:`Member` spawns +onto an :class:`~ai_functions.protocols.Coordinator` under a +:class:`~ai_functions.types.ThreadId` parent and hands back a +:class:`Spawned`; the workspace's ``lead`` is a +:class:`~ai_functions.handle.ThreadHandle`; hook tools are rebuilt per cycle +against a :class:`~ai_functions.types.ThreadContext`. + +Invariants: + T1 — ``Member`` is a name and three verbs; ``retire`` is idempotent. + + T4 — ``Workspace`` has identity, not value equality, and a fresh + ``run_id`` per construction. ``Workspace.calls`` counts each member's + asks so the core can mint deterministic call ids. + + T6 — every ``TeamHook`` method is optional; the protocol is not runtime + checkable on purpose, so ``isinstance`` cannot be used where ``getattr`` + is meant. + + T8 (the seam half) — ``Described`` is the one seam the lead's member-tool + text reads: a member that says what it is for exposes ``description``, and + the core appends it to the generic tool text; a member without the seam, or + with ``None`` or blank, keeps the generic text byte for byte. The core reads + the seam by ``isinstance``, never a bare attribute name. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +from ai_functions.types import TokenUsage + +from .core import CallContext +from .types import RunId, TranscriptEntry, Verdict + +if TYPE_CHECKING: + from strands.types.tools import AgentTool + + from ...ai_thread.config import ThreadKwargs + from ...handle import ThreadHandle + from ...protocols import Coordinator + from ...types import ThreadContext, ThreadId + +# ── The member shape ── + +@runtime_checkable +class Spawned(Protocol): + """Synchronous methods and properties; calls finish before returning. + + What ``Member.spawn`` hands back: something the core can name by ``id``. + + The core reads only ``.id`` off a spawned member — to make it the parent + of nothing and the child of the lead — so this is the whole of the + promise. A :class:`~ai_functions.handle.ThreadHandle` satisfies it as + written; so does a synthetic handle for a member with no single thread of its own. + """ + + @property + def id(self) -> ThreadId: + """The thread this member now lives on, or a synthetic id when no single thread does.""" + ... + +@runtime_checkable +class Member(Protocol): + """Async methods; await each operation. Sync providers use ``as_async``. + + A member: something with a name that spawns, answers once, and retires. + + Three verbs and nothing else, because that is the whole of what the core + does to a member — it stands one up, asks it one question, and takes it + down in the ``finally``. Anything richer would be a contract the library + cannot honor for every member shape it wants to accept (T1). + + Deliberately a protocol rather than a base class: the members worth + having already exist and already differ. A typed method on an agent + adapts by holding the thread its own ``spawn`` returns and routing ``ask`` + to one parameter; a member of your own can run a whole team per ask. Each + joins the lead's wire the same way — one ``request: str``, one answer. + + ``spawn`` returns a :class:`Spawned` rather than a concrete handle: the + core reads only ``.id`` off it, and demanding a + :class:`~ai_functions.handle.ThreadHandle` would exclude a member shape + that has no single thread. + + The thread a member spawns should carry none of the runtime's + ``list_threads`` / ``send_message``, because a member reachable around the + lead's typed tool is work the team never sees (C12): the adapters in + :mod:`~ai_functions.experimental.teams.members` spawn with them off (TB1). + """ + + name: str + """Unique within a team; each member becomes a lead tool named after it.""" + + async def spawn(self, coordinator: Coordinator, *, parent_id: ThreadId | None = None) -> Spawned: + """Put this member on a live thread as a child of ``parent_id``. + + Args: + coordinator: The coordinator every thread of this run spawns onto. + parent_id: The lead's thread, so the run is one subtree in one event log. + + Returns: + Something with an ``id``; the core reads nothing else off it. + """ + ... + + async def ask(self, request: str) -> Any: + """Run one cycle with ``request`` and return its result. + + Args: + request: The whole of what the lead tells the member; it never sees the lead's conversation. + + Returns: + Whatever the member produced; the lead reads text as is and anything else as its JSON (C6). + """ + ... + + async def retire(self) -> None: + """Tear the thread down. Idempotent — an unwind loop must not crash mid-unwind.""" + ... + +@runtime_checkable +class Equippable(Protocol): + """Synchronous methods and properties; calls finish before returning. + + The optional seam hook-contributed member tools ride in on. + + The core equips each member that exposes this with one composed + ``config_hook`` between construction and spawn, without the caller + writing any wiring. Members without the seam are skipped, not refused: + a mixed cast is legitimate, and a member that cannot take tools is a fact + rather than a fault. The fact is recorded: the run's transcript names each + skipped member in an ``unequipped`` note (C2). + """ + + def equip(self, config_hook: Callable[[ThreadContext], ThreadKwargs]) -> None: + """Attach a per-cycle ``config_hook``, applied at the next ``spawn``. + + Args: + config_hook: Rebuilt every cycle against that cycle's context; its + ``tools`` patch is what the runtime installs for the cycle. + """ + ... + +@runtime_checkable +class Described(Protocol): + """Synchronous methods and properties; calls finish before returning. + + The optional seam the lead's member-tool text rides in on. + + A member that can say what it is for exposes ``description``; the core + appends it to the generic "Ask your team member ..." text of the tool the + lead calls it through, so the lead can pick between two members by what + they do rather than by name alone. ``None`` or blank keeps the generic + text unchanged. Members without the seam are skipped, not refused, like + ``Equippable``: an in-process typed member is described by its signature + and docstring already (T8). + """ + + description: str | None + """What the lead reads about this member on its tool; ``None`` keeps the generic text.""" + +# ── What hooks see ── + +@dataclass(eq=False, kw_only=True) +class Workspace: + """One run's shared state, handed to every hook method. + + A single argument rather than a growing signature, so a new field here + breaks no existing hook. ``data`` is the cross-hook scratch space the + run's report publishes; ``transcript`` is the audit trail the core itself + writes member calls and revise rounds into, and hooks may append their + own :class:`~ai_functions.experimental.teams.types.Note` entries. + ``members`` are live once ``on_assemble`` runs — spawning precedes every + hook call. ``lead`` is the lead's live thread handle, set as soon as the + lead's thread exists — before any hook runs — because a hook that opens a + delivery channel to the lead must open it before the lead's first cycle. + The lead's thread is registered but *not running* until after the + ``on_request`` fold. + + Identity, not value: two workspaces are never equal unless they are the + same object, so per-run hook state may key by the workspace itself, and + every construction mints a fresh ``run_id`` for the places a reference + cannot reach (T4). + """ + + team: Any + """The composition that opened this run. Typed loosely because the state plane must not import it.""" + request: str + """What the team was asked, before any ``on_request`` rewrite.""" + coordinator: Coordinator + """The coordinator every thread of this run lives on.""" + members: Sequence[Member] + """The cast, in team order; live from ``on_assemble`` onward.""" + run_id: RunId = ... + """Fresh per construction; the identity of this run in any journal.""" + data: dict[str, Any] = ... + """Cross-hook scratch space; published verbatim as the run's ``hooks_data``.""" + transcript: list[TranscriptEntry] = ... + """The audit trail: what the wire carried, in order.""" + lead: ThreadHandle[..., Any] | None = ... + """The lead's live handle; set before any hook runs, ``None`` only before that.""" + parent_call: CallContext | None = ... + cancelled: bool = ... + cancel_reason: str | None = ... + usage: TokenUsage | None = ... + fault: BaseException | None = ... + """The error the run is going to raise, set before ``on_closed`` (C14): what faulted the run, else the first + teardown error, else an earlier hook's ``on_closed`` error; ``None`` for a run that answered or was cancelled. + + A hook settling a record at ``on_closed`` reads the outcome here: a teardown error does not propagate until + every ``on_closed`` has run, so ``sys.exception()`` there cannot see it. + """ + calls: dict[str, int] = ... + """How many times each member has been asked this run; what ``MemberCall.call_id`` counts from (T4).""" + + async def ask(self, member: Member, request: str) -> Any: + """Ask from a hook through the run's ``before_ask`` checks and cancellation, preserving the typed answer. + + Hooks use this method for their own member asks. A direct ``member.ask`` + bypasses the team's ``before_ask`` checks and cancellation. The hook owns its transcript records. + """ + ... + +class TeamHook(Protocol): + """What a hook may implement. Every method is optional — absence means skip. + + Checked with ``getattr(hook, name, None)`` rather than ``isinstance``, + deliberately: a hook implements the two methods it needs and nothing + else, and an introspection-safe absence check is what lets a debugger's + ``hasattr`` probe never detonate a guard. This protocol is therefore *not* + ``runtime_checkable`` — an ``isinstance`` against it raises, which is the + point (T6). A protocol cannot mark a method optional, so a two-method hook + does not satisfy this one; ``Team(hooks=...)`` takes :data:`PartialTeamHook`, + one protocol per method here, and this class documents the full shape. + + ``tools_for_lead`` and ``tools_for_member`` are **synchronous**, because + the runtime calls the one ``config_hook`` synchronously inside a cycle and + documents "must not block". The other six methods (``on_assemble``, + ``on_request``, ``on_answer``, ``on_teardown``, ``on_closed``, + ``before_ask``) may be sync or async and the core awaits whichever it + finds. + + Three conventions a concrete hook should keep: per-run state resets on + ``Workspace`` identity, so one hook instance may serve many runs; a + model-fixable failure returns as ``"error: ..."`` text through the + ordinary channel, never raises, and never reads as approval; a + persistence failure raises. + """ + + def on_assemble(self, work: Workspace) -> Any: + """Members are live; the lead has not run. Hooks that brief members start here.""" + ... + + def on_request(self, work: Workspace, request: str) -> str | Awaitable[str]: + """Rewrite the request the lead will be asked. Folded left across hooks in order.""" + ... + + def tools_for_lead(self, work: Workspace, ctx: ThreadContext) -> Sequence[AgentTool]: + """Extra tools for the lead's cycle, rebuilt every cycle against ``ctx``. Sync.""" + ... + + def tools_for_member(self, work: Workspace, member: Member, ctx: ThreadContext) -> Sequence[AgentTool]: + """Extra tools for one member's cycle, rebuilt every cycle against ``ctx``. Sync.""" + ... + + def on_answer(self, work: Workspace, answer: Any) -> Verdict | Awaitable[Verdict]: + """Review the lead's answer. ``Accept`` passes it on; ``Revise`` re-runs the lead.""" + ... + + def on_teardown(self, work: Workspace) -> Any: + """The run is over (successfully or not); the cast is still live. Clean up.""" + ... + + def on_closed(self, work: Workspace) -> Any: + """Every member and the lead have been retired; nothing of the run is live. Release what they used. + + The step after ``on_teardown``: a hook that closes a resource a member + held open does it here and not in ``on_teardown``, where the member is + still alive. ``work.usage`` holds the run's final total and + ``work.fault`` what the run is going to raise (C13, C14). + """ + ... + + def before_ask(self, work: Workspace, member: Member) -> Any: + """Sync or awaitable check before every member ask, the lead's and a hook's. + + Raise ``StopRun`` to end the live run with a cancelled report. + Persistence errors must propagate; they never authorize an ask. + """ + ... + +# ── One protocol per hook method: what ``Team(hooks=...)`` accepts ── + +class _OnAssemble(Protocol): + def on_assemble(self, work: Workspace) -> Any: ... + +class _OnRequest(Protocol): + def on_request(self, work: Workspace, request: str) -> str | Awaitable[str]: ... + +class _ToolsForLead(Protocol): + def tools_for_lead(self, work: Workspace, ctx: ThreadContext) -> Sequence[AgentTool]: ... + +class _ToolsForMember(Protocol): + def tools_for_member(self, work: Workspace, member: Member, ctx: ThreadContext) -> Sequence[AgentTool]: ... + +class _OnAnswer(Protocol): + def on_answer(self, work: Workspace, answer: Any) -> Verdict | Awaitable[Verdict]: ... + +class _OnTeardown(Protocol): + def on_teardown(self, work: Workspace) -> Any: ... + +class _OnClosed(Protocol): + def on_closed(self, work: Workspace) -> Any: ... + +class _BeforeAsk(Protocol): + def before_ask(self, work: Workspace, member: Member) -> Any: ... + +PartialTeamHook = ( + _OnAssemble | _OnRequest | _ToolsForLead | _ToolsForMember | _OnAnswer | _OnTeardown | _OnClosed | _BeforeAsk +) +"""What a team's ``hooks`` holds: an object implementing at least one :class:`TeamHook` method. + +The core reaches every method with ``getattr`` and skips an absent one (T6), so a hook implements the methods it needs. +A union is satisfied by any one member: a type checker accepts the object once one of its methods matches the signature +:class:`TeamHook` documents and does not check its other methods, so a mistyped second method fails only at run time. An +object with none of the methods is rejected. +""" diff --git a/spec/ai_functions/experimental/teams/core.pyi b/spec/ai_functions/experimental/teams/core.pyi new file mode 100644 index 00000000..45538081 --- /dev/null +++ b/spec/ai_functions/experimental/teams/core.pyi @@ -0,0 +1,382 @@ +"""``Team`` — a lead over live members, with hooks for everything that is not that. + +The core owns exactly what every team needs — spawn the members, run the lead +with the members as typed tools, retire everybody — and everything else arrives +as a :class:`~ai_functions.experimental.teams.TeamHook`. A hook can contribute +tools to the lead or the members, rewrite the request, and review the answer +through a bounded Accept/Revise loop the core drives. The bare team is the +whole API for the common case:: + + team = Team( + lead=decide, + members=[FunctionMember(read, name="left"), FunctionMember(read, name="right")], + ) + run = await team.run("who is right") + print(run.answer) + +**The one hard mechanical constraint.** The runtime resolves exactly one +``config_hook`` per cycle and its ``tools`` patch *replaces* the compiled +tools. ``AIThread._run_cycle`` calls ``self._config.config_hook(ctx)`` once at +the top of the cycle, pops the ``config_hook`` key out of the patch and applies +the rest with ``dataclasses.replace``, so a ``tools`` key stands in for the +compiled ``tools=`` wholesale; ``ThreadConfig.config_hook`` documents that +replace-not-merge semantics; and the same method appends the coordinator tools +*after* the patch, only when the cycle's ``coordinator_tools_enabled`` holds, +which the core turns off on every thread it spawns (C12). +So the core owns the single hook on each thread and folds every contribution +into it: the lead's own hook and tools are recomposed first, then the member +tools, then every ``tools_for_lead``; a member's own ``tools=`` override is +recomposed by ``FunctionMember.spawn``'s wrapper. Nothing else may set +``config_hook`` on a thread the core manages — ``FunctionMember.equip`` refuses a +member that carries its own. + +Invariants: + C1 — one run is one subtree, torn down on every path. The lead's thread + exists first, every member spawns as its child, and the ``finally`` runs + every ``on_teardown`` and then retires everybody even when a hook or the + lead raised; a teardown hook's own error resurfaces only when nothing + else is already propagating. + + C2 — one ``config_hook`` per managed thread. The core composes every tool + source into that one hook and never stacks a second; the lead's own hook + and ``tools=`` survive the fold, and so do a member's own tools. A member + without the ``Equippable`` seam is skipped by the fold, not refused, and + the transcript names it in an ``unequipped`` note with the hooks whose + tools it goes without. + + C3 — the answer loop terminates. Every restart spends one round of some + hook's budget, so the loop is bounded by the sum of the caps in play; + budgets persist across restarts, the cap is read off the latest verdict, + and exhaustion passes the answer through and records ``revise_cap`` once + per hook. + + C4 — a verdict that is neither ``Accept`` nor ``Revise`` is a wiring bug + and raises :class:`WiringError` naming the hook, never a silent pass. + + C5 — member names are unique on the lead's wire; a collision is refused + at construction, before anything spawns. + + C6 — the transcript is written by the wire that carried the event, as + the typed records of :mod:`~ai_functions.experimental.teams.types`, and + always under the member's real name whatever the tool had to be called. + A member's answer is recorded as the text the lead read: a string as is, + anything else formatted as a Strands tool formats its result (a model as + its JSON). + + C7 — NoUnreviewedShip. After any ``Revise`` by hook k the walk restarts + at hook 1 and re-consults hooks 1..k-1 on the revised answer before any + ``Accept`` can ship it; an answer ships only after one full uninterrupted + pass of ``Accept`` and cap-exhausted pass-throughs. In the code this is + the ``break`` out of the ``for`` and the ``while revised`` walking again + in ``_answer_loop``. The rejected alternative — per-hook loops rather + than a restart chain — ships an accepted-then-mutated answer, and + ``tests/teams/test_teams_invariants.py`` pins the code to the restart + chain. A cap-exhausted pass-through in that final pass is still a hook + that did not approve; whether the run may call itself a success anyway + is C8's question. + + C8 — an exhausted mandatory review is a blocked result, never ordinary + success. When the final pass carries a cap-exhausted pass-through whose + latest verdict said ``mandatory=True``, ``Team.run`` returns + ``TeamRun(blocked=True, blocked_by=(