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

mcp_agent/{workflows/llm → llm}/augmented_llm.py

@@ -2,32 +2,14 @@ from abc import abstractmethod
 from typing import (
     TYPE_CHECKING,
     Any,
+    Generic,
     List,
     Optional,
     Type,
+    TypeVar,
     cast,
 )
 
-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.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,
@@ -35,21 +17,45 @@ from mcp.types import (
     PromptMessage,
     TextContent,
 )
+from pydantic_core import from_json
 from rich.text import Text
 
 from mcp_agent.context_dependent import ContextDependent
-from mcp_agent.core.exceptions import ModelConfigError, PromptExitError
+from mcp_agent.core.exceptions import PromptExitError
+from mcp_agent.core.prompt import Prompt
 from mcp_agent.core.request_params import RequestParams
 from mcp_agent.event_progress import ProgressAction
+from mcp_agent.llm.memory import Memory, SimpleMemory
+from mcp_agent.llm.sampling_format_converter import (
+    BasicFormatConverter,
+    ProviderFormatConverter,
+)
+from mcp_agent.logging.logger import get_logger
+from mcp_agent.mcp.interfaces import (
+    AugmentedLLMProtocol,
+    ModelT,
+)
 from mcp_agent.mcp.mcp_aggregator import MCPAggregator
+from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
+from mcp_agent.mcp.prompt_render import render_multipart_message
+from mcp_agent.mcp.prompt_serialization import multipart_messages_to_delimited_format
 from mcp_agent.ui.console_display import ConsoleDisplay
-from mcp_agent.workflows.llm.memory import Memory, SimpleMemory
+
+# Define type variables locally
+MessageParamT = TypeVar("MessageParamT")
+MessageT = TypeVar("MessageT")
+
+# Forward reference for type annotations
+if TYPE_CHECKING:
+    from mcp_agent.agents.agent import Agent
+    from mcp_agent.context import Context
+
 
 # TODO -- move this to a constant
 HUMAN_INPUT_TOOL_NAME = "__human_input__"
 
 
-class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, MessageT]):
+class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[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.
@@ -66,7 +72,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         instruction: str | None = None,
         name: str | None = None,
         request_params: RequestParams | None = None,
-        type_converter: Type[ProviderFormatConverter[MessageParamT, MessageT]] = BasicFormatConverter,
+        type_converter: Type[
+            ProviderFormatConverter[MessageParamT, MessageT]
+        ] = BasicFormatConverter,
         context: Optional["Context"] = None,
         **kwargs: dict[str, Any],
     ) -> None:
@@ -81,10 +89,14 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         self.logger = get_logger(__name__)
         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.name = agent.name if agent else name
+        self.instruction = agent.instruction if agent else instruction
+
+        # memory contains provider specific API types.
         self.history: Memory[MessageParamT] = SimpleMemory[MessageParamT]()
 
+        self.message_history: List[PromptMessageMultipart] = []
+
         # Initialize the display component
         self.display = ConsoleDisplay(config=self.context.config)
 
@@ -93,56 +105,82 @@ 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")
 
-    @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"""
+    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,
+        )
 
-    @abstractmethod
-    async def generate_str(
+    async def structured(
         self,
-        message: str | MessageParamT | List[MessageParamT],
+        prompt: List[PromptMessageMultipart],
+        model: Type[ModelT],
         request_params: RequestParams | None = None,
-    ) -> str:
-        """Request an LLM generation and return the string representation of the result"""
+    ) -> ModelT | None:
+        """Apply the prompt and return the result as a Pydantic model, or None if coercion fails"""
+        try:
+            result: PromptMessageMultipart = await self.generate(prompt, request_params)
+            json_data = from_json(result.first_text(), allow_partial=True)
+            validated_model = model.model_validate(json_data)
+            return cast("ModelT", validated_model)
+        except Exception as e:
+            logger = get_logger(__name__)
+            logger.error(f"Failed to parse structured response: {str(e)}")
+            return None
 
-    @abstractmethod
-    async def generate_structured(
+    async def generate(
         self,
-        message: str | MessageParamT | List[MessageParamT],
-        response_model: Type[ModelT],
+        multipart_messages: List[PromptMessageMultipart],
         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:
+    ) -> PromptMessageMultipart:
         """
-        Return the configured model (legacy support)
+        Create a completion with the LLM using the provided messages.
         """
-        if request_params and request_params.model:
-            return request_params.model
+        if multipart_messages[-1].first_text().startswith("***SAVE_HISTORY"):
+            parts: list[str] = multipart_messages[-1].first_text().split(" ", 1)
+            filename: str = (
+                parts[1].strip() if len(parts) > 1 else f"{self.name or 'assistant'}_prompts.txt"
+            )
+            await self._save_history(filename)
+            self.show_user_message(
+                f"History saved to {filename}", model=self.default_request_params.model, chat_turn=0
+            )
+            return Prompt.assistant(f"History saved to {filename}")
 
-        raise ModelConfigError("Internal Error: Model is not configured correctly")
+        self.message_history.extend(multipart_messages)
 
-    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,
+        if multipart_messages[-1].role == "user":
+            self.show_user_message(
+                render_multipart_message(multipart_messages[-1]),
+                model=self.default_request_params.model,
+                chat_turn=self.chat_turn(),
+            )
+
+        assistant_response: PromptMessageMultipart = await self._apply_prompt_provider_specific(
+            multipart_messages, request_params
         )
 
-    def _merge_request_params(self, default_params: RequestParams, provided_params: RequestParams) -> RequestParams:
+        self.message_history.append(assistant_response)
+        return assistant_response
+
+    def chat_turn(self) -> int:
+        """Return the current chat turn number"""
+        return 1 + sum(1 for message in self.message_history if message.role == "assistant")
+
+    def _merge_request_params(
+        self, default_params: RequestParams, provided_params: RequestParams
+    ) -> RequestParams:
         """Merge default and provided request parameters"""
 
         merged = default_params.model_dump()
@@ -176,24 +214,13 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         return default_request_params
 
     @classmethod
-    def convert_message_to_message_param(cls, message: MessageT, **kwargs: dict[str, Any]) -> 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 cast("MessageParamT", 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) -> None:
         """Display a tool result in a formatted panel."""
         self.display.show_tool_result(result)
@@ -208,10 +235,12 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
 
     async def show_assistant_message(
         self,
-        message_text: str | Text,
+        message_text: str | Text | None,
         highlight_namespaced_tool: str = "",
         title: str = "ASSISTANT",
     ) -> None:
+        if message_text is None:
+            message_text = Text("No content to display", style="dim green italic")
         """Display an assistant message in a formatted panel."""
         await self.display.show_assistant_message(
             message_text,
@@ -225,11 +254,15 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         """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
 
@@ -264,7 +297,9 @@ 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
@@ -283,66 +318,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
                 ],
             )
 
-    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)  # type: ignore
-                    if text_parts:
-                        return "\n".join(text_parts)
-
-        # For objects with content attribute
-        if hasattr(message, "content"):
-            content = message.content  # type: ignore
-            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)  # type: ignore
-        if result != str(message):
-            return result
-
-        # Additional handling for output-specific formats
-        if hasattr(message, "content"):
-            content = getattr(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) -> 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:
@@ -435,33 +413,28 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         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)
+        result = await self._apply_prompt_provider_specific(multipart_messages, None)
+        return result.first_text()
 
-    async def apply_prompt(
-        self,
-        multipart_messages: List["PromptMessageMultipart"],
-        request_params: RequestParams | None = None,
-    ) -> str:
+    async def _save_history(self, filename: str) -> None:
         """
-        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
+        Save the Message History to a file in a simple delimeted format.
         """
-        # Delegate to the provider-specific implementation
-        return await self._apply_prompt_template_provider_specific(multipart_messages, request_params)
+        # Convert to delimited format
+        delimited_content = multipart_messages_to_delimited_format(
+            self.message_history,
+        )
+
+        # Write to file
+        with open(filename, "w", encoding="utf-8") as f:
+            f.write("\n\n".join(delimited_content))
 
-    # this shouln't need to be very big...
-    async def _apply_prompt_template_provider_specific(
+    @abstractmethod
+    async def _apply_prompt_provider_specific(
         self,
         multipart_messages: List["PromptMessageMultipart"],
         request_params: RequestParams | None = None,
-    ) -> str:
+    ) -> PromptMessageMultipart:
         """
         Provider-specific implementation of apply_prompt_template.
         This default implementation handles basic text content for any LLM type.
@@ -475,101 +448,3 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             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.from_multipart()
-                    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 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")
-                    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.from_multipart()
-                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}"  # ignore # type: ignore
-                        )
-                    else:
-                        assistant_text_parts.append(
-                            f"[Resource Type: {mime_type}]\n{content.resource.text}"  # type ignore # type: ignore
-                        )
-                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
-
-
-#####################################
-### NEW INTERFACE METHODS BELOW   ###
-#####################################