mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-04 16:38:49 +00:00
92d6fca323
Move extract_documents() to nanobot.utils.document as a reusable helper and call it once in AgentLoop._process_message, the single entry point for all message processing (API + all channels). This replaces the previous API-only _extract_documents() in server.py, ensuring Telegram, Feishu, Slack, WeChat, and all other channels also benefit from automatic document text extraction. Adds a configurable max_file_size guard (default 50 MB) to skip oversized files gracefully, preventing unbounded memory/CPU usage from channel-downloaded attachments. - server.py: removed _extract_documents and related imports - document.py: added extract_documents() with size limit - loop.py: calls extract_documents() at the top of _process_message - Tests updated: 70 related tests pass Made-with: Cursor
295 lines
11 KiB
Python
295 lines
11 KiB
Python
"""OpenAI-compatible HTTP API server for a fixed nanobot session.
|
|
|
|
Provides /v1/chat/completions and /v1/models endpoints.
|
|
All requests route to a single persistent API session.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import mimetypes
|
|
import re
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from aiohttp import web
|
|
from loguru import logger
|
|
|
|
from nanobot.config.paths import get_media_dir
|
|
from nanobot.utils.helpers import safe_filename
|
|
from nanobot.utils.runtime import EMPTY_FINAL_RESPONSE_MESSAGE
|
|
|
|
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10 MB
|
|
_DATA_URL_RE = re.compile(r"^data:([^;]+);base64,(.+)$", re.DOTALL)
|
|
|
|
|
|
class _FileSizeExceeded(Exception):
|
|
"""Raised when an uploaded file exceeds the size limit."""
|
|
|
|
API_SESSION_KEY = "api:default"
|
|
API_CHAT_ID = "default"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Response helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _error_json(status: int, message: str, err_type: str = "invalid_request_error") -> web.Response:
|
|
return web.json_response(
|
|
{"error": {"message": message, "type": err_type, "code": status}},
|
|
status=status,
|
|
)
|
|
|
|
|
|
def _chat_completion_response(content: str, model: str) -> dict[str, Any]:
|
|
return {
|
|
"id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": content},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
|
}
|
|
|
|
|
|
def _response_text(value: Any) -> str:
|
|
"""Normalize process_direct output to plain assistant text."""
|
|
if value is None:
|
|
return ""
|
|
if hasattr(value, "content"):
|
|
return str(getattr(value, "content") or "")
|
|
return str(value)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Upload helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _save_base64_data_url(data_url: str, media_dir: Path) -> str | None:
|
|
"""Decode a data:...;base64,... URL and save to disk."""
|
|
m = _DATA_URL_RE.match(data_url)
|
|
if not m:
|
|
return None
|
|
mime_type, b64_payload = m.group(1), m.group(2)
|
|
try:
|
|
raw = base64.b64decode(b64_payload)
|
|
except Exception:
|
|
return None
|
|
if len(raw) > MAX_FILE_SIZE:
|
|
raise _FileSizeExceeded(
|
|
f"File exceeds {MAX_FILE_SIZE // (1024 * 1024)}MB limit"
|
|
)
|
|
ext = mimetypes.guess_extension(mime_type) or ".bin"
|
|
filename = f"{uuid.uuid4().hex[:12]}{ext}"
|
|
dest = media_dir / safe_filename(filename)
|
|
dest.write_bytes(raw)
|
|
return str(dest)
|
|
|
|
|
|
def _parse_json_content(body: dict) -> tuple[str, list[str]]:
|
|
"""Parse JSON request body. Returns (text, media_paths)."""
|
|
messages = body.get("messages")
|
|
if not isinstance(messages, list) or len(messages) != 1:
|
|
raise ValueError("Only a single user message is supported")
|
|
message = messages[0]
|
|
if not isinstance(message, dict) or message.get("role") != "user":
|
|
raise ValueError("Only a single user message is supported")
|
|
|
|
user_content = message.get("content", "")
|
|
media_dir = get_media_dir("api")
|
|
media_paths: list[str] = []
|
|
|
|
if isinstance(user_content, list):
|
|
text_parts: list[str] = []
|
|
for part in user_content:
|
|
if not isinstance(part, dict):
|
|
continue
|
|
if part.get("type") == "text":
|
|
text_parts.append(part.get("text", ""))
|
|
elif part.get("type") == "image_url":
|
|
url = part.get("image_url", {}).get("url", "")
|
|
if url.startswith("data:"):
|
|
saved = _save_base64_data_url(url, media_dir)
|
|
if saved:
|
|
media_paths.append(saved)
|
|
text = " ".join(text_parts)
|
|
elif isinstance(user_content, str):
|
|
text = user_content
|
|
else:
|
|
raise ValueError("Invalid content format")
|
|
|
|
return text, media_paths
|
|
|
|
|
|
async def _parse_multipart(request: web.Request) -> tuple[str, list[str], str | None]:
|
|
"""Parse multipart/form-data. Returns (text, media_paths, session_id)."""
|
|
media_dir = get_media_dir("api")
|
|
reader = await request.multipart()
|
|
text = ""
|
|
session_id = None
|
|
media_paths: list[str] = []
|
|
|
|
while True:
|
|
part = await reader.next()
|
|
if part is None:
|
|
break
|
|
if part.name == "message":
|
|
text = (await part.read()).decode("utf-8")
|
|
elif part.name == "session_id":
|
|
session_id = (await part.read()).decode("utf-8").strip()
|
|
elif part.name == "files":
|
|
raw = await part.read()
|
|
if len(raw) > MAX_FILE_SIZE:
|
|
raise _FileSizeExceeded(f"File '{part.filename}' exceeds {MAX_FILE_SIZE // (1024*1024)}MB limit")
|
|
filename = safe_filename(part.filename or f"{uuid.uuid4().hex[:12]}.bin")
|
|
dest = media_dir / filename
|
|
dest.write_bytes(raw)
|
|
media_paths.append(str(dest))
|
|
|
|
if not text:
|
|
text = "请分析上传的文件"
|
|
|
|
return text, media_paths, session_id
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Route handlers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
async def handle_chat_completions(request: web.Request) -> web.Response:
|
|
"""POST /v1/chat/completions — supports JSON and multipart/form-data."""
|
|
content_type = request.content_type or ""
|
|
if not isinstance(content_type, str):
|
|
content_type = ""
|
|
|
|
agent_loop = request.app["agent_loop"]
|
|
timeout_s: float = request.app.get("request_timeout", 120.0)
|
|
model_name: str = request.app.get("model_name", "nanobot")
|
|
|
|
try:
|
|
if content_type.startswith("multipart/"):
|
|
text, media_paths, session_id = await _parse_multipart(request)
|
|
else:
|
|
try:
|
|
body = await request.json()
|
|
except Exception:
|
|
return _error_json(400, "Invalid JSON body")
|
|
if body.get("stream", False):
|
|
return _error_json(400, "stream=true is not supported yet. Set stream=false or omit it.")
|
|
if (requested_model := body.get("model")) and requested_model != model_name:
|
|
return _error_json(400, f"Only configured model '{model_name}' is available")
|
|
text, media_paths = _parse_json_content(body)
|
|
session_id = body.get("session_id")
|
|
except ValueError as e:
|
|
return _error_json(400, str(e))
|
|
except _FileSizeExceeded as e:
|
|
return _error_json(413, str(e), err_type="invalid_request_error")
|
|
except Exception:
|
|
logger.exception("Error parsing upload")
|
|
return _error_json(413, "File too large or invalid upload")
|
|
|
|
session_key = f"api:{session_id}" if session_id else API_SESSION_KEY
|
|
session_locks: dict[str, asyncio.Lock] = request.app["session_locks"]
|
|
session_lock = session_locks.setdefault(session_key, asyncio.Lock())
|
|
|
|
logger.info("API request session_key={} media={} text={}", session_key, len(media_paths), text[:80])
|
|
|
|
_FALLBACK = EMPTY_FINAL_RESPONSE_MESSAGE
|
|
|
|
try:
|
|
async with session_lock:
|
|
try:
|
|
response = await asyncio.wait_for(
|
|
agent_loop.process_direct(
|
|
content=text,
|
|
media=media_paths if media_paths else None,
|
|
session_key=session_key,
|
|
channel="api",
|
|
chat_id=API_CHAT_ID,
|
|
),
|
|
timeout=timeout_s,
|
|
)
|
|
response_text = _response_text(response)
|
|
|
|
if not response_text or not response_text.strip():
|
|
logger.warning("Empty response for session {}, retrying", session_key)
|
|
retry_response = await asyncio.wait_for(
|
|
agent_loop.process_direct(
|
|
content=text,
|
|
media=media_paths if media_paths else None,
|
|
session_key=session_key,
|
|
channel="api",
|
|
chat_id=API_CHAT_ID,
|
|
),
|
|
timeout=timeout_s,
|
|
)
|
|
response_text = _response_text(retry_response)
|
|
if not response_text or not response_text.strip():
|
|
logger.warning("Empty response after retry, using fallback")
|
|
response_text = _FALLBACK
|
|
|
|
except asyncio.TimeoutError:
|
|
return _error_json(504, f"Request timed out after {timeout_s}s")
|
|
except Exception:
|
|
logger.exception("Error processing request for session {}", session_key)
|
|
return _error_json(500, "Internal server error", err_type="server_error")
|
|
except Exception:
|
|
logger.exception("Unexpected API lock error for session {}", session_key)
|
|
return _error_json(500, "Internal server error", err_type="server_error")
|
|
|
|
return web.json_response(_chat_completion_response(response_text, model_name))
|
|
|
|
|
|
async def handle_models(request: web.Request) -> web.Response:
|
|
"""GET /v1/models"""
|
|
model_name = request.app.get("model_name", "nanobot")
|
|
return web.json_response({
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"id": model_name,
|
|
"object": "model",
|
|
"created": 0,
|
|
"owned_by": "nanobot",
|
|
}
|
|
],
|
|
})
|
|
|
|
|
|
async def handle_health(request: web.Request) -> web.Response:
|
|
"""GET /health"""
|
|
return web.json_response({"status": "ok"})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# App factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def create_app(agent_loop, model_name: str = "nanobot", request_timeout: float = 120.0) -> web.Application:
|
|
"""Create the aiohttp application.
|
|
|
|
Args:
|
|
agent_loop: An initialized AgentLoop instance.
|
|
model_name: Model name reported in responses.
|
|
request_timeout: Per-request timeout in seconds.
|
|
"""
|
|
app = web.Application(client_max_size=20 * 1024 * 1024) # 20MB for base64 images
|
|
app["agent_loop"] = agent_loop
|
|
app["model_name"] = model_name
|
|
app["request_timeout"] = request_timeout
|
|
app["session_locks"] = {} # per-user locks, keyed by session_key
|
|
|
|
app.router.add_post("/v1/chat/completions", handle_chat_completions)
|
|
app.router.add_get("/v1/models", handle_models)
|
|
app.router.add_get("/health", handle_health)
|
|
return app
|