agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/adk/_modules/tasks.py

@@ -0,0 +1,124 @@
+from datetime import timedelta
+
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.services.adk.tasks import TasksService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+from agentex.lib.core.temporal.activities.adk.tasks_activities import (
+    DeleteTaskParams,
+    GetTaskParams,
+    TasksActivityName,
+)
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.task import Task
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import in_temporal_workflow
+
+logger = make_logger(__name__)
+
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class TasksModule:
+    """
+    Module for managing tasks in Agentex.
+    Provides high-level async methods for retrieving, listing, and deleting tasks.
+    """
+
+    def __init__(
+        self,
+        tasks_service: TasksService | None = None,
+    ):
+        if tasks_service is None:
+            agentex_client = AsyncAgentex()
+            tracer = AsyncTracer(agentex_client)
+            self._tasks_service = TasksService(
+                agentex_client=agentex_client, tracer=tracer
+            )
+        else:
+            self._tasks_service = tasks_service
+
+    async def get(
+        self,
+        *,
+        task_id: str | None = None,
+        task_name: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> Task:
+        """
+        Get a task by ID or name.
+        Args:
+            task_id: The ID of the task to retrieve.
+            task_name: The name of the task to retrieve.
+        Returns:
+            The task entry.
+        """
+        params = GetTaskParams(
+            task_id=task_id,
+            task_name=task_name,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=TasksActivityName.GET_TASK,
+                request=params,
+                response_type=Task,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._tasks_service.get_task(
+                task_id=task_id,
+                task_name=task_name,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def delete(
+        self,
+        *,
+        task_id: str | None = None,
+        task_name: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> Task:
+        """
+        Delete a task by ID or name.
+        Args:
+            task_id: The ID of the task to delete.
+            task_name: The name of the task to delete.
+        Returns:
+            The deleted task entry.
+        """
+        params = DeleteTaskParams(
+            task_id=task_id,
+            task_name=task_name,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=TasksActivityName.DELETE_TASK,
+                request=params,
+                response_type=Task,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._tasks_service.delete_task(
+                task_id=task_id,
+                task_name=task_name,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )

agentex/lib/adk/_modules/tracing.py

@@ -0,0 +1,194 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from typing import Any
+
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.services.adk.tracing import TracingService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+from agentex.lib.core.temporal.activities.adk.tracing_activities import (
+    EndSpanParams,
+    StartSpanParams,
+    TracingActivityName,
+)
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.span import Span
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.model_utils import BaseModel
+from agentex.lib.utils.temporal import in_temporal_workflow
+
+logger = make_logger(__name__)
+
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class TracingModule:
+    """
+    Module for managing tracing and span operations in Agentex.
+    Provides high-level async methods for starting, ending, and managing spans for distributed tracing.
+    """
+
+    def __init__(self, tracing_service: TracingService | None = None):
+        """
+        Initialize the tracing interface.
+
+        Args:
+            tracing_activities (Optional[TracingActivities]): Optional pre-configured tracing activities. If None, will be auto-initialized.
+        """
+        if tracing_service is None:
+            agentex_client = AsyncAgentex()
+            tracer = AsyncTracer(agentex_client)
+            self._tracing_service = TracingService(tracer=tracer)
+        else:
+            self._tracing_service = tracing_service
+
+    @asynccontextmanager
+    async def span(
+        self,
+        trace_id: str,
+        name: str,
+        input: list[Any] | dict[str, Any] | BaseModel | None = None,
+        data: list[Any] | dict[str, Any] | BaseModel | None = None,
+        parent_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> AsyncGenerator[Span | None, None]:
+        """
+        Async context manager for creating and automatically closing a span.
+        Yields the started span object. The span is automatically ended when the context exits.
+
+        If trace_id is falsy, acts as a no-op context manager.
+
+        Args:
+            trace_id (str): The trace ID for the span.
+            name (str): The name of the span.
+            input (Union[List, Dict, BaseModel]): The input for the span.
+            parent_id (Optional[str]): The parent span ID for the span.
+            data (Optional[Union[List, Dict, BaseModel]]): The data for the span.
+            start_to_close_timeout (timedelta): The start to close timeout for the span.
+            heartbeat_timeout (timedelta): The heartbeat timeout for the span.
+            retry_policy (RetryPolicy): The retry policy for the span.
+
+        Returns:
+            AsyncGenerator[Optional[Span], None]: An async generator that yields the started span object.
+        """
+        if not trace_id:
+            yield None
+            return
+
+        span: Span | None = await self.start_span(
+            trace_id=trace_id,
+            name=name,
+            input=input,
+            parent_id=parent_id,
+            data=data,
+            start_to_close_timeout=start_to_close_timeout,
+            heartbeat_timeout=heartbeat_timeout,
+            retry_policy=retry_policy,
+        )
+        try:
+            yield span
+        finally:
+            if span:
+                await self.end_span(
+                    trace_id=trace_id,
+                    span=span,
+                    start_to_close_timeout=start_to_close_timeout,
+                    heartbeat_timeout=heartbeat_timeout,
+                    retry_policy=retry_policy,
+                )
+
+    async def start_span(
+        self,
+        trace_id: str,
+        name: str,
+        input: list[Any] | dict[str, Any] | BaseModel | None = None,
+        parent_id: str | None = None,
+        data: list[Any] | dict[str, Any] | BaseModel | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=1),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> Span | None:
+        """
+        Start a new span in the trace.
+
+        Args:
+            trace_id (str): The trace ID for the span.
+            name (str): The name of the span.
+            input (Union[List, Dict, BaseModel]): The input for the span.
+            parent_id (Optional[str]): The parent span ID for the span.
+            data (Optional[Union[List, Dict, BaseModel]]): The data for the span.
+            start_to_close_timeout (timedelta): The start to close timeout for the span.
+            heartbeat_timeout (timedelta): The heartbeat timeout for the span.
+            retry_policy (RetryPolicy): The retry policy for the span.
+
+        Returns:
+            Span: The started span object.
+        """
+        params = StartSpanParams(
+            trace_id=trace_id,
+            parent_id=parent_id,
+            name=name,
+            input=input,
+            data=data,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=TracingActivityName.START_SPAN,
+                request=params,
+                response_type=Span,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._tracing_service.start_span(
+                trace_id=trace_id,
+                name=name,
+                input=input,
+                parent_id=parent_id,
+                data=data,
+            )
+
+    async def end_span(
+        self,
+        trace_id: str,
+        span: Span,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=1),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> Span:
+        """
+        End an existing span in the trace.
+
+        Args:
+            trace_id (str): The trace ID for the span.
+            span (Span): The span to end.
+            start_to_close_timeout (timedelta): The start to close timeout for the span.
+            heartbeat_timeout (timedelta): The heartbeat timeout for the span.
+            retry_policy (RetryPolicy): The retry policy for the span.
+
+        Returns:
+            Span: The ended span object.
+        """
+        params = EndSpanParams(
+            trace_id=trace_id,
+            span=span,
+        )
+        if in_temporal_workflow():
+            return await ActivityHelpers.execute_activity(
+                activity_name=TracingActivityName.END_SPAN,
+                request=params,
+                response_type=Span,
+                start_to_close_timeout=start_to_close_timeout,
+                retry_policy=retry_policy,
+                heartbeat_timeout=heartbeat_timeout,
+            )
+        else:
+            return await self._tracing_service.end_span(
+                trace_id=trace_id,
+                span=span,
+            )

agentex/lib/adk/providers/__init__.py

@@ -0,0 +1,9 @@
+from agentex.lib.adk.providers._modules.litellm import LiteLLMModule
+from agentex.lib.adk.providers._modules.openai import OpenAIModule
+from agentex.lib.adk.providers._modules.sgp import SGPModule
+
+openai = OpenAIModule()
+litellm = LiteLLMModule()
+sgp = SGPModule()
+
+__all__ = ["openai", "litellm", "sgp"]

agentex/lib/adk/providers/_modules/litellm.py

@@ -0,0 +1,232 @@
+from collections.abc import AsyncGenerator
+from datetime import timedelta
+
+from temporalio.common import RetryPolicy
+
+from agentex import AsyncAgentex
+from agentex.lib.core.adapters.llm.adapter_litellm import LiteLLMGateway
+from agentex.lib.core.adapters.streams.adapter_redis import RedisEventStreamRepository
+from agentex.lib.core.services.adk.providers.litellm import LiteLLMService
+from agentex.lib.core.services.adk.streaming import StreamingService
+from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers
+from agentex.lib.core.temporal.activities.adk.providers.litellm_activities import (
+    ChatCompletionAutoSendParams,
+    ChatCompletionParams,
+    ChatCompletionStreamAutoSendParams,
+    LiteLLMActivityName,
+)
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.lib.types.llm_messages import Completion, LLMConfig
+from agentex.types.task_message import TaskMessage
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import in_temporal_workflow
+
+logger = make_logger(__name__)
+
+# Default retry policy for all LiteLLM operations
+DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
+
+
+class LiteLLMModule:
+    """
+    Module for managing LiteLLM agent operations in Agentex.
+    Provides high-level methods for chat completion, streaming, agentic streaming.
+    """
+
+    def __init__(
+        self,
+        litellm_service: LiteLLMService | None = None,
+    ):
+        if litellm_service is None:
+            # Create default service
+            agentex_client = AsyncAgentex()
+            stream_repository = RedisEventStreamRepository()
+            streaming_service = StreamingService(
+                agentex_client=agentex_client,
+                stream_repository=stream_repository,
+            )
+            litellm_gateway = LiteLLMGateway()
+            tracer = AsyncTracer(agentex_client)
+            self._litellm_service = LiteLLMService(
+                agentex_client=agentex_client,
+                llm_gateway=litellm_gateway,
+                streaming_service=streaming_service,
+                tracer=tracer,
+            )
+        else:
+            self._litellm_service = litellm_service
+
+    async def chat_completion(
+        self,
+        llm_config: LLMConfig,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=120),
+        heartbeat_timeout: timedelta = timedelta(seconds=120),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> Completion:
+        """
+        Perform a chat completion using LiteLLM.
+
+        Args:
+            llm_config (LLMConfig): The configuration for the LLM.
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            Completion: An OpenAI compatible Completion object
+        """
+        if in_temporal_workflow():
+            params = ChatCompletionParams(
+                trace_id=trace_id, parent_span_id=parent_span_id, llm_config=llm_config
+            )
+            return await ActivityHelpers.execute_activity(
+                activity_name=LiteLLMActivityName.CHAT_COMPLETION,
+                request=params,
+                response_type=Completion,
+                start_to_close_timeout=start_to_close_timeout,
+                heartbeat_timeout=heartbeat_timeout,
+                retry_policy=retry_policy,
+            )
+        else:
+            return await self._litellm_service.chat_completion(
+                llm_config=llm_config,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def chat_completion_auto_send(
+        self,
+        task_id: str,
+        llm_config: LLMConfig,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=120),
+        heartbeat_timeout: timedelta = timedelta(seconds=120),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> TaskMessage | None:
+        """
+        Chat completion with automatic TaskMessage creation.
+
+        Args:
+            task_id (str): The ID of the task.
+            llm_config (LLMConfig): The configuration for the LLM (must have stream=False).
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            TaskMessage: The final TaskMessage
+        """
+        if in_temporal_workflow():
+            # Use streaming activity with stream=False for non-streaming auto-send
+            params = ChatCompletionAutoSendParams(
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+                task_id=task_id,
+                llm_config=llm_config,
+            )
+            return await ActivityHelpers.execute_activity(
+                activity_name=LiteLLMActivityName.CHAT_COMPLETION_AUTO_SEND,
+                request=params,
+                response_type=TaskMessage,
+                start_to_close_timeout=start_to_close_timeout,
+                heartbeat_timeout=heartbeat_timeout,
+                retry_policy=retry_policy,
+            )
+        else:
+            return await self._litellm_service.chat_completion_auto_send(
+                task_id=task_id,
+                llm_config=llm_config,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )
+
+    async def chat_completion_stream(
+        self,
+        llm_config: LLMConfig,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> AsyncGenerator[Completion, None]:
+        """
+        Stream chat completion chunks using LiteLLM.
+
+        DEFAULT: Returns raw streaming chunks for manual handling.
+
+        NOTE: This method does NOT work in Temporal workflows!
+        Temporal activities cannot return generators. Use chat_completion_stream_auto_send() instead.
+
+        Args:
+            llm_config (LLMConfig): The configuration for the LLM (must have stream=True).
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            AsyncGenerator[Completion, None]: Generator yielding completion chunks
+
+        Raises:
+            ValueError: If called from within a Temporal workflow
+        """
+        # Delegate to service - it handles temporal workflow checks
+        async for chunk in self._litellm_service.chat_completion_stream(
+            llm_config=llm_config,
+            trace_id=trace_id,
+            parent_span_id=parent_span_id,
+        ):
+            yield chunk
+
+    async def chat_completion_stream_auto_send(
+        self,
+        task_id: str,
+        llm_config: LLMConfig,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=120),
+        heartbeat_timeout: timedelta = timedelta(seconds=120),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> TaskMessage | None:
+        """
+        Stream chat completion with automatic TaskMessage creation and streaming.
+
+        Args:
+            task_id (str): The ID of the task to run the agent for.
+            llm_config (LLMConfig): The configuration for the LLM (must have stream=True).
+            trace_id (Optional[str]): The trace ID for tracing.
+            parent_span_id (Optional[str]): The parent span ID for tracing.
+            start_to_close_timeout (timedelta): The start to close timeout.
+            heartbeat_timeout (timedelta): The heartbeat timeout.
+            retry_policy (RetryPolicy): The retry policy.
+
+        Returns:
+            TaskMessage: The final TaskMessage after streaming is complete
+        """
+        if in_temporal_workflow():
+            params = ChatCompletionStreamAutoSendParams(
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+                task_id=task_id,
+                llm_config=llm_config,
+            )
+            return await ActivityHelpers.execute_activity(
+                activity_name=LiteLLMActivityName.CHAT_COMPLETION_STREAM_AUTO_SEND,
+                request=params,
+                response_type=TaskMessage,
+                start_to_close_timeout=start_to_close_timeout,
+                heartbeat_timeout=heartbeat_timeout,
+                retry_policy=retry_policy,
+            )
+        else:
+            return await self._litellm_service.chat_completion_stream_auto_send(
+                task_id=task_id,
+                llm_config=llm_config,
+                trace_id=trace_id,
+                parent_span_id=parent_span_id,
+            )