augment17/claude-code-backend
0
1"""2Claude Code Backend — Agentic coding backend powered by NVIDIA NIM models.3Exposes an OpenAI-compatible /v1/chat/completions endpoint with built-in4tools for file operations and bash execution.5 6Architecture:7 Space 1 (better-chatbot) --> this backend --> NVIDIA NIM API8 9The agentic loop:10 1. Receive user message from Space 111 2. Send to NIM model with tool definitions12 3. If model returns tool_calls, execute them and loop13 4. If model returns text, stream it back to Space 114 5. Persist conversation in Postgres15"""16 17import os18import shutil19import threading20import json21import uuid22import subprocess23import asyncio24import time25import re26import collections27from pathlib import Path28from typing import AsyncIterator, Optional, List, Dict, Any29from pydantic import BaseModel30import contextvars31import urllib.parse32 33workspace_var = contextvars.ContextVar("workspace_dir", default="/tmp/workspace")34 35# --- Ultimate Agent Brain Imports ---36from second_brain import SecondBrainWrapper37from survival_watchdog import SurvivalWatchdog, get_metrics38from swarm_llm import swarm39from helix_state import helix_db40from context_engine import ContextEngine41 42# Instantiate singletons for the orchestrator43SPACE_NAME = os.environ.get("SPACE_NAME", "space2-cerebrum")44brain = SecondBrainWrapper(space_name=SPACE_NAME)45context_engine = ContextEngine(brain)46watchdog = SurvivalWatchdog()47 48 49from fastapi import FastAPI, Request, Header, HTTPException50from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse51from fastapi.middleware.cors import CORSMiddleware52from openai import AsyncOpenAI53import anyio54import asyncpg55 56# ---------------------------------------------------------------------------57# Globals & Activity Logs58# ---------------------------------------------------------------------------59activity_logs = collections.deque(maxlen=100)60MODEL_STATUSES = {}61ACTIVE_SESSIONS = set()62 63def log_activity(msg: str):64 timestamp = time.strftime("%H:%M:%S")65 log_line = f"[{timestamp}] {msg}"66 activity_logs.append(log_line)67 print(log_line)68 69# ---------------------------------------------------------------------------70# Configuration71# ---------------------------------------------------------------------------72 73NIM_API_KEY = os.environ.get("NVIDIA_NIM_API_KEY", "")74BACKEND_API_KEY = os.environ.get("BACKEND_API_KEY", "")75DATABASE_URL = os.environ.get("DATABASE_URL", "")76WORKSPACE_DIR = os.environ.get("WORKSPACE_DIR", "/tmp/workspace")77BACKUP_GIT_REPO = os.environ.get("BACKUP_GIT_REPO", "")78MAX_TOOL_ROUNDS = int(os.environ.get("MAX_TOOL_ROUNDS", "10"))79 80# NIM models that reliably support tool/function calling81TOOL_CAPABLE_MODELS = {82 "nvidia/nemotron-3-ultra-550b-a55b": "Nemotron 3 Ultra 550B (Agentic)",83 "z-ai/glm-5.1": "GLM 5.1 (Agentic)",84 "moonshotai/kimi-k2.6": "Kimi K2.6 (Agentic)",85 "minimaxai/minimax-m3": "MiniMax M3 (Agentic)",86 "stepfun-ai/step-3.7-flash": "Step 3.7 Flash (Agentic)",87 "minimaxai/minimax-m2.7": "MiniMax M2.7 (Agentic)",88 "meta/llama-3.1-70b-instruct": "Llama 3.1 70B (Agentic)",89 "meta/llama-3.1-405b-instruct": "Llama 3.1 405B (Agentic)",90 "qwen/qwen2.5-coder-32b-instruct": "Qwen 2.5 Coder 32B (Agentic)",91 "nvidia/llama-3.1-nemotron-70b-instruct": "Nemotron 70B (Agentic)",92 "meta/llama-3.3-70b-instruct": "Llama 3.3 70B (Agentic)",93}94 95# All models (tool-capable get agentic mode, others get plain chat)96ALL_MODELS = {97 **TOOL_CAPABLE_MODELS,98 "deepseek-ai/deepseek-r1": "DeepSeek R1 (Chat only)",99 "mistralai/mistral-large-2-instruct": "Mistral Large 2 (Chat only)",100}101 102RECOMMENDED_MODEL = "nvidia/llama-3.1-nemotron-70b-instruct"103 104# Ensure workspace exists105Path(WORKSPACE_DIR).mkdir(parents=True, exist_ok=True)106 107# ---------------------------------------------------------------------------108# NIM Client109# ---------------------------------------------------------------------------110 111nim_client = AsyncOpenAI(112 base_url="https://integrate.api.nvidia.com/v1",113 api_key=NIM_API_KEY,114)115 116# ---------------------------------------------------------------------------117# Rate Limiting & Multi-Provider Setup118# ---------------------------------------------------------------------------119 120MISTRAL_API_KEY = os.environ.get("MISTRAL_API_KEY", "YyZD5l1SCwL83hJIKsRVF5I7Oklcxu4k")121mistral_client = AsyncOpenAI(122 base_url="https://api.mistral.ai/v1",123 api_key=MISTRAL_API_KEY if MISTRAL_API_KEY else "dummy_key",124) if MISTRAL_API_KEY else None125 126class MultiProviderRateLimiter:127 def __init__(self):128 self.nim_limit = 40129 self.nim_window = 60130 self.nim_calls = []131 self.mistral_last_call = 0.0132 self.lock = asyncio.Lock()133 134 async def wait_for_mistral(self):135 async with self.lock:136 now = time.time()137 elapsed = now - self.mistral_last_call138 if elapsed < 1.0:139 await asyncio.sleep(1.0 - elapsed)140 self.mistral_last_call = time.time()141 142 async def wait_for_nim(self):143 async with self.lock:144 now = time.time()145 self.nim_calls = [t for t in self.nim_calls if now - t < self.nim_window]146 if len(self.nim_calls) >= self.nim_limit - 2:147 sleep_time = self.nim_window - (now - self.nim_calls[0])148 print(f"[RateLimiter] Approaching NIM rate limit (40 RPM). Sleeping {sleep_time:.2f}s...")149 await asyncio.sleep(sleep_time)150 self.nim_calls.append(time.time())151 152rate_limiter = MultiProviderRateLimiter()153 154# ---------------------------------------------------------------------------155# Tool Definitions (OpenAI function calling format)156# ---------------------------------------------------------------------------157 158TOOLS = [159 {160 "type": "function",161 "function": {162 "name": "read_file",163 "description": "Read the contents of a file. Use this to inspect existing code, configs, or any text file.",164 "parameters": {165 "type": "object",166 "properties": {167 "path": {168 "type": "string",169 "description": "Relative path to the file from the workspace root"170 }171 },172 "required": ["path"]173 }174 }175 },176 {177 "type": "function",178 "function": {179 "name": "write_file",180 "description": "Write content to a file. Creates the file if it doesn't exist, overwrites if it does. Creates parent directories automatically.",181 "parameters": {182 "type": "object",183 "properties": {184 "path": {185 "type": "string",186 "description": "Relative path to the file from the workspace root"187 },188 "content": {189 "type": "string",190 "description": "The full content to write to the file"191 }192 },193 "required": ["path", "content"]194 }195 }196 },197 {198 "type": "function",199 "function": {200 "name": "run_bash",201 "description": "Execute a bash command in the workspace directory. Use for installing packages, running scripts, git operations, etc. Commands run with a 30 second timeout.",202 "parameters": {203 "type": "object",204 "properties": {205 "command": {206 "type": "string",207 "description": "The bash command to execute"208 }209 },210 "required": ["command"]211 }212 }213 },214 {215 "type": "function",216 "function": {217 "name": "list_directory",218 "description": "List files and directories in a given path. Shows file sizes and directory markers.",219 "parameters": {220 "type": "object",221 "properties": {222 "path": {223 "type": "string",224 "description": "Relative path to the directory from workspace root. Use '.' for the workspace root."225 }226 },227 "required": ["path"]228 }229 }230 },231 {232 "type": "function",233 "function": {234 "name": "grep_search",235 "description": "Search for a pattern in files within the workspace. Returns matching lines with file paths and line numbers.",236 "parameters": {237 "type": "object",238 "properties": {239 "pattern": {240 "type": "string",241 "description": "The search pattern (supports basic regex)"242 },243 "path": {244 "type": "string",245 "description": "Directory or file to search in, relative to workspace root. Defaults to '.'",246 }247 },248 "required": ["pattern"]249 }250 }251 },252 {253 "type": "function",254 "function": {255 "name": "web_search",256 "description": "Search the web for up-to-date information, news, papers, or documentation.",257 "parameters": {258 "type": "object",259 "properties": {260 "query": {261 "type": "string",262 "description": "The search query (be specific and detailed)"263 }264 },265 "required": ["query"]266 }267 }268 },269 {270 "type": "function",271 "function": {272 "name": "web_read",273 "description": "Read the clean markdown content of any webpage URL to get detailed context, articles, documentation, or code.",274 "parameters": {275 "type": "object",276 "properties": {277 "url": {278 "type": "string",279 "description": "The absolute URL of the webpage to read"280 }281 },282 "required": ["url"]283 }284 }285 },286]287 288# ---------------------------------------------------------------------------289# Tool Execution290# ---------------------------------------------------------------------------291 292def _safe_path(rel_path: str) -> Path:293 """Resolve a relative path safely within the workspace."""294 workspace = Path(workspace_var.get()).resolve()295 target = (workspace / rel_path).resolve()296 # Prevent path traversal297 if not str(target).startswith(str(workspace)):298 raise ValueError(f"Path traversal detected: {rel_path}")299 return target300 301 302def repair_arguments(func_name: str, args: dict) -> tuple[dict, list[str]]:303 notes = []304 repaired_args = dict(args)305 306 # 1. Nesting extraction (e.g. {"path": {"path": "file.txt"}})307 for key in list(repaired_args.keys()):308 val = repaired_args[key]309 if isinstance(val, dict) and key in val:310 repaired_args[key] = val[key]311 notes.append(f"Flattened nested parameter '{key}'")312 313 # 2. Markdown stripping from bash command314 if func_name == "run_bash" and "command" in repaired_args:315 cmd = repaired_args["command"]316 if isinstance(cmd, str):317 pattern = r"```(?:bash)?\s*(.*?)\s*```"318 match = re.search(pattern, cmd, re.DOTALL)319 if match:320 repaired_args["command"] = match.group(1).strip()321 notes.append("Stripped markdown code blocks from bash command")322 323 # 3. Stringified array conversion324 for key, val in repaired_args.items():325 if isinstance(val, str) and val.strip().startswith("[") and val.strip().endswith("]"):326 try:327 parsed_arr = json.loads(val)328 if isinstance(parsed_arr, list):329 repaired_args[key] = parsed_arr330 notes.append(f"Converted stringified array for parameter '{key}' to native array")331 except:332 pass333 334 # 4. Optional empty objects replacing Null335 for key in list(repaired_args.keys()):336 if repaired_args[key] == {}:337 repaired_args[key] = None338 notes.append(f"Replaced empty object for parameter '{key}' with null")339 340 return repaired_args, notes341 342 343# ---------------------------------------------------------------------------344# Search and MCP Helpers345# ---------------------------------------------------------------------------346 347def sanitize_function_name(name: str) -> str:348 sanitized = re.sub(r'[^a-zA-Z0-9_\.\-]', '_', name)349 if not re.match(r'^[a-zA-Z_]', sanitized):350 sanitized = '_' + sanitized351 if len(sanitized) > 124:352 sanitized = sanitized[:124]353 return sanitized354 355def create_mcp_tool_id(server_name: str, tool_name: str) -> str:356 san_server = sanitize_function_name(server_name)357 san_tool = sanitize_function_name(tool_name)358 max_len = 124359 sep = "_"360 if len(san_server) + len(san_tool) + len(sep) > max_len:361 total = len(san_server) + len(san_tool)362 server_portion = int((len(san_server) / total) * (max_len - len(sep)))363 tool_portion = max_len - len(sep) - server_portion364 return f"{san_server[:server_portion]}{sep}{san_tool[:tool_portion]}"365 return f"{san_server}{sep}{san_tool}"366 367async def get_active_mcp_tools():368 """369 Queries `mcp_server` table and returns:370 1. List of OpenAI function definitions to append to dynamic tools list.371 2. Dictionary mapping `mcp_tool_id` to `(server_name, original_tool_name, config)`.372 """373 mcp_tools_list = []374 mcp_mapping = {}375 if not db_pool:376 return mcp_tools_list, mcp_mapping377 try:378 async with db_pool.acquire() as conn:379 rows = await conn.fetch("SELECT name, config, tool_info FROM mcp_server WHERE enabled = true")380 for row in rows:381 server_name = row["name"]382 config_raw = row["config"]383 if isinstance(config_raw, str):384 config = json.loads(config_raw)385 else:386 config = config_raw387 tool_info_raw = row["tool_info"]388 if not tool_info_raw:389 continue390 if isinstance(tool_info_raw, str):391 tool_info = json.loads(tool_info_raw)392 else:393 tool_info = tool_info_raw394 for tool in tool_info:395 tool_name = tool.get("name")396 description = tool.get("description", "")397 input_schema = tool.get("inputSchema", {})398 tool_id = create_mcp_tool_id(server_name, tool_name)399 mcp_mapping[tool_id] = (server_name, tool_name, config)400 mcp_tools_list.append({401 "type": "function",402 "function": {403 "name": tool_id,404 "description": f"[from MCP server: {server_name}] {description}",405 "parameters": input_schema406 }407 })408 except Exception as e:409 log_activity(f"[MCP Database Query Warning] Failed to load MCP tools: {e}")410 return mcp_tools_list, mcp_mapping411 412async def execute_web_search(query: str) -> str:413 tavily_key = os.environ.get("TAVILY_API_KEY", "")414 if tavily_key:415 try:416 async with httpx.AsyncClient(timeout=15.0) as client:417 r = await client.post("https://api.tavily.com/search", json={418 "api_key": tavily_key,419 "query": query,420 "search_depth": "basic",421 "max_results": 5422 })423 if r.status_code == 200:424 data = r.json()425 results = []426 for item in data.get("results", []):427 results.append(f"Title: {item.get('title')}\nURL: {item.get('url')}\nContent: {item.get('content')}\n---")428 return "\n".join(results) if results else "No results found."429 except Exception as e:430 log_activity(f"[Tavily Error] {e}")431 432 exa_key = os.environ.get("EXA_API_KEY", "")433 if exa_key:434 try:435 async with httpx.AsyncClient(timeout=15.0) as client:436 r = await client.post("https://api.exa.ai/search", headers={437 "x-api-key": exa_key,438 "Content-Type": "application/json"439 }, json={440 "query": query,441 "numResults": 5,442 "text": True443 })444 if r.status_code == 200:445 data = r.json()446 results = []447 for item in data.get("results", []):448 results.append(f"Title: {item.get('title')}\nURL: {item.get('url')}\nContent: {item.get('text', '')[:2000]}\n---")449 return "\n".join(results) if results else "No results found."450 except Exception as e:451 log_activity(f"[Exa Error] {e}")452 453 try:454 escaped_query = urllib.parse.quote(query)455 async with httpx.AsyncClient(headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}, timeout=10.0) as client:456 r = await client.get(f"https://html.duckduckgo.com/html/?q={escaped_query}")457 if r.status_code == 200:458 try:459 from bs4 import BeautifulSoup460 soup = BeautifulSoup(r.text, 'html.parser')461 results = []462 for a in soup.find_all('a', class_='result__snippet')[:5]:463 title_el = a.find_previous('a', class_='result__url')464 title = title_el.text.strip() if title_el else "No Title"465 url = title_el['href'] if title_el and 'href' in title_el.attrs else ""466 snippet = a.text.strip()467 results.append(f"Title: {title}\nURL: {url}\nSnippet: {snippet}\n---")468 return "\n".join(results) if results else "No results found."469 except Exception:470 snippets = re.findall(r'<a class="result__snippet"[^>]*>(.*?)</a>', r.text, re.DOTALL)471 results = []472 for s in snippets[:5]:473 clean_s = re.sub(r'<[^>]*>', '', s).strip()474 results.append(f"Snippet: {clean_s}\n---")475 return "\n".join(results) if results else "No results found."476 except Exception as e:477 log_activity(f"[DDG Scrape Error] {e}")478 return "Error: Web search failed. Setup API keys (TAVILY_API_KEY, EXA_API_KEY) for best results."479 480async def execute_web_read(url: str) -> str:481 jina_url = f"https://r.jina.ai/{url}"482 try:483 async with httpx.AsyncClient(timeout=20.0) as client:484 r = await client.get(jina_url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"})485 if r.status_code == 200:486 content = r.text487 if len(content) > 30000:488 content = content[:30000] + "\n\n[Truncated - webpage content is extremely long]"489 return content490 except Exception as e:491 log_activity(f"[Jina Reader Error] {e}")492 493 try:494 async with httpx.AsyncClient(timeout=15.0) as client:495 r = await client.get(url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"})496 if r.status_code == 200:497 html = r.text498 clean = re.sub(r'<script.*?>.*?</script>', '', html, flags=re.DOTALL | re.IGNORECASE)499 clean = re.sub(r'<style.*?>.*?</style>', '', clean, flags=re.DOTALL | re.IGNORECASE)500 clean = re.sub(r'<.*?>', ' ', clean)501 clean = re.sub(r'\s+', ' ', clean).strip()502 if len(clean) > 15000:503 clean = clean[:15000] + "\n\n[Truncated]"504 return clean505 except Exception as e:506 log_activity(f"[Direct Scrape Error] {e}")507 return f"Error: Failed to read URL {url}."508 509 510async def execute_tool(name: str, arguments: dict, mcp_mapping: dict = None) -> str:511 """Execute a tool and return its output as a string asynchronously."""512 try:513 if name == "read_file":514 path = _safe_path(arguments["path"])515 if not path.exists():516 return f"Error: File not found: {arguments['path']}"517 if not path.is_file():518 return f"Error: Not a file: {arguments['path']}"519 content = path.read_text(encoding="utf-8", errors="replace")520 if len(content) > 50000:521 return content[:50000] + f"\n\n[Truncated — file is {len(content)} chars]"522 return content523 524 elif name == "write_file":525 path = _safe_path(arguments["path"])526 path.parent.mkdir(parents=True, exist_ok=True)527 path.write_text(arguments["content"], encoding="utf-8")528 return f"Successfully wrote {len(arguments['content'])} chars to {arguments['path']}"529 530 elif name == "run_bash":531 command = arguments["command"]532 # Safety: block dangerous commands533 blocked = ["rm -rf /", "mkfs", "dd if=", ":(){", "fork bomb"]534 if any(b in command.lower() for b in blocked):535 return "Error: Command blocked for safety reasons"536 537 # ASYNC SUBPROCESS - This prevents the FastAPI server from freezing!538 process = await asyncio.create_subprocess_shell(539 command,540 stdout=asyncio.subprocess.PIPE,541 stderr=asyncio.subprocess.PIPE,542 cwd=workspace_var.get(),543 env={**os.environ, "HOME": "/tmp", "PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")},544 )545 546 try:547 stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=30)548 output = ""549 if stdout:550 output += stdout.decode('utf-8', errors='replace')551 if stderr:552 output += ("\n" if output else "") + f"[stderr] {stderr.decode('utf-8', errors='replace')}"553 if process.returncode != 0:554 output += f"\n[exit code: {process.returncode}]"555 if not output:556 output = "[command completed with no output]"557 except asyncio.TimeoutError:558 try:559 process.kill()560 except Exception:561 pass562 await process.communicate()563 return "Error: Command timed out after 30 seconds"564 565 if len(output) > 20000:566 output = output[:20000] + f"\n\n[Truncated — output is {len(output)} chars]"567 return output568 569 elif name == "list_directory":570 path = _safe_path(arguments.get("path", "."))571 if not path.exists():572 return f"Error: Directory not found: {arguments.get('path', '.')}"573 if not path.is_dir():574 return f"Error: Not a directory: {arguments.get('path', '.')}"575 entries = []576 for item in sorted(path.iterdir()):577 if item.is_dir():578 entries.append(f" 📁 {item.name}/")579 else:580 size = item.stat().st_size581 if size < 1024:582 size_str = f"{size}B"583 elif size < 1024 * 1024:584 size_str = f"{size/1024:.1f}KB"585 else:586 size_str = f"{size/(1024*1024):.1f}MB"587 entries.append(f" 📄 {item.name} ({size_str})")588 return f"Contents of {arguments.get('path', '.')}:\n" + "\n".join(entries) if entries else "Empty directory"589 590 elif name == "grep_search":591 pattern = arguments["pattern"]592 search_path = arguments.get("path", ".")593 path = _safe_path(search_path)594 595 process = await asyncio.create_subprocess_exec(596 "grep", "-rn", "--include=*", pattern, str(path),597 stdout=asyncio.subprocess.PIPE,598 stderr=asyncio.subprocess.PIPE,599 cwd=workspace_var.get(),600 )601 try:602 stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=10)603 output = stdout.decode('utf-8', errors='replace') if stdout else "No matches found"604 except asyncio.TimeoutError:605 try:606 process.kill()607 except Exception:608 pass609 await process.communicate()610 return "Error: Grep search timed out"611 612 if len(output) > 10000:613 output = output[:10000] + "\n\n[Truncated]"614 return output615 616 elif name == "web_search":617 query = arguments.get("query")618 return await execute_web_search(query)619 620 elif name == "web_read":621 url = arguments.get("url")622 return await execute_web_read(url)623 624 elif mcp_mapping and name in mcp_mapping:625 server_name, original_tool_name, config = mcp_mapping[name]626 try:627 # Ensure registered on Space 4628 await apost_json(f"{SPACE4_URL}/api/mcp/register", {629 "name": server_name,630 "config": config631 })632 # Call tool on Space 4633 call_res = await apost_json(f"{SPACE4_URL}/api/mcp/call", {634 "serverName": server_name,635 "toolName": original_tool_name,636 "arguments": arguments637 })638 if isinstance(call_res, dict) and "error" in call_res:639 return f"Error from MCP server {server_name}: {call_res['error']}"640 if isinstance(call_res, dict) and "content" in call_res:641 texts = []642 for item in call_res["content"]:643 if isinstance(item, dict) and item.get("type") == "text":644 texts.append(item.get("text", ""))645 return "\n".join(texts)646 return str(call_res)647 except Exception as e:648 return f"Error executing MCP tool {name}: {str(e)}"649 650 else:651 return f"Error: Unknown tool: {name}"652 653 except Exception as e:654 return f"Error executing {name}: {str(e)}"655 656 657# ---------------------------------------------------------------------------658# Database (Session Persistence)659# ---------------------------------------------------------------------------660 661db_pool: Optional[asyncpg.Pool] = None662 663 664async def init_db():665 """Initialize database connection pool and create tables."""666 global db_pool667 if not DATABASE_URL:668 return669 try:670 db_pool = await asyncpg.create_pool(671 DATABASE_URL, 672 ssl="require", 673 min_size=1, 674 max_size=3,675 max_inactive_connection_lifetime=300676 )677 async with db_pool.acquire() as conn:678 await conn.execute("""679 CREATE TABLE IF NOT EXISTS agent_session_entries (680 id BIGSERIAL PRIMARY KEY,681 project_key TEXT NOT NULL,682 session_id TEXT NOT NULL,683 subpath TEXT,684 entry JSONB NOT NULL,685 created_at TIMESTAMPTZ NOT NULL DEFAULT now()686 );687 CREATE INDEX IF NOT EXISTS idx_session_key ON agent_session_entries (project_key, session_id, subpath, id);688 CREATE INDEX IF NOT EXISTS idx_project_session ON agent_session_entries (project_key, session_id);689 690 CREATE TABLE IF NOT EXISTS eternity_projects (691 id SERIAL PRIMARY KEY,692 project_name VARCHAR(100) UNIQUE NOT NULL,693 goal TEXT NOT NULL,694 deadline TIMESTAMPTZ NOT NULL,695 current_mode VARCHAR(20) NOT NULL DEFAULT 'build',696 is_active BOOLEAN NOT NULL DEFAULT true,697 priority VARCHAR(20) NOT NULL DEFAULT 'low',698 roadmap JSONB DEFAULT '[]',699 latest_brief TEXT,700 created_at TIMESTAMPTZ NOT NULL DEFAULT now()701 );702 """)703 except Exception as e:704 print(f"[DB] Warning: Could not initialize database: {e}")705 db_pool = None706 707 708async def save_message(session_id: str, role: str, content: str = None, 709 tool_calls: list = None, tool_call_id: str = None):710 """Save a message to the session store using the unified schema."""711 if not db_pool:712 return713 msg = {"role": role}714 if content is not None:715 msg["content"] = content716 if tool_calls:717 msg["tool_calls"] = tool_calls718 if tool_call_id:719 msg["tool_call_id"] = tool_call_id720 721 try:722 async with db_pool.acquire() as conn:723 await conn.execute(724 "INSERT INTO agent_session_entries (project_key, session_id, subpath, entry) VALUES ($1, $2, $3, $4)",725 "fastapi-completions",726 session_id,727 None,728 json.dumps(msg)729 )730 except Exception as e:731 print(f"[DB] Warning: Could not save message: {e}")732 733 734async def load_session(session_id: str) -> list:735 """Load conversation history from the session store using the unified schema."""736 if not db_pool:737 return []738 try:739 async with db_pool.acquire() as conn:740 rows = await conn.fetch(741 "SELECT entry FROM agent_session_entries WHERE project_key = $1 AND session_id = $2 AND subpath IS NOT DISTINCT FROM $3 ORDER BY id",742 "fastapi-completions",743 session_id,744 None745 )746 return [json.loads(row["entry"]) for row in rows]747 except Exception as e:748 print(f"[DB] Warning: Could not load session: {e}")749 return []750 751 752# ---------------------------------------------------------------------------753# SSE Chunk Formatting (OpenAI delta format)754# ---------------------------------------------------------------------------755 756def make_chunk(request_id: str, model: str, content: str = "", finish_reason: str = None) -> str:757 """Create an OpenAI-compatible SSE chunk."""758 delta = {}759 if content:760 delta["content"] = content761 if finish_reason and not content:762 delta = {}763 764 chunk = {765 "id": f"chatcmpl-{request_id}",766 "object": "chat.completion.chunk",767 "created": int(time.time()),768 "model": model,769 "choices": [{770 "index": 0,771 "delta": delta,772 "finish_reason": finish_reason,773 }],774 }775 return f"data: {json.dumps(chunk)}\n\n"776 777 778# ---------------------------------------------------------------------------779# FastAPI Application780# ---------------------------------------------------------------------------781 782app = FastAPI(title="Claude Code Backend", version="1.0.0")783app.add_middleware(784 CORSMiddleware,785 allow_origins=["*"],786 allow_methods=["*"],787 allow_headers=["*"],788)789 790 791def auth(authorization: str = None):792 """Verify bearer token."""793 if not BACKEND_API_KEY:794 return # No auth configured795 expected = f"Bearer {BACKEND_API_KEY}"796 if authorization != expected:797 raise HTTPException(status_code=401, detail="Unauthorized")798 799 800async def check_models_health():801 # Mark all models as statically ONLINE to save API RPM quotas802 for model in ALL_MODELS.keys():803 MODEL_STATUSES[model] = {"status": "ONLINE (Stable)", "latency": "Fast", "raw_latency": 0.1}804 log_activity("[Health Check] Zero-cost status check complete: all models marked ONLINE (No API calls made).")805 806@app.on_event("startup")807async def startup():808 await init_db()809 Path(WORKSPACE_DIR).mkdir(parents=True, exist_ok=True)810 # Initialize statuses for all models as ONLINE by default811 for model_id, display_name in ALL_MODELS.items():812 MODEL_STATUSES[model_id] = {"status": "ONLINE", "latency": "0.10s", "raw_latency": 0.10}813 log_activity(f"FastAPI backend started. Workspace: {WORKSPACE_DIR}")814 815 816# ---------------------------------------------------------------------------817# /v1/chat/completions — Main endpoint818# ---------------------------------------------------------------------------819 820AGENTIC_SYSTEM_PROMPT = """You are an expert coding assistant with access to tools for file operations and command execution. 821 822When the user asks you to create, edit, or debug code:8231. Use `list_directory` and `read_file` to understand the current state8242. Use `write_file` to create or modify files8253. Use `run_bash` to execute commands (install packages, run scripts, test code)8264. Use `grep_search` to find patterns in code827 828IMPORTANT RULES:829- Always use tools to take action. Do NOT just describe what to do — actually DO it.830- After writing code, run it to verify it works.831- If a command fails, read the error and fix it.832- Work in the /tmp/workspace directory.833- Be concise in your explanations, but thorough in your tool usage.834- You are running in a multi-round tool-execution loop. Do NOT output redundant greetings (e.g. "Hey saumil!", "Let me check...", "Ready to build...") in intermediate rounds. Only output them in your first response if appropriate. Be direct and proceed with tool calls.835"""836 837 838def compact_history(messages: list) -> list:839 """840 Bulletproof conversation history compaction.841 Guarantees that the total context length stays well below the model's token limits.842 """843 if not messages:844 return []845 846 # Find and preserve the system prompt847 system_msg = None848 if messages[0].get("role") == "system":849 system_msg = messages[0]850 start_idx = 1851 else:852 start_idx = 0853 854 other_messages = messages[start_idx:]855 856 # 1. Truncate individually massive messages (e.g. file contents or massive bash outputs)857 # Even recent messages should be compacted if they are ridiculously large!858 compacted_others = []859 for msg in other_messages:860 role = msg.get("role")861 content = msg.get("content") or ""862 msg_copy = dict(msg)863 864 if len(content) > 15000:865 msg_copy["content"] = content[:5000] + f"\n\n[... Truncated {len(content) - 10000} characters to prevent model context limits from overflowing ...]\n\n" + content[-5000:]866 867 compacted_others.append(msg_copy)868 869 # 2. Enforce total character size budget (max ~300,000 characters / ~75,000 tokens)870 # Iterate backwards from newest to oldest871 final_list = []872 total_chars = 0873 max_budget = 300000874 875 for msg in reversed(compacted_others):876 msg_len = len(msg.get("content") or "")877 if total_chars + msg_len < max_budget or len(final_list) < 2:878 final_list.append(msg)879 total_chars += msg_len880 else:881 # Drop older messages once we exceed context window budget882 pass883 884 # Reverse back to chronological order885 final_list.reverse()886 887 if system_msg:888 final_list.insert(0, system_msg)889 890 if len(messages) != len(final_list):891 log_activity(f"[Auto-Compaction] Sliced context window from {len(messages)} down to {len(final_list)} messages (Total chars: {total_chars})")892 893 return final_list894 895 896@app.post("/v1/chat/completions")897async def chat_completions(request: Request, authorization: str = Header(None)):898 auth(authorization)899 body = await request.json()900 901 requested_model = body.get("model", "meta/llama-3.1-70b-instruct")902 messages = body.get("messages", [])903 stream = body.get("stream", False)904 session_id = body.get("session_id") or str(uuid.uuid4())905 906 # Route Mistral queries natively to the Mistral API907 client = nim_client908 if "mistral" in requested_model.lower():909 if mistral_client:910 client = mistral_client911 if requested_model == "mistralai/mistral-large-2-instruct":912 requested_model = "mistral-large-latest"913 914 is_agentic = requested_model in TOOL_CAPABLE_MODELS915 request_id = str(uuid.uuid4())[:8]916 917 ACTIVE_SESSIONS.add(session_id)918 log_activity(f"Session [{session_id[:6]}] connected. Model: {requested_model} | Provider: {'Mistral' if client == mistral_client else 'NIM'}")919 920 # Build message history921 final_messages = []922 923 # Add agentic system prompt for tool-capable models924 if is_agentic:925 # Check if there's already a system message926 has_system = any(m.get("role") == "system" for m in messages)927 if has_system:928 # Prepend agentic prompt to existing system message929 for m in messages:930 if m["role"] == "system":931 final_messages.append({932 "role": "system",933 "content": AGENTIC_SYSTEM_PROMPT + "\n\nAdditional instructions:\n" + m["content"]934 })935 else:936 final_messages.append(m)937 else:938 final_messages.append({"role": "system", "content": AGENTIC_SYSTEM_PROMPT})939 final_messages.extend(messages)940 else:941 final_messages = list(messages)942 943 # Perform auto-compaction before executing agent loops944 final_messages = compact_history(final_messages)945 946 # Save the user's message to DB947 user_msg = next((m for m in reversed(messages) if m.get("role") == "user"), None)948 if user_msg:949 await save_message(session_id, "user", user_msg.get("content", ""))950 951 session_folder = session_id[:8] if session_id else "default"952 active_workspace = os.path.join(WORKSPACE_DIR, session_folder)953 os.makedirs(active_workspace, exist_ok=True)954 955 if not stream:956 # Non-streaming: simple completion957 workspace_token = workspace_var.set(active_workspace)958 try:959 mcp_tools, mcp_mapping = await get_active_mcp_tools()960 active_tools = list(TOOLS) + mcp_tools961 kwargs = {"model": requested_model, "messages": final_messages}962 if is_agentic:963 kwargs["tools"] = active_tools964 kwargs["tool_choice"] = "auto"965 async with completions_semaphore:966 response = await client.chat.completions.create(**kwargs)967 content = response.choices[0].message.content or ""968 await save_message(session_id, "assistant", content)969 ACTIVE_SESSIONS.discard(session_id)970 log_activity(f"Session [{session_id[:6]}] finished (non-streaming)")971 return JSONResponse({972 "id": f"chatcmpl-{request_id}",973 "object": "chat.completion",974 "created": int(time.time()),975 "model": requested_model,976 "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],977 })978 except Exception as e:979 ACTIVE_SESSIONS.discard(session_id)980 return JSONResponse({"error": {"message": str(e), "type": "internal_error"}}, status_code=500)981 finally:982 workspace_var.reset(workspace_token)983 984 # Streaming + agentic loop985 async def generate() -> AsyncIterator[str]:986 nonlocal final_messages987 workspace_token = workspace_var.set(active_workspace)988 try:989 mcp_tools, mcp_mapping = await get_active_mcp_tools()990 active_tools = list(TOOLS) + mcp_tools991 async with completions_semaphore:992 for round_num in range(MAX_TOOL_ROUNDS + 1):993 # Perform auto-compaction before calling NIM API994 final_messages = compact_history(final_messages)995 kwargs = {"model": requested_model, "messages": final_messages, "stream": True}996 if is_agentic:997 kwargs["tools"] = active_tools998 kwargs["tool_choice"] = "auto"999 1000 # Collect streamed response1001 full_content = ""1002 tool_calls_raw = {} # index -> {id, name, arguments_str}1003 1004 async for chunk in await client.chat.completions.create(**kwargs):1005 choice = chunk.choices[0] if chunk.choices else None1006 if not choice:1007 continue1008 delta = choice.delta1009 1010 # Stream text content to client1011 if delta and delta.content:1012 full_content += delta.content1013 yield make_chunk(request_id, requested_model, delta.content)1014 1015 # Collect tool calls1016 if delta and delta.tool_calls:1017 for tc in delta.tool_calls:1018 idx = tc.index1019 if idx not in tool_calls_raw:1020 tool_calls_raw[idx] = {1021 "id": tc.id or f"call_{uuid.uuid4().hex[:8]}",1022 "name": tc.function.name if tc.function and tc.function.name else "",1023 "arguments": ""1024 }1025 if tc.function and tc.function.name:1026 tool_calls_raw[idx]["name"] = tc.function.name1027 if tc.id:1028 tool_calls_raw[idx]["id"] = tc.id1029 if tc.function and tc.function.arguments:1030 tool_calls_raw[idx]["arguments"] += tc.function.arguments1031 1032 # Check for finish1033 if choice.finish_reason == "stop":1034 break1035 if choice.finish_reason == "tool_calls":1036 break1037 1038 # If no tool calls, we're done1039 if not tool_calls_raw:1040 await save_message(session_id, "assistant", full_content)1041 yield make_chunk(request_id, requested_model, finish_reason="stop")1042 yield "data: [DONE]\n\n"1043 return1044 1045 # Execute tool calls1046 tool_calls_list = []1047 for idx in sorted(tool_calls_raw.keys()):1048 tc = tool_calls_raw[idx]1049 tool_calls_list.append({1050 "id": tc["id"],1051 "type": "function",1052 "function": {"name": tc["name"], "arguments": tc["arguments"]}1053 })1054 1055 # Add assistant message with tool calls to history1056 assistant_msg = {"role": "assistant", "content": full_content or None, "tool_calls": tool_calls_list}1057 final_messages.append(assistant_msg)1058 1059 # Execute each tool and add results1060 for tc in tool_calls_list:1061 func_name = tc["function"]["name"]1062 raw_args_str = tc["function"]["arguments"]1063 try:1064 func_args = json.loads(raw_args_str)1065 except json.JSONDecodeError:1066 # Attempt raw JSON repair1067 repaired_str = raw_args_str.strip()1068 if not repaired_str.startswith("{"):1069 repaired_str = "{" + repaired_str1070 if not repaired_str.endswith("}"):1071 repaired_str = repaired_str + "}"1072 try:1073 func_args = json.loads(repaired_str)1074 log_activity(f"Auto-fixed invalid JSON string for tool: {func_name}")1075 except:1076 func_args = {}1077 1078 # Perform semantic repairs1079 repaired_args, repair_notes = repair_arguments(func_name, func_args)1080 1081 # Log activity1082 log_activity(f"Tool execution: {func_name} args={repaired_args}")1083 if repair_notes:1084 for note in repair_notes:1085 log_activity(f"[Tool Repair] {note}")1086 1087 # Show tool execution to user1088 yield make_chunk(request_id, requested_model, f"\n\n🔧 **{func_name}**")1089 if repair_notes:1090 yield make_chunk(request_id, requested_model, " *(Auto-Repaired)*")1091 if func_name == "run_bash" and "command" in repaired_args:1092 yield make_chunk(request_id, requested_model, f": `{repaired_args['command']}`\n")1093 elif func_name == "read_file" and "path" in repaired_args:1094 yield make_chunk(request_id, requested_model, f": `{repaired_args['path']}`\n")1095 elif func_name == "write_file" and "path" in repaired_args:1096 yield make_chunk(request_id, requested_model, f": `{repaired_args['path']}`\n")1097 elif func_name == "list_directory":1098 yield make_chunk(request_id, requested_model, f": `{repaired_args.get('path', '.')}`\n")1099 elif func_name == "grep_search":1100 yield make_chunk(request_id, requested_model, f": `{repaired_args.get('pattern', '')}`\n")1101 elif func_name == "web_search" and "query" in repaired_args:1102 yield make_chunk(request_id, requested_model, f": `{repaired_args['query']}`\n")1103 elif func_name == "web_read" and "url" in repaired_args:1104 yield make_chunk(request_id, requested_model, f": `{repaired_args['url']}`\n")1105 elif mcp_mapping and func_name in mcp_mapping:1106 yield make_chunk(request_id, requested_model, f": calling tool\n")1107 else:1108 yield make_chunk(request_id, requested_model, "\n")1109 1110 # Execute the tool1111 result = await execute_tool(func_name, repaired_args, mcp_mapping)1112 1113 # Append teaching note if repaired1114 if repair_notes:1115 result += f"\n\n[SYSTEM REPAIR NOTE: The harness automatically fixed formatting issues: {', '.join(repair_notes)}. Please strictly follow the tool's JSON schema in subsequent calls without these wrapping/formatting errors.]"1116 1117 # Show truncated result to user1118 preview = result[:500] + ("..." if len(result) > 500 else "")1119 yield make_chunk(request_id, requested_model, f"```\n{preview}\n```\n")1120 1121 # Add tool result to message history1122 final_messages.append({1123 "role": "tool",1124 "tool_call_id": tc["id"],1125 "content": result,1126 })1127 1128 await save_message(session_id, "tool", result, tool_call_id=tc["id"])1129 1130 # Continue the agentic loop (model processes tool results)1131 1132 # If we hit max rounds, finish1133 yield make_chunk(request_id, requested_model, "\n\n⚠️ Reached maximum tool call rounds.")1134 yield make_chunk(request_id, requested_model, finish_reason="stop")1135 yield "data: [DONE]\n\n"1136 1137 except Exception as e:1138 error_msg = f"\n\n❌ Error: {str(e)}"1139 yield make_chunk(request_id, requested_model, error_msg)1140 yield make_chunk(request_id, requested_model, finish_reason="stop")1141 yield "data: [DONE]\n\n"1142 finally:1143 workspace_var.reset(workspace_token)1144 1145 return StreamingResponse(1146 generate(),1147 media_type="text/event-stream",1148 headers={1149 "Cache-Control": "no-cache",1150 "X-Accel-Buffering": "no",1151 "Connection": "keep-alive",1152 },1153 )1154 1155 1156# ---------------------------------------------------------------------------1157# /v1/models — Model listing1158# ---------------------------------------------------------------------------1159 1160@app.get("/v1/models")1161async def list_models(authorization: str = Header(None)):1162 auth(authorization)1163 models = []1164 for model_id, display_name in ALL_MODELS.items():1165 models.append({1166 "id": model_id,1167 "object": "model",1168 "created": 1700000000,1169 "owned_by": "nvidia-nim",1170 "permission": [],1171 "root": model_id,1172 "parent": None,1173 })1174 return {"object": "list", "data": models}1175 1176 1177# ---------------------------------------------------------------------------1178# /health — Health check1179# ---------------------------------------------------------------------------1180 1181@app.get("/api/workspace/tree")1182async def get_workspace_tree():1183 def build_tree(current_path: Path, relative_to: Path) -> dict:1184 name = current_path.name1185 try:1186 rel_path = str(current_path.relative_to(relative_to)).replace("\\", "/")1187 except ValueError:1188 rel_path = ""1189 if rel_path == ".":1190 rel_path = ""1191 1192 if current_path.is_dir():1193 children = []1194 try:1195 for child in sorted(current_path.iterdir(), key=lambda x: (not x.is_dir(), x.name)):1196 if child.name in [".git", "node_modules", ".next", "__pycache__", ".agents", ".gemini"]:1197 continue1198 children.append(build_tree(child, relative_to))1199 except Exception:1200 pass