316 lines
11 KiB
Python
316 lines
11 KiB
Python
"""
|
||
Code Agent - Handles file operations, code execution, and testing
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import requests
|
||
import subprocess
|
||
import webbrowser
|
||
import time
|
||
from pathlib import Path
|
||
from datetime import datetime
|
||
from typing import Dict, Any
|
||
from dotenv import load_dotenv
|
||
|
||
load_dotenv()
|
||
|
||
# ============================================================================
|
||
# CONFIGURATION
|
||
# ============================================================================
|
||
|
||
LLM_SERVER = os.getenv("LLM_SERVER_URL", "http://localhost:8000")
|
||
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
|
||
LLM_FAST = os.getenv("LLM_FAST_MODEL", "phi")
|
||
LLM_POWER = os.getenv("LLM_POWER_MODEL", "deepseek-coder:33b")
|
||
|
||
# ============================================================================
|
||
# CODE AGENT CLASS
|
||
# ============================================================================
|
||
|
||
class CodeAgent:
|
||
"""Agent that modifies files, executes code, and tests in browser"""
|
||
|
||
def __init__(self, work_dir: str = "./workspace"):
|
||
self.work_dir = Path(work_dir)
|
||
self.work_dir.mkdir(exist_ok=True)
|
||
print(f"📁 Workspace: {self.work_dir.absolute()}")
|
||
|
||
# ========================================================================
|
||
# FILE OPERATIONS
|
||
# ========================================================================
|
||
|
||
def read_file(self, file_path: str) -> str:
|
||
"""Read file from workspace"""
|
||
file = self.work_dir / file_path
|
||
|
||
if not file.exists():
|
||
return f"❌ File not found: {file_path}"
|
||
|
||
try:
|
||
with open(file, 'r', encoding='utf-8') as f:
|
||
content = f.read()
|
||
print(f"✅ Read: {file_path} ({len(content)} bytes)")
|
||
return content
|
||
except Exception as e:
|
||
return f"❌ Error reading {file_path}: {e}"
|
||
|
||
def write_file(self, file_path: str, content: str) -> str:
|
||
"""Write/modify file in workspace"""
|
||
file = self.work_dir / file_path
|
||
file.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
try:
|
||
with open(file, 'w', encoding='utf-8') as f:
|
||
f.write(content)
|
||
print(f"✅ Wrote: {file_path}")
|
||
return f"File saved: {file_path}"
|
||
except Exception as e:
|
||
return f"❌ Error writing {file_path}: {e}"
|
||
|
||
def delete_file(self, file_path: str) -> str:
|
||
"""Delete file from workspace"""
|
||
file = self.work_dir / file_path
|
||
|
||
if not file.exists():
|
||
return f"❌ File not found: {file_path}"
|
||
|
||
try:
|
||
file.unlink()
|
||
return f"✅ Deleted: {file_path}"
|
||
except Exception as e:
|
||
return f"❌ Error deleting {file_path}: {e}"
|
||
|
||
def list_files(self) -> str:
|
||
"""List all files in workspace"""
|
||
files = list(self.work_dir.rglob("*"))
|
||
file_list = "\n".join([
|
||
f" {f.relative_to(self.work_dir)}"
|
||
for f in files if f.is_file()
|
||
])
|
||
return f"📁 Files in workspace:\n{file_list}" if file_list else "📁 No files"
|
||
|
||
# ========================================================================
|
||
# CODE EXECUTION
|
||
# ========================================================================
|
||
|
||
def execute_python(self, code: str, file_name: str = "exec.py") -> str:
|
||
"""Execute Python code"""
|
||
script_path = self.work_dir / file_name
|
||
|
||
try:
|
||
# Save script
|
||
with open(script_path, 'w', encoding='utf-8') as f:
|
||
f.write(code)
|
||
|
||
# Execute
|
||
result = subprocess.run(
|
||
["python3", str(script_path)],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=30,
|
||
cwd=self.work_dir
|
||
)
|
||
|
||
output = result.stdout + result.stderr
|
||
print(f"🐍 Python executed: {file_name}")
|
||
return output[:1000] # Limit output
|
||
|
||
except subprocess.TimeoutExpired:
|
||
return "❌ Timeout: Script took more than 30 seconds"
|
||
except Exception as e:
|
||
return f"❌ Error: {e}"
|
||
|
||
def execute_node(self, code: str, file_name: str = "exec.js") -> str:
|
||
"""Execute Node.js code"""
|
||
script_path = self.work_dir / file_name
|
||
|
||
try:
|
||
with open(script_path, 'w', encoding='utf-8') as f:
|
||
f.write(code)
|
||
|
||
result = subprocess.run(
|
||
["node", str(script_path)],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=30,
|
||
cwd=self.work_dir
|
||
)
|
||
|
||
output = result.stdout + result.stderr
|
||
print(f"📟 Node.js executed: {file_name}")
|
||
return output[:1000]
|
||
|
||
except FileNotFoundError:
|
||
return "❌ Node.js not installed"
|
||
except Exception as e:
|
||
return f"❌ Error: {e}"
|
||
|
||
def execute_shell(self, command: str) -> str:
|
||
"""Execute shell command (whitelist safe commands only)"""
|
||
allowed_commands = [
|
||
'ls', 'pwd', 'mkdir', 'rm', 'cp', 'mv',
|
||
'cat', 'grep', 'find', 'du', 'df',
|
||
'git', 'npm', 'pip', 'python'
|
||
]
|
||
|
||
cmd_name = command.split()[0] if command else ""
|
||
if cmd_name not in allowed_commands:
|
||
return f"❌ Command not allowed: {cmd_name}"
|
||
|
||
try:
|
||
result = subprocess.run(
|
||
command.split(),
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=30,
|
||
cwd=self.work_dir
|
||
)
|
||
output = result.stdout + result.stderr
|
||
return output[:500]
|
||
except Exception as e:
|
||
return f"❌ Error: {e}"
|
||
|
||
# ========================================================================
|
||
# TESTING & BROWSER
|
||
# ========================================================================
|
||
|
||
def test_html(self, html_file: str) -> str:
|
||
"""Open HTML file in browser for testing"""
|
||
file_path = self.work_dir / html_file
|
||
|
||
if not file_path.exists():
|
||
return f"❌ File not found: {html_file}"
|
||
|
||
try:
|
||
url = f"file:///{file_path.absolute()}".replace("\\", "/")
|
||
webbrowser.open(url)
|
||
print(f"🌐 Opened in browser: {html_file}")
|
||
return f"✅ HTML opened: {html_file}"
|
||
except Exception as e:
|
||
return f"❌ Error: {e}"
|
||
|
||
def open_browser(self, url: str) -> str:
|
||
"""Open URL in browser"""
|
||
try:
|
||
webbrowser.open(url)
|
||
return f"✅ Browser opened: {url}"
|
||
except Exception as e:
|
||
return f"❌ Error: {e}"
|
||
|
||
# ========================================================================
|
||
# GIT OPERATIONS
|
||
# ========================================================================
|
||
|
||
def git_commit(self, message: str = None) -> str:
|
||
"""Auto-commit changes"""
|
||
try:
|
||
if not message:
|
||
message = f"Auto-commit from agent - {datetime.now().isoformat()}"
|
||
|
||
# Git add
|
||
result = subprocess.run(
|
||
['git', 'add', '.'],
|
||
cwd=self.work_dir,
|
||
capture_output=True,
|
||
text=True
|
||
)
|
||
|
||
if result.returncode != 0:
|
||
return f"❌ Git add failed: {result.stderr}"
|
||
|
||
# Git commit
|
||
result = subprocess.run(
|
||
['git', 'commit', '-m', message],
|
||
cwd=self.work_dir,
|
||
capture_output=True,
|
||
text=True
|
||
)
|
||
|
||
if result.returncode != 0:
|
||
return f"ℹ️ Nothing to commit"
|
||
|
||
return f"✅ Committed: {message}"
|
||
except Exception as e:
|
||
return f"❌ Error: {e}"
|
||
|
||
def git_push(self) -> str:
|
||
"""Push to remote repository"""
|
||
try:
|
||
result = subprocess.run(
|
||
['git', 'push'],
|
||
cwd=self.work_dir,
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=30
|
||
)
|
||
|
||
if result.returncode == 0:
|
||
return "✅ Pushed to remote"
|
||
else:
|
||
return f"❌ Push failed: {result.stderr}"
|
||
except Exception as e:
|
||
return f"❌ Error: {e}"
|
||
|
||
# ========================================================================
|
||
# LLM INTERACTION
|
||
# ========================================================================
|
||
|
||
def call_llm(self, prompt: str, model: str = "phi") -> str:
|
||
"""Call LLM server"""
|
||
try:
|
||
response = requests.post(
|
||
f"{OLLAMA_URL}/api/generate",
|
||
json={
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"stream": False
|
||
},
|
||
timeout=300
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
return response.json()["response"]
|
||
else:
|
||
return f"❌ LLM Error: {response.status_code}"
|
||
|
||
except requests.exceptions.ConnectionError:
|
||
return f"❌ Cannot connect to {OLLAMA_URL}"
|
||
except Exception as e:
|
||
return f"❌ Error: {e}"
|
||
|
||
# ============================================================================
|
||
# GET AGENT TOOLS
|
||
# ============================================================================
|
||
|
||
def get_agent_tools(agent: CodeAgent) -> Dict[str, Any]:
|
||
"""Return dictionary of available tools"""
|
||
return {
|
||
"read_file": agent.read_file,
|
||
"write_file": agent.write_file,
|
||
"delete_file": agent.delete_file,
|
||
"execute_python": agent.execute_python,
|
||
"execute_node": agent.execute_node,
|
||
"execute_shell": agent.execute_shell,
|
||
"test_html": agent.test_html,
|
||
"open_browser": agent.open_browser,
|
||
"git_commit": agent.git_commit,
|
||
"git_push": agent.git_push,
|
||
"list_files": agent.list_files,
|
||
"call_llm": agent.call_llm,
|
||
}
|
||
|
||
# ============================================================================
|
||
# MAIN (for testing)
|
||
# ============================================================================
|
||
|
||
if __name__ == "__main__":
|
||
agent = CodeAgent()
|
||
|
||
# Test
|
||
print("\n🧪 Testing Agent...")
|
||
print(agent.list_files())
|
||
|
||
# Create test file
|
||
agent.write_file("test.txt", "Hello from agent!")
|
||
print(agent.read_file("test.txt")) |