fast-agent-mcp 0.2.58__py3-none-any.whl → 0.3.1__py3-none-any.whl

mcp_agent/llm/augmented_llm.py → fast_agent/llm/fastagent_llm.py

@@ -10,44 +10,32 @@ from typing import (
     Tuple,
     Type,
     TypeVar,
-    Union,
     cast,
 )
 
+from mcp import Tool
 from mcp.types import (
-    CallToolRequest,
-    CallToolResult,
     GetPromptResult,
     PromptMessage,
-    TextContent,
 )
 from openai import NotGiven
 from openai.lib._parsing import type_to_response_format_param as _type_to_response_format
 from pydantic_core import from_json
-from rich.text import Text
-
-from mcp_agent.context_dependent import ContextDependent
-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.model_database import ModelDatabase
-from mcp_agent.llm.provider_types import Provider
-from mcp_agent.llm.sampling_format_converter import (
-    BasicFormatConverter,
-    ProviderFormatConverter,
-)
-from mcp_agent.llm.usage_tracking import TurnUsage, UsageAccumulator
-from mcp_agent.logging.logger import get_logger
-from mcp_agent.mcp.helpers.content_helpers import get_text
-from mcp_agent.mcp.interfaces import (
-    AugmentedLLMProtocol,
+
+from fast_agent.context_dependent import ContextDependent
+from fast_agent.core.logging.logger import get_logger
+from fast_agent.core.prompt import Prompt
+from fast_agent.event_progress import ProgressAction
+from fast_agent.interfaces import (
+    FastAgentLLMProtocol,
     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 fast_agent.llm.memory import Memory, SimpleMemory
+from fast_agent.llm.model_database import ModelDatabase
+from fast_agent.llm.provider_types import Provider
+from fast_agent.llm.usage_tracking import TurnUsage, UsageAccumulator
+from fast_agent.mcp.helpers.content_helpers import get_text
+from fast_agent.types import PromptMessageExtended, RequestParams
 
 # Define type variables locally
 MessageParamT = TypeVar("MessageParamT")
@@ -55,15 +43,11 @@ MessageT = TypeVar("MessageT")
 
 # Forward reference for type annotations
 if TYPE_CHECKING:
-    from mcp_agent.agents.agent import Agent
-    from mcp_agent.context import Context
+    from fast_agent.context import Context
 
 
-# TODO -- move this to a constant
-HUMAN_INPUT_TOOL_NAME = "__human_input__"
-
 # Context variable for storing MCP metadata
-_mcp_metadata_var: ContextVar[Dict[str, Any] | None] = ContextVar('mcp_metadata', default=None)
+_mcp_metadata_var: ContextVar[Dict[str, Any] | None] = ContextVar("mcp_metadata", default=None)
 
 
 def deep_merge(dict1: Dict[Any, Any], dict2: Dict[Any, Any]) -> Dict[Any, Any]:
@@ -89,7 +73,7 @@ def deep_merge(dict1: Dict[Any, Any], dict2: Dict[Any, Any]) -> Dict[Any, Any]:
     return dict1
 
 
-class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT, MessageT]):
+class FastAgentLLM(ContextDependent, FastAgentLLMProtocol, Generic[MessageParamT, MessageT]):
     # Common parameter names used across providers
     PARAM_MESSAGES = "messages"
     PARAM_MODEL = "model"
@@ -107,42 +91,28 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
     BASE_EXCLUDE_FIELDS = {PARAM_METADATA}
 
     """
-    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.
-    Our current models can actively use these capabilities—generating their own search queries,
-    selecting appropriate tools, and determining what information to retain.
+    Implementation of the Llm Protocol - intended be subclassed for Provider
+    or behaviour specific reasons. Contains convenience and template methods.
     """
 
-    provider: Provider | None = None
-
     def __init__(
         self,
         provider: Provider,
-        agent: Optional["Agent"] = None,
-        server_names: List[str] | None = None,
         instruction: str | None = None,
         name: str | None = None,
         request_params: RequestParams | None = None,
-        type_converter: Type[
-            ProviderFormatConverter[MessageParamT, MessageT]
-        ] = BasicFormatConverter,
         context: Optional["Context"] = None,
         model: Optional[str] = None,
         api_key: Optional[str] = None,
         **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.
-        If an agent is provided, all other properties are optional
 
         Args:
-            agent: Optional Agent that owns this LLM
-            server_names: List of MCP server names to connect to
+            provider: LLM API Provider
             instruction: System prompt for the LLM
-            name: Optional name identifier for the LLM
-            request_params: RequestParams to configure LLM behavior
-            type_converter: Provider-specific format converter class
+            name: Name for the LLM (usually attached Agent name)
+            request_params: RequestParams to configure LLM behaviour
             context: Application context
             model: Optional model name override
             **kwargs: Additional provider-specific parameters
@@ -152,24 +122,18 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
         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.name = agent.name if agent else name
-        self.instruction = agent.instruction if agent else instruction
-        self.provider = provider
+        self.name: str = name or "fast-agent"
+        self.instruction = instruction
+        self._provider = provider
         # memory contains provider specific API types.
         self.history: Memory[MessageParamT] = SimpleMemory[MessageParamT]()
 
-        self._message_history: List[PromptMessageMultipart] = []
+        self._message_history: List[PromptMessageExtended] = []
 
         # Initialize the display component
-        if self.context.config and self.context.config.logger.use_legacy_display:
-            from mcp_agent.ui.console_display_legacy import ConsoleDisplay
-        else:
-            from mcp_agent.ui.console_display import ConsoleDisplay
-        self.display = ConsoleDisplay(config=self.context.config)
+        from fast_agent.ui.console_display import ConsoleDisplay
 
-        # Tool call counter for current turn
-        self._current_turn_tool_calls = 0
+        self.display = ConsoleDisplay(config=self.context.config)
 
         # Initialize default parameters, passing model info
         model_kwargs = kwargs.copy()
@@ -183,13 +147,15 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
                 self.default_request_params, self._init_request_params
             )
 
-        self.type_converter = type_converter
+        # Cache effective model name for type-safe access
+        self._model_name: Optional[str] = getattr(self.default_request_params, "model", None)
+
         self.verb = kwargs.get("verb")
 
         self._init_api_key = api_key
 
         # Initialize usage tracking
-        self.usage_accumulator = UsageAccumulator()
+        self._usage_accumulator = UsageAccumulator()
 
     def _initialize_default_params(self, kwargs: dict) -> RequestParams:
         """Initialize default parameters for the LLM.
@@ -209,43 +175,47 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
 
     async def generate(
         self,
-        multipart_messages: List[Union[PromptMessageMultipart, PromptMessage]],
+        messages: List[PromptMessageExtended],
         request_params: RequestParams | None = None,
-    ) -> PromptMessageMultipart:
+        tools: List[Tool] | None = None,
+    ) -> PromptMessageExtended:
         """
-        Create a completion with the LLM using the provided messages.
-        """
-        # note - check changes here are mirrored in structured(). i've thought hard about
-        # a strategy to reduce duplication etc, but aiming for simple but imperfect for the moment
+        Generate a completion using normalized message lists.
+
+        This is the primary LLM interface that works directly with
+        List[PromptMessageExtended] for efficient internal usage.
 
-        # Convert PromptMessage to PromptMessageMultipart if needed
-        if multipart_messages and isinstance(multipart_messages[0], PromptMessage):
-            multipart_messages = PromptMessageMultipart.to_multipart(multipart_messages)
+        Args:
+            messages: List of PromptMessageExtended objects
+            request_params: Optional parameters to configure the LLM request
+            tools: Optional list of tools available to the LLM
 
+        Returns:
+            A PromptMessageExtended containing the Assistant response
+        """
         # TODO -- create a "fast-agent" control role rather than magic strings
 
-        if multipart_messages[-1].first_text().startswith("***SAVE_HISTORY"):
-            parts: list[str] = multipart_messages[-1].first_text().split(" ", 1)
+        if messages[-1].first_text().startswith("***SAVE_HISTORY"):
+            parts: list[str] = 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}")
 
-        self._precall(multipart_messages)
+        self._precall(messages)
 
         # Store MCP metadata in context variable
         final_request_params = self.get_request_params(request_params)
         if final_request_params.mcp_metadata:
             _mcp_metadata_var.set(final_request_params.mcp_metadata)
 
-        assistant_response: PromptMessageMultipart = await self._apply_prompt_provider_specific(
-            multipart_messages, request_params
+        assistant_response: PromptMessageExtended = await self._apply_prompt_provider_specific(
+            messages, request_params, tools
         )
 
+        self.usage_accumulator.count_tools(len(assistant_response.tool_calls or {}))
+
         # add generic error and termination reason handling/rollback
         self._message_history.append(assistant_response)
 
@@ -254,10 +224,11 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
     @abstractmethod
     async def _apply_prompt_provider_specific(
         self,
-        multipart_messages: List["PromptMessageMultipart"],
+        multipart_messages: List["PromptMessageExtended"],
         request_params: RequestParams | None = None,
+        tools: List[Tool] | None = None,
         is_template: bool = False,
-    ) -> PromptMessageMultipart:
+    ) -> PromptMessageExtended:
         """
         Provider-specific implementation of apply_prompt_template.
         This default implementation handles basic text content for any LLM type.
@@ -265,7 +236,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
         multimodal content appropriately.
 
         Args:
-            multipart_messages: List of PromptMessageMultipart objects parsed from the prompt template
+            multipart_messages: List of PromptMessageExtended objects parsed from the prompt template
 
         Returns:
             String representation of the assistant's response if generated,
@@ -274,25 +245,36 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
 
     async def structured(
         self,
-        multipart_messages: List[Union[PromptMessageMultipart, PromptMessage]],
+        messages: List[PromptMessageExtended],
         model: Type[ModelT],
         request_params: RequestParams | None = None,
-    ) -> Tuple[ModelT | None, PromptMessageMultipart]:
-        """Return a structured response from the LLM using the provided messages."""
+    ) -> Tuple[ModelT | None, PromptMessageExtended]:
+        """
+        Generate a structured response using normalized message lists.
+
+        This is the primary LLM interface for structured output that works directly with
+        List[PromptMessageExtended] for efficient internal usage.
+
+        Args:
+            messages: List of PromptMessageExtended objects
+            model: The Pydantic model class to parse the response into
+            request_params: Optional parameters to configure the LLM request
+
+        Returns:
+            Tuple of (parsed model instance or None, assistant response message)
+        """
 
-        # Convert PromptMessage to PromptMessageMultipart if needed
-        if multipart_messages and isinstance(multipart_messages[0], PromptMessage):
-            multipart_messages = PromptMessageMultipart.to_multipart(multipart_messages)
+        self._precall(messages)
 
-        self._precall(multipart_messages)
-        
         # Store MCP metadata in context variable
         final_request_params = self.get_request_params(request_params)
+
+        # TODO -- this doesn't need to go here anymore.
         if final_request_params.mcp_metadata:
             _mcp_metadata_var.set(final_request_params.mcp_metadata)
-            
+
         result, assistant_response = await self._apply_prompt_provider_specific_structured(
-            multipart_messages, model, request_params
+            messages, model, request_params
         )
 
         self._message_history.append(assistant_response)
@@ -339,10 +321,10 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
 
     async def _apply_prompt_provider_specific_structured(
         self,
-        multipart_messages: List[PromptMessageMultipart],
+        multipart_messages: List[PromptMessageExtended],
         model: Type[ModelT],
         request_params: RequestParams | None = None,
-    ) -> Tuple[ModelT | None, PromptMessageMultipart]:
+    ) -> Tuple[ModelT | None, PromptMessageExtended]:
         """Base class attempts to parse JSON - subclasses can use provider specific functionality"""
 
         request_params = self.get_request_params(request_params)
@@ -352,14 +334,14 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             if schema is not NotGiven:
                 request_params.response_format = schema
 
-        result: PromptMessageMultipart = await self._apply_prompt_provider_specific(
+        result: PromptMessageExtended = await self._apply_prompt_provider_specific(
             multipart_messages, request_params
         )
         return self._structured_from_multipart(result, model)
 
     def _structured_from_multipart(
-        self, message: PromptMessageMultipart, model: Type[ModelT]
-    ) -> Tuple[ModelT | None, PromptMessageMultipart]:
+        self, message: PromptMessageExtended, model: Type[ModelT]
+    ) -> Tuple[ModelT | None, PromptMessageExtended]:
         """Parse the content of a PromptMessage and return the structured model and message itself"""
         try:
             text = get_text(message.content[-1]) or ""
@@ -371,15 +353,10 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             logger.warning(f"Failed to parse structured response: {str(e)}")
             return None, message
 
-    def _precall(self, multipart_messages: List[PromptMessageMultipart]) -> None:
+    def _precall(self, multipart_messages: List[PromptMessageExtended]) -> None:
         """Pre-call hook to modify the message before sending it to the provider."""
+        # Ensure all messages are PromptMessageExtended before extending history
         self._message_history.extend(multipart_messages)
-        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(),
-            )
 
     def chat_turn(self) -> int:
         """Return the current chat turn number"""
@@ -471,108 +448,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
         # Many LLM implementations will allow the same type for input and output messages
         return cast("MessageParamT", message)
 
-    def show_tool_result(self, result: CallToolResult) -> None:
-        """Display a tool result in a formatted panel."""
-        self.display.show_tool_result(result, name=self.name)
-
-    def show_tool_call(self, available_tools, tool_name, tool_args) -> None:
-        """Display a tool call in a formatted panel."""
-        self._current_turn_tool_calls += 1
-        self.display.show_tool_call(available_tools, tool_name, tool_args, name=self.name)
-
-    def _reset_turn_tool_calls(self) -> None:
-        """Reset tool call counter for new turn."""
-        self._current_turn_tool_calls = 0
-
     def _finalize_turn_usage(self, turn_usage: "TurnUsage") -> None:
         """Set tool call count on TurnUsage and add to accumulator."""
-        turn_usage.set_tool_calls(self._current_turn_tool_calls)
-        self.usage_accumulator.add_turn(turn_usage)
-
-    async def show_assistant_message(
-        self,
-        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,
-            aggregator=self.aggregator,
-            highlight_namespaced_tool=highlight_namespaced_tool,
-            title=title,
-            name=self.name,
-        )
-
-    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:
-        """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:
-        """Called after a tool execution. Can modify the result before it's returned."""
-        return result
-
-    async def call_tool(
-        self,
-        request: CallToolRequest,
-        tool_call_id: str | None = None,
-    ) -> CallToolResult:
-        """Call a tool with the given parameters and optional ID"""
-
-        try:
-            preprocess = await self.pre_tool_call(
-                tool_call_id=tool_call_id,
-                request=request,
-            )
-
-            if isinstance(preprocess, bool):
-                if not preprocess:
-                    return CallToolResult(
-                        isError=True,
-                        content=[
-                            TextContent(
-                                type="text",
-                                text=f"Error: Tool '{request.params.name}' was not allowed to run.",
-                            )
-                        ],
-                    )
-            else:
-                request = preprocess
-
-            tool_name = request.params.name
-            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
-            )
-
-            if isinstance(postprocess, CallToolResult):
-                result = postprocess
-
-            return result
-        except PromptExitError:
-            raise
-        except Exception as e:
-            return CallToolResult(
-                isError=True,
-                content=[
-                    TextContent(
-                        type="text",
-                        text=f"Error executing tool '{request.params.name}': {str(e)}",
-                    )
-                ],
-            )
+        self._usage_accumulator.add_turn(turn_usage)
 
     def _log_chat_progress(
         self, chat_turn: Optional[int] = None, model: Optional[str] = None
@@ -661,7 +539,6 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             description=description,
             message_count=message_count,
             agent_name=self.name,
-            aggregator=self.aggregator,
             arguments=arguments,
         )
 
@@ -679,7 +556,7 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             String representation of the assistant's response if generated,
             or the last assistant message in the prompt
         """
-        from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
+        from fast_agent.types import PromptMessageExtended
 
         # Check if we have any messages
         if not prompt_result.messages:
@@ -696,8 +573,8 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             arguments=arguments,
         )
 
-        # Convert to PromptMessageMultipart objects
-        multipart_messages = PromptMessageMultipart.parse_get_prompt_result(prompt_result)
+        # Convert to PromptMessageExtended objects
+        multipart_messages = PromptMessageExtended.parse_get_prompt_result(prompt_result)
 
         # Delegate to the provider-specific implementation
         result = await self._apply_prompt_provider_specific(
@@ -712,21 +589,21 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
         Uses JSON format for .json files (MCP SDK compatible format) and
         delimited text format for other extensions.
         """
-        from mcp_agent.mcp.prompt_serialization import save_messages_to_file
+        from fast_agent.mcp.prompt_serialization import save_messages_to_file
 
         # Save messages using the unified save function that auto-detects format
         save_messages_to_file(self._message_history, filename)
 
     @property
-    def message_history(self) -> List[PromptMessageMultipart]:
+    def message_history(self) -> List[PromptMessageExtended]:
         """
-        Return the agent's message history as PromptMessageMultipart objects.
+        Return the agent's message history as PromptMessageExtended objects.
 
         This history can be used to transfer state between agents or for
         analysis and debugging purposes.
 
         Returns:
-            List of PromptMessageMultipart objects representing the conversation history
+            List of PromptMessageExtended objects representing the conversation history
         """
         return self._message_history
 
@@ -734,11 +611,15 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
         if self._init_api_key:
             return self._init_api_key
 
-        from mcp_agent.llm.provider_key_manager import ProviderKeyManager
+        from fast_agent.llm.provider_key_manager import ProviderKeyManager
 
         assert self.provider
         return ProviderKeyManager.get_api_key(self.provider.value, self.context.config)
 
+    @property
+    def usage_accumulator(self):
+        return self._usage_accumulator
+
     def get_usage_summary(self) -> dict:
         """
         Get a summary of usage statistics for this LLM instance.
@@ -747,4 +628,32 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             Dictionary containing usage statistics including tokens, cache metrics,
             and context window utilization.
         """
-        return self.usage_accumulator.get_summary()
+        return self._usage_accumulator.get_summary()
+
+    @property
+    def provider(self) -> Provider:
+        """
+        Return the LLM provider type.
+
+        Returns:
+            The Provider enum value representing the LLM provider
+        """
+        return self._provider
+
+    @property
+    def model_name(self) -> str | None:
+        """Return the effective model name, if set."""
+        return self._model_name
+
+    @property
+    def model_info(self):
+        """Return resolved model information with capabilities.
+
+        Uses a lightweight resolver backed by the ModelDatabase and provides
+        text/document/vision flags, context window, etc.
+        """
+        from fast_agent.llm.model_info import ModelInfo
+
+        if not self._model_name:
+            return None
+        return ModelInfo.from_name(self._model_name, self._provider)