544 lines
20 KiB
Python
544 lines
20 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 json
|
||
import os
|
||
import re
|
||
from pathlib import Path
|
||
import asyncio
|
||
import pprint
|
||
|
||
import requests
|
||
from dotenv import load_dotenv
|
||
#from mcp_server_adapter import MCPToolAdapter # Import from current directory for easier testing without package structure
|
||
from backend.agent.mcp_server_adapter import MCPToolAdapter
|
||
|
||
# ── mcp server initialization ────────────────────────────────────────────────────────────────
|
||
adapter = MCPToolAdapter()
|
||
print("MCPToolAdapter created. Listing all tools from servers...")
|
||
asyncio.run(adapter.initialize_all_servers())
|
||
print("listed tools from all servers")
|
||
|
||
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
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# Tool dispatching and result handling
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
|
||
def build_all_tool_description() -> str:
|
||
"""Get relevant tools from the MCP servers based on the query."""
|
||
|
||
all_tools = adapter.get_all_tools()
|
||
print(f"Building tool description for {len(all_tools)} tools.")
|
||
|
||
descriptions = []
|
||
for tool in all_tools:
|
||
pprint.pprint(f"{tool}")
|
||
descriptions.append(f"- {tool['tool_name']}: {tool['tool_description']}")
|
||
|
||
return "\n".join(descriptions)
|
||
|
||
async def dispatch_tool(tool_name: str, arguments: dict) -> str:
|
||
"""Call a tool by name with the given arguments using the MCP adapter."""
|
||
|
||
if tool_name == "done":
|
||
# Handle the "done" tool locally since it's not an MCP tool
|
||
summary = arguments.get("summary", "Task completed.")
|
||
return f"DONE: {summary}"
|
||
|
||
try:
|
||
print(f"Trying to call tool '{tool_name}' in dispatch_tool through MCPToolAdapter...")
|
||
result = await adapter.call_tool(tool_name, arguments)
|
||
|
||
print(f"Raw result from tool '{tool_name}': {result}")
|
||
|
||
if result.isError:
|
||
texts = [block.text for block in result.content if block.type == "text"]
|
||
return f"Tool error: {' '.join(texts)}"
|
||
|
||
texts = [block.text for block in result.content if block.type == "text"]
|
||
return "\n".join(texts)
|
||
|
||
except Exception as e:
|
||
return f"Error calling tool '{tool_name}': {e}"
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# 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.
|
||
You solve tasks by interacting with a workspace through a
|
||
dynamic set of tools via the Model Context Protocol (MCP).
|
||
|
||
<capabilities>
|
||
You have access to a workspace where you can manage files,
|
||
analyze code, and execute Python scripts.
|
||
You can call tools to interact with the workspace and get feedback.
|
||
You can write and read files, list directory contents, search for patterns,
|
||
validate Python syntax, and run Python code.
|
||
You can access web search and page fetching tools to gather information from the internet.
|
||
You can use these capabilities to iteratively work towards completing the user's task.
|
||
</capabilities>
|
||
|
||
<tools>
|
||
{build_all_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>
|
||
"""
|
||
|
||
|
||
# ═════════════════════════════════════════════════════════════════════════════
|
||
# 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 ""
|
||
|
||
# Drop the oldest messages first (index 2 onwards) until we are under the limit.
|
||
# The system prompt (0) and original task (1) are never dropped.
|
||
while tail and sum(len(m["content"]) for m in head + tail) > MAX_HISTORY_CHARS:
|
||
tail.pop(0)
|
||
|
||
# Inject a reminder so the agent doesn't lose track of its goal after trimming.
|
||
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 _repair_json_strings(text: str) -> str:
|
||
"""
|
||
Replace unescaped control characters (newline, tab, carriage return)
|
||
inside JSON string values with their proper escape sequences.
|
||
|
||
LLMs frequently emit literal newlines inside long string values, which
|
||
is invalid JSON. This function fixes that without touching structural
|
||
whitespace outside strings.
|
||
"""
|
||
result: list[str] = []
|
||
in_string = False
|
||
escape = False
|
||
_escapes = {'\n': '\\n', '\r': '\\r', '\t': '\\t'}
|
||
for ch in text:
|
||
if escape:
|
||
result.append(ch)
|
||
escape = False
|
||
continue
|
||
if ch == '\\' and in_string:
|
||
result.append(ch)
|
||
escape = True
|
||
continue
|
||
if ch == '"':
|
||
in_string = not in_string
|
||
result.append(ch)
|
||
continue
|
||
if in_string and ch in _escapes:
|
||
result.append(_escapes[ch])
|
||
continue
|
||
result.append(ch)
|
||
return ''.join(result)
|
||
|
||
|
||
def extract_json(text: str) -> str:
|
||
"""
|
||
Extract and repair a JSON object or array from an LLM response that may
|
||
contain extra prose, markdown code fences, or unescaped control characters.
|
||
|
||
Strategy:
|
||
1. Strip markdown ```json ... ``` or ``` ... ``` fences.
|
||
2. Find the first '{' or '[' and extract to the matching closing bracket.
|
||
3. Repair unescaped newlines/tabs inside string values.
|
||
|
||
Returns the cleaned JSON string, or the original text as a fallback
|
||
(so json.loads can raise a meaningful error with context).
|
||
"""
|
||
|
||
if text is None:
|
||
return ""
|
||
|
||
# 1. Strip markdown fences
|
||
fenced = re.sub(r"```(?:json)?\s*([\s\S]*?)\s*```", r"\1", text.strip())
|
||
if fenced != text.strip():
|
||
return _repair_json_strings(fenced.strip())
|
||
|
||
# 2. Find first JSON container and extract to matching close
|
||
extracted = text
|
||
for start_char, end_char in [('{', '}'), ('[', ']')]:
|
||
idx = text.find(start_char)
|
||
if idx == -1:
|
||
continue
|
||
depth = 0
|
||
in_string = False
|
||
escape = False
|
||
for i, ch in enumerate(text[idx:], start=idx):
|
||
if escape:
|
||
escape = False
|
||
continue
|
||
if ch == '\\' and in_string:
|
||
escape = True
|
||
continue
|
||
if ch == '"':
|
||
in_string = not in_string
|
||
continue
|
||
if in_string:
|
||
continue
|
||
if ch == start_char:
|
||
depth += 1
|
||
elif ch == end_char:
|
||
depth -= 1
|
||
if depth == 0:
|
||
extracted = text[idx: i + 1]
|
||
break
|
||
break
|
||
|
||
# 3. Repair unescaped control characters inside string values
|
||
return _repair_json_strings(extracted)
|
||
|
||
|
||
def _strip_code_fences(text: str) -> str:
|
||
"""Remove markdown code fences (```json ... ```) from a string."""
|
||
if text is None:
|
||
return ""
|
||
|
||
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")
|
||
|
||
#async def _call_api(self, messages: list) -> str:
|
||
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, # low temperature → deterministic, more reliable tool calls
|
||
"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
|
||
|
||
async 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)
|
||
cleaned = extract_json(raw)
|
||
action = json.loads(cleaned)
|
||
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
|
||
|
||
async 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 = await dispatch_tool(tool_name, arguments)
|
||
result = truncate_result(result)
|
||
|
||
# Wrap the tool output in an XML tag so the LLM can easily find it.
|
||
# Append a <replan> tag on errors to force the agent to reconsider
|
||
# its plan rather than blindly retrying the same failing action.
|
||
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
|
||
|
||
def main():
|
||
"""Example of how to use the CodingAgent in a simple loop."""
|
||
agent = CodingAgent()
|
||
task = "Write a Python function that returns the nth Fibonacci number."
|
||
agent.start_task(task)
|
||
|
||
if agent.pending_action:
|
||
print(f"Initial proposed action: {agent.pending_action['action']}")
|
||
|
||
|
||
while not agent.is_done:
|
||
action = asyncio.run(agent.propose_next_action())
|
||
print(f"Proposed action: {action}")
|
||
|
||
if action["tool"] == "done":
|
||
print("Task completed.")
|
||
break
|
||
else:
|
||
user_feedback = input("Approve this action? (y/n) ")
|
||
if user_feedback.lower() == "y":
|
||
result = asyncio.run(agent.approve())
|
||
print(f"Tool result: {result}")
|
||
elif user_feedback.lower() == "n":
|
||
feedback = input("Enter feedback for the agent: ")
|
||
agent.reject(feedback)
|
||
|
||
|
||
if result["is_done"]:
|
||
print("Task completed.")
|
||
break
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
|