solace-agent-mesh 1.0.9__py3-none-any.whl → 1.3.0__py3-none-any.whl

solace_agent_mesh/agent/adk/callbacks.py

@@ -16,9 +16,6 @@ from google.adk.artifacts import BaseArtifactService
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
-from ...common.a2a_protocol import (
-    A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY,
-)
 from google.genai import types as adk_types
 from google.adk.tools.mcp_tool import MCPTool
 from solace_ai_connector.common.log import log
@@ -48,14 +45,14 @@ from ...common.utils.embeds import (
 
 from ...common.utils.embeds.modifiers import MODIFIER_IMPLEMENTATIONS
 
-if TYPE_CHECKING:
-    from ..sac.component import SamAgentComponent
-
-from ...common.types import (
-    TaskStatusUpdateEvent,
-    TaskStatus,
-    TaskState,
-    Message as A2AMessage,
+from ...common import a2a
+from ...common.a2a.types import ContentPart
+from ...common.data_parts import (
+    AgentProgressUpdateData,
+    ArtifactCreationProgressData,
+    LlmInvocationData,
+    ToolInvocationStartData,
+    ToolResultData,
 )
 
 from ...agent.utils.artifact_helpers import (
@@ -68,6 +65,7 @@ from ..tools.builtin_artifact_tools import _internal_create_artifact
 from ...agent.adk.tool_wrapper import ADKToolWrapper
 
 # Import the new parser and its events
+from pydantic import BaseModel
 from ...agent.adk.stream_parser import (
     FencedBlockStreamParser,
     BlockStartedEvent,
@@ -78,6 +76,44 @@ from ...agent.adk.stream_parser import (
     ARTIFACT_BLOCK_DELIMITER_CLOSE,
 )
 
+A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY = "temp:llm_stream_chunks_processed"
+
+if TYPE_CHECKING:
+    from ..sac.component import SamAgentComponent
+
+
+async def _publish_data_part_status_update(
+    host_component: "SamAgentComponent",
+    a2a_context: Dict[str, Any],
+    data_part_model: BaseModel,
+):
+    """Helper to construct and publish a TaskStatusUpdateEvent with a DataPart."""
+    logical_task_id = a2a_context.get("logical_task_id")
+    context_id = a2a_context.get("contextId")
+
+    status_update_event = a2a.create_data_signal_event(
+        task_id=logical_task_id,
+        context_id=context_id,
+        signal_data=data_part_model,
+        agent_name=host_component.agent_name,
+    )
+
+    loop = host_component.get_async_loop()
+    if loop and loop.is_running():
+        asyncio.run_coroutine_threadsafe(
+            host_component._publish_status_update_with_buffer_flush(
+                status_update_event,
+                a2a_context,
+                skip_buffer_flush=False,
+            ),
+            loop,
+        )
+    else:
+        log.error(
+            "%s Async loop not available. Cannot publish status update.",
+            host_component.log_identifier,
+        )
+
 
 async def process_artifact_blocks_callback(
     callback_context: CallbackContext,
@@ -132,8 +168,11 @@ async def process_artifact_blocks_callback(
                         )
                         filename = event.params.get("filename", "unknown_artifact")
                         if a2a_context:
-                            await host_component._publish_agent_status_signal_update(
-                                f"Receiving artifact `{filename}`...", a2a_context
+                            progress_data = AgentProgressUpdateData(
+                                status_text=f"Receiving artifact `{filename}`..."
+                            )
+                            await _publish_data_part_status_update(
+                                host_component, a2a_context, progress_data
                             )
                         params_str = " ".join(
                             [f'{k}="{v}"' for k, v in event.params.items()]
@@ -149,10 +188,14 @@ async def process_artifact_blocks_callback(
                         )
                         params = parser._block_params
                         filename = params.get("filename", "unknown_artifact")
-                        status_message = f"Creating artifact `{filename}` ({event.buffered_size}B saved)..."
                         if a2a_context:
-                            await host_component._publish_agent_status_signal_update(
-                                status_message, a2a_context
+                            progress_data = ArtifactCreationProgressData(
+                                filename=filename,
+                                bytes_saved=event.buffered_size,
+                                artifact_chunk=event.chunk,
+                            )
+                            await _publish_data_part_status_update(
+                                host_component, a2a_context, progress_data
                             )
 
                     elif isinstance(event, BlockCompletedEvent):
@@ -217,6 +260,7 @@ async def process_artifact_blocks_callback(
                             tool_config=None,  # No specific config for this internal tool
                             tool_name="_internal_create_artifact",
                             origin="internal",
+                            resolution_type="early",
                         )
                         save_result = await wrapped_creator(**kwargs_for_call)
 
@@ -694,8 +738,7 @@ async def manage_large_mcp_tool_responses_callback(
 
     if needs_saving_as_artifact and (
         save_result
-        and save_result.status
-        in [McpSaveStatus.SUCCESS, McpSaveStatus.PARTIAL_SUCCESS]
+        and save_result.status in [McpSaveStatus.SUCCESS, McpSaveStatus.PARTIAL_SUCCESS]
     ):
         if needs_truncation_for_llm:
             final_llm_response_dict["status"] = "processed_saved_and_truncated"
@@ -1345,20 +1388,10 @@ def solace_llm_invocation_callback(
         )
         return None
 
-    try:
-        from ...common.a2a_protocol import A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY
-
-        callback_context.state[A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY] = False
-        log.debug(
-            "%s Reset %s to False.", log_identifier, A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY
-        )
-    except Exception as e_flag_reset:
-        log.error(
-            "%s Error resetting %s: %s",
-            log_identifier,
-            A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY,
-            e_flag_reset,
-        )
+    callback_context.state[A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY] = False
+    log.debug(
+        "%s Reset %s to False.", log_identifier, A2A_LLM_STREAM_CHUNKS_PROCESSED_KEY
+    )
 
     try:
         a2a_context = callback_context.state.get("a2a_context")
@@ -1369,28 +1402,15 @@ def solace_llm_invocation_callback(
             )
             return None
 
-        agent_name = host_component.get_config("agent_name", "unknown_agent")
         logical_task_id = a2a_context.get("logical_task_id")
-
-        llm_invocation_metadata = {
-            "type": "llm_invocation",
-            "data": llm_request.model_dump(exclude_none=True),
-        }
-        a2a_message = A2AMessage(
-            role="agent",
-            parts=[],
-            metadata=llm_invocation_metadata,
-        )
-        task_status = TaskStatus(
-            state=TaskState.WORKING,
-            message=a2a_message,
-            timestamp=datetime.now(timezone.utc),
-        )
-        status_update_event = TaskStatusUpdateEvent(
-            id=logical_task_id,
-            status=task_status,
-            final=False,
-            metadata={"agent_name": agent_name},
+        context_id = a2a_context.get("contextId")
+
+        llm_data = LlmInvocationData(request=llm_request.model_dump(exclude_none=True))
+        status_update_event = a2a.create_data_signal_event(
+            task_id=logical_task_id,
+            context_id=context_id,
+            signal_data=llm_data,
+            agent_name=host_component.agent_name,
         )
 
         loop = host_component.get_async_loop()
@@ -1454,20 +1474,23 @@ def solace_llm_response_callback(
         agent_name = host_component.get_config("agent_name", "unknown_agent")
         logical_task_id = a2a_context.get("logical_task_id")
 
-        llm_response_metadata = {
+        llm_response_data = {
             "type": "llm_response",
             "data": llm_response.model_dump(exclude_none=True),
         }
-        a2a_message = A2AMessage(role="agent", parts=[], metadata=llm_response_metadata)
-        task_status = TaskStatus(
-            state=TaskState.WORKING,
-            message=a2a_message,
-            timestamp=datetime.now(timezone.utc),
+        # 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)
+        a2a_message = a2a.create_agent_parts_message(
+            parts=[data_part],
+            task_id=logical_task_id,
+            context_id=a2a_context.get("contextId"),
         )
-        status_update_event = TaskStatusUpdateEvent(
-            id=logical_task_id,
-            status=task_status,
-            final=True,
+        status_update_event = a2a.create_status_update(
+            task_id=logical_task_id,
+            context_id=a2a_context.get("contextId"),
+            message=a2a_message,
+            is_final=False,
             metadata={"agent_name": agent_name},
         )
         loop = host_component.get_async_loop()
@@ -1529,8 +1552,6 @@ def notify_tool_invocation_start_callback(
         )
         return
 
-    logical_task_id = a2a_context.get("logical_task_id")
-
     try:
         serializable_args = {}
         for k, v in args.items():
@@ -1540,57 +1561,97 @@ def notify_tool_invocation_start_callback(
             except TypeError:
                 serializable_args[k] = str(v)
 
-        a2a_message_parts = []
-        message_metadata = {
-            "type": "tool_invocation_start",
-            "data": {
-                "tool_name": tool.name,
-                "tool_args": serializable_args,
-                "function_call_id": tool_context.function_call_id,
-            },
-        }
+        tool_data = ToolInvocationStartData(
+            tool_name=tool.name,
+            tool_args=serializable_args,
+            function_call_id=tool_context.function_call_id,
+        )
+        asyncio.run_coroutine_threadsafe(
+            _publish_data_part_status_update(host_component, a2a_context, tool_data),
+            host_component.get_async_loop(),
+        )
+        log.debug(
+            "%s Scheduled tool_invocation_start notification.",
+            log_identifier,
+        )
 
-        a2a_message = A2AMessage(
-            role="agent", parts=a2a_message_parts, metadata=message_metadata
+    except Exception as e:
+        log.exception(
+            "%s Error publishing tool_invocation_start status update: %s",
+            log_identifier,
+            e,
         )
 
-        task_status = TaskStatus(
-            state=TaskState.WORKING,
-            message=a2a_message,
-            timestamp=datetime.now(timezone.utc),
+    return None
+
+
+def notify_tool_execution_result_callback(
+    tool: BaseTool,
+    args: Dict[str, Any],
+    tool_context: ToolContext,
+    tool_response: Any,
+    host_component: "SamAgentComponent",
+) -> None:
+    """
+    ADK after_tool_callback to send an A2A status message with the result
+    of a tool's execution.
+    """
+    log_identifier = f"[Callback:NotifyToolResult:{tool.name}]"
+    log.debug("%s Triggered for tool '%s'", log_identifier, tool.name)
+
+    if not host_component:
+        log.error(
+            "%s Host component instance not provided. Cannot send notification.",
+            log_identifier,
         )
+        return
 
-        status_update_event = TaskStatusUpdateEvent(
-            id=logical_task_id,
-            status=task_status,
-            final=False,
-            metadata={"agent_name": host_component.get_config("agent_name")},
+    a2a_context = tool_context.state.get("a2a_context")
+    if not a2a_context:
+        log.error(
+            "%s a2a_context not found in tool_context.state. Cannot send notification.",
+            log_identifier,
         )
+        return
 
-        loop = host_component.get_async_loop()
-        if loop and loop.is_running():
+    if tool.is_long_running and not tool_response:
+        log.debug(
+            "%s Tool is long-running and is not yet complete. Don't notify its completion",
+            log_identifier,
+        )
+        return
 
-            asyncio.run_coroutine_threadsafe(
-                host_component._publish_status_update_with_buffer_flush(
-                    status_update_event,
-                    a2a_context,
-                    skip_buffer_flush=False,
-                ),
-                loop,
-            )
-            log.debug(
-                "%s Scheduled tool_invocation_start notification with buffer flush.",
-                log_identifier,
-            )
+    try:
+        # Attempt to make the response JSON serializable
+        serializable_response = tool_response
+        if hasattr(tool_response, "model_dump"):
+            serializable_response = tool_response.model_dump(exclude_none=True)
         else:
-            log.error(
-                "%s Async loop not available. Cannot publish tool_invocation_start notification.",
-                log_identifier,
-            )
+            try:
+                # A simple check to see if it can be dumped.
+                # This isn't perfect but catches many non-serializable types.
+                json.dumps(tool_response)
+            except (TypeError, OverflowError):
+                serializable_response = str(tool_response)
+
+        tool_data = ToolResultData(
+            tool_name=tool.name,
+            result_data=serializable_response,
+            function_call_id=tool_context.function_call_id,
+        )
+        asyncio.run_coroutine_threadsafe(
+            _publish_data_part_status_update(host_component, a2a_context, tool_data),
+            host_component.get_async_loop(),
+        )
+        log.debug(
+            "%s Scheduled tool_result notification for function call ID %s.",
+            log_identifier,
+            tool_context.function_call_id,
+        )
 
     except Exception as e:
         log.exception(
-            "%s Error publishing tool_invocation_start status update: %s",
+            "%s Error publishing tool_result status update: %s",
             log_identifier,
             e,
         )

solace_agent_mesh/agent/adk/embed_resolving_mcp_toolset.py

@@ -40,11 +40,6 @@ class EmbedResolvingMCPTool(MCPTool):
         self._original_mcp_tool = original_mcp_tool
         self._tool_config = tool_config or {}
 
-        log.info(
-            "Created EmbedResolvingMCPTool for '%s' with embed resolution capabilities",
-            self.name,
-        )
-
     async def _resolve_embeds_recursively(
         self,
         data: Any,
@@ -287,10 +282,6 @@ class EmbedResolvingMCPToolset(MCPToolset):
 
         # Wrap each tool with embed resolution capability
         embed_resolving_tools = []
-        log.info(
-            "EmbedResolvingMCPToolset: Wrapping %d MCP tools with embed resolution",
-            len(original_tools),
-        )
 
         for tool in original_tools:
             # Get tool-specific config
@@ -304,13 +295,4 @@ class EmbedResolvingMCPToolset(MCPToolset):
             )
             embed_resolving_tools.append(embed_resolving_tool)
 
-            log.info(
-                "EmbedResolvingMCPToolset: Successfully wrapped MCP tool '%s'",
-                tool.name,
-            )
-
-        log.info(
-            "EmbedResolvingMCPToolset: Completed wrapping %d MCP tools",
-            len(embed_resolving_tools),
-        )
         return embed_resolving_tools

solace_agent_mesh/agent/adk/models/models_llm.txt

@@ -1,94 +1,143 @@
-Here is the developer guide for the `models` directory.
+# DEVELOPER GUIDE for models directory
 
 ## Quick Summary
-This directory contains concrete implementations of the `BaseLlm` interface, acting as wrappers or clients for various Large Language Model APIs. These classes translate the ADK's standard `LlmRequest` into the format required by the specific LLM backend and parse the backend's response back into a standard `LlmResponse`.
+This directory contains concrete implementations of the `BaseLlm` interface, providing wrappers for various Large Language Model APIs. These classes translate the ADK's standard `LlmRequest` into provider-specific formats and parse responses back into standard `LlmResponse` objects.
 
 ## Files Overview
-- `lite_llm.py`: An LLM client that uses the `litellm` library to support a wide variety of models from different providers.
+- `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
 
 ### lite_llm.py
-**Purpose:** Provides the `LiteLlm` class, a `BaseLlm` implementation that interfaces with hundreds of LLM models through the `litellm` library. This allows developers to use models from providers like OpenAI, Anthropic, Vertex AI, etc., by simply changing the model string. Environment variables required by the target model provider must be set.
+**Purpose:** Provides the `LiteLlm` class, a `BaseLlm` implementation that interfaces with hundreds of LLM models through the `litellm` library. Supports models from OpenAI, Anthropic, Vertex AI, and many other providers by simply changing the model string.
 
-**Import:** `from google.adk.models.lite_llm import LiteLlm`
+**Import:** `from solace_agent_mesh.agent.adk.models.lite_llm import LiteLlm`
 
 **Classes:**
-- `LiteLlm(model: str, **kwargs)` - Wrapper around `litellm` that can be used with any model it supports.
-  - `__init__(self, model: str, **kwargs)` - Initializes the `LiteLlm` client.
-    - **model**: The name of the model as recognized by `litellm` (e.g., `"vertex_ai/gemini-1.5-pro-preview-0409"`, `"claude-3-opus-20240229"`, `"gpt-4-turbo"`).
-    - **\*\*kwargs**: Additional arguments to pass directly to the `litellm.completion` or `litellm.acompletion` API on every call.
-  - `async generate_content_async(llm_request: LlmRequest, stream: bool = False) -> AsyncGenerator[LlmResponse, None]` - Sends a request to the configured LLM and yields one or more responses.
-    - **llm_request**: The request object containing conversation history, tool definitions, and configuration.
-    - **stream**: If `True`, yields partial responses as they become available. If `False`, yields a single, complete response.
-    - **Returns**: An async generator that yields `LlmResponse` objects.
-  - `supported_models() -> list[str]` - Provides a list of supported models. For `LiteLlm`, this returns an empty list because `litellm` supports a vast and dynamic set of models. Refer to the `litellm` documentation for a complete list.
+- `LiteLlm(model: str, **kwargs)` - Wrapper around `litellm` supporting any model it recognizes
+  - `generate_content_async(llm_request: LlmRequest, stream: bool = False) -> AsyncGenerator[LlmResponse, None]` - Generates content asynchronously with optional streaming
+  - `supported_models() -> list[str]` - Returns list of supported models (empty for LiteLlm due to dynamic model support)
+  - `model: str` - The name of the LiteLlm model
+  - `llm_client: LiteLLMClient` - The LLM client instance used for API calls
+
+- `LiteLLMClient()` - Internal client providing completion methods for better testability
+  - `acompletion(model, messages, tools, **kwargs) -> Union[ModelResponse, CustomStreamWrapper]` - Asynchronous completion call
+  - `completion(model, messages, tools, stream=False, **kwargs) -> Union[ModelResponse, CustomStreamWrapper]` - Synchronous completion call
+
+- `FunctionChunk(BaseModel)` - Represents a function call chunk in streaming responses
+  - `id: Optional[str]` - Function call ID
+  - `name: Optional[str]` - Function name
+  - `args: Optional[str]` - Function arguments as JSON string
+  - `index: Optional[int]` - Index of the function call
+
+- `TextChunk(BaseModel)` - Represents a text chunk in streaming responses
+  - `text: str` - The text content
+
+- `UsageMetadataChunk(BaseModel)` - Represents token usage information
+  - `prompt_tokens: int` - Number of tokens in the prompt
+  - `completion_tokens: int` - Number of tokens in the completion
+  - `total_tokens: int` - Total number of tokens used
+
+**Functions:**
+- `_content_to_message_param(content: types.Content) -> Union[Message, list[Message]]` - Converts ADK Content to litellm Message format
+- `_get_content(parts: Iterable[types.Part]) -> Union[OpenAIMessageContent, str]` - Converts parts to litellm content format
+- `_function_declaration_to_tool_param(function_declaration: types.FunctionDeclaration) -> dict` - Converts function declarations to OpenAPI spec format
+- `_model_response_to_generate_content_response(response: ModelResponse) -> LlmResponse` - Converts litellm response to LlmResponse
 
 **Usage Examples:**
 ```python
 import asyncio
 import os
-from google.adk.models.lite_llm import LiteLlm
-from google.adk.models.llm_request import LlmRequest, LlmConfig
-from google.adk.models.types import Content, Part
-
-# Example using a Vertex AI model via litellm.
-# Set environment variables required by the provider.
-# os.environ["VERTEXAI_PROJECT"] = "your-gcp-project-id"
-# os.environ["VERTEXAI_LOCATION"] = "your-gcp-location"
+from solace_agent_mesh.agent.adk.models.lite_llm import LiteLlm
+from solace_agent_mesh.agent.adk.models.llm_request import LlmRequest, LlmConfig
+from google.genai.types import Content, Part
 
-# Example using an OpenAI model via litellm.
-# os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
+# Set environment variables for your chosen provider
+# For OpenAI:
+# os.environ["OPENAI_API_KEY"] = "your-api-key"
+# For Vertex AI:
+# os.environ["VERTEXAI_PROJECT"] = "your-project-id"
+# os.environ["VERTEXAI_LOCATION"] = "your-location"
 
 async def main():
-    # Instantiate the LiteLlm client with a specific model.
-    # Any additional kwargs are passed to litellm's completion call.
+    # Initialize LiteLlm with a specific model
     llm = LiteLlm(
         model="gpt-4-turbo",
-        temperature=0.5,
-        max_tokens=150
+        temperature=0.7,
+        max_completion_tokens=150
     )
-
-    # Construct a request to the LLM
+    
+    # Create a request
     request = LlmRequest(
         contents=[
             Content(
                 role="user",
-                parts=[Part.from_text("Why is the sky blue?")]
+                parts=[Part.from_text("Explain quantum computing in simple terms")]
             )
         ],
         config=LlmConfig(
-            # The temperature set in the constructor can be overridden here
-            temperature=0.7
+            temperature=0.5,
+            max_output_tokens=200
         )
     )
-
-    print("--- Non-streaming example ---")
-    # Get a single, complete response
+    
+    # Non-streaming generation
+    print("=== Non-streaming ===")
     async for response in llm.generate_content_async(request, stream=False):
-        if response.text:
-            print(response.text)
+        print(f"Response: {response.text}")
         if response.usage_metadata:
-            print(f"Token usage: {response.usage_metadata.total_token_count}")
-
-    print("\n--- Streaming example ---")
-    # Get a stream of partial responses
-    full_response_text = ""
+            print(f"Tokens used: {response.usage_metadata.total_token_count}")
+    
+    # Streaming generation
+    print("\n=== Streaming ===")
     async for response in llm.generate_content_async(request, stream=True):
         if response.text:
             print(response.text, end="", flush=True)
-            full_response_text += response.text
         if response.usage_metadata:
-            # Usage metadata is typically sent with the final chunk
-            print(f"\nToken usage: {response.usage_metadata.total_token_count}")
+            print(f"\nTotal tokens: {response.usage_metadata.total_token_count}")
+
+# Example with function calling
+async def function_calling_example():
+    from google.genai.types import FunctionDeclaration, Schema, Type, Tool
+    
+    # Define a function for the LLM to call
+    get_weather_func = FunctionDeclaration(
+        name="get_weather",
+        description="Get current weather for a location",
+        parameters=Schema(
+            type=Type.OBJECT,
+            properties={
+                "location": Schema(type=Type.STRING, description="City name"),
+                "unit": Schema(type=Type.STRING, description="Temperature unit")
+            },
+            required=["location"]
+        )
+    )
+    
+    llm = LiteLlm(model="gpt-4-turbo")
+    
+    request = LlmRequest(
+        contents=[
+            Content(
+                role="user", 
+                parts=[Part.from_text("What's the weather like in Tokyo?")]
+            )
+        ],
+        config=LlmConfig(
+            tools=[Tool(function_declarations=[get_weather_func])]
+        )
+    )
+    
+    async for response in llm.generate_content_async(request):
+        if response.function_calls:
+            for func_call in response.function_calls:
+                print(f"Function called: {func_call.name}")
+                print(f"Arguments: {func_call.args}")
 
 if __name__ == "__main__":
-    # To run this example, you need to have the necessary environment variables set
-    # for the model you choose (e.g., OPENAI_API_KEY).
-    # You would also need to install the required provider packages, e.g.,
-    # pip install litellm[openai]
-    #
-    # Since this is an async function, we run it with asyncio.
-    # asyncio.run(main())
-    pass
-```
+    asyncio.run(main())
+    # asyncio.run(function_calling_example())
+```
+
+# content_hash: 12789ad2e16cd9ea5a81abdd68258d9ef30520bed5c51ba8d00ea66014191964