Refactored AnthropicAugmentedLLM and GoogleNativeAugmentedLLM

pyproject.toml

@@ -57,6 +57,7 @@ dev = [
     "ruff>=0.8.4",
     "tomli>=2.2.1",
     "pytest>=7.4.0",
+    "pytest-mock",
     "pytest-asyncio>=0.21.1",
     "pytest-cov",
 ]
@@ -107,6 +108,7 @@ dev = [
     "pytest>=7.4.0",
     "pytest-asyncio>=0.21.1",
     "pytest-cov>=6.1.1",
+    "pytest-mock",
     "ipdb>=0.13.13",
 ]
 

src/mcp_agent/llm/augmented_llm.py

@@ -85,6 +85,14 @@ def deep_merge(dict1: Dict[Any, Any], dict2: Dict[Any, Any]) -> Dict[Any, Any]:
 
 
 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.
+    Our current models can actively use these capabilities—generating their own search queries,
+    selecting appropriate tools, and determining what information to retain.
+    """
+
+
     # Common parameter names used across providers
     PARAM_MESSAGES = "messages"
     PARAM_MODEL = "model"
@@ -100,12 +108,10 @@ class AugmentedLLM(ContextDependent, AugmentedLLMProtocol, Generic[MessageParamT
     # 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.
-    Our current models can actively use these capabilities—generating their own search queries,
-    selecting appropriate tools, and determining what information to retain.
-    """
+    class _Actions:
+        STOP = "Stop"  # Making the actions available like so: self.ACTIONS.STOP
+        CONTINUE_WITH_TOOLS = "CONTINUE_WITH_TOOLS"
+    ACTIONS = _Actions()
 
     provider: Provider | None = None
 

src/mcp_agent/llm/providers/augmented_llm_azure.py

@@ -99,7 +99,7 @@ def _api_key(self):
             return "AzureCredential"
         return super()._api_key()
 
-    def _openai_client(self) -> AsyncOpenAI:
+    def _initialize_client(self) -> AsyncOpenAI:
         """
         Returns an AzureOpenAI client, handling both API Key and DefaultAzureCredential.
         """