solace-agent-mesh 1.5.1__py3-none-any.whl → 1.6.1__py3-none-any.whl

solace_agent_mesh/agent/adk/callbacks.py

@@ -1094,11 +1094,6 @@ If a plan is created:
             e_last_call,
         )
 
-    if host_component.get_config("inject_current_time", True):
-        current_time = datetime.now(timezone.utc).strftime("%A, %d %b %Y %H:%M:%S UTC")
-        instruction = f"Current time {current_time}."
-        injected_instructions.append(instruction)
-
     if injected_instructions:
         combined_instructions = "\n\n---\n\n".join(injected_instructions)
         if llm_request.config is None:

solace_agent_mesh/agent/adk/models/lite_llm.py

@@ -53,6 +53,7 @@ from typing_extensions import override
 from google.adk.models.base_llm import BaseLlm
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
+from .oauth2_token_manager import OAuth2ClientCredentialsTokenManager
 
 logger = logging.getLogger("google_adk." + __name__)
 
@@ -479,6 +480,7 @@ def _message_to_generate_content_response(
 
 def _get_completion_inputs(
     llm_request: LlmRequest,
+    cache_strategy: str = "5m",
 ) -> Tuple[
     List[Message],
     Optional[List[Dict]],
@@ -489,6 +491,7 @@ def _get_completion_inputs(
 
     Args:
       llm_request: The LlmRequest to convert.
+      cache_strategy: Cache strategy to apply ("none", "5m", "1h").
 
     Returns:
       The litellm inputs (message list, tool dictionary and response format).
@@ -501,16 +504,32 @@ def _get_completion_inputs(
         elif message_param_or_list:  # Ensure it's not None before appending
             messages.append(message_param_or_list)
 
-    if llm_request.config.system_instruction:
+    if llm_request.config and llm_request.config.system_instruction:
+        # Build system instruction content with optional cache control
+        system_content = {
+            "type": "text",
+            "text": llm_request.config.system_instruction,
+        }
+
+        # Add cache control based on strategy
+        # LiteLLM translates this to provider-specific format (Anthropic, OpenAI, Bedrock, Deepseek)
+        if cache_strategy == "5m":
+            # 5-minute ephemeral cache (Anthropic default)
+            system_content["cache_control"] = {"type": "ephemeral"}
+        elif cache_strategy == "1h":
+            # 1-hour extended cache (Anthropic extended)
+            system_content["cache_control"] = {"type": "ephemeral", "ttl": "1h"}
+        # For "none", no cache_control is added
+
         messages.insert(
             0,
             ChatCompletionDeveloperMessage(
                 role="developer",
-                content=llm_request.config.system_instruction,
+                content=[system_content],
             ),
         )
 
-    # 2. Convert tool declarations
+    # 2. Convert tool declarations with caching support
     tools: Optional[List[Dict]] = None
     if (
         llm_request.config
@@ -522,6 +541,16 @@ def _get_completion_inputs(
             for tool in llm_request.config.tools[0].function_declarations
         ]
 
+        # Enable tool caching via LiteLLM's generic interface
+        # LiteLLM handles provider-specific translation (Anthropic, OpenAI, Bedrock, Deepseek)
+        # Tools are stable because peer agents are alphabetically sorted (component.py)
+        if tools and cache_strategy != "none":
+            # Add cache_control to the LAST tool (required by caching providers)
+            if cache_strategy == "5m":
+                tools[-1]["cache_control"] = {"type": "ephemeral"}
+            elif cache_strategy == "1h":
+                tools[-1]["cache_control"] = {"type": "ephemeral", "ttl": "1h"}
+
     # 3. Handle response format
     response_format: Optional[types.SchemaUnion] = None
     if llm_request.config and llm_request.config.response_schema:
@@ -595,7 +624,7 @@ def _build_request_log(req: LlmRequest) -> str:
 
     function_decls: list[types.FunctionDeclaration] = cast(
         list[types.FunctionDeclaration],
-        req.config.tools[0].function_declarations if req.config.tools else [],
+        req.config.tools[0].function_declarations if req.config and req.config.tools else [],
     )
     function_logs = (
         [_build_function_declaration_log(func_decl) for func_decl in function_decls]
@@ -616,7 +645,7 @@ def _build_request_log(req: LlmRequest) -> str:
 LLM Request:
 -----------------------------------------------------------
 System Instruction:
-{req.config.system_instruction}
+{req.config.system_instruction if req.config else None}
 -----------------------------------------------------------
 Contents:
 {_NEW_LINE.join(contents_logs)}
@@ -654,16 +683,42 @@ class LiteLlm(BaseLlm):
     """The LLM client to use for the model."""
 
     _additional_args: Dict[str, Any] = None
+    _oauth_token_manager: Optional[OAuth2ClientCredentialsTokenManager] = None
+    _cache_strategy: str = "5m"  # Default to 5-minute ephemeral cache
 
-    def __init__(self, model: str, **kwargs):
+    def __init__(self, model: str, cache_strategy: str = "5m", **kwargs):
         """Initializes the LiteLlm class.
 
         Args:
           model: The name of the LiteLlm model.
+          cache_strategy: Cache strategy to use. Options: "none", "5m" (ephemeral), "1h" (extended).
+                         Defaults to "5m" for backward compatibility.
           **kwargs: Additional arguments to pass to the litellm completion api.
+                   Can include OAuth configuration parameters.
         """
         super().__init__(model=model, **kwargs)
-        self._additional_args = kwargs
+        self._additional_args = kwargs.copy()
+
+        # Validate and store cache strategy
+        valid_strategies = ["none", "5m", "1h"]
+        if cache_strategy not in valid_strategies:
+            logger.warning(
+                "Invalid cache_strategy '%s'. Valid options are: %s. Defaulting to '5m'.",
+                cache_strategy,
+                valid_strategies,
+            )
+            cache_strategy = "5m"
+        self._cache_strategy = cache_strategy
+        logger.info("LiteLlm initialized with cache strategy: %s", self._cache_strategy)
+
+        # Extract OAuth configuration if present
+        oauth_config = self._extract_oauth_config(self._additional_args)
+        if oauth_config:
+            self._oauth_token_manager = OAuth2ClientCredentialsTokenManager(**oauth_config)
+            logger.info("OAuth2 token manager initialized for model: %s", model)
+        else:
+            self._oauth_token_manager = None
+
         # preventing generation call with llm_client
         # and overriding messages, tools and stream which are managed internally
         self._additional_args.pop("llm_client", None)
@@ -672,6 +727,48 @@ class LiteLlm(BaseLlm):
         # public api called from runner determines to stream or not
         self._additional_args.pop("stream", None)
 
+    def _extract_oauth_config(self, kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Extract OAuth configuration from kwargs.
+
+        Args:
+            kwargs: Keyword arguments that may contain OAuth parameters
+
+        Returns:
+            OAuth configuration dictionary or None if no OAuth config found
+        """
+        oauth_params = [
+            "oauth_token_url",
+            "oauth_client_id",
+            "oauth_client_secret",
+            "oauth_scope",
+            "oauth_ca_cert",
+            "oauth_token_refresh_buffer_seconds",
+            "oauth_max_retries"
+        ]
+
+        oauth_config = {}
+        for param in oauth_params:
+            if param in kwargs:
+                # Map parameter names to OAuth2ClientCredentialsTokenManager constructor
+                if param == "oauth_ca_cert":
+                    oauth_config["ca_cert_path"] = kwargs.pop(param)
+                elif param == "oauth_token_refresh_buffer_seconds":
+                    oauth_config["refresh_buffer_seconds"] = kwargs.pop(param)
+                elif param == "oauth_max_retries":
+                    oauth_config["max_retries"] = kwargs.pop(param)
+                else:
+                    # Remove oauth_ prefix for the token manager
+                    key = param.replace("oauth_", "")
+                    oauth_config[key] = kwargs.pop(param)
+
+        # Return config only if we have the required parameters
+        if "token_url" in oauth_config and "client_id" in oauth_config and "client_secret" in oauth_config:
+            return oauth_config
+        elif oauth_config:
+            logger.warning("Incomplete OAuth configuration found, missing required parameters")
+
+        return None
+
     async def generate_content_async(
         self, llm_request: LlmRequest, stream: bool = False
     ) -> AsyncGenerator[LlmResponse, None]:
@@ -693,7 +790,7 @@ class LiteLlm(BaseLlm):
         logger.debug(_build_request_log(llm_request))
 
         messages, tools, response_format, generation_params = _get_completion_inputs(
-            llm_request
+            llm_request, self._cache_strategy
         )
         completion_args = {
             "model": self.model,
@@ -704,6 +801,24 @@ class LiteLlm(BaseLlm):
         }
         completion_args.update(self._additional_args)
 
+        # Inject OAuth token if OAuth is configured
+        if self._oauth_token_manager:
+            try:
+                access_token = await self._oauth_token_manager.get_token()
+                # Inject Bearer token via extra_headers
+                extra_headers = completion_args.get("extra_headers", {})
+                extra_headers["Authorization"] = f"Bearer {access_token}"
+                completion_args["extra_headers"] = extra_headers
+                logger.debug("OAuth token injected into request headers")
+            except Exception as e:
+                logger.error("Failed to get OAuth token: %s", str(e))
+                # Check if we have a fallback API key
+                if "api_key" in completion_args:
+                    logger.info("Falling back to API key authentication")
+                else:
+                    logger.error("No fallback authentication available")
+                    raise
+
         if generation_params:
             completion_args.update(generation_params)
 

solace_agent_mesh/agent/adk/models/oauth2_token_manager.py

@@ -0,0 +1,245 @@
+"""OAuth 2.0 Client Credentials Token Manager.
+
+This module provides OAuth 2.0 Client Credentials flow implementation for LLM authentication.
+It handles token acquisition, caching, and automatic refresh with proper error handling.
+"""
+
+import asyncio
+import logging
+import random
+import time
+from typing import Any, Dict, Optional
+
+import httpx
+
+from solace_agent_mesh.common.utils.in_memory_cache import InMemoryCache
+
+logger = logging.getLogger(__name__)
+
+
+class OAuth2ClientCredentialsTokenManager:
+    """Manages OAuth 2.0 Client Credentials tokens with caching and automatic refresh.
+    
+    This class implements the OAuth 2.0 Client Credentials flow as defined in RFC 6749.
+    It provides thread-safe token management with automatic refresh before expiration
+    and integrates with the existing InMemoryCache for token storage.
+    
+    Attributes:
+        token_url: OAuth 2.0 token endpoint URL
+        client_id: OAuth client identifier
+        client_secret: OAuth client secret
+        scope: OAuth scope (optional)
+        ca_cert_path: Path to custom CA certificate (optional)
+        refresh_buffer_seconds: Seconds before expiry to refresh token
+    """
+
+    def __init__(
+        self,
+        token_url: str,
+        client_id: str,
+        client_secret: str,
+        scope: Optional[str] = None,
+        ca_cert_path: Optional[str] = None,
+        refresh_buffer_seconds: int = 300,
+        max_retries: int = 3,
+    ):
+        """Initialize the OAuth2 Client Credentials Token Manager.
+        
+        Args:
+            token_url: OAuth 2.0 token endpoint URL
+            client_id: OAuth client identifier
+            client_secret: OAuth client secret
+            scope: OAuth scope (optional, space-separated string)
+            ca_cert_path: Path to custom CA certificate file (optional)
+            refresh_buffer_seconds: Seconds before actual expiry to refresh token
+            max_retries: Maximum number of retry attempts for token requests
+            
+        Raises:
+            ValueError: If required parameters are missing or invalid
+        """
+        if not token_url:
+            raise ValueError("token_url is required")
+        if not client_id:
+            raise ValueError("client_id is required")
+        if not client_secret:
+            raise ValueError("client_secret is required")
+        if refresh_buffer_seconds < 0:
+            raise ValueError("refresh_buffer_seconds must be non-negative")
+            
+        self.token_url = token_url
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.scope = scope
+        self.ca_cert_path = ca_cert_path
+        self.refresh_buffer_seconds = refresh_buffer_seconds
+        self.max_retries = max_retries
+        
+        # Thread-safe token access
+        self._lock = asyncio.Lock()
+        
+        # Token cache using existing InMemoryCache singleton
+        self._cache = InMemoryCache()
+        
+        # Cache key for this token manager instance
+        self._cache_key = f"oauth_token_{hash((token_url, client_id))}"
+        
+        logger.info(
+            "OAuth2ClientCredentialsTokenManager initialized for endpoint: %s",
+            token_url
+        )
+
+    async def get_token(self) -> str:
+        """Get a valid OAuth 2.0 access token.
+        
+        This method checks the cache first and returns a cached token if it's still valid.
+        If no token exists or the token is expired/near expiry, it fetches a new token.
+        
+        Returns:
+            Valid OAuth 2.0 access token
+            
+        Raises:
+            httpx.HTTPError: If token request fails
+            ValueError: If token response is invalid
+        """
+        async with self._lock:
+            # Check if we have a cached token
+            cached_token_data = self._cache.get(self._cache_key)
+            
+            if cached_token_data and not self._is_token_expired(cached_token_data):
+                logger.debug("Using cached OAuth token")
+                return cached_token_data["access_token"]
+            
+            # Fetch new token
+            logger.info("Fetching new OAuth token from %s", self.token_url)
+            token_data = await self._fetch_token()
+            
+            # Cache the token with TTL
+            expires_in = token_data.get("expires_in", 3600)  # Default 1 hour
+            cache_ttl = max(expires_in - self.refresh_buffer_seconds, 60)  # Min 1 minute
+            
+            self._cache.set(self._cache_key, token_data, ttl=cache_ttl)
+            
+            logger.info("OAuth token cached with TTL: %d seconds", cache_ttl)
+            return token_data["access_token"]
+
+    def _is_token_expired(self, token_data: Dict[str, Any]) -> bool:
+        """Check if a token is expired or near expiry.
+        
+        Args:
+            token_data: Token data dictionary with 'expires_at' timestamp
+            
+        Returns:
+            True if token is expired or near expiry, False otherwise
+        """
+        if "expires_at" not in token_data:
+            return True
+            
+        current_time = time.time()
+        expires_at = token_data["expires_at"]
+        
+        # Consider token expired if it expires within the buffer time
+        return current_time >= (expires_at - self.refresh_buffer_seconds)
+
+    async def _fetch_token(self) -> Dict[str, Any]:
+        """Fetch a new OAuth 2.0 access token from the token endpoint.
+
+        Implements retry logic with exponential backoff for transient failures.
+
+        Returns:
+            Token data dictionary containing access_token, expires_in, etc.
+
+        Raises:
+            httpx.HTTPError: If HTTP request fails after all retries
+            ValueError: If response is invalid or missing required fields
+        """
+        # Prepare request payload
+        payload = {
+            "grant_type": "client_credentials",
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+        }
+        
+        if self.scope:
+            payload["scope"] = self.scope
+        
+        # Configure HTTP client with SSL settings
+        verify = True
+        if self.ca_cert_path:
+            verify = self.ca_cert_path
+            
+        headers = {
+            "Content-Type": "application/x-www-form-urlencoded",
+            "Accept": "application/json",
+        }
+        
+        last_exception = None
+
+        for attempt in range(self.max_retries + 1):
+            try:
+                async with httpx.AsyncClient(verify=verify) as client:
+                    response = await client.post(
+                        self.token_url,
+                        data=payload,
+                        headers=headers,
+                        timeout=30.0,
+                    )
+                    response.raise_for_status()
+
+                    token_data = response.json()
+
+                    # Validate response
+                    if "access_token" not in token_data:
+                        raise ValueError("Token response missing 'access_token' field")
+
+                    # Add expiration timestamp for cache management
+                    expires_in = token_data.get("expires_in", 3600)
+                    token_data["expires_at"] = time.time() + expires_in
+
+                    logger.info("Successfully fetched OAuth token, expires in %d seconds", expires_in)
+                    return token_data
+
+            except httpx.HTTPStatusError as e:
+                last_exception = e
+                # Don't retry on 4xx errors (client errors)
+                if 400 <= e.response.status_code < 500:
+                    logger.error(
+                        "OAuth token request failed with client error %d: %s",
+                        e.response.status_code,
+                        e.response.text
+                    )
+                    raise
+
+                logger.warning(
+                    "OAuth token request failed with status %d (attempt %d/%d): %s",
+                    e.response.status_code,
+                    attempt + 1,
+                    self.max_retries + 1,
+                    e.response.text
+                )
+
+            except httpx.RequestError as e:
+                last_exception = e
+                logger.warning(
+                    "OAuth token request failed (attempt %d/%d): %s",
+                    attempt + 1,
+                    self.max_retries + 1,
+                    str(e)
+                )
+
+            except Exception as e:
+                last_exception = e
+                logger.error("Unexpected error during OAuth token fetch: %s", str(e))
+                raise
+
+            # Exponential backoff with jitter for retries
+            if attempt < self.max_retries:
+                delay = (2 ** attempt) + random.uniform(0, 1)
+                logger.info("Retrying OAuth token request in %.2f seconds", delay)
+                await asyncio.sleep(delay)
+
+        # All retries exhausted
+        logger.error("OAuth token request failed after %d attempts", self.max_retries + 1)
+        if last_exception:
+            raise last_exception
+        else:
+            raise RuntimeError("OAuth token request failed after all retries")