Remove custom ag-ui from app backend.
This commit is contained in:
Yiorgis Gozadinos
parent
26bc02ef14
commit
f2fb64cbca
2 changed files with 74 additions and 199 deletions
|
|
@ -1,17 +1,14 @@
|
|||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ag_ui.core import EventType, StateSnapshotEvent
|
||||
from pydantic import BaseModel
|
||||
from pydantic_ai import Agent, RunContext, format_as_xml
|
||||
from pydantic_ai import Agent, RunContext, ToolReturn, format_as_xml
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models import SearchResult
|
||||
from haiku.rag.utils import get_model
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||
|
||||
|
||||
class CitationInfo(BaseModel):
|
||||
"""Citation info for frontend display."""
|
||||
|
|
@ -74,7 +71,6 @@ class ChatDeps:
|
|||
|
||||
client: HaikuRAG
|
||||
config: AppConfig
|
||||
agui_emitter: "AGUIEmitter | None" = None
|
||||
search_results: list[SearchResult] | None = None
|
||||
session_state: ChatSessionState | None = None
|
||||
|
||||
|
|
@ -135,7 +131,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
ctx: RunContext[ChatDeps],
|
||||
query: str,
|
||||
document_name: str | None = None,
|
||||
) -> str:
|
||||
) -> ToolReturn:
|
||||
"""Search the knowledge base for relevant documents.
|
||||
|
||||
Use this when you need to find documents or explore the knowledge base.
|
||||
|
|
@ -147,12 +143,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
"""
|
||||
from search_agent import SearchAgent
|
||||
|
||||
if ctx.deps.agui_emitter:
|
||||
msg = f"Searching: {query}"
|
||||
if document_name:
|
||||
msg += f" (in {document_name})"
|
||||
ctx.deps.agui_emitter.log(msg)
|
||||
|
||||
# Build context from conversation history
|
||||
context = None
|
||||
if ctx.deps.session_state and ctx.deps.session_state.qa_history:
|
||||
|
|
@ -169,7 +159,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
ctx.deps.search_results = results
|
||||
|
||||
if not results:
|
||||
return "No results found."
|
||||
return ToolReturn(return_value="No results found.")
|
||||
|
||||
# Build citation infos for frontend display
|
||||
citation_infos = [
|
||||
|
|
@ -186,23 +176,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
for i, r in enumerate(results)
|
||||
]
|
||||
|
||||
# Emit search results as citations
|
||||
if ctx.deps.agui_emitter:
|
||||
ctx.deps.agui_emitter.update_state(
|
||||
ChatSessionState(
|
||||
session_id=(
|
||||
ctx.deps.session_state.session_id
|
||||
if ctx.deps.session_state
|
||||
else ""
|
||||
),
|
||||
citations=citation_infos,
|
||||
qa_history=(
|
||||
ctx.deps.session_state.qa_history
|
||||
if ctx.deps.session_state
|
||||
else []
|
||||
),
|
||||
)
|
||||
)
|
||||
# Build new state with citations
|
||||
new_state = ChatSessionState(
|
||||
session_id=(
|
||||
ctx.deps.session_state.session_id if ctx.deps.session_state else ""
|
||||
),
|
||||
citations=citation_infos,
|
||||
qa_history=(
|
||||
ctx.deps.session_state.qa_history if ctx.deps.session_state else []
|
||||
),
|
||||
)
|
||||
|
||||
# Return detailed results for the agent to present
|
||||
result_lines = []
|
||||
|
|
@ -219,14 +202,23 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
line += f"\n {snippet}"
|
||||
result_lines.append(line)
|
||||
|
||||
return f"Found {len(results)} results:\n\n" + "\n\n".join(result_lines)
|
||||
return ToolReturn(
|
||||
return_value=f"Found {len(results)} results:\n\n"
|
||||
+ "\n\n".join(result_lines),
|
||||
metadata=[
|
||||
StateSnapshotEvent(
|
||||
type=EventType.STATE_SNAPSHOT,
|
||||
snapshot=new_state.model_dump(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
@agent.tool
|
||||
async def ask(
|
||||
ctx: RunContext[ChatDeps],
|
||||
question: str,
|
||||
document_name: str | None = None,
|
||||
) -> str:
|
||||
) -> ToolReturn:
|
||||
"""Answer a specific question using the knowledge base.
|
||||
|
||||
Use this for direct questions that need a focused answer with citations.
|
||||
|
|
@ -241,12 +233,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
from haiku.rag.graph.research.models import Citation, SearchAnswer
|
||||
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
|
||||
|
||||
if ctx.deps.agui_emitter:
|
||||
msg = f"Answering: {question}"
|
||||
if document_name:
|
||||
msg += f" (in {document_name})"
|
||||
ctx.deps.agui_emitter.log(msg)
|
||||
|
||||
# Build filter from document_name
|
||||
doc_filter = build_document_filter(document_name) if document_name else None
|
||||
|
||||
|
|
@ -290,8 +276,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
search_filter=doc_filter,
|
||||
max_concurrency=ctx.deps.config.research.max_concurrency,
|
||||
)
|
||||
# Don't pass agui_emitter to research graph - its state model differs from ChatSessionState
|
||||
# The ask tool handles final state emission with citations
|
||||
deps = ResearchDeps(
|
||||
client=ctx.deps.client,
|
||||
)
|
||||
|
|
@ -323,23 +307,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
)
|
||||
ctx.deps.session_state.qa_history.append(qa_response)
|
||||
|
||||
# Emit updated state with citations AND accumulated qa_history
|
||||
if ctx.deps.agui_emitter:
|
||||
ctx.deps.agui_emitter.update_state(
|
||||
ChatSessionState(
|
||||
session_id=(
|
||||
ctx.deps.session_state.session_id
|
||||
if ctx.deps.session_state
|
||||
else ""
|
||||
),
|
||||
citations=citation_infos,
|
||||
qa_history=(
|
||||
ctx.deps.session_state.qa_history
|
||||
if ctx.deps.session_state
|
||||
else []
|
||||
),
|
||||
)
|
||||
)
|
||||
# Build new state with citations AND accumulated qa_history
|
||||
new_state = ChatSessionState(
|
||||
session_id=(
|
||||
ctx.deps.session_state.session_id if ctx.deps.session_state else ""
|
||||
),
|
||||
citations=citation_infos,
|
||||
qa_history=(
|
||||
ctx.deps.session_state.qa_history if ctx.deps.session_state else []
|
||||
),
|
||||
)
|
||||
|
||||
# Format answer with citation references and confidence
|
||||
answer_text = result.answer
|
||||
|
|
@ -347,7 +324,15 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citation_infos)))
|
||||
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
|
||||
|
||||
return answer_text
|
||||
return ToolReturn(
|
||||
return_value=answer_text,
|
||||
metadata=[
|
||||
StateSnapshotEvent(
|
||||
type=EventType.STATE_SNAPSHOT,
|
||||
snapshot=new_state.model_dump(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
@agent.tool
|
||||
async def get_document(
|
||||
|
|
@ -361,9 +346,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
Args:
|
||||
query: The document title or URI to look up
|
||||
"""
|
||||
if ctx.deps.agui_emitter:
|
||||
ctx.deps.agui_emitter.log(f"Fetching document: {query}")
|
||||
|
||||
# Try exact URI match first
|
||||
doc = await ctx.deps.client.get_document_by_uri(query)
|
||||
|
||||
|
|
|
|||
|
|
@ -3,57 +3,19 @@ import os
|
|||
from pathlib import Path
|
||||
|
||||
from agent import ChatDeps, ChatSessionState, QAResponse, create_chat_agent
|
||||
from anyio import (
|
||||
EndOfStream,
|
||||
create_memory_object_stream,
|
||||
create_task_group,
|
||||
move_on_after,
|
||||
)
|
||||
from anyio.streams.memory import MemoryObjectSendStream
|
||||
from dotenv import load_dotenv
|
||||
from pydantic_ai.messages import (
|
||||
ModelMessage,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
TextPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.ui import SSE_CONTENT_TYPE
|
||||
from pydantic_ai.ui.ag_ui import AGUIAdapter
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse, StreamingResponse
|
||||
from starlette.responses import JSONResponse, Response, StreamingResponse
|
||||
from starlette.routing import Route
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import load_yaml_config
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||
from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event
|
||||
|
||||
|
||||
def convert_messages_to_history(
|
||||
messages: list[dict[str, str]],
|
||||
) -> list[ModelMessage]:
|
||||
"""Convert AG-UI/CopilotKit messages to pydantic-ai message history.
|
||||
|
||||
Skips the last message since it will be passed as user_prompt to agent.run().
|
||||
"""
|
||||
history: list[ModelMessage] = []
|
||||
|
||||
# Skip the last message - it will be the current user prompt
|
||||
for msg in messages[:-1]:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "user":
|
||||
history.append(ModelRequest(parts=[UserPromptPart(content=content)]))
|
||||
elif role == "assistant":
|
||||
history.append(ModelResponse(parts=[TextPart(content=content)]))
|
||||
# Skip other roles (system, tool, etc.) for now
|
||||
|
||||
return history
|
||||
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
|
@ -103,109 +65,40 @@ def get_client(effective_db_path: Path) -> HaikuRAG:
|
|||
return _client_cache[path_key]
|
||||
|
||||
|
||||
async def stream_chat(request: Request) -> StreamingResponse:
|
||||
async def stream_chat(request: Request) -> Response:
|
||||
"""Chat streaming endpoint with AG-UI protocol."""
|
||||
body = await request.json()
|
||||
logger.info(f"Received request: {list(body.keys())}")
|
||||
input_data = RunAgentInput(**body)
|
||||
body = await request.body()
|
||||
logger.info("Received chat request")
|
||||
|
||||
user_message = ""
|
||||
message_history: list[ModelMessage] = []
|
||||
if input_data.messages:
|
||||
user_message = input_data.messages[-1].get("content", "")
|
||||
message_history = convert_messages_to_history(input_data.messages)
|
||||
# Parse request to build run_input
|
||||
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
|
||||
run_input = AGUIAdapter.build_run_input(body)
|
||||
|
||||
send_stream, receive_stream = create_memory_object_stream[str]()
|
||||
# Restore qa_history from incoming state
|
||||
initial_qa_history: list[QAResponse] = []
|
||||
state = getattr(run_input, "state", None)
|
||||
if state and "qa_history" in state:
|
||||
initial_qa_history = [QAResponse(**qa) for qa in state.get("qa_history", [])]
|
||||
|
||||
async def run_agent_with_streaming(
|
||||
send_stream: MemoryObjectSendStream[str],
|
||||
) -> None:
|
||||
"""Execute agent and forward events to stream."""
|
||||
async with send_stream:
|
||||
try:
|
||||
# Create emitter for streaming
|
||||
emitter: AGUIEmitter = AGUIEmitter(
|
||||
thread_id=input_data.thread_id,
|
||||
run_id=input_data.run_id,
|
||||
use_deltas=True,
|
||||
)
|
||||
# Build deps with session state
|
||||
thread_id = getattr(run_input, "thread_id", None)
|
||||
deps = ChatDeps(
|
||||
client=get_client(db_path),
|
||||
config=Config,
|
||||
session_state=ChatSessionState(
|
||||
session_id=thread_id or "",
|
||||
qa_history=initial_qa_history,
|
||||
),
|
||||
)
|
||||
|
||||
# Get client
|
||||
effective_db_path = db_path
|
||||
if input_data.config and input_data.config.get("db_path"):
|
||||
effective_db_path = Path(input_data.config["db_path"])
|
||||
client = get_client(effective_db_path)
|
||||
|
||||
# Parse incoming state to restore qa_history
|
||||
initial_qa_history: list[QAResponse] = []
|
||||
if input_data.state and "qa_history" in input_data.state:
|
||||
initial_qa_history = [
|
||||
QAResponse(**qa)
|
||||
for qa in input_data.state.get("qa_history", [])
|
||||
]
|
||||
|
||||
# Create initial state with restored history
|
||||
initial_state = ChatSessionState(
|
||||
session_id=input_data.thread_id or "",
|
||||
qa_history=initial_qa_history,
|
||||
)
|
||||
|
||||
# Create deps with session state
|
||||
deps = ChatDeps(
|
||||
client=client,
|
||||
config=Config,
|
||||
agui_emitter=emitter,
|
||||
session_state=initial_state,
|
||||
)
|
||||
emitter.start_run(initial_state=initial_state)
|
||||
|
||||
# Forward events
|
||||
async def forward_events():
|
||||
async for event in emitter:
|
||||
event_type = event.get("type")
|
||||
logger.debug(f"AG-UI event: {event_type}")
|
||||
await send_stream.send(format_sse_event(event))
|
||||
|
||||
# Run agent and forward concurrently
|
||||
async with create_task_group() as tg:
|
||||
tg.start_soon(forward_events)
|
||||
|
||||
result = await chat_agent.run(
|
||||
user_message, deps=deps, message_history=message_history
|
||||
)
|
||||
emitter.log(result.output)
|
||||
emitter.finish_run(result.output)
|
||||
await emitter.close()
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Error executing agent")
|
||||
try:
|
||||
await send_stream.send(
|
||||
format_sse_event({"type": "RUN_ERROR", "message": str(e)})
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def event_generator():
|
||||
"""Generate SSE events with heartbeat to keep connection alive."""
|
||||
async with create_task_group() as tg:
|
||||
tg.start_soon(run_agent_with_streaming, send_stream)
|
||||
async with receive_stream:
|
||||
while True:
|
||||
try:
|
||||
# Wait for event with timeout, send heartbeat if nothing received
|
||||
with move_on_after(15): # 15 second timeout
|
||||
event_str = await receive_stream.receive()
|
||||
yield event_str
|
||||
continue
|
||||
# No event received within timeout - send SSE comment as heartbeat
|
||||
yield ": heartbeat\n\n"
|
||||
except EndOfStream:
|
||||
break
|
||||
# Use AGUIAdapter for streaming
|
||||
adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept)
|
||||
event_stream = adapter.run_stream(deps=deps)
|
||||
sse_event_stream = adapter.encode_stream(event_stream)
|
||||
|
||||
return StreamingResponse(
|
||||
event_generator(),
|
||||
media_type="text/event-stream",
|
||||
sse_event_stream,
|
||||
media_type=accept,
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
|
|
|
|||
Loading…
Reference in a new issue