remdb 0.3.180__py3-none-any.whl → 0.3.258__py3-none-any.whl

rem/api/routers/chat/child_streaming.py

@@ -0,0 +1,393 @@
+"""
+Child Agent Event Handling.
+
+Handles events from child agents during multi-agent orchestration.
+
+Event Flow:
+```
+Parent Agent (Orchestrator)
+      │
+      ▼
+  ask_agent tool
+      │
+      ├──────────────────────────────────┐
+      ▼                                  │
+  Child Agent (intake_diverge)           │
+      │                                  │
+      ├── child_tool_start ──────────────┼──► Event Sink (Queue)
+      ├── child_content ─────────────────┤
+      └── child_tool_result ─────────────┘
+                                         │
+                                         ▼
+                            drain_child_events()
+                                         │
+                                         ├── SSE to client
+                                         └── DB persistence
+```
+
+IMPORTANT: When child_content is streamed, parent text output should be SKIPPED
+to prevent content duplication.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import uuid
+from typing import TYPE_CHECKING, Any, AsyncGenerator
+
+from loguru import logger
+
+from .streaming_utils import StreamingState, build_content_chunk
+from .sse_events import MetadataEvent, ToolCallEvent, format_sse_event
+from ....services.session import SessionMessageStore
+from ....settings import settings
+from ....utils.date_utils import to_iso, utc_now
+
+if TYPE_CHECKING:
+    from ....agentic.context import AgentContext
+
+
+async def handle_child_tool_start(
+    state: StreamingState,
+    child_agent: str,
+    tool_name: str,
+    arguments: dict | str | None,
+    session_id: str | None,
+    user_id: str | None,
+) -> AsyncGenerator[str, None]:
+    """
+    Handle child_tool_start event.
+
+    Actions:
+    1. Log the tool call
+    2. Emit SSE event
+    3. Save to database (with tool_arguments in metadata for consistency with parent)
+    """
+    full_tool_name = f"{child_agent}:{tool_name}"
+    tool_id = f"call_{uuid.uuid4().hex[:8]}"
+
+    # Normalize arguments - may come as JSON string from ToolCallPart.args
+    if isinstance(arguments, str):
+        try:
+            arguments = json.loads(arguments)
+        except json.JSONDecodeError:
+            arguments = None
+    elif not isinstance(arguments, dict):
+        arguments = None
+
+    # 1. LOG
+    logger.info(f"🔧 {full_tool_name}")
+
+    # 2. EMIT SSE
+    yield format_sse_event(ToolCallEvent(
+        tool_name=full_tool_name,
+        tool_id=tool_id,
+        status="started",
+        arguments=arguments,
+    ))
+
+    # 3. SAVE TO DB - content contains args as JSON (pydantic_messages.py parses it)
+    if session_id and settings.postgres.enabled:
+        try:
+            store = SessionMessageStore(
+                user_id=user_id or settings.test.effective_user_id
+            )
+            tool_msg = {
+                "role": "tool",
+                # Content is the tool call args as JSON - this is what the agent sees on reload
+                # and what pydantic_messages.py parses for ToolCallPart.args
+                "content": json.dumps(arguments) if arguments else "",
+                "timestamp": to_iso(utc_now()),
+                "tool_call_id": tool_id,
+                "tool_name": full_tool_name,
+            }
+            await store.store_session_messages(
+                session_id=session_id,
+                messages=[tool_msg],
+                user_id=user_id,
+                compress=False,
+            )
+        except Exception as e:
+            logger.warning(f"Failed to save child tool call: {e}")
+
+
+def handle_child_content(
+    state: StreamingState,
+    child_agent: str,
+    content: str,
+) -> str | None:
+    """
+    Handle child_content event.
+
+    CRITICAL: Sets state.child_content_streamed = True
+    This flag is used to skip parent text output and prevent duplication.
+
+    Returns:
+        SSE chunk or None if content is empty
+    """
+    if not content:
+        return None
+
+    # Track that child content was streamed
+    # Parent text output should be SKIPPED when this is True
+    state.child_content_streamed = True
+    state.responding_agent = child_agent
+
+    return build_content_chunk(state, content)
+
+
+async def handle_child_tool_result(
+    state: StreamingState,
+    child_agent: str,
+    result: Any,
+    message_id: str | None,
+    session_id: str | None,
+    agent_schema: str | None,
+) -> AsyncGenerator[str, None]:
+    """
+    Handle child_tool_result event.
+
+    Actions:
+    1. Log metadata if present
+    2. Emit metadata event if present
+    3. Emit tool completion event
+    """
+    # Check for metadata registration
+    if isinstance(result, dict) and result.get("_metadata_event"):
+        risk = result.get("risk_level", "")
+        conf = result.get("confidence", "")
+        logger.info(f"📊 {child_agent} metadata: risk={risk}, confidence={conf}")
+
+        # Update responding agent from child
+        if result.get("agent_schema"):
+            state.responding_agent = result.get("agent_schema")
+
+        # Build extra dict with risk fields
+        extra_data = {}
+        if risk:
+            extra_data["risk_level"] = risk
+
+        yield format_sse_event(MetadataEvent(
+            message_id=message_id,
+            session_id=session_id,
+            agent_schema=agent_schema,
+            responding_agent=state.responding_agent,
+            confidence=result.get("confidence"),
+            extra=extra_data if extra_data else None,
+        ))
+
+    # Emit tool completion
+    # Preserve full result for dict/list types (needed for frontend)
+    if isinstance(result, (dict, list)):
+        result_for_sse = result
+    else:
+        result_for_sse = str(result) if result else None
+
+    yield format_sse_event(ToolCallEvent(
+        tool_name=f"{child_agent}:tool",
+        tool_id=f"call_{uuid.uuid4().hex[:8]}",
+        status="completed",
+        result=result_for_sse,
+    ))
+
+
+async def drain_child_events(
+    event_sink: asyncio.Queue,
+    state: StreamingState,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    message_id: str | None = None,
+    agent_schema: str | None = None,
+) -> AsyncGenerator[str, None]:
+    """
+    Drain all pending child events from the event sink.
+
+    This is called during tool execution to process events
+    pushed by child agents via ask_agent.
+
+    IMPORTANT: When child_content events are processed, this sets
+    state.child_content_streamed = True. Callers should check this
+    flag and skip parent text output to prevent duplication.
+    """
+    while not event_sink.empty():
+        try:
+            child_event = event_sink.get_nowait()
+            async for chunk in process_child_event(
+                child_event, state, session_id, user_id, message_id, agent_schema
+            ):
+                yield chunk
+        except Exception as e:
+            logger.warning(f"Error processing child event: {e}")
+
+
+async def process_child_event(
+    child_event: dict,
+    state: StreamingState,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    message_id: str | None = None,
+    agent_schema: str | None = None,
+) -> AsyncGenerator[str, None]:
+    """Process a single child event and yield SSE chunks."""
+    event_type = child_event.get("type", "")
+    child_agent = child_event.get("agent_name", "child")
+
+    if event_type == "child_tool_start":
+        async for chunk in handle_child_tool_start(
+            state=state,
+            child_agent=child_agent,
+            tool_name=child_event.get("tool_name", "tool"),
+            arguments=child_event.get("arguments"),
+            session_id=session_id,
+            user_id=user_id,
+        ):
+            yield chunk
+
+    elif event_type == "child_content":
+        chunk = handle_child_content(
+            state=state,
+            child_agent=child_agent,
+            content=child_event.get("content", ""),
+        )
+        if chunk:
+            yield chunk
+
+    elif event_type == "child_tool_result":
+        async for chunk in handle_child_tool_result(
+            state=state,
+            child_agent=child_agent,
+            result=child_event.get("result"),
+            message_id=message_id,
+            session_id=session_id,
+            agent_schema=agent_schema,
+        ):
+            yield chunk
+
+
+async def stream_with_child_events(
+    tools_stream,
+    child_event_sink: asyncio.Queue,
+    state: StreamingState,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    message_id: str | None = None,
+    agent_schema: str | None = None,
+) -> AsyncGenerator[tuple[str, Any], None]:
+    """
+    Multiplex tool events with child events using asyncio.wait().
+
+    This is the key fix for child agent streaming - instead of draining
+    the queue synchronously during tool event iteration, we concurrently
+    listen to both sources and yield events as they arrive.
+
+    Yields:
+        Tuples of (event_type, event_data) where event_type is either
+        "tool" or "child", allowing the caller to handle each appropriately.
+    """
+    tool_iter = tools_stream.__aiter__()
+
+    # Create initial tasks
+    pending_tool: asyncio.Task | None = None
+    pending_child: asyncio.Task | None = None
+
+    try:
+        pending_tool = asyncio.create_task(tool_iter.__anext__())
+    except StopAsyncIteration:
+        # No tool events, just drain any remaining child events
+        while not child_event_sink.empty():
+            try:
+                child_event = child_event_sink.get_nowait()
+                yield ("child", child_event)
+            except asyncio.QueueEmpty:
+                break
+        return
+
+    # Start listening for child events with a short timeout
+    pending_child = asyncio.create_task(
+        _get_child_event_with_timeout(child_event_sink, timeout=0.05)
+    )
+
+    try:
+        while True:
+            # Wait for either source to produce an event
+            tasks = {t for t in [pending_tool, pending_child] if t is not None}
+            if not tasks:
+                break
+
+            done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+
+            for task in done:
+                try:
+                    result = task.result()
+                except asyncio.TimeoutError:
+                    # Child queue timeout - restart listener
+                    if task is pending_child:
+                        pending_child = asyncio.create_task(
+                            _get_child_event_with_timeout(child_event_sink, timeout=0.05)
+                        )
+                    continue
+                except StopAsyncIteration:
+                    # Tool stream exhausted
+                    if task is pending_tool:
+                        pending_tool = None
+                        # Final drain of any remaining child events
+                        if pending_child:
+                            pending_child.cancel()
+                            try:
+                                await pending_child
+                            except asyncio.CancelledError:
+                                pass
+                        while not child_event_sink.empty():
+                            try:
+                                child_event = child_event_sink.get_nowait()
+                                yield ("child", child_event)
+                            except asyncio.QueueEmpty:
+                                break
+                        return
+                    continue
+
+                if task is pending_child and result is not None:
+                    # Got a child event
+                    yield ("child", result)
+                    # Restart child listener
+                    pending_child = asyncio.create_task(
+                        _get_child_event_with_timeout(child_event_sink, timeout=0.05)
+                    )
+                elif task is pending_tool:
+                    # Got a tool event
+                    yield ("tool", result)
+                    # Get next tool event
+                    try:
+                        pending_tool = asyncio.create_task(tool_iter.__anext__())
+                    except StopAsyncIteration:
+                        pending_tool = None
+                elif task is pending_child and result is None:
+                    # Timeout with no event - restart listener
+                    pending_child = asyncio.create_task(
+                        _get_child_event_with_timeout(child_event_sink, timeout=0.05)
+                    )
+    finally:
+        # Cleanup any pending tasks
+        for task in [pending_tool, pending_child]:
+            if task and not task.done():
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
+
+async def _get_child_event_with_timeout(
+    queue: asyncio.Queue, timeout: float = 0.05
+) -> dict | None:
+    """
+    Get an event from the queue with a timeout.
+
+    Returns None on timeout (no event available).
+    This allows the multiplexer to check for tool events regularly.
+    """
+    try:
+        return await asyncio.wait_for(queue.get(), timeout=timeout)
+    except asyncio.TimeoutError:
+        return None

rem/api/routers/chat/completions.py

@@ -16,11 +16,11 @@ IMPORTANT: Session IDs MUST be UUIDs. Non-UUID session IDs will cause message
     kubectl port-forward -n observability svc/otel-collector-collector 4318:4318
 
     # Terminal 2: Phoenix UI - view traces at http://localhost:6006
-    kubectl port-forward -n siggy svc/phoenix 6006:6006
+    kubectl port-forward -n rem svc/phoenix 6006:6006
 
 2. Get Phoenix API Key (REQUIRED for feedback->Phoenix sync):
 
-    export PHOENIX_API_KEY=$(kubectl get secret -n siggy rem-phoenix-api-key \\
+    export PHOENIX_API_KEY=$(kubectl get secret -n rem rem-phoenix-api-key \\
       -o jsonpath='{.data.PHOENIX_API_KEY}' | base64 -d)
 
 3. Start API with OTEL and Phoenix enabled:
@@ -70,7 +70,7 @@ OTEL Architecture
 =================
 
     REM API --[OTLP/HTTP]--> OTEL Collector --[relay]--> Phoenix
-             (port 4318)    (k8s: observability)         (k8s: siggy)
+             (port 4318)    (k8s: observability)         (k8s: rem)
 
 Environment Variables:
     OTEL__ENABLED=true              Enable OTEL tracing (required for trace capture)
@@ -164,7 +164,7 @@ from .models import (
     ChatCompletionUsage,
     ChatMessage,
 )
-from .streaming import stream_openai_response, stream_openai_response_with_save, stream_simulator_response
+from .streaming import stream_openai_response, stream_openai_response_with_save, stream_simulator_response, save_user_message
 
 router = APIRouter(prefix="/api/v1", tags=["chat"])
 
@@ -215,7 +215,7 @@ async def ensure_session_with_metadata(
     Merges request metadata with existing session metadata.
 
     Args:
-        session_id: Session identifier (maps to Session.name)
+        session_id: Session UUID from X-Session-Id header
         user_id: User identifier
         tenant_id: Tenant identifier
         is_eval: Whether this is an evaluation session
@@ -228,12 +228,8 @@ async def ensure_session_with_metadata(
     try:
         repo = Repository(Session, table_name="sessions")
 
-        # Try to load existing session by name (session_id is the name field)
-        existing_list = await repo.find(
-            filters={"name": session_id, "tenant_id": tenant_id},
-            limit=1,
-        )
-        existing = existing_list[0] if existing_list else None
+        # Look up session by UUID (id field)
+        existing = await repo.get_by_id(session_id)
 
         if existing:
             # Merge metadata if provided
@@ -254,9 +250,10 @@ async def ensure_session_with_metadata(
                 await repo.upsert(existing)
                 logger.debug(f"Updated session {session_id} (eval={is_eval}, metadata keys={list(merged_metadata.keys())})")
         else:
-            # Create new session
+            # Create new session with the provided UUID as the id
             session = Session(
-                name=session_id,
+                id=session_id,  # Use the provided UUID as session id
+                name=session_id,  # Default name to UUID, can be updated later with LLM-generated name
                 mode=SessionMode.EVALUATION if is_eval else SessionMode.NORMAL,
                 user_id=user_id,
                 tenant_id=tenant_id,
@@ -503,16 +500,51 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
             logger.error(f"Failed to transcribe audio: {e}")
             # Fall through with original content (will likely fail at agent)
 
-    # Use ContextBuilder to construct complete message list with:
-    # 1. System context hint (date + user profile)
-    # 2. Session history (if session_id provided)
-    # 3. New messages from request body (transcribed if audio)
+    # Use ContextBuilder to construct context and basic messages
+    # Note: We load session history separately for proper pydantic-ai message_history
     context, messages = await ContextBuilder.build_from_headers(
         headers=dict(request.headers),
         new_messages=new_messages,
         user_id=temp_context.user_id,  # From JWT token (source of truth)
     )
 
+    # Load raw session history for proper pydantic-ai message_history format
+    # This enables proper tool call/return pairing for LLM API compatibility
+    from ....services.session import SessionMessageStore, session_to_pydantic_messages, audit_session_history
+    from ....agentic.schema import get_system_prompt
+
+    pydantic_message_history = None
+    if context.session_id and settings.postgres.enabled:
+        try:
+            store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
+            raw_session_history = await store.load_session_messages(
+                session_id=context.session_id,
+                user_id=context.user_id,
+                compress_on_load=False,  # Don't compress - we need full data for reconstruction
+            )
+            if raw_session_history:
+                # CRITICAL: Extract and pass the agent's system prompt
+                # pydantic-ai only auto-adds system prompts when message_history is empty
+                # When we pass message_history, we must include the system prompt ourselves
+                agent_system_prompt = get_system_prompt(agent_schema) if agent_schema else None
+                pydantic_message_history = session_to_pydantic_messages(
+                    raw_session_history,
+                    system_prompt=agent_system_prompt,
+                )
+                logger.debug(f"Converted {len(raw_session_history)} session messages to {len(pydantic_message_history)} pydantic-ai messages (with system prompt)")
+
+                # Audit session history if enabled (for debugging)
+                audit_session_history(
+                    session_id=context.session_id,
+                    agent_name=schema_name or "default",
+                    prompt=body.messages[-1].content if body.messages else "",
+                    raw_session_history=raw_session_history,
+                    pydantic_messages_count=len(pydantic_message_history),
+                )
+        except Exception as e:
+            logger.warning(f"Failed to load session history for message_history: {e}")
+            # Fall back to old behavior (concatenated prompt)
+
     logger.info(f"Built context with {len(messages)} total messages (includes history + user context)")
 
     # Ensure session exists with metadata and eval mode if applicable
@@ -533,33 +565,30 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
         model_override=body.model,  # type: ignore[arg-type]
     )
 
-    # Combine all messages into single prompt for agent
-    # ContextBuilder already assembled: system context + history + new messages
-    prompt = "\n".join(msg.content for msg in messages)
+    # Build the prompt for the agent
+    # If we have proper message_history, use just the latest user message as prompt
+    # Otherwise, fall back to concatenating all messages (legacy behavior)
+    if pydantic_message_history:
+        # Use the latest user message as the prompt, with history passed separately
+        user_prompt = body.messages[-1].content if body.messages else ""
+        prompt = user_prompt
+        logger.debug(f"Using message_history with {len(pydantic_message_history)} messages")
+    else:
+        # Legacy: Combine all messages into single prompt for agent
+        prompt = "\n".join(msg.content for msg in messages)
 
     # Generate OpenAI-compatible request ID
     request_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
 
     # Streaming mode
     if body.stream:
-        # Save user message before streaming starts
-        if settings.postgres.enabled and context.session_id:
-            user_message = {
-                "role": "user",
-                "content": body.messages[-1].content if body.messages else "",
-                "timestamp": datetime.utcnow().isoformat(),
-            }
-            try:
-                store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
-                await store.store_session_messages(
-                    session_id=context.session_id,
-                    messages=[user_message],
-                    user_id=context.user_id,
-                    compress=False,  # User messages are typically short
-                )
-                logger.debug(f"Saved user message to session {context.session_id}")
-            except Exception as e:
-                logger.error(f"Failed to save user message: {e}", exc_info=True)
+        # Save user message before streaming starts (using shared utility)
+        if context.session_id:
+            await save_user_message(
+                session_id=context.session_id,
+                user_id=context.user_id,
+                content=body.messages[-1].content if body.messages else "",
+            )
 
         return StreamingResponse(
             stream_openai_response_with_save(
@@ -570,6 +599,8 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
                 agent_schema=schema_name,
                 session_id=context.session_id,
                 user_id=context.user_id,
+                agent_context=context,  # Pass context for multi-agent support
+                message_history=pydantic_message_history,  # Native pydantic-ai message history
             ),
             media_type="text/event-stream",
             headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
@@ -592,10 +623,16 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
         ) as span:
             # Capture trace context from the span we just created
             trace_id, span_id = get_current_trace_context()
-            result = await agent.run(prompt)
+            if pydantic_message_history:
+                result = await agent.run(prompt, message_history=pydantic_message_history)
+            else:
+                result = await agent.run(prompt)
     else:
         # No tracer available, run without tracing
-        result = await agent.run(prompt)
+        if pydantic_message_history:
+            result = await agent.run(prompt, message_history=pydantic_message_history)
+        else:
+            result = await agent.run(prompt)
 
     # Determine content format based on response_format request
     if body.response_format and body.response_format.type == "json_object":

rem/api/routers/chat/sse_events.py

@@ -321,7 +321,11 @@ class MetadataEvent(BaseModel):
     # Agent info
     agent_schema: str | None = Field(
         default=None,
-        description="Name of the agent schema used for this response (e.g., 'rem', 'query-assistant')"
+        description="Name of the top-level agent schema (e.g., 'rem', 'intake')"
+    )
+    responding_agent: str | None = Field(
+        default=None,
+        description="Name of the agent that produced this response (may differ from agent_schema if delegated via ask_agent)"
     )
 
     # Session info
@@ -409,9 +413,9 @@ class ToolCallEvent(BaseModel):
         default=None,
         description="Tool arguments (for 'started' status)"
     )
-    result: str | None = Field(
+    result: str | dict[str, Any] | list[Any] | None = Field(
         default=None,
-        description="Tool result summary (for 'completed' status)"
+        description="Tool result - full dict/list for structured data, string for simple results"
     )
     error: str | None = Field(
         default=None,