agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/core/services/adk/agent_task_tracker.py

@@ -0,0 +1,85 @@
+from agentex import AsyncAgentex
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.lib.utils.logging import make_logger
+from agentex.types.agent_task_tracker import AgentTaskTracker
+
+logger = make_logger(__name__)
+
+
+class AgentTaskTrackerService:
+    def __init__(
+        self, agentex_client: AsyncAgentex, tracer: AsyncTracer,
+    ):
+        self._agentex_client = agentex_client
+        self._tracer = tracer
+
+    async def get_agent_task_tracker(
+        self,
+        tracker_id: str,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> AgentTaskTracker:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="get_agent_task_tracker",
+            input={"tracker_id": tracker_id},
+        ) as span:
+            tracker = await self._agentex_client.tracker.retrieve(
+                tracker_id
+            )
+            if span:
+                span.output = tracker.model_dump()
+            return tracker
+
+    async def get_by_task_and_agent(
+        self,
+        task_id: str,
+        agent_id: str,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> AgentTaskTracker | None:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="get_by_task_and_agent",
+            input={"task_id": task_id, "agent_id": agent_id},
+        ) as span:
+            trackers = await self._agentex_client.tracker.list(
+                task_id=task_id,
+                agent_id=agent_id,
+            )
+            tracker = trackers[0] if trackers else None
+            if span:
+                span.output = tracker.model_dump() if tracker else None
+            return tracker
+
+    async def update_agent_task_tracker(
+        self,
+        tracker_id: str,
+        last_processed_event_id: str | None = None,
+        status: str | None = None,
+        status_reason: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> AgentTaskTracker:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="update_agent_task_tracker",
+            input={
+                "tracker_id": tracker_id,
+                "last_processed_event_id": last_processed_event_id,
+                "status": status,
+                "status_reason": status_reason,
+            },
+        ) as span:
+            tracker = await self._agentex_client.tracker.update(
+                tracker_id=tracker_id,
+                last_processed_event_id=last_processed_event_id,
+                status=status,
+                status_reason=status_reason,
+            )
+            if span:
+                span.output = tracker.model_dump()
+            return tracker

agentex/lib/core/services/adk/agents.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+from agentex import AsyncAgentex
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.agent import Agent
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+
+logger = make_logger(__name__)
+
+
+class AgentsService:
+    def __init__(
+        self,
+        agentex_client: AsyncAgentex,
+        tracer: AsyncTracer,
+    ):
+        self._agentex_client = agentex_client
+        self._tracer = tracer
+
+    async def get_agent(
+        self,
+        agent_id: Optional[str] = None,
+        agent_name: Optional[str] = None,
+        trace_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+    ) -> Agent:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="get_agent",
+            input={"agent_id": agent_id, "agent_name": agent_name},
+        ) as span:
+            heartbeat_if_in_workflow("get agent")
+            if agent_id:
+                agent = await self._agentex_client.agents.retrieve(agent_id=agent_id)
+            elif agent_name:
+                agent = await self._agentex_client.agents.retrieve_by_name(agent_name=agent_name)
+            else:
+                raise ValueError("Either agent_id or agent_name must be provided")
+            if span:
+                span.output = agent.model_dump()
+            return agent

agentex/lib/core/services/adk/events.py

@@ -0,0 +1,61 @@
+from agentex import AsyncAgentex
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.event import Event
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class EventsService:
+    def __init__(
+        self, agentex_client: AsyncAgentex, tracer: AsyncTracer
+    ):
+        self._agentex_client = agentex_client
+        self._tracer = tracer
+
+    async def get_event(
+        self,
+        event_id: str,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> Event | None:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="get_event",
+            input={"event_id": event_id},
+        ) as span:
+            event = await self._agentex_client.events.retrieve(event_id=event_id)
+            if span:
+                span.output = event.model_dump()
+            return event
+
+    async def list_events(
+        self,
+        task_id: str,
+        agent_id: str,
+        last_processed_event_id: str | None = None,
+        limit: int | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> list[Event]:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="list_events",
+            input={
+                "task_id": task_id,
+                "agent_id": agent_id,
+                "last_processed_event_id": last_processed_event_id,
+                "limit": limit,
+            },
+        ) as span:
+            events = await self._agentex_client.events.list(
+                task_id=task_id,
+                agent_id=agent_id,
+                last_processed_event_id=last_processed_event_id,
+                limit=limit,
+            )
+            if span:
+                span.output = [event.model_dump() for event in events]
+            return events

agentex/lib/core/services/adk/messages.py

@@ -0,0 +1,164 @@
+import asyncio
+from typing import Any, Coroutine, cast
+
+from agentex import AsyncAgentex
+from agentex.lib.core.services.adk.streaming import StreamingService
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.lib.types.task_message_updates import StreamTaskMessageFull, TaskMessageUpdate
+from agentex.types.task_message import TaskMessage, TaskMessageContent
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+from agentex.types.task_message_content_param import TaskMessageContentParam
+
+logger = make_logger(__name__)
+
+
+class MessagesService:
+    def __init__(
+        self,
+        agentex_client: AsyncAgentex,
+        streaming_service: StreamingService,
+        tracer: AsyncTracer,
+    ):
+        self._agentex_client = agentex_client
+        self._streaming_service = streaming_service
+        self._tracer = tracer
+
+    async def create_message(
+        self,
+        task_id: str,
+        content: TaskMessageContent,
+        emit_updates: bool = True,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> TaskMessage:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="create_message",
+            input={"task_id": task_id, "message": content},
+        ) as span:
+            heartbeat_if_in_workflow("create message")
+            task_message = await self._agentex_client.messages.create(
+                task_id=task_id,
+                content=content.model_dump(),
+            )
+            if emit_updates:
+                await self._emit_updates([task_message])
+            if span:
+                span.output = task_message.model_dump()
+            return task_message
+
+    async def update_message(
+        self,
+        task_id: str,
+        message_id: str,
+        content: TaskMessageContent,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> TaskMessage:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="update_message",
+            input={
+                "task_id": task_id,
+                "message_id": message_id,
+                "message": content,
+            },
+        ) as span:
+            heartbeat_if_in_workflow("update message")
+            task_message = await self._agentex_client.messages.update(
+                task_id=task_id,
+                message_id=message_id,
+                content=content.model_dump(),
+            )
+            if span:
+                span.output = task_message.model_dump()
+            return task_message
+
+    async def create_messages_batch(
+        self,
+        task_id: str,
+        contents: list[TaskMessageContent],
+        emit_updates: bool = True,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> list[TaskMessage]:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="create_messages_batch",
+            input={"task_id": task_id, "messages": contents},
+        ) as span:
+            heartbeat_if_in_workflow("create messages batch")
+            task_messages = await self._agentex_client.messages.batch.create(
+                task_id=task_id,
+                contents=[content.model_dump() for content in contents],
+            )
+            if emit_updates:
+                await self._emit_updates(task_messages)
+            if span:
+                span.output = [task_message.model_dump() for task_message in task_messages]
+            return task_messages
+
+    async def update_messages_batch(
+        self,
+        task_id: str,
+        updates: dict[str, TaskMessageContent],
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> list[TaskMessage]:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="update_messages_batch",
+            input={"task_id": task_id, "updates": updates},
+        ) as span:
+            heartbeat_if_in_workflow("update messages batch")
+            task_messages = await self._agentex_client.messages.batch.update(
+                task_id=task_id,
+                updates={
+                    message_id: content.model_dump()
+                    for message_id, content in updates.items()
+                },
+            )
+            if span:
+                span.output = [task_message.model_dump() for task_message in task_messages]
+            return task_messages
+
+    async def list_messages(
+        self,
+        task_id: str,
+        limit: int | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> list[TaskMessage]:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="list_messages",
+            input={"task_id": task_id, "limit": limit},
+        ) as span:
+            heartbeat_if_in_workflow("list messages")
+            task_messages = await self._agentex_client.messages.list(
+                task_id=task_id,
+                limit=limit,
+            )
+            if span:
+                span.output = [task_message.model_dump() for task_message in task_messages]
+            return task_messages
+
+    async def _emit_updates(self, task_messages: list[TaskMessage]) -> None:
+        stream_update_handlers: list[Coroutine[Any, Any, TaskMessageUpdate | None]] = []
+        for task_message in task_messages:
+            stream_update_handler = self._streaming_service.stream_update(
+                update=StreamTaskMessageFull(
+                    type="full",
+                    parent_task_message=task_message,
+                    content=task_message.content,
+                )
+            )
+            stream_update_handlers.append(stream_update_handler)
+
+        await asyncio.gather(*stream_update_handlers)

agentex/lib/core/services/adk/providers/litellm.py

@@ -0,0 +1,256 @@
+from collections.abc import AsyncGenerator
+
+from agentex import AsyncAgentex
+from agentex.lib.core.adapters.llm.adapter_litellm import LiteLLMGateway
+from agentex.lib.core.services.adk.streaming import StreamingService
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.lib.types.llm_messages import (
+    Completion,
+    LLMConfig,
+)
+from agentex.lib.types.task_message_updates import (
+    StreamTaskMessageDelta,
+    StreamTaskMessageFull,
+    TextDelta,
+)
+from agentex.types.task_message import TaskMessage
+from agentex.types.task_message_content import TextContent
+from agentex.lib.utils import logging
+from agentex.lib.utils.completions import concat_completion_chunks
+from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+
+logger = logging.make_logger(__name__)
+
+
+class LiteLLMService:
+    def __init__(
+        self,
+        agentex_client: AsyncAgentex,
+        streaming_service: StreamingService,
+        tracer: AsyncTracer,
+        llm_gateway: LiteLLMGateway | None = None,
+    ):
+        self.agentex_client = agentex_client
+        self.llm_gateway = llm_gateway
+        self.streaming_service = streaming_service
+        self.tracer = tracer
+
+    async def chat_completion(
+        self,
+        llm_config: LLMConfig,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> Completion:
+        trace = self.tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="chat_completion",
+            input=llm_config.model_dump(),
+        ) as span:
+            heartbeat_if_in_workflow("chat completion")
+            if self.llm_gateway is None:
+                raise ValueError("LLM Gateway is not set")
+            completion = await self.llm_gateway.acompletion(**llm_config.model_dump())
+            if span:
+                span.output = completion.model_dump()
+            return completion
+
+    async def chat_completion_auto_send(
+        self,
+        task_id: str,
+        llm_config: LLMConfig,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> TaskMessage | None:
+        """
+        Chat completion with automatic TaskMessage creation. This does not stream the completion. To stream use chat_completion_stream_auto_send.
+
+        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).
+
+        Returns:
+            TaskMessage: A TaskMessage object
+        """
+
+        if llm_config.stream:
+            raise ValueError(
+                "LLM config must not have stream=True. To stream use `chat_completion_stream` or `chat_completion_stream_auto_send`."
+            )
+
+        if self.llm_gateway is None:
+            raise ValueError("LLM Gateway is not set")
+
+        trace = self.tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="chat_completion_auto_send",
+            input=llm_config.model_dump(),
+        ) as span:
+            heartbeat_if_in_workflow("chat completion auto send")
+
+            async with self.streaming_service.streaming_task_message_context(
+                task_id=task_id,
+                initial_content=TextContent(
+                    author="agent",
+                    content="",
+                    format="markdown",
+                ),
+            ) as streaming_context:
+                completion = await self.llm_gateway.acompletion(**llm_config.model_dump())
+                if (
+                    completion.choices
+                    and len(completion.choices) > 0
+                    and completion.choices[0].message
+                ):
+                    final_content = TextContent(
+                        author="agent",
+                        content=completion.choices[0].message.content or "",
+                        format="markdown",
+                    )
+                    await streaming_context.stream_update(
+                        update=StreamTaskMessageFull(
+                            parent_task_message=streaming_context.task_message,
+                            content=final_content,
+                        ),
+                    )
+                else:
+                    raise ValueError("No completion message returned from LLM")
+
+            if span:
+                if streaming_context.task_message:
+                    span.output = streaming_context.task_message.model_dump()
+            return streaming_context.task_message if streaming_context.task_message else None
+
+    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.
+
+        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.
+
+        Returns:
+            AsyncGenerator[Completion, None]: Generator yielding completion chunks
+
+        Raises:
+            ValueError: If called from within a Temporal workflow or if stream=False
+        """
+        if not llm_config.stream:
+            raise ValueError("LLM config must have stream=True for streaming")
+
+        if self.llm_gateway is None:
+            raise ValueError("LLM Gateway is not set")
+
+        trace = self.tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="chat_completion_stream",
+            input=llm_config.model_dump(),
+        ) as span:
+            # Direct streaming outside temporal - yield each chunk as it comes
+            chunks: list[Completion] = []
+            async for chunk in self.llm_gateway.acompletion_stream(
+                **llm_config.model_dump()
+            ):
+                chunks.append(chunk)
+                yield chunk
+            if span:
+                span.output = concat_completion_chunks(chunks).model_dump()
+
+    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,
+    ) -> 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).
+
+        Returns:
+            TaskMessage: A TaskMessage object
+        """
+        heartbeat_if_in_workflow("chat completion stream")
+
+        if self.llm_gateway is None:
+            raise ValueError("LLM Gateway is not set")
+
+        if not llm_config.stream:
+            llm_config.stream = True
+
+        trace = self.tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="chat_completion_stream",
+            input=llm_config.model_dump(),
+        ) as span:
+            # Use streaming context manager
+            async with self.streaming_service.streaming_task_message_context(
+                task_id=task_id,
+                initial_content=TextContent(
+                    author="agent",
+                    content="",
+                    format="markdown",
+                ),
+            ) as streaming_context:
+                # Get the streaming response
+                chunks = []
+                async for response in self.llm_gateway.acompletion_stream(
+                    **llm_config.model_dump()
+                ):
+                    heartbeat_if_in_workflow("chat completion streaming")
+                    if (
+                        response.choices
+                        and len(response.choices) > 0
+                        and response.choices[0].delta
+                    ):
+                        delta = response.choices[0].delta.content
+                        if delta:
+                            # Stream the chunk via the context manager
+                            await streaming_context.stream_update(
+                                update=StreamTaskMessageDelta(
+                                    parent_task_message=streaming_context.task_message,
+                                    delta=TextDelta(text_delta=delta),
+                                ),
+                            )
+                            heartbeat_if_in_workflow("content chunk streamed")
+
+                        # Store the chunk for final message assembly
+                        chunks.append(response)
+
+                # Update the final message content
+                complete_message = concat_completion_chunks(chunks)
+                if (
+                    complete_message
+                    and complete_message.choices
+                    and complete_message.choices[0].message
+                ):
+                    final_content = TextContent(
+                        author="agent",
+                        content=complete_message.choices[0].message.content or "",
+                    )
+                    await streaming_context.stream_update(
+                        update=StreamTaskMessageFull(
+                            parent_task_message=streaming_context.task_message,
+                            content=final_content,
+                        ),
+                    )
+
+                heartbeat_if_in_workflow("chat completion stream complete")
+
+            if span:
+                if streaming_context.task_message:
+                    span.output = streaming_context.task_message.model_dump()
+
+        return streaming_context.task_message if streaming_context.task_message else None