solace-agent-mesh 0.2.3__py3-none-any.whl → 1.0.1__py3-none-any.whl

solace_agent_mesh/agent/adk/models/models_llm.txt

@@ -0,0 +1,94 @@
+Here is the developer guide for the `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`.
+
+## Files Overview
+- `lite_llm.py`: An LLM client that uses the `litellm` library to support a wide variety of models from different providers.
+
+## 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.
+
+**Import:** `from google.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.
+
+**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"
+
+# Example using an OpenAI model via litellm.
+# os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
+
+async def main():
+    # Instantiate the LiteLlm client with a specific model.
+    # Any additional kwargs are passed to litellm's completion call.
+    llm = LiteLlm(
+        model="gpt-4-turbo",
+        temperature=0.5,
+        max_tokens=150
+    )
+
+    # Construct a request to the LLM
+    request = LlmRequest(
+        contents=[
+            Content(
+                role="user",
+                parts=[Part.from_text("Why is the sky blue?")]
+            )
+        ],
+        config=LlmConfig(
+            # The temperature set in the constructor can be overridden here
+            temperature=0.7
+        )
+    )
+
+    print("--- Non-streaming example ---")
+    # Get a single, complete response
+    async for response in llm.generate_content_async(request, stream=False):
+        if response.text:
+            print(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 = ""
+    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}")
+
+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
+```

solace_agent_mesh/agent/adk/runner.py

@@ -0,0 +1,353 @@
+"""
+Manages the asynchronous execution of the ADK Runner.
+"""
+
+import asyncio
+from google.adk.agents.invocation_context import LlmCallsLimitExceededError
+
+
+class TaskCancelledError(Exception):
+    """Raised when an ADK task is cancelled via external signal."""
+
+    pass
+
+
+from typing import Dict, Any, TYPE_CHECKING
+
+from solace_ai_connector.common.log import log
+
+from google.adk.sessions import Session as ADKSession
+from google.adk.agents import RunConfig
+from google.genai import types as adk_types
+from google.adk.events import Event as ADKEvent
+from google.adk.events.event_actions import EventActions
+
+from ...common.types import CancelTaskRequest, TaskIdParams
+
+if TYPE_CHECKING:
+    from ..sac.component import SamAgentComponent
+    from ..sac.task_execution_context import TaskExecutionContext
+
+
+async def run_adk_async_task_thread_wrapper(
+    component: "SamAgentComponent",
+    adk_session: ADKSession,
+    adk_content: adk_types.Content,
+    run_config: RunConfig,
+    a2a_context: Dict[str, Any],
+):
+    """
+    Wrapper to run the async ADK task.
+    Calls component finalization methods upon completion or error.
+
+    Args:
+        component: The SamAgentComponent instance.
+        adk_session: The ADK session to use (from component.session_service).
+        adk_content: The input content for the ADK agent.
+        run_config: The ADK run configuration.
+        a2a_context: The context dictionary for this specific A2A request.
+    """
+    logical_task_id = a2a_context.get("logical_task_id", "unknown_task")
+    is_paused = False
+    exception_to_finalize_with = None
+    task_context = None
+    try:
+        with component.active_tasks_lock:
+            task_context = component.active_tasks.get(logical_task_id)
+
+        if not task_context:
+            log.error(
+                "%s TaskExecutionContext not found for task %s. Cannot start ADK runner.",
+                component.log_identifier,
+                logical_task_id,
+            )
+            return
+
+        task_context.flush_streaming_buffer()
+        log.debug(
+            "%s Cleared streaming text buffer before starting ADK task %s.",
+            component.log_identifier,
+            logical_task_id,
+        )
+
+        if adk_session and component.session_service:
+            context_setting_invocation_id = logical_task_id
+            context_setting_event = ADKEvent(
+                invocation_id=context_setting_invocation_id,
+                author="A2A_Host_System",
+                content=adk_types.Content(
+                    parts=[
+                        adk_types.Part(text="Initializing A2A context for task run.")
+                    ]
+                ),
+                actions=EventActions(state_delta={"a2a_context": a2a_context}),
+                branch=None,
+            )
+            try:
+                await component.session_service.append_event(
+                    session=adk_session, event=context_setting_event
+                )
+                log.info(
+                    "%s Appended context-setting event to ADK session %s (via component.session_service) for task %s.",
+                    component.log_identifier,
+                    adk_session.id,
+                    logical_task_id,
+                )
+            except Exception as e_append:
+                log.error(
+                    "%s Failed to append context-setting event for task %s: %s. Tool scope filtering might not work if state is not persisted.",
+                    component.log_identifier,
+                    logical_task_id,
+                    e_append,
+                    exc_info=True,
+                )
+        else:
+            log.warning(
+                "%s Could not inject a2a_context into ADK session state via event for task %s (session or session_service invalid). Tool scope filtering might not work.",
+                component.log_identifier,
+                logical_task_id,
+            )
+
+        is_paused = await run_adk_async_task(
+            component,
+            task_context,
+            adk_session,
+            adk_content,
+            run_config,
+            a2a_context,
+        )
+
+        log.debug(
+            "%s ADK task %s awaited and completed (Paused: %s).",
+            component.log_identifier,
+            logical_task_id,
+            is_paused,
+        )
+
+    except TaskCancelledError as tce:
+        exception_to_finalize_with = tce
+        log.info(
+            "%s Task %s was cancelled. Propagating to peers before scheduling finalization. Message: %s",
+            component.log_identifier,
+            logical_task_id,
+            tce,
+        )
+        sub_tasks_to_cancel = task_context.peer_sub_tasks if task_context else []
+
+        if sub_tasks_to_cancel:
+            log.info(
+                "%s Propagating cancellation to %d peer sub-task(s) for main task %s.",
+                component.log_identifier,
+                len(sub_tasks_to_cancel),
+                logical_task_id,
+            )
+            for sub_task_info in sub_tasks_to_cancel:
+                try:
+                    sub_task_id = sub_task_info.get("sub_task_id")
+                    target_peer_agent_name = sub_task_info.get("peer_agent_name")
+                    if not sub_task_id or not target_peer_agent_name:
+                        log.warning(
+                            "%s Incomplete sub-task info found, cannot cancel: %s",
+                            component.log_identifier,
+                            sub_task_info,
+                        )
+                        continue
+                    peer_cancel_params = TaskIdParams(id=sub_task_id)
+                    peer_cancel_request = CancelTaskRequest(params=peer_cancel_params)
+                    peer_cancel_user_props = {"clientId": component.agent_name}
+                    peer_request_topic = component._get_agent_request_topic(
+                        target_peer_agent_name
+                    )
+                    component._publish_a2a_message(
+                        payload=peer_cancel_request.model_dump(exclude_none=True),
+                        topic=peer_request_topic,
+                        user_properties=peer_cancel_user_props,
+                    )
+                except Exception as e_peer_cancel:
+                    log.error(
+                        "%s Failed to send CancelTaskRequest for sub-task %s: %s",
+                        component.log_identifier,
+                        sub_task_info.get("sub_task_id"),
+                        e_peer_cancel,
+                        exc_info=True,
+                    )
+    except LlmCallsLimitExceededError as llm_limit_e:
+        exception_to_finalize_with = llm_limit_e
+        log.warning(
+            "%s LLM call limit exceeded for task %s: %s. Scheduling finalization.",
+            component.log_identifier,
+            logical_task_id,
+            llm_limit_e,
+        )
+    except Exception as e:
+        exception_to_finalize_with = e
+        log.exception(
+            "%s Exception in ADK runner for task %s: %s. Scheduling finalization.",
+            component.log_identifier,
+            logical_task_id,
+            e,
+        )
+
+    loop = component.get_async_loop()
+    if loop and loop.is_running():
+        log.debug(
+            "%s Scheduling finalize_task_with_cleanup for task %s.",
+            component.log_identifier,
+            logical_task_id,
+        )
+        asyncio.run_coroutine_threadsafe(
+            component.finalize_task_with_cleanup(
+                a2a_context, is_paused, exception_to_finalize_with
+            ),
+            loop,
+        )
+    else:
+        log.error(
+            "%s Async loop not available. Cannot schedule finalization for task %s.",
+            component.log_identifier,
+            logical_task_id,
+        )
+
+        log.debug(
+            "%s ADK runner for task %s finished.",
+            component.log_identifier,
+            logical_task_id,
+        )
+
+
+async def run_adk_async_task(
+    component: "SamAgentComponent",
+    task_context: "TaskExecutionContext",
+    adk_session: ADKSession,
+    adk_content: adk_types.Content,
+    run_config: RunConfig,
+    a2a_context: Dict[str, Any],
+) -> bool:
+    """
+    Runs the ADK Runner asynchronously and calls component methods to process
+    intermediate events and finalize the task.
+    Returns:
+        bool: True if the task is paused for a long-running tool, False otherwise.
+    """
+    logical_task_id = a2a_context.get("logical_task_id", "unknown_task")
+    event_loop_stored = False
+    current_loop = asyncio.get_running_loop()
+    is_paused = False
+
+    adk_event_generator = component.runner.run_async(
+        user_id=adk_session.user_id,
+        session_id=adk_session.id,
+        new_message=adk_content,
+        run_config=run_config,
+    )
+
+    try:
+        while True:
+            next_event_task = asyncio.create_task(adk_event_generator.__anext__())
+            cancel_wait_task = asyncio.create_task(
+                task_context.cancellation_event.wait()
+            )
+
+            done, pending = await asyncio.wait(
+                {next_event_task, cancel_wait_task},
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            if cancel_wait_task in done:
+                next_event_task.cancel()
+                try:
+                    await next_event_task
+                except asyncio.CancelledError:
+                    log.debug(
+                        "%s Suppressed CancelledError for next_event_task after signal.",
+                        component.log_identifier,
+                    )
+                log.info(
+                    "%s Task %s cancellation detected while awaiting ADK event.",
+                    component.log_identifier,
+                    logical_task_id,
+                )
+                raise TaskCancelledError(
+                    f"Task {logical_task_id} was cancelled by signal."
+                )
+
+            for task in pending:
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    log.debug(
+                        "%s Suppressed CancelledError for lingering task after event.",
+                        component.log_identifier,
+                    )
+
+            try:
+                event = await next_event_task
+            except StopAsyncIteration:
+                break
+
+            if event.long_running_tool_ids:
+                is_paused = True
+
+            if not event_loop_stored and event.invocation_id:
+                task_context.set_event_loop(current_loop)
+                a2a_context["invocation_id"] = event.invocation_id
+                event_loop_stored = True
+
+            try:
+                await component.process_and_publish_adk_event(event, a2a_context)
+            except Exception as process_err:
+                log.exception(
+                    "%s Error processing intermediate ADK event %s for task %s: %s",
+                    component.log_identifier,
+                    event.id,
+                    logical_task_id,
+                    process_err,
+                )
+
+            if task_context.is_cancelled():
+                raise TaskCancelledError(
+                    f"Task {logical_task_id} was cancelled after processing ADK event {event.id}."
+                )
+
+            if event.content and event.content.parts:
+                for part in event.content.parts:
+                    if part.function_response:
+                        if part.function_response.name.startswith("peer_"):
+                            pass
+
+    except TaskCancelledError:
+        raise
+    except Exception as e:
+        log.exception(
+            "%s Unexpected error in ADK runner loop for task %s: %s",
+            component.log_identifier,
+            logical_task_id,
+            e,
+        )
+        raise
+
+    if task_context.is_cancelled():
+        log.info(
+            "%s Task %s cancellation detected before finalization.",
+            component.log_identifier,
+            logical_task_id,
+        )
+        raise TaskCancelledError(
+            f"Task {logical_task_id} was cancelled before finalization."
+        )
+
+    if is_paused:
+        log.info(
+            "%s ADK run completed by invoking a long-running tool. Task %s will remain open, awaiting peer response.",
+            component.log_identifier,
+            logical_task_id,
+        )
+        return True
+
+    log.debug(
+        "%s ADK run_async completed for task %s. Returning to wrapper for finalization.",
+        component.log_identifier,
+        logical_task_id,
+    )
+    return False

solace_agent_mesh/agent/adk/services.py

@@ -0,0 +1,240 @@
+"""
+Initializes ADK Services based on configuration.
+"""
+
+import os
+import re
+from typing import Dict
+
+from solace_ai_connector.common.log import log
+
+from google.adk.sessions import (
+    BaseSessionService,
+    InMemorySessionService,
+    DatabaseSessionService,
+    VertexAiSessionService,
+)
+from google.adk.artifacts import (
+    BaseArtifactService,
+    InMemoryArtifactService,
+    GcsArtifactService,
+)
+from google.adk.memory import (
+    BaseMemoryService,
+    InMemoryMemoryService,
+    VertexAiRagMemoryService,
+)
+
+from .filesystem_artifact_service import FilesystemArtifactService
+
+try:
+    from tests.integration.infrastructure.artifact_service.service import (
+        TestInMemoryArtifactService,
+    )
+except ImportError:
+    TestInMemoryArtifactService = None
+
+
+def _sanitize_for_path(identifier: str) -> str:
+    """Sanitizes a string to be safe for use as a directory name."""
+    if not identifier:
+        return "_invalid_scope_"
+    sanitized = re.sub(r'[\\/*?:"<>|]', "_", identifier)
+    sanitized = re.sub(r"_+", "_", sanitized)
+    sanitized = sanitized.strip("_ ")
+    if not sanitized:
+        return "_empty_scope_"
+    return sanitized
+
+
+def initialize_session_service(component) -> BaseSessionService:
+    """Initializes the ADK Session Service based on configuration."""
+    config: Dict = component.get_config("session_service", {})
+    service_type = config.get("type", "memory").lower()
+    log.info(
+        "%s Initializing Session Service of type: %s",
+        component.log_identifier,
+        service_type,
+    )
+
+    if service_type == "memory":
+        return InMemorySessionService()
+    elif service_type == "database":
+        db_url = config.get("db_url")
+        if not db_url:
+            raise ValueError(
+                f"{component.log_identifier} 'db_url' is required for database session service."
+            )
+        try:
+            return DatabaseSessionService(db_url=db_url)
+        except ImportError:
+            log.error(
+                "%s SQLAlchemy not installed. Please install 'google-adk[database]' or 'sqlalchemy'.",
+                component.log_identifier,
+            )
+            raise
+    elif service_type == "vertex":
+        project = os.environ.get("GOOGLE_CLOUD_PROJECT")
+        location = os.environ.get("GOOGLE_CLOUD_LOCATION")
+        if not project or not location:
+            raise ValueError(
+                f"{component.log_identifier} GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION env vars required for vertex session service."
+            )
+        return VertexAiSessionService(project=project, location=location)
+    else:
+        raise ValueError(
+            f"{component.log_identifier} Unsupported session service type: {service_type}"
+        )
+
+
+def initialize_artifact_service(component) -> BaseArtifactService:
+    """Initializes the ADK Artifact Service based on configuration."""
+    config: Dict = component.get_config("artifact_service", {"type": "memory"})
+    service_type = config.get("type", "memory").lower()
+    log.info(
+        "%s Initializing Artifact Service of type: %s",
+        component.log_identifier,
+        service_type,
+    )
+
+    if service_type == "memory":
+        return InMemoryArtifactService()
+    elif service_type == "gcs":
+        bucket_name = config.get("bucket_name")
+        if not bucket_name:
+            raise ValueError(
+                f"{component.log_identifier} 'bucket_name' is required for GCS artifact service."
+            )
+        try:
+            gcs_args = {
+                k: v for k, v in config.items() if k not in ["type", "bucket_name"]
+            }
+            return GcsArtifactService(bucket_name=bucket_name, **gcs_args)
+        except ImportError:
+            log.error(
+                "%s google-cloud-storage not installed. Please install 'google-adk[gcs]' or 'google-cloud-storage'.",
+                component.log_identifier,
+            )
+            raise
+    elif service_type == "filesystem":
+        base_path = config.get("base_path")
+        if not base_path:
+            raise ValueError(
+                f"{component.log_identifier} 'base_path' is required for filesystem artifact service."
+            )
+
+        artifact_scope = config.get("artifact_scope", "namespace").lower()
+        scope_identifier_raw = None
+
+        if artifact_scope == "app":
+            app_instance = component.get_app()
+            if not app_instance or not app_instance.name:
+                raise ValueError(
+                    f"{component.log_identifier} Cannot determine app name for 'app' scope."
+                )
+            scope_identifier_raw = app_instance.name
+            log.info(
+                "%s Using 'app' scope for filesystem artifacts: %s",
+                component.log_identifier,
+                scope_identifier_raw,
+            )
+        elif artifact_scope == "namespace":
+            scope_identifier_raw = component.get_config("namespace")
+            log.info(
+                "%s Using 'namespace' scope for filesystem artifacts: %s",
+                component.log_identifier,
+                scope_identifier_raw,
+            )
+        elif artifact_scope == "custom":
+            scope_identifier_raw = config.get("artifact_scope_value")
+            if not scope_identifier_raw:
+                raise ValueError(
+                    f"{component.log_identifier} 'artifact_scope_value' is required when artifact_scope is 'custom'."
+                )
+            log.info(
+                "%s Using 'custom' scope for filesystem artifacts: %s",
+                component.log_identifier,
+                scope_identifier_raw,
+            )
+        else:
+            raise ValueError(
+                f"{component.log_identifier} Invalid 'artifact_scope' value: {artifact_scope}"
+            )
+
+        if not scope_identifier_raw:
+            raise ValueError(
+                f"{component.log_identifier} Failed to determine scope identifier for filesystem artifacts."
+            )
+
+        scope_identifier_sanitized = _sanitize_for_path(scope_identifier_raw)
+        log.info(
+            "%s Sanitized scope identifier: %s",
+            component.log_identifier,
+            scope_identifier_sanitized,
+        )
+
+        try:
+            return FilesystemArtifactService(
+                base_path=base_path, scope_identifier=scope_identifier_sanitized
+            )
+        except Exception as e:
+            log.error(
+                "%s Failed to initialize FilesystemArtifactService: %s",
+                component.log_identifier,
+                e,
+            )
+            raise
+    elif service_type == "test_in_memory":
+        if TestInMemoryArtifactService is None:
+            log.error(
+                "%s TestInMemoryArtifactService is configured but could not be imported. "
+                "Ensure test infrastructure is in PYTHONPATH if running tests, or check configuration.",
+                component.log_identifier,
+            )
+            raise ImportError("TestInMemoryArtifactService not available.")
+        log.info(
+            "%s Using TestInMemoryArtifactService for testing.",
+            component.log_identifier,
+        )
+        return TestInMemoryArtifactService()
+    else:
+        raise ValueError(
+            f"{component.log_identifier} Unsupported artifact service type: {service_type}"
+        )
+
+
+def initialize_memory_service(component) -> BaseMemoryService:
+    """Initializes the ADK Memory Service based on configuration."""
+    config: Dict = component.get_config("memory_service", {"type": "memory"})
+    service_type = config.get("type", "memory").lower()
+    log.info(
+        "%s Initializing Memory Service of type: %s",
+        component.log_identifier,
+        service_type,
+    )
+
+    if service_type == "memory":
+        return InMemoryMemoryService()
+    elif service_type == "vertex_rag":
+        try:
+            rag_args = {
+                k: v for k, v in config.items() if k not in ["type", "default_behavior"]
+            }
+            return VertexAiRagMemoryService(**rag_args)
+        except ImportError:
+            log.error(
+                "%s google-cloud-aiplatform not installed. Please install 'google-adk[vertex]' or 'google-cloud-aiplatform'.",
+                component.log_identifier,
+            )
+            raise
+        except TypeError as e:
+            log.error(
+                "%s Error initializing VertexAiRagMemoryService: %s. Check config params.",
+                component.log_identifier,
+                e,
+            )
+            raise
+    else:
+        raise ValueError(
+            f"{component.log_identifier} Unsupported memory service type: {service_type}"
+        )