chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,338 @@
|
||||
"""FastAPI server: Agent Spec agent on LangGraph over AG-UI, + durable memory.
|
||||
|
||||
We hand-roll the AG-UI streaming route (copied from the adapter's thin
|
||||
`add_agentspec_fastapi_endpoint`) so we can persist each exchange to Oracle
|
||||
Agent Memory — the adapter exposes no post-run hook. Persistence runs as a
|
||||
background task once the run finishes (off the SSE critical path, so the stream
|
||||
closes at RUN_FINISHED); it is fully server-side, and the frontend just streams
|
||||
from /run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import html
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from ag_ui.core import EventType, RunAgentInput, RunErrorEvent
|
||||
from ag_ui.encoder import EventEncoder
|
||||
from ag_ui_agentspec.agent import AgentSpecAgent
|
||||
from ag_ui_agentspec.agentspec_tracing_exporter import EVENT_QUEUE
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from .agent import build_agent_json
|
||||
from .memory import get_memory
|
||||
from .reconcile import reconcile_durable_memories
|
||||
from .tools import DEMO_USER_ID, TOOL_REGISTRY
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
# ── Multi-turn fix (upstream adapter workaround) ──────────────────────────────
|
||||
# The ag_ui_agentspec LangGraph runner checkpoints history per thread_id and, on
|
||||
# each turn, tries to append only the client messages whose ids aren't already in
|
||||
# the checkpoint (filter_only_new_messages). But CopilotKit re-sends the *full*
|
||||
# history with ids that never match the checkpoint's, so a second copy of the
|
||||
# assistant(tool_calls)/tool block gets appended; OpenAI then rejects the malformed
|
||||
# sequence on the next turn (400: "a message with role 'tool' must be a response to
|
||||
# a preceeding message with 'tool_calls'"), breaking every follow-up after a server
|
||||
# tool runs. See docs/known-issues/agentspec-multiturn-toolcall-correlation.md.
|
||||
#
|
||||
# Since the client already sends the full, valid history each turn, we replace the
|
||||
# adapter's incremental merge with a full-history *replace*: clear the checkpoint's
|
||||
# messages (RemoveMessage) and use the client's history verbatim. Drop this once the
|
||||
# upstream adapter records ToolExecutionRequests so the ids correlate.
|
||||
from langchain_core.messages import RemoveMessage # noqa: E402
|
||||
from langgraph.graph.message import REMOVE_ALL_MESSAGES # noqa: E402
|
||||
import ag_ui_agentspec.runtimes.langgraph_runner as _lg_runner # noqa: E402
|
||||
|
||||
|
||||
def _repair_dangling_tool_calls(messages: list[dict]) -> list[dict]:
|
||||
"""Synthesize a tool result for any assistant tool_call that has no response.
|
||||
|
||||
book_flight is a client-side HITL tool: calling it interrupts the run and emits
|
||||
an assistant message with a tool_call, then waits for the UI to return a result
|
||||
when the traveler clicks Confirm/Cancel. If they instead send another chat
|
||||
message, that tool_call is left unanswered — and because we forward the client's
|
||||
full history verbatim, OpenAI rejects the next turn (400: "tool_call_ids did not
|
||||
have response messages"). This is the inverse of the duplicate-tool-block issue
|
||||
the history replace already handles (see the comment above).
|
||||
|
||||
For each assistant tool_call with no real tool result, insert a synthetic
|
||||
"not completed" tool result directly after the assistant message so the sequence
|
||||
is valid and the model can answer the new question. In this app the only tool
|
||||
that can dangle is the book_flight HITL — server tools resolve within the run —
|
||||
so the synthetic content is phrased for that case.
|
||||
|
||||
Assumes CopilotKit's normal ordering, where a real tool result immediately
|
||||
follows its assistant tool_calls message: this repairs *missing* results, not a
|
||||
result that has been re-ordered away from its originating call.
|
||||
"""
|
||||
# tool_call_ids that already have a REAL result somewhere in the history.
|
||||
answered = {
|
||||
m["tool_call_id"]
|
||||
for m in messages
|
||||
if m.get("role") == "tool" and m.get("tool_call_id")
|
||||
}
|
||||
repaired: list[dict] = []
|
||||
for m in messages:
|
||||
repaired.append(m)
|
||||
if m.get("role") != "assistant" or not m.get("tool_calls"):
|
||||
continue
|
||||
# De-dupe within THIS message only — a second assistant message carrying the
|
||||
# same unanswered id still needs its own result, so `answered` is never
|
||||
# mutated here (mutating it was the original bug: it suppressed the repair the
|
||||
# next occurrence needed).
|
||||
synthesized: set[str] = set()
|
||||
for tc in m["tool_calls"]:
|
||||
tc_id = tc.get("id")
|
||||
if not tc_id:
|
||||
# Can't synthesize a result without an id; surface it rather than
|
||||
# silently leave a dangling call that 400s on the next turn.
|
||||
print("[history] warning: assistant tool_call has no id; cannot repair")
|
||||
continue
|
||||
if tc_id in answered or tc_id in synthesized:
|
||||
continue
|
||||
synthesized.add(tc_id)
|
||||
name = (tc.get("function") or {}).get("name") or "the requested action"
|
||||
repaired.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tc_id,
|
||||
"content": f"{name} was not completed — the traveler continued without confirming.",
|
||||
}
|
||||
)
|
||||
return repaired
|
||||
|
||||
|
||||
async def _replace_history_with_client(_agent, _thread_id, input_messages):
|
||||
"""Replace the checkpoint's messages with the client's full history each turn,
|
||||
repairing any dangling tool_call (e.g. an abandoned book_flight HITL) first."""
|
||||
if not input_messages:
|
||||
return input_messages
|
||||
return [RemoveMessage(id=REMOVE_ALL_MESSAGES), *_repair_dangling_tool_calls(input_messages)]
|
||||
|
||||
|
||||
_lg_runner.filter_only_new_messages = _replace_history_with_client
|
||||
|
||||
# ── Oracle checkpointer injection (Plan §3, Option A) ─────────────────────────
|
||||
# ag_ui_agentspec's load_agent_spec hardcodes checkpointer=MemorySaver(); we
|
||||
# replace it so the LangGraph graph is compiled with our flag-gated checkpointer
|
||||
# (AsyncOracleSaver when LANGGRAPH_CHECKPOINTER=oracle, else MemorySaver). The
|
||||
# underlying pyagentspec AgentSpecLoader already accepts a checkpointer; only the
|
||||
# convenience wrapper needed patching. Drop this once the upstream adapter takes a
|
||||
# checkpointer param (Plan §3, Option B). AgentSpecAgent.__init__ resolves the name
|
||||
# from ag_ui_agentspec.agent, so we rebind both module namespaces.
|
||||
import ag_ui_agentspec.agent as _agent_mod # noqa: E402
|
||||
import ag_ui_agentspec.agentspecloader as _asl_mod # noqa: E402
|
||||
from pyagentspec.adapters.langgraph import AgentSpecLoader as _LGLoader # noqa: E402
|
||||
|
||||
from .checkpointer import resolve_checkpointer, init_checkpointer, close_checkpointer # noqa: E402
|
||||
|
||||
_orig_load_agent_spec = _agent_mod.load_agent_spec
|
||||
|
||||
|
||||
def _load_agent_spec_with_checkpointer(
|
||||
runtime, agent_spec_json, tool_registry=None, components_registry=None
|
||||
):
|
||||
if runtime != "langgraph":
|
||||
return _orig_load_agent_spec(
|
||||
runtime, agent_spec_json, tool_registry, components_registry
|
||||
)
|
||||
return _LGLoader(
|
||||
tool_registry=tool_registry, checkpointer=resolve_checkpointer()
|
||||
).load_json(agent_spec_json, components_registry)
|
||||
|
||||
|
||||
_agent_mod.load_agent_spec = _load_agent_spec_with_checkpointer
|
||||
_asl_mod.load_agent_spec = _load_agent_spec_with_checkpointer
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _lifespan(_app: FastAPI):
|
||||
# Build the Oracle checkpointer (if LANGGRAPH_CHECKPOINTER=oracle) before the
|
||||
# lazy agent build so resolve_checkpointer() sees an initialised saver. No-op
|
||||
# under the default `memory` flag.
|
||||
await init_checkpointer()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Drain in-flight background persists on shutdown so a graceful stop doesn't
|
||||
# drop the last turn's memory write. Loop rather than a single gather: a
|
||||
# request finishing during the drain can add a task after the snapshot, so
|
||||
# re-check until the set is empty. Persists are serialized (one at a time).
|
||||
while _PERSIST_TASKS:
|
||||
await asyncio.gather(*list(_PERSIST_TASKS), return_exceptions=True)
|
||||
await close_checkpointer()
|
||||
|
||||
|
||||
app = FastAPI(title="Oracle Concierge Agent", lifespan=_lifespan)
|
||||
app.add_middleware(
|
||||
CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]
|
||||
)
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _get_agentspec_agent() -> AgentSpecAgent:
|
||||
"""Build the agent once, on first request. Construction eagerly resolves the
|
||||
LLM (ChatOpenAI), which needs OPENAI_API_KEY, so we defer it out of import."""
|
||||
return AgentSpecAgent(
|
||||
build_agent_json(), runtime="langgraph", tool_registry=TOOL_REGISTRY
|
||||
)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> dict[str, str]:
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
def _last_user_message(messages: list) -> str:
|
||||
for message in reversed(messages):
|
||||
if getattr(message, "role", None) == "user":
|
||||
return getattr(message, "content", "") or ""
|
||||
return ""
|
||||
|
||||
|
||||
def _clean_assistant_text(parts: list[str]) -> str:
|
||||
"""Assemble the streamed assistant deltas into the text we persist to memory.
|
||||
|
||||
The agentspec exporter HTML-escapes every TEXT_MESSAGE_CHUNK delta for safe
|
||||
transport to the browser (agentspec_tracing_exporter._escape_html: & < > ->
|
||||
& < >). We must reverse that before persisting, or Oracle Agent
|
||||
Memory stores corrupted facts like "fares < $700" and recall/extraction
|
||||
operate on the mangled text. Join first, then unescape, so an entity split
|
||||
across two delta boundaries (e.g. "&l" + "t;") is still decoded correctly.
|
||||
The streamed copy yielded to the client is untouched — only the persisted
|
||||
copy is unescaped here.
|
||||
"""
|
||||
return html.unescape("".join(parts))
|
||||
|
||||
|
||||
# Background persistence tasks are tracked here so the event loop keeps a strong
|
||||
# reference until each finishes — a bare fire-and-forget task can be garbage
|
||||
# collected mid-flight (see the asyncio.create_task docs).
|
||||
_PERSIST_TASKS: set[asyncio.Task] = set()
|
||||
# Serialize background persists: only one extraction + reconciliation runs at a
|
||||
# time. The old await made the client wait on stream-close before sending the
|
||||
# next turn, which serialized persists for free; now that the stream closes at
|
||||
# RUN_FINISHED, overlapping turns could otherwise run reconcile's read-modify-
|
||||
# write concurrently (racing on which durable fact "wins") and exhaust the small
|
||||
# Oracle connection pool. Background persists queue on this lock instead.
|
||||
_PERSIST_LOCK = asyncio.Lock()
|
||||
|
||||
|
||||
async def _persist_serialized(user_text: str, assistant_text: str) -> None:
|
||||
async with _PERSIST_LOCK:
|
||||
await asyncio.to_thread(_persist_sync, user_text, assistant_text)
|
||||
|
||||
|
||||
def _on_persist_done(task: asyncio.Task) -> None:
|
||||
"""Drop the task ref and surface any failure. The write happens off the
|
||||
request path, so a silently-dropped task exception would make a lost write
|
||||
invisible. _persist_sync swallows its own DB/LLM errors; this catches
|
||||
cancellation (loop shutdown) and scheduling failures that would vanish."""
|
||||
_PERSIST_TASKS.discard(task)
|
||||
if task.cancelled():
|
||||
print("[persist] warning: background persist cancelled before completing")
|
||||
return
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
print(f"[persist] warning: background persist task failed ({exc!r})")
|
||||
|
||||
|
||||
def _spawn_persist(user_text: str, assistant_text: str) -> None:
|
||||
"""Run persistence off the request's critical path.
|
||||
|
||||
_persist_sync makes two LLM calls (memory extraction + reconciliation) plus
|
||||
DB writes — ~2-13s in practice. Awaiting it in the SSE generator's finally
|
||||
held the HTTP stream open that whole time *after* RUN_FINISHED, so the client
|
||||
(which ends its run/loading state on stream-close, not on RUN_FINISHED) showed
|
||||
a multi-second lag once the reply had already finished. Spawning it as a
|
||||
tracked, serialized background task lets the stream close at RUN_FINISHED; the
|
||||
write still lands a few seconds later, well before a human starts the next turn.
|
||||
"""
|
||||
task = asyncio.create_task(_persist_serialized(user_text, assistant_text))
|
||||
_PERSIST_TASKS.add(task)
|
||||
task.add_done_callback(_on_persist_done)
|
||||
|
||||
|
||||
def _persist_sync(user_text: str, assistant_text: str) -> None:
|
||||
"""Persist the exchange for memory extraction, then supersede stale facts.
|
||||
|
||||
Both turns are stored: extraction reliably distills durable facts from a full
|
||||
user+assistant exchange, whereas a lone user turn often yields only a raw
|
||||
"message" record and no extracted fact. The agent's echoes land as "message"
|
||||
records too, but recall_memory filters those out (see tools.DURABLE_RECORD_TYPES),
|
||||
so they never re-assert a stale preference. Reconciliation then deletes
|
||||
outdated/duplicate *durable* facts so an updated preference wins on recall.
|
||||
"""
|
||||
exchange: list[dict[str, str]] = []
|
||||
if user_text:
|
||||
exchange.append({"role": "user", "content": user_text})
|
||||
if assistant_text:
|
||||
exchange.append({"role": "assistant", "content": assistant_text})
|
||||
if not exchange:
|
||||
return
|
||||
try:
|
||||
memory = get_memory()
|
||||
thread = memory.create_thread(user_id=DEMO_USER_ID)
|
||||
thread.add_messages(exchange) # triggers automatic memory extraction
|
||||
# Delete outdated/duplicate durable facts so the newest preference wins on recall.
|
||||
reconcile_durable_memories(DEMO_USER_ID)
|
||||
except Exception as exc: # persistence is best-effort; a DB blip must not 500 the run
|
||||
# Mirror recall_memory's graceful degradation: the chat reply already streamed
|
||||
# successfully, so swallow + log rather than raising out of the SSE generator's
|
||||
# finally (which surfaced as "Exception in ASGI application" while the DB was down).
|
||||
print(f"[persist] warning: memory persist failed, degrading gracefully ({exc})")
|
||||
|
||||
|
||||
@app.post("/run")
|
||||
async def run_endpoint(input_data: RunAgentInput, request: Request):
|
||||
"""Stream the Agent Spec run over AG-UI, then persist the turn in the background.
|
||||
|
||||
The event_generator mirrors the adapter's endpoint.py: a per-request queue is
|
||||
set into EVENT_QUEUE, the run is spawned as a task, and events are drained to
|
||||
SSE. We additionally collect the assistant's text deltas and, once the stream
|
||||
closes, spawn persistence as a background task (off the critical path).
|
||||
"""
|
||||
encoder = EventEncoder(accept=request.headers.get("accept"))
|
||||
user_text = _last_user_message(input_data.messages)
|
||||
|
||||
async def event_generator():
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
token = EVENT_QUEUE.set(queue)
|
||||
assistant_parts: list[str] = []
|
||||
|
||||
async def run_and_close():
|
||||
try:
|
||||
await _get_agentspec_agent().run(input_data)
|
||||
except Exception as exc: # surface failures to the client
|
||||
queue.put_nowait(RunErrorEvent(message=repr(exc)))
|
||||
finally:
|
||||
queue.put_nowait(None)
|
||||
|
||||
try:
|
||||
asyncio.create_task(run_and_close())
|
||||
while True:
|
||||
item = await queue.get()
|
||||
if item is None:
|
||||
break
|
||||
if item.type in (EventType.RUN_STARTED, EventType.RUN_FINISHED):
|
||||
item.thread_id = input_data.thread_id
|
||||
item.run_id = input_data.run_id
|
||||
if item.type == EventType.TEXT_MESSAGE_CHUNK:
|
||||
assistant_parts.append(getattr(item, "delta", "") or "")
|
||||
yield encoder.encode(item)
|
||||
except Exception as exc:
|
||||
yield encoder.encode(RunErrorEvent(message=str(exc)))
|
||||
finally:
|
||||
EVENT_QUEUE.reset(token)
|
||||
# Persist off the critical path so the SSE stream closes at RUN_FINISHED
|
||||
# instead of blocking on memory extraction + reconciliation. The write
|
||||
# still lands shortly after, so the next session can recall it.
|
||||
_spawn_persist(user_text, _clean_assistant_text(assistant_parts))
|
||||
|
||||
return StreamingResponse(event_generator(), media_type=encoder.get_content_type())
|
||||
Reference in New Issue
Block a user