266 lines
8.3 KiB
Python
266 lines
8.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import uuid
|
||
|
|
from collections.abc import AsyncIterator
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Callable
|
||
|
|
|
||
|
|
_JSONRPC_VERSION = "2.0"
|
||
|
|
_NANOBOT_API_STREAM_LIMIT = 2 * 1024 * 1024
|
||
|
|
NANOBOT_API_SOCKET = Path(
|
||
|
|
os.getenv("NANOBOT_API_SOCKET", str(Path.home() / ".nanobot" / "api.sock"))
|
||
|
|
).expanduser()
|
||
|
|
|
||
|
|
|
||
|
|
class _NanobotApiError(RuntimeError):
|
||
|
|
def __init__(self, code: int, message: str) -> None:
|
||
|
|
super().__init__(message)
|
||
|
|
self.code = code
|
||
|
|
|
||
|
|
|
||
|
|
def _jsonrpc_request(request_id: str, method: str, params: dict[str, Any]) -> dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"jsonrpc": _JSONRPC_VERSION,
|
||
|
|
"id": request_id,
|
||
|
|
"method": method,
|
||
|
|
"params": params,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _jsonrpc_notification(method: str, params: dict[str, Any]) -> dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"jsonrpc": _JSONRPC_VERSION,
|
||
|
|
"method": method,
|
||
|
|
"params": params,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
async def _open_nanobot_api_socket(
|
||
|
|
socket_path: Path = NANOBOT_API_SOCKET,
|
||
|
|
*,
|
||
|
|
stream_limit: int = _NANOBOT_API_STREAM_LIMIT,
|
||
|
|
) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
|
||
|
|
if not socket_path.exists():
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Nanobot API socket not found at {socket_path}. "
|
||
|
|
"Enable channels.api and start `nanobot gateway`."
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
return await asyncio.open_unix_connection(
|
||
|
|
path=str(socket_path),
|
||
|
|
limit=stream_limit,
|
||
|
|
)
|
||
|
|
except OSError as exc:
|
||
|
|
raise RuntimeError(f"failed to connect to Nanobot API socket: {exc}") from exc
|
||
|
|
|
||
|
|
|
||
|
|
async def _send_nanobot_api_request(
|
||
|
|
method: str,
|
||
|
|
params: dict[str, Any],
|
||
|
|
*,
|
||
|
|
timeout_seconds: float,
|
||
|
|
socket_path: Path = NANOBOT_API_SOCKET,
|
||
|
|
stream_limit: int = _NANOBOT_API_STREAM_LIMIT,
|
||
|
|
) -> Any:
|
||
|
|
request_id = str(uuid.uuid4())
|
||
|
|
reader, writer = await _open_nanobot_api_socket(socket_path, stream_limit=stream_limit)
|
||
|
|
try:
|
||
|
|
writer.write(
|
||
|
|
(
|
||
|
|
json.dumps(_jsonrpc_request(request_id, method, params), ensure_ascii=False) + "\n"
|
||
|
|
).encode("utf-8")
|
||
|
|
)
|
||
|
|
await writer.drain()
|
||
|
|
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
deadline = loop.time() + timeout_seconds
|
||
|
|
|
||
|
|
while True:
|
||
|
|
remaining = deadline - loop.time()
|
||
|
|
if remaining <= 0:
|
||
|
|
raise RuntimeError(f"timed out waiting for Nanobot API response to {method}")
|
||
|
|
|
||
|
|
try:
|
||
|
|
line = await asyncio.wait_for(reader.readline(), timeout=remaining)
|
||
|
|
except ValueError as exc:
|
||
|
|
raise RuntimeError(
|
||
|
|
"Nanobot API response exceeded the configured stream limit"
|
||
|
|
) from exc
|
||
|
|
if not line:
|
||
|
|
raise RuntimeError("Nanobot API socket closed before responding")
|
||
|
|
|
||
|
|
try:
|
||
|
|
message = json.loads(line.decode("utf-8", errors="replace"))
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
continue
|
||
|
|
if not isinstance(message, dict):
|
||
|
|
continue
|
||
|
|
if message.get("jsonrpc") != _JSONRPC_VERSION:
|
||
|
|
continue
|
||
|
|
if "method" in message:
|
||
|
|
continue
|
||
|
|
if str(message.get("id", "")).strip() != request_id:
|
||
|
|
continue
|
||
|
|
if "error" in message:
|
||
|
|
error = message.get("error", {})
|
||
|
|
if isinstance(error, dict):
|
||
|
|
raise _NanobotApiError(
|
||
|
|
int(error.get("code", -32000)),
|
||
|
|
str(error.get("message", "unknown Nanobot API error")),
|
||
|
|
)
|
||
|
|
raise _NanobotApiError(-32000, str(error))
|
||
|
|
return message.get("result")
|
||
|
|
finally:
|
||
|
|
writer.close()
|
||
|
|
await writer.wait_closed()
|
||
|
|
|
||
|
|
|
||
|
|
async def _run_nanobot_agent_turn(
|
||
|
|
content: str,
|
||
|
|
*,
|
||
|
|
chat_id: str,
|
||
|
|
metadata: dict[str, Any] | None = None,
|
||
|
|
timeout_seconds: float = 90.0,
|
||
|
|
notification_to_event: Callable[[dict[str, Any]], dict[str, Any] | None],
|
||
|
|
socket_path: Path = NANOBOT_API_SOCKET,
|
||
|
|
stream_limit: int = _NANOBOT_API_STREAM_LIMIT,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
final_message = ""
|
||
|
|
saw_activity = False
|
||
|
|
async for typed_event in _stream_nanobot_agent_turn(
|
||
|
|
content,
|
||
|
|
chat_id=chat_id,
|
||
|
|
metadata=metadata,
|
||
|
|
timeout_seconds=timeout_seconds,
|
||
|
|
notification_to_event=notification_to_event,
|
||
|
|
socket_path=socket_path,
|
||
|
|
stream_limit=stream_limit,
|
||
|
|
):
|
||
|
|
event_type = typed_event.get("type")
|
||
|
|
if event_type == "message":
|
||
|
|
saw_activity = True
|
||
|
|
if (
|
||
|
|
not bool(typed_event.get("is_progress", False))
|
||
|
|
and typed_event.get("role") == "nanobot"
|
||
|
|
):
|
||
|
|
final_message = str(typed_event.get("content", ""))
|
||
|
|
continue
|
||
|
|
|
||
|
|
if event_type == "agent_state":
|
||
|
|
saw_activity = True
|
||
|
|
|
||
|
|
return {
|
||
|
|
"status": "ok",
|
||
|
|
"message": final_message,
|
||
|
|
"completed": saw_activity,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
async def _stream_nanobot_agent_turn(
|
||
|
|
content: str,
|
||
|
|
*,
|
||
|
|
chat_id: str,
|
||
|
|
metadata: dict[str, Any] | None = None,
|
||
|
|
timeout_seconds: float = 60.0,
|
||
|
|
idle_grace_seconds: float = 1.0,
|
||
|
|
notification_to_event: Callable[[dict[str, Any]], dict[str, Any] | None],
|
||
|
|
socket_path: Path = NANOBOT_API_SOCKET,
|
||
|
|
stream_limit: int = _NANOBOT_API_STREAM_LIMIT,
|
||
|
|
) -> AsyncIterator[dict[str, Any]]:
|
||
|
|
reader, writer = await _open_nanobot_api_socket(socket_path, stream_limit=stream_limit)
|
||
|
|
saw_final_message = False
|
||
|
|
saw_activity = False
|
||
|
|
try:
|
||
|
|
writer.write(
|
||
|
|
(
|
||
|
|
json.dumps(
|
||
|
|
_jsonrpc_notification(
|
||
|
|
"message.send",
|
||
|
|
{
|
||
|
|
"content": content,
|
||
|
|
"chat_id": chat_id,
|
||
|
|
"metadata": metadata or {},
|
||
|
|
},
|
||
|
|
),
|
||
|
|
ensure_ascii=False,
|
||
|
|
)
|
||
|
|
+ "\n"
|
||
|
|
).encode("utf-8")
|
||
|
|
)
|
||
|
|
await writer.drain()
|
||
|
|
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
deadline = loop.time() + timeout_seconds
|
||
|
|
|
||
|
|
while True:
|
||
|
|
remaining = deadline - loop.time()
|
||
|
|
if remaining <= 0:
|
||
|
|
raise RuntimeError("Timed out waiting for a Nanobot response.")
|
||
|
|
|
||
|
|
timeout = min(remaining, idle_grace_seconds) if saw_final_message else remaining
|
||
|
|
try:
|
||
|
|
line = await asyncio.wait_for(reader.readline(), timeout=timeout)
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
if saw_final_message:
|
||
|
|
break
|
||
|
|
raise RuntimeError("Timed out waiting for Nanobot to start responding")
|
||
|
|
except ValueError as exc:
|
||
|
|
raise RuntimeError(
|
||
|
|
"Nanobot API response exceeded the configured stream limit"
|
||
|
|
) from exc
|
||
|
|
|
||
|
|
if not line:
|
||
|
|
break
|
||
|
|
|
||
|
|
try:
|
||
|
|
obj = json.loads(line.decode("utf-8", errors="replace").strip())
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
continue
|
||
|
|
if not isinstance(obj, dict):
|
||
|
|
continue
|
||
|
|
|
||
|
|
typed_event = notification_to_event(obj)
|
||
|
|
if typed_event is None:
|
||
|
|
continue
|
||
|
|
|
||
|
|
event_type = typed_event.get("type")
|
||
|
|
if event_type == "message":
|
||
|
|
saw_activity = True
|
||
|
|
if (
|
||
|
|
not bool(typed_event.get("is_progress", False))
|
||
|
|
and typed_event.get("role") == "nanobot"
|
||
|
|
):
|
||
|
|
saw_final_message = True
|
||
|
|
|
||
|
|
yield typed_event
|
||
|
|
|
||
|
|
if (
|
||
|
|
saw_final_message
|
||
|
|
and event_type == "agent_state"
|
||
|
|
and str(typed_event.get("state", "")).strip().lower() == "idle"
|
||
|
|
):
|
||
|
|
break
|
||
|
|
|
||
|
|
if not saw_activity:
|
||
|
|
raise RuntimeError("Timed out waiting for Nanobot to start responding")
|
||
|
|
finally:
|
||
|
|
writer.close()
|
||
|
|
await writer.wait_closed()
|
||
|
|
|
||
|
|
|
||
|
|
NanobotApiError = _NanobotApiError
|
||
|
|
send_nanobot_api_request = _send_nanobot_api_request
|
||
|
|
run_nanobot_agent_turn = _run_nanobot_agent_turn
|
||
|
|
stream_nanobot_agent_turn = _stream_nanobot_agent_turn
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"NANOBOT_API_SOCKET",
|
||
|
|
"NanobotApiError",
|
||
|
|
"run_nanobot_agent_turn",
|
||
|
|
"send_nanobot_api_request",
|
||
|
|
"stream_nanobot_agent_turn",
|
||
|
|
]
|