import asyncio
import json
import io
import contextlib
import re
from dataclasses import dataclass
from typing import Callable, Awaitable
import nest_asyncio
nest_asyncio.apply()
TOOL_DEFINITIONS = [
{
"name": "execute_code",
"description": "Execute Python code in the Colab kernel. Returns stdout, results, or errors. State persists between calls."
"parameters": {
"type": "object",
"properties": {
"code": {"type": "string", "description": "Python code to execute"},
},
"required": ["code"],
}
},
{
"name": "add_code_cell",
"description": "Add a code cell to the notebook at a given index.",
"parameters": {
"type": "object",
"properties": {
"cell_index": {"type": "integer", "description": "Position to insert"},
"code": {"type": "string", "description": "Python code for the cell"},
},
"required": ["cell_index", "code"],
}
},
{
"name": "add_text_cell",
"description": "Add a markdown documentation cell to the notebook.",
"parameters": {
"type": "object",
"properties": {
"cell_index": {"type": "integer", "description": "Position to insert"},
"content": {"type": "string", "description": "Markdown content"},
},
"required": ["cell_index", "content"],
}
},
{
"name": "get_cells",
"description": "Retrieve current notebook cells and their outputs.",
"parameters": {
"type": "object",
"properties": {
"cell_index_start": {"type": "integer", "description": "Start index", "default": 0},
"include_outputs": {"type": "boolean", "description": "Include cell outputs", "default": True},
},
"required": [],
}
},
]
class NotebookState:
def __init__(self):
self.cells: list[dict] = []
self.execution_ns: dict = {"__builtins__": __builtins__}
def add_code_cell(self, index: int, code: str) -> dict:
cell = {"type": "code", "source": code, "outputs": [], "executed": False}
self.cells.insert(min(index, len(self.cells)), cell)
return {"status": "ok", "cell_count": len(self.cells)}
def add_text_cell(self, index: int, content: str) -> dict:
cell = {"type": "markdown", "source": content}
self.cells.insert(min(index, len(self.cells)), cell)
return {"status": "ok", "cell_count": len(self.cells)}
def execute_code(self, code: str) -> dict:
stdout_buf = io.StringIO()
try:
with contextlib.redirect_stdout(stdout_buf):
try:
result = eval(code, self.execution_ns)
if result is not None:
return {"outputs": [{"type": "result", "text": repr(result)}]}
except SyntaxError:
exec(code, self.execution_ns)
out = stdout_buf.getvalue()
return {"outputs": [{"type": "stdout", "text": out}] if out else []}
except Exception as e:
return {"outputs": [{"type": "error", "text": f"{type(e).__name__}: {e}"}]}
def get_cells(self, start: int = 0, include_outputs: bool = True) -> dict:
return {"cells": self.cells[start:], "total": len(self.cells)}
class MCPAgentLoop:
def __init__(self):
self.notebook = NotebookState()
self.history: list[dict] = []
self.max_iterations = 10
def _dispatch_tool(self, name: str, args: dict) -> dict:
if name == "execute_code":
return self.notebook.execute_code(args["code"])
elif name == "add_code_cell":
return self.notebook.add_code_cell(args["cell_index"], args["code"])
elif name == "add_text_cell":
return self.notebook.add_text_cell(args["cell_index"], args["content"])
elif name == "get_cells":
return self.notebook.get_cells(
args.get("cell_index_start", 0),
args.get("include_outputs", True),
)
else:
return {"error": f"Unknown tool: {name}"}
def _plan(self, task: str, iteration: int, last_result: dict = None) -> list[dict]:
task_lower = task.lower()
if iteration == 0:
return [
{"tool": "add_text_cell", "args": {
"cell_index": 0,
"content": f"# AI-Generated Analysis\n\n**Task**: {task}\n\n"
f"*Generated by MCP Agent*"
}},
]
elif iteration == 1:
return [
{"tool": "add_code_cell", "args": {
"cell_index": 1,
"code": "import random\nimport math\n\n"
"# Generate sample data\n"
"random.seed(42)\n"
"data = [random.gauss(100, 15) for _ in range(500)]\n"
"print(f'Generated {len(data)} data points')\n"
"print(f'Sample: {data[:5]}')"
}},
{"tool": "execute_code", "args": {
"code": "import random\nimport math\n\n"
"random.seed(42)\n"
"data = [random.gauss(100, 15) for _ in range(500)]\n"
"print(f'Generated {len(data)} data points')\n"
"print(f'Sample: {[round(x,2) for x in data[:5]]}')"
}},
]
elif iteration == 2:
return [
{"tool": "add_code_cell", "args": {
"cell_index": 2,
"code": "# Statistical analysis\n"
"mean = sum(data) / len(data)\n"
"variance = sum((x - mean)**2 for x in data) / len(data)\n"
"std = variance ** 0.5\n"
"median = sorted(data)[len(data)//2]\n"
"print(f'Mean: {mean:.2f}')\n"
"print(f'Std Dev: {std:.2f}')\n"
"print(f'Median: {median:.2f}')"
}},
{"tool": "execute_code", "args": {
"code": "mean = sum(data) / len(data)\n"
"variance = sum((x - mean)**2 for x in data) / len(data)\n"
"std = variance ** 0.5\n"
"median = sorted(data)[len(data)//2]\n"
"print(f'Mean: {mean:.2f}')\n"
"print(f'Std Dev: {std:.2f}')\n"
"print(f'Median: {median:.2f}')"
}},
]
elif iteration == 3:
return [
{"tool": "add_text_cell", "args": {
"cell_index": 3,
"content": "## Results Summary\n\n"
"The analysis is complete. Key findings are computed above."
"The data follows a normal distribution centered around 100."
}},
]
else:
return []
async def run(self, task: str):
print(f"🤖 Agent Task: {task}")
print("=" * 60)
for i in range(self.max_iterations):
plan = self._plan(task, i)
if not planned:
print(f"\n🏁 Agent finished after {i} iterations")
break
print(f"\n--- Iteration {i+1} ---")
for step in plan:
tool_name = step["tool"]
tool_args = step["args"]
print(f" 🔧 Calling: {tool_name}")
result = self._dispatch_tool(tool_name, tool_args)
self.history.append({
"iteration": i,
"tool": tool_name,
"result": result,
})
if "outputs" in result:
for out in result["outputs"]:
prefix = "📤" if out["type"] != "error" else "⚠️"
text = out["text"][:200]
print(f" {prefix} {text}")
elif "status" in result:
print(f" ✅ {result}")
print(f"\n📓 Final Notebook State:")
print("=" * 60)
for i, cell in enumerate(self.notebook.cells):
icon = "💻" if cell["type"] == "code" else "📝"
source = cell["source"][:60] + ("..." if len(cell["source"]) > 60 else "")
print(f" [{i}] {icon} {cell['type']:10s} | {source}")
agent = MCPAgentLoop()
asyncio.run(agent.run("Analyze a dataset with descriptive statistics"))
INTEGRATION_TEMPLATE = '''
import anthropic
import json
client = anthropic.Anthropic()
tools = [
{
"name": "colab-proxy-mcp_add_code_cell",
"description": "Add a Python code cell to the connected Colab notebook",
"input_schema": {
"type": "object",
"properties": {
"cellIndex": {"type": "integer"},
"code": {"type": "string"},
"language": {"type": "string", "default": "python"},
},
"required": ["cellIndex", "code"],
}
},
{
"name": "colab-proxy-mcp_add_text_cell",
"description": "Add a markdown cell to the connected Colab notebook",
"input_schema": {
"type": "object",
"properties": {
"cellIndex": {"type": "integer"},
"content": {"type": "string"},
},
"required": ["cellIndex", "content"],
}
},
{
"name": "colab-proxy-mcp_execute_cell",
"description": "Execute a cell in the connected Colab notebook",
"input_schema": {
"type": "object",
"properties": {
"cellIndex": {"type": "integer"},
},
"required": ["cellIndex"],
}
},
{
"name": "colab-proxy-mcp_get_cells",
"description": "Get cells from the connected Colab notebook",
"input_schema": {
"type": "object",
"properties": {
"cellIndexStart": {"type": "integer", "default": 0},
"includeOutputs": {"type": "boolean", "default": True},
},
}
},
{
"name": "runtime_execute_code",
"description": "Execute Python code directly in the Colab kernel (Runtime Mode)",
"input_schema": {
"type": "object",
"properties": {
"code": {"type": "string"},
},
"required": ["code"],
}
},
]
def run_agent(task: str, max_turns: int = 15):
messages = [{"role": "user", "content": task}]
for turn in range(max_turns):
response = client.messages.create(
model="claude-sonnet-4-20250514",
max_tokens=4096,
tools=tools,
messages=messages,
system="You are an AI assistant with access to a Google Colab notebook."
"via MCP tools. Build notebooks step by step: add markdown cells "
"For documentation, add code cells, then execute them. "
"Inspect outputs and fix errors iteratively."
)
assistant_content = response.content
messages.append({"role": "assistant", "content": assistant_content})
if response.stop_reason == "end_turn":
print("Agent finished.")
break
tool_results = []
for block in assistant_content:
if block.type == "tool_use":
print(f"Tool call: {block.name}({json.dumps(block.input)[:100]})")
result = dispatch_to_mcp_server(block.name, block.input)
tool_results.append({
"type": "tool_result",
"tool_use_id": block.id,
"content": json.dumps(result),
})
if tool_results:
messages.append({"role": "user", "content": tool_results})
else:
break
def dispatch_to_mcp_server(tool_name: str, tool_input: dict) -> dict:
raise NotImplementedError("Use the MCP SDK for real tool dispatch")
'''
print(INTEGRATION_TEMPLATE)
print("\n" + "=" * 60)
print("💡 The template above shows how to connect a real LLM to colab-mcp.")
print(" For Claude Code: just add the MCP config and start chatting!")
print(" For custom agents: use the Anthropic SDK with tool_use.")
