atlas-chat 0.1.0__py3-none-any.whl

atlas/application/chat/elicitation_manager.py

@@ -0,0 +1,191 @@
+"""
+Elicitation manager for handling user input requests during tool execution.
+
+This manager coordinates between MCP servers requesting user input (via ctx.elicit())
+and the frontend UI collecting that input. It provides a synchronization mechanism
+where tool execution pauses until the user responds.
+"""
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ElicitationRequest:
+    """Represents a pending elicitation request awaiting user response."""
+    elicitation_id: str
+    tool_call_id: str
+    tool_name: str
+    message: str
+    response_schema: Dict[str, Any]
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    future: asyncio.Future = field(default_factory=lambda: asyncio.get_event_loop().create_future())
+
+    async def wait_for_response(self, timeout: float = 300.0) -> Dict[str, Any]:
+        """
+        Wait for the user to respond to the elicitation request.
+
+        Args:
+            timeout: Maximum time to wait in seconds (default 5 minutes)
+
+        Returns:
+            Dict with 'action' and optionally 'data' keys
+
+        Raises:
+            asyncio.TimeoutError: If timeout is reached
+        """
+        return await asyncio.wait_for(self.future, timeout=timeout)
+
+
+class ElicitationManager:
+    """
+    Manages elicitation requests and responses.
+
+    Provides synchronization between:
+    - MCP servers requesting input via ctx.elicit()
+    - Frontend UI collecting user responses
+    """
+
+    def __init__(self):
+        """Initialize the elicitation manager."""
+        self._pending_requests: Dict[str, ElicitationRequest] = {}
+        self._lock = asyncio.Lock()
+
+    def create_elicitation_request(
+        self,
+        elicitation_id: str,
+        tool_call_id: str,
+        tool_name: str,
+        message: str,
+        response_schema: Dict[str, Any]
+    ) -> ElicitationRequest:
+        """
+        Create a new elicitation request.
+
+        Args:
+            elicitation_id: Unique identifier for this elicitation
+            tool_call_id: ID of the tool call that requested input
+            tool_name: Name of the tool requesting input
+            message: Prompt message to display to user
+            response_schema: JSON schema defining expected response structure
+
+        Returns:
+            ElicitationRequest object that can be awaited for response
+        """
+        request = ElicitationRequest(
+            elicitation_id=elicitation_id,
+            tool_call_id=tool_call_id,
+            tool_name=tool_name,
+            message=message,
+            response_schema=response_schema
+        )
+        self._pending_requests[elicitation_id] = request
+        logger.info(
+            f"Created elicitation request: id={elicitation_id}, "
+            f"tool={tool_name}, message='{message[:50]}...'"
+        )
+        return request
+
+    def handle_elicitation_response(
+        self,
+        elicitation_id: str,
+        action: str,
+        data: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """
+        Handle an elicitation response from the user.
+
+        Args:
+            elicitation_id: ID of the elicitation being responded to
+            action: User action - "accept", "decline", or "cancel"
+            data: Optional response data (present when action is "accept")
+
+        Returns:
+            True if response was handled, False if request not found
+        """
+        request = self._pending_requests.get(elicitation_id)
+        if not request:
+            logger.warning(f"Received response for unknown elicitation: {elicitation_id}")
+            return False
+
+        response = {
+            "action": action,
+            "data": data
+        }
+
+        if not request.future.done():
+            request.future.set_result(response)
+            logger.info(
+                f"Elicitation response received: id={elicitation_id}, "
+                f"action={action}, has_data={data is not None}"
+            )
+        else:
+            logger.warning(
+                f"Elicitation response ignored (already resolved): {elicitation_id}"
+            )
+
+        return True
+
+    def cleanup_request(self, elicitation_id: str) -> None:
+        """
+        Clean up a completed elicitation request.
+
+        Args:
+            elicitation_id: ID of the request to clean up
+        """
+        if elicitation_id in self._pending_requests:
+            del self._pending_requests[elicitation_id]
+            logger.debug(f"Cleaned up elicitation request: {elicitation_id}")
+
+    def get_pending_request(self, elicitation_id: str) -> Optional[ElicitationRequest]:
+        """
+        Get a pending elicitation request by ID.
+
+        Args:
+            elicitation_id: ID of the request to retrieve
+
+        Returns:
+            ElicitationRequest if found, None otherwise
+        """
+        return self._pending_requests.get(elicitation_id)
+
+    def get_all_pending_requests(self) -> Dict[str, ElicitationRequest]:
+        """
+        Get all pending elicitation requests.
+
+        Returns:
+            Dictionary mapping elicitation IDs to requests
+        """
+        return dict(self._pending_requests)
+
+    def cancel_all_requests(self) -> None:
+        """Cancel all pending elicitation requests."""
+        for request in self._pending_requests.values():
+            if not request.future.done():
+                request.future.set_exception(
+                    asyncio.CancelledError("Elicitation cancelled")
+                )
+        self._pending_requests.clear()
+        logger.info("Cancelled all pending elicitation requests")
+
+
+# Global singleton instance
+_elicitation_manager: Optional[ElicitationManager] = None
+
+
+def get_elicitation_manager() -> ElicitationManager:
+    """
+    Get the global elicitation manager singleton.
+
+    Returns:
+        Global ElicitationManager instance
+    """
+    global _elicitation_manager
+    if _elicitation_manager is None:
+        _elicitation_manager = ElicitationManager()
+    return _elicitation_manager

atlas/application/chat/events/__init__.py

@@ -0,0 +1 @@
+"""Event modules for chat application."""

atlas/application/chat/events/agent_event_relay.py

@@ -0,0 +1,112 @@
+"""Agent event relay - maps AgentEvents to EventPublisher calls."""
+
+import logging
+from typing import Any, Awaitable, Callable, Optional
+
+from atlas.interfaces.events import EventPublisher
+
+from ..agent.protocols import AgentEvent
+
+logger = logging.getLogger(__name__)
+
+# Constants
+UNKNOWN_TOOL_NAME = "unknown"
+
+
+class AgentEventRelay:
+    """
+    Translates agent loop events to UI update events.
+
+    Maps AgentEvent instances to appropriate EventPublisher method calls,
+    providing a clean separation between agent logic and UI transport.
+    """
+
+    def __init__(
+        self,
+        event_publisher: EventPublisher,
+        artifact_processor: Optional[Callable[[Any], Awaitable[None]]] = None,
+    ):
+        """
+        Initialize agent event relay.
+
+        Args:
+            event_publisher: Publisher for sending UI updates
+            artifact_processor: Optional callback for processing tool artifacts
+        """
+        self.event_publisher = event_publisher
+        self.artifact_processor = artifact_processor
+
+    async def handle_event(self, evt: AgentEvent) -> None:
+        """
+        Handle an agent event and relay it to the UI.
+
+        Args:
+            evt: Agent event to handle
+        """
+        et = evt.type
+        p = evt.payload or {}
+
+        # Map event types to publisher calls
+        if et == "agent_start":
+            await self.event_publisher.publish_agent_update(
+                update_type="agent_start",
+                max_steps=p.get("max_steps"),
+                strategy=p.get("strategy"),
+            )
+
+        elif et == "agent_turn_start":
+            await self.event_publisher.publish_agent_update(
+                update_type="agent_turn_start",
+                step=p.get("step"),
+            )
+
+        elif et == "agent_reason":
+            await self.event_publisher.publish_agent_update(
+                update_type="agent_reason",
+                message=p.get("message"),
+                step=p.get("step"),
+            )
+
+        elif et == "agent_request_input":
+            await self.event_publisher.publish_agent_update(
+                update_type="agent_request_input",
+                question=p.get("question"),
+                step=p.get("step"),
+            )
+
+        elif et == "agent_tool_start":
+            await self.event_publisher.publish_tool_start(
+                tool_name=p.get("tool", UNKNOWN_TOOL_NAME),
+            )
+
+        elif et == "agent_tool_complete":
+            await self.event_publisher.publish_tool_complete(
+                tool_name=p.get("tool", UNKNOWN_TOOL_NAME),
+                result=p.get("result"),
+            )
+
+        elif et == "agent_tool_results":
+            # Delegate artifact processing to external handler
+            if self.artifact_processor:
+                results = p.get("results") or []
+                if results:
+                    await self.artifact_processor(results)
+
+        elif et == "agent_observe":
+            await self.event_publisher.publish_agent_update(
+                update_type="agent_observe",
+                message=p.get("message"),
+                step=p.get("step"),
+            )
+
+        elif et == "agent_completion":
+            await self.event_publisher.publish_agent_update(
+                update_type="agent_completion",
+                steps=p.get("steps"),
+            )
+
+        elif et == "agent_error":
+            await self.event_publisher.publish_agent_update(
+                update_type="agent_error",
+                message=p.get("message"),
+            )

atlas/application/chat/modes/__init__.py

@@ -0,0 +1 @@
+"""Mode runner modules for different chat execution modes."""

atlas/application/chat/modes/agent.py

@@ -0,0 +1,125 @@
+"""Agent mode runner - handles LLM calls with agent loop execution."""
+
+import logging
+from typing import Any, Awaitable, Callable, Dict, List, Optional
+
+from atlas.domain.messages.models import Message, MessageRole, ToolResult
+from atlas.domain.sessions.models import Session
+from atlas.interfaces.events import EventPublisher
+
+from ..agent import AgentLoopFactory
+from ..agent.protocols import AgentContext
+from ..events.agent_event_relay import AgentEventRelay
+from ..utilities import event_notifier
+
+logger = logging.getLogger(__name__)
+
+# Type hint for the update callback
+UpdateCallback = Callable[[Dict[str, Any]], Awaitable[None]]
+
+
+class AgentModeRunner:
+    """
+    Runner for agent mode.
+
+    Executes agent loops with event streaming and artifact processing.
+    """
+
+    def __init__(
+        self,
+        agent_loop_factory: AgentLoopFactory,
+        event_publisher: EventPublisher,
+        artifact_processor: Optional[Callable[[Session, List[ToolResult], Optional[UpdateCallback]], Awaitable[None]]] = None,
+        default_strategy: str = "think-act",
+    ):
+        """
+        Initialize agent mode runner.
+
+        Args:
+            agent_loop_factory: Factory for creating agent loops
+            event_publisher: Event publisher for UI updates
+            artifact_processor: Optional callback for processing tool artifacts
+            default_strategy: Default agent loop strategy
+        """
+        self.agent_loop_factory = agent_loop_factory
+        self.event_publisher = event_publisher
+        self.artifact_processor = artifact_processor
+        self.default_strategy = default_strategy
+
+    async def run(
+        self,
+        session: Session,
+        model: str,
+        messages: List[Dict[str, Any]],
+        selected_tools: Optional[List[str]],
+        selected_data_sources: Optional[List[str]],
+        max_steps: int,
+        temperature: float = 0.7,
+        agent_loop_strategy: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Execute agent mode.
+
+        Args:
+            session: Current chat session
+            model: LLM model to use
+            messages: Message history
+            selected_tools: Optional list of tools to make available
+            selected_data_sources: Optional list of data sources
+            max_steps: Maximum number of agent steps
+            temperature: LLM temperature parameter
+            agent_loop_strategy: Strategy name (react, think-act). Falls back to default.
+
+        Returns:
+            Response dictionary
+        """
+        # Get agent loop from factory based on strategy
+        strategy = agent_loop_strategy or self.default_strategy
+        agent_loop = self.agent_loop_factory.create(strategy)
+
+        # Build agent context
+        agent_context = AgentContext(
+            session_id=session.id,
+            user_email=session.user_email,
+            files=session.context.get("files", {}),
+            history=session.history,
+        )
+
+        # Artifact processor wrapper for handling tool results
+        async def process_artifacts(results):
+            if self.artifact_processor:
+                await self.artifact_processor(session, results, None)
+
+        # Create event relay to map AgentEvents to UI updates
+        event_relay = AgentEventRelay(
+            event_publisher=self.event_publisher,
+            artifact_processor=process_artifacts,
+        )
+
+        # Run the loop
+        result = await agent_loop.run(
+            model=model,
+            messages=messages,
+            context=agent_context,
+            selected_tools=selected_tools,
+            data_sources=selected_data_sources,
+            max_steps=max_steps,
+            temperature=temperature,
+            event_handler=event_relay.handle_event,
+        )
+
+        # Append final message
+        assistant_message = Message(
+            role=MessageRole.ASSISTANT,
+            content=result.final_answer,
+            metadata={"agent_mode": True, "steps": result.steps},
+        )
+        session.history.add_message(assistant_message)
+
+        # Completion update
+        await self.event_publisher.publish_agent_update(
+            update_type="agent_completion",
+            steps=result.steps
+        )
+
+        return event_notifier.create_chat_response(result.final_answer)

atlas/application/chat/modes/plain.py

@@ -0,0 +1,74 @@
+"""Plain mode runner - handles simple LLM calls without tools or RAG."""
+
+import logging
+from typing import Any, Dict, List
+
+from atlas.domain.messages.models import Message, MessageRole
+from atlas.domain.sessions.models import Session
+from atlas.interfaces.events import EventPublisher
+from atlas.interfaces.llm import LLMProtocol
+
+from ..utilities import event_notifier
+
+logger = logging.getLogger(__name__)
+
+
+class PlainModeRunner:
+    """
+    Runner for plain LLM mode.
+
+    Executes simple LLM calls without tools or RAG integration.
+    """
+
+    def __init__(
+        self,
+        llm: LLMProtocol,
+        event_publisher: EventPublisher,
+    ):
+        """
+        Initialize plain mode runner.
+
+        Args:
+            llm: LLM protocol implementation
+            event_publisher: Event publisher for UI updates
+        """
+        self.llm = llm
+        self.event_publisher = event_publisher
+
+    async def run(
+        self,
+        session: Session,
+        model: str,
+        messages: List[Dict[str, str]],
+        temperature: float = 0.7,
+    ) -> Dict[str, Any]:
+        """
+        Execute plain LLM mode.
+
+        Args:
+            session: Current chat session
+            model: LLM model to use
+            messages: Message history
+            temperature: LLM temperature parameter
+
+        Returns:
+            Response dictionary
+        """
+        # Call LLM
+        response_content = await self.llm.call_plain(model, messages, temperature=temperature)
+
+        # Add assistant message to history
+        assistant_message = Message(
+            role=MessageRole.ASSISTANT,
+            content=response_content
+        )
+        session.history.add_message(assistant_message)
+
+        # Publish events
+        await self.event_publisher.publish_chat_response(
+            message=response_content,
+            has_pending_tools=False,
+        )
+        await self.event_publisher.publish_response_complete()
+
+        return event_notifier.create_chat_response(response_content)

atlas/application/chat/modes/rag.py

@@ -0,0 +1,81 @@
+"""RAG mode runner - handles LLM calls with RAG integration."""
+
+import logging
+from typing import Any, Dict, List
+
+from atlas.domain.messages.models import Message, MessageRole
+from atlas.domain.sessions.models import Session
+from atlas.interfaces.events import EventPublisher
+from atlas.interfaces.llm import LLMProtocol
+
+from ..utilities import event_notifier
+
+logger = logging.getLogger(__name__)
+
+
+class RagModeRunner:
+    """
+    Runner for RAG mode.
+
+    Executes LLM calls with Retrieval-Augmented Generation integration.
+    """
+
+    def __init__(
+        self,
+        llm: LLMProtocol,
+        event_publisher: EventPublisher,
+    ):
+        """
+        Initialize RAG mode runner.
+
+        Args:
+            llm: LLM protocol implementation
+            event_publisher: Event publisher for UI updates
+        """
+        self.llm = llm
+        self.event_publisher = event_publisher
+
+    async def run(
+        self,
+        session: Session,
+        model: str,
+        messages: List[Dict[str, str]],
+        data_sources: List[str],
+        user_email: str,
+        temperature: float = 0.7,
+    ) -> Dict[str, Any]:
+        """
+        Execute RAG mode.
+
+        Args:
+            session: Current chat session
+            model: LLM model to use
+            messages: Message history
+            data_sources: List of data sources to query
+            user_email: User email for authorization
+            temperature: LLM temperature parameter
+
+        Returns:
+            Response dictionary
+        """
+        # Call LLM with RAG
+        response_content = await self.llm.call_with_rag(
+            model, messages, data_sources, user_email, temperature=temperature
+        )
+
+        # Add assistant message to history
+        assistant_message = Message(
+            role=MessageRole.ASSISTANT,
+            content=response_content,
+            metadata={"data_sources": data_sources}
+        )
+        session.history.add_message(assistant_message)
+
+        # Publish events
+        await self.event_publisher.publish_chat_response(
+            message=response_content,
+            has_pending_tools=False,
+        )
+        await self.event_publisher.publish_response_complete()
+
+        return event_notifier.create_chat_response(response_content)