agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/core/temporal/activities/__init__.py

@@ -0,0 +1,207 @@
+from scale_gp import SGPClient, SGPClientError
+
+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.acp.acp import ACPService
+from agentex.lib.core.services.adk.agent_task_tracker import AgentTaskTrackerService
+from agentex.lib.core.services.adk.events import EventsService
+from agentex.lib.core.services.adk.messages import MessagesService
+from agentex.lib.core.services.adk.providers.litellm import LiteLLMService
+from agentex.lib.core.services.adk.providers.openai import OpenAIService
+from agentex.lib.core.services.adk.providers.sgp import SGPService
+from agentex.lib.core.services.adk.state import StateService
+from agentex.lib.core.services.adk.streaming import StreamingService
+from agentex.lib.core.services.adk.tasks import TasksService
+from agentex.lib.core.services.adk.tracing import TracingService
+from agentex.lib.core.services.adk.utils.templating import TemplatingService
+from agentex.lib.core.temporal.activities.adk.acp.acp_activities import ACPActivities
+from agentex.lib.core.temporal.activities.adk.agent_task_tracker_activities import (
+    AgentTaskTrackerActivities,
+)
+from agentex.lib.core.temporal.activities.adk.events_activities import EventsActivities
+from agentex.lib.core.temporal.activities.adk.messages_activities import MessagesActivities
+from agentex.lib.core.temporal.activities.adk.providers.litellm_activities import (
+    LiteLLMActivities,
+)
+from agentex.lib.core.temporal.activities.adk.providers.openai_activities import (
+    OpenAIActivities,
+)
+from agentex.lib.core.temporal.activities.adk.providers.sgp_activities import SGPActivities
+from agentex.lib.core.temporal.activities.adk.state_activities import StateActivities
+from agentex.lib.core.temporal.activities.adk.streaming_activities import (
+    StreamingActivities,
+)
+from agentex.lib.core.temporal.activities.adk.tasks_activities import TasksActivities
+from agentex.lib.core.temporal.activities.adk.tracing_activities import TracingActivities
+from agentex.lib.core.temporal.activities.adk.utils.templating_activities import (
+    TemplatingActivities,
+)
+from agentex.lib.core.tracing import AsyncTracer
+
+
+def get_all_activities(sgp_client=None):
+    """
+    Returns a list of all standard activity functions that can be directly passed to worker.run().
+
+    Args:
+        sgp_client: Optional SGP client instance. If not provided, SGP activities will not be included.
+
+    Returns:
+        list: A list of activity functions ready to be passed to worker.run()
+    """
+    # Initialize common dependencies
+    try:
+        sgp_client = SGPClient()
+    except SGPClientError:
+        sgp_client = None
+
+    llm_gateway = LiteLLMGateway()
+    stream_repository = RedisEventStreamRepository()
+    agentex_client = AsyncAgentex()
+    tracer = AsyncTracer(agentex_client)
+
+    # Services
+
+    ## ADK
+    streaming_service = StreamingService(
+        agentex_client=agentex_client,
+        stream_repository=stream_repository,
+    )
+    messages_service = MessagesService(
+        agentex_client=agentex_client,
+        streaming_service=streaming_service,
+        tracer=tracer,
+    )
+    events_service = EventsService(
+        agentex_client=agentex_client,
+        tracer=tracer,
+    )
+    agent_task_tracker_service = AgentTaskTrackerService(
+        agentex_client=agentex_client,
+        tracer=tracer,
+    )
+    state_service = StateService(
+        agentex_client=agentex_client,
+        tracer=tracer,
+    )
+    tasks_service = TasksService(
+        agentex_client=agentex_client,
+        tracer=tracer,
+    )
+    tracing_service = TracingService(
+        tracer=tracer,
+    )
+
+    ## ACP
+    acp_service = ACPService(
+        agentex_client=agentex_client,
+        tracer=tracer,
+    )
+
+    ## Providers
+    litellm_service = LiteLLMService(
+        agentex_client=agentex_client,
+        llm_gateway=llm_gateway,
+        streaming_service=streaming_service,
+        tracer=tracer,
+    )
+    openai_service = OpenAIService(
+        agentex_client=agentex_client,
+        streaming_service=streaming_service,
+        tracer=tracer,
+    )
+    sgp_service = None
+    if sgp_client is not None:
+        sgp_service = SGPService(
+            sgp_client=sgp_client,
+            tracer=tracer,
+        )
+
+    ## Utils
+    templating_service = TemplatingService(
+        tracer=tracer,
+    )
+
+    # ADK
+
+    ## Core activities
+    messages_activities = MessagesActivities(messages_service=messages_service)
+    events_activities = EventsActivities(events_service=events_service)
+    agent_task_tracker_activities = AgentTaskTrackerActivities(
+        agent_task_tracker_service=agent_task_tracker_service
+    )
+    state_activities = StateActivities(state_service=state_service)
+    streaming_activities = StreamingActivities(streaming_service=streaming_service)
+    tasks_activities = TasksActivities(tasks_service=tasks_service)
+    tracing_activities = TracingActivities(tracing_service=tracing_service)
+
+    ## ACP
+    acp_activities = ACPActivities(acp_service=acp_service)
+
+    ## Providers
+    litellm_activities = LiteLLMActivities(litellm_service=litellm_service)
+    openai_activities = OpenAIActivities(openai_service=openai_service)
+    if sgp_client is not None:
+        sgp_activities = SGPActivities(sgp_service=sgp_service)
+    else:
+        sgp_activities = None
+
+    ## Utils
+    templating_activities = TemplatingActivities(templating_service=templating_service)
+
+    # Build list of standard activities
+    activities = [
+        # Core activities
+        ## Messages activities
+        messages_activities.create_message,
+        messages_activities.update_message,
+        messages_activities.create_messages_batch,
+        messages_activities.update_messages_batch,
+        messages_activities.list_messages,
+        ## Events activities
+        events_activities.get_event,
+        events_activities.list_events,
+        ## Agent Task Tracker activities
+        agent_task_tracker_activities.get_agent_task_tracker,
+        agent_task_tracker_activities.get_agent_task_tracker_by_task_and_agent,
+        agent_task_tracker_activities.update_agent_task_tracker,
+        ## State activities
+        state_activities.create_state,
+        state_activities.get_state,
+        state_activities.update_state,
+        state_activities.delete_state,
+        ## Streaming activities
+        streaming_activities.stream_update,
+        ## Tasks activities
+        tasks_activities.get_task,
+        tasks_activities.delete_task,
+        ## Tracing activities
+        tracing_activities.start_span,
+        tracing_activities.end_span,
+        # ACP activities
+        acp_activities.task_create,
+        acp_activities.message_send,
+        acp_activities.event_send,
+        acp_activities.task_cancel,
+        # Providers
+        ## LiteLLM activities
+        litellm_activities.chat_completion,
+        litellm_activities.chat_completion_auto_send,
+        litellm_activities.chat_completion_stream_auto_send,
+        ## OpenAI activities
+        openai_activities.run_agent,
+        openai_activities.run_agent_auto_send,
+        openai_activities.run_agent_streamed_auto_send,
+        # Utils
+        templating_activities.render_jinja,
+    ]
+
+    # SGP activities
+    if sgp_client is not None:
+        sgp_all_activities = [
+            sgp_activities.download_file_content,
+        ]
+        activities.extend(sgp_all_activities)
+
+    return activities

agentex/lib/core/temporal/activities/activity_helpers.py

@@ -0,0 +1,37 @@
+from datetime import timedelta
+from typing import Any, TypeVar
+
+from pydantic import TypeAdapter
+from temporalio import workflow
+from temporalio.common import RetryPolicy
+
+from agentex.lib.utils.model_utils import BaseModel
+
+T = TypeVar("T", bound="BaseModel")
+
+
+class ActivityHelpers:
+    @staticmethod
+    async def execute_activity(
+        activity_name: str,
+        request: BaseModel | str | int | float | bool | dict[str, Any] | list[Any],
+        response_type: Any,
+        start_to_close_timeout: timedelta | None = None,
+        heartbeat_timeout: timedelta | None = None,
+        retry_policy: RetryPolicy | None = None,
+    ) -> Any:
+        if start_to_close_timeout is None:
+            start_to_close_timeout = timedelta(seconds=10)
+        if retry_policy is None:
+            retry_policy = RetryPolicy(maximum_attempts=0)
+
+        response = await workflow.execute_activity(
+            activity=activity_name,
+            arg=request,
+            start_to_close_timeout=start_to_close_timeout,
+            retry_policy=retry_policy,
+            heartbeat_timeout=heartbeat_timeout,
+        )
+
+        adapter = TypeAdapter(response_type)
+        return adapter.validate_python(response)

agentex/lib/core/temporal/activities/adk/acp/acp_activities.py

@@ -0,0 +1,86 @@
+from enum import Enum
+from typing import Any
+
+from temporalio import activity
+
+from agentex.lib.core.services.adk.acp.acp import ACPService
+from agentex.types.event import Event
+from agentex.types.task_message import TaskMessage
+from agentex.types.task_message_content import TaskMessageContent
+from agentex.types.task import Task
+from agentex.lib.types.tracing import BaseModelWithTraceParams
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class ACPActivityName(str, Enum):
+    TASK_CREATE = "task-create"
+    MESSAGE_SEND = "message-send"
+    EVENT_SEND = "event-send"
+    TASK_CANCEL = "task-cancel"
+
+
+class TaskCreateParams(BaseModelWithTraceParams):
+    name: str | None = None
+    agent_id: str | None = None
+    agent_name: str | None = None
+    params: dict[str, Any] | None = None
+
+
+class MessageSendParams(BaseModelWithTraceParams):
+    agent_id: str | None = None
+    agent_name: str | None = None
+    task_id: str | None = None
+    content: TaskMessageContent
+
+
+class EventSendParams(BaseModelWithTraceParams):
+    agent_id: str | None = None
+    agent_name: str | None = None
+    task_id: str | None = None
+    content: TaskMessageContent
+
+
+class TaskCancelParams(BaseModelWithTraceParams):
+    task_id: str | None = None
+    task_name: str | None = None
+
+
+class ACPActivities:
+    def __init__(self, acp_service: ACPService):
+        self._acp_service = acp_service
+
+    @activity.defn(name=ACPActivityName.TASK_CREATE)
+    async def task_create(self, params: TaskCreateParams) -> Task:
+        return await self._acp_service.task_create(
+            name=params.name,
+            agent_id=params.agent_id,
+            agent_name=params.agent_name,
+            params=params.params,
+        )
+
+    @activity.defn(name=ACPActivityName.MESSAGE_SEND)
+    async def message_send(self, params: MessageSendParams) -> TaskMessage:
+        return await self._acp_service.message_send(
+            agent_id=params.agent_id,
+            agent_name=params.agent_name,
+            task_id=params.task_id,
+            content=params.content,
+        )
+
+    @activity.defn(name=ACPActivityName.EVENT_SEND)
+    async def event_send(self, params: EventSendParams) -> Event:
+        return await self._acp_service.event_send(
+            agent_id=params.agent_id,
+            agent_name=params.agent_name,
+            task_id=params.task_id,
+            content=params.content,
+        )
+
+    @activity.defn(name=ACPActivityName.TASK_CANCEL)
+    async def task_cancel(self, params: TaskCancelParams) -> Task:
+        return await self._acp_service.task_cancel(
+            task_id=params.task_id,
+            task_name=params.task_name,
+        )

agentex/lib/core/temporal/activities/adk/agent_task_tracker_activities.py

@@ -0,0 +1,76 @@
+from enum import Enum
+
+from temporalio import activity
+
+from agentex.lib.core.services.adk.agent_task_tracker import AgentTaskTrackerService
+from agentex.lib.types.tracing import BaseModelWithTraceParams
+from agentex.lib.utils.logging import make_logger
+from agentex.types.agent_task_tracker import AgentTaskTracker
+
+logger = make_logger(__name__)
+
+
+class AgentTaskTrackerActivityName(str, Enum):
+    GET_AGENT_TASK_TRACKER = "get-agent-task-tracker"
+    GET_AGENT_TASK_TRACKER_BY_TASK_AND_AGENT = (
+        "get-agent-task-tracker-by-task-and-agent"
+    )
+    UPDATE_AGENT_TASK_TRACKER = "update-agent-task-tracker"
+
+
+class GetAgentTaskTrackerParams(BaseModelWithTraceParams):
+    tracker_id: str
+
+
+class GetAgentTaskTrackerByTaskAndAgentParams(BaseModelWithTraceParams):
+    task_id: str
+    agent_id: str
+
+
+class UpdateAgentTaskTrackerParams(BaseModelWithTraceParams):
+    tracker_id: str
+    last_processed_event_id: str | None
+    status: str | None
+    status_reason: str | None
+
+
+class AgentTaskTrackerActivities:
+    def __init__(self, agent_task_tracker_service: AgentTaskTrackerService):
+        self._agent_task_tracker_service = agent_task_tracker_service
+
+    @activity.defn(name=AgentTaskTrackerActivityName.GET_AGENT_TASK_TRACKER)
+    async def get_agent_task_tracker(
+        self, params: GetAgentTaskTrackerParams
+    ) -> AgentTaskTracker:
+        return await self._agent_task_tracker_service.get_agent_task_tracker(
+            tracker_id=params.tracker_id,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )
+
+    @activity.defn(
+        name=AgentTaskTrackerActivityName.GET_AGENT_TASK_TRACKER_BY_TASK_AND_AGENT
+    )
+    async def get_agent_task_tracker_by_task_and_agent(
+        self,
+        params: GetAgentTaskTrackerByTaskAndAgentParams,
+    ) -> AgentTaskTracker | None:
+        return await self._agent_task_tracker_service.get_by_task_and_agent(
+            task_id=params.task_id,
+            agent_id=params.agent_id,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )
+
+    @activity.defn(name=AgentTaskTrackerActivityName.UPDATE_AGENT_TASK_TRACKER)
+    async def update_agent_task_tracker(
+        self, params: UpdateAgentTaskTrackerParams
+    ) -> AgentTaskTracker:
+        return await self._agent_task_tracker_service.update_agent_task_tracker(
+            tracker_id=params.tracker_id,
+            last_processed_event_id=params.last_processed_event_id,
+            status=params.status,
+            status_reason=params.status_reason,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )

agentex/lib/core/temporal/activities/adk/agents_activities.py

@@ -0,0 +1,35 @@
+from enum import Enum
+from typing import Optional
+
+from agentex.lib.core.services.adk.agents import AgentsService
+from agentex.types.agent import Agent
+from temporalio import activity
+
+from agentex.lib.types.tracing import BaseModelWithTraceParams
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class AgentsActivityName(str, Enum):
+    GET_AGENT = "get-agent"
+
+
+class GetAgentParams(BaseModelWithTraceParams):
+    agent_id: Optional[str] = None
+    agent_name: Optional[str] = None
+
+
+class AgentsActivities:
+    def __init__(self, agents_service: AgentsService):
+        self._agents_service = agents_service
+
+    @activity.defn(name=AgentsActivityName.GET_AGENT)
+    async def get_agent(self, params: GetAgentParams) -> Agent | None:
+        return await self._agents_service.get_agent(
+            agent_id=params.agent_id,
+            agent_name=params.agent_name,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )
+

agentex/lib/core/temporal/activities/adk/events_activities.py

@@ -0,0 +1,50 @@
+from enum import Enum
+
+from temporalio import activity
+
+from agentex.lib.core.services.adk.events import EventsService
+from agentex.types.event import Event
+from agentex.lib.types.tracing import BaseModelWithTraceParams
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class EventsActivityName(str, Enum):
+    GET_EVENT = "get-event"
+    LIST_EVENTS = "list-events"
+
+
+class GetEventParams(BaseModelWithTraceParams):
+    event_id: str
+
+
+class ListEventsParams(BaseModelWithTraceParams):
+    task_id: str
+    agent_id: str
+    last_processed_event_id: str | None = None
+    limit: int | None = None
+
+
+class EventsActivities:
+    def __init__(self, events_service: EventsService):
+        self._events_service = events_service
+
+    @activity.defn(name=EventsActivityName.GET_EVENT)
+    async def get_event(self, params: GetEventParams) -> Event | None:
+        return await self._events_service.get_event(
+            event_id=params.event_id,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )
+
+    @activity.defn(name=EventsActivityName.LIST_EVENTS)
+    async def list_events(self, params: ListEventsParams) -> list[Event]:
+        return await self._events_service.list_events(
+            task_id=params.task_id,
+            agent_id=params.agent_id,
+            last_processed_event_id=params.last_processed_event_id,
+            limit=params.limit,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )

agentex/lib/core/temporal/activities/adk/messages_activities.py

@@ -0,0 +1,94 @@
+from enum import Enum
+
+from temporalio import activity
+
+from agentex.lib.core.services.adk.messages import MessagesService
+from agentex.types.task_message import TaskMessage
+from agentex.types.task_message_content import TaskMessageContent
+from agentex.lib.types.tracing import BaseModelWithTraceParams
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class MessagesActivityName(str, Enum):
+    CREATE_MESSAGE = "create-message"
+    UPDATE_MESSAGE = "update-message"
+    CREATE_MESSAGES_BATCH = "create-messages-batch"
+    UPDATE_MESSAGES_BATCH = "update-messages-batch"
+    LIST_MESSAGES = "list-messages"
+
+
+class CreateMessageParams(BaseModelWithTraceParams):
+    task_id: str
+    content: TaskMessageContent
+    emit_updates: bool = True
+
+
+class UpdateMessageParams(BaseModelWithTraceParams):
+    task_id: str
+    message_id: str
+    content: TaskMessageContent
+
+
+class CreateMessagesBatchParams(BaseModelWithTraceParams):
+    task_id: str
+    contents: list[TaskMessageContent]
+    emit_updates: bool = True
+
+
+class UpdateMessagesBatchParams(BaseModelWithTraceParams):
+    task_id: str
+    updates: dict[str, TaskMessageContent]
+
+
+class ListMessagesParams(BaseModelWithTraceParams):
+    task_id: str
+    limit: int | None = None
+
+
+class MessagesActivities:
+    def __init__(self, messages_service: MessagesService):
+        self._messages_service = messages_service
+
+    @activity.defn(name=MessagesActivityName.CREATE_MESSAGE)
+    async def create_message(self, params: CreateMessageParams) -> TaskMessage:
+        return await self._messages_service.create_message(
+            task_id=params.task_id,
+            content=params.content,
+            emit_updates=params.emit_updates,
+        )
+
+    @activity.defn(name=MessagesActivityName.UPDATE_MESSAGE)
+    async def update_message(self, params: UpdateMessageParams) -> TaskMessage:
+        return await self._messages_service.update_message(
+            task_id=params.task_id,
+            message_id=params.message_id,
+            content=params.content,
+        )
+
+    @activity.defn(name=MessagesActivityName.CREATE_MESSAGES_BATCH)
+    async def create_messages_batch(
+        self, params: CreateMessagesBatchParams
+    ) -> list[TaskMessage]:
+        return await self._messages_service.create_messages_batch(
+            task_id=params.task_id,
+            contents=params.contents,
+            emit_updates=params.emit_updates,
+        )
+
+    @activity.defn(name=MessagesActivityName.UPDATE_MESSAGES_BATCH)
+    async def update_messages_batch(
+        self, params: UpdateMessagesBatchParams
+    ) -> list[TaskMessage]:
+        return await self._messages_service.update_messages_batch(
+            task_id=params.task_id,
+            updates=params.updates,
+        )
+
+    @activity.defn(name=MessagesActivityName.LIST_MESSAGES)
+    async def list_messages(self, params: ListMessagesParams) -> list[TaskMessage]:
+        return await self._messages_service.list_messages(
+            task_id=params.task_id,
+            limit=params.limit,
+        )

agentex/lib/core/temporal/activities/adk/providers/litellm_activities.py

@@ -0,0 +1,71 @@
+from enum import Enum
+
+from temporalio import activity
+
+from agentex.lib.core.services.adk.providers.litellm import LiteLLMService
+from agentex.lib.types.llm_messages import Completion, LLMConfig
+from agentex.types.task_message import TaskMessage
+from agentex.lib.types.tracing import BaseModelWithTraceParams
+from agentex.lib.utils import logging
+
+logger = logging.make_logger(__name__)
+
+
+class LiteLLMActivityName(str, Enum):
+    CHAT_COMPLETION = "chat-completion"
+    CHAT_COMPLETION_AUTO_SEND = "chat-completion-auto-send"
+    # Note: CHAT_COMPLETION_STREAM is not supported in Temporal due to generator limitations
+    CHAT_COMPLETION_STREAM_AUTO_SEND = "chat-completion-stream-auto-send"
+
+
+class ChatCompletionParams(BaseModelWithTraceParams):
+    llm_config: LLMConfig
+
+
+class ChatCompletionAutoSendParams(BaseModelWithTraceParams):
+    task_id: str
+    llm_config: LLMConfig
+
+
+class ChatCompletionStreamAutoSendParams(BaseModelWithTraceParams):
+    task_id: str
+    llm_config: LLMConfig
+
+
+class LiteLLMActivities:
+    def __init__(self, litellm_service: LiteLLMService):
+        self._litellm_service = litellm_service
+
+    @activity.defn(name=LiteLLMActivityName.CHAT_COMPLETION)
+    async def chat_completion(self, params: ChatCompletionParams) -> Completion:
+        return await self._litellm_service.chat_completion(
+            llm_config=params.llm_config,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )
+
+    @activity.defn(name=LiteLLMActivityName.CHAT_COMPLETION_AUTO_SEND)
+    async def chat_completion_auto_send(self, params: ChatCompletionAutoSendParams) -> TaskMessage | None:
+        """
+        Activity for non-streaming chat completion with automatic TaskMessage creation.
+        """
+        return await self._litellm_service.chat_completion_auto_send(
+            task_id=params.task_id,
+            llm_config=params.llm_config,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )
+
+    @activity.defn(name=LiteLLMActivityName.CHAT_COMPLETION_STREAM_AUTO_SEND)
+    async def chat_completion_stream_auto_send(
+        self, params: ChatCompletionStreamAutoSendParams
+    ) -> TaskMessage | None:
+        """
+        Activity for streaming chat completion with automatic TaskMessage creation.
+        """
+        return await self._litellm_service.chat_completion_stream_auto_send(
+            task_id=params.task_id,
+            llm_config=params.llm_config,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
+        )