fast-agent-mcp 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

mcp_agent/workflows/llm/augmented_llm.py

@@ -1,7 +1,6 @@
 from abc import abstractmethod
 
 from typing import (
-    Any,
     Generic,
     List,
     Optional,
@@ -9,18 +8,30 @@ from typing import (
     Type,
     TypeVar,
     TYPE_CHECKING,
-    Union,
 )
 
+from mcp import CreateMessageResult, SamplingMessage
+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 pydantic import Field
 
 from mcp.types import (
     CallToolRequest,
     CallToolResult,
     CreateMessageRequestParams,
-    CreateMessageResult,
     ModelPreferences,
-    SamplingMessage,
     PromptMessage,
     TextContent,
     GetPromptResult,
@@ -34,22 +45,10 @@ from mcp_agent.workflows.llm.llm_selector import ModelSelector
 from mcp_agent.ui.console_display import ConsoleDisplay
 from rich.text import Text
 
-if TYPE_CHECKING:
-    from mcp_agent.agents.agent import Agent
-    from mcp_agent.context import Context
-
-MessageParamT = TypeVar("MessageParamT")
-"""A type representing an input message to an LLM."""
-
-MessageT = TypeVar("MessageT")
-"""A type representing an output message from an LLM."""
 
 ModelT = TypeVar("ModelT")
 """A type representing a structured output message from an LLM."""
 
-# TODO: saqadri - SamplingMessage is fairly limiting - consider extending
-MCPMessageParam = SamplingMessage
-MCPMessageResult = CreateMessageResult
 
 # TODO -- move this to a constant
 HUMAN_INPUT_TOOL_NAME = "__human_input__"
@@ -218,25 +217,10 @@ class AugmentedLLMProtocol(Protocol, Generic[MessageParamT, MessageT]):
     ) -> ModelT:
         """Request a structured LLM generation and return the result as a Pydantic model."""
 
-
-class ProviderToMCPConverter(Protocol, Generic[MessageParamT, MessageT]):
-    """Conversions between LLM provider and MCP types"""
-
-    @classmethod
-    def to_mcp_message_result(cls, result: MessageT) -> MCPMessageResult:
-        """Convert an LLM response to an MCP message result type."""
-
-    @classmethod
-    def from_mcp_message_result(cls, result: MCPMessageResult) -> MessageT:
-        """Convert an MCP message result to an LLM response type."""
-
-    @classmethod
-    def to_mcp_message_param(cls, param: MessageParamT) -> MCPMessageParam:
-        """Convert an LLM input to an MCP message (SamplingMessage) type."""
-
-    @classmethod
-    def from_mcp_message_param(cls, param: MCPMessageParam) -> MessageParamT:
-        """Convert an MCP message (SamplingMessage) to an LLM input type."""
+    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]):
@@ -259,7 +243,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         instruction: str | None = None,
         name: str | None = None,
         request_params: RequestParams | None = None,
-        type_converter: Type[ProviderToMCPConverter[MessageParamT, MessageT]] = None,
+        type_converter: Type[SamplingFormatConverter[MessageParamT, MessageT]] = None,
         context: Optional["Context"] = None,
         **kwargs,
     ):
@@ -337,6 +321,10 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
     ) -> ModelT:
         """Request a structured LLM generation and return the result as a Pydantic model."""
 
+    # aysnc def generate2_str(self, prompt: PromptMessageMultipart, request_params: RequestParams | None = None) -> List[MessageT]:
+    #     """Request an LLM generation, which may run multiple iterations, and return the result"""
+    #     return None
+
     async def select_model(
         self, request_params: RequestParams | None = None
     ) -> str | None:
@@ -381,10 +369,6 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         merged.update(provided_params.model_dump(exclude_unset=True))
         final_params = RequestParams(**merged)
 
-        # self.logger.debug(
-        #     "Final merged params:", extra={"params": final_params.model_dump()}
-        # )
-
         return final_params
 
     def get_request_params(
@@ -411,24 +395,24 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
 
         return default_request_params
 
-    def to_mcp_message_result(self, result: MessageT) -> MCPMessageResult:
+    def to_mcp_message_result(self, result: MessageT) -> CreateMessageResult:
         """Convert an LLM response to an MCP message result type."""
-        return self.type_converter.to_mcp_message_result(result)
+        return self.type_converter.to_sampling_result(result)
 
-    def from_mcp_message_result(self, result: MCPMessageResult) -> MessageT:
+    def from_mcp_message_result(self, result: CreateMessageResult) -> MessageT:
         """Convert an MCP message result to an LLM response type."""
-        return self.type_converter.from_mcp_message_result(result)
+        return self.type_converter.from_sampling_result(result)
 
-    def to_mcp_message_param(self, param: MessageParamT) -> MCPMessageParam:
+    def to_mcp_message_param(self, param: MessageParamT) -> SamplingMessage:
         """Convert an LLM input to an MCP message (SamplingMessage) type."""
-        return self.type_converter.to_mcp_message_param(param)
+        return self.type_converter.to_sampling_message(param)
 
-    def from_mcp_message_param(self, param: MCPMessageParam) -> MessageParamT:
+    def from_mcp_message_param(self, param: SamplingMessage) -> MessageParamT:
         """Convert an MCP message (SamplingMessage) to an LLM input type."""
-        return self.type_converter.from_mcp_message_param(param)
+        return self.type_converter.from_sampling_message(param)
 
     def from_mcp_prompt_message(self, message: PromptMessage) -> MessageParamT:
-        return self.type_converter.from_mcp_prompt_message(message)
+        return self.type_converter.from_prompt_message(message)
 
     @classmethod
     def convert_message_to_message_param(
@@ -567,7 +551,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                             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
@@ -575,7 +559,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                 return content
             elif hasattr(content, "text"):
                 return content.text
-                
+
         # Default fallback
         return str(message)
 
@@ -588,7 +572,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         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
@@ -600,7 +584,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                         text_parts.append(block.text)
                 if text_parts:
                     return "\n".join(text_parts)
-        
+
         # Default fallback
         return str(message)
 
@@ -650,7 +634,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
     ):
         """
         Display information about a loaded prompt template.
-        
+
         Args:
             prompt_name: The name of the prompt
             description: Optional description of the prompt
@@ -679,13 +663,13 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             prompt_name: The name of the prompt being applied
 
         Returns:
-            String representation of the assistant's response if generated, 
+            String representation of the assistant's response if generated,
             or the last assistant message in the prompt
         """
-        prompt_messages: List[PromptMessage] = prompt_result.messages
-        
+        from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
+
         # Check if we have any messages
-        if not prompt_messages:
+        if not prompt_result.messages:
             return "Prompt contains no messages"
 
         # Extract arguments if they were stored in the result
@@ -695,105 +679,135 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         await self.show_prompt_loaded(
             prompt_name=prompt_name,
             description=prompt_result.description,
-            message_count=len(prompt_messages),
+            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)
+
+    async def _apply_prompt_template_provider_specific(
+        self, multipart_messages: List["PromptMessageMultipart"]
+    ) -> 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 = prompt_messages[-1]
-        
+        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")
-            
+            self.logger.debug(
+                "Last message in prompt is from user, generating assistant response"
+            )
+
             # Add all but the last message to history
-            if len(prompt_messages) > 1:
-                previous_messages = prompt_messages[:-1]
+            if len(multipart_messages) > 1:
+                previous_messages = multipart_messages[:-1]
                 converted = []
+
+                # Fallback generic method for all LLM types
                 for msg in previous_messages:
-                    converted.append(self.type_converter.from_mcp_prompt_message(msg))
+                    # 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)
-            
-            # Extract the user's question and generate a response
-            user_content = last_message.content
-            user_text = user_content.text if hasattr(user_content, "text") else str(user_content)
-            
+
+            # 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")
-            
+            self.logger.debug(
+                "Last message in prompt is from assistant, returning it directly"
+            )
+
             # Convert and add all messages to history
             converted = []
-            for msg in prompt_messages:
-                converted.append(self.type_converter.from_mcp_prompt_message(msg))
-            self.history.extend(converted, is_prompt=True)
-            
-            # Return the assistant's message
-            content = last_message.content
-            return content.text if hasattr(content, "text") else str(content)
-        
-
 
-class PassthroughLLM(AugmentedLLM):
-    """
-    A specialized LLM implementation that simply passes through input messages without modification.
-
-    This is useful for cases where you need an object with the AugmentedLLM interface
-    but want to preserve the original message without any processing, such as in a
-    parallel workflow where no fan-in aggregation is needed.
-    """
-
-    def __init__(self, name: str = "Passthrough", context=None, **kwargs):
-        super().__init__(name=name, context=context, **kwargs)
+            # 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)
+                    )
 
-    async def generate(
-        self,
-        message: Union[str, MessageParamT, List[MessageParamT]],
-        request_params: Optional[RequestParams] = None,
-    ) -> Union[List[MessageT], Any]:
-        """Simply return the input message as is."""
-        # Return in the format expected by the caller
-        return [message] if isinstance(message, list) else message
+            self.history.extend(converted, is_prompt=True)
 
-    async def generate_str(
-        self,
-        message: Union[str, MessageParamT, List[MessageParamT]],
-        request_params: Optional[RequestParams] = None,
-    ) -> str:
-        """Return the input message as a string."""
-        self.show_user_message(message, model="fastagent-passthrough", chat_turn=0)
-        await self.show_assistant_message(message, title="ASSISTANT/PASSTHROUGH")
+            # 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)
+            )
 
-        return str(message)
+            # 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]"
 
-    async def generate_structured(
-        self,
-        message: Union[str, MessageParamT, List[MessageParamT]],
-        response_model: Type[ModelT],
-        request_params: Optional[RequestParams] = None,
-    ) -> ModelT:
-        """
-        Return the input message as the requested model type.
-        This is a best-effort implementation - it may fail if the
-        message cannot be converted to the requested model.
-        """
-        if isinstance(message, response_model):
-            return message
-        elif isinstance(message, dict):
-            return response_model(**message)
-        elif isinstance(message, str):
-            try:
-                # Try to parse as JSON if it's a string
-                import json
-
-                data = json.loads(message)
-                return response_model(**data)
-            except:  # noqa: E722
-                raise ValueError(
-                    f"Cannot convert message of type {type(message)} to {response_model}"
-                )
-        else:
-            raise ValueError(
-                f"Cannot convert message of type {type(message)} to {response_model}"
-            )
+            return result