diff --git a/backend/src/agent.py b/backend/src/agent.py index 9cf34d6..4a34fa1 100644 --- a/backend/src/agent.py +++ b/backend/src/agent.py @@ -146,9 +146,10 @@ class Agent: else: return 'user' - # Returns all resource values matching the kind and key - def get_resources(self, kind, key): - pass # TODO lol + # Returns all resource values matching the kind, and the key when given + def get_resources(self, kind, key=None): + return [r for r in self.resources + if r.kind == kind and (key is None or r.key == key)] class DebugAgent(Agent): diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index 71de467..d0b262e 100644 --- a/backend/src/agents/expert.py +++ b/backend/src/agents/expert.py @@ -24,6 +24,10 @@ class ExpertAgent(Agent): async def run(self, run): opts = {} + # Inject remotes as tools + if 'remotes' in run.input: + run.input['tools'] = run.input.get('tools', []) + self.app.getRemoteTools(self.user_id, run.input['remotes']) + # Is this roon using *Agent's* tools, as opposed to run tools using_agent_tools = False if 'tools' not in run.input: @@ -49,11 +53,16 @@ class ExpertAgent(Agent): # Collect inference output. These are resused across calls tool_calls = [] output_stream = None + reasoning_stream = None inference_end = asyncio.Event() def on_output(stream): nonlocal output_stream output_stream = stream + def on_reasoning(stream): + nonlocal reasoning_stream + reasoning_stream = stream run.on('output', on_output) + run.on('reasoning', on_reasoning) run.on('inference-end', lambda event: inference_end.set()) run.on('tool-call', lambda tc: print(f"[expert] tool call: {tc}") or tool_calls.append(tc)) @@ -74,7 +83,13 @@ class ExpertAgent(Agent): await inference_end.wait() - messages.append(inference.LlamaMessage('assistant', output_stream and await output_stream.content_wait() or "")) + content = await output_stream.content_wait(); + reasoning = await reasoning_stream.content_wait(); + + if content or reasoning: + msg = inference.LlamaMessage('assistant', content or "") + if reasoning: msg.reasoning = reasoning + messages.append(msg) print(f"[expert] collected tool calls: {len(tool_calls)}, using? {using_agent_tools}") if using_agent_tools and tool_calls: @@ -89,7 +104,7 @@ class ExpertAgent(Agent): tool_calls = [] # Buffer time to let openrouter rest? await asyncio.sleep(2) - else: + elif content or not reasoning: print("No tool calls to execute, finish") break diff --git a/backend/src/app.py b/backend/src/app.py index 8b8260d..b0542e0 100644 --- a/backend/src/app.py +++ b/backend/src/app.py @@ -19,6 +19,7 @@ import agents as agents_mod import temporary as tmp from server import start_server from chat import ChatView, Message, MESSAGE_PAGE_SIZE +from remote import RemoteConnection, RemoteToolSpec, parse_function_specs from app_mixins import ChatMixin @@ -26,7 +27,8 @@ class App(ChatMixin): def __init__(self, data_directory, replicate_token=None, openrouter_token=None): self.data_directory = data_directory self.user_connections = {} - self.remote_connections = {} + # Connected remotes: (user_id, remotename) -> RemoteConnection. + self.remotes = {} self.streams = {} self.event_listeners = [] # active websocket tasks, for manual cancelling upon server shutdown @@ -81,15 +83,71 @@ class App(ChatMixin): if not self.user_connections[user_id]: del self.user_connections[user_id] - def add_remote_connection(self, user_id, remotename, channel): - if user_id not in self.remote_connections: - self.remote_connections[user_id] = {} - user_connections = self.remote_connections[user_id] + # --- Remotes --- - if remotename not in user_connections: - user_connections[remotename] = [] - if channel not in user_connections[remotename]: - user_connections[remotename].append(channel) + def connectRemote(self, user_id, name, function_specs, send): + """Connect a remote served by a live session's send channel. + Remotes are scoped to the user: connecting a name that is already + connected fails (no overriding). The remote stays connected until + disconnectRemote — the 'remote-disconnect' action or the close of + the serving session.""" + if not name: + raise Exception("remotename is required") + if (user_id, name) in self.remotes: + raise Exception(f"remote '{name}' is already connected") + connection = RemoteConnection( + user_id, name, parse_function_specs(function_specs), send) + self.remotes[(user_id, name)] = connection + self.emit_event("remote-connected", {"remotename": name, "remote": connection.describe()}) + return connection + + def disconnectRemote(self, user_id, name): + """Disconnect a remote of this user; its pending calls fail. + Idempotent: returns the connection, or None when no remote is + connected under that name (used by both the manual disconnect + action and session teardown).""" + connection = self.remotes.pop((user_id, name), None) + if connection is not None: + connection.close() + self.emit_event("remote-disconnected", {"remotename": name}) + return connection + + def listRemotes(self, user_id): + """All connected remotes of this user.""" + return [conn.describe() + for key, conn in self.remotes.items() if key[0] == user_id] + + async def callRemote(self, user_id, name, function, args): + """Invoke a function of a connected remote and await the result. + Raises on: unknown remote, unknown function, timeout, or an error + reported by the remote.""" + connection = self.remotes.get((user_id, name)) + if connection is None: + raise Exception(f"no connected remote named '{name}'") + return await connection.call(function, args) + + def getRemoteTools(self, user_id, remote_specs): + """Expand run-input remote specs into agent tools. + + Entries are remote names, or {'name': ..., 'prefix': ...} to prefix + the tool names. Every function of the remote becomes one tool named + after the function itself, prefixed only when requested. Remotes + that are not connected are skipped with a warning: they can vanish + at any time and the run should proceed without them. + """ + tools = [] + for spec in remote_specs or []: + if isinstance(spec, str): + name, prefix = spec, '' + else: + name, prefix = spec.get('name'), spec.get('prefix', '') + connection = self.remotes.get((user_id, name)) + if connection is None: + print(f"WARNING: remote '{name}' requested by run input is not connected, skipping") + continue + for fn in connection.functions.values(): + tools.append(RemoteToolSpec(connection, fn, name=prefix + fn.name)) + return tools # --- Agent Methods (Asynchronous) --- diff --git a/backend/src/inference.py b/backend/src/inference.py index 26834f1..99c7b13 100644 --- a/backend/src/inference.py +++ b/backend/src/inference.py @@ -15,19 +15,22 @@ CONTEXT_WINDOW = 10_000 # TODO: Rename to ChatMessage. Add methods for constructing from instruct/openai/etc formatted objects. class LlamaMessage: - def __init__(self, role, content): + def __init__(self, role, content, reasoning=None): if role not in VALID_ROLES: raise Exception("message role must be one of: " + ",".join(VALID_ROLES)) self.role = role self.content = content + self.reasoning = reasoning def to_dict(self): """Serialize to the dict format expected by llama-cpp / ollama.""" d = {"role": self.role, "content": self.content} + if self.reasoning: + d["reasoning"] = self.reasoning return d def __repr__(self): - return f"LlamaMessage(role={self.role!r}, content={self.content!r})" + return f"LlamaMessage(role={self.role!r}, content={self.content!r}, reasoning={self.reasoning!r})" class Models: llama_8b = "meta/meta-llama-3-8b-instruct" diff --git a/backend/src/remote.py b/backend/src/remote.py new file mode 100644 index 0000000..1a6f937 --- /dev/null +++ b/backend/src/remote.py @@ -0,0 +1,169 @@ +import asyncio +import itertools +import json + +# Seconds to wait for a remote to answer a call before failing it +REMOTE_CALL_TIMEOUT = 30 + + +class RemoteError(Exception): + """A remote call could not be completed: unknown function, timeout, + error reported by the remote, or the remote disconnected.""" + + +class RemoteFunction: + """One function exposed by a remote, parsed from the 'tools' list of + its remote-connect message.""" + + def __init__(self, name, description, parameters): + self.name = name + self.description = description or "" + # Full OpenAI tool-parameter schema dict ({'type', 'properties', + # 'required'}), kept as sent by the remote + self.parameters = parameters or { + "type": "object", "properties": {}, "required": [], + } + + +def parse_function_specs(specs): + """Parse the 'tools' list of a remote-connect message into a + name -> RemoteFunction dict.""" + functions = {} + for spec in specs or []: + if not isinstance(spec, dict) or not spec.get('name'): + continue + functions[spec['name']] = RemoteFunction( + name=spec['name'], + description=spec.get('description', ''), + parameters=spec.get('parameters'), + ) + return functions + + +class RemoteConnection: + """A remote connected by a session via the 'remote-connect' action. + + Owns the RPC state for the protocol spoken over the serving session's + websocket: + backend -> remote: {'action': 'remote-call', 'requestId', 'function', 'args'} + remote -> backend: {'action': 'remote-call-result', 'requestId', 'result' | 'error'} + + Lives in App.remotes under (user_id, name) from App.connectRemote until + App.disconnectRemote (manual action or serving session close). + """ + + def __init__(self, user_id, name, functions, send): + self.user_id = user_id + self.name = name + self.functions = functions # dict[str, RemoteFunction] + self.send = send # async callable(dict) over the ws + self.pending_calls = {} # requestId -> asyncio.Future + self._request_ids = itertools.count() + + def describe(self): + """Shape served by the remote-list action. A connection is in the + registry iff it is connected, so connected is always true. + 'agentIds' is kept for frontend shape compatibility and is always + empty: agents receive remotes per-run via run inputs, never links.""" + return { + "name": self.name, + "connected": True, + "agentIds": [], + "functions": [{ + "name": fn.name, + "description": fn.description, + "parameters": fn.parameters, + } for fn in self.functions.values()], + } + + async def call(self, function, args, timeout=REMOTE_CALL_TIMEOUT): + """Invoke a remote function and await the result. Raises RemoteError + on any failure; never returns a partial result.""" + if function not in self.functions: + raise RemoteError(f"remote '{self.name}' has no function '{function}'") + if self.send is None: + raise RemoteError(f"remote '{self.name}' is disconnected") + + # '#' is unique among one user's remotes, so results can + # be routed even when a session serves several remotes + request_id = f"{self.name}#{next(self._request_ids)}" + future = asyncio.get_running_loop().create_future() + self.pending_calls[request_id] = future + try: + await self.send({ + "action": "remote-call", + "requestId": request_id, + "function": function, + "args": args, + }) + return await asyncio.wait_for(future, timeout) + except asyncio.TimeoutError: + raise RemoteError( + f"remote call '{function}' on '{self.name}' timed out after {timeout}s") + finally: + self.pending_calls.pop(request_id, None) + + def resolve(self, request_id, result): + """Complete a pending call with the remote's result. No-op for + unknown request ids (results are fanned out over a session's + connections, see the remote-call-result handler).""" + future = self.pending_calls.get(request_id) + if future is not None and not future.done(): + future.set_result(result) + + def reject(self, request_id, error): + """Fail a pending call with an error reported by the remote. No-op + for unknown request ids.""" + future = self.pending_calls.get(request_id) + if future is not None and not future.done(): + future.set_exception(RemoteError(str(error))) + + def close(self): + """Fail all pending calls; called by App.disconnectRemote on manual + disconnect and on session close. Callers waiting in call() get + RemoteError.""" + self.send = None + for future in self.pending_calls.values(): + if not future.done(): + future.set_exception(RemoteError(f"remote '{self.name}' disconnected")) + self.pending_calls.clear() + + +class RemoteToolSpec: + """Duck-typed ToolSpec exposing one remote function to an agent run. + + Presents exactly the surface used by expert.py's tool loop and run + serialization: .name, .description, .to_openai_dict(), and async + .execute_async(**args). Execution delegates to RemoteConnection.call; + failures come back as error strings so the model sees them as tool + results and can recover. + """ + + def __init__(self, connection, function, name=None): + self.connection = connection + self.function = function + self.name = name or function.name + self.description = function.description + + def to_openai_dict(self): + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.function.parameters, + }, + } + + async def execute_async(self, **kwargs): + try: + result = await self.connection.call(self.function.name, kwargs) + except RemoteError as err: + result = {"error": str(err)} + # ToolCall.result is string content; remote results are JSON values + if not isinstance(result, str): + result = json.dumps(result) + return result + + def __repr__(self): + return f"RemoteToolSpec(name={self.name!r}, remote={self.connection.name!r})" diff --git a/backend/src/server.py b/backend/src/server.py index 06a3a6c..63ef5eb 100644 --- a/backend/src/server.py +++ b/backend/src/server.py @@ -127,7 +127,7 @@ async def start_server(app, shutdown_event): await websocket.send(json.dumps({'error': 'invalid auth'})) return - user_id = app.verify_user_credentials(auth_req['username'], auth_req['password']) + user_id = await app.verify_user_credentials(auth_req['username'], auth_req['password']) if user_id is None: print("Invalid Credentials"); await websocket.send(json.dumps({'error': 'invalid credentials'})) @@ -150,13 +150,13 @@ async def start_server(app, shutdown_event): ws_task = asyncio.current_task() app.ws_tasks.append(ws_task) + session_connections = {} # remotename -> RemoteConnection served by this socket + stop_listening = None try: - # TODO: remotes aren't separate connections - # instead, any connection can create a remotes and listen for requests sent to them - if 'remotename' in auth_req: - app.add_remote_connection(user_id, authreq['remotename'], authreq['tools'], send_queue) - else: - app.add_user_connection(user_id, send_queue) + # All sessions are identical after authentication: every session + # is a user connection, and any session can additionally connect + # remotes ('remote-connect' action) and serve calls on them. + app.add_user_connection(user_id, send_queue) # Listen for events and broadcast them async def on_event(name, data): @@ -258,6 +258,52 @@ async def start_server(app, shutdown_event): elif action == "cancel-agent-run": await app.cancel_agent_run(body['runId']) + elif action == "remote-list": + await send({'remotes': app.listRemotes(user_id), 'nonce': nonce}) + + elif action == "remote-connect": + # Connect a remote served by THIS session's channel. + # Expected failures (duplicate name) are replied, not + # raised, so the session stays alive. + try: + connection = app.connectRemote(user_id, body['remotename'], body.get('functions', []), send) + session_connections[body['remotename']] = connection + await send({'message': "remote connected", 'remotename': body['remotename'], 'nonce': nonce}) + except Exception as e: + await send({'error': str(e), 'nonce': nonce}) + + elif action == "remote-disconnect": + # Manual disconnect, allowed from any session of the user + connection = app.disconnectRemote(user_id, body['remotename']) + session_connections.pop(body['remotename'], None) + if connection is None: + await send({'error': "no connected remote named '" + str(body['remotename']) + "'", 'nonce': nonce}) + else: + await send({'message': "remote disconnected", 'remotename': body['remotename'], 'nonce': nonce}) + + elif action == "remote-call": + # Invoke a function of one of the user's connected remotes. + # Expected failures (unknown remote/function, timeout, + # remote-side error) are replied, not raised. + try: + result = await app.callRemote(user_id, body['remotename'], body['function'], body.get('args', {})) + await send({'result': result, 'nonce': nonce}) + except Exception as e: + await send({'error': str(e), 'nonce': nonce}) + + elif action == "remote-call-result": + # A session serving remotes replies to a downstream call. + # Fields live at the message top level, not in 'body'. + # requestIds are '#' and resolve/reject + # are no-ops for unknown ids, so fanning out over this + # session's connections routes the reply correctly. + if 'error' in req: + for conn in session_connections.values(): + conn.reject(req.get('requestId'), req['error']) + else: + for conn in session_connections.values(): + conn.resolve(req.get('requestId'), req.get('result')) + elif action == "create-chat": chat_id = await app.create_chat( user_id=user_id, @@ -351,6 +397,15 @@ async def start_server(app, shutdown_event): print("User logged out:", user_id) if ws_task in app.ws_tasks: app.ws_tasks.remove(ws_task) + # Unregister the event listener: the send closure would + # otherwise keep queueing events for a dead socket forever + if stop_listening: + stop_listening() + # All remotes served by this session are disconnected with it. + # disconnectRemote is idempotent, so remotes that were already + # disconnected manually (possibly from another session) are safe. + for remotename in list(session_connections): + app.disconnectRemote(user_id, remotename) app.remove_user_connection(user_id, send_queue) # https://github.com/hrgarber/openai-compatible-api/blob/main/app/main.py diff --git a/docs/remotes.md b/docs/remotes.md new file mode 100644 index 0000000..dabd62b --- /dev/null +++ b/docs/remotes.md @@ -0,0 +1,115 @@ +# Remotes + +## Overview + +A **session** is one authenticated `/ws` connection. Authentication is +uniform — token or username/password — and every session is a user +connection that receives event broadcasts and can send any action. + +A **remote** is a named set of functions served by a session: + +- Connected by that session with the `remote-connect` action (backed by the + `App.connectRemote` method), stored in `App.remotes` scoped to the user. +- **Lifetime**: connected until `remote-disconnect` (manual, allowed from any + session of the same user) or until the serving session closes — session + close disconnects every remote it connected. Nothing is persisted. +- **No overriding**: connecting a name that is already connected for that + user fails with an error reply. + +**Local Remotes** — purely user-side constructs, entirely implemented and +stored in the frontend (`Remotes.vue`). They never touch the backend. + +Agents never link to remotes. A prompt passes remote names in +`options.remotes`; `ExpertAgent.run` expands them into run-scoped tools. +Every function of a remote becomes a tool named exactly after the function, +prefixed only when the run input asks for it. + +This also implements the existing `server.py` TODO ("remotes aren't separate +connections — any connection can create a remotes and listen for requests +sent to them"). + +The frontend is already implemented and out of scope; the backend serves the +two actions it uses (`remote-list`, `remote-call`) with the shapes below. + +## Wire protocol + +Every message in both directions uses the existing action envelope. There is +no separate `type`-keyed protocol. + +Upstream (session → backend): + +| action | payload | success reply | error reply | +|---|---|---|---| +| `remote-connect` | body `{remotename, tools}` | `{'message': 'remote connected', 'remotename', 'nonce'}` | `{'error': "remote 'x' is already connected", 'nonce'}` | +| `remote-disconnect` | body `{remotename}` | `{'message': 'remote disconnected', 'remotename', 'nonce'}` | `{'error': "no connected remote named 'x'", 'nonce'}` | +| `remote-list` | — | `{'remotes': [...], 'nonce'}` | — | +| `remote-call` | body `{remotename, function, args}` | `{'result': , 'nonce'}` | `{'error': '...', 'nonce'}` | +| `remote-call-result` | top-level `{requestId, result}` or `{requestId, error}` | none | none | + +`tools` entries (function specs) and the `remote-list` entries share one +shape: `{name, description, parameters}` with `parameters` a full OpenAI +tool-parameter schema dict. + +Downstream (backend → the session serving the remote): + +~~~json +{"action": "remote-call", "requestId": "weather-service#0", "function": "get_forecast", "args": {"city": "Rome"}} +~~~ + +Result routing: requestIds are `'#'`, unique among a user's +remotes. A `remote-call-result` is fanned out over the receiving session's +own connections; `resolve`/`reject` ignore unknown ids, so one session can +serve several remotes without ambiguity. Fan-out is session-scoped, so +identical id strings belonging to another user's remote can never be +resolved from this session. + +Example connect + call round trip: + +~~~json +{"action": "remote-connect", "nonce": "c1", "body": {"remotename": "weather-service", "tools": [{"name": "get_forecast", "description": "Retrieve weather forecast for a location.", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "City name."}}, "required": ["city"]}}]}} +~~~ + +~~~json +{"action": "remote-call", "nonce": "c2", "body": {"remotename": "weather-service", "function": "get_forecast", "args": {"city": "Rome"}}} +~~~ + +~~~json +{"action": "remote-call-result", "requestId": "weather-service#0", "result": {"forecast": "sunny", "city": "Rome"}} +~~~ + +Calls not answered within `REMOTE_CALL_TIMEOUT` (30s, constant in +`remote.py`) fail; calls pending when the remote disconnects fail +immediately. All failures surface as `RemoteError`, which becomes an error +tool-result string for the model — never an exception that kills a run, and +never a killed caller session. + +## The prompt contract + +~~~json +{ + "action": "send-agent-prompt", + "body": { + "agentId": 7, + "instruction": "What's the weather in Rome? Use your tools.", + "options": { "remotes": ["weather-service", { "name": "code-repo-indexer", "prefix": "repo_" }] } + } +} +~~~ + +- `Agent.prompt` already splats options into the run input + (`input={"instruction": ..., "input": ..., **kwargs}`), so `options.remotes` + lands in `run.input["remotes"]`. No server or app change needed. +- Each entry is a remote name or `{"name": ..., "prefix": ...}`. Every + function of the remote becomes one tool named `prefix + function_name` + (empty prefix by default). +- Collisions: the agent's native tools are listed before remote tools and + the expert lookup takes the first match, so native tools win. Pass a + `prefix` to disambiguate. +- Remotes are resolved once, at run start. A remote not connected at that + moment is skipped with a console warning and the run proceeds (remotes + disconnect independently of prompts). If a remote disconnects mid-run, + its pending calls fail and the model sees error tool results. +- Existing expert rule unchanged: if the run input already carries its own + `tools`, those replace the agent toolset entirely and remotes are not in + play. +- Users only see their own remotes — the registry key includes `user_id`. \ No newline at end of file diff --git a/frontend/src/views/Remotes.vue b/frontend/src/views/Remotes.vue index 8ce912c..00e3989 100644 --- a/frontend/src/views/Remotes.vue +++ b/frontend/src/views/Remotes.vue @@ -1,11 +1,11 @@