Cyber Catalyst Team commited on
Commit
da93542
·
0 Parent(s):

Initial: Agentic backend with NIM + tools

Browse files
Files changed (4) hide show
  1. Dockerfile +28 -0
  2. README.md +14 -0
  3. backend.py +668 -0
  4. requirements.txt +5 -0
Dockerfile ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.11-slim
2
+
3
+ # Install system dependencies
4
+ RUN apt-get update && apt-get install -y --no-install-recommends \
5
+ git bash curl && \
6
+ rm -rf /var/lib/apt/lists/*
7
+
8
+ # Create non-root user (HF requirement: uid 1000)
9
+ RUN useradd -m -u 1000 user
10
+ USER user
11
+ ENV HOME=/home/user \
12
+ PATH=/home/user/.local/bin:$PATH
13
+
14
+ WORKDIR $HOME/app
15
+
16
+ # Install Python dependencies
17
+ COPY --chown=user requirements.txt .
18
+ RUN pip install --no-cache-dir --user -r requirements.txt
19
+
20
+ # Copy application
21
+ COPY --chown=user backend.py .
22
+
23
+ # Create workspace directory
24
+ RUN mkdir -p /tmp/workspace
25
+
26
+ EXPOSE 7860
27
+
28
+ CMD ["python", "backend.py"]
README.md ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: claude-code-backend
3
+ emoji: 🤖
4
+ colorFrom: blue
5
+ colorTo: purple
6
+ sdk: docker
7
+ app_port: 7860
8
+ ---
9
+
10
+ # Claude Code Backend
11
+
12
+ Agentic coding backend powered by NVIDIA NIM models. Provides Claude-Code-like capabilities (Read, Write, Edit, Bash, Grep) through an OpenAI-compatible API.
13
+
14
+ Used as a backend provider for [better-chatbot](https://huggingface.co/spaces/augment17/better-chatbot).
backend.py ADDED
@@ -0,0 +1,668 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Claude Code Backend — Agentic coding backend powered by NVIDIA NIM models.
3
+ Exposes an OpenAI-compatible /v1/chat/completions endpoint with built-in
4
+ tools for file operations and bash execution.
5
+
6
+ Architecture:
7
+ Space 1 (better-chatbot) --> this backend --> NVIDIA NIM API
8
+
9
+ The agentic loop:
10
+ 1. Receive user message from Space 1
11
+ 2. Send to NIM model with tool definitions
12
+ 3. If model returns tool_calls, execute them and loop
13
+ 4. If model returns text, stream it back to Space 1
14
+ 5. Persist conversation in Postgres
15
+ """
16
+
17
+ import os
18
+ import json
19
+ import uuid
20
+ import subprocess
21
+ import asyncio
22
+ import time
23
+ from pathlib import Path
24
+ from typing import AsyncIterator, Optional
25
+
26
+ from fastapi import FastAPI, Request, Header, HTTPException
27
+ from fastapi.responses import StreamingResponse, JSONResponse
28
+ from fastapi.middleware.cors import CORSMiddleware
29
+ from openai import AsyncOpenAI
30
+ import asyncpg
31
+
32
+ # ---------------------------------------------------------------------------
33
+ # Configuration
34
+ # ---------------------------------------------------------------------------
35
+
36
+ NIM_API_KEY = os.environ.get("NVIDIA_NIM_API_KEY", "")
37
+ BACKEND_API_KEY = os.environ.get("BACKEND_API_KEY", "")
38
+ DATABASE_URL = os.environ.get("DATABASE_URL", "")
39
+ WORKSPACE_DIR = os.environ.get("WORKSPACE_DIR", "/tmp/workspace")
40
+ MAX_TOOL_ROUNDS = int(os.environ.get("MAX_TOOL_ROUNDS", "10"))
41
+
42
+ # NIM models that reliably support tool/function calling
43
+ TOOL_CAPABLE_MODELS = {
44
+ "meta/llama-3.1-70b-instruct": "Llama 3.1 70B (Agentic)",
45
+ "meta/llama-3.1-405b-instruct": "Llama 3.1 405B (Agentic)",
46
+ "qwen/qwen2.5-coder-32b-instruct": "Qwen 2.5 Coder 32B (Agentic)",
47
+ "nvidia/llama-3.1-nemotron-70b-instruct": "Nemotron 70B (Agentic)",
48
+ "meta/llama-3.3-70b-instruct": "Llama 3.3 70B (Agentic)",
49
+ }
50
+
51
+ # All models (tool-capable get agentic mode, others get plain chat)
52
+ ALL_MODELS = {
53
+ **TOOL_CAPABLE_MODELS,
54
+ "deepseek-ai/deepseek-r1": "DeepSeek R1 (Chat only)",
55
+ "mistralai/mistral-large-2-instruct": "Mistral Large 2 (Chat only)",
56
+ }
57
+
58
+ # Ensure workspace exists
59
+ Path(WORKSPACE_DIR).mkdir(parents=True, exist_ok=True)
60
+
61
+ # ---------------------------------------------------------------------------
62
+ # NIM Client
63
+ # ---------------------------------------------------------------------------
64
+
65
+ nim_client = AsyncOpenAI(
66
+ base_url="https://integrate.api.nvidia.com/v1",
67
+ api_key=NIM_API_KEY,
68
+ )
69
+
70
+ # ---------------------------------------------------------------------------
71
+ # Tool Definitions (OpenAI function calling format)
72
+ # ---------------------------------------------------------------------------
73
+
74
+ TOOLS = [
75
+ {
76
+ "type": "function",
77
+ "function": {
78
+ "name": "read_file",
79
+ "description": "Read the contents of a file. Use this to inspect existing code, configs, or any text file.",
80
+ "parameters": {
81
+ "type": "object",
82
+ "properties": {
83
+ "path": {
84
+ "type": "string",
85
+ "description": "Relative path to the file from the workspace root"
86
+ }
87
+ },
88
+ "required": ["path"]
89
+ }
90
+ }
91
+ },
92
+ {
93
+ "type": "function",
94
+ "function": {
95
+ "name": "write_file",
96
+ "description": "Write content to a file. Creates the file if it doesn't exist, overwrites if it does. Creates parent directories automatically.",
97
+ "parameters": {
98
+ "type": "object",
99
+ "properties": {
100
+ "path": {
101
+ "type": "string",
102
+ "description": "Relative path to the file from the workspace root"
103
+ },
104
+ "content": {
105
+ "type": "string",
106
+ "description": "The full content to write to the file"
107
+ }
108
+ },
109
+ "required": ["path", "content"]
110
+ }
111
+ }
112
+ },
113
+ {
114
+ "type": "function",
115
+ "function": {
116
+ "name": "run_bash",
117
+ "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.",
118
+ "parameters": {
119
+ "type": "object",
120
+ "properties": {
121
+ "command": {
122
+ "type": "string",
123
+ "description": "The bash command to execute"
124
+ }
125
+ },
126
+ "required": ["command"]
127
+ }
128
+ }
129
+ },
130
+ {
131
+ "type": "function",
132
+ "function": {
133
+ "name": "list_directory",
134
+ "description": "List files and directories in a given path. Shows file sizes and directory markers.",
135
+ "parameters": {
136
+ "type": "object",
137
+ "properties": {
138
+ "path": {
139
+ "type": "string",
140
+ "description": "Relative path to the directory from workspace root. Use '.' for the workspace root."
141
+ }
142
+ },
143
+ "required": ["path"]
144
+ }
145
+ }
146
+ },
147
+ {
148
+ "type": "function",
149
+ "function": {
150
+ "name": "grep_search",
151
+ "description": "Search for a pattern in files within the workspace. Returns matching lines with file paths and line numbers.",
152
+ "parameters": {
153
+ "type": "object",
154
+ "properties": {
155
+ "pattern": {
156
+ "type": "string",
157
+ "description": "The search pattern (supports basic regex)"
158
+ },
159
+ "path": {
160
+ "type": "string",
161
+ "description": "Directory or file to search in, relative to workspace root. Defaults to '.'",
162
+ }
163
+ },
164
+ "required": ["pattern"]
165
+ }
166
+ }
167
+ },
168
+ ]
169
+
170
+ # ---------------------------------------------------------------------------
171
+ # Tool Execution
172
+ # ---------------------------------------------------------------------------
173
+
174
+ def _safe_path(rel_path: str) -> Path:
175
+ """Resolve a relative path safely within the workspace."""
176
+ workspace = Path(WORKSPACE_DIR).resolve()
177
+ target = (workspace / rel_path).resolve()
178
+ # Prevent path traversal
179
+ if not str(target).startswith(str(workspace)):
180
+ raise ValueError(f"Path traversal detected: {rel_path}")
181
+ return target
182
+
183
+
184
+ def execute_tool(name: str, arguments: dict) -> str:
185
+ """Execute a tool and return its output as a string."""
186
+ try:
187
+ if name == "read_file":
188
+ path = _safe_path(arguments["path"])
189
+ if not path.exists():
190
+ return f"Error: File not found: {arguments['path']}"
191
+ if not path.is_file():
192
+ return f"Error: Not a file: {arguments['path']}"
193
+ content = path.read_text(encoding="utf-8", errors="replace")
194
+ if len(content) > 50000:
195
+ return content[:50000] + f"\n\n[Truncated — file is {len(content)} chars]"
196
+ return content
197
+
198
+ elif name == "write_file":
199
+ path = _safe_path(arguments["path"])
200
+ path.parent.mkdir(parents=True, exist_ok=True)
201
+ path.write_text(arguments["content"], encoding="utf-8")
202
+ return f"Successfully wrote {len(arguments['content'])} chars to {arguments['path']}"
203
+
204
+ elif name == "run_bash":
205
+ command = arguments["command"]
206
+ # Safety: block dangerous commands
207
+ blocked = ["rm -rf /", "mkfs", "dd if=", ":(){", "fork bomb"]
208
+ if any(b in command.lower() for b in blocked):
209
+ return "Error: Command blocked for safety reasons"
210
+ result = subprocess.run(
211
+ ["bash", "-c", command],
212
+ cwd=WORKSPACE_DIR,
213
+ capture_output=True,
214
+ text=True,
215
+ timeout=30,
216
+ env={**os.environ, "HOME": "/tmp", "PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")},
217
+ )
218
+ output = ""
219
+ if result.stdout:
220
+ output += result.stdout
221
+ if result.stderr:
222
+ output += ("\n" if output else "") + f"[stderr] {result.stderr}"
223
+ if result.returncode != 0:
224
+ output += f"\n[exit code: {result.returncode}]"
225
+ if not output:
226
+ output = "[command completed with no output]"
227
+ # Truncate very long outputs
228
+ if len(output) > 20000:
229
+ output = output[:20000] + f"\n\n[Truncated — output is {len(output)} chars]"
230
+ return output
231
+
232
+ elif name == "list_directory":
233
+ path = _safe_path(arguments.get("path", "."))
234
+ if not path.exists():
235
+ return f"Error: Directory not found: {arguments.get('path', '.')}"
236
+ if not path.is_dir():
237
+ return f"Error: Not a directory: {arguments.get('path', '.')}"
238
+ entries = []
239
+ for item in sorted(path.iterdir()):
240
+ if item.is_dir():
241
+ entries.append(f" 📁 {item.name}/")
242
+ else:
243
+ size = item.stat().st_size
244
+ if size < 1024:
245
+ size_str = f"{size}B"
246
+ elif size < 1024 * 1024:
247
+ size_str = f"{size/1024:.1f}KB"
248
+ else:
249
+ size_str = f"{size/(1024*1024):.1f}MB"
250
+ entries.append(f" 📄 {item.name} ({size_str})")
251
+ return f"Contents of {arguments.get('path', '.')}:\n" + "\n".join(entries) if entries else "Empty directory"
252
+
253
+ elif name == "grep_search":
254
+ pattern = arguments["pattern"]
255
+ search_path = arguments.get("path", ".")
256
+ path = _safe_path(search_path)
257
+ result = subprocess.run(
258
+ ["grep", "-rn", "--include=*", pattern, str(path)],
259
+ capture_output=True,
260
+ text=True,
261
+ timeout=10,
262
+ cwd=WORKSPACE_DIR,
263
+ )
264
+ output = result.stdout if result.stdout else "No matches found"
265
+ if len(output) > 10000:
266
+ output = output[:10000] + "\n\n[Truncated]"
267
+ return output
268
+
269
+ else:
270
+ return f"Error: Unknown tool: {name}"
271
+
272
+ except subprocess.TimeoutExpired:
273
+ return "Error: Command timed out after 30 seconds"
274
+ except ValueError as e:
275
+ return f"Error: {str(e)}"
276
+ except Exception as e:
277
+ return f"Error executing {name}: {str(e)}"
278
+
279
+
280
+ # ---------------------------------------------------------------------------
281
+ # Database (Session Persistence)
282
+ # ---------------------------------------------------------------------------
283
+
284
+ db_pool: Optional[asyncpg.Pool] = None
285
+
286
+
287
+ async def init_db():
288
+ """Initialize database connection pool and create tables."""
289
+ global db_pool
290
+ if not DATABASE_URL:
291
+ return
292
+ try:
293
+ db_pool = await asyncpg.create_pool(DATABASE_URL, ssl="require", min_size=1, max_size=5)
294
+ async with db_pool.acquire() as conn:
295
+ await conn.execute("""
296
+ CREATE TABLE IF NOT EXISTS agent_sessions (
297
+ id BIGSERIAL PRIMARY KEY,
298
+ session_id TEXT NOT NULL,
299
+ role TEXT NOT NULL,
300
+ content TEXT,
301
+ tool_calls JSONB,
302
+ tool_call_id TEXT,
303
+ created_at TIMESTAMPTZ DEFAULT NOW()
304
+ );
305
+ CREATE INDEX IF NOT EXISTS idx_agent_sessions_sid ON agent_sessions(session_id);
306
+ """)
307
+ except Exception as e:
308
+ print(f"[DB] Warning: Could not initialize database: {e}")
309
+ db_pool = None
310
+
311
+
312
+ async def save_message(session_id: str, role: str, content: str = None,
313
+ tool_calls: list = None, tool_call_id: str = None):
314
+ """Save a message to the session store."""
315
+ if not db_pool:
316
+ return
317
+ try:
318
+ async with db_pool.acquire() as conn:
319
+ await conn.execute(
320
+ "INSERT INTO agent_sessions (session_id, role, content, tool_calls, tool_call_id) VALUES ($1, $2, $3, $4, $5)",
321
+ session_id, role, content,
322
+ json.dumps(tool_calls) if tool_calls else None,
323
+ tool_call_id,
324
+ )
325
+ except Exception as e:
326
+ print(f"[DB] Warning: Could not save message: {e}")
327
+
328
+
329
+ async def load_session(session_id: str) -> list:
330
+ """Load conversation history from the session store."""
331
+ if not db_pool:
332
+ return []
333
+ try:
334
+ async with db_pool.acquire() as conn:
335
+ rows = await conn.fetch(
336
+ "SELECT role, content, tool_calls, tool_call_id FROM agent_sessions WHERE session_id = $1 ORDER BY id",
337
+ session_id,
338
+ )
339
+ messages = []
340
+ for row in rows:
341
+ msg = {"role": row["role"]}
342
+ if row["content"]:
343
+ msg["content"] = row["content"]
344
+ if row["tool_calls"]:
345
+ msg["tool_calls"] = json.loads(row["tool_calls"])
346
+ if row["tool_call_id"]:
347
+ msg["tool_call_id"] = row["tool_call_id"]
348
+ messages.append(msg)
349
+ return messages
350
+ except Exception as e:
351
+ print(f"[DB] Warning: Could not load session: {e}")
352
+ return []
353
+
354
+
355
+ # ---------------------------------------------------------------------------
356
+ # SSE Chunk Formatting (OpenAI delta format)
357
+ # ---------------------------------------------------------------------------
358
+
359
+ def make_chunk(request_id: str, model: str, content: str = "", finish_reason: str = None) -> str:
360
+ """Create an OpenAI-compatible SSE chunk."""
361
+ delta = {}
362
+ if content:
363
+ delta["content"] = content
364
+ if finish_reason and not content:
365
+ delta = {}
366
+
367
+ chunk = {
368
+ "id": f"chatcmpl-{request_id}",
369
+ "object": "chat.completion.chunk",
370
+ "created": int(time.time()),
371
+ "model": model,
372
+ "choices": [{
373
+ "index": 0,
374
+ "delta": delta,
375
+ "finish_reason": finish_reason,
376
+ }],
377
+ }
378
+ return f"data: {json.dumps(chunk)}\n\n"
379
+
380
+
381
+ # ---------------------------------------------------------------------------
382
+ # FastAPI Application
383
+ # ---------------------------------------------------------------------------
384
+
385
+ app = FastAPI(title="Claude Code Backend", version="1.0.0")
386
+ app.add_middleware(
387
+ CORSMiddleware,
388
+ allow_origins=["*"],
389
+ allow_methods=["*"],
390
+ allow_headers=["*"],
391
+ )
392
+
393
+
394
+ def auth(authorization: str = None):
395
+ """Verify bearer token."""
396
+ if not BACKEND_API_KEY:
397
+ return # No auth configured
398
+ expected = f"Bearer {BACKEND_API_KEY}"
399
+ if authorization != expected:
400
+ raise HTTPException(status_code=401, detail="Unauthorized")
401
+
402
+
403
+ @app.on_event("startup")
404
+ async def startup():
405
+ await init_db()
406
+ Path(WORKSPACE_DIR).mkdir(parents=True, exist_ok=True)
407
+ print(f"[Backend] Started. Workspace: {WORKSPACE_DIR}")
408
+ print(f"[Backend] Tool-capable models: {list(TOOL_CAPABLE_MODELS.keys())}")
409
+ print(f"[Backend] DB connected: {db_pool is not None}")
410
+
411
+
412
+ # ---------------------------------------------------------------------------
413
+ # /v1/chat/completions — Main endpoint
414
+ # ---------------------------------------------------------------------------
415
+
416
+ AGENTIC_SYSTEM_PROMPT = """You are an expert coding assistant with access to tools for file operations and command execution.
417
+
418
+ When the user asks you to create, edit, or debug code:
419
+ 1. Use `list_directory` and `read_file` to understand the current state
420
+ 2. Use `write_file` to create or modify files
421
+ 3. Use `run_bash` to execute commands (install packages, run scripts, test code)
422
+ 4. Use `grep_search` to find patterns in code
423
+
424
+ IMPORTANT RULES:
425
+ - Always use tools to take action. Do NOT just describe what to do — actually DO it.
426
+ - After writing code, run it to verify it works.
427
+ - If a command fails, read the error and fix it.
428
+ - Work in the /tmp/workspace directory.
429
+ - Be concise in your explanations, but thorough in your tool usage.
430
+ """
431
+
432
+
433
+ @app.post("/v1/chat/completions")
434
+ async def chat_completions(request: Request, authorization: str = Header(None)):
435
+ auth(authorization)
436
+ body = await request.json()
437
+
438
+ requested_model = body.get("model", "meta/llama-3.1-70b-instruct")
439
+ messages = body.get("messages", [])
440
+ stream = body.get("stream", False)
441
+ session_id = body.get("session_id") or str(uuid.uuid4())
442
+
443
+ is_agentic = requested_model in TOOL_CAPABLE_MODELS
444
+ request_id = str(uuid.uuid4())[:8]
445
+
446
+ # Build message history
447
+ final_messages = []
448
+
449
+ # Add agentic system prompt for tool-capable models
450
+ if is_agentic:
451
+ # Check if there's already a system message
452
+ has_system = any(m.get("role") == "system" for m in messages)
453
+ if has_system:
454
+ # Prepend agentic prompt to existing system message
455
+ for m in messages:
456
+ if m["role"] == "system":
457
+ final_messages.append({
458
+ "role": "system",
459
+ "content": AGENTIC_SYSTEM_PROMPT + "\n\nAdditional instructions:\n" + m["content"]
460
+ })
461
+ else:
462
+ final_messages.append(m)
463
+ else:
464
+ final_messages.append({"role": "system", "content": AGENTIC_SYSTEM_PROMPT})
465
+ final_messages.extend(messages)
466
+ else:
467
+ final_messages = list(messages)
468
+
469
+ # Save the user's message to DB
470
+ user_msg = next((m for m in reversed(messages) if m.get("role") == "user"), None)
471
+ if user_msg:
472
+ await save_message(session_id, "user", user_msg.get("content", ""))
473
+
474
+ if not stream:
475
+ # Non-streaming: simple completion
476
+ try:
477
+ kwargs = {"model": requested_model, "messages": final_messages}
478
+ if is_agentic:
479
+ kwargs["tools"] = TOOLS
480
+ kwargs["tool_choice"] = "auto"
481
+ response = await nim_client.chat.completions.create(**kwargs)
482
+ content = response.choices[0].message.content or ""
483
+ await save_message(session_id, "assistant", content)
484
+ return JSONResponse({
485
+ "id": f"chatcmpl-{request_id}",
486
+ "object": "chat.completion",
487
+ "created": int(time.time()),
488
+ "model": requested_model,
489
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
490
+ })
491
+ except Exception as e:
492
+ return JSONResponse({"error": {"message": str(e), "type": "internal_error"}}, status_code=500)
493
+
494
+ # Streaming + agentic loop
495
+ async def generate() -> AsyncIterator[str]:
496
+ nonlocal final_messages
497
+ try:
498
+ for round_num in range(MAX_TOOL_ROUNDS + 1):
499
+ kwargs = {"model": requested_model, "messages": final_messages, "stream": True}
500
+ if is_agentic:
501
+ kwargs["tools"] = TOOLS
502
+ kwargs["tool_choice"] = "auto"
503
+
504
+ # Collect streamed response
505
+ full_content = ""
506
+ tool_calls_raw = {} # index -> {id, name, arguments_str}
507
+
508
+ async for chunk in await nim_client.chat.completions.create(**kwargs):
509
+ choice = chunk.choices[0] if chunk.choices else None
510
+ if not choice:
511
+ continue
512
+ delta = choice.delta
513
+
514
+ # Stream text content to client
515
+ if delta and delta.content:
516
+ full_content += delta.content
517
+ yield make_chunk(request_id, requested_model, delta.content)
518
+
519
+ # Collect tool calls
520
+ if delta and delta.tool_calls:
521
+ for tc in delta.tool_calls:
522
+ idx = tc.index
523
+ if idx not in tool_calls_raw:
524
+ tool_calls_raw[idx] = {
525
+ "id": tc.id or f"call_{uuid.uuid4().hex[:8]}",
526
+ "name": tc.function.name if tc.function and tc.function.name else "",
527
+ "arguments": ""
528
+ }
529
+ if tc.function and tc.function.name:
530
+ tool_calls_raw[idx]["name"] = tc.function.name
531
+ if tc.id:
532
+ tool_calls_raw[idx]["id"] = tc.id
533
+ if tc.function and tc.function.arguments:
534
+ tool_calls_raw[idx]["arguments"] += tc.function.arguments
535
+
536
+ # Check for finish
537
+ if choice.finish_reason == "stop":
538
+ break
539
+ if choice.finish_reason == "tool_calls":
540
+ break
541
+
542
+ # If no tool calls, we're done
543
+ if not tool_calls_raw:
544
+ await save_message(session_id, "assistant", full_content)
545
+ yield make_chunk(request_id, requested_model, finish_reason="stop")
546
+ yield "data: [DONE]\n\n"
547
+ return
548
+
549
+ # Execute tool calls
550
+ tool_calls_list = []
551
+ for idx in sorted(tool_calls_raw.keys()):
552
+ tc = tool_calls_raw[idx]
553
+ tool_calls_list.append({
554
+ "id": tc["id"],
555
+ "type": "function",
556
+ "function": {"name": tc["name"], "arguments": tc["arguments"]}
557
+ })
558
+
559
+ # Add assistant message with tool calls to history
560
+ assistant_msg = {"role": "assistant", "content": full_content or None, "tool_calls": tool_calls_list}
561
+ final_messages.append(assistant_msg)
562
+
563
+ # Execute each tool and add results
564
+ for tc in tool_calls_list:
565
+ func_name = tc["function"]["name"]
566
+ try:
567
+ func_args = json.loads(tc["function"]["arguments"])
568
+ except json.JSONDecodeError:
569
+ func_args = {}
570
+
571
+ # Show tool execution to user
572
+ yield make_chunk(request_id, requested_model, f"\n\n🔧 **{func_name}**")
573
+ if func_name == "run_bash" and "command" in func_args:
574
+ yield make_chunk(request_id, requested_model, f": `{func_args['command']}`\n")
575
+ elif func_name == "read_file" and "path" in func_args:
576
+ yield make_chunk(request_id, requested_model, f": `{func_args['path']}`\n")
577
+ elif func_name == "write_file" and "path" in func_args:
578
+ yield make_chunk(request_id, requested_model, f": `{func_args['path']}`\n")
579
+ elif func_name == "list_directory":
580
+ yield make_chunk(request_id, requested_model, f": `{func_args.get('path', '.')}`\n")
581
+ elif func_name == "grep_search":
582
+ yield make_chunk(request_id, requested_model, f": `{func_args.get('pattern', '')}`\n")
583
+ else:
584
+ yield make_chunk(request_id, requested_model, "\n")
585
+
586
+ # Execute the tool
587
+ result = execute_tool(func_name, func_args)
588
+
589
+ # Show truncated result to user
590
+ preview = result[:500] + ("..." if len(result) > 500 else "")
591
+ yield make_chunk(request_id, requested_model, f"```\n{preview}\n```\n")
592
+
593
+ # Add tool result to message history
594
+ final_messages.append({
595
+ "role": "tool",
596
+ "tool_call_id": tc["id"],
597
+ "content": result,
598
+ })
599
+
600
+ await save_message(session_id, "tool", result, tool_call_id=tc["id"])
601
+
602
+ # Continue the agentic loop (model processes tool results)
603
+
604
+ # If we hit max rounds, finish
605
+ yield make_chunk(request_id, requested_model, "\n\n⚠️ Reached maximum tool call rounds.")
606
+ yield make_chunk(request_id, requested_model, finish_reason="stop")
607
+ yield "data: [DONE]\n\n"
608
+
609
+ except Exception as e:
610
+ error_msg = f"\n\n❌ Error: {str(e)}"
611
+ yield make_chunk(request_id, requested_model, error_msg)
612
+ yield make_chunk(request_id, requested_model, finish_reason="stop")
613
+ yield "data: [DONE]\n\n"
614
+
615
+ return StreamingResponse(
616
+ generate(),
617
+ media_type="text/event-stream",
618
+ headers={
619
+ "Cache-Control": "no-cache",
620
+ "X-Accel-Buffering": "no",
621
+ "Connection": "keep-alive",
622
+ },
623
+ )
624
+
625
+
626
+ # ---------------------------------------------------------------------------
627
+ # /v1/models — Model listing
628
+ # ---------------------------------------------------------------------------
629
+
630
+ @app.get("/v1/models")
631
+ async def list_models(authorization: str = Header(None)):
632
+ auth(authorization)
633
+ models = []
634
+ for model_id, display_name in ALL_MODELS.items():
635
+ models.append({
636
+ "id": model_id,
637
+ "object": "model",
638
+ "created": 1700000000,
639
+ "owned_by": "nvidia-nim",
640
+ "permission": [],
641
+ "root": model_id,
642
+ "parent": None,
643
+ })
644
+ return {"object": "list", "data": models}
645
+
646
+
647
+ # ---------------------------------------------------------------------------
648
+ # /health — Health check
649
+ # ---------------------------------------------------------------------------
650
+
651
+ @app.get("/health")
652
+ async def health():
653
+ return {
654
+ "status": "ok",
655
+ "workspace": WORKSPACE_DIR,
656
+ "workspace_exists": Path(WORKSPACE_DIR).exists(),
657
+ "db_connected": db_pool is not None,
658
+ "models_count": len(ALL_MODELS),
659
+ }
660
+
661
+
662
+ # ---------------------------------------------------------------------------
663
+ # Entrypoint
664
+ # ---------------------------------------------------------------------------
665
+
666
+ if __name__ == "__main__":
667
+ import uvicorn
668
+ uvicorn.run(app, host="0.0.0.0", port=7860)
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ fastapi==0.115.12
2
+ uvicorn[standard]==0.34.2
3
+ openai==1.86.0
4
+ asyncpg==0.30.0
5
+ anyio==4.9.0