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

mcp_agent/workflows/llm/augmented_llm.py

@@ -1,195 +1,54 @@
 from abc import abstractmethod
-
 from typing import (
-    Generic,
+    TYPE_CHECKING,
+    Any,
     List,
     Optional,
-    Protocol,
     Type,
-    TypeVar,
-    TYPE_CHECKING,
+    cast,
 )
 
-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 mcp.types import (
     CallToolRequest,
     CallToolResult,
+    GetPromptResult,
     PromptMessage,
     TextContent,
-    GetPromptResult,
 )
+from rich.text import Text
 
 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.mcp.mcp_aggregator import MCPAggregator
 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 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
@@ -198,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__(
@@ -210,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.
@@ -222,15 +78,11 @@ 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
@@ -241,9 +93,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
 
         # 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.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")
@@ -273,13 +123,11 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
     ) -> 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:
+    async def select_model(self, request_params: RequestParams | None = None) -> str | None:
         """
         Return the configured model (legacy support)
         """
-        if request_params.model:
+        if request_params and request_params.model:
             return request_params.model
 
         raise ModelConfigError("Internal Error: Model is not configured correctly")
@@ -294,9 +142,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             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()
@@ -330,12 +176,10 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         return default_request_params
 
     @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:
         """
@@ -350,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)
 
@@ -367,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,
@@ -377,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
 
@@ -412,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.",
                             )
                         ],
                     )
@@ -423,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
@@ -463,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"):
@@ -484,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 = []
@@ -503,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:
@@ -522,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,
@@ -531,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.
@@ -546,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.
 
@@ -565,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
@@ -599,14 +432,10 @@ 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
-        )
+        return await self._apply_prompt_template_provider_specific(multipart_messages, None)
 
     async def apply_prompt(
         self,
@@ -625,10 +454,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             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
-        )
+        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"],
@@ -652,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:
@@ -664,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)
 
@@ -677,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")
@@ -692,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 = []
@@ -702,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)
 
@@ -723,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
@@ -740,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   ###
+#####################################