May / utils.py
Mayank2027's picture
Create utils.py
3b9084c verified
Raw
History Blame Contribute Delete
5.67 kB
# File 4: utils.py
with open(f"{base_dir}/utils.py", "w") as f:
f.write('''"""Shared utilities for Swarm CLI"""
import json
import os
import hashlib
import datetime
from typing import Any, Dict, List, Optional
from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax
from rich.table import Table
from rich.tree import Tree
console = Console()
class Logger:
"""Activity logger"""
@staticmethod
def log(action: str, details: Dict[str, Any], agent: str = "system"):
"""Log an action"""
entry = {
"timestamp": datetime.datetime.now().isoformat(),
"agent": agent,
"action": action,
"details": details
}
# Print to console
console.print(f"[dim][{entry['timestamp']}] {agent}: {action}[/dim]")
# Append to log file
try:
with open("swarm.log", "a", encoding="utf-8") as f:
f.write(json.dumps(entry) + "\n")
except:
pass
return entry
class Memory:
"""Simple JSON-based memory store"""
def __init__(self, filepath: str = "memory.json"):
self.filepath = filepath
self.data = self._load()
def _load(self) -> Dict:
if os.path.exists(self.filepath):
try:
with open(self.filepath, "r", encoding="utf-8") as f:
return json.load(f)
except:
return {"facts": [], "conversations": [], "tools_used": []}
return {"facts": [], "conversations": [], "tools_used": []}
def save(self):
with open(self.filepath, "w", encoding="utf-8") as f:
json.dump(self.data, f, indent=2, ensure_ascii=False)
def add_fact(self, key: str, value: Any):
self.data["facts"].append({
"key": key,
"value": value,
"timestamp": datetime.datetime.now().isoformat()
})
self.save()
def get_facts(self, query: str = "") -> List[Dict]:
if not query:
return self.data["facts"][-20:]
return [f for f in self.data["facts"] if query.lower() in f["key"].lower()]
def add_conversation(self, role: str, content: str):
self.data["conversations"].append({
"role": role,
"content": content,
"timestamp": datetime.datetime.now().isoformat()
})
self.data["conversations"] = self.data["conversations"][-100:]
self.save()
def get_conversation_history(self, limit: int = 10) -> List[Dict]:
return self.data["conversations"][-limit:]
def record_tool_use(self, tool_name: str, success: bool, details: str = ""):
self.data["tools_used"].append({
"tool": tool_name,
"success": success,
"details": details,
"timestamp": datetime.datetime.now().isoformat()
})
self.save()
class UI:
"""Rich UI helpers"""
@staticmethod
def header(text: str):
console.print(Panel(text, style="bold cyan", border_style="cyan"))
@staticmethod
def agent_message(agent: str, message: str):
colors = {
"orchestrator": "green",
"planner": "blue",
"executor": "yellow",
"reviewer": "magenta",
"researcher": "cyan",
"coder": "bright_green",
"system": "white"
}
color = colors.get(agent.lower(), "white")
console.print(f"[bold {color}]{agent.upper()}:[/bold {color}] {message}")
@staticmethod
def code_block(code: str, language: str = "python"):
syntax = Syntax(code, language, theme="monokai", line_numbers=True)
console.print(Panel(syntax, border_style="dim"))
@staticmethod
def table(headers: List[str], rows: List[List[str]], title: str = ""):
table = Table(title=title, show_header=True, header_style="bold magenta")
for h in headers:
table.add_column(h)
for row in rows:
table.add_row(*row)
console.print(table)
@staticmethod
def tree(label: str, items: List[str]):
tree = Tree(f"[bold]{label}[/bold]")
for item in items:
tree.add(item)
console.print(tree)
@staticmethod
def confirm(prompt: str) -> bool:
"""Ask for user confirmation"""
response = console.input(f"[bold yellow]{prompt} (y/n): [/bold yellow]").lower().strip()
return response in ("y", "yes", "yeah", "sure")
@staticmethod
def success(message: str):
console.print(f"[bold green]✓ {message}[/bold green]")
@staticmethod
def error(message: str):
console.print(f"[bold red]✗ {message}[/bold red]")
@staticmethod
def warning(message: str):
console.print(f"[bold yellow]⚠ {message}[/bold yellow]")
@staticmethod
def info(message: str):
console.print(f"[dim]{message}[/dim]")
def hash_string(text: str) -> str:
"""Generate MD5 hash of string"""
return hashlib.md5(text.encode()).hexdigest()[:8]
def truncate(text: str, max_length: int = 500) -> str:
"""Truncate text with ellipsis"""
if len(text) <= max_length:
return text
return text[:max_length] + "... [truncated]"
def format_size(size_bytes: int) -> str:
"""Format bytes to human readable"""
for unit in ["B", "KB", "MB", "GB"]:
if size_bytes < 1024:
return f"{size_bytes:.1f} {unit}"
size_bytes /= 1024
return f"{size_bytes:.1f} TB"
''')
print("utils.py done")