feat: send agent prompt -> prompt agent

This commit is contained in:
Shivering Ocelot committed 2026-10-10 12:38:01 -05:00
1 parent 161dfd507f
commit 0dee6a392f
7 files changed
+199 -105

No files matched your search

+37 -28
View File
@@ -5,7 +5,6 @@ import asyncio
import traceback
from temporary import generate_tmp_id
from stream import Stream
# Tool related classes must be moved to a different module
@@ -66,7 +65,10 @@ class Agent:
agent_run.close()
raise NotImplementedError("Subclasses must implement run()")
# Deprecate this. This is just run with extra baggage.
# Async generator: constructs a run, executes it, and yields every message
# the run emits (raw inference.Message objects: role + content) in arrival
# order. No projection: no delta collapsing, no tool filtering, no dropping
# of reasoning/usage. Consumers route or accumulate as they see fit.
async def prompt(self, instruction, input=None, **kwargs):
run = AgentRun.create(
self.app,
@@ -79,39 +81,46 @@ class Agent:
},
)
done_event = asyncio.Event()
out_stream = Stream(self.app)
tool_calls = []
tool_names = None
if 'tools' in kwargs:
tool_names = [t['function']['name'] for t in kwargs['tools'] if isinstance(t, dict)]
queue = asyncio.Queue()
def on_chunk(chunk):
out_stream.feed(chunk)
def on_inference(msg):
kwargs = {'content': msg.content}
queue.put_nowait(('message', msg))
if msg.role == 'assistant':
msg.content.listen(on_chunk, lambda: None)
elif msg.role == 'tool' and tool_names:
if msg.content.tool_name in tool_names:
tool_calls.append(msg.content)
def on_done(ev):
done_event.set()
out_stream.close()
def on_done(_data):
queue.put_nowait(('done', None))
# Register listeners before scheduling the task. AgentRun.on replays
# past events of that name, so this both avoids a race and guarantees
# no early messages are lost.
run.on('inference', on_inference)
run.on('done', on_done)
await self.run(run)
if tool_names:
await done_event.wait()
return {
'output': out_stream,
'tool_calls': tool_calls,
}
else:
return out_stream
run.task = asyncio.create_task(self.run(run))
def on_task_done(task):
# Safety net: an agent that raises (or forgets run.close()) would
# otherwise leave the consumer awaiting forever.
queue.put_nowait(('task', task))
run.task.add_done_callback(on_task_done)
try:
while True:
kind, payload = await queue.get()
if kind == 'message':
yield payload
elif kind == 'done':
break
else: # task
if payload.cancelled():
break
exc = payload.exception()
if exc:
raise exc
break
finally:
if not run.task.done():
run.task.cancel()
async def received_message(self, message):
+2 -1
View File
@@ -47,7 +47,8 @@ class App(ChatMixin, AgentMixin):
]
from features.calc import Calc
self.tool_classes = [Calc]
from features.web import WebSearch, WebFetch
self.tool_classes = [Calc, WebSearch, WebFetch]
# --- Connection Management (Synchronous) ---
+8 -3
View File
@@ -180,10 +180,15 @@ class AgentMixin:
runs = [r for r in runs if r.user_id == user_id]
return [r.serialize() for r in runs]
async def send_agent_prompt(self, agent_id, instruction, input=None, options={}):
async def prompt_agent(self, agent_id, instruction, input=None, options={}):
# Mirror agent.prompt's generator so run ownership and lifecycle stay
# inside the mixin: callers just iterate and receive the same raw
# inference.Message objects. get_agent_control is awaited lazily on first
# iteration, keeping the control object alive for the generator's
# lifetime.
control = await self.get_agent_control(agent_id)
response = await control.prompt(instruction, input, **options)
return response
async for msg in control.prompt(instruction, input, **options):
yield msg
async def call_agent_action(self, agent_id, action, args):
control = await self.get_agent_control(agent_id)
+61
View File
@@ -0,0 +1,61 @@
"""Web tools: online search and page fetching.
Both tools subclass :class:`~backend.src.tool.Tool` and rely on the base
class's automatic ``ToolSpec`` generation (docstring -> description,
``run`` signature -> parameters).
"""
import re
import urllib.request
from html import unescape
from tool import Tool
def _strip_html(html: str) -> str:
"""Light HTML filter: drop scripts/styles, strip tags, collapse whitespace."""
html = re.sub(
r"<(script|style|noscript)[^>]*>.*?</\1>",
" ",
html,
flags=re.DOTALL | re.IGNORECASE,
)
text = re.sub(r"<[^>]+>", " ", html)
text = unescape(text)
return re.sub(r"\s+", " ", text).strip()
class WebSearch(Tool):
"""Search the web with DuckDuckGo. Returns titles, URLs and snippets."""
def run(self, query: str, max_results: int = 5):
# `ddgs` is the current package name; `duckduckgo_search` is the legacy one.
try:
from ddgs import DDGS
except ImportError: # pragma: no cover - fallback for older installs
from duckduckgo_search import DDGS
with DDGS() as ddgs:
results = list(ddgs.text(query, max_results=max_results))
return [
{
"title": r.get("title", ""),
"url": r.get("href", ""),
"snippet": r.get("body", ""),
}
for r in results
]
class WebFetch(Tool):
"""Fetch a web page and return its text content, lightly filtered."""
def run(self, url: str):
req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
with urllib.request.urlopen(req, timeout=15) as resp:
html = resp.read().decode("utf-8", errors="replace")
return _strip_html(html)
__all__ = ["WebSearch", "WebFetch"]
+48 -43
View File
@@ -2,7 +2,6 @@ import json, time, traceback, asyncio, uuid
from quart import Response, request
from quart_auth import login_required, current_user
from tool import ToolSpec, ToolParameter, ToolCall
from agent import AgentRun
from utils import encode_json
# ---------------------------------------------------------------------------
@@ -62,43 +61,49 @@ async def execute_deferred(tc):
tc.result = await tc.tool.execute_async(**tc.arguments)
def attach_run_handlers(run, tool_names):
"""Wire the run's output events into a single queue.
def messages_to_queue(generator, tool_names):
"""Drive a prompt generator into a single queue.
The queue receives ('content'|'reasoning'|'tool_call'|'deferred_call', data)
tuples and terminates with ('eof', None). Tool calls are routed (client vs
native/deferred) here so both response modes share the same logic.
native/deferred) here so both response modes share the same logic. The
returned ``state`` dict tracks whether any tool call was observed, which
drives ``finish_reason``.
"""
queue = asyncio.Queue()
loop = asyncio.get_event_loop()
state = {'tool_call': False, 'error': None}
async def on_content(chunk):
await queue.put(('content', chunk))
async def on_reasoning(chunk):
await queue.put(('reasoning', chunk))
async def on_eof(): pass
def on_eof(): pass
def on_inference(msg):
# assistant and reasoning are assumed to have a Stream
# tool is assumed to have a ToolCall
if msg.role == 'assistant':
msg.content.listen(on_content, on_eof)
elif msg.role == 'reasoning':
msg.content.listen(on_reasoning, on_eof)
elif msg.role == 'tool':
deferred = route_tool_call(msg.content, tool_names)
msg_type = 'deferred_call' if deferred else 'tool_call'
loop.create_task(queue.put((msg_type, msg.content)))
elif msg.role == 'usage':
loop.create_task(queue.put(('usage', msg.content)))
async def drain():
try:
async for msg in generator:
# assistant and reasoning are assumed to have a Stream
# tool is assumed to have a ToolCall
if msg.role == 'assistant':
msg.content.listen(on_content, on_eof)
elif msg.role == 'reasoning':
msg.content.listen(on_reasoning, on_eof)
elif msg.role == 'tool':
deferred = route_tool_call(msg.content, tool_names)
msg_type = 'deferred_call' if deferred else 'tool_call'
state['tool_call'] = True
await queue.put((msg_type, msg.content))
elif msg.role == 'usage':
await queue.put(('usage', msg.content))
except Exception as e:
# Surface the failure to the consumer via the eof drain.
state['error'] = e
finally:
await queue.put(('eof', None))
def on_done(output):
loop.create_task(queue.put(('eof', None)))
run.on('inference', on_inference)
run.on('done', on_done)
return queue
loop.create_task(drain())
return queue, state
def format_chunk(agent_id, delta, finish_reason=None, usage=None):
@@ -242,25 +247,23 @@ def register(server, app):
# tool and gets deferred.
tool_names = {spec.name for spec in tools}
# --- Get agent control and create AgentRun directly ---
# --- Get agent control and drive the run through the shared prompt generator ---
control = await app.get_agent_control(agent_id)
run = AgentRun.create(
app,
agent_id=agent_id,
user_id=user_id,
input={
"messages": messages,
"tools": tools,
# Never execute tools inside the run: emit them instead so
# the handler can send client tools back and defer native
# ones.
"tool_policy": "ignore",
},
# The OpenAI inputs ride as kwargs so they land in run.input, which
# ExpertAgent.run reads ("tools", "tool_policy", "messages").
generator = control.prompt(
None,
input=None,
messages=messages,
tools=tools,
# Never execute tools inside the run: emit them instead so
# the handler can send client tools back and defer native
# ones.
tool_policy="ignore",
)
queue = attach_run_handlers(run, tool_names)
await control.run(run)
queue, state = messages_to_queue(generator, tool_names)
usage = {'completion_tokens': 0, 'prompt_tokens': 0, 'total_tokens': 0, 'cost': 0}
@@ -303,10 +306,10 @@ def register(server, app):
usage = data
elif msg_type == 'eof':
if state['error']:
raise state['error']
# Determine finish_reason — was there a tool call?
finish_reason = "tool_calls" if any(
e['name'] == 'tool_call' for e in run.events
) else "stop"
finish_reason = "tool_calls" if state['tool_call'] else "stop"
yield f"data: {json.dumps(format_chunk(agent_id, {}, finish_reason))}\n\n"
break
@@ -334,6 +337,8 @@ def register(server, app):
while True:
msg_type, data = await queue.get()
if msg_type == 'eof':
if state['error']:
raise state['error']
break
elif msg_type == 'content':
content += data
+16 -24
View File
@@ -1,7 +1,6 @@
import json, asyncio, traceback
from quart import websocket
from utils import encode_json
from stream import Stream # isinstance check in send-agent-prompt
def register(server, app, auth_manager):
@server.websocket('/ws')
@@ -298,37 +297,30 @@ def register(server, app, auth_manager):
)
await send({'message': "chat thread created", 'threadId': thread_id, 'nonce': nonce})
elif action == "send-agent-prompt":
elif action == "prompt-agent":
if 'instruction' not in body and 'input' not in body:
await send({'error': "instruction or input required"})
await send({'error': "instruction or input required", 'nonce': nonce})
continue
# Shape of response changes based on options lol
response = await app.send_agent_prompt(
# Multiple-result action, same shape as read-stream: every
# generated message is its own packet carrying the request
# nonce, so the client callback runs once per message. The
# terminal eof packet resolves the action. Streams referenced
# by open assistant/reasoning messages are read by the
# client via the existing read-stream action.
generator = app.prompt_agent(
await agent_id_from_body(),
body.get('instruction'),
body.get('input'),
body.get('options', {})
)
print("prompt response", repr(response))
tool_calls = None
# Extract from dict
if isinstance(response, dict):
tool_calls = response['tool_calls']
response = response['output']
if isinstance(response, Stream):
if body.get('stream'):
response = {"streamId": response.id}
else:
response = await response.content_wait()
packet_out = {'message': "prompt sent", 'response': response, 'nonce': nonce}
if tool_calls:
packet_out['toolCalls'] = tool_calls
await send(packet_out)
try:
async for msg in generator:
await send({'message': 'agent-message', 'data': msg, 'nonce': nonce})
finally:
# Disconnect / cancellation point: closes the run task.
await generator.aclose()
await send({'eof': True, 'nonce': nonce})
elif action == "read-stream":
stream_id = body['streamId']