atlas-chat 0.1.0__py3-none-any.whl

atlas/application/chat/modes/tools.py

@@ -0,0 +1,179 @@
+"""Tools mode runner - handles LLM calls with tool 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 atlas.interfaces.llm import LLMProtocol
+from atlas.interfaces.tools import ToolManagerProtocol
+from atlas.modules.prompts.prompt_provider import PromptProvider
+
+from ..preprocessors.message_builder import build_session_context
+from ..utilities import error_handler, event_notifier, tool_executor
+
+logger = logging.getLogger(__name__)
+
+# Type hint for the update callback
+UpdateCallback = Callable[[Dict[str, Any]], Awaitable[None]]
+
+
+class ToolsModeRunner:
+    """
+    Runner for tools mode.
+
+    Executes LLM calls with tool integration, including tool execution
+    and artifact processing.
+    """
+
+    def __init__(
+        self,
+        llm: LLMProtocol,
+        tool_manager: ToolManagerProtocol,
+        event_publisher: EventPublisher,
+        prompt_provider: Optional[PromptProvider] = None,
+        artifact_processor: Optional[Callable[[Session, List[ToolResult], Optional[UpdateCallback]], Awaitable[None]]] = None,
+        config_manager=None,
+    ):
+        """
+        Initialize tools mode runner.
+
+        Args:
+            llm: LLM protocol implementation
+            tool_manager: Tool manager for tool execution
+            event_publisher: Event publisher for UI updates
+            prompt_provider: Optional prompt provider
+            artifact_processor: Optional callback for processing tool artifacts
+            config_manager: Optional config manager for approval settings
+        """
+        self.llm = llm
+        self.tool_manager = tool_manager
+        self.event_publisher = event_publisher
+        self.prompt_provider = prompt_provider
+        self.artifact_processor = artifact_processor
+        self.config_manager = config_manager
+        self.skip_approval = False
+
+        # Verify event_publisher has send_json for elicitation support
+        if hasattr(event_publisher, 'send_json'):
+            logger.debug(f"ToolsModeRunner initialized with event_publisher that has send_json: {type(event_publisher)}")
+        else:
+            logger.warning(f"ToolsModeRunner initialized with event_publisher WITHOUT send_json: {type(event_publisher)}")
+
+    async def run(
+        self,
+        session: Session,
+        model: str,
+        messages: List[Dict[str, Any]],
+        selected_tools: List[str],
+        selected_data_sources: Optional[List[str]] = None,
+        user_email: Optional[str] = None,
+        tool_choice_required: bool = False,
+        update_callback: Optional[UpdateCallback] = None,
+        temperature: float = 0.7,
+    ) -> Dict[str, Any]:
+        """
+        Execute tools mode.
+
+        Args:
+            session: Current chat session
+            model: LLM model to use
+            messages: Message history
+            selected_tools: List of tools to make available
+            selected_data_sources: Optional list of data sources (for RAG+tools)
+            user_email: Optional user email for authorization
+            tool_choice_required: Whether tool use is required
+            update_callback: Optional callback for streaming updates
+            temperature: LLM temperature parameter
+
+        Returns:
+            Response dictionary
+        """
+        # Resolve tool schemas
+        tools_schema = await error_handler.safe_get_tools_schema(self.tool_manager, selected_tools)
+
+        # Call LLM with tools (and RAG if provided)
+        llm_response = await error_handler.safe_call_llm_with_tools(
+            llm_caller=self.llm,
+            model=model,
+            messages=messages,
+            tools_schema=tools_schema,
+            data_sources=selected_data_sources,
+            user_email=user_email,
+            tool_choice=("required" if tool_choice_required else "auto"),
+            temperature=temperature,
+        )
+
+        # No tool calls -> treat as plain content
+        if not llm_response or not llm_response.has_tool_calls():
+            content = llm_response.content if llm_response else ""
+            assistant_message = Message(role=MessageRole.ASSISTANT, content=content)
+            session.history.add_message(assistant_message)
+
+            await self.event_publisher.publish_chat_response(
+                message=content,
+                has_pending_tools=False,
+            )
+            await self.event_publisher.publish_response_complete()
+
+            return event_notifier.create_chat_response(content)
+
+        # Execute tool workflow
+        session_context = build_session_context(session)
+
+        # Ensure update_callback is never None (critical for elicitation)
+        effective_callback = update_callback
+        if effective_callback is None:
+            effective_callback = self._get_send_json()
+            logger.debug("Tools mode: update_callback was None, using event_publisher.send_json fallback")
+
+        if effective_callback is None:
+            logger.warning("Tools mode: No update callback available - elicitation will not work!")
+
+        final_response, tool_results = await tool_executor.execute_tools_workflow(
+            llm_response=llm_response,
+            messages=messages,
+            model=model,
+            session_context=session_context,
+            tool_manager=self.tool_manager,
+            llm_caller=self.llm,
+            prompt_provider=self.prompt_provider,
+            update_callback=effective_callback,
+            config_manager=self.config_manager,
+            skip_approval=self.skip_approval,
+            user_email=user_email,
+        )
+
+        # Process artifacts if handler provided
+        if self.artifact_processor:
+            await self.artifact_processor(session, tool_results, effective_callback)
+
+        # Add final assistant message to history
+        assistant_message = Message(
+            role=MessageRole.ASSISTANT,
+            content=final_response,
+            metadata={
+                "tools": selected_tools,
+                **({"data_sources": selected_data_sources} if selected_data_sources else {}),
+            },
+        )
+        session.history.add_message(assistant_message)
+
+        # Emit final chat response
+        await self.event_publisher.publish_chat_response(
+            message=final_response,
+            has_pending_tools=False,
+        )
+        await self.event_publisher.publish_response_complete()
+
+        return event_notifier.create_chat_response(final_response)
+
+    def _get_send_json(self) -> Optional[UpdateCallback]:
+        """Get send_json callback from event publisher if available."""
+        if hasattr(self.event_publisher, 'send_json'):
+            callback = self.event_publisher.send_json
+            logger.debug(f"_get_send_json: event_publisher.send_json = {callback is not None}")
+            return callback
+        logger.warning(f"_get_send_json: event_publisher does not have send_json method. Type: {type(self.event_publisher)}")
+        return None

atlas/application/chat/orchestrator.py

@@ -0,0 +1,213 @@
+"""Chat orchestrator - coordinates the full chat request flow."""
+
+import logging
+from typing import Any, Dict, List, Optional
+from uuid import UUID
+
+from atlas.domain.errors import SessionNotFoundError
+from atlas.domain.messages.models import Message, MessageRole
+from atlas.interfaces.events import EventPublisher
+from atlas.interfaces.llm import LLMProtocol
+from atlas.interfaces.sessions import SessionRepository
+from atlas.interfaces.tools import ToolManagerProtocol
+from atlas.modules.prompts.prompt_provider import PromptProvider
+
+from .modes.agent import AgentModeRunner
+from .modes.plain import PlainModeRunner
+from .modes.rag import RagModeRunner
+from .modes.tools import ToolsModeRunner
+from .policies.tool_authorization import ToolAuthorizationService
+from .preprocessors.message_builder import MessageBuilder
+from .preprocessors.prompt_override_service import PromptOverrideService
+from .utilities import file_processor
+
+logger = logging.getLogger(__name__)
+
+
+class ChatOrchestrator:
+    """
+    Orchestrates the full chat request flow.
+
+    Coordinates preprocessing, policy checks, mode selection, and execution.
+    Provides clean separation between request handling and business logic.
+    """
+
+    def __init__(
+        self,
+        llm: LLMProtocol,
+        event_publisher: EventPublisher,
+        session_repository: SessionRepository,
+        tool_manager: Optional[ToolManagerProtocol] = None,
+        prompt_provider: Optional[PromptProvider] = None,
+        file_manager: Optional[Any] = None,
+        artifact_processor: Optional[Any] = None,
+        plain_mode: Optional[PlainModeRunner] = None,
+        rag_mode: Optional[RagModeRunner] = None,
+        tools_mode: Optional[ToolsModeRunner] = None,
+        agent_mode: Optional[AgentModeRunner] = None,
+    ):
+        """
+        Initialize chat orchestrator.
+
+        Args:
+            llm: LLM protocol implementation
+            event_publisher: Event publisher for UI updates
+            session_repository: Session storage repository
+            tool_manager: Optional tool manager
+            prompt_provider: Optional prompt provider
+            file_manager: Optional file manager
+            artifact_processor: Optional artifact processor callback
+            plain_mode: Optional pre-configured plain mode runner
+            rag_mode: Optional pre-configured RAG mode runner
+            tools_mode: Optional pre-configured tools mode runner
+            agent_mode: Optional pre-configured agent mode runner
+        """
+        self.llm = llm
+        self.event_publisher = event_publisher
+        self.session_repository = session_repository
+        self.tool_manager = tool_manager
+        self.prompt_provider = prompt_provider
+        self.file_manager = file_manager
+
+        # Initialize services
+        self.tool_authorization = ToolAuthorizationService(tool_manager=tool_manager)
+        self.prompt_override = PromptOverrideService(tool_manager=tool_manager)
+        self.message_builder = MessageBuilder(prompt_provider=prompt_provider)
+
+        # Initialize or use provided mode runners
+        self.plain_mode = plain_mode or PlainModeRunner(
+            llm=llm,
+            event_publisher=event_publisher,
+        )
+        self.rag_mode = rag_mode or RagModeRunner(
+            llm=llm,
+            event_publisher=event_publisher,
+        )
+        self.tools_mode = tools_mode or ToolsModeRunner(
+            llm=llm,
+            tool_manager=tool_manager,
+            event_publisher=event_publisher,
+            prompt_provider=prompt_provider,
+            artifact_processor=artifact_processor,
+        )
+        self.agent_mode = agent_mode
+
+    async def execute(
+        self,
+        session_id: UUID,
+        content: str,
+        model: str,
+        user_email: Optional[str] = None,
+        selected_tools: Optional[List[str]] = None,
+        selected_prompts: Optional[List[str]] = None,
+        selected_data_sources: Optional[List[str]] = None,
+        only_rag: bool = False,
+        tool_choice_required: bool = False,
+        agent_mode: bool = False,
+        temperature: float = 0.7,
+        files: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Execute a chat request through the full pipeline.
+
+        Args:
+            session_id: Session identifier
+            content: User message content
+            model: LLM model to use
+            user_email: Optional user email
+            selected_tools: Optional list of tools
+            selected_prompts: Optional list of MCP prompts
+            selected_data_sources: Optional list of data sources
+            only_rag: Whether to use only RAG (no tools)
+            tool_choice_required: Whether tool use is required
+            agent_mode: Whether to use agent mode
+            temperature: LLM temperature
+            files: Optional files to attach
+            **kwargs: Additional parameters
+
+        Returns:
+            Response dictionary
+        """
+        # Get session from repository
+        session = await self.session_repository.get(session_id)
+        if not session:
+            raise SessionNotFoundError(f"Session {session_id} not found")
+
+        # Add user message to history
+        user_message = Message(
+            role=MessageRole.USER,
+            content=content,
+            metadata={"model": model}
+        )
+        session.history.add_message(user_message)
+        session.update_timestamp()
+
+        # Handle file ingestion
+        update_callback = kwargs.get("update_callback")
+        logger.debug(f"Orchestrator.execute: update_callback present = {update_callback is not None}")
+        session.context = await file_processor.handle_session_files(
+            session_context=session.context,
+            user_email=user_email,
+            files_map=files,
+            file_manager=self.file_manager,
+            update_callback=update_callback
+        )
+
+        # Build messages with history and files manifest
+        messages = await self.message_builder.build_messages(
+            session=session,
+            include_files_manifest=True
+        )
+
+        # Apply MCP prompt override
+        messages = await self.prompt_override.apply_prompt_override(
+            messages=messages,
+            selected_prompts=selected_prompts
+        )
+
+        # Route to appropriate mode
+        if agent_mode and self.agent_mode:
+            return await self.agent_mode.run(
+                session=session,
+                model=model,
+                messages=messages,
+                selected_tools=selected_tools,
+                selected_data_sources=selected_data_sources,
+                max_steps=kwargs.get("agent_max_steps", 30),
+                temperature=temperature,
+                agent_loop_strategy=kwargs.get("agent_loop_strategy"),
+            )
+        elif selected_tools and not only_rag:
+            # Apply tool authorization
+            selected_tools = await self.tool_authorization.filter_authorized_tools(
+                selected_tools=selected_tools,
+                user_email=user_email
+            )
+            return await self.tools_mode.run(
+                session=session,
+                model=model,
+                messages=messages,
+                selected_tools=selected_tools,
+                selected_data_sources=selected_data_sources,
+                user_email=user_email,
+                tool_choice_required=tool_choice_required,
+                update_callback=update_callback,
+                temperature=temperature,
+            )
+        elif selected_data_sources:
+            return await self.rag_mode.run(
+                session=session,
+                model=model,
+                messages=messages,
+                data_sources=selected_data_sources,
+                user_email=user_email,
+                temperature=temperature,
+            )
+        else:
+            return await self.plain_mode.run(
+                session=session,
+                model=model,
+                messages=messages,
+                temperature=temperature,
+            )

atlas/application/chat/policies/__init__.py

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

atlas/application/chat/policies/tool_authorization.py

@@ -0,0 +1,99 @@
+"""Tool authorization policy - filters tools based on user access control."""
+
+import logging
+from typing import Any, List, Optional
+
+from atlas.core.auth import is_user_in_group
+
+logger = logging.getLogger(__name__)
+
+
+class ToolAuthorizationService:
+    """
+    Service that filters selected tools based on user authorization.
+
+    Enforces MCP tool access control lists (ACLs) by checking:
+    - User authorization to MCP servers
+    - Special cases (e.g., canvas_canvas tool is always allowed)
+    """
+
+    def __init__(self, tool_manager: Optional[Any] = None):
+        """
+        Initialize the tool authorization service.
+
+        Args:
+            tool_manager: Optional tool manager with server configuration
+        """
+        self.tool_manager = tool_manager
+
+    async def filter_authorized_tools(
+        self,
+        selected_tools: List[str],
+        user_email: Optional[str] = None,
+    ) -> List[str]:
+        """
+        Filter tools to only those the user is authorized to use.
+
+        Args:
+            selected_tools: List of tool names (format: "server_toolname")
+            user_email: Email of the user making the request
+
+        Returns:
+            Filtered list of authorized tool names
+        """
+        if not selected_tools or not self.tool_manager:
+            return selected_tools or []
+
+        try:
+            user = user_email or ""
+
+            # Get authorized servers for this user
+            authorized_servers = await self._get_authorized_servers(user)
+
+            # Filter tools by server prefix
+            filtered_tools: List[str] = []
+            for tool in selected_tools:
+                # Special case: canvas_canvas is always allowed
+                if tool == "canvas_canvas":
+                    filtered_tools.append(tool)
+                    continue
+
+                # Check if tool belongs to an authorized server
+                if isinstance(tool, str) and "_" in tool:
+                    # Match against authorized servers by checking if tool name starts with server_
+                    # This handles server names that contain underscores (e.g., "pptx_generator")
+                    matched_server = None
+                    for auth_server in authorized_servers:
+                        if tool.startswith(f"{auth_server}_"):
+                            matched_server = auth_server
+                            break
+
+                    if matched_server:
+                        filtered_tools.append(tool)
+
+            return filtered_tools
+
+        except Exception:
+            logger.debug(
+                "Tool ACL filtering failed; proceeding with original selection",
+                exc_info=True
+            )
+            return selected_tools
+
+    async def _get_authorized_servers(self, user: str) -> List[str]:
+        """
+        Get list of MCP servers the user is authorized to access.
+
+        Args:
+            user: User email
+
+        Returns:
+            List of authorized server names
+        """
+        # Use tool_manager's authorization method if available
+        if hasattr(self.tool_manager, "get_authorized_servers"):
+            return await self.tool_manager.get_authorized_servers(user, is_user_in_group)  # type: ignore[attr-defined]
+
+        # If no authorization method available, return empty list (no authorized servers)
+        logger.warning("Tool manager has no get_authorized_servers method for user %s", user)
+        return []

atlas/application/chat/preprocessors/__init__.py

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

atlas/application/chat/preprocessors/message_builder.py

@@ -0,0 +1,92 @@
+"""Message builder - constructs messages with history and files manifest."""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from atlas.domain.sessions.models import Session
+from atlas.modules.prompts.prompt_provider import PromptProvider
+
+from ..utilities import file_processor
+
+logger = logging.getLogger(__name__)
+
+
+def build_session_context(session: Session) -> Dict[str, Any]:
+    """
+    Build session context dictionary from session.
+
+    Args:
+        session: Chat session
+
+    Returns:
+        Session context dictionary
+    """
+    return {
+        "session_id": session.id,
+        "user_email": session.user_email,
+        "files": session.context.get("files", {}),
+        **session.context
+    }
+
+
+class MessageBuilder:
+    """
+    Service that builds complete message arrays for LLM calls.
+
+    Combines conversation history with files manifest and system prompt.
+    """
+
+    def __init__(self, prompt_provider: Optional[PromptProvider] = None):
+        """
+        Initialize message builder.
+
+        Args:
+            prompt_provider: Optional prompt provider for loading system prompt
+        """
+        self.prompt_provider = prompt_provider
+
+    async def build_messages(
+        self,
+        session: Session,
+        include_files_manifest: bool = True,
+        include_system_prompt: bool = True,
+    ) -> List[Dict[str, Any]]:
+        """
+        Build messages array from session history and context.
+
+        Args:
+            session: Current chat session
+            include_files_manifest: Whether to append files manifest
+            include_system_prompt: Whether to prepend system prompt
+
+        Returns:
+            List of messages ready for LLM call
+        """
+        messages = []
+
+        # Optionally add system prompt at the beginning
+        if include_system_prompt and self.prompt_provider:
+            system_prompt = self.prompt_provider.get_system_prompt(
+                user_email=session.user_email
+            )
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+                logger.debug(f"Added system prompt (len={len(system_prompt)})")
+
+        # Get conversation history from session
+        history_messages = session.history.get_messages_for_llm()
+        messages.extend(history_messages)
+
+        # Optionally add files manifest
+        if include_files_manifest:
+            session_context = build_session_context(session)
+            files_in_context = session_context.get("files", {})
+            logger.debug(f"Session has {len(files_in_context)} files: {list(files_in_context.keys())}")
+            files_manifest = file_processor.build_files_manifest(session_context)
+            if files_manifest:
+                logger.debug(f"Adding files manifest to messages: {files_manifest['content'][:100]}")
+                messages.append(files_manifest)
+            else:
+                logger.warning("No files manifest generated despite include_files_manifest=True")
+
+        return messages

atlas/application/chat/preprocessors/prompt_override_service.py

@@ -0,0 +1,104 @@
+"""Prompt override service - handles MCP system prompt injection."""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PromptOverrideService:
+    """
+    Service that handles MCP prompt override injection.
+
+    Retrieves MCP-provided prompts and injects them as system messages,
+    applying only the first valid prompt found.
+    """
+
+    def __init__(self, tool_manager: Optional[Any] = None):
+        """
+        Initialize the prompt override service.
+
+        Args:
+            tool_manager: Optional tool manager with prompt retrieval capability
+        """
+        self.tool_manager = tool_manager
+
+    async def apply_prompt_override(
+        self,
+        messages: List[Dict[str, Any]],
+        selected_prompts: Optional[List[str]] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Apply MCP prompt override if selected prompts are provided.
+
+        Only the first valid prompt is applied, prepended as a system message.
+
+        Args:
+            messages: Current message history
+            selected_prompts: List of prompt keys (format: "server_promptname")
+
+        Returns:
+            Messages with prompt override prepended (if applicable)
+        """
+        if not selected_prompts or not self.tool_manager:
+            return messages
+
+        try:
+            # Iterate in order; when found, fetch prompt content and inject
+            for key in selected_prompts:
+                if not isinstance(key, str) or "_" not in key:
+                    continue
+
+                server, prompt_name = key.split("_", 1)
+
+                # Retrieve prompt from MCP
+                try:
+                    prompt_obj = await self.tool_manager.get_prompt(server, prompt_name)
+                    prompt_text = self._extract_prompt_text(prompt_obj)
+
+                    if prompt_text:
+                        # Prepend as system message override
+                        messages = [{"role": "system", "content": prompt_text}] + messages
+                        logger.info(
+                            "Applied MCP system prompt override (len=%d)",
+                            len(prompt_text),
+                        )
+                        break  # apply only one
+
+                except Exception:
+                    logger.debug("Failed retrieving MCP prompt %s", key, exc_info=True)
+
+        except Exception:
+            logger.debug(
+                "Prompt override injection skipped due to non-fatal error",
+                exc_info=True
+            )
+
+        return messages
+
+    def _extract_prompt_text(self, prompt_obj: Any) -> Optional[str]:
+        """
+        Extract text content from various MCP prompt object formats.
+
+        Args:
+            prompt_obj: Prompt object from MCP (could be string or structured object)
+
+        Returns:
+            Extracted prompt text, or None if extraction failed
+        """
+        # Simple string case
+        if isinstance(prompt_obj, str):
+            return prompt_obj
+
+        # FastMCP PromptMessage-like: may have 'content' list with text entries
+        if hasattr(prompt_obj, "content"):
+            content_field = getattr(prompt_obj, "content")
+
+            # content could be list of objects with 'text'
+            if isinstance(content_field, list) and content_field:
+                first = content_field[0]
+                if hasattr(first, "text") and isinstance(first.text, str):
+                    return first.text
+
+        # Fallback: string dump
+        return str(prompt_obj)