| import asyncio |
| import json |
| import os |
| import time |
| from typing import Any |
| from urllib.parse import urlparse |
|
|
| import requests |
| from fastapi import FastAPI, Request |
| from fastapi.responses import JSONResponse |
|
|
| app = FastAPI() |
|
|
| MODAL_ENDPOINT = os.environ["MODAL_ENDPOINT"] |
| MODAL_KEY = os.environ.get("MODAL_KEY", "") |
| MODAL_SECRET = os.environ.get("MODAL_SECRET", "") |
| WEBHOOK_SECRET = os.environ.get("WEBHOOK_SECRET", "") |
| MAX_TOOL_ROUNDS = int(os.environ.get("MAX_TOOL_ROUNDS", "5")) |
| SESSION_TTL = int(os.environ.get("SESSION_TTL", "3600")) |
| MAX_SESSIONS = int(os.environ.get("MAX_SESSIONS", "1000")) |
| MAX_HISTORY = int(os.environ.get("MAX_HISTORY", "50")) |
|
|
| sessions: dict[int, list[dict]] = {} |
| session_timestamps: dict[int, float] = {} |
|
|
| SYSTEM_PROMPT = """You are a helpful AI assistant with access to tools. Use them when needed to answer questions accurately. |
| |
| Available tools: |
| - web_search: Search the web for information |
| - web_fetch: Fetch content from a URL |
| - browser_navigate: Navigate to a URL and return page content |
| |
| When you need to use a tool, output a JSON tool call in this format: |
| {"tool": "tool_name", "args": {"arg1": "value1"}} |
| |
| After receiving tool results, continue reasoning and provide a final answer. |
| When you have enough information, provide your final answer directly without tool calls.""" |
|
|
| BLOCKED_HOSTS = { |
| "169.254.169.254", |
| "metadata.google.internal", |
| "100.100.100.200", |
| "localhost", |
| "127.0.0.1", |
| "0.0.0.0", |
| } |
|
|
|
|
| def _is_safe_url(url: str) -> bool: |
| try: |
| parsed = urlparse(url) |
| if parsed.scheme not in ("http", "https"): |
| return False |
| hostname = parsed.hostname or "" |
| if hostname in BLOCKED_HOSTS: |
| return False |
| if hostname.startswith("10.") or hostname.startswith("172."): |
| parts = hostname.split(".") |
| if len(parts) == 4: |
| try: |
| if parts[0] == "172" and 16 <= int(parts[1]) <= 31: |
| return False |
| if parts[0] == "10": |
| return False |
| except ValueError: |
| pass |
| if hostname.startswith("192.168."): |
| return False |
| if hostname.startswith("fc") or hostname.startswith("fd"): |
| return False |
| return True |
| except Exception: |
| return False |
|
|
|
|
| def _strip_code_fences(text: str) -> str: |
| stripped = text.strip() |
| if stripped.startswith("```"): |
| first_newline = stripped.index("\n") if "\n" in stripped else len(stripped) |
| stripped = stripped[first_newline + 1 :] |
| if stripped.endswith("```"): |
| stripped = stripped[:-3] |
| return stripped.strip() |
|
|
|
|
| def execute_tool(name: str, args: dict[str, Any]) -> str: |
| if name == "web_search": |
| return _web_search(args.get("query", "")) |
| elif name == "web_fetch": |
| return _web_fetch(args.get("url", "")) |
| elif name == "browser_navigate": |
| return _browser_navigate(args.get("url", "")) |
| return f"Unknown tool: {name}" |
|
|
|
|
| def _web_search(query: str) -> str: |
| try: |
| resp = requests.get( |
| "https://api.duckduckgo.com/", |
| params={"q": query, "format": "json"}, |
| timeout=10, |
| ) |
| data = resp.json() |
| results = [] |
| for r in data.get("RelatedTopics", [])[:5]: |
| if isinstance(r, dict) and "Text" in r: |
| results.append(r["Text"]) |
| return "\n".join(results) if results else "No results found" |
| except Exception: |
| return "Search error" |
|
|
|
|
| def _web_fetch(url: str) -> str: |
| if not _is_safe_url(url): |
| return "Error: URL not allowed" |
| try: |
| resp = requests.get(url, timeout=15, headers={"User-Agent": "Mozilla/5.0"}) |
| return resp.text[:5000] |
| except Exception: |
| return "Fetch error" |
|
|
|
|
| def _browser_navigate(url: str) -> str: |
| if not _is_safe_url(url): |
| return "Error: URL not allowed" |
| try: |
| from playwright.sync_api import sync_playwright |
|
|
| with sync_playwright() as p: |
| browser = p.chromium.launch(headless=True) |
| page = browser.new_page() |
| page.goto(url, timeout=30000, wait_until="domcontentloaded") |
| content = page.content() |
| browser.close() |
| return content[:5000] |
| except Exception: |
| return "Browser error" |
|
|
|
|
| def _cleanup_sessions(): |
| now = time.time() |
| expired = [cid for cid, ts in session_timestamps.items() if now - ts > SESSION_TTL] |
| for cid in expired: |
| sessions.pop(cid, None) |
| session_timestamps.pop(cid, None) |
|
|
|
|
| def get_session(chat_id: int) -> list[dict]: |
| _cleanup_sessions() |
| if len(sessions) >= MAX_SESSIONS and chat_id not in sessions: |
| oldest = min(session_timestamps, key=session_timestamps.get) |
| sessions.pop(oldest, None) |
| session_timestamps.pop(oldest, None) |
| session_timestamps[chat_id] = time.time() |
| return sessions.setdefault(chat_id, []) |
|
|
|
|
| def call_modal(messages: list[dict]) -> str: |
| headers = {"Content-Type": "application/json"} |
| if MODAL_KEY: |
| headers["Modal-Key"] = MODAL_KEY |
| headers["Modal-Secret"] = MODAL_SECRET |
|
|
| resp = requests.post( |
| MODAL_ENDPOINT, |
| headers=headers, |
| json={ |
| "model": "llm", |
| "messages": messages, |
| "max_tokens": 4096, |
| }, |
| timeout=300, |
| ) |
| resp.raise_for_status() |
| return resp.json()["choices"][0]["message"]["content"] |
|
|
|
|
| def agentic_loop(chat_id: int, user_message: str) -> str: |
| history = get_session(chat_id) |
| history.append({"role": "user", "content": user_message}) |
|
|
| if len(history) > MAX_HISTORY: |
| history[:] = history[-MAX_HISTORY:] |
|
|
| messages = [{"role": "system", "content": SYSTEM_PROMPT}] + history |
|
|
| for _ in range(MAX_TOOL_ROUNDS): |
| content = call_modal(messages) |
|
|
| cleaned = _strip_code_fences(content) |
| try: |
| parsed = json.loads(cleaned) |
| if isinstance(parsed, dict) and "tool" in parsed: |
| tool_name = parsed["tool"] |
| tool_args = parsed.get("args", {}) |
| result = execute_tool(tool_name, tool_args) |
|
|
| messages.append({"role": "assistant", "content": content}) |
| messages.append( |
| {"role": "tool", "content": result[:3000], "name": tool_name} |
| ) |
| continue |
| except (json.JSONDecodeError, TypeError): |
| pass |
|
|
| history.append({"role": "assistant", "content": content}) |
| return content |
|
|
| history.append({"role": "assistant", "content": "Max tool rounds reached."}) |
| return "Max tool rounds reached." |
|
|
|
|
| def _handle_webhook(data: dict) -> dict: |
| chat_id = data.get("chat_id") |
| message = data.get("message") |
|
|
| if not chat_id or not message: |
| return {"error": "chat_id and message required"} |
|
|
| try: |
| response = agentic_loop(chat_id, message) |
| return {"response": response} |
| except Exception: |
| return {"error": "Internal error"} |
|
|
|
|
| @app.post("/webhook") |
| async def webhook(request: Request): |
| if WEBHOOK_SECRET: |
| auth = request.headers.get("Authorization", "") |
| if auth != f"Bearer {WEBHOOK_SECRET}": |
| return JSONResponse({"error": "Unauthorized"}, status_code=401) |
|
|
| data = await request.json() |
| result = await asyncio.to_thread(_handle_webhook, data) |
|
|
| if "error" in result: |
| return JSONResponse(result, status_code=500) |
| return result |
|
|
|
|
| @app.get("/health") |
| async def health(): |
| return {"status": "ok"} |
|
|
