509 lines
19 KiB
Python
509 lines
19 KiB
Python
"""
|
||
Coding Agent
|
||
============
|
||
Autonomous AI coding agent based on the Plan→Act→Observe→Fix→Done loop.
|
||
Structure follows ex05_coding_agent_solution.py from the course.
|
||
|
||
Instead of a blocking CLI loop (input()), the CodingAgent class exposes
|
||
step-by-step methods so Streamlit can drive the loop via session_state:
|
||
|
||
agent.start_task(task) # initialise
|
||
action = agent.propose_next_action() # ask LLM → returns action, does NOT execute
|
||
result = agent.approve() # execute pending action
|
||
agent.reject(feedback) # skip action, inject user feedback
|
||
"""
|
||
|
||
import ast
|
||
import inspect
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import requests
|
||
from dotenv import load_dotenv
|
||
|
||
load_dotenv()
|
||
|
||
# ── Workspace ────────────────────────────────────────────────────────────────
|
||
# Two levels up from backend/agent/ → project root → workspace/
|
||
WORKSPACE = Path(__file__).resolve().parents[2] / "workspace"
|
||
WORKSPACE.mkdir(exist_ok=True)
|
||
|
||
# ── Agent limits ─────────────────────────────────────────────────────────────
|
||
MAX_ITERATIONS = 100
|
||
MAX_RESULT_LENGTH = 10_000
|
||
MAX_HISTORY_CHARS = 80_000
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# PART A – TOOL FUNCTIONS
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
#
|
||
# Each tool is a plain Python function decorated with @register_tool.
|
||
# The decorator adds the function to TOOL_REGISTRY so the dispatcher
|
||
# can call it by name at runtime.
|
||
|
||
TOOL_REGISTRY: dict[str, callable] = {}
|
||
|
||
|
||
def register_tool(func):
|
||
"""Decorator – adds a function to the global tool registry."""
|
||
TOOL_REGISTRY[func.__name__] = func
|
||
return func
|
||
|
||
|
||
@register_tool
|
||
def read_file(path: str) -> str:
|
||
"""Read a .py or .txt file from the workspace and return its contents."""
|
||
target = (WORKSPACE / path).resolve()
|
||
if not str(target).startswith(str(WORKSPACE.resolve())):
|
||
return "ERROR: path is outside the workspace."
|
||
if not target.exists():
|
||
return f"ERROR: file '{path}' not found."
|
||
if target.suffix not in (".py", ".txt"):
|
||
return f"ERROR: can only read .py and .txt files, got '{target.suffix}'."
|
||
return target.read_text()
|
||
|
||
|
||
@register_tool
|
||
def write_file(path: str, content: str) -> str:
|
||
"""Write content to a .py or .txt file in the workspace."""
|
||
target = (WORKSPACE / path).resolve()
|
||
if not str(target).startswith(str(WORKSPACE.resolve())):
|
||
return "ERROR: path is outside the workspace."
|
||
if target.suffix not in (".py", ".txt"):
|
||
return f"ERROR: can only write .py and .txt files, got '{target.suffix}'."
|
||
target.parent.mkdir(parents=True, exist_ok=True)
|
||
target.write_text(content)
|
||
return f"OK: wrote {len(content)} chars to {path}."
|
||
|
||
|
||
@register_tool
|
||
def list_files(file_glob: str = "*") -> str:
|
||
"""List files in the workspace matching the glob pattern."""
|
||
found = sorted(WORKSPACE.glob(file_glob))
|
||
found = [f.relative_to(WORKSPACE) for f in found if f.is_file()]
|
||
if not found:
|
||
return f"No files matching '{file_glob}' in workspace."
|
||
return "\n".join(str(f) for f in found)
|
||
|
||
|
||
@register_tool
|
||
def grep_search(pattern: str, file_glob: str = "*.py") -> str:
|
||
"""Search for a pattern in workspace files and return matching lines with line numbers."""
|
||
matches = []
|
||
for filepath in sorted(WORKSPACE.glob(file_glob)):
|
||
if filepath.suffix not in (".py", ".txt"):
|
||
continue
|
||
try:
|
||
lines = filepath.read_text().splitlines()
|
||
except Exception:
|
||
continue
|
||
for i, line in enumerate(lines, 1):
|
||
if pattern in line:
|
||
rel = filepath.relative_to(WORKSPACE)
|
||
matches.append(f"{rel}:{i}: {line}")
|
||
if not matches:
|
||
return f"No matches for '{pattern}' in {file_glob}."
|
||
return "\n".join(matches)
|
||
|
||
|
||
@register_tool
|
||
def run_python(path: str) -> str:
|
||
"""Execute a Python file in the workspace and return stdout and stderr."""
|
||
target = (WORKSPACE / path).resolve()
|
||
if not str(target).startswith(str(WORKSPACE.resolve())):
|
||
return "ERROR: path is outside the workspace."
|
||
if not target.exists():
|
||
return f"ERROR: file '{path}' not found."
|
||
result = subprocess.run(
|
||
[sys.executable, str(target)],
|
||
capture_output=True, text=True, timeout=30,
|
||
cwd=str(WORKSPACE),
|
||
)
|
||
output = ""
|
||
if result.stdout:
|
||
output += f"STDOUT:\n{result.stdout}"
|
||
if result.stderr:
|
||
output += f"STDERR:\n{result.stderr}"
|
||
output += f"\nExit code: {result.returncode}"
|
||
return output.strip()
|
||
|
||
|
||
@register_tool
|
||
def validate_python(path: str) -> str:
|
||
"""Check whether a Python file has valid syntax using ast.parse."""
|
||
target = (WORKSPACE / path).resolve()
|
||
if not str(target).startswith(str(WORKSPACE.resolve())):
|
||
return "ERROR: path is outside the workspace."
|
||
if not target.exists():
|
||
return f"ERROR: file '{path}' not found."
|
||
source = target.read_text()
|
||
try:
|
||
ast.parse(source)
|
||
return "OK: syntax is valid."
|
||
except SyntaxError as e:
|
||
return f"SYNTAX ERROR: {e}"
|
||
|
||
|
||
@register_tool
|
||
def done(summary: str) -> str:
|
||
"""Signal that the agent has finished its task."""
|
||
return f"DONE: {summary}"
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# PART B – TOOL DISPATCHER
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
def build_tool_description() -> str:
|
||
"""Auto-generate tool descriptions from function signatures and docstrings."""
|
||
lines = []
|
||
for name, func in TOOL_REGISTRY.items():
|
||
sig = inspect.signature(func)
|
||
params = []
|
||
for pname, param in sig.parameters.items():
|
||
if param.default is inspect.Parameter.empty:
|
||
params.append(f'"{pname}": "<value>"')
|
||
else:
|
||
params.append(f'"{pname}": "<optional, default={param.default!r}>"')
|
||
param_str = ", ".join(params)
|
||
doc = (func.__doc__ or "").strip().split("\n")[0]
|
||
lines.append(f" - {name}({{{param_str}}}): {doc}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def dispatch_tool(tool_name: str, arguments: dict) -> str:
|
||
"""Call a tool by name with the given arguments."""
|
||
if tool_name not in TOOL_REGISTRY:
|
||
return f"ERROR: unknown tool '{tool_name}'. Available: {list(TOOL_REGISTRY.keys())}"
|
||
func = TOOL_REGISTRY[tool_name]
|
||
try:
|
||
return func(**arguments)
|
||
except TypeError as e:
|
||
return f"ERROR calling {tool_name}: {e}"
|
||
except Exception as e:
|
||
return f"ERROR in {tool_name}: {type(e).__name__}: {e}"
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# PART C – SYSTEM PROMPT
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
SYSTEM_PROMPT = f"""\
|
||
You are a coding agent that helps users with Python programming tasks.
|
||
You work inside a workspace directory and have access to tools.
|
||
|
||
<capabilities>
|
||
You can:
|
||
- Read .py and .txt files from the workspace
|
||
- Write .py and .txt files to the workspace
|
||
- List files in the workspace
|
||
- Search for patterns in files using grep
|
||
- Execute Python files and see their output
|
||
- Validate Python syntax using ast.parse
|
||
- Signal completion when the task is done
|
||
</capabilities>
|
||
|
||
<tools>
|
||
{build_tool_description()}
|
||
</tools>
|
||
|
||
<workflow>
|
||
For every user request, follow this workflow:
|
||
1. PLAN: Think about what steps are needed. List them in "thought".
|
||
2. ACT: Choose ONE tool to call for the current step.
|
||
3. OBSERVE: You will receive the tool output. Analyse it carefully.
|
||
4. REPLAN: If the result was unexpected, revise your plan in "thought".
|
||
5. REPEAT: Go back to step 2 if more work is needed.
|
||
6. DONE: Call the "done" tool when the task is complete.
|
||
</workflow>
|
||
|
||
<response_format>
|
||
You MUST respond with a JSON object every time:
|
||
{{
|
||
"thought": "<your reasoning about what to do next>",
|
||
"tool": "<tool name>",
|
||
"arguments": {{ <tool arguments> }}
|
||
}}
|
||
|
||
Example:
|
||
{{
|
||
"thought": "I need to write fibonacci.py first, then validate and run it.",
|
||
"tool": "write_file",
|
||
"arguments": {{"path": "fibonacci.py", "content": "def fib(n): ..."}}
|
||
}}
|
||
</response_format>
|
||
|
||
<rules>
|
||
- Always plan before acting.
|
||
- Call exactly ONE tool per response.
|
||
- After writing code, ALWAYS validate it with validate_python.
|
||
- After validation passes, run it with run_python to verify correctness.
|
||
- If an error occurs, analyse it and try to fix it (up to 3 retries).
|
||
- Stay within the workspace directory.
|
||
- When the task is fully complete, call the "done" tool.
|
||
- If you receive a <human_message>, acknowledge it and adjust your plan.
|
||
- If you receive a <replan> tag, revise your plan before choosing the next tool.
|
||
</rules>
|
||
"""
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# PART D – CODING AGENT CLASS
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
def truncate_result(result: str) -> str:
|
||
"""Truncate a tool result that exceeds MAX_RESULT_LENGTH."""
|
||
if len(result) <= MAX_RESULT_LENGTH:
|
||
return result
|
||
half = MAX_RESULT_LENGTH // 2
|
||
return (
|
||
result[:half]
|
||
+ f"\n\n... [TRUNCATED – {len(result)} chars total] ...\n\n"
|
||
+ result[-half:]
|
||
)
|
||
|
||
|
||
def trim_messages(messages: list) -> list:
|
||
"""Drop old messages when history exceeds MAX_HISTORY_CHARS.
|
||
Always keeps the system prompt (index 0) and original task (index 1).
|
||
"""
|
||
total = sum(len(m["content"]) for m in messages)
|
||
if total <= MAX_HISTORY_CHARS:
|
||
return messages
|
||
|
||
head = messages[:2]
|
||
tail = messages[2:]
|
||
original_task = messages[1]["content"] if len(messages) > 1 else ""
|
||
|
||
while tail and sum(len(m["content"]) for m in head + tail) > MAX_HISTORY_CHARS:
|
||
tail.pop(0)
|
||
|
||
reminder = {
|
||
"role": "user",
|
||
"content": (
|
||
"<system_note>Earlier conversation history was trimmed to fit the context window. "
|
||
f"REMINDER – your original task was:\n{original_task}\n"
|
||
"Continue working towards completing this task. "
|
||
"Do NOT start over or redo work already completed.</system_note>"
|
||
),
|
||
}
|
||
return head + [reminder] + tail
|
||
|
||
|
||
def _strip_code_fences(text: str) -> str:
|
||
"""Remove markdown code fences (```json ... ```) from a string."""
|
||
text = text.strip()
|
||
if text.startswith("```"):
|
||
lines = text.split("\n")
|
||
end = -1 if lines[-1].strip() == "```" else len(lines)
|
||
text = "\n".join(lines[1:end])
|
||
return text.strip()
|
||
|
||
|
||
class CodingAgent:
|
||
"""Step-by-step coding agent for use in Streamlit.
|
||
|
||
Usage:
|
||
agent = CodingAgent()
|
||
agent.start_task("Write a fibonacci function")
|
||
|
||
action = agent.propose_next_action()
|
||
# → {"thought": "...", "tool": "write_file", "arguments": {...}}
|
||
# Show action in UI, wait for user input
|
||
|
||
result = agent.approve()
|
||
# → {"tool": "write_file", "result": "OK: ...", "is_done": False}
|
||
|
||
agent.reject("Please use recursion instead")
|
||
# → agent replans on next propose_next_action()
|
||
"""
|
||
|
||
def __init__(self):
|
||
self.messages: list = []
|
||
self.pending_action: dict | None = None
|
||
self.is_done: bool = False
|
||
self.iteration: int = 0
|
||
self._setup_api()
|
||
|
||
def _setup_api(self):
|
||
"""Load API configuration from environment variables."""
|
||
self.api_url = f"http://{os.getenv('HOST')}:{os.getenv('PORT')}/v1/chat/completions"
|
||
self.api_key = os.getenv("API_KEY")
|
||
self.model = os.getenv("MODEL")
|
||
|
||
def _call_api(self, messages: list) -> str:
|
||
"""Make a raw API call and return the response content string."""
|
||
headers = {"Content-Type": "application/json"}
|
||
if self.api_key and self.api_key != "EMPTY":
|
||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||
|
||
payload = {
|
||
"model": self.model,
|
||
"messages": messages,
|
||
"temperature": 0.2,
|
||
"max_tokens": 4096,
|
||
"stream": False,
|
||
}
|
||
|
||
response = requests.post(self.api_url, headers=headers, json=payload, timeout=60)
|
||
|
||
if response.status_code != 200:
|
||
raise Exception(f"API Error {response.status_code}: {response.text}")
|
||
|
||
data = response.json()
|
||
if "choices" in data and len(data["choices"]) > 0:
|
||
return data["choices"][0]["message"]["content"]
|
||
raise Exception("Invalid API response format")
|
||
|
||
# ── Public interface ──────────────────────────────────────────────────────
|
||
|
||
def start_task(self, task: str) -> None:
|
||
"""Initialise the agent with a new task. Resets all state."""
|
||
self.messages = [
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{"role": "user", "content": task},
|
||
]
|
||
self.pending_action = None
|
||
self.is_done = False
|
||
self.iteration = 0
|
||
|
||
def propose_next_action(self) -> dict:
|
||
"""Ask the LLM what to do next.
|
||
|
||
Returns the parsed action dict without executing anything.
|
||
The action is stored internally as pending_action until
|
||
approve() or reject() is called.
|
||
|
||
Returns:
|
||
{"thought": str, "tool": str, "arguments": dict}
|
||
"""
|
||
if self.is_done:
|
||
return {"thought": "Task already completed.", "tool": "done", "arguments": {}}
|
||
|
||
if self.iteration >= MAX_ITERATIONS:
|
||
return {"thought": "Max iterations reached.", "tool": "done",
|
||
"arguments": {"summary": "Stopped: max iterations reached."}}
|
||
|
||
self.iteration += 1
|
||
self.messages = trim_messages(self.messages)
|
||
|
||
try:
|
||
raw = self._call_api(self.messages)
|
||
raw = _strip_code_fences(raw)
|
||
action = json.loads(raw)
|
||
except json.JSONDecodeError:
|
||
action = {
|
||
"thought": "Could not parse LLM response as JSON.",
|
||
"tool": "done",
|
||
"arguments": {"summary": "Stopped: JSON parse error."},
|
||
}
|
||
raw = json.dumps(action)
|
||
except Exception as e:
|
||
action = {
|
||
"thought": f"API call failed: {e}",
|
||
"tool": "done",
|
||
"arguments": {"summary": f"Stopped: {e}"},
|
||
}
|
||
raw = json.dumps(action)
|
||
|
||
self.pending_action = {"raw": raw, "action": action}
|
||
return action
|
||
|
||
def approve(self) -> dict:
|
||
"""Execute the pending action and return the result.
|
||
|
||
Returns:
|
||
{"tool": str, "arguments": dict, "result": str, "is_done": bool}
|
||
"""
|
||
if not self.pending_action:
|
||
raise Exception("No pending action. Call propose_next_action() first.")
|
||
|
||
raw = self.pending_action["raw"]
|
||
action = self.pending_action["action"]
|
||
tool_name = action.get("tool", "")
|
||
arguments = action.get("arguments", {})
|
||
|
||
# Append the assistant message to history
|
||
self.messages.append({"role": "assistant", "content": raw})
|
||
self.pending_action = None
|
||
|
||
# Handle completion
|
||
if tool_name == "done":
|
||
self.is_done = True
|
||
return {
|
||
"tool": "done",
|
||
"arguments": arguments,
|
||
"result": arguments.get("summary", "Task completed."),
|
||
"is_done": True,
|
||
}
|
||
|
||
# Execute the tool
|
||
result = dispatch_tool(tool_name, arguments)
|
||
result = truncate_result(result)
|
||
|
||
# Build feedback – nudge agent to replan on errors
|
||
feedback = f'<tool_result tool="{tool_name}">\n{result}\n</tool_result>'
|
||
if result.startswith("ERROR") or result.startswith("SYNTAX ERROR"):
|
||
feedback += (
|
||
"\n\n<replan>The tool returned an error. "
|
||
"Re-examine your plan: what went wrong and what should you do differently? "
|
||
"State your revised plan in your next thought.</replan>"
|
||
)
|
||
|
||
self.messages.append({"role": "user", "content": feedback})
|
||
|
||
return {
|
||
"tool": tool_name,
|
||
"arguments": arguments,
|
||
"result": result,
|
||
"is_done": False,
|
||
}
|
||
|
||
def follow_up(self, question: str) -> None:
|
||
"""Continue a completed task with a follow-up question.
|
||
|
||
Resets the done-flag and appends the user's question to the existing
|
||
conversation, so the agent retains full context of what was already done.
|
||
Call propose_next_action() afterwards to continue the loop.
|
||
"""
|
||
self.is_done = False
|
||
self.messages.append({
|
||
"role": "user",
|
||
"content": (
|
||
f"<human_message>{question}</human_message>\n"
|
||
"<replan>The user has a follow-up question or correction regarding "
|
||
"the task you just completed. Review what you already did and "
|
||
"address their question accordingly.</replan>"
|
||
),
|
||
})
|
||
|
||
def reject(self, feedback: str) -> None:
|
||
"""Reject the pending action and inject user feedback.
|
||
|
||
The pending action is NOT executed. On the next call to
|
||
propose_next_action() the agent will replan based on the feedback.
|
||
"""
|
||
if not self.pending_action:
|
||
return
|
||
|
||
raw = self.pending_action["raw"]
|
||
action = self.pending_action["action"]
|
||
tool_name = action.get("tool", "unknown")
|
||
|
||
self.messages.append({"role": "assistant", "content": raw})
|
||
self.messages.append({
|
||
"role": "user",
|
||
"content": (
|
||
f"<human_message>{feedback}</human_message>\n"
|
||
"<replan>The user has given you guidance BEFORE you executed "
|
||
f"your proposed action ({tool_name}). Do NOT proceed with that action. "
|
||
"Revise your plan to incorporate their feedback and state your "
|
||
"updated plan in your next thought.</replan>"
|
||
),
|
||
})
|
||
self.pending_action = None
|