fast-agent-mcp 0.1.12__py3-none-any.whl → 0.2.0__py3-none-any.whl

mcp_agent/workflows/llm/augmented_llm.py

@@ -1,753 +0,0 @@
-from abc import abstractmethod
-
-from typing import (
-    Generic,
-    List,
-    Optional,
-    Protocol,
-    Type,
-    TypeVar,
-    TYPE_CHECKING,
-)
-
-from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
-from mcp_agent.workflows.llm.sampling_format_converter import (
-    SamplingFormatConverter,
-    MessageParamT,
-    MessageT,
-)
-
-# Forward reference for type annotations
-if TYPE_CHECKING:
-    from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
-    from mcp_agent.agents.agent import Agent
-    from mcp_agent.context import Context
-
-
-
-from mcp.types import (
-    CallToolRequest,
-    CallToolResult,
-    PromptMessage,
-    TextContent,
-    GetPromptResult,
-)
-
-from mcp_agent.context_dependent import ContextDependent
-from mcp_agent.core.exceptions import ModelConfigError, PromptExitError
-from mcp_agent.core.request_params import RequestParams
-from mcp_agent.event_progress import ProgressAction
-
-try:
-    from mcp_agent.mcp.mcp_aggregator import MCPAggregator
-except ImportError:
-    # For testing purposes
-    class MCPAggregator:
-        pass
-
-
-from mcp_agent.ui.console_display import ConsoleDisplay
-from rich.text import Text
-
-
-ModelT = TypeVar("ModelT")
-"""A type representing a structured output message from an LLM."""
-
-
-# TODO -- move this to a constant
-HUMAN_INPUT_TOOL_NAME = "__human_input__"
-
-
-class Memory(Protocol, Generic[MessageParamT]):
-    """
-    Simple memory management for storing past interactions in-memory.
-    """
-
-    # TODO: saqadri - add checkpointing and other advanced memory capabilities
-
-    def __init__(self): ...
-
-    def extend(
-        self, messages: List[MessageParamT], is_prompt: bool = False
-    ) -> None: ...
-
-    def set(self, messages: List[MessageParamT], is_prompt: bool = False) -> None: ...
-
-    def append(self, message: MessageParamT, is_prompt: bool = False) -> None: ...
-
-    def get(self, include_history: bool = True) -> List[MessageParamT]: ...
-
-    def clear(self, clear_prompts: bool = False) -> None: ...
-
-
-class SimpleMemory(Memory, Generic[MessageParamT]):
-    """
-    Simple memory management for storing past interactions in-memory.
-
-    Maintains both prompt messages (which are always included) and
-    generated conversation history (which is included based on use_history setting).
-    """
-
-    def __init__(self):
-        self.history: List[MessageParamT] = []
-        self.prompt_messages: List[MessageParamT] = []  # Always included
-
-    def extend(self, messages: List[MessageParamT], is_prompt: bool = False):
-        """
-        Add multiple messages to history.
-
-        Args:
-            messages: Messages to add
-            is_prompt: If True, add to prompt_messages instead of regular history
-        """
-        if is_prompt:
-            self.prompt_messages.extend(messages)
-        else:
-            self.history.extend(messages)
-
-    def set(self, messages: List[MessageParamT], is_prompt: bool = False):
-        """
-        Replace messages in history.
-
-        Args:
-            messages: Messages to set
-            is_prompt: If True, replace prompt_messages instead of regular history
-        """
-        if is_prompt:
-            self.prompt_messages = messages.copy()
-        else:
-            self.history = messages.copy()
-
-    def append(self, message: MessageParamT, is_prompt: bool = False):
-        """
-        Add a single message to history.
-
-        Args:
-            message: Message to add
-            is_prompt: If True, add to prompt_messages instead of regular history
-        """
-        if is_prompt:
-            self.prompt_messages.append(message)
-        else:
-            self.history.append(message)
-
-    def get(self, include_history: bool = True) -> List[MessageParamT]:
-        """
-        Get all messages in memory.
-
-        Args:
-            include_history: If True, include regular history messages
-                             If False, only return prompt messages
-
-        Returns:
-            Combined list of prompt messages and optionally history messages
-        """
-        if include_history:
-            return self.prompt_messages + self.history
-        else:
-            return self.prompt_messages.copy()
-
-    def clear(self, clear_prompts: bool = False):
-        """
-        Clear history and optionally prompt messages.
-
-        Args:
-            clear_prompts: If True, also clear prompt messages
-        """
-        self.history = []
-        if clear_prompts:
-            self.prompt_messages = []
-
-
-class AugmentedLLMProtocol(Protocol, Generic[MessageParamT, MessageT]):
-    """Protocol defining the interface for augmented LLMs"""
-
-    async def generate(
-        self,
-        message: str | MessageParamT | List[MessageParamT],
-        request_params: RequestParams | None = None,
-    ) -> List[MessageT]:
-        """Request an LLM generation, which may run multiple iterations, and return the result"""
-
-    async def generate_str(
-        self,
-        message: str | MessageParamT | List[MessageParamT],
-        request_params: RequestParams | None = None,
-    ) -> str:
-        """Request an LLM generation and return the string representation of the result"""
-
-    async def generate_structured(
-        self,
-        message: str | MessageParamT | List[MessageParamT],
-        response_model: Type[ModelT],
-        request_params: RequestParams | None = None,
-    ) -> ModelT:
-        """Request a structured LLM generation and return the result as a Pydantic model."""
-
-    async def generate_prompt(
-        self, prompt: PromptMessageMultipart, request_params: RequestParams | None
-    ) -> str:
-        """Request an LLM generation and return a string representation of the result"""
-
-
-class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, MessageT]):
-    """
-    The basic building block of agentic systems is an LLM enhanced with augmentations
-    such as retrieval, tools, and memory provided from a collection of MCP servers.
-    Our current models can actively use these capabilities—generating their own search queries,
-    selecting appropriate tools, and determining what information to retain.
-    """
-
-    # TODO: saqadri - add streaming support (e.g. generate_stream)
-    # TODO: saqadri - consider adding middleware patterns for pre/post processing of messages, for now we have pre/post_tool_call
-
-    provider: str | None = None
-
-    def __init__(
-        self,
-        agent: Optional["Agent"] = None,
-        server_names: List[str] | None = None,
-        instruction: str | None = None,
-        name: str | None = None,
-        request_params: RequestParams | None = None,
-        type_converter: Type[SamplingFormatConverter[MessageParamT, MessageT]] = None,
-        context: Optional["Context"] = None,
-        **kwargs,
-    ):
-        """
-        Initialize the LLM with a list of server names and an instruction.
-        If a name is provided, it will be used to identify the LLM.
-        If an agent is provided, all other properties are optional
-        """
-        # Extract request_params before super() call
-        self._init_request_params = request_params
-        super().__init__(context=context, **kwargs)
-
-        self.executor = self.context.executor
-        self.aggregator = (
-            agent if agent is not None else MCPAggregator(server_names or [])
-        )
-        self.name = name or (agent.name if agent else None)
-        self.instruction = instruction or (
-            agent.instruction if agent and isinstance(agent.instruction, str) else None
-        )
-        self.history: Memory[MessageParamT] = SimpleMemory[MessageParamT]()
-
-        # Initialize the display component
-        self.display = ConsoleDisplay(config=self.context.config)
-
-        # Initialize default parameters
-        self.default_request_params = self._initialize_default_params(kwargs)
-
-        # Merge with provided params if any
-        if self._init_request_params:
-            self.default_request_params = self._merge_request_params(
-                self.default_request_params, self._init_request_params
-            )
-
-        self.type_converter = type_converter
-        self.verb = kwargs.get("verb")
-
-    @abstractmethod
-    async def generate(
-        self,
-        message: str | MessageParamT | List[MessageParamT],
-        request_params: RequestParams | None = None,
-    ) -> List[MessageT]:
-        """Request an LLM generation, which may run multiple iterations, and return the result"""
-
-    @abstractmethod
-    async def generate_str(
-        self,
-        message: str | MessageParamT | List[MessageParamT],
-        request_params: RequestParams | None = None,
-    ) -> str:
-        """Request an LLM generation and return the string representation of the result"""
-
-    @abstractmethod
-    async def generate_structured(
-        self,
-        message: str | MessageParamT | List[MessageParamT],
-        response_model: Type[ModelT],
-        request_params: RequestParams | None = None,
-    ) -> ModelT:
-        """Request a structured LLM generation and return the result as a Pydantic model."""
-
-    async def select_model(
-        self, request_params: RequestParams | None = None
-    ) -> str | None:
-        """
-        Return the configured model (legacy support)
-        """
-        if request_params.model:
-            return request_params.model
-
-        raise ModelConfigError("Internal Error: Model is not configured correctly")
-
-    def _initialize_default_params(self, kwargs: dict) -> RequestParams:
-        """Initialize default parameters for the LLM.
-        Should be overridden by provider implementations to set provider-specific defaults."""
-        return RequestParams(
-            systemPrompt=self.instruction,
-            parallel_tool_calls=True,
-            max_iterations=10,
-            use_history=True,
-        )
-
-    def _merge_request_params(
-        self, default_params: RequestParams, provided_params: RequestParams
-    ) -> RequestParams:
-        """Merge default and provided request parameters"""
-
-        merged = default_params.model_dump()
-        merged.update(provided_params.model_dump(exclude_unset=True))
-        final_params = RequestParams(**merged)
-
-        return final_params
-
-    def get_request_params(
-        self,
-        request_params: RequestParams | None = None,
-        default: RequestParams | None = None,
-    ) -> RequestParams:
-        """
-        Get request parameters with merged-in defaults and overrides.
-        Args:
-            request_params: The request parameters to use as overrides.
-            default: The default request parameters to use as the base.
-                If unspecified, self.default_request_params will be used.
-        """
-        # Start with the defaults
-        default_request_params = default or self.default_request_params
-
-        if not default_request_params:
-            default_request_params = self._initialize_default_params({})
-
-        # If user provides overrides, merge them with defaults
-        if request_params:
-            return self._merge_request_params(default_request_params, request_params)
-
-        return default_request_params
-
-    @classmethod
-    def convert_message_to_message_param(
-        cls, message: MessageT, **kwargs
-    ) -> MessageParamT:
-        """Convert a response object to an input parameter object to allow LLM calls to be chained."""
-        # Many LLM implementations will allow the same type for input and output messages
-        return message
-
-    async def get_last_message(self) -> MessageParamT | None:
-        """
-        Return the last message generated by the LLM or None if history is empty.
-        This is useful for prompt chaining workflows where the last message from one LLM is used as input to another.
-        """
-        history = self.history.get()
-        return history[-1] if history else None
-
-    async def get_last_message_str(self) -> str | None:
-        """Return the string representation of the last message generated by the LLM or None if history is empty."""
-        last_message = await self.get_last_message()
-        return self.message_param_str(last_message) if last_message else None
-
-    def show_tool_result(self, result: CallToolResult):
-        """Display a tool result in a formatted panel."""
-        self.display.show_tool_result(result)
-
-    def show_oai_tool_result(self, result):
-        """Display a tool result in a formatted panel."""
-        self.display.show_oai_tool_result(result)
-
-    def show_tool_call(self, available_tools, tool_name, tool_args):
-        """Display a tool call in a formatted panel."""
-        self.display.show_tool_call(available_tools, tool_name, tool_args)
-
-    async def show_assistant_message(
-        self,
-        message_text: str | Text,
-        highlight_namespaced_tool: str = "",
-        title: str = "ASSISTANT",
-    ):
-        """Display an assistant message in a formatted panel."""
-        await self.display.show_assistant_message(
-            message_text,
-            aggregator=self.aggregator,
-            highlight_namespaced_tool=highlight_namespaced_tool,
-            title=title,
-            name=self.name,
-        )
-
-    def show_user_message(self, message, model: str | None, chat_turn: int):
-        """Display a user message in a formatted panel."""
-        self.display.show_user_message(message, model, chat_turn, name=self.name)
-
-    async def pre_tool_call(
-        self, tool_call_id: str | None, request: CallToolRequest
-    ) -> CallToolRequest | bool:
-        """Called before a tool is executed. Return False to prevent execution."""
-        return request
-
-    async def post_tool_call(
-        self, tool_call_id: str | None, request: CallToolRequest, result: CallToolResult
-    ) -> CallToolResult:
-        """Called after a tool execution. Can modify the result before it's returned."""
-        return result
-
-    async def call_tool(
-        self,
-        request: CallToolRequest,
-        tool_call_id: str | None = None,
-    ) -> CallToolResult:
-        """Call a tool with the given parameters and optional ID"""
-
-        try:
-            preprocess = await self.pre_tool_call(
-                tool_call_id=tool_call_id,
-                request=request,
-            )
-
-            if isinstance(preprocess, bool):
-                if not preprocess:
-                    return CallToolResult(
-                        isError=True,
-                        content=[
-                            TextContent(
-                                text=f"Error: Tool '{request.params.name}' was not allowed to run."
-                            )
-                        ],
-                    )
-            else:
-                request = preprocess
-
-            tool_name = request.params.name
-            tool_args = request.params.arguments
-            result = await self.aggregator.call_tool(tool_name, tool_args)
-
-            postprocess = await self.post_tool_call(
-                tool_call_id=tool_call_id, request=request, result=result
-            )
-
-            if isinstance(postprocess, CallToolResult):
-                result = postprocess
-
-            return result
-        except PromptExitError:
-            raise
-        except Exception as e:
-            return CallToolResult(
-                isError=True,
-                content=[
-                    TextContent(
-                        type="text",
-                        text=f"Error executing tool '{request.params.name}': {str(e)}",
-                    )
-                ],
-            )
-
-    def message_param_str(self, message: MessageParamT) -> str:
-        """
-        Convert an input message to a string representation.
-        Tries to extract just the content when possible.
-        """
-        if isinstance(message, dict):
-            # For dictionary format messages
-            if "content" in message:
-                content = message["content"]
-                # Handle both string and structured content formats
-                if isinstance(content, str):
-                    return content
-                elif isinstance(content, list) and content:
-                    # Try to extract text from content parts
-                    text_parts = []
-                    for part in content:
-                        if isinstance(part, dict) and "text" in part:
-                            text_parts.append(part["text"])
-                        elif hasattr(part, "text"):
-                            text_parts.append(part.text)
-                    if text_parts:
-                        return "\n".join(text_parts)
-
-        # For objects with content attribute
-        if hasattr(message, "content"):
-            content = message.content
-            if isinstance(content, str):
-                return content
-            elif hasattr(content, "text"):
-                return content.text
-
-        # Default fallback
-        return str(message)
-
-    def message_str(self, message: MessageT) -> str:
-        """
-        Convert an output message to a string representation.
-        Tries to extract just the content when possible.
-        """
-        # First try to use the same method for consistency
-        result = self.message_param_str(message)
-        if result != str(message):
-            return result
-
-        # Additional handling for output-specific formats
-        if hasattr(message, "content"):
-            content = message.content
-            if isinstance(content, list):
-                # Extract text from content blocks
-                text_parts = []
-                for block in content:
-                    if hasattr(block, "text") and block.text:
-                        text_parts.append(block.text)
-                if text_parts:
-                    return "\n".join(text_parts)
-
-        # Default fallback
-        return str(message)
-
-    def _log_chat_progress(
-        self, chat_turn: Optional[int] = None, model: Optional[str] = None
-    ):
-        """Log a chat progress event"""
-        # Determine action type based on verb
-        if hasattr(self, "verb") and self.verb:
-            # Use verb directly regardless of type
-            act = self.verb
-        else:
-            act = ProgressAction.CHATTING
-
-        data = {
-            "progress_action": act,
-            "model": model,
-            "agent_name": self.name,
-            "chat_turn": chat_turn if chat_turn is not None else None,
-        }
-        self.logger.debug("Chat in progress", data=data)
-
-    def _log_chat_finished(self, model: Optional[str] = None):
-        """Log a chat finished event"""
-        data = {
-            "progress_action": ProgressAction.READY,
-            "model": model,
-            "agent_name": self.name,
-        }
-        self.logger.debug("Chat finished", data=data)
-
-    def _convert_prompt_messages(
-        self, prompt_messages: List[PromptMessage]
-    ) -> List[MessageParamT]:
-        """
-        Convert prompt messages to this LLM's specific message format.
-        To be implemented by concrete LLM classes.
-        """
-        raise NotImplementedError("Must be implemented by subclass")
-
-    async def show_prompt_loaded(
-        self,
-        prompt_name: str,
-        description: Optional[str] = None,
-        message_count: int = 0,
-        arguments: Optional[dict[str, str]] = None,
-    ):
-        """
-        Display information about a loaded prompt template.
-
-        Args:
-            prompt_name: The name of the prompt
-            description: Optional description of the prompt
-            message_count: Number of messages in the prompt
-            arguments: Optional dictionary of arguments passed to the prompt
-        """
-        await self.display.show_prompt_loaded(
-            prompt_name=prompt_name,
-            description=description,
-            message_count=message_count,
-            agent_name=self.name,
-            aggregator=self.aggregator,
-            arguments=arguments,
-        )
-
-    async def apply_prompt_template(
-        self, prompt_result: GetPromptResult, prompt_name: str
-    ) -> str:
-        """
-        Apply a prompt template by adding it to the conversation history.
-        If the last message in the prompt is from a user, automatically
-        generate an assistant response.
-
-        Args:
-            prompt_result: The GetPromptResult containing prompt messages
-            prompt_name: The name of the prompt being applied
-
-        Returns:
-            String representation of the assistant's response if generated,
-            or the last assistant message in the prompt
-        """
-        from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
-
-        # Check if we have any messages
-        if not prompt_result.messages:
-            return "Prompt contains no messages"
-
-        # Extract arguments if they were stored in the result
-        arguments = getattr(prompt_result, "arguments", None)
-
-        # Display information about the loaded prompt
-        await self.show_prompt_loaded(
-            prompt_name=prompt_name,
-            description=prompt_result.description,
-            message_count=len(prompt_result.messages),
-            arguments=arguments,
-        )
-
-        # Convert to PromptMessageMultipart objects
-        multipart_messages = PromptMessageMultipart.parse_get_prompt_result(
-            prompt_result
-        )
-
-        # Delegate to the provider-specific implementation
-        return await self._apply_prompt_template_provider_specific(
-            multipart_messages, None
-        )
-
-    async def apply_prompt(
-        self,
-        multipart_messages: List["PromptMessageMultipart"],
-        request_params: RequestParams | None = None,
-    ) -> str:
-        """
-        Apply a list of PromptMessageMultipart messages directly to the LLM.
-        This is a cleaner interface to _apply_prompt_template_provider_specific.
-
-        Args:
-            multipart_messages: List of PromptMessageMultipart objects
-            request_params: Optional parameters to configure the LLM request
-
-        Returns:
-            String representation of the assistant's response
-        """
-        # Delegate to the provider-specific implementation
-        return await self._apply_prompt_template_provider_specific(
-            multipart_messages, request_params
-        )
-
-    async def _apply_prompt_template_provider_specific(
-        self,
-        multipart_messages: List["PromptMessageMultipart"],
-        request_params: RequestParams | None = None,
-    ) -> str:
-        """
-        Provider-specific implementation of apply_prompt_template.
-        This default implementation handles basic text content for any LLM type.
-        Provider-specific subclasses should override this method to handle
-        multimodal content appropriately.
-
-        Args:
-            multipart_messages: List of PromptMessageMultipart objects parsed from the prompt template
-
-        Returns:
-            String representation of the assistant's response if generated,
-            or the last assistant message in the prompt
-        """
-        # Check the last message role
-        last_message = multipart_messages[-1]
-
-        if last_message.role == "user":
-            # For user messages: Add all previous messages to history, then generate response to the last one
-            self.logger.debug(
-                "Last message in prompt is from user, generating assistant response"
-            )
-
-            # Add all but the last message to history
-            if len(multipart_messages) > 1:
-                previous_messages = multipart_messages[:-1]
-                converted = []
-
-                # Fallback generic method for all LLM types
-                for msg in previous_messages:
-                    # Convert each PromptMessageMultipart to individual PromptMessages
-                    prompt_messages = msg.to_prompt_messages()
-                    for prompt_msg in prompt_messages:
-                        converted.append(
-                            self.type_converter.from_prompt_message(prompt_msg)
-                        )
-
-                self.history.extend(converted, is_prompt=True)
-
-            # For generic LLMs, extract text and describe non-text content
-            user_text_parts = []
-            for content in last_message.content:
-                if content.type == "text":
-                    user_text_parts.append(content.text)
-                elif content.type == "resource" and hasattr(content.resource, "text"):
-                    user_text_parts.append(content.resource.text)
-                elif content.type == "image":
-                    # Add a placeholder for images
-                    mime_type = getattr(content, "mimeType", "image/unknown")
-                    user_text_parts.append(f"[Image: {mime_type}]")
-
-            user_text = "\n".join(user_text_parts) if user_text_parts else ""
-            if not user_text:
-                # Fallback to original method if we couldn't extract text
-                user_text = str(last_message.content)
-
-            return await self.generate_str(user_text)
-        else:
-            # For assistant messages: Add all messages to history and return the last one
-            self.logger.debug(
-                "Last message in prompt is from assistant, returning it directly"
-            )
-
-            # Convert and add all messages to history
-            converted = []
-
-            # Fallback to the original method for all LLM types
-            for msg in multipart_messages:
-                # Convert each PromptMessageMultipart to individual PromptMessages
-                prompt_messages = msg.to_prompt_messages()
-                for prompt_msg in prompt_messages:
-                    converted.append(
-                        self.type_converter.from_prompt_message(prompt_msg)
-                    )
-
-            self.history.extend(converted, is_prompt=True)
-
-            # Return the assistant's message with proper handling of different content types
-            assistant_text_parts = []
-            has_non_text_content = False
-
-            for content in last_message.content:
-                if content.type == "text":
-                    assistant_text_parts.append(content.text)
-                elif content.type == "resource" and hasattr(content.resource, "text"):
-                    # Add resource text with metadata
-                    mime_type = getattr(content.resource, "mimeType", "text/plain")
-                    uri = getattr(content.resource, "uri", "")
-                    if uri:
-                        assistant_text_parts.append(
-                            f"[Resource: {uri}, Type: {mime_type}]\n{content.resource.text}"
-                        )
-                    else:
-                        assistant_text_parts.append(
-                            f"[Resource Type: {mime_type}]\n{content.resource.text}"
-                        )
-                elif content.type == "image":
-                    # Note the presence of images
-                    mime_type = getattr(content, "mimeType", "image/unknown")
-                    assistant_text_parts.append(f"[Image: {mime_type}]")
-                    has_non_text_content = True
-                else:
-                    # Other content types
-                    assistant_text_parts.append(f"[Content of type: {content.type}]")
-                    has_non_text_content = True
-
-            # Join all parts with double newlines for better readability
-            result = (
-                "\n\n".join(assistant_text_parts)
-                if assistant_text_parts
-                else str(last_message.content)
-            )
-
-            # Add a note if non-text content was present
-            if has_non_text_content:
-                result += "\n\n[Note: This message contained non-text content that may not be fully represented in text format]"
-
-            return result