diff --git a/examples/coding-agent/README.md b/examples/coding-agent/README.md new file mode 100644 index 000000000..d37c5c097 --- /dev/null +++ b/examples/coding-agent/README.md @@ -0,0 +1,74 @@ +# Coding agent + +A minimal coding agent harness: an LLM that reads files, writes files, and runs +shell commands in a loop until it has finished a task. + +The point of this example is the *loop*. A tool-calling assistant runs one tool +and returns to the user. An agent keeps going on its own -- it reads a file, +sees what's in it, edits it, runs the tests, sees them fail, and tries again -- +without a human in between. That is a cycle, and expressing cycles clearly is +what Burr is for. + +![State machine](statemachine.png) + +## The loop + +| Action | What it does | +| --- | --- | +| `human_input` | Takes the task and seeds the message history | +| `create_prompt` | A seam for shaping the prompt before each call -- add a file tree, a summary of earlier steps, retrieved context | +| `call_llm` | Asks the model what to do next: call a tool, or finish | +| `read_file` / `write_file` / `list_files` / `run_bash` | One bound action per tool, so each shows up separately in the Burr UI | +| `respond` | Surfaces the answer, or reports that the step budget ran out | + +Two transitions leave `call_llm` for `respond`: one when the model answers +without requesting a tool, and one when `steps` reaches `max_steps`. Both exits +are listed before the tool transitions, because Burr takes the first condition +that matches. An agent that decides its own next step can loop forever, so the +budget is not optional. + +Tool results are appended to `messages` as `role: "tool"` entries carrying the +`tool_call_id` they answer. That is what lets the model see what its last action +actually returned. + +## Running it + + pip install -r requirements.txt + python application.py + +With no API key set, the example uses a scripted client that replays a fixed +sequence of tool calls, so it runs immediately and in CI. Set `OPENAI_API_KEY` +to use a real model instead. + +Regenerating `statemachine.png` needs the Graphviz binary (`apt-get install +graphviz`), not just the Python package. + +To watch a run in the Burr UI: + + burr + +## Safety + +`run_bash` executes shell commands with the permissions of the process running +the agent. File paths are confined to the workspace directory, but a shell +command can simply `cd` out of it. **This is a teaching example, not a sandbox.** +Run it against a directory you do not mind losing, ideally inside a container. + +Adding a permission step before tool execution -- prompting for approval, or +checking an allowlist -- is the natural next thing to build. + +## Known simplifications + +- One tool call is executed per iteration; if the model requests several, the + extras are dropped. +- `create_prompt` is a passthrough. It exists to show where prompt construction + belongs, not because it does anything yet. +- Message history grows without bound. A longer-running agent needs summarisation + or truncation. + +## Files + +- [application.py](application.py) -- the state machine, actions, and LLM clients +- [tools.py](tools.py) -- the four tools +- [notebook.ipynb](notebook.ipynb) -- the same example, walked through step by step +- [requirements.txt](requirements.txt) -- the environment \ No newline at end of file diff --git a/examples/coding-agent/__init__.py b/examples/coding-agent/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/examples/coding-agent/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/examples/coding-agent/application.py b/examples/coding-agent/application.py new file mode 100644 index 000000000..5a66c0ca9 --- /dev/null +++ b/examples/coding-agent/application.py @@ -0,0 +1,271 @@ +"""Tools the coding agent can call. Each takes typed args and returns a dict.""" +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import inspect +import json +import os +from typing import Callable, Optional + +from burr.core import State, action, expr, when +from burr.core.application import ApplicationBuilder + +from tools import list_files, read_file, run_bash, write_file + +TOOLS = { + "list_files": list_files, + "read_file": read_file, + "write_file": write_file, + "run_bash": run_bash, +} + +TYPE_MAP = {str: "string", int: "integer", float: "number", bool: "boolean"} + +OPENAI_TOOLS = [ + { + "type": "function", + "function": { + "name": name, + "description": fn.__doc__ or name, + "parameters": { + "type": "object", + "properties": { + p.name: { + "type": TYPE_MAP.get(p.annotation, "string"), + "description": p.name, + } + for p in inspect.signature(fn).parameters.values() + }, + "required": [ + p.name + for p in inspect.signature(fn).parameters.values() + if p.default is inspect.Parameter.empty + ], + }, + }, + } + for name, fn in TOOLS.items() +] + +SYSTEM_PROMPT = ( + "You are a coding agent working inside a project directory. " + "Use the tools to inspect and modify files, and to run commands. " + "Work one step at a time: look before you edit, and verify changes by running them. " + "When the task is complete, reply with a short summary and request no tool." +) + + +# --- LLM clients ----------------------------------------------------------- +# Both return the same shape: +# {"content": str | None, "tool_calls": [{"id", "name", "args"}]} + + +class OpenAIClient: + """Calls OpenAI's chat completions API with tool calling enabled.""" + + def __init__(self, model: str = "gpt-4o"): + self.model = model + + def __call__(self, messages: list[dict]) -> dict: + import openai + + response = openai.chat.completions.create( + model=self.model, messages=messages, tools=OPENAI_TOOLS + ) + message = response.choices[0].message + calls = [ + {"id": c.id, "name": c.function.name, "args": json.loads(c.function.arguments)} + for c in (message.tool_calls or []) + ] + return {"content": message.content, "tool_calls": calls} + + +class ScriptedClient: + """Replays a fixed list of responses. Lets the example run without an API key.""" + + def __init__(self, responses: list[dict]): + self.responses = list(responses) + self.index = 0 + + def __call__(self, messages: list[dict]) -> dict: + if self.index >= len(self.responses): + return {"content": "Script exhausted.", "tool_calls": []} + response = self.responses[self.index] + self.index += 1 + return response + + +DEFAULT_SCRIPT = [ + {"content": None, "tool_calls": [{"id": "c1", "name": "list_files", "args": {}}]}, + { + "content": None, + "tool_calls": [ + {"id": "c2", "name": "write_file", + "args": {"path": "hello.py", "contents": "print('hello from the agent')\n"}} + ], + }, + {"content": None, "tool_calls": [{"id": "c3", "name": "run_bash", + "args": {"command": "python hello.py"}}]}, + {"content": "Created hello.py and confirmed it runs.", "tool_calls": []}, +] + + +def get_client(): + """Real client when OPENAI_API_KEY is set, otherwise the scripted stand-in.""" + if os.environ.get("OPENAI_API_KEY"): + return OpenAIClient() + return ScriptedClient(DEFAULT_SCRIPT) + + +# --- Actions --------------------------------------------------------------- + + +@action(reads=[], writes=["task", "messages", "steps", "done", "final_answer"]) +def human_input(state: State, task: str) -> State: + """Takes a task from the user and starts a fresh run.""" + return state.update( + task=task, + messages=[ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": task}, + ], + steps=0, + done=False, + final_answer=None, + ) + + +@action(reads=["messages"], writes=["messages"]) +def create_prompt(state: State) -> State: + """Seam for shaping the prompt before each call -- add file trees, summaries, etc.""" + return state.update(messages=state["messages"]) + + +@action( + reads=["messages", "steps"], + writes=["messages", "next_tool", "next_args", "last_tool_call_id", + "steps", "done", "final_answer"], +) +def call_llm(state: State, client: Callable) -> State: + """Asks the model what to do next: call a tool, or finish.""" + result = client(state["messages"]) + calls = result["tool_calls"] + + if not calls: + return state.update( + messages=state["messages"] + [{"role": "assistant", "content": result["content"]}], + next_tool=None, + next_args={}, + last_tool_call_id=None, + steps=state["steps"] + 1, + done=True, + final_answer=result["content"], + ) + + # One tool per iteration keeps the graph readable; extras are dropped. + call = calls[0] + assistant_message = { + "role": "assistant", + "content": result["content"], + "tool_calls": [ + { + "id": call["id"], + "type": "function", + "function": {"name": call["name"], "arguments": json.dumps(call["args"])}, + } + ], + } + return state.update( + messages=state["messages"] + [assistant_message], + next_tool=call["name"], + next_args=call["args"], + last_tool_call_id=call["id"], + steps=state["steps"] + 1, + done=False, + final_answer=None, + ) + + +@action(reads=["next_args", "last_tool_call_id", "messages"], writes=["messages"]) +def execute_tool(state: State, tool_function: Callable) -> State: + """Runs one tool and feeds its result back to the model.""" + result = tool_function(**state["next_args"]) + return state.update( + messages=state["messages"] + + [ + { + "role": "tool", + "tool_call_id": state["last_tool_call_id"], + "content": json.dumps(result), + } + ] + ) + + +@action(reads=["final_answer", "steps", "max_steps"], writes=["final_answer"]) +def respond(state: State) -> State: + """Surfaces the answer, or explains that the budget ran out.""" + if state["final_answer"]: + return state.update(final_answer=state["final_answer"]) + return state.update( + final_answer=f"Stopped after {state['steps']} steps without finishing the task." + ) + + +def application(app_id: Optional[str] = None, max_steps: int = 15, client: Callable = None): + """Builds the coding agent application.""" + client = client or get_client() + return ( + ApplicationBuilder() + .with_actions( + human_input, + create_prompt, + respond, + call_llm=call_llm.bind(client=client), + read_file=execute_tool.bind(tool_function=read_file), + write_file=execute_tool.bind(tool_function=write_file), + list_files=execute_tool.bind(tool_function=list_files), + run_bash=execute_tool.bind(tool_function=run_bash), + ) + .with_transitions( + ("human_input", "create_prompt"), + ("create_prompt", "call_llm"), + ("call_llm", "respond", when(done=True)), + ("call_llm", "respond", expr("steps>=max_steps")), + ("call_llm", "read_file", when(next_tool="read_file")), + ("call_llm", "write_file", when(next_tool="write_file")), + ("call_llm", "list_files", when(next_tool="list_files")), + ("call_llm", "run_bash", when(next_tool="run_bash")), + (["read_file", "write_file", "list_files", "run_bash"], "call_llm"), + ("respond", "human_input"), + ) + .with_state(max_steps=max_steps, steps=0, messages=[], done=False, final_answer=None) + .with_identifiers(app_id=app_id) + .with_entrypoint("human_input") + .with_tracker(project="demo_coding_agent") + .build() + ) + + +if __name__ == "__main__": + app = application() + app.visualize(output_file_path="./statemachine.png") + _, _, state = app.run( + halt_after=["respond"], + inputs={"task": "Create a hello.py that prints a greeting, then run it."}, + ) + print(state["final_answer"]) \ No newline at end of file diff --git a/examples/coding-agent/notebook.ipynb b/examples/coding-agent/notebook.ipynb new file mode 100644 index 000000000..ce354a34f --- /dev/null +++ b/examples/coding-agent/notebook.ipynb @@ -0,0 +1,716 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "# Licensed to the Apache Software Foundation (ASF) under one\n", + "# or more contributor license agreements. See the NOTICE file\n", + "# distributed with this work for additional information\n", + "# regarding copyright ownership. The ASF licenses this file\n", + "# to you under the Apache License, Version 2.0 (the\n", + "# \"License\"); you may not use this file except in compliance\n", + "# with the License. You may obtain a copy of the License at\n", + "#\n", + "# http://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing,\n", + "# software distributed under the License is distributed on an\n", + "# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n", + "# KIND, either express or implied. See the License for the\n", + "# specific language governing permissions and limitations\n", + "# under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Coding agent\n", + "\n", + "A minimal coding agent harness: an LLM that reads files, writes files, and runs\n", + "shell commands in a loop until the task is done.\n", + "\n", + "The point is the **loop**. A tool-calling assistant runs one tool and hands back\n", + "to the user. An agent keeps going by itself -- read a file, edit it, run the\n", + "tests, see them fail, try again -- with no human in between. That is a cycle,\n", + "and cycles are what Burr is for.\n", + "\n", + "This notebook runs without an API key: a scripted client stands in for the model." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Imports" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "import inspect\n", + "import json\n", + "import os\n", + "import subprocess\n", + "from typing import Callable, Optional\n", + "\n", + "from burr.core import State, action, expr, when\n", + "from burr.core.application import ApplicationBuilder" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# The tools\n", + "\n", + "Four functions. Each takes primitive, type-annotated arguments and returns a\n", + "dict -- and never raises. An exception would kill the state machine; an\n", + "`{\"error\": ...}` dict goes back to the model, which can then try something else.\n", + "\n", + "Paths are confined to a workspace directory. Note that `run_bash` can `cd` out\n", + "of it: this is a teaching example, not a sandbox." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"Tools the coding agent can call. Each takes typed args and returns a dict.\"\"\"\n", + "# Licensed to the Apache Software Foundation (ASF) under one\n", + "# or more contributor license agreements. See the NOTICE file\n", + "# distributed with this work for additional information\n", + "# regarding copyright ownership. The ASF licenses this file\n", + "# to you under the Apache License, Version 2.0 (the\n", + "# \"License\"); you may not use this file except in compliance\n", + "# with the License. You may obtain a copy of the License at\n", + "#\n", + "# http://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing,\n", + "# software distributed under the License is distributed on an\n", + "# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n", + "# KIND, either express or implied. See the License for the\n", + "# specific language governing permissions and limitations\n", + "# under the License.\n", + "\n", + "import inspect\n", + "import json\n", + "import os\n", + "from typing import Callable, Optional\n", + "\n", + "import openai\n", + "import requests\n", + "\n", + "from burr.core import State, action, when\n", + "from burr.core.application import ApplicationBuilder\n", + "\n", + "import os\n", + "import subprocess\n", + "\n", + "# Everything is confined to this directory. See README on why this is not a sandbox.\n", + "WORKSPACE = os.environ.get(\"CODING_AGENT_WORKSPACE\", \"./workspace\")\n", + "\n", + "\n", + "def _resolve(path: str) -> str:\n", + " \"\"\"Resolve a path inside the workspace, refusing anything that escapes it.\"\"\"\n", + " root = os.path.abspath(WORKSPACE)\n", + " full = os.path.abspath(os.path.join(root, path))\n", + " if not full.startswith(root + os.sep) and full != root:\n", + " raise ValueError(f\"path escapes workspace: {path}\")\n", + " return full\n", + "\n", + "\n", + "def list_files(directory: str = \".\") -> dict:\n", + " \"\"\"Lists the files and directories at the given path inside the workspace.\"\"\"\n", + " try:\n", + " target = _resolve(directory)\n", + " return {\"entries\": sorted(os.listdir(target))}\n", + " except (ValueError, OSError) as e:\n", + " return {\"error\": str(e)}\n", + "\n", + "\n", + "def read_file(path: str) -> dict:\n", + " \"\"\"Reads a text file from the workspace and returns its contents.\"\"\"\n", + " try:\n", + " with open(_resolve(path)) as f:\n", + " return {\"path\": path, \"contents\": f.read()}\n", + " except (ValueError, OSError) as e:\n", + " return {\"error\": str(e)}\n", + "\n", + "\n", + "def write_file(path: str, contents: str) -> dict:\n", + " \"\"\"Writes text to a file in the workspace, creating or overwriting it.\"\"\"\n", + " try:\n", + " full = _resolve(path)\n", + " os.makedirs(os.path.dirname(full), exist_ok=True)\n", + " with open(full, \"w\") as f:\n", + " f.write(contents)\n", + " return {\"path\": path, \"bytes_written\": len(contents)}\n", + " except (ValueError, OSError) as e:\n", + " return {\"error\": str(e)}\n", + "\n", + "\n", + "def run_bash(command: str) -> dict:\n", + " \"\"\"Runs a shell command in the workspace and returns its output.\"\"\"\n", + " try:\n", + " proc = subprocess.run(\n", + " command, shell=True, cwd=os.path.abspath(WORKSPACE),\n", + " capture_output=True, text=True, timeout=30,\n", + " )\n", + " return {\n", + " \"exit_code\": proc.returncode,\n", + " \"stdout\": proc.stdout[-4000:],\n", + " \"stderr\": proc.stderr[-4000:],\n", + " }\n", + " except subprocess.TimeoutExpired:\n", + " return {\"error\": \"command timed out after 30s\"}\n", + " " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Describing the tools to the model\n", + "\n", + "The schema is derived from signatures and docstrings, so the tools stay the single source of truth. Only parameters without defaults are marked required." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "dfc10623", + "metadata": {}, + "outputs": [], + "source": [ + "TOOLS = {\n", + " \"list_files\": list_files,\n", + " \"read_file\": read_file,\n", + " \"write_file\": write_file,\n", + " \"run_bash\": run_bash,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "TYPE_MAP = {str: \"string\", int: \"integer\", float: \"number\", bool: \"boolean\"}\n", + "\n", + "OPENAI_TOOLS = [\n", + " {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": name,\n", + " \"description\": fn.__doc__ or name,\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " p.name: {\n", + " \"type\": TYPE_MAP.get(p.annotation, \"string\"),\n", + " \"description\": p.name,\n", + " }\n", + " for p in inspect.signature(fn).parameters.values()\n", + " },\n", + " \"required\": [\n", + " p.name\n", + " for p in inspect.signature(fn).parameters.values()\n", + " if p.default is inspect.Parameter.empty\n", + " ],\n", + " },\n", + " },\n", + " }\n", + " for name, fn in TOOLS.items()\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "SYSTEM_PROMPT = (\n", + " \"You are a coding agent working inside a project directory. \"\n", + " \"Use the tools to inspect and modify files, and to run commands. \"\n", + " \"Work one step at a time: look before you edit, and verify changes by running them. \"\n", + " \"When the task is complete, reply with a short summary and request no tool.\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Two clients\n", + "\n", + "Both return the same shape, so the rest of the harness doesn't care which is in\n", + "use. The scripted client replays a fixed sequence, which is what lets this\n", + "notebook -- and CI -- run with no API key." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "class OpenAIClient:\n", + " \"\"\"Calls OpenAI's chat completions API with tool calling enabled.\"\"\"\n", + "\n", + " def __init__(self, model: str = \"gpt-4o\"):\n", + " self.model = model\n", + "\n", + " def __call__(self, messages: list[dict]) -> dict:\n", + " import openai\n", + "\n", + " response = openai.chat.completions.create(\n", + " model=self.model, messages=messages, tools=OPENAI_TOOLS\n", + " )\n", + " message = response.choices[0].message\n", + " calls = [\n", + " {\"id\": c.id, \"name\": c.function.name, \"args\": json.loads(c.function.arguments)}\n", + " for c in (message.tool_calls or [])\n", + " ]\n", + " return {\"content\": message.content, \"tool_calls\": calls}\n", + "\n", + "\n", + "class ScriptedClient:\n", + " \"\"\"Replays a fixed list of responses. Lets the example run without an API key.\"\"\"\n", + "\n", + " def __init__(self, responses: list[dict]):\n", + " self.responses = list(responses)\n", + " self.index = 0\n", + "\n", + " def __call__(self, messages: list[dict]) -> dict:\n", + " if self.index >= len(self.responses):\n", + " return {\"content\": \"Script exhausted.\", \"tool_calls\": []}\n", + " response = self.responses[self.index]\n", + " self.index += 1\n", + " return response\n", + "\n", + "\n", + "DEFAULT_SCRIPT = [\n", + " {\"content\": None, \"tool_calls\": [{\"id\": \"c1\", \"name\": \"list_files\", \"args\": {}}]},\n", + " {\n", + " \"content\": None,\n", + " \"tool_calls\": [\n", + " {\"id\": \"c2\", \"name\": \"write_file\",\n", + " \"args\": {\"path\": \"hello.py\", \"contents\": \"print('hello from the agent')\\n\"}}\n", + " ],\n", + " },\n", + " {\"content\": None, \"tool_calls\": [{\"id\": \"c3\", \"name\": \"run_bash\",\n", + " \"args\": {\"command\": \"python hello.py\"}}]},\n", + " {\"content\": \"Created hello.py and confirmed it runs.\", \"tool_calls\": []},\n", + "]\n", + "\n", + "\n", + "def get_client():\n", + " \"\"\"Real client when OPENAI_API_KEY is set, otherwise the scripted stand-in.\"\"\"\n", + " if os.environ.get(\"OPENAI_API_KEY\"):\n", + " return OpenAIClient()\n", + " return ScriptedClient(DEFAULT_SCRIPT)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# The actions\n", + "\n", + "`call_llm` is the interesting one. It appends the model's reply to `messages`,\n", + "and either sets `done` (no tool requested) or records which tool to run next.\n", + "Tool results come back as `role: \"tool\"` messages carrying the `tool_call_id`\n", + "they answer -- that is how the model sees what its last action returned." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "@action(reads=[], writes=[\"task\", \"messages\", \"steps\", \"done\", \"final_answer\"])\n", + "def human_input(state: State, task: str) -> State:\n", + " \"\"\"Takes a task from the user and starts a fresh run.\"\"\"\n", + " return state.update(\n", + " task=task,\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n", + " {\"role\": \"user\", \"content\": task},\n", + " ],\n", + " steps=0,\n", + " done=False,\n", + " final_answer=None,\n", + " )\n", + "\n", + "\n", + "@action(reads=[\"messages\"], writes=[\"messages\"])\n", + "def create_prompt(state: State) -> State:\n", + " \"\"\"Seam for shaping the prompt before each call -- add file trees, summaries, etc.\"\"\"\n", + " return state.update(messages=state[\"messages\"])\n", + "\n", + "\n", + "@action(\n", + " reads=[\"messages\", \"steps\"],\n", + " writes=[\"messages\", \"next_tool\", \"next_args\", \"last_tool_call_id\",\n", + " \"steps\", \"done\", \"final_answer\"],\n", + ")\n", + "def call_llm(state: State, client: Callable) -> State:\n", + " \"\"\"Asks the model what to do next: call a tool, or finish.\"\"\"\n", + " result = client(state[\"messages\"])\n", + " calls = result[\"tool_calls\"]\n", + "\n", + " if not calls:\n", + " return state.update(\n", + " messages=state[\"messages\"] + [{\"role\": \"assistant\", \"content\": result[\"content\"]}],\n", + " next_tool=None,\n", + " next_args={},\n", + " last_tool_call_id=None,\n", + " steps=state[\"steps\"] + 1,\n", + " done=True,\n", + " final_answer=result[\"content\"],\n", + " )\n", + "\n", + " # One tool per iteration keeps the graph readable; extras are dropped.\n", + " call = calls[0]\n", + " assistant_message = {\n", + " \"role\": \"assistant\",\n", + " \"content\": result[\"content\"],\n", + " \"tool_calls\": [\n", + " {\n", + " \"id\": call[\"id\"],\n", + " \"type\": \"function\",\n", + " \"function\": {\"name\": call[\"name\"], \"arguments\": json.dumps(call[\"args\"])},\n", + " }\n", + " ],\n", + " }\n", + " return state.update(\n", + " messages=state[\"messages\"] + [assistant_message],\n", + " next_tool=call[\"name\"],\n", + " next_args=call[\"args\"],\n", + " last_tool_call_id=call[\"id\"],\n", + " steps=state[\"steps\"] + 1,\n", + " done=False,\n", + " final_answer=None,\n", + " )\n", + "\n", + "\n", + "@action(reads=[\"next_args\", \"last_tool_call_id\", \"messages\"], writes=[\"messages\"])\n", + "def execute_tool(state: State, tool_function: Callable) -> State:\n", + " \"\"\"Runs one tool and feeds its result back to the model.\"\"\"\n", + " result = tool_function(**state[\"next_args\"])\n", + " return state.update(\n", + " messages=state[\"messages\"]\n", + " + [\n", + " {\n", + " \"role\": \"tool\",\n", + " \"tool_call_id\": state[\"last_tool_call_id\"],\n", + " \"content\": json.dumps(result),\n", + " }\n", + " ]\n", + " )\n", + "\n", + "\n", + "@action(reads=[\"final_answer\", \"steps\", \"max_steps\"], writes=[\"final_answer\"])\n", + "def respond(state: State) -> State:\n", + " \"\"\"Surfaces the answer, or explains that the budget ran out.\"\"\"\n", + " if state[\"final_answer\"]:\n", + " return state.update(final_answer=state[\"final_answer\"])\n", + " return state.update(\n", + " final_answer=f\"Stopped after {state['steps']} steps without finishing the task.\"\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Wiring the graph\n", + "\n", + "Two transitions leave `call_llm` for `respond`: the model finished, or the step\n", + "budget ran out. Both are listed **before** the tool transitions, because Burr\n", + "takes the first matching condition. An agent that picks its own next step can\n", + "loop forever, so the budget is not optional.\n", + "\n", + "Each tool is a separately bound action, so the Burr UI shows exactly which one\n", + "ran at each step." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [], + "source": [ + "def application(app_id: Optional[str] = None, max_steps: int = 15, client: Callable = None):\n", + " \"\"\"Builds the coding agent application.\"\"\"\n", + " client = client or get_client()\n", + " return (\n", + " ApplicationBuilder()\n", + " .with_actions(\n", + " human_input,\n", + " create_prompt,\n", + " respond,\n", + " call_llm=call_llm.bind(client=client),\n", + " read_file=execute_tool.bind(tool_function=read_file),\n", + " write_file=execute_tool.bind(tool_function=write_file),\n", + " list_files=execute_tool.bind(tool_function=list_files),\n", + " run_bash=execute_tool.bind(tool_function=run_bash),\n", + " )\n", + " .with_transitions(\n", + " (\"human_input\", \"create_prompt\"),\n", + " (\"create_prompt\", \"call_llm\"),\n", + " (\"call_llm\", \"respond\", when(done=True)),\n", + " (\"call_llm\", \"respond\", expr(\"steps>=max_steps\")),\n", + " (\"call_llm\", \"read_file\", when(next_tool=\"read_file\")),\n", + " (\"call_llm\", \"write_file\", when(next_tool=\"write_file\")),\n", + " (\"call_llm\", \"list_files\", when(next_tool=\"list_files\")),\n", + " (\"call_llm\", \"run_bash\", when(next_tool=\"run_bash\")),\n", + " ([\"read_file\", \"write_file\", \"list_files\", \"run_bash\"], \"call_llm\"),\n", + " (\"respond\", \"human_input\"),\n", + " )\n", + " .with_state(max_steps=max_steps, steps=0, messages=[], done=False, final_answer=None)\n", + " .with_identifiers(app_id=app_id)\n", + " .with_entrypoint(\"human_input\")\n", + " .with_tracker(project=\"demo_coding_agent\")\n", + " .build()\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Running it" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "image/svg+xml": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "%3\n", + "\n", + "\n", + "\n", + "human_input\n", + "\n", + "human_input\n", + "\n", + "\n", + "\n", + "create_prompt\n", + "\n", + "create_prompt\n", + "\n", + "\n", + "\n", + "human_input->create_prompt\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "input__task\n", + "\n", + "input: task\n", + "\n", + "\n", + "\n", + "input__task->human_input\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "call_llm\n", + "\n", + "call_llm\n", + "\n", + "\n", + "\n", + "create_prompt->call_llm\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "respond\n", + "\n", + "respond\n", + "\n", + "\n", + "\n", + "respond->human_input\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "call_llm->respond\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "call_llm->respond\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "read_file\n", + "\n", + "read_file\n", + "\n", + "\n", + "\n", + "call_llm->read_file\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "write_file\n", + "\n", + "write_file\n", + "\n", + "\n", + "\n", + "call_llm->write_file\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "list_files\n", + "\n", + "list_files\n", + "\n", + "\n", + "\n", + "call_llm->list_files\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "run_bash\n", + "\n", + "run_bash\n", + "\n", + "\n", + "\n", + "call_llm->run_bash\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "read_file->call_llm\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "write_file->call_llm\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "list_files->call_llm\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "run_bash->call_llm\n", + "\n", + "\n", + "\n", + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "app = application()\n", + "app.visualize()" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Created hello.py and confirmed it runs.\n" + ] + } + ], + "source": [ + "_, _, state = app.run(\n", + " halt_after=[\"respond\"],\n", + " inputs={\"task\": \"Create a hello.py that prints a greeting, then run it.\"},\n", + ")\n", + "print(state[\"final_answer\"])" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv (3.12.1.final.0)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.1" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/coding-agent/requirements.txt b/examples/coding-agent/requirements.txt new file mode 100644 index 000000000..40d0c6a61 --- /dev/null +++ b/examples/coding-agent/requirements.txt @@ -0,0 +1,19 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +burr +openai \ No newline at end of file diff --git a/examples/coding-agent/statemachine.png b/examples/coding-agent/statemachine.png new file mode 100644 index 000000000..701ea129c Binary files /dev/null and b/examples/coding-agent/statemachine.png differ diff --git a/examples/coding-agent/tools.py b/examples/coding-agent/tools.py new file mode 100644 index 000000000..81c6d45eb --- /dev/null +++ b/examples/coding-agent/tools.py @@ -0,0 +1,89 @@ +"""Tools the coding agent can call. Each takes typed args and returns a dict.""" +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import inspect +import json +import os +from typing import Callable, Optional + +import openai +import requests + +from burr.core import State, action, when +from burr.core.application import ApplicationBuilder + +import os +import subprocess + +# Everything is confined to this directory. See README on why this is not a sandbox. +WORKSPACE = os.environ.get("CODING_AGENT_WORKSPACE", "./workspace") + + +def _resolve(path: str) -> str: + """Resolve a path inside the workspace, refusing anything that escapes it.""" + root = os.path.abspath(WORKSPACE) + full = os.path.abspath(os.path.join(root, path)) + if not full.startswith(root + os.sep) and full != root: + raise ValueError(f"path escapes workspace: {path}") + return full + + +def list_files(directory: str = ".") -> dict: + """Lists the files and directories at the given path inside the workspace.""" + try: + target = _resolve(directory) + return {"entries": sorted(os.listdir(target))} + except (ValueError, OSError) as e: + return {"error": str(e)} + + +def read_file(path: str) -> dict: + """Reads a text file from the workspace and returns its contents.""" + try: + with open(_resolve(path)) as f: + return {"path": path, "contents": f.read()} + except (ValueError, OSError) as e: + return {"error": str(e)} + + +def write_file(path: str, contents: str) -> dict: + """Writes text to a file in the workspace, creating or overwriting it.""" + try: + full = _resolve(path) + os.makedirs(os.path.dirname(full), exist_ok=True) + with open(full, "w") as f: + f.write(contents) + return {"path": path, "bytes_written": len(contents)} + except (ValueError, OSError) as e: + return {"error": str(e)} + + +def run_bash(command: str) -> dict: + """Runs a shell command in the workspace and returns its output.""" + try: + proc = subprocess.run( + command, shell=True, cwd=os.path.abspath(WORKSPACE), + capture_output=True, text=True, timeout=30, + ) + return { + "exit_code": proc.returncode, + "stdout": proc.stdout[-4000:], + "stderr": proc.stderr[-4000:], + } + except subprocess.TimeoutExpired: + return {"error": "command timed out after 30s"} \ No newline at end of file