576 lines
23 KiB
Python
576 lines
23 KiB
Python
|
import asyncio, json, os, subprocess, sys, tempfile, re, textwrap
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Optional
|
||
|
|
from urllib.parse import urlparse, quote
|
||
|
|
|
||
|
|
import uvicorn
|
||
|
|
import aiohttp
|
||
|
|
import aiofiles
|
||
|
|
from fastapi import FastAPI, HTTPException, Request
|
||
|
|
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
|
||
|
|
from fastapi.staticfiles import StaticFiles
|
||
|
|
|
||
|
|
app = FastAPI(title="OpenDAN Web")
|
||
|
|
|
||
|
|
LLAMA_CHAT = "http://127.0.0.1:8081"
|
||
|
|
LLAMA_EMBED = "http://127.0.0.1:8082"
|
||
|
|
MYAI = Path.home() / "myai"
|
||
|
|
MAX_HISTORY = 50
|
||
|
|
conversations: dict[str, list[dict]] = {}
|
||
|
|
proposals: dict[str, dict] = {}
|
||
|
|
|
||
|
|
TOOL_DEFINITIONS = [
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "system.shell.exec",
|
||
|
|
"description": "Execute a shell command in Linux bash. Use for running programs, system administration, file operations.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"command": {"type": "string", "description": "The bash command to execute"}},
|
||
|
|
"required": ["command"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "system.code_interpreter",
|
||
|
|
"description": "Execute Python code in an isolated environment. Use for data analysis, calculations, scripting.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"code": {"type": "string", "description": "Python code to execute"}},
|
||
|
|
"required": ["code"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "agent.workspace.read_file",
|
||
|
|
"description": "Read a file from the user's filesystem.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"path": {"type": "string", "description": "Absolute path to the file"}},
|
||
|
|
"required": ["path"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "agent.workspace.write_file",
|
||
|
|
"description": "Write content to a file. Creates parent directories if needed.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"path": {"type": "string", "description": "Absolute path to the file"},
|
||
|
|
"content": {"type": "string", "description": "Content to write"}
|
||
|
|
},
|
||
|
|
"required": ["path", "content"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "agent.workspace.list_dir",
|
||
|
|
"description": "List files and directories in a given path.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"path": {"type": "string", "description": "Directory path to list"}},
|
||
|
|
"required": ["path"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "web.search",
|
||
|
|
"description": "Search the web using DuckDuckGo. Use for finding current information.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {"query": {"type": "string", "description": "Search query"}},
|
||
|
|
"required": ["query"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "system.now",
|
||
|
|
"description": "Get the current date and time.",
|
||
|
|
"parameters": {"type": "object", "properties": {}}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
]
|
||
|
|
|
||
|
|
TOOL_IMPLS = {}
|
||
|
|
|
||
|
|
async def tool_shell_exec(command: str) -> dict:
|
||
|
|
try:
|
||
|
|
proc = await asyncio.create_subprocess_shell(
|
||
|
|
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||
|
|
cwd=str(MYAI))
|
||
|
|
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=30)
|
||
|
|
out = stdout.decode(errors="replace").strip()
|
||
|
|
err = stderr.decode(errors="replace").strip()
|
||
|
|
return {"status": "success" if proc.returncode == 0 else "error",
|
||
|
|
"stdout": out[:10000], "stderr": err[:5000],
|
||
|
|
"returncode": proc.returncode}
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
return {"status": "error", "error": "Command timed out after 30s"}
|
||
|
|
except Exception as e:
|
||
|
|
return {"status": "error", "error": str(e)}
|
||
|
|
|
||
|
|
async def tool_code_interpreter(code: str) -> dict:
|
||
|
|
tmp_dir = MYAI / "tmp_code"
|
||
|
|
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
script = tmp_dir / "_web_exec.py"
|
||
|
|
async with aiofiles.open(script, "w") as f:
|
||
|
|
await f.write(code)
|
||
|
|
try:
|
||
|
|
proc = await asyncio.create_subprocess_exec(
|
||
|
|
sys.executable, str(script),
|
||
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||
|
|
cwd=str(tmp_dir))
|
||
|
|
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=30)
|
||
|
|
out = stdout.decode(errors="replace").strip()
|
||
|
|
err = stderr.decode(errors="replace").strip()
|
||
|
|
return {"status": "success" if proc.returncode == 0 else "error",
|
||
|
|
"stdout": out[:10000], "stderr": err[:5000],
|
||
|
|
"returncode": proc.returncode}
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
return {"status": "error", "error": "Execution timed out after 30s"}
|
||
|
|
except Exception as e:
|
||
|
|
return {"status": "error", "error": str(e)}
|
||
|
|
finally:
|
||
|
|
if script.exists():
|
||
|
|
script.unlink()
|
||
|
|
|
||
|
|
async def tool_read_file(path: str) -> dict:
|
||
|
|
p = Path(path).expanduser().resolve()
|
||
|
|
if not p.exists():
|
||
|
|
return {"status": "error", "error": f"File not found: {path}"}
|
||
|
|
if not p.is_file():
|
||
|
|
return {"status": "error", "error": f"Not a file: {path}"}
|
||
|
|
try:
|
||
|
|
async with aiofiles.open(p, "r") as f:
|
||
|
|
content = await f.read()
|
||
|
|
return {"status": "success", "content": content[:50000]}
|
||
|
|
except Exception as e:
|
||
|
|
return {"status": "error", "error": str(e)}
|
||
|
|
|
||
|
|
async def tool_write_file(path: str, content: str) -> dict:
|
||
|
|
p = Path(path).expanduser().resolve()
|
||
|
|
p.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
try:
|
||
|
|
async with aiofiles.open(p, "w") as f:
|
||
|
|
await f.write(content)
|
||
|
|
return {"status": "success", "path": str(p)}
|
||
|
|
except Exception as e:
|
||
|
|
return {"status": "error", "error": str(e)}
|
||
|
|
|
||
|
|
async def tool_list_dir(path: str) -> dict:
|
||
|
|
p = Path(path).expanduser().resolve()
|
||
|
|
if not p.exists():
|
||
|
|
return {"status": "error", "error": f"Path not found: {path}"}
|
||
|
|
if not p.is_dir():
|
||
|
|
return {"status": "error", "error": f"Not a directory: {path}"}
|
||
|
|
try:
|
||
|
|
items = []
|
||
|
|
for entry in sorted(p.iterdir()):
|
||
|
|
items.append({"name": entry.name, "type": "dir" if entry.is_dir() else "file",
|
||
|
|
"size": entry.stat().st_size if entry.is_file() else 0})
|
||
|
|
return {"status": "success", "path": str(p), "items": items}
|
||
|
|
except Exception as e:
|
||
|
|
return {"status": "error", "error": str(e)}
|
||
|
|
|
||
|
|
async def tool_web_search(query: str) -> dict:
|
||
|
|
try:
|
||
|
|
loop = asyncio.get_event_loop()
|
||
|
|
results = await loop.run_in_executor(None, _search_ddg, query)
|
||
|
|
return {"status": "success", "query": query, "results": results}
|
||
|
|
except Exception as e:
|
||
|
|
return {"status": "error", "error": str(e)}
|
||
|
|
|
||
|
|
def _search_ddg(query: str, max_results: int = 5) -> list[dict]:
|
||
|
|
from curl_cffi import requests
|
||
|
|
from bs4 import BeautifulSoup
|
||
|
|
r = requests.get(
|
||
|
|
"https://lite.duckduckgo.com/lite/",
|
||
|
|
params={"q": query, "kp": "-2"},
|
||
|
|
impersonate="chrome",
|
||
|
|
timeout=15
|
||
|
|
)
|
||
|
|
soup = BeautifulSoup(r.text, "lxml")
|
||
|
|
rows = soup.select("tr")
|
||
|
|
out = []
|
||
|
|
i = 0
|
||
|
|
while i < len(rows) and len(out) < max_results:
|
||
|
|
a = rows[i].select_one("a.result-link")
|
||
|
|
if not a:
|
||
|
|
i += 1
|
||
|
|
continue
|
||
|
|
title = a.get_text(strip=True)
|
||
|
|
raw_url = a.get("href", "")
|
||
|
|
snippet = ""
|
||
|
|
if i + 1 < len(rows):
|
||
|
|
snip_td = rows[i + 1].select_one("td.result-snippet")
|
||
|
|
if snip_td:
|
||
|
|
snippet = snip_td.get_text(strip=True)
|
||
|
|
out.append({"title": title[:200], "snippet": snippet[:500], "url": raw_url[:300]})
|
||
|
|
i += 3
|
||
|
|
if not out:
|
||
|
|
out.append({"title": "No results", "snippet": f"No search results for: {query}", "url": ""})
|
||
|
|
return out
|
||
|
|
|
||
|
|
async def tool_system_now() -> dict:
|
||
|
|
from datetime import datetime
|
||
|
|
return {"status": "success", "datetime": datetime.now().isoformat()}
|
||
|
|
|
||
|
|
TOOL_IMPLS = {
|
||
|
|
"system.shell.exec": tool_shell_exec,
|
||
|
|
"system.code_interpreter": tool_code_interpreter,
|
||
|
|
"agent.workspace.read_file": tool_read_file,
|
||
|
|
"agent.workspace.write_file": tool_write_file,
|
||
|
|
"agent.workspace.list_dir": tool_list_dir,
|
||
|
|
"web.search": tool_web_search,
|
||
|
|
"system.now": tool_system_now,
|
||
|
|
}
|
||
|
|
|
||
|
|
async def call_llm(messages: list[dict], tools: list[dict] = None,
|
||
|
|
tool_choice: str = "auto", model: str = None,
|
||
|
|
max_tokens: int = 4096, temperature: float = 0.7) -> dict:
|
||
|
|
body = {"messages": messages, "max_tokens": max_tokens, "temperature": temperature}
|
||
|
|
if model:
|
||
|
|
body["model"] = model
|
||
|
|
if tools:
|
||
|
|
body["tools"] = tools
|
||
|
|
body["tool_choice"] = tool_choice
|
||
|
|
async with aiohttp.ClientSession() as sess:
|
||
|
|
async with sess.post(f"{LLAMA_CHAT}/v1/chat/completions",
|
||
|
|
json=body, timeout=aiohttp.ClientTimeout(total=120)) as resp:
|
||
|
|
if resp.status != 200:
|
||
|
|
txt = await resp.text()
|
||
|
|
raise HTTPException(status_code=resp.status, detail=txt[:500])
|
||
|
|
return await resp.json()
|
||
|
|
|
||
|
|
async def stream_llm(messages: list[dict], tools: list[dict] = None,
|
||
|
|
tool_choice: str = "auto", model: str = None):
|
||
|
|
body = {"messages": messages, "max_tokens": 4096, "temperature": 0.7,
|
||
|
|
"stream": True}
|
||
|
|
if model:
|
||
|
|
body["model"] = model
|
||
|
|
if tools:
|
||
|
|
body["tools"] = tools
|
||
|
|
body["tool_choice"] = tool_choice
|
||
|
|
async with aiohttp.ClientSession() as sess:
|
||
|
|
async with sess.post(f"{LLAMA_CHAT}/v1/chat/completions",
|
||
|
|
json=body, timeout=aiohttp.ClientTimeout(total=120)) as resp:
|
||
|
|
if resp.status != 200:
|
||
|
|
yield f"data: {json.dumps({'type': 'error', 'detail': await resp.text()})}\n\n"
|
||
|
|
return
|
||
|
|
buffer = ""
|
||
|
|
async for chunk in resp.content.iter_chunks():
|
||
|
|
data = chunk[0].decode(errors="replace")
|
||
|
|
buffer += data
|
||
|
|
while "\n" in buffer:
|
||
|
|
line, buffer = buffer.split("\n", 1)
|
||
|
|
line = line.strip()
|
||
|
|
if not line:
|
||
|
|
continue
|
||
|
|
if line.startswith("data: "):
|
||
|
|
payload = line[6:]
|
||
|
|
if payload == "[DONE]":
|
||
|
|
yield f"data: {json.dumps({'type': 'done'})}\n\n"
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
obj = json.loads(payload)
|
||
|
|
yield f"data: {json.dumps({'type': 'token', 'data': obj})}\n\n"
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def get_thread(thread_id: str) -> list[dict]:
|
||
|
|
if thread_id not in conversations:
|
||
|
|
conversations[thread_id] = [
|
||
|
|
{"role": "system", "content": textwrap.dedent("""\
|
||
|
|
You are an AI assistant with full system access.
|
||
|
|
Use tools to answer questions requiring data or actions.
|
||
|
|
Never describe what tools you would use — just use them.
|
||
|
|
For simple greetings or chat, reply directly without tools.
|
||
|
|
After your final answer, suggest 3 brief follow-up questions
|
||
|
|
the user might want to ask next. Wrap them as a JSON object
|
||
|
|
on its own line: {"suggestions": ["q1?", "q2?", "q3?"]}.
|
||
|
|
""").strip()}
|
||
|
|
]
|
||
|
|
return conversations[thread_id]
|
||
|
|
|
||
|
|
def trim_thread(thread: list[dict]):
|
||
|
|
if len(thread) > MAX_HISTORY:
|
||
|
|
system = [m for m in thread if m["role"] == "system"]
|
||
|
|
rest = [m for m in thread if m["role"] != "system"]
|
||
|
|
thread[:] = system + rest[-(MAX_HISTORY - len(system)):]
|
||
|
|
|
||
|
|
@app.get("/health")
|
||
|
|
async def health():
|
||
|
|
async with aiohttp.ClientSession() as sess:
|
||
|
|
try:
|
||
|
|
async with sess.get(f"{LLAMA_CHAT}/health",
|
||
|
|
timeout=aiohttp.ClientTimeout(total=3)) as r:
|
||
|
|
chat = r.status == 200
|
||
|
|
except:
|
||
|
|
chat = False
|
||
|
|
try:
|
||
|
|
async with sess.get(f"{LLAMA_EMBED}/health",
|
||
|
|
timeout=aiohttp.ClientTimeout(total=3)) as r:
|
||
|
|
embed = r.status == 200
|
||
|
|
except:
|
||
|
|
embed = False
|
||
|
|
return {"status": "ok", "chat": chat, "embeddings": embed, "tools": list(TOOL_IMPLS.keys())}
|
||
|
|
|
||
|
|
@app.get("/", response_class=HTMLResponse)
|
||
|
|
async def index():
|
||
|
|
p = Path(__file__).parent / "templates" / "index.html"
|
||
|
|
return HTMLResponse(p.read_text())
|
||
|
|
|
||
|
|
@app.post("/v1/chat")
|
||
|
|
async def chat(request: Request):
|
||
|
|
body = await request.json()
|
||
|
|
msg = body.get("message", "").strip()
|
||
|
|
thread_id = body.get("thread_id", "default")
|
||
|
|
if not msg:
|
||
|
|
raise HTTPException(400, "message is required")
|
||
|
|
thread = get_thread(thread_id)
|
||
|
|
thread.append({"role": "user", "content": msg})
|
||
|
|
trim_thread(thread)
|
||
|
|
return await _process_tools(thread, thread_id)
|
||
|
|
|
||
|
|
@app.post("/v1/chat/stream")
|
||
|
|
async def chat_stream(request: Request):
|
||
|
|
body = await request.json()
|
||
|
|
thread_id = body.get("thread_id", "default")
|
||
|
|
proposal_id = body.get("proposal_id")
|
||
|
|
decisions = body.get("decisions")
|
||
|
|
if proposal_id and decisions is not None:
|
||
|
|
return StreamingResponse(_execute_proposal_stream(proposal_id, decisions, thread_id),
|
||
|
|
media_type="text/event-stream")
|
||
|
|
msg = body.get("message", "").strip()
|
||
|
|
if not msg:
|
||
|
|
raise HTTPException(400, "message is required")
|
||
|
|
thread = get_thread(thread_id)
|
||
|
|
thread.append({"role": "user", "content": msg})
|
||
|
|
trim_thread(thread)
|
||
|
|
return StreamingResponse(_process_stream(thread, thread_id),
|
||
|
|
media_type="text/event-stream")
|
||
|
|
|
||
|
|
@app.delete("/v1/conversation/{thread_id}")
|
||
|
|
async def clear_conversation(thread_id: str):
|
||
|
|
conversations.pop(thread_id, None)
|
||
|
|
return {"status": "ok"}
|
||
|
|
|
||
|
|
@app.get("/v1/tools")
|
||
|
|
async def list_tools():
|
||
|
|
return TOOL_DEFINITIONS
|
||
|
|
|
||
|
|
@app.post("/v1/chat/propose")
|
||
|
|
async def chat_propose(request: Request):
|
||
|
|
body = await request.json()
|
||
|
|
msg = body.get("message", "").strip()
|
||
|
|
thread_id = body.get("thread_id", "default")
|
||
|
|
if not msg:
|
||
|
|
raise HTTPException(400, "message is required")
|
||
|
|
thread = get_thread(thread_id)
|
||
|
|
temp_msgs = thread + [{"role": "user", "content": msg}]
|
||
|
|
resp = await call_llm(temp_msgs, tools=TOOL_DEFINITIONS, tool_choice="auto")
|
||
|
|
msg_obj = resp["choices"][0]["message"]
|
||
|
|
tool_calls = msg_obj.get("tool_calls")
|
||
|
|
if not tool_calls:
|
||
|
|
thread.append({"role": "user", "content": msg})
|
||
|
|
content = msg_obj.get("content", "") or ""
|
||
|
|
thread.append({"role": "assistant", "content": content})
|
||
|
|
trim_thread(thread)
|
||
|
|
return {"response": content, "thread_id": thread_id, "tool_calls": []}
|
||
|
|
pid = f"{thread_id}_{len(thread)}_{int(asyncio.get_event_loop().time()*1000000)}"
|
||
|
|
raw = []
|
||
|
|
for tc in tool_calls:
|
||
|
|
try:
|
||
|
|
a = json.loads(tc["function"]["arguments"])
|
||
|
|
except:
|
||
|
|
a = {}
|
||
|
|
raw.append({"tool_call_id": tc["id"], "name": tc["function"]["name"], "args": a})
|
||
|
|
proposals[pid] = {
|
||
|
|
"thread_id": thread_id, "user_content": msg,
|
||
|
|
"assistant_content": msg_obj.get("content") or "",
|
||
|
|
"tool_calls": tool_calls,
|
||
|
|
}
|
||
|
|
return {"proposal_id": pid, "tool_calls": raw, "content": msg_obj.get("content") or "",
|
||
|
|
"thread_id": thread_id}
|
||
|
|
|
||
|
|
async def _process_tools(thread: list[dict], thread_id: str, depth: int = 0):
|
||
|
|
if depth > 5:
|
||
|
|
return {"response": "Tool call depth exceeded (max 5).", "thread_id": thread_id}
|
||
|
|
resp = await call_llm(thread, tools=TOOL_DEFINITIONS, tool_choice="auto")
|
||
|
|
msg = resp["choices"][0]["message"]
|
||
|
|
tool_calls = msg.get("tool_calls")
|
||
|
|
if tool_calls:
|
||
|
|
thread.append({"role": "assistant", "content": msg.get("content") or "",
|
||
|
|
"tool_calls": tool_calls})
|
||
|
|
results = []
|
||
|
|
for tc in tool_calls:
|
||
|
|
fn = tc["function"]
|
||
|
|
name = fn["name"]
|
||
|
|
try:
|
||
|
|
args = json.loads(fn["arguments"])
|
||
|
|
except:
|
||
|
|
args = {}
|
||
|
|
impl = TOOL_IMPLS.get(name)
|
||
|
|
if impl:
|
||
|
|
result = await impl(**args)
|
||
|
|
else:
|
||
|
|
result = {"status": "error", "error": f"Unknown tool: {name}"}
|
||
|
|
results.append({"tool_call_id": tc["id"], "name": name, "result": result})
|
||
|
|
thread.append({"role": "tool", "tool_call_id": tc["id"],
|
||
|
|
"content": json.dumps(result)})
|
||
|
|
trim_thread(thread)
|
||
|
|
final = await _process_tools(thread, thread_id, depth + 1)
|
||
|
|
final["tool_results"] = results
|
||
|
|
return final
|
||
|
|
else:
|
||
|
|
content = msg.get("content", "") or ""
|
||
|
|
thread.append({"role": "assistant", "content": content})
|
||
|
|
trim_thread(thread)
|
||
|
|
return {"response": content, "thread_id": thread_id}
|
||
|
|
|
||
|
|
async def _process_stream(thread: list[dict], thread_id: str, depth: int = 0):
|
||
|
|
if depth > 5:
|
||
|
|
yield f"data: {json.dumps({'type': 'error', 'detail': 'Tool call depth exceeded (max 5)'})}\n\n"
|
||
|
|
return
|
||
|
|
collected_content = ""
|
||
|
|
collected_tool_calls = None
|
||
|
|
async for event in stream_llm(thread, tools=TOOL_DEFINITIONS, tool_choice="auto"):
|
||
|
|
yield event
|
||
|
|
try:
|
||
|
|
data = json.loads(event[6:].strip()) if event.startswith("data: ") else None
|
||
|
|
except:
|
||
|
|
data = None
|
||
|
|
if data is None:
|
||
|
|
continue
|
||
|
|
if data.get("type") == "done":
|
||
|
|
break
|
||
|
|
obj = data.get("data", {})
|
||
|
|
delta = obj.get("choices", [{}])[0].get("delta", {})
|
||
|
|
collected_content += delta.get("content", "") or ""
|
||
|
|
if "tool_calls" in delta:
|
||
|
|
if collected_tool_calls is None:
|
||
|
|
collected_tool_calls = []
|
||
|
|
for tc in delta["tool_calls"]:
|
||
|
|
idx = tc.get("index", len(collected_tool_calls))
|
||
|
|
while len(collected_tool_calls) <= idx:
|
||
|
|
collected_tool_calls.append({"id": "", "function": {"name": "", "arguments": ""}})
|
||
|
|
if tc.get("type"):
|
||
|
|
collected_tool_calls[idx]["type"] = tc["type"]
|
||
|
|
if tc.get("id"):
|
||
|
|
collected_tool_calls[idx]["id"] = tc["id"]
|
||
|
|
if tc.get("function", {}).get("name"):
|
||
|
|
collected_tool_calls[idx]["function"]["name"] += tc["function"]["name"]
|
||
|
|
if tc.get("function", {}).get("arguments"):
|
||
|
|
collected_tool_calls[idx]["function"]["arguments"] += tc["function"]["arguments"]
|
||
|
|
if collected_tool_calls:
|
||
|
|
msg = {"role": "assistant", "content": collected_content,
|
||
|
|
"tool_calls": collected_tool_calls}
|
||
|
|
thread.append(msg)
|
||
|
|
results = []
|
||
|
|
for tc in collected_tool_calls:
|
||
|
|
fn = tc["function"]
|
||
|
|
name = fn["name"]
|
||
|
|
try:
|
||
|
|
args = json.loads(fn["arguments"])
|
||
|
|
except:
|
||
|
|
args = {}
|
||
|
|
impl = TOOL_IMPLS.get(name)
|
||
|
|
yield f"data: {json.dumps({'type': 'tool_start', 'name': name, 'args': args})}\n\n"
|
||
|
|
if impl:
|
||
|
|
result = await impl(**args)
|
||
|
|
else:
|
||
|
|
result = {"status": "error", "error": f"Unknown tool: {name}"}
|
||
|
|
results.append({"tool_call_id": tc["id"], "name": name, "result": result})
|
||
|
|
yield f"data: {json.dumps({'type': 'tool_result', 'name': name, 'result': result})}\n\n"
|
||
|
|
thread.append({"role": "tool", "tool_call_id": tc["id"],
|
||
|
|
"content": json.dumps(result)})
|
||
|
|
trim_thread(thread)
|
||
|
|
async for event in _process_stream(thread, thread_id, depth + 1):
|
||
|
|
yield event
|
||
|
|
else:
|
||
|
|
if collected_content:
|
||
|
|
display, suggestions = _extract_suggestions(collected_content)
|
||
|
|
if display != collected_content:
|
||
|
|
yield f"data: {json.dumps({'type': 'suggestions', 'suggestions': suggestions})}\n\n"
|
||
|
|
thread.append({"role": "assistant", "content": collected_content})
|
||
|
|
trim_thread(thread)
|
||
|
|
|
||
|
|
def _extract_suggestions(text: str) -> tuple[str, list[str]]:
|
||
|
|
m = re.search(r'\{\s*"suggestions"\s*:\s*\[(.*?)\]\s*\}', text, re.DOTALL)
|
||
|
|
if not m:
|
||
|
|
return text, []
|
||
|
|
try:
|
||
|
|
obj = json.loads("{" + m.group(0) + "}")
|
||
|
|
suggestions = obj.get("suggestions", [])[:3]
|
||
|
|
rest = text[:m.start()].rstrip() + text[m.end():]
|
||
|
|
return rest, suggestions
|
||
|
|
except:
|
||
|
|
return text, []
|
||
|
|
|
||
|
|
async def _execute_proposal_stream(proposal_id: str, decisions: list[dict], thread_id: str):
|
||
|
|
if proposal_id not in proposals:
|
||
|
|
yield f"data: {json.dumps({'type': 'error', 'detail': 'Proposal not found'})}\n\n"
|
||
|
|
return
|
||
|
|
prop = proposals.pop(proposal_id)
|
||
|
|
thread = get_thread(thread_id)
|
||
|
|
thread.append({"role": "user", "content": prop["user_content"]})
|
||
|
|
tool_calls = prop["tool_calls"]
|
||
|
|
assistant_content = prop["assistant_content"]
|
||
|
|
thread.append({"role": "assistant", "content": assistant_content, "tool_calls": tool_calls})
|
||
|
|
decision_map = {d["tool_call_id"]: d.get("approved", False) for d in decisions}
|
||
|
|
for tc in tool_calls:
|
||
|
|
tc_id = tc["id"]
|
||
|
|
fn = tc["function"]
|
||
|
|
name = fn["name"]
|
||
|
|
try:
|
||
|
|
args = json.loads(fn["arguments"]) if isinstance(fn["arguments"], str) else fn["arguments"]
|
||
|
|
except:
|
||
|
|
args = {}
|
||
|
|
yield f"data: {json.dumps({'type': 'tool_start', 'name': name, 'args': args, 'tool_call_id': tc_id})}\n\n"
|
||
|
|
if decision_map.get(tc_id, False):
|
||
|
|
impl = TOOL_IMPLS.get(name)
|
||
|
|
if impl:
|
||
|
|
result = await impl(**args)
|
||
|
|
else:
|
||
|
|
result = {"status": "error", "error": f"Unknown tool: {name}"}
|
||
|
|
else:
|
||
|
|
result = {"status": "skipped", "reason": "Rejected by user"}
|
||
|
|
yield f"data: {json.dumps({'type': 'tool_result', 'name': name, 'result': result, 'tool_call_id': tc_id})}\n\n"
|
||
|
|
thread.append({"role": "tool", "tool_call_id": tc_id, "content": json.dumps(result)})
|
||
|
|
trim_thread(thread)
|
||
|
|
collected_content = ""
|
||
|
|
async for event in stream_llm(thread, tools=None):
|
||
|
|
yield event
|
||
|
|
try:
|
||
|
|
data = json.loads(event[6:].strip()) if event.startswith("data: ") else None
|
||
|
|
except:
|
||
|
|
data = None
|
||
|
|
if data and data.get("type") == "done":
|
||
|
|
pass
|
||
|
|
elif data and data.get("type") == "token":
|
||
|
|
delta = data.get("data", {}).get("choices", [{}])[0].get("delta", {})
|
||
|
|
collected_content += delta.get("content", "") or ""
|
||
|
|
if collected_content:
|
||
|
|
display, suggestions = _extract_suggestions(collected_content)
|
||
|
|
if display != collected_content:
|
||
|
|
yield f"data: {json.dumps({'type': 'suggestions', 'suggestions': suggestions})}\n\n"
|
||
|
|
thread.append({"role": "assistant", "content": collected_content})
|
||
|
|
trim_thread(thread)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
uvicorn.run(app, host="0.0.0.0", port=8080)
|