fast-agent-mcp 0.2.16__py3-none-any.whl → 0.2.18__py3-none-any.whl

mcp_agent/llm/augmented_llm.py

@@ -18,6 +18,8 @@ from mcp.types import (
     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
 
@@ -58,6 +60,20 @@ HUMAN_INPUT_TOOL_NAME = "__human_input__"
 
 
 class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT, MessageT]):
+    # Common parameter names used across providers
+    PARAM_MESSAGES = "messages"
+    PARAM_MODEL = "model"
+    PARAM_MAX_TOKENS = "maxTokens"
+    PARAM_SYSTEM_PROMPT = "systemPrompt"
+    PARAM_STOP_SEQUENCES = "stopSequences"
+    PARAM_PARALLEL_TOOL_CALLS = "parallel_tool_calls"
+    PARAM_METADATA = "metadata"
+    PARAM_USE_HISTORY = "use_history"
+    PARAM_MAX_ITERATIONS = "max_iterations"
+
+    # Base set of fields that should always be excluded
+    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.
@@ -141,26 +157,6 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             use_history=True,
         )
 
-    async def structured(
-        self,
-        prompt: List[PromptMessageMultipart],
-        model: Type[ModelT],
-        request_params: RequestParams | None = None,
-    ) -> Tuple[ModelT | None, PromptMessageMultipart]:
-        """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)
-            final_generation = get_text(result.content[-1]) or ""
-            await self.show_assistant_message(final_generation)
-            json_data = from_json(final_generation, allow_partial=True)
-            validated_model = model.model_validate(json_data)
-
-            return cast("ModelT", validated_model), Prompt.assistant(json_data)
-        except Exception as e:
-            logger = get_logger(__name__)
-            logger.error(f"Failed to parse structured response: {str(e)}")
-            return None, Prompt.assistant(f"Failed to parse structured response: {str(e)}")
-
     async def generate(
         self,
         multipart_messages: List[PromptMessageMultipart],
@@ -169,6 +165,12 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
         """
         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
+
+        # We never expect this for structured() calls - this is for interactive use - developers
+        # can do this programatically
+        # 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)
             filename: str = (
@@ -180,26 +182,174 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             )
             return Prompt.assistant(f"History saved to {filename}")
 
-        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(),
-            )
+        self._precall(multipart_messages)
 
         assistant_response: PromptMessageMultipart = await self._apply_prompt_provider_specific(
             multipart_messages, request_params
         )
 
+        # add generic error and termination reason handling/rollback
         self._message_history.append(assistant_response)
         return assistant_response
 
+    @abstractmethod
+    async def _apply_prompt_provider_specific(
+        self,
+        multipart_messages: List["PromptMessageMultipart"],
+        request_params: RequestParams | None = None,
+        is_template: bool = False,
+    ) -> PromptMessageMultipart:
+        """
+        Provider-specific implementation of apply_prompt_template.
+        This default implementation handles basic text content for any LLM type.
+        Provider-specific subclasses should override this method to handle
+        multimodal content appropriately.
+
+        Args:
+            multipart_messages: List of PromptMessageMultipart objects parsed from the prompt template
+
+        Returns:
+            String representation of the assistant's response if generated,
+            or the last assistant message in the prompt
+        """
+
+    async def structured(
+        self,
+        multipart_messages: List[PromptMessageMultipart],
+        model: Type[ModelT],
+        request_params: RequestParams | None = None,
+    ) -> Tuple[ModelT | None, PromptMessageMultipart]:
+        """Return a structured response from the LLM using the provided messages."""
+        self._precall(multipart_messages)
+        result, assistant_response = await self._apply_prompt_provider_specific_structured(
+            multipart_messages, model, request_params
+        )
+
+        self._message_history.append(assistant_response)
+        return result, assistant_response
+
+    @staticmethod
+    def model_to_response_format(
+        model: Type[Any],
+    ) -> Any:
+        """
+        Convert a pydantic model to the appropriate response format schema.
+        This allows for reuse in multiple provider implementations.
+
+        Args:
+            model: The pydantic model class to convert to a schema
+
+        Returns:
+            Provider-agnostic schema representation or NotGiven if conversion fails
+        """
+        return _type_to_response_format(model)
+
+    @staticmethod
+    def model_to_schema_str(
+        model: Type[Any],
+    ) -> str:
+        """
+        Convert a pydantic model to a schema string representation.
+        This provides a simpler interface for provider implementations
+        that need a string representation.
+
+        Args:
+            model: The pydantic model class to convert to a schema
+
+        Returns:
+            Schema as a string, or empty string if conversion fails
+        """
+        import json
+
+        try:
+            schema = model.model_json_schema()
+            return json.dumps(schema)
+        except Exception:
+            return ""
+
+    async def _apply_prompt_provider_specific_structured(
+        self,
+        multipart_messages: List[PromptMessageMultipart],
+        model: Type[ModelT],
+        request_params: RequestParams | None = None,
+    ) -> Tuple[ModelT | None, PromptMessageMultipart]:
+        """Base class attempts to parse JSON - subclasses can use provider specific functionality"""
+
+        request_params = self.get_request_params(request_params)
+
+        if not request_params.response_format:
+            schema = self.model_to_response_format(model)
+            if schema is not NotGiven:
+                request_params.response_format = schema
+
+        result: PromptMessageMultipart = 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]:
+        """Parse the content of a PromptMessage and return the structured model and message itself"""
+        try:
+            text = get_text(message.content[-1]) or ""
+            json_data = from_json(text, allow_partial=True)
+            validated_model = model.model_validate(json_data)
+            return cast("ModelT", validated_model), message
+        except ValueError as e:
+            logger = get_logger(__name__)
+            logger.warning(f"Failed to parse structured response: {str(e)}")
+            return None, message
+
+    def _precall(self, multipart_messages: List[PromptMessageMultipart]) -> None:
+        """Pre-call hook to modify the message before sending it to the provider."""
+        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"""
         return 1 + sum(1 for message in self._message_history if message.role == "assistant")
 
+    def prepare_provider_arguments(
+        self,
+        base_args: dict,
+        request_params: RequestParams,
+        exclude_fields: set | None = None,
+    ) -> dict:
+        """
+        Prepare arguments for provider API calls by merging request parameters.
+
+        Args:
+            base_args: Base arguments dictionary with provider-specific required parameters
+            params: The RequestParams object containing all parameters
+            exclude_fields: Set of field names to exclude from params. If None, uses BASE_EXCLUDE_FIELDS.
+
+        Returns:
+            Complete arguments dictionary with all applicable parameters
+        """
+        # Start with base arguments
+        arguments = base_args.copy()
+
+        # Use provided exclude_fields or fall back to base exclusions
+        exclude_fields = exclude_fields or self.BASE_EXCLUDE_FIELDS.copy()
+
+        # Add all fields from params that aren't explicitly excluded
+        params_dict = request_params.model_dump(exclude=exclude_fields)
+        for key, value in params_dict.items():
+            if value is not None and key not in arguments:
+                arguments[key] = value
+
+        # Finally, add any metadata fields as a last layer of overrides
+        if request_params.metadata:
+            arguments.update(request_params.metadata)
+
+        return arguments
+
     def _merge_request_params(
         self, default_params: RequestParams, provided_params: RequestParams
     ) -> RequestParams:
@@ -214,7 +364,6 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
     def get_request_params(
         self,
         request_params: RequestParams | None = None,
-        default: RequestParams | None = None,
     ) -> RequestParams:
         """
         Get request parameters with merged-in defaults and overrides.
@@ -223,17 +372,12 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
             default: The default request parameters to use as the base.
                 If unspecified, self.default_request_params will be used.
         """
-        # Start with the defaults
-        default_request_params = default or self.default_request_params
-
-        if not default_request_params:
-            default_request_params = self._initialize_default_params({})
 
         # If user provides overrides, merge them with defaults
         if request_params:
-            return self._merge_request_params(default_request_params, request_params)
+            return self._merge_request_params(self.default_request_params, request_params)
 
-        return default_request_params
+        return self.default_request_params.model_copy()
 
     @classmethod
     def convert_message_to_message_param(
@@ -435,7 +579,9 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
         multipart_messages = PromptMessageMultipart.parse_get_prompt_result(prompt_result)
 
         # Delegate to the provider-specific implementation
-        result = await self._apply_prompt_provider_specific(multipart_messages, None)
+        result = await self._apply_prompt_provider_specific(
+            multipart_messages, None, is_template=True
+        )
         return result.first_text()
 
     async def _save_history(self, filename: str) -> None:
@@ -450,26 +596,6 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
         # Save messages using the unified save function that auto-detects format
         save_messages_to_file(self._message_history, filename)
 
-    @abstractmethod
-    async def _apply_prompt_provider_specific(
-        self,
-        multipart_messages: List["PromptMessageMultipart"],
-        request_params: RequestParams | None = None,
-    ) -> PromptMessageMultipart:
-        """
-        Provider-specific implementation of apply_prompt_template.
-        This default implementation handles basic text content for any LLM type.
-        Provider-specific subclasses should override this method to handle
-        multimodal content appropriately.
-
-        Args:
-            multipart_messages: List of PromptMessageMultipart objects parsed from the prompt template
-
-        Returns:
-            String representation of the assistant's response if generated,
-            or the last assistant message in the prompt
-        """
-
     @property
     def message_history(self) -> List[PromptMessageMultipart]:
         """

mcp_agent/llm/augmented_llm_passthrough.py

@@ -143,7 +143,6 @@ class PassthroughLLM(AugmentedLLM):
     ) -> PromptMessageMultipart:
         last_message = multipart_messages[-1]
 
-        # TODO -- improve when we support Audio/Multimodal gen
         if self.is_tool_call(last_message):
             result = Prompt.assistant(await self.generate_str(last_message.first_text()))
             await self.show_assistant_message(result.first_text())
@@ -158,6 +157,7 @@ class PassthroughLLM(AugmentedLLM):
             await self.show_assistant_message(self._fixed_response)
             return Prompt.assistant(self._fixed_response)
         else:
+            # TODO -- improve when we support Audio/Multimodal gen models e.g. gemini . This should really just return the input as "assistant"...
             concatenated: str = "\n".join(message.all_text() for message in multipart_messages)
             await self.show_assistant_message(concatenated)
             return Prompt.assistant(concatenated)

mcp_agent/llm/augmented_llm_playback.py

@@ -1,9 +1,11 @@
-from typing import Any, List
+from typing import Any, List, Type
 
+from mcp_agent.core.exceptions import ModelConfigError
 from mcp_agent.core.prompt import Prompt
 from mcp_agent.llm.augmented_llm import RequestParams
 from mcp_agent.llm.augmented_llm_passthrough import PassthroughLLM
 from mcp_agent.llm.provider_types import Provider
+from mcp_agent.mcp.interfaces import ModelT
 from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
 from mcp_agent.mcp.prompts.prompt_helpers import MessageContent
 
@@ -82,3 +84,21 @@ class PlaybackLLM(PassthroughLLM):
         )
 
         return response
+
+    async def structured(
+        self,
+        multipart_messages: List[PromptMessageMultipart],
+        model: Type[ModelT],
+        request_params: RequestParams | None = None,
+    ) -> tuple[ModelT | None, PromptMessageMultipart]:
+        """
+        Handle structured requests by returning the next assistant message.
+        """
+
+        if -1 == self._current_index:
+            raise ModelConfigError("Use generate() to load playback history")
+
+        return self._structured_from_multipart(
+            self._get_next_assistant_message(),
+            model,
+        )

mcp_agent/llm/memory.py

@@ -19,7 +19,7 @@ class Memory(Protocol, Generic[MessageParamT]):
 
     def append(self, message: MessageParamT, is_prompt: bool = False) -> None: ...
 
-    def get(self, include_history: bool = True) -> List[MessageParamT]: ...
+    def get(self, include_completion_history: bool = True) -> List[MessageParamT]: ...
 
     def clear(self, clear_prompts: bool = False) -> None: ...
 
@@ -75,7 +75,7 @@ class SimpleMemory(Memory, Generic[MessageParamT]):
         else:
             self.history.append(message)
 
-    def get(self, include_history: bool = True) -> List[MessageParamT]:
+    def get(self, include_completion_history: bool = True) -> List[MessageParamT]:
         """
         Get all messages in memory.
 
@@ -86,7 +86,7 @@ class SimpleMemory(Memory, Generic[MessageParamT]):
         Returns:
             Combined list of prompt messages and optionally history messages
         """
-        if include_history:
+        if include_completion_history:
             return self.prompt_messages + self.history
         else:
             return self.prompt_messages.copy()

mcp_agent/llm/model_factory.py

@@ -12,6 +12,7 @@ from mcp_agent.llm.provider_types import Provider
 from mcp_agent.llm.providers.augmented_llm_anthropic import AnthropicAugmentedLLM
 from mcp_agent.llm.providers.augmented_llm_deepseek import DeepSeekAugmentedLLM
 from mcp_agent.llm.providers.augmented_llm_generic import GenericAugmentedLLM
+from mcp_agent.llm.providers.augmented_llm_google import GoogleAugmentedLLM
 from mcp_agent.llm.providers.augmented_llm_openai import OpenAIAugmentedLLM
 from mcp_agent.llm.providers.augmented_llm_openrouter import OpenRouterAugmentedLLM
 from mcp_agent.mcp.interfaces import AugmentedLLMProtocol
@@ -107,6 +108,7 @@ class ModelFactory:
         Provider.FAST_AGENT: PassthroughLLM,
         Provider.DEEPSEEK: DeepSeekAugmentedLLM,
         Provider.GENERIC: GenericAugmentedLLM,
+        Provider.GOOGLE: GoogleAugmentedLLM,  # type: ignore
         Provider.OPENROUTER: OpenRouterAugmentedLLM,
     }
 
@@ -161,7 +163,7 @@ class ModelFactory:
         Creates a factory function that follows the attach_llm protocol.
 
         Args:
-            model_string: The model specification string (e.g. "gpt-4o.high")
+            model_string: The model specification string (e.g. "gpt-4.1")
             request_params: Optional parameters to configure LLM behavior
 
         Returns:

mcp_agent/llm/provider_key_manager.py

@@ -14,6 +14,7 @@ PROVIDER_ENVIRONMENT_MAP: Dict[str, str] = {
     "anthropic": "ANTHROPIC_API_KEY",
     "openai": "OPENAI_API_KEY",
     "deepseek": "DEEPSEEK_API_KEY",
+    "google": "GOOGLE_API_KEY",
     "openrouter": "OPENROUTER_API_KEY",
     "generic": "GENERIC_API_KEY",
 }

mcp_agent/llm/provider_types.py

@@ -11,6 +11,7 @@ class Provider(Enum):
     ANTHROPIC = "anthropic"
     OPENAI = "openai"
     FAST_AGENT = "fast-agent"
+    GOOGLE = "google"
     DEEPSEEK = "deepseek"
     GENERIC = "generic"
-    OPENROUTER = "openrouter"
+    OPENROUTER = "openrouter"

mcp_agent/llm/providers/augmented_llm_anthropic.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Tuple, Type
 
 from mcp.types import EmbeddedResource, ImageContent, TextContent
 
@@ -10,6 +10,7 @@ from mcp_agent.llm.providers.multipart_converter_anthropic import (
 from mcp_agent.llm.providers.sampling_converter_anthropic import (
     AnthropicSamplingConverter,
 )
+from mcp_agent.mcp.interfaces import ModelT
 from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
 
 if TYPE_CHECKING:
@@ -50,6 +51,19 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
     selecting appropriate tools, and determining what information to retain.
     """
 
+    # Anthropic-specific parameter exclusions
+    ANTHROPIC_EXCLUDE_FIELDS = {
+        AugmentedLLM.PARAM_MESSAGES,
+        AugmentedLLM.PARAM_MODEL,
+        AugmentedLLM.PARAM_SYSTEM_PROMPT,
+        AugmentedLLM.PARAM_STOP_SEQUENCES,
+        AugmentedLLM.PARAM_MAX_TOKENS,
+        AugmentedLLM.PARAM_METADATA,
+        AugmentedLLM.PARAM_USE_HISTORY,
+        AugmentedLLM.PARAM_MAX_ITERATIONS,
+        AugmentedLLM.PARAM_PARALLEL_TOOL_CALLS,
+    }
+
     def __init__(self, *args, **kwargs) -> None:
         # Initialize logger - keep it simple without name reference
         self.logger = get_logger(__name__)
@@ -73,7 +87,7 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
         assert self.context.config
         return self.context.config.anthropic.base_url if self.context.config.anthropic else None
 
-    async def generate_internal(
+    async def _anthropic_completion(
         self,
         message_param,
         request_params: RequestParams | None = None,
@@ -100,7 +114,7 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
 
         # Always include prompt messages, but only include conversation history
         # if use_history is True
-        messages.extend(self.history.get(include_history=params.use_history))
+        messages.extend(self.history.get(include_completion_history=params.use_history))
 
         messages.append(message_param)
 
@@ -120,7 +134,8 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
 
         for i in range(params.max_iterations):
             self._log_chat_progress(self.chat_turn(), model=model)
-            arguments = {
+            # Create base arguments dictionary
+            base_args = {
                 "model": model,
                 "messages": messages,
                 "system": self.instruction or params.systemPrompt,
@@ -129,10 +144,12 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
             }
 
             if params.maxTokens is not None:
-                arguments["max_tokens"] = params.maxTokens
+                base_args["max_tokens"] = params.maxTokens
 
-            if params.metadata:
-                arguments = {**arguments, **params.metadata}
+            # Use the base class method to prepare all arguments with Anthropic-specific exclusions
+            arguments = self.prepare_provider_arguments(
+                base_args, params, self.ANTHROPIC_EXCLUDE_FIELDS
+            )
 
             self.logger.debug(f"{arguments}")
 
@@ -265,7 +282,7 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
         # Keep the prompt messages separate
         if params.use_history:
             # Get current prompt messages
-            prompt_messages = self.history.get(include_history=False)
+            prompt_messages = self.history.get(include_completion_history=False)
 
             # Calculate new conversation messages (excluding prompts)
             new_messages = messages[len(prompt_messages) :]
@@ -288,7 +305,7 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
         Override this method to use a different LLM.
 
         """
-        res = await self.generate_internal(
+        res = await self._anthropic_completion(
             message_param=message_param,
             request_params=request_params,
         )
@@ -298,6 +315,7 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
         self,
         multipart_messages: List["PromptMessageMultipart"],
         request_params: RequestParams | None = None,
+        is_template: bool = False,
     ) -> PromptMessageMultipart:
         # Check the last message role
         last_message = multipart_messages[-1]
@@ -310,7 +328,7 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
         for msg in messages_to_add:
             converted.append(AnthropicConverter.convert_to_anthropic(msg))
 
-        self.history.extend(converted, is_prompt=True)
+        self.history.extend(converted, is_prompt=is_template)
 
         if last_message.role == "user":
             self.logger.debug("Last message in prompt is from user, generating assistant response")
@@ -321,6 +339,28 @@ class AnthropicAugmentedLLM(AugmentedLLM[MessageParam, Message]):
             self.logger.debug("Last message in prompt is from assistant, returning it directly")
             return last_message
 
+    async def _apply_prompt_provider_specific_structured(
+        self,
+        multipart_messages: List[PromptMessageMultipart],
+        model: Type[ModelT],
+        request_params: RequestParams | None = None,
+    ) -> Tuple[ModelT | None, PromptMessageMultipart]:  # noqa: F821
+        request_params = self.get_request_params(request_params)
+
+        # TODO - convert this to use Tool Calling convention for Anthropic Structured outputs
+        multipart_messages[-1].add_text(
+            """YOU MUST RESPOND IN THE FOLLOWING FORMAT:
+            {schema}
+            RESPOND ONLY WITH THE JSON, NO PREAMBLE, CODE FENCES OR 'properties' ARE PERMISSABLE """.format(
+                schema=model.model_json_schema()
+            )
+        )
+
+        result: PromptMessageMultipart = await self._apply_prompt_provider_specific(
+            multipart_messages, request_params
+        )
+        return self._structured_from_multipart(result, model)
+
     @classmethod
     def convert_message_to_message_param(cls, message: Message, **kwargs) -> MessageParam:
         """Convert a response object to an input parameter object to allow LLM calls to be chained."""

mcp_agent/llm/providers/augmented_llm_deepseek.py

@@ -1,4 +1,3 @@
-
 from mcp_agent.core.request_params import RequestParams
 from mcp_agent.llm.provider_types import Provider
 from mcp_agent.llm.providers.augmented_llm_openai import OpenAIAugmentedLLM
@@ -9,10 +8,7 @@ DEFAULT_DEEPSEEK_MODEL = "deepseekchat"  # current Deepseek only has two type mo
 
 class DeepSeekAugmentedLLM(OpenAIAugmentedLLM):
     def __init__(self, *args, **kwargs) -> None:
-        kwargs["provider_name"] = "Deepseek"  # Set provider name in kwargs
-        super().__init__(
-            *args, provider=Provider.DEEPSEEK, **kwargs
-        )  # Properly pass args and kwargs to parent
+        super().__init__(*args, provider=Provider.DEEPSEEK, **kwargs)
 
     def _initialize_default_params(self, kwargs: dict) -> RequestParams:
         """Initialize Deepseek-specific default parameters"""

mcp_agent/llm/providers/augmented_llm_google.py

@@ -0,0 +1,30 @@
+from mcp_agent.core.request_params import RequestParams
+from mcp_agent.llm.provider_types import Provider
+from mcp_agent.llm.providers.augmented_llm_openai import OpenAIAugmentedLLM
+
+GOOGLE_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai"
+DEFAULT_GOOGLE_MODEL = "gemini-2.0-flash"
+
+
+class GoogleAugmentedLLM(OpenAIAugmentedLLM):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, provider=Provider.GOOGLE, **kwargs)
+
+    def _initialize_default_params(self, kwargs: dict) -> RequestParams:
+        """Initialize Google OpenAI Compatibility default parameters"""
+        chosen_model = kwargs.get("model", DEFAULT_GOOGLE_MODEL)
+
+        return RequestParams(
+            model=chosen_model,
+            systemPrompt=self.instruction,
+            parallel_tool_calls=False,
+            max_iterations=10,
+            use_history=True,
+        )
+
+    def _base_url(self) -> str:
+        base_url = None
+        if self.context.config and self.context.config.google:
+            base_url = self.context.config.google.base_url
+
+        return base_url if base_url else GOOGLE_BASE_URL