Files

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)