remdb 0.3.172__py3-none-any.whl → 0.3.223__py3-none-any.whl

rem/agentic/README.md

@@ -716,11 +716,271 @@ curl -X POST http://localhost:8000/api/v1/chat/completions \
 
 See `rem/api/README.md` for full SSE event protocol documentation.
 
+## Multi-Agent Orchestration
+
+REM supports hierarchical agent orchestration where agents can delegate work to other agents via the `ask_agent` tool. This enables complex workflows with specialized agents.
+
+### Architecture
+
+```mermaid
+sequenceDiagram
+    participant User
+    participant API as Chat API
+    participant Orchestrator as Orchestrator Agent
+    participant EventSink as Event Sink (Queue)
+    participant Child as Child Agent
+    participant DB as PostgreSQL
+
+    User->>API: POST /chat/completions (stream=true)
+    API->>API: Create event sink (asyncio.Queue)
+    API->>Orchestrator: agent.iter(prompt)
+
+    loop Streaming Loop
+        Orchestrator->>API: PartDeltaEvent (text)
+        API->>User: SSE: data: {"delta": {"content": "..."}}
+    end
+
+    Orchestrator->>Orchestrator: Decides to call ask_agent
+    Orchestrator->>API: ToolCallPart (ask_agent)
+    API->>User: SSE: event: tool_call
+
+    API->>Child: ask_agent("child_name", input)
+    Child->>EventSink: push_event(child_content)
+    EventSink->>API: Consume child events
+    API->>User: SSE: data: {"delta": {"content": "..."}}
+
+    Child->>Child: Completes
+    Child-->>Orchestrator: Return result
+
+    Orchestrator->>API: Final response
+    API->>DB: Save tool calls
+    API->>DB: Save assistant message
+    API->>User: SSE: data: [DONE]
+```
+
+### Event Sink Pattern
+
+When an agent delegates to a child via `ask_agent`, the child's streaming events need to bubble up to the parent's stream. This is achieved through an **event sink** pattern using Python's `ContextVar`:
+
+```python
+# context.py
+from contextvars import ContextVar
+
+_parent_event_sink: ContextVar["asyncio.Queue | None"] = ContextVar(
+    "parent_event_sink", default=None
+)
+
+async def push_event(event: Any) -> bool:
+    """Push event to parent's event sink if available."""
+    sink = _parent_event_sink.get()
+    if sink is not None:
+        await sink.put(event)
+        return True
+    return False
+```
+
+The streaming controller sets up the event sink before agent execution:
+
+```python
+# streaming.py
+child_event_sink: asyncio.Queue = asyncio.Queue()
+set_event_sink(child_event_sink)
+
+async for node in agent.iter(prompt):
+    # Process agent events...
+
+    # Consume any child events that arrived
+    while not child_event_sink.empty():
+        child_event = child_event_sink.get_nowait()
+        if child_event["type"] == "child_content":
+            yield format_sse_content_delta(child_event["content"])
+```
+
+### ask_agent Tool Implementation
+
+The `ask_agent` tool in `mcp_router/tools.py` uses Pydantic AI's streaming iteration:
+
+```python
+async def ask_agent(agent_name: str, input_text: str, ...):
+    """Delegate work to another agent."""
+
+    # Load and create child agent
+    schema = await load_agent_schema_async(agent_name, user_id)
+    child_agent = await create_agent(context=context, agent_schema_override=schema)
+
+    # Stream child agent with event proxying
+    async with child_agent.iter(prompt) as agent_run:
+        async for node in agent_run:
+            if Agent.is_model_request_node(node):
+                async with node.stream(agent_run.ctx) as request_stream:
+                    async for event in request_stream:
+                        if isinstance(event, PartDeltaEvent):
+                            # Push content to parent's event sink
+                            await push_event({
+                                "type": "child_content",
+                                "agent_name": agent_name,
+                                "content": event.delta.content_delta,
+                            })
+
+    return agent_run.result
+```
+
+### Pydantic AI Features Used
+
+#### 1. Streaming Iteration (`agent.iter()`)
+
+Unlike `agent.run()` which blocks until completion, `agent.iter()` provides fine-grained control over the execution flow:
+
+```python
+async with agent.iter(prompt) as agent_run:
+    async for node in agent_run:
+        if Agent.is_model_request_node(node):
+            # Model is generating - stream the response
+            async with node.stream(agent_run.ctx) as stream:
+                async for event in stream:
+                    if isinstance(event, PartStartEvent):
+                        # Tool call starting
+                    elif isinstance(event, PartDeltaEvent):
+                        # Content chunk
+        elif Agent.is_call_tools_node(node):
+            # Tools are being executed
+            async with node.stream(agent_run.ctx) as stream:
+                async for event in stream:
+                    if isinstance(event, FunctionToolResultEvent):
+                        # Tool completed
+```
+
+#### 2. Node Types
+
+- **`ModelRequestNode`**: The model is generating a response (text or tool calls)
+- **`CallToolsNode`**: Tools are being executed
+- **`End`**: Agent execution complete
+
+#### 3. Event Types
+
+- **`PartStartEvent`**: A new part (text or tool call) is starting
+- **`PartDeltaEvent`**: Content chunk for streaming text
+- **`FunctionToolResultEvent`**: Tool execution completed with result
+- **`ToolCallPart`**: Metadata about a tool call (name, arguments)
+- **`TextPart`**: Text content
+
+### Message Persistence
+
+All messages are persisted to PostgreSQL for session continuity:
+
+```python
+# streaming.py - after agent completes
+async def save_session_messages(...):
+    store = SessionMessageStore(user_id=user_id)
+
+    # Save each tool call as a tool message
+    for tool_call in tool_calls:
+        await store.save_message(
+            session_id=session_id,
+            role="tool",
+            content=tool_call.result,
+            tool_name=tool_call.name,
+            tool_call_id=tool_call.id,
+        )
+
+    # Save the final assistant response
+    await store.save_message(
+        session_id=session_id,
+        role="assistant",
+        content=accumulated_content,
+    )
+```
+
+Messages are stored with:
+- **Embeddings**: For semantic search across conversation history
+- **Compression**: Long conversations are summarized to manage context window
+- **Session isolation**: Each session maintains its own message history
+
+### Testing Multi-Agent Systems
+
+#### Integration Tests
+
+Real end-to-end tests without mocking are in `tests/integration/test_ask_agent_streaming.py`:
+
+```python
+class TestAskAgentStreaming:
+    async def test_ask_agent_streams_and_saves(self, session_id, user_id):
+        """Test delegation via ask_agent."""
+        # Uses test_orchestrator which always delegates to test_responder
+        agent = await create_agent(context=context, agent_schema_override=schema)
+
+        chunks = []
+        async for chunk in stream_openai_response_with_save(
+            agent=agent,
+            prompt="Hello, please delegate this",
+            ...
+        ):
+            chunks.append(chunk)
+
+        # Verify streaming worked
+        assert len(content_chunks) > 0
+
+        # Verify persistence
+        messages = await store.load_session_messages(session_id)
+        assert len([m for m in messages if m["role"] == "assistant"]) == 1
+        assert len([m for m in messages if m["tool_name"] == "ask_agent"]) >= 1
+
+    async def test_multi_turn_saves_all_assistant_messages(self, session_id, user_id):
+        """Test that each turn saves its own assistant message.
+
+        This catches scoping bugs like accumulated_content not being
+        properly scoped per-turn.
+        """
+        turn_prompts = [
+            "Hello, how are you?",
+            "Tell me something interesting",
+            "Thanks for chatting!",
+        ]
+
+        for prompt in turn_prompts:
+            async for chunk in stream_openai_response_with_save(...):
+                pass
+
+        # Each turn should save an assistant message
+        messages = await store.load_session_messages(session_id)
+        assistant_msgs = [m for m in messages if m["role"] == "assistant"]
+        assert len(assistant_msgs) == 3
+```
+
+#### Test Agent Schemas
+
+Test agents are defined in `tests/data/schemas/agents/`:
+
+- **`test_orchestrator.yaml`**: Always delegates via `ask_agent`
+- **`test_responder.yaml`**: Simple agent that responds directly
+
+```yaml
+# test_orchestrator.yaml
+type: object
+description: |
+  You are a TEST ORCHESTRATOR that ALWAYS delegates to another agent.
+  Call ask_agent with agent_name="test_responder" on EVERY turn.
+json_schema_extra:
+  kind: agent
+  name: test_orchestrator
+  tools:
+    - name: ask_agent
+      mcp_server: rem
+```
+
+#### Running Integration Tests
+
+```bash
+# Run individually (recommended due to async isolation)
+POSTGRES__CONNECTION_STRING="postgresql://rem:rem@localhost:5050/rem" \
+  uv run pytest tests/integration/test_ask_agent_streaming.py::TestAskAgentStreaming::test_multi_turn_saves_all_assistant_messages -v -s
+```
+
 ## Future Work
 
 - [ ] Phoenix evaluator integration
 - [ ] Agent schema registry (load schemas by URI)
 - [ ] Schema validation and versioning
-- [ ] Multi-turn conversation management
-- [ ] Agent composition (agents calling agents)
+- [x] Multi-turn conversation management
+- [x] Agent composition (agents calling agents)
 - [ ] Alternative provider implementations (if needed)

rem/agentic/context.py

@@ -22,14 +22,153 @@ Key Design Pattern:
 - Enables session tracking across API, CLI, and test execution
 - Supports header-based configuration override (model, schema URI)
 - Clean separation: context (who/what) vs agent (how)
+
+Multi-Agent Context Propagation:
+- ContextVar (_current_agent_context) threads context through nested agent calls
+- Parent context is automatically available to child agents via get_current_context()
+- Use agent_context_scope() context manager for scoped context setting
+- Child agents inherit user_id, tenant_id, session_id, is_eval from parent
 """
 
+import asyncio
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Any, Generator
+
 from loguru import logger
 from pydantic import BaseModel, Field
 
 from ..settings import settings
 
 
+# Thread-local context for current agent execution
+# This enables context propagation through nested agent calls (multi-agent)
+_current_agent_context: ContextVar["AgentContext | None"] = ContextVar(
+    "current_agent_context", default=None
+)
+
+# Event sink for streaming child agent events to parent
+# When set, child agents (via ask_agent) should push their events here
+# for the parent's streaming loop to proxy to the client
+_parent_event_sink: ContextVar["asyncio.Queue | None"] = ContextVar(
+    "parent_event_sink", default=None
+)
+
+
+def get_current_context() -> "AgentContext | None":
+    """
+    Get the current agent context from context var.
+
+    Used by MCP tools (like ask_agent) to inherit context from parent agent.
+    Returns None if no context is set (e.g., direct CLI invocation without context).
+
+    Example:
+        # In an MCP tool
+        parent_context = get_current_context()
+        if parent_context:
+            # Inherit user_id, session_id, etc. from parent
+            child_context = parent_context.child_context(agent_schema_uri="child-agent")
+    """
+    return _current_agent_context.get()
+
+
+def set_current_context(ctx: "AgentContext | None") -> None:
+    """
+    Set the current agent context.
+
+    Called by streaming layer before agent execution.
+    Should be cleared (set to None) after execution completes.
+    """
+    _current_agent_context.set(ctx)
+
+
+@contextmanager
+def agent_context_scope(ctx: "AgentContext") -> Generator["AgentContext", None, None]:
+    """
+    Context manager for scoped context setting.
+
+    Automatically restores previous context when exiting scope.
+    Safe for nested agent calls - each level preserves its parent's context.
+
+    Example:
+        context = AgentContext(user_id="user-123")
+        with agent_context_scope(context):
+            # Context is available via get_current_context()
+            result = await agent.run(...)
+        # Previous context (or None) is restored
+    """
+    previous = _current_agent_context.get()
+    _current_agent_context.set(ctx)
+    try:
+        yield ctx
+    finally:
+        _current_agent_context.set(previous)
+
+
+# =============================================================================
+# Event Sink for Streaming Multi-Agent Delegation
+# =============================================================================
+
+
+def get_event_sink() -> "asyncio.Queue | None":
+    """
+    Get the parent's event sink for streaming child events.
+
+    Used by ask_agent to push child agent events to the parent's stream.
+    Returns None if not in a streaming context.
+    """
+    return _parent_event_sink.get()
+
+
+def set_event_sink(sink: "asyncio.Queue | None") -> None:
+    """Set the event sink for child agents to push events to."""
+    _parent_event_sink.set(sink)
+
+
+@contextmanager
+def event_sink_scope(sink: "asyncio.Queue") -> Generator["asyncio.Queue", None, None]:
+    """
+    Context manager for scoped event sink setting.
+
+    Used by streaming layer to set up event proxying before tool execution.
+    Child agents (via ask_agent) will push their events to this sink.
+
+    Example:
+        event_queue = asyncio.Queue()
+        with event_sink_scope(event_queue):
+            # ask_agent will push child events to event_queue
+            async for event in tools_stream:
+                ...
+            # Also consume from event_queue
+    """
+    previous = _parent_event_sink.get()
+    _parent_event_sink.set(sink)
+    try:
+        yield sink
+    finally:
+        _parent_event_sink.set(previous)
+
+
+async def push_event(event: Any) -> bool:
+    """
+    Push an event to the parent's event sink (if available).
+
+    Used by ask_agent to proxy child agent events to the parent's stream.
+    Returns True if event was pushed, False if no sink available.
+
+    Args:
+        event: Any streaming event (ToolCallEvent, content chunk, etc.)
+
+    Returns:
+        True if event was pushed to sink, False otherwise
+    """
+    sink = _parent_event_sink.get()
+    if sink is not None:
+        await sink.put(event)
+        return True
+    return False
+
+
 class AgentContext(BaseModel):
     """
     Session and configuration context for agent execution.
@@ -85,6 +224,40 @@ class AgentContext(BaseModel):
 
     model_config = {"populate_by_name": True}
 
+    def child_context(
+        self,
+        agent_schema_uri: str | None = None,
+        model_override: str | None = None,
+    ) -> "AgentContext":
+        """
+        Create a child context for nested agent calls.
+
+        Inherits user_id, tenant_id, session_id, is_eval from parent.
+        Allows overriding agent_schema_uri and default_model for the child.
+
+        Args:
+            agent_schema_uri: Agent schema for the child agent (required for lineage)
+            model_override: Optional model override for child agent
+
+        Returns:
+            New AgentContext for the child agent
+
+        Example:
+            parent_context = get_current_context()
+            child_context = parent_context.child_context(
+                agent_schema_uri="sentiment-analyzer"
+            )
+            agent = await create_agent(context=child_context)
+        """
+        return AgentContext(
+            user_id=self.user_id,
+            tenant_id=self.tenant_id,
+            session_id=self.session_id,
+            default_model=model_override or self.default_model,
+            agent_schema_uri=agent_schema_uri or self.agent_schema_uri,
+            is_eval=self.is_eval,
+        )
+
     @staticmethod
     def get_user_id_or_default(
         user_id: str | None,

rem/agentic/context_builder.py

@@ -217,11 +217,21 @@ class ContextBuilder:
                 )
 
                 # Convert to ContextMessage format
+                # For tool messages, wrap content with clear markers so the agent
+                # can see previous tool results when the prompt is concatenated
                 for msg_dict in session_history:
+                    role = msg_dict["role"]
+                    content = msg_dict["content"]
+
+                    if role == "tool":
+                        # Wrap tool results with clear markers for visibility
+                        tool_name = msg_dict.get("tool_name", "unknown")
+                        content = f"[TOOL RESULT: {tool_name}]\n{content}\n[/TOOL RESULT]"
+
                     messages.append(
                         ContextMessage(
-                            role=msg_dict["role"],
-                            content=msg_dict["content"],
+                            role=role,
+                            content=content,
                         )
                     )
 

rem/agentic/mcp/tool_wrapper.py

@@ -116,7 +116,7 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
     the artificial MCP distinction between tools and resources.
 
     Supports both:
-    - Concrete URIs: "rem://schemas" -> tool with no parameters
+    - Concrete URIs: "rem://agents" -> tool with no parameters
     - Template URIs: "patient-profile://field/{field_key}" -> tool with field_key parameter
 
     Args:
@@ -131,7 +131,7 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
 
     Example:
         # Concrete URI -> no-param tool
-        tool = create_resource_tool("rem://schemas", "List all agent schemas")
+        tool = create_resource_tool("rem://agents", "List all agent schemas")
 
         # Template URI -> parameterized tool
         tool = create_resource_tool("patient-profile://field/{field_key}", "Get field definition", mcp_server=mcp)
@@ -161,6 +161,11 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
         param_desc = ", ".join(template_vars)
         description = f"{description}\n\nParameters: {param_desc}"
 
+    # Capture mcp_server reference at tool creation time (for closure)
+    # This ensures the correct server is used even if called later
+    _captured_mcp_server = mcp_server
+    _captured_uri = uri  # Also capture URI for consistent logging
+
     if template_vars:
         # Template URI -> create parameterized tool
         async def wrapper(**kwargs: Any) -> str:
@@ -168,13 +173,17 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
             import asyncio
             import inspect
 
+            logger.debug(f"Resource tool invoked: uri={_captured_uri}, kwargs={kwargs}, mcp_server={'set' if _captured_mcp_server else 'None'}")
+
             # Try to resolve from MCP server's resource templates first
-            if mcp_server is not None:
+            if _captured_mcp_server is not None:
                 try:
                     # Get resource templates from MCP server
-                    templates = await mcp_server.get_resource_templates()
-                    if uri in templates:
-                        template = templates[uri]
+                    templates = await _captured_mcp_server.get_resource_templates()
+                    logger.debug(f"MCP server templates: {list(templates.keys())}")
+                    if _captured_uri in templates:
+                        template = templates[_captured_uri]
+                        logger.debug(f"Found template for {_captured_uri}, calling fn with kwargs={kwargs}")
                         # Call the template's underlying function directly
                         # The fn expects the template variables as kwargs
                         fn_result = template.fn(**kwargs)
@@ -184,17 +193,22 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
                         if isinstance(fn_result, str):
                             return fn_result
                         return json.dumps(fn_result, indent=2)
+                    else:
+                        logger.warning(f"Template {_captured_uri} not found in MCP server templates: {list(templates.keys())}")
                 except Exception as e:
-                    logger.warning(f"Failed to resolve resource {uri} from MCP server: {e}")
+                    logger.warning(f"Failed to resolve resource {_captured_uri} from MCP server: {e}", exc_info=True)
+            else:
+                logger.warning(f"No MCP server provided for resource tool {_captured_uri}, using fallback")
 
             # Fallback: substitute template variables and use load_resource
-            resolved_uri = uri
+            resolved_uri = _captured_uri
             for var in template_vars:
                 if var in kwargs:
                     resolved_uri = resolved_uri.replace(f"{{{var}}}", str(kwargs[var]))
                 else:
                     return json.dumps({"error": f"Missing required parameter: {var}"})
 
+            logger.debug(f"Using fallback load_resource for resolved URI: {resolved_uri}")
             from rem.api.mcp_router.resources import load_resource
             result = await load_resource(resolved_uri)
             if isinstance(result, str):
@@ -208,7 +222,7 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
         wrapper.__annotations__ = {var: str for var in template_vars}
         wrapper.__annotations__['return'] = str
 
-        logger.info(f"Built parameterized resource tool: {func_name} (uri: {uri}, params: {template_vars})")
+        logger.info(f"Built parameterized resource tool: {func_name} (uri: {uri}, params: {template_vars}, mcp_server={'provided' if mcp_server else 'None'})")
     else:
         # Concrete URI -> no-param tool
         async def wrapper(**kwargs: Any) -> str:
@@ -219,12 +233,16 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
             if kwargs:
                 logger.warning(f"Resource tool {func_name} called with unexpected kwargs: {list(kwargs.keys())}")
 
+            logger.debug(f"Concrete resource tool invoked: uri={_captured_uri}, mcp_server={'set' if _captured_mcp_server else 'None'}")
+
             # Try to resolve from MCP server's resources first
-            if mcp_server is not None:
+            if _captured_mcp_server is not None:
                 try:
-                    resources = await mcp_server.get_resources()
-                    if uri in resources:
-                        resource = resources[uri]
+                    resources = await _captured_mcp_server.get_resources()
+                    logger.debug(f"MCP server resources: {list(resources.keys())}")
+                    if _captured_uri in resources:
+                        resource = resources[_captured_uri]
+                        logger.debug(f"Found resource for {_captured_uri}")
                         # Call the resource's underlying function
                         fn_result = resource.fn()
                         if inspect.iscoroutine(fn_result):
@@ -232,12 +250,17 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
                         if isinstance(fn_result, str):
                             return fn_result
                         return json.dumps(fn_result, indent=2)
+                    else:
+                        logger.warning(f"Resource {_captured_uri} not found in MCP server resources: {list(resources.keys())}")
                 except Exception as e:
-                    logger.warning(f"Failed to resolve resource {uri} from MCP server: {e}")
+                    logger.warning(f"Failed to resolve resource {_captured_uri} from MCP server: {e}", exc_info=True)
+            else:
+                logger.warning(f"No MCP server provided for resource tool {_captured_uri}, using fallback")
 
             # Fallback to load_resource
+            logger.debug(f"Using fallback load_resource for URI: {_captured_uri}")
             from rem.api.mcp_router.resources import load_resource
-            result = await load_resource(uri)
+            result = await load_resource(_captured_uri)
             if isinstance(result, str):
                 return result
             return json.dumps(result, indent=2)
@@ -245,6 +268,6 @@ def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None) -> T
         wrapper.__name__ = func_name
         wrapper.__doc__ = description
 
-        logger.info(f"Built resource tool: {func_name} (uri: {uri})")
+        logger.info(f"Built resource tool: {func_name} (uri: {uri}, mcp_server={'provided' if mcp_server else 'None'})")
 
     return Tool(wrapper)