solace-agent-mesh 1.4.12__py3-none-any.whl → 1.5.1__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."""
 
@@ -1405,6 +1407,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 +1455,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 +1486,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/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

@@ -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/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,

solace_agent_mesh/agent/adk/setup.py

@@ -2,13 +2,14 @@
 Handles ADK Agent and Runner initialization, including tool loading and callback assignment.
 """
 
+import logging
 from typing import Dict, List, Optional, Union, Callable, Tuple, Set, Any, TYPE_CHECKING, Type
 import functools
 import inspect
-from solace_ai_connector.common.log import log
 from solace_ai_connector.common.utils import import_module
 from ...common.utils.type_utils import is_subclass_by_name
 
+
 from .app_llm_agent import AppLlmAgent
 from .tool_wrapper import ADKToolWrapper
 from .embed_resolving_mcp_toolset import EmbedResolvingMCPToolset
@@ -46,6 +47,7 @@ from ..tools.tool_config_types import (
 from ...agent.adk import callbacks as adk_callbacks
 from ...agent.adk.models.lite_llm import LiteLlm
 
+log = logging.getLogger(__name__)
 
 # Define a clear return type for all tool-loading helpers
 ToolLoadingResult = Tuple[List[Union[BaseTool, Callable]], List[BuiltinTool], List[Callable]]

solace_agent_mesh/agent/adk/tool_wrapper.py

@@ -2,13 +2,12 @@
 Defines the ADKToolWrapper, a consolidated wrapper for ADK tools.
 """
 
+import logging
 import asyncio
 import functools
 import inspect
 from typing import Callable, Dict, List, Optional, Literal
 
-from solace_ai_connector.common.log import log
-
 from ...common.utils.embeds import (
     resolve_embeds_in_string,
     evaluate_embed,
@@ -17,6 +16,7 @@ from ...common.utils.embeds import (
     EMBED_DELIMITER_OPEN,
 )
 
+log = logging.getLogger(__name__)
 
 class ADKToolWrapper:
     """

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