fast-agent-mcp 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

mcp_agent/workflows/llm/augmented_llm.py

@@ -21,17 +21,18 @@ from mcp.types import (
     CreateMessageResult,
     ModelPreferences,
     SamplingMessage,
+    PromptMessage,
     TextContent,
+    GetPromptResult,
 )
 
 from mcp_agent.context_dependent import ContextDependent
 from mcp_agent.core.exceptions import PromptExitError
 from mcp_agent.event_progress import ProgressAction
-from mcp_agent.mcp.mcp_aggregator import MCPAggregator, SEP
+from mcp_agent.mcp.mcp_aggregator import MCPAggregator
 from mcp_agent.workflows.llm.llm_selector import ModelSelector
-from rich.panel import Panel
+from mcp_agent.ui.console_display import ConsoleDisplay
 from rich.text import Text
-from mcp_agent import console
 
 if TYPE_CHECKING:
     from mcp_agent.agents.agent import Agent
@@ -63,39 +64,96 @@ class Memory(Protocol, Generic[MessageParamT]):
 
     def __init__(self): ...
 
-    def extend(self, messages: List[MessageParamT]) -> None: ...
+    def extend(
+        self, messages: List[MessageParamT], is_prompt: bool = False
+    ) -> None: ...
 
-    def set(self, messages: List[MessageParamT]) -> None: ...
+    def set(self, messages: List[MessageParamT], is_prompt: bool = False) -> None: ...
 
-    def append(self, message: MessageParamT) -> None: ...
+    def append(self, message: MessageParamT, is_prompt: bool = False) -> None: ...
 
-    def get(self) -> List[MessageParamT]: ...
+    def get(self, include_history: bool = True) -> List[MessageParamT]: ...
 
-    def clear(self) -> None: ...
+    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.
 
-    def extend(self, messages: List[MessageParamT]):
-        self.history.extend(messages)
+        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]):
-        self.history = messages.copy()
+    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
 
-    def append(self, message: MessageParamT):
-        self.history.append(message)
+        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 get(self) -> List[MessageParamT]:
-        return self.history
+    def clear(self, clear_prompts: bool = False):
+        """
+        Clear history and optionally prompt messages.
 
-    def clear(self):
+        Args:
+            clear_prompts: If True, also clear prompt messages
+        """
         self.history = []
+        if clear_prompts:
+            self.prompt_messages = []
 
 
 class RequestParams(CreateMessageRequestParams):
@@ -224,6 +282,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         )
         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,
@@ -315,14 +376,6 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         self, default_params: RequestParams, provided_params: RequestParams
     ) -> RequestParams:
         """Merge default and provided request parameters"""
-        # Log parameter merging if debug logging is enabled
-        # self.context.config.logger.debug(
-        #     "Merging provided request params with defaults",
-        #     extra={
-        #         "defaults": default_params.model_dump(),
-        #         "provided": provided_params.model_dump(),
-        #     },
-        # )
 
         merged = default_params.model_dump()
         merged.update(provided_params.model_dump(exclude_unset=True))
@@ -374,6 +427,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         """Convert an MCP message (SamplingMessage) to an LLM input type."""
         return self.type_converter.from_mcp_message_param(param)
 
+    def from_mcp_prompt_message(self, message: PromptMessage) -> MessageParamT:
+        return self.type_converter.from_mcp_prompt_message(message)
+
     @classmethod
     def convert_message_to_message_param(
         cls, message: MessageT, **kwargs
@@ -397,113 +453,15 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
 
     def show_tool_result(self, result: CallToolResult):
         """Display a tool result in a formatted panel."""
-
-        if not self.context.config.logger.show_tools:
-            return
-
-        if result.isError:
-            style = "red"
-        else:
-            style = "magenta"
-
-        panel = Panel(
-            Text(
-                str(result.content), overflow="..."
-            ),  # TODO support multi-model/multi-part responses
-            title="[TOOL RESULT]",
-            title_align="right",
-            style=style,
-            border_style="bold white",
-            padding=(1, 2),
-        )
-
-        if self.context.config.logger.truncate_tools:
-            if len(str(result.content)) > 360:
-                panel.height = 8
-
-        console.console.print(panel)
-        console.console.print("\n")
+        self.display.show_tool_result(result)
 
     def show_oai_tool_result(self, result):
         """Display a tool result in a formatted panel."""
-
-        if not self.context.config.logger.show_tools:
-            return
-
-        panel = Panel(
-            Text(str(result), overflow="..."),  # TODO update openai support
-            title="[TOOL RESULT]",
-            title_align="right",
-            style="magenta",
-            border_style="bold white",
-            padding=(1, 2),
-        )
-
-        if self.context.config.logger.truncate_tools:
-            if len(str(result)) > 360:
-                panel.height = 8
-
-        console.console.print(panel)
-        console.console.print("\n")
+        self.display.show_oai_tool_result(result)
 
     def show_tool_call(self, available_tools, tool_name, tool_args):
         """Display a tool call in a formatted panel."""
-
-        if not self.context.config.logger.show_tools:
-            return
-
-        display_tool_list = Text()
-        for display_tool in available_tools:
-            # Handle both OpenAI and Anthropic tool formats
-            # TODO -- this should really be using the ToolCall abstraction and converting at the concrete layer??
-            if isinstance(display_tool, dict):
-                if "function" in display_tool:
-                    # OpenAI format
-                    tool_call_name = display_tool["function"]["name"]
-                else:
-                    # Anthropic format
-                    tool_call_name = display_tool["name"]
-            else:
-                # Handle potential object format (e.g., Pydantic models)
-                tool_call_name = (
-                    display_tool.function.name
-                    if hasattr(display_tool, "function")
-                    else display_tool.name
-                )
-
-            parts = (
-                tool_call_name.split(SEP)
-                if SEP in tool_call_name
-                else [tool_call_name, tool_call_name]
-            )
-            if tool_name.split(SEP)[0] == parts[0]:
-                if tool_call_name == tool_name:
-                    style = "magenta"
-                else:
-                    style = "dim white"
-
-                shortened_name = (
-                    parts[1] if len(parts[1]) <= 12 else parts[1][:11] + "…"
-                )
-                display_tool_list.append(f"[{shortened_name}] ", style)
-
-        panel = Panel(
-            Text(str(tool_args), overflow="ellipsis"),
-            title="[TOOL CALL]",
-            title_align="left",
-            style="magenta",
-            border_style="bold white",
-            subtitle=display_tool_list,
-            subtitle_align="left",
-            padding=(1, 2),
-        )
-
-        if self.context.config.logger.truncate_tools:
-            if len(str(tool_args)) > 360:
-                panel.height = 8
-
-        console.console.print(panel)
-        console.console.print("\n")
+        self.display.show_tool_call(available_tools, tool_name, tool_args)
 
     async def show_assistant_message(
         self,
@@ -512,62 +470,17 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         title: str = "ASSISTANT",
     ):
         """Display an assistant message in a formatted panel."""
-
-        if not self.context.config.logger.show_chat:
-            return
-
-        mcp_server_name = (
-            highlight_namespaced_tool.split(SEP)
-            if SEP in highlight_namespaced_tool
-            else [highlight_namespaced_tool]
-        )
-
-        display_server_list = Text()
-
-        tools = await self.aggregator.list_tools()
-        if any(tool.name == HUMAN_INPUT_TOOL_NAME for tool in tools.tools):
-            style = (
-                "green"
-                if highlight_namespaced_tool == HUMAN_INPUT_TOOL_NAME
-                else "dim white"
-            )
-            display_server_list.append("[human] ", style)
-
-        for server_name in await self.aggregator.list_servers():
-            style = "green" if server_name == mcp_server_name[0] else "dim white"
-            display_server_list.append(f"[{server_name}] ", style)
-
-        panel = Panel(
+        await self.display.show_assistant_message(
             message_text,
-            title=f"[{title}]{f' ({self.name})' if self.name else ''}",
-            title_align="left",
-            style="green",
-            border_style="bold white",
-            padding=(1, 2),
-            subtitle=display_server_list,
-            subtitle_align="left",
+            aggregator=self.aggregator,
+            highlight_namespaced_tool=highlight_namespaced_tool,
+            title=title,
+            name=self.name,
         )
-        console.console.print(panel)
-        console.console.print("\n")
 
     def show_user_message(self, message, model: str | None, chat_turn: int):
         """Display a user message in a formatted panel."""
-
-        if not self.context.config.logger.show_chat:
-            return
-
-        panel = Panel(
-            message,
-            title=f"{f'({self.name}) [USER]' if self.name else '[USER]'}",
-            title_align="right",
-            style="blue",
-            border_style="bold white",
-            padding=(1, 2),
-            subtitle=Text(f"{model} turn {chat_turn}", style="dim white"),
-            subtitle_align="left",
-        )
-        console.console.print(panel)
-        console.console.print("\n")
+        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
@@ -633,11 +546,62 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
             )
 
     def message_param_str(self, message: MessageParamT) -> str:
-        """Convert an input message to a string representation."""
+        """
+        Convert an input message to a string representation.
+        Tries to extract just the content when possible.
+        """
+        if isinstance(message, dict):
+            # For dictionary format messages
+            if "content" in message:
+                content = message["content"]
+                # Handle both string and structured content formats
+                if isinstance(content, str):
+                    return content
+                elif isinstance(content, list) and content:
+                    # Try to extract text from content parts
+                    text_parts = []
+                    for part in content:
+                        if isinstance(part, dict) and "text" in part:
+                            text_parts.append(part["text"])
+                        elif hasattr(part, "text"):
+                            text_parts.append(part.text)
+                    if text_parts:
+                        return "\n".join(text_parts)
+                    
+        # For objects with content attribute
+        if hasattr(message, "content"):
+            content = message.content
+            if isinstance(content, str):
+                return content
+            elif hasattr(content, "text"):
+                return content.text
+                
+        # Default fallback
         return str(message)
 
     def message_str(self, message: MessageT) -> str:
-        """Convert an output message to a string representation."""
+        """
+        Convert an output message to a string representation.
+        Tries to extract just the content when possible.
+        """
+        # First try to use the same method for consistency
+        result = self.message_param_str(message)
+        if result != str(message):
+            return result
+            
+        # Additional handling for output-specific formats
+        if hasattr(message, "content"):
+            content = message.content
+            if isinstance(content, list):
+                # Extract text from content blocks
+                text_parts = []
+                for block in content:
+                    if hasattr(block, "text") and block.text:
+                        text_parts.append(block.text)
+                if text_parts:
+                    return "\n".join(text_parts)
+        
+        # Default fallback
         return str(message)
 
     def _log_chat_progress(
@@ -668,6 +632,108 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol[MessageParamT, Message
         }
         self.logger.debug("Chat finished", data=data)
 
+    def _convert_prompt_messages(
+        self, prompt_messages: List[PromptMessage]
+    ) -> List[MessageParamT]:
+        """
+        Convert prompt messages to this LLM's specific message format.
+        To be implemented by concrete LLM classes.
+        """
+        raise NotImplementedError("Must be implemented by subclass")
+
+    async def show_prompt_loaded(
+        self,
+        prompt_name: str,
+        description: Optional[str] = None,
+        message_count: int = 0,
+        arguments: Optional[dict[str, str]] = None,
+    ):
+        """
+        Display information about a loaded prompt template.
+        
+        Args:
+            prompt_name: The name of the prompt
+            description: Optional description of the prompt
+            message_count: Number of messages in the prompt
+            arguments: Optional dictionary of arguments passed to the prompt
+        """
+        await self.display.show_prompt_loaded(
+            prompt_name=prompt_name,
+            description=description,
+            message_count=message_count,
+            agent_name=self.name,
+            aggregator=self.aggregator,
+            arguments=arguments,
+        )
+
+    async def apply_prompt_template(
+        self, prompt_result: GetPromptResult, prompt_name: str
+    ) -> str:
+        """
+        Apply a prompt template by adding it to the conversation history.
+        If the last message in the prompt is from a user, automatically
+        generate an assistant response.
+
+        Args:
+            prompt_result: The GetPromptResult containing prompt messages
+            prompt_name: The name of the prompt being applied
+
+        Returns:
+            String representation of the assistant's response if generated, 
+            or the last assistant message in the prompt
+        """
+        prompt_messages: List[PromptMessage] = prompt_result.messages
+        
+        # Check if we have any messages
+        if not prompt_messages:
+            return "Prompt contains no messages"
+
+        # Extract arguments if they were stored in the result
+        arguments = getattr(prompt_result, "arguments", None)
+
+        # Display information about the loaded prompt
+        await self.show_prompt_loaded(
+            prompt_name=prompt_name,
+            description=prompt_result.description,
+            message_count=len(prompt_messages),
+            arguments=arguments,
+        )
+        
+        # Check the last message role
+        last_message = prompt_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(prompt_messages) > 1:
+                previous_messages = prompt_messages[:-1]
+                converted = []
+                for msg in previous_messages:
+                    converted.append(self.type_converter.from_mcp_prompt_message(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)
+            
+            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 = []
+            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):
     """

mcp_agent/workflows/llm/augmented_llm_anthropic.py

@@ -39,8 +39,10 @@ from mcp_agent.workflows.llm.augmented_llm import (
 )
 from mcp_agent.core.exceptions import ProviderKeyError
 from mcp_agent.logging.logger import get_logger
+from mcp.types import PromptMessage
 from rich.text import Text
 
+_logger = get_logger(__name__)
 DEFAULT_ANTHROPIC_MODEL = "claude-3-7-sonnet-latest"
 
 
@@ -94,8 +96,9 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
                 "Please check that your API key is valid and not expired.",
             ) from e
 
-        if params.use_history:
-            messages.extend(self.history.get())
+        # Always include prompt messages, but only include conversation history 
+        # if use_history is True
+        messages.extend(self.history.get(include_history=params.use_history))
 
         if isinstance(message, str):
             messages.append({"role": "user", "content": message})
@@ -287,8 +290,17 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
                         )
                     )
 
+        # Only save the new conversation messages to history if use_history is true
+        # Keep the prompt messages separate
         if params.use_history:
-            self.history.set(messages)
+            # Get current prompt messages
+            prompt_messages = self.history.get(include_history=False)
+            
+            # Calculate new conversation messages (excluding prompts)
+            new_messages = messages[len(prompt_messages):]
+            
+            # Update conversation history
+            self.history.set(new_messages)
 
         self._log_chat_finished(model=model)
 
@@ -525,6 +537,40 @@ class AnthropicMCPTypeConverter(ProviderToMCPConverter[MessageParam, Message]):
             **typed_dict_extras(param, ["role", "content"]),
         )
 
+    @classmethod
+    def from_mcp_prompt_message(cls, message: PromptMessage) -> MessageParam:
+        """Convert an MCP PromptMessage to an Anthropic MessageParam."""
+
+        # Extract content text
+        content_text = (
+            message.content.text
+            if hasattr(message.content, "text")
+            else str(message.content)
+        )
+
+        # Extract extras for flexibility
+        extras = message.model_dump(exclude={"role", "content"})
+
+        # Handle based on role
+        if message.role == "user":
+            return {"role": "user", "content": content_text, **extras}
+        elif message.role == "assistant":
+            return {
+                "role": "assistant",
+                "content": [{"type": "text", "text": content_text}],
+                **extras,
+            }
+        else:
+            # Fall back to user for any unrecognized role, including "system"
+            _logger.warning(
+                f"Unsupported role '{message.role}' in PromptMessage. Falling back to 'user' role."
+            )
+            return {
+                "role": "user",
+                "content": f"[{message.role.upper()}] {content_text}",
+                **extras,
+            }
+
 
 def mcp_content_to_anthropic_content(
     content: TextContent | ImageContent | EmbeddedResource,