agentex-sdk 0.1.0a6__py3-none-any.whl

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

@@ -0,0 +1,99 @@
+import base64
+import os
+import tempfile
+
+from scale_gp import SGPClient
+
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.lib.types.files import FileContentResponse
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+
+logger = make_logger(__name__)
+
+
+class SGPService:
+    def __init__(self, sgp_client: SGPClient, tracer: AsyncTracer):
+        self.sgp_client = sgp_client
+        self.tracer = tracer
+
+    async def download_file_content(
+        self,
+        file_id: str,
+        filename: str,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> FileContentResponse:
+        """
+        Download file content from SGP.
+
+        Args:
+            file_id: The ID of the file to download.
+            filename: The filename of the file to download.
+            trace_id: The trace ID for tracing.
+            parent_span_id: The parent span ID for tracing.
+
+        Returns:
+            FileContentResponse with mime_type and base64_content for constructing LLM input.
+        """
+        trace = self.tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="download_file_content",
+            input={"file_id": file_id, "filename": filename},
+        ) as span:
+            logger.info(f"Downloading file content for file_id: {file_id}")
+            heartbeat_if_in_workflow("downloading file content")
+
+            # Get the SGP response
+            response = self.sgp_client.beta.files.content(file_id)
+            heartbeat_if_in_workflow("file content downloaded")
+
+            # Determine mime type based on file extension
+            mime_type = "application/pdf"  # Default
+            file_extension = os.path.splitext(filename)[1].lower()
+            if file_extension:
+                if file_extension == ".pdf":
+                    mime_type = "application/pdf"
+                elif file_extension in [".doc", ".docx"]:
+                    mime_type = "application/msword"
+                elif file_extension in [".txt", ".text"]:
+                    mime_type = "text/plain"
+                elif file_extension in [".png"]:
+                    mime_type = "image/png"
+                elif file_extension in [".jpg", ".jpeg"]:
+                    mime_type = "image/jpeg"
+
+            # Use a named temporary file - simpler approach
+            with tempfile.NamedTemporaryFile(suffix=file_extension) as temp_file:
+                heartbeat_if_in_workflow(f"saving to temp file: {temp_file.name}")
+
+                # Use write_to_file method if available
+                if hasattr(response, "write_to_file"):
+                    response.write_to_file(temp_file.name)
+                else:
+                    # Fallback to direct writing
+                    content_bytes = response.read()
+                    temp_file.write(content_bytes)
+                    temp_file.flush()
+
+                # Seek to beginning of file for reading
+                temp_file.seek(0)
+
+                # Read the file in binary mode - exactly like the example
+                data = temp_file.read()
+
+                # Encode to base64
+                base64_content = base64.b64encode(data).decode("utf-8")
+
+                result = FileContentResponse(
+                    mime_type=mime_type, base64_content=base64_content
+                )
+
+            # Record metadata for tracing
+            span.output = {
+                "file_id": file_id,
+                "mime_type": result.mime_type,
+                "content_size": len(result.base64_content),
+            }
+            return result

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

@@ -0,0 +1,120 @@
+from typing import Any, Dict
+
+from agentex import AsyncAgentex
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.state import State
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class StateService:
+    def __init__(
+        self, agentex_client: AsyncAgentex, tracer: AsyncTracer
+    ):
+        self._agentex_client = agentex_client
+        self._tracer = tracer
+
+    async def create_state(
+        self,
+        task_id: str,
+        agent_id: str,
+        state: dict[str, Any],
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> State:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="create_state",
+            input={"task_id": task_id, "agent_id": agent_id, "state": state},
+        ) as span:
+            state_model = await self._agentex_client.states.create(
+                task_id=task_id,
+                agent_id=agent_id,
+                state=state,
+            )
+            if span:
+                span.output = state_model.model_dump()
+            return state_model
+
+    async def get_state(
+        self,
+        state_id: str | None = None,
+        task_id: str | None = None,
+        agent_id: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> State | None:
+        trace = self._tracer.trace(trace_id) if self._tracer else None
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="get_state",
+            input={
+                "state_id": state_id,
+                "task_id": task_id,
+                "agent_id": agent_id,
+            },
+        ) as span:
+            if state_id:
+                state = await self._agentex_client.states.retrieve(state_id=state_id)
+            elif task_id and agent_id:
+                states = await self._agentex_client.states.list(
+                    task_id=task_id,
+                    agent_id=agent_id,
+                )
+                state = states[0] if states else None
+            else:
+                raise ValueError(
+                    "Must provide either state_id or both task_id and agent_id"
+                )
+            if span:
+                span.output = state.model_dump() if state else None
+            return state
+
+    async def update_state(
+        self,
+        state_id: str,
+        task_id: str,
+        agent_id: str,
+        state: Dict[str, object],
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> State:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="update_state",
+            input={
+                "state_id": state_id,
+                "task_id": task_id,
+                "agent_id": agent_id,
+                "state": state,
+            },
+        ) as span:
+            state_model = await self._agentex_client.states.update(
+                state_id=state_id,
+                task_id=task_id,
+                agent_id=agent_id,
+                state=state,
+            )
+            if span:
+                span.output = state_model.model_dump()
+            return state_model
+
+    async def delete_state(
+        self,
+        state_id: str,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> State:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="delete_state",
+            input={"state_id": state_id},
+        ) as span:
+            state = await self._agentex_client.states.delete(state_id)
+            if span:
+                span.output = state.model_dump()
+            return state

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

@@ -0,0 +1,262 @@
+import json
+from typing import Literal, cast
+
+from agentex import AsyncAgentex
+from agentex.lib.core.adapters.streams.port import EventStreamRepository
+from agentex.lib.types.task_message_updates import (
+    TaskMessageDelta, 
+    TaskMessageUpdate,
+    TextDelta,
+    DataDelta,
+    ToolRequestDelta,
+    ToolResponseDelta,
+    StreamTaskMessage,
+    StreamTaskMessageStart,
+    StreamTaskMessageDelta,
+    StreamTaskMessageFull,
+    StreamTaskMessageDone,
+)
+from agentex.lib.utils.logging import make_logger
+from agentex.types.data_content import DataContent
+from agentex.types.task_message import (
+    TaskMessage,
+    TaskMessageContent,
+)
+from agentex.types.task_message_content_param import TaskMessageContentParam
+from agentex.types.text_content import TextContent
+from agentex.types.tool_request_content import ToolRequestContent
+from agentex.types.tool_response_content import ToolResponseContent
+
+logger = make_logger(__name__)
+
+
+def _get_stream_topic(task_id: str) -> str:
+    return f"task:{task_id}"
+
+
+class DeltaAccumulator:
+    def __init__(self):
+        self._accumulated_deltas: list[TaskMessageDelta] = []
+        self._delta_type: Literal["text", "data", "tool_request", "tool_response"] | None = None
+
+    def add_delta(self, delta: TaskMessageDelta):
+        if self._delta_type is None:
+            if delta.type == "text":
+                self._delta_type = "text"
+            elif delta.type == "data":
+                self._delta_type = "data"
+            elif delta.type == "tool_request":
+                self._delta_type = "tool_request"
+            elif delta.type == "tool_response":
+                self._delta_type = "tool_response"
+            else:
+                raise ValueError(f"Unknown delta type: {delta.type}")
+        else:
+            if self._delta_type != delta.type:
+                raise ValueError(
+                    f"Delta type mismatch: {self._delta_type} != {delta.type}"
+                )
+
+        self._accumulated_deltas.append(delta)
+
+    def convert_to_content(self) -> TaskMessageContent:
+        if self._delta_type == "text":
+            # Type assertion: we know all deltas are TextDelta when _delta_type is TEXT
+            text_deltas = [delta for delta in self._accumulated_deltas if isinstance(delta, TextDelta)]
+            text_content_str = "".join(
+                [delta.text_delta or "" for delta in text_deltas]
+            )
+            return TextContent(
+                author="agent",
+                content=text_content_str,
+            )
+        elif self._delta_type == "data":
+            # Type assertion: we know all deltas are DataDelta when _delta_type is DATA
+            data_deltas = [delta for delta in self._accumulated_deltas if isinstance(delta, DataDelta)]
+            data_content_str = "".join(
+                [delta.data_delta or "" for delta in data_deltas]
+            )
+            try:
+                data = json.loads(data_content_str)
+            except json.JSONDecodeError as e:
+                raise ValueError(
+                    f"Accumulated data content is not valid JSON: {data_content_str}"
+                ) from e
+            return DataContent(
+                author="agent",
+                data=data,
+            )
+        elif self._delta_type == "tool_request":
+            # Type assertion: we know all deltas are ToolRequestDelta when _delta_type is TOOL_REQUEST
+            tool_request_deltas = [delta for delta in self._accumulated_deltas if isinstance(delta, ToolRequestDelta)]
+            arguments_content_str = "".join(
+                [delta.arguments_delta or "" for delta in tool_request_deltas]
+            )
+            try:
+                arguments = json.loads(arguments_content_str)
+            except json.JSONDecodeError as e:
+                raise ValueError(
+                    f"Accumulated tool request arguments is not valid JSON: {arguments_content_str}"
+                ) from e
+            return ToolRequestContent(
+                author="agent",
+                tool_call_id=tool_request_deltas[0].tool_call_id,
+                name=tool_request_deltas[0].name,
+                arguments=arguments,
+            )
+        elif self._delta_type == "tool_response":
+            # Type assertion: we know all deltas are ToolResponseDelta when _delta_type is TOOL_RESPONSE
+            tool_response_deltas = [delta for delta in self._accumulated_deltas if isinstance(delta, ToolResponseDelta)]
+            tool_response_content_str = "".join(
+                [delta.tool_response_delta or "" for delta in tool_response_deltas]
+            )
+            return ToolResponseContent(
+                author="agent",
+                tool_call_id=tool_response_deltas[0].tool_call_id,
+                name=tool_response_deltas[0].name,
+                content=tool_response_content_str,
+            )
+        else:
+            raise ValueError(f"Unknown delta type: {self._delta_type}")
+
+
+class StreamingTaskMessageContext:
+    def __init__(
+        self,
+        task_id: str,
+        initial_content: TaskMessageContent,
+        agentex_client: AsyncAgentex,
+        streaming_service: "StreamingService",
+    ):
+        self.task_id = task_id
+        self.initial_content = initial_content
+        self.task_message: TaskMessage | None = None
+        self._agentex_client = agentex_client
+        self._streaming_service = streaming_service
+        self._is_closed = False
+        self._delta_accumulator = DeltaAccumulator()
+
+    async def __aenter__(self) -> "StreamingTaskMessageContext":
+        return await self.open()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return await self.close()
+
+    async def open(self) -> "StreamingTaskMessageContext":
+        self._is_closed = False
+
+        self.task_message = await self._agentex_client.messages.create(
+            task_id=self.task_id,
+            content=self.initial_content.model_dump(),
+            streaming_status="IN_PROGRESS",
+        )
+
+        # Send the START event
+        start_event = StreamTaskMessageStart(
+            parent_task_message=self.task_message,
+            content=self.initial_content,
+        )
+        await self._streaming_service.stream_update(start_event)
+
+        return self
+
+    async def close(self) -> TaskMessage:
+        """Close the streaming context."""
+        if not self.task_message:
+            raise ValueError("Context not properly initialized - no task message")
+
+        if self._is_closed:
+            return self.task_message  # Already done
+
+        # Send the DONE event
+        done_event = StreamTaskMessageDone(parent_task_message=self.task_message)
+        await self._streaming_service.stream_update(done_event)
+
+        # Update the task message with the final content
+        if self._delta_accumulator._accumulated_deltas:
+            self.task_message.content = self._delta_accumulator.convert_to_content()
+
+        await self._agentex_client.messages.update(
+            task_id=self.task_id,
+            message_id=self.task_message.id,
+            content=self.task_message.content.model_dump(),
+            streaming_status="DONE",
+        )
+
+        # Mark the context as done
+        self._is_closed = True
+        return self.task_message
+
+    async def stream_update(
+        self, update: StreamTaskMessage
+    ) -> StreamTaskMessage | None:
+        """Stream an update to the repository."""
+        if self._is_closed:
+            raise ValueError("Context is already done")
+
+        if not self.task_message:
+            raise ValueError("Context not properly initialized - no task message")
+
+        if isinstance(update, StreamTaskMessageDelta):
+            if update.delta is not None:
+                self._delta_accumulator.add_delta(update.delta)
+
+        result = await self._streaming_service.stream_update(update)
+
+        if isinstance(update, StreamTaskMessageDone):
+            await self.close()
+            return update
+        elif isinstance(update, StreamTaskMessageFull):
+            await self._agentex_client.messages.update(
+                task_id=self.task_id,
+                message_id=update.parent_task_message.id,
+                content=update.content.model_dump(),
+                streaming_status="DONE",
+            )
+            self._is_closed = True
+        return result
+
+
+class StreamingService:
+    def __init__(
+        self,
+        agentex_client: AsyncAgentex,
+        stream_repository: EventStreamRepository,
+    ):
+        self._agentex_client = agentex_client
+        self._stream_repository = stream_repository
+
+    def streaming_task_message_context(
+        self,
+        task_id: str,
+        initial_content: TaskMessageContent,
+    ) -> StreamingTaskMessageContext:
+        return StreamingTaskMessageContext(
+            task_id=task_id,
+            initial_content=initial_content,
+            agentex_client=self._agentex_client,
+            streaming_service=self,
+        )
+
+    async def stream_update(
+        self, update: TaskMessageUpdate
+    ) -> TaskMessageUpdate | None:
+        """
+        Stream an update to the repository.
+
+        Args:
+            update: The update to stream
+
+        Returns:
+            True if event was streamed successfully, False otherwise
+        """
+        stream_topic = _get_stream_topic(update.parent_task_message.task_id)
+
+        try:
+            await self._stream_repository.send_event(
+                topic=stream_topic, event=update.model_dump(mode="json")  # type: ignore
+            )
+            return update
+        except Exception as e:
+            logger.exception(f"Failed to stream event: {e}")
+            return None

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

@@ -0,0 +1,69 @@
+from agentex import AsyncAgentex
+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 heartbeat_if_in_workflow
+
+logger = make_logger(__name__)
+
+
+class TasksService:
+    def __init__(
+        self,
+        agentex_client: AsyncAgentex,
+        tracer: AsyncTracer,
+    ):
+        self._agentex_client = agentex_client
+        self._tracer = tracer
+
+    async def get_task(
+        self,
+        task_id: str | None = None,
+        task_name: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> Task:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="get_task",
+            input={"task_id": task_id, "task_name": task_name},
+        ) as span:
+            heartbeat_if_in_workflow("get task")
+            if not task_id and not task_name:
+                raise ValueError("Either task_id or task_name must be provided.")
+            if task_id:
+                task_model = await self._agentex_client.tasks.retrieve(task_id=task_id)
+            elif task_name:
+                task_model = await self._agentex_client.tasks.retrieve_by_name(task_name=task_name)
+            else:
+                raise ValueError("Either task_id or task_name must be provided.")
+            if span:
+                span.output = task_model.model_dump()
+            return task_model
+
+    async def delete_task(
+        self,
+        task_id: str | None = None,
+        task_name: str | None = None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> Task:
+        trace = self._tracer.trace(trace_id) if self._tracer else None
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="delete_task",
+            input={"task_id": task_id, "task_name": task_name},
+        ) as span:
+            heartbeat_if_in_workflow("delete task")
+            if not task_id and not task_name:
+                raise ValueError("Either task_id or task_name must be provided.")
+            if task_id:
+                task_model = await self._agentex_client.tasks.delete(task_id=task_id)
+            elif task_name:
+                task_model = await self._agentex_client.tasks.delete_by_name(task_name=task_name)
+            else:
+                raise ValueError("Either task_id or task_name must be provided.")
+            if span:
+                span.output = task_model.model_dump()
+            return task_model

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

@@ -0,0 +1,36 @@
+from typing import Any
+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 heartbeat_if_in_workflow
+
+logger = make_logger(__name__)
+
+
+class TracingService:
+    def __init__(self, tracer: AsyncTracer):
+        self._tracer = tracer
+
+    async def start_span(
+        self,
+        trace_id: str,
+        name: str,
+        parent_id: str | None = None,
+        input: list[Any] | dict[str, Any] | BaseModel | None = None,
+        data: list[Any] | dict[str, Any] | BaseModel | None = None,
+    ) -> Span | None:
+        trace = self._tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_id,
+            name=name,
+            input=input or {},
+            data=data,
+        ) as span:
+            heartbeat_if_in_workflow("start span")
+            return span if span else None
+
+    async def end_span(self, trace_id: str, span: Span) -> Span:
+        trace = self._tracer.trace(trace_id)
+        await trace.end_span(span)
+        return span

agentex/lib/core/services/adk/utils/templating.py

@@ -0,0 +1,58 @@
+from datetime import datetime
+from typing import Any
+
+from jinja2 import BaseLoader, Environment
+
+from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+
+# Create a Jinja environment
+JINJA_ENV = Environment(
+    loader=BaseLoader(),
+    trim_blocks=True,
+    lstrip_blocks=True,
+    extensions=["jinja2.ext.do"],
+)
+
+
+class TemplatingService:
+    def __init__(self, tracer: AsyncTracer | None = None):
+        self.tracer = tracer
+
+    async def render_jinja(
+        self,
+        template: str,
+        variables: dict[str, Any],
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ) -> str:
+        """
+        Activity that renders a Jinja template with the provided data.
+
+        Args:
+            template: The template string to render.
+            variables: The variables to render the template with.
+            trace_id: The trace ID for tracing.
+            parent_span_id: The parent span ID for tracing.
+
+        Returns:
+            The rendered template as a string
+        """
+        trace = self.tracer.trace(trace_id)
+        async with trace.span(
+            parent_id=parent_span_id,
+            name="render_jinja",
+            input={"template": template, "variables": variables},
+        ) as span:
+            heartbeat_if_in_workflow("render jinja")
+            global_variables = {
+                "datetime": datetime,
+            }
+            jinja_template = JINJA_ENV.from_string(template, globals=global_variables)
+            try:
+                rendered_template = jinja_template.render(variables)
+                if span:
+                    span.output = {"jinja_output": rendered_template}
+                return rendered_template
+            except Exception as e:
+                raise ValueError(f"Error rendering Jinja template: {str(e)}") from e