solace-agent-mesh 1.5.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

solace_agent_mesh/agent/adk/callbacks.py

@@ -4,6 +4,7 @@ Includes dynamic instruction injection, artifact metadata injection,
 embed resolution, and logging.
 """
 
+import logging
 import json
 import asyncio
 import uuid
@@ -18,7 +19,7 @@ from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
 from google.genai import types as adk_types
 from google.adk.tools.mcp_tool import MCPTool
-from solace_ai_connector.common.log import log
+
 from .intelligent_mcp_callbacks import (
     save_mcp_response_as_artifact_intelligent,
     McpSaveStatus,
@@ -76,6 +77,8 @@ from ...agent.adk.stream_parser import (
     ARTIFACT_BLOCK_DELIMITER_CLOSE,
 )
 
+log = logging.getLogger(__name__)
+
 A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY = "temp:llm_stream_chunks_processed"
 
 if TYPE_CHECKING:
@@ -775,17 +778,16 @@ def _generate_fenced_artifact_instruction() -> str:
 To create an artifact from content you generate (like code, a report, or a document), you MUST use a special `save_artifact` block. This is the only reliable way to ensure your content is saved correctly.
 
 **Syntax:**
-```
 {open_delim}save_artifact: filename="your_filename.ext" mime_type="text/plain" description="A brief description."
 The full content you want to save goes here.
 It can span multiple lines.
 {close_delim}
-```
 
 - **Rules:**
   - The parameters `filename` and `mime_type` are required. `description` is optional but recommended.
   - All parameter values **MUST** be enclosed in double quotes.
   - You **MUST NOT** use double quotes `"` inside the parameter values (e.g., within the description string). Use single quotes or rephrase instead.
+  - Do not surround a save_artifact block with '```' (triple backticks). This will create rendering issues.
 
 The system will automatically save the content and give you a confirmation in the next turn."""
 
@@ -1092,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:
@@ -1484,32 +1481,32 @@ def solace_llm_response_callback(
             "type": "llm_response",
             "data": llm_response.model_dump(exclude_none=True),
         }
-        
+
         # Extract and record token usage
         if llm_response.usage_metadata:
             usage = llm_response.usage_metadata
             model_name = callback_context.state.get("model_name", "unknown")
-            
+
             usage_dict = {
                 "input_tokens": usage.prompt_token_count,
                 "output_tokens": usage.candidates_token_count,
                 "model": model_name,
             }
-            
+
             # Check for cached tokens (provider-specific)
             cached_tokens = 0
-            if hasattr(usage, 'prompt_tokens_details') and usage.prompt_tokens_details:
-                cached_tokens = getattr(usage.prompt_tokens_details, 'cached_tokens', 0)
+            if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
+                cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0)
                 if cached_tokens > 0:
                     usage_dict["cached_input_tokens"] = cached_tokens
-            
+
             # Add to response data
             llm_response_data["usage"] = usage_dict
-            
+
             # Record in task context for aggregation
             with host_component.active_tasks_lock:
                 task_context = host_component.active_tasks.get(logical_task_id)
-            
+
             if task_context:
                 task_context.record_token_usage(
                     input_tokens=usage.prompt_token_count,
@@ -1526,7 +1523,7 @@ def solace_llm_response_callback(
                     cached_tokens,
                     model_name,
                 )
-        
+
         # This signal doesn't have a dedicated Pydantic model, so we create the
         # DataPart directly and use the lower-level helpers.
         data_part = a2a.create_data_part(data=llm_response_data)

solace_agent_mesh/agent/adk/embed_resolving_mcp_toolset.py

@@ -2,6 +2,7 @@
 Custom MCPToolset that resolves embeds in tool parameters before calling MCP tools.
 """
 
+import logging
 import asyncio
 from typing import Dict, List, Optional, Any
 
@@ -12,7 +13,7 @@ from google.adk.tools.mcp_tool.mcp_session_manager import (
     StreamableHTTPConnectionParams,
 )
 from google.adk.tools.tool_context import ToolContext
-from solace_ai_connector.common.log import log
+
 
 from ..utils.context_helpers import get_original_session_id
 from ...common.utils.embeds import (
@@ -23,6 +24,7 @@ from ...common.utils.embeds import (
     EMBED_DELIMITER_OPEN,
 )
 
+log = logging.getLogger(__name__)
 
 class EmbedResolvingMCPTool(MCPTool):
     """

solace_agent_mesh/agent/adk/intelligent_mcp_callbacks.py

@@ -5,6 +5,7 @@ This module contains the refactored MCP callback functions that use intelligent
 content processing to save MCP tool responses as appropriately typed artifacts.
 """
 
+import logging
 import json
 import uuid
 from datetime import datetime, timezone
@@ -13,7 +14,6 @@ from enum import Enum
 from pydantic import BaseModel
 
 from google.adk.tools import ToolContext, BaseTool
-from solace_ai_connector.common.log import log
 
 from .mcp_content_processor import MCPContentProcessor, MCPContentProcessorConfig
 from ...agent.utils.artifact_helpers import (
@@ -22,6 +22,7 @@ from ...agent.utils.artifact_helpers import (
 )
 from ...agent.utils.context_helpers import get_original_session_id
 
+log = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
     from ...agent.sac.component import SamAgentComponent

solace_agent_mesh/agent/adk/mcp_content_processor.py

@@ -12,6 +12,7 @@ Supports:
 - Resource content with URI-based filename extraction
 """
 
+import logging
 import base64
 import csv
 import json
@@ -23,7 +24,7 @@ from io import StringIO
 from typing import Any, Dict, List, Optional, Tuple, Union
 from urllib.parse import urlparse
 
-from solace_ai_connector.common.log import log
+log = logging.getLogger(__name__)
 
 from ...common.utils.mime_helpers import (
     is_text_based_mime_type,

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")

solace_agent_mesh/agent/adk/runner.py

@@ -2,6 +2,7 @@
 Manages the asynchronous execution of the ADK Runner.
 """
 
+import logging
 import asyncio
 
 from google.adk.agents.invocation_context import LlmCallsLimitExceededError
@@ -21,10 +22,11 @@ from google.adk.events import Event as ADKEvent
 from google.adk.events.event_actions import EventActions
 from google.adk.sessions import Session as ADKSession
 from google.genai import types as adk_types
-from solace_ai_connector.common.log import log
 
 from ...common import a2a
 
+log = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from ..sac.component import SamAgentComponent
     from ..sac.task_execution_context import TaskExecutionContext

solace_agent_mesh/agent/adk/services.py

@@ -2,13 +2,14 @@
 Initializes ADK Services based on configuration.
 """
 
+import logging
 import os
 import re
 from typing import Dict, Optional, List, Any
 from typing_extensions import override
 
 from google.genai import types as adk_types
-from solace_ai_connector.common.log import log
+
 
 from google.adk.sessions import (
     BaseSessionService,
@@ -29,6 +30,8 @@ from google.adk.memory import (
 
 from .artifacts.filesystem_artifact_service import FilesystemArtifactService
 
+log = logging.getLogger(__name__)
+
 try:
     from sam_test_infrastructure.artifact_service.service import (
         TestInMemoryArtifactService,