fast-agent-mcp 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

mcp_agent/workflows/llm/augmented_llm.py

@@ -1,228 +1,54 @@
 from abc import abstractmethod
-
 from typing import (
-    Generic,
+    TYPE_CHECKING,
+    Any,
     List,
     Optional,
-    Protocol,
     Type,
-    TypeVar,
-    TYPE_CHECKING,
+    cast,
 )
 
-from mcp import CreateMessageResult, SamplingMessage
-from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
-from mcp_agent.workflows.llm.sampling_format_converter import (
-    SamplingFormatConverter,
+from mcp_agent.logging.logger import get_logger
+from mcp_agent.mcp.interfaces import (
+    AugmentedLLMProtocol,
     MessageParamT,
     MessageT,
+    ModelT,
+)
+from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
+from mcp_agent.workflows.llm.sampling_format_converter import (
+    BasicFormatConverter,
+    ProviderFormatConverter,
 )
 
 # 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_agent.mcp.prompt_message_multipart import PromptMessageMultipart
 
 
-from pydantic import Field
-
 from mcp.types import (
     CallToolRequest,
     CallToolResult,
-    CreateMessageRequestParams,
-    ModelPreferences,
+    GetPromptResult,
     PromptMessage,
     TextContent,
-    GetPromptResult,
 )
+from rich.text import Text
 
 from mcp_agent.context_dependent import ContextDependent
-from mcp_agent.core.exceptions import PromptExitError
+from mcp_agent.core.exceptions import ModelConfigError, PromptExitError
+from mcp_agent.core.request_params import RequestParams
 from mcp_agent.event_progress import ProgressAction
 from mcp_agent.mcp.mcp_aggregator import MCPAggregator
-from mcp_agent.workflows.llm.llm_selector import ModelSelector
 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."""
-
+from mcp_agent.workflows.llm.memory import Memory, SimpleMemory
 
 # 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 RequestParams(CreateMessageRequestParams):
-    """
-    Parameters to configure the AugmentedLLM 'generate' requests.
-    """
-
-    messages: None = Field(exclude=True, default=None)
-    """
-    Ignored. 'messages' are removed from CreateMessageRequestParams 
-    to avoid confusion with the 'message' parameter on 'generate' method.
-    """
-
-    maxTokens: int = 2048
-    """The maximum number of tokens to sample, as requested by the server."""
-
-    model: str | None = None
-    """
-    The model to use for the LLM generation.
-    If specified, this overrides the 'modelPreferences' selection criteria.
-    """
-
-    use_history: bool = True
-    """
-    Include the message history in the generate request.
-    """
-
-    max_iterations: int = 10
-    """
-    The maximum number of iterations to run the LLM for.
-    """
-
-    parallel_tool_calls: bool = True
-    """
-    Whether to allow multiple tool calls per iteration.
-    Also known as multi-step tool use.
-    """
-
-
-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
@@ -231,9 +57,6 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
     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__(
@@ -243,10 +66,10 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         instruction: str | None = None,
         name: str | None = None,
         request_params: RequestParams | None = None,
-        type_converter: Type[SamplingFormatConverter[MessageParamT, MessageT]] = None,
+        type_converter: Type[ProviderFormatConverter[MessageParamT, MessageT]] = BasicFormatConverter,
         context: Optional["Context"] = None,
-        **kwargs,
-    ):
+        **kwargs: dict[str, Any],
+    ) -> None:
         """
         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.
@@ -255,44 +78,23 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         # Extract request_params before super() call
         self._init_request_params = request_params
         super().__init__(context=context, **kwargs)
-
+        self.logger = get_logger(__name__)
         self.executor = self.context.executor
-        self.aggregator = (
-            agent if agent is not None else MCPAggregator(server_names or [])
-        )
+        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.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)
 
-        # Set initial model preferences
-        self.model_preferences = ModelPreferences(
-            costPriority=0.3,
-            speedPriority=0.4,
-            intelligencePriority=0.3,
-        )
-
         # Initialize default parameters
         self.default_request_params = self._initialize_default_params(kwargs)
 
-        # Update model preferences from default params
-        if self.default_request_params and self.default_request_params.modelPreferences:
-            self.model_preferences = self.default_request_params.modelPreferences
-
         # 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
-            )
-            # Update model preferences again if they changed in the merge
-            if self.default_request_params.modelPreferences:
-                self.model_preferences = self.default_request_params.modelPreferences
+            self.default_request_params = self._merge_request_params(self.default_request_params, self._init_request_params)
 
-        self.model_selector = self.context.model_selector
         self.type_converter = type_converter
         self.verb = kwargs.get("verb")
 
@@ -321,48 +123,26 @@ 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:
+    async def select_model(self, request_params: RequestParams | None = None) -> str | None:
         """
-        Select an LLM based on the request parameters.
-        If a model is specified in the request, it will override the model selection criteria.
+        Return the configured model (legacy support)
         """
-        model_preferences = self.model_preferences
-        if request_params is not None:
-            model_preferences = request_params.modelPreferences or model_preferences
-            model = request_params.model
-            if model:
-                return model
-
-        ## TODO -- can't have been tested, returns invalid model strings (e.g. claude-35-sonnet)
-        if not self.model_selector:
-            self.model_selector = ModelSelector()
-
-        model_info = self.model_selector.select_best_model(
-            model_preferences=model_preferences, provider=self.provider
-        )
+        if request_params and request_params.model:
+            return request_params.model
 
-        return model_info.name
+        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(
-            modelPreferences=self.model_preferences,
             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:
+    def _merge_request_params(self, default_params: RequestParams, provided_params: RequestParams) -> RequestParams:
         """Merge default and provided request parameters"""
 
         merged = default_params.model_dump()
@@ -395,32 +175,11 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
 
         return default_request_params
 
-    def to_mcp_message_result(self, result: MessageT) -> CreateMessageResult:
-        """Convert an LLM response to an MCP message result type."""
-        return self.type_converter.to_sampling_result(result)
-
-    def from_mcp_message_result(self, result: CreateMessageResult) -> MessageT:
-        """Convert an MCP message result to an LLM response type."""
-        return self.type_converter.from_sampling_result(result)
-
-    def to_mcp_message_param(self, param: MessageParamT) -> SamplingMessage:
-        """Convert an LLM input to an MCP message (SamplingMessage) type."""
-        return self.type_converter.to_sampling_message(param)
-
-    def from_mcp_message_param(self, param: SamplingMessage) -> MessageParamT:
-        """Convert an MCP message (SamplingMessage) to an LLM input type."""
-        return self.type_converter.from_sampling_message(param)
-
-    def from_mcp_prompt_message(self, message: PromptMessage) -> MessageParamT:
-        return self.type_converter.from_prompt_message(message)
-
     @classmethod
-    def convert_message_to_message_param(
-        cls, message: MessageT, **kwargs
-    ) -> MessageParamT:
+    def convert_message_to_message_param(cls, message: MessageT, **kwargs: dict[str, Any]) -> 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
+        return cast("MessageParamT", message)
 
     async def get_last_message(self) -> MessageParamT | None:
         """
@@ -435,15 +194,15 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         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):
+    def show_tool_result(self, result: CallToolResult) -> None:
         """Display a tool result in a formatted panel."""
         self.display.show_tool_result(result)
 
-    def show_oai_tool_result(self, result):
+    def show_oai_tool_result(self, result: str) -> None:
         """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):
+    def show_tool_call(self, available_tools, tool_name, tool_args) -> None:
         """Display a tool call in a formatted panel."""
         self.display.show_tool_call(available_tools, tool_name, tool_args)
 
@@ -452,7 +211,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         message_text: str | Text,
         highlight_namespaced_tool: str = "",
         title: str = "ASSISTANT",
-    ):
+    ) -> None:
         """Display an assistant message in a formatted panel."""
         await self.display.show_assistant_message(
             message_text,
@@ -462,19 +221,15 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             name=self.name,
         )
 
-    def show_user_message(self, message, model: str | None, chat_turn: int):
+    def show_user_message(self, message, model: str | None, chat_turn: int) -> None:
         """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:
+    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:
+    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
 
@@ -497,7 +252,8 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                         isError=True,
                         content=[
                             TextContent(
-                                text=f"Error: Tool '{request.params.name}' was not allowed to run."
+                                type="text",
+                                text=f"Error: Tool '{request.params.name}' was not allowed to run.",
                             )
                         ],
                     )
@@ -508,9 +264,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             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
-            )
+            postprocess = await self.post_tool_call(tool_call_id=tool_call_id, request=request, result=result)
 
             if isinstance(postprocess, CallToolResult):
                 result = postprocess
@@ -548,13 +302,13 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                         if isinstance(part, dict) and "text" in part:
                             text_parts.append(part["text"])
                         elif hasattr(part, "text"):
-                            text_parts.append(part.text)
+                            text_parts.append(part.text)  # type: ignore
                     if text_parts:
                         return "\n".join(text_parts)
 
         # For objects with content attribute
         if hasattr(message, "content"):
-            content = message.content
+            content = message.content  # type: ignore
             if isinstance(content, str):
                 return content
             elif hasattr(content, "text"):
@@ -569,13 +323,13 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         Tries to extract just the content when possible.
         """
         # First try to use the same method for consistency
-        result = self.message_param_str(message)
+        result = self.message_param_str(message)  # type: ignore
         if result != str(message):
             return result
 
         # Additional handling for output-specific formats
         if hasattr(message, "content"):
-            content = message.content
+            content = getattr(message, "content")
             if isinstance(content, list):
                 # Extract text from content blocks
                 text_parts = []
@@ -588,9 +342,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         # Default fallback
         return str(message)
 
-    def _log_chat_progress(
-        self, chat_turn: Optional[int] = None, model: Optional[str] = None
-    ):
+    def _log_chat_progress(self, chat_turn: Optional[int] = None, model: Optional[str] = None) -> None:
         """Log a chat progress event"""
         # Determine action type based on verb
         if hasattr(self, "verb") and self.verb:
@@ -607,7 +359,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         }
         self.logger.debug("Chat in progress", data=data)
 
-    def _log_chat_finished(self, model: Optional[str] = None):
+    def _log_chat_finished(self, model: Optional[str] = None) -> None:
         """Log a chat finished event"""
         data = {
             "progress_action": ProgressAction.READY,
@@ -616,9 +368,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         }
         self.logger.debug("Chat finished", data=data)
 
-    def _convert_prompt_messages(
-        self, prompt_messages: List[PromptMessage]
-    ) -> List[MessageParamT]:
+    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.
@@ -631,7 +381,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         description: Optional[str] = None,
         message_count: int = 0,
         arguments: Optional[dict[str, str]] = None,
-    ):
+    ) -> None:
         """
         Display information about a loaded prompt template.
 
@@ -650,9 +400,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             arguments=arguments,
         )
 
-    async def apply_prompt_template(
-        self, prompt_result: GetPromptResult, prompt_name: str
-    ) -> str:
+    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
@@ -684,15 +432,35 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         )
 
         # Convert to PromptMessageMultipart objects
-        multipart_messages = PromptMessageMultipart.parse_get_prompt_result(
-            prompt_result
-        )
+        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)
+        return await self._apply_prompt_template_provider_specific(multipart_messages, request_params)
 
+    # this shouln't need to be very big...
     async def _apply_prompt_template_provider_specific(
-        self, multipart_messages: List["PromptMessageMultipart"]
+        self,
+        multipart_messages: List["PromptMessageMultipart"],
+        request_params: RequestParams | None = None,
     ) -> str:
         """
         Provider-specific implementation of apply_prompt_template.
@@ -712,9 +480,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
 
         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(multipart_messages) > 1:
@@ -724,11 +490,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                 # Fallback generic method for all LLM types
                 for msg in previous_messages:
                     # Convert each PromptMessageMultipart to individual PromptMessages
-                    prompt_messages = msg.to_prompt_messages()
+                    prompt_messages = msg.from_multipart()
                     for prompt_msg in prompt_messages:
-                        converted.append(
-                            self.type_converter.from_prompt_message(prompt_msg)
-                        )
+                        converted.append(self.type_converter.from_prompt_message(prompt_msg))
 
                 self.history.extend(converted, is_prompt=True)
 
@@ -737,8 +501,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             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 == "resource" and getattr(content, "resource", None) is not None:
+                    if hasattr(content.resource, "text"):
+                        user_text_parts.append(content.resource.text)  # type: ignore
                 elif content.type == "image":
                     # Add a placeholder for images
                     mime_type = getattr(content, "mimeType", "image/unknown")
@@ -752,9 +517,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             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 = []
@@ -762,11 +525,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             # 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()
+                prompt_messages = msg.from_multipart()
                 for prompt_msg in prompt_messages:
-                    converted.append(
-                        self.type_converter.from_prompt_message(prompt_msg)
-                    )
+                    converted.append(self.type_converter.from_prompt_message(prompt_msg))
 
             self.history.extend(converted, is_prompt=True)
 
@@ -783,11 +544,11 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                     uri = getattr(content.resource, "uri", "")
                     if uri:
                         assistant_text_parts.append(
-                            f"[Resource: {uri}, Type: {mime_type}]\n{content.resource.text}"
+                            f"[Resource: {uri}, Type: {mime_type}]\n{content.resource.text}"  # ignore # type: ignore
                         )
                     else:
                         assistant_text_parts.append(
-                            f"[Resource Type: {mime_type}]\n{content.resource.text}"
+                            f"[Resource Type: {mime_type}]\n{content.resource.text}"  # type ignore # type: ignore
                         )
                 elif content.type == "image":
                     # Note the presence of images
@@ -800,14 +561,15 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                     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)
-            )
+            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
+
+
+#####################################
+### NEW INTERFACE METHODS BELOW   ###
+#####################################