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

solace_agent_mesh/agent/adk/callbacks.py

@@ -1405,6 +1405,12 @@ def solace_llm_invocation_callback(
         logical_task_id = a2a_context.get("logical_task_id")
         context_id = a2a_context.get("contextId")
 
+        # Store model name in callback state for later use in response callback
+        model_name = host_component.model_config
+        if isinstance(model_name, dict):
+            model_name = model_name.get("model", "unknown")
+        callback_context.state["model_name"] = model_name
+
         llm_data = LlmInvocationData(request=llm_request.model_dump(exclude_none=True))
         status_update_event = a2a.create_data_signal_event(
             task_id=logical_task_id,
@@ -1447,8 +1453,8 @@ def solace_llm_response_callback(
     host_component: "SamAgentComponent",
 ) -> Optional[LlmResponse]:
     """
-    ADK after_model_callback to send a Solace message with the LLM's response,
-    using the host_component's process_and_publish_adk_event method.
+    ADK after_model_callback to send a Solace message with the LLM's response
+    and token usage information.
     """
     log_identifier = "[Callback:SolaceLLMResponse]"
     if llm_response.partial:  # Don't send partial responses for this notification
@@ -1478,6 +1484,49 @@ 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 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,
+                    output_tokens=usage.candidates_token_count,
+                    model=model_name,
+                    source="agent",
+                    cached_input_tokens=cached_tokens,
+                )
+                log.debug(
+                    "%s Recorded token usage: input=%d, output=%d, cached=%d, model=%s",
+                    log_identifier,
+                    usage.prompt_token_count,
+                    usage.candidates_token_count,
+                    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/models/lite_llm.py

@@ -700,6 +700,7 @@ class LiteLlm(BaseLlm):
             "messages": messages,
             "tools": tools,
             "response_format": response_format,
+            "stream_options": {"include_usage": True},
         }
         completion_args.update(self._additional_args)
 

solace_agent_mesh/agent/adk/models/models_llm.txt

@@ -5,7 +5,6 @@ This directory contains concrete implementations of the `BaseLlm` interface, pro
 
 ## Files Overview
 - `lite_llm.py` - LLM client using the `litellm` library to support hundreds of models from different providers
-- `models_llm.txt` - Documentation file containing developer guide content
 
 ## Developer API Reference
 
@@ -140,4 +139,4 @@ if __name__ == "__main__":
     # asyncio.run(function_calling_example())
 ```
 
-# content_hash: 12789ad2e16cd9ea5a81abdd68258d9ef30520bed5c51ba8d00ea66014191964
+# content_hash: 97abaea0097e0d6fe18f1d747ed54ced092c0a96829dc5be8448b087ac5f7c20

solace_agent_mesh/agent/adk/services.py

@@ -271,28 +271,43 @@ def initialize_artifact_service(component) -> BaseArtifactService:
         try:
             from .artifacts.s3_artifact_service import S3ArtifactService
 
+            # Whitelist of valid parameters for the boto3 S3 client.
             valid_boto3_params = [
                 "aws_access_key_id",
                 "aws_secret_access_key",
                 "aws_session_token",
                 "region_name",
-                "config"
+                "endpoint_url",
+                "config",
             ]
 
-            s3_config = {k: v for k, v in config.items() if k in valid_boto3_params}
-
-            endpoint_url = config.get("endpoint_url") or os.environ.get("S3_ENDPOINT_URL") or "https://s3.amazonaws.com"
-            s3_config["endpoint_url"] = endpoint_url
-
-            aws_access_key_id = config.get("aws_access_key_id") or os.environ.get("AWS_ACCESS_KEY_ID")
-            aws_secret_access_key = config.get("aws_secret_access_key") or os.environ.get("AWS_SECRET_ACCESS_KEY")
-
-            if aws_access_key_id:
-                s3_config["aws_access_key_id"] = aws_access_key_id
-            if aws_secret_access_key:
-                s3_config["aws_secret_access_key"] = aws_secret_access_key
-
-            concrete_service = S3ArtifactService(bucket_name=bucket_name, **s3_config)
+            s3_config = {}
+
+            # Explicitly map the 'region' from our config to 'region_name' for boto3.
+            if config.get("region"):
+                s3_config["region_name"] = config.get("region")
+
+            # Copy any other valid parameters from the config.
+            for key in valid_boto3_params:
+                if key in config and config[key] is not None:
+                    s3_config[key] = config[key]
+
+            # Set credentials from environment variables as a fallback.
+            if "endpoint_url" not in s3_config:
+                s3_config["endpoint_url"] = os.environ.get("S3_ENDPOINT_URL")
+            if "aws_access_key_id" not in s3_config:
+                env_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
+                if env_access_key is not None:
+                    s3_config["aws_access_key_id"] = env_access_key
+            if "aws_secret_access_key" not in s3_config:
+                env_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
+                if env_secret_key is not None:
+                    s3_config["aws_secret_access_key"] = env_secret_key
+
+            # Filter out any keys that ended up with a None value.
+            s3_config_cleaned = {k: v for k, v in s3_config.items() if v is not None}
+
+            concrete_service = S3ArtifactService(bucket_name=bucket_name, **s3_config_cleaned)
         except ImportError as e:
             log.error(
                 "%s S3 dependencies not available: %s",

solace_agent_mesh/agent/adk/setup.py

@@ -22,6 +22,8 @@ from google.adk.models.llm_response import LlmResponse
 from google.adk.tools.mcp_tool.mcp_session_manager import (
     SseServerParams,
     StdioConnectionParams,
+    StreamableHTTPServerParams,
+
 )
 
 from mcp import StdioServerParameters
@@ -566,6 +568,8 @@ async def _load_mcp_tool(component: "SamAgentComponent", tool_config: Dict) -> T
 
     elif connection_type == "sse":
         connection_params = SseServerParams(**connection_args)
+    elif connection_type == "streamable-http":
+        connection_params = StreamableHTTPServerParams(**connection_args)
     else:
         raise ValueError(f"Unsupported MCP connection type: {connection_type}")
 

solace_agent_mesh/agent/agent_llm.txt

@@ -375,4 +375,4 @@ async def my_business_tool(
     # Process business logic
     result =
 
-# content_hash: 7e177b117b6361dbdd3eaeaa246c2173ac9cfaada01eb78219f2b8cd8b13d816
+# content_hash: 596f8253b415eaca3335f7660837dee9a090b3cab0b6e0a4daf92ca9a16d0123