agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/core/temporal/workers/worker.py

@@ -0,0 +1,162 @@
+import dataclasses
+import datetime
+import os
+import uuid
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+from aiohttp import web
+from temporalio.client import Client
+from temporalio.converter import (
+    AdvancedJSONEncoder,
+    CompositePayloadConverter,
+    DataConverter,
+    DefaultPayloadConverter,
+    JSONPlainPayloadConverter,
+    JSONTypeConverter,
+    _JSONTypeConverterUnhandled,
+)
+from temporalio.runtime import OpenTelemetryConfig, Runtime, TelemetryConfig
+from temporalio.worker import (
+    UnsandboxedWorkflowRunner,
+    Worker,
+)
+
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class DateTimeJSONEncoder(AdvancedJSONEncoder):
+    def default(self, o: Any) -> Any:
+        if isinstance(o, datetime.datetime):
+            return o.isoformat()
+        return super().default(o)
+
+
+class DateTimeJSONTypeConverter(JSONTypeConverter):
+    def to_typed_value(
+        self, hint: type, value: Any
+    ) -> Any | None | _JSONTypeConverterUnhandled:
+        if hint == datetime.datetime:
+            return datetime.datetime.fromisoformat(value)
+        return JSONTypeConverter.Unhandled
+
+
+class DateTimePayloadConverter(CompositePayloadConverter):
+    def __init__(self) -> None:
+        json_converter = JSONPlainPayloadConverter(
+            encoder=DateTimeJSONEncoder,
+            custom_type_converters=[DateTimeJSONTypeConverter()],
+        )
+        super().__init__(
+            *[
+                c if not isinstance(c, JSONPlainPayloadConverter) else json_converter
+                for c in DefaultPayloadConverter.default_encoding_payload_converters
+            ]
+        )
+
+
+custom_data_converter = dataclasses.replace(
+    DataConverter.default,
+    payload_converter_class=DateTimePayloadConverter,
+)
+
+
+async def get_temporal_client(temporal_address: str, metrics_url: str = None) -> Client:
+    if not metrics_url:
+        client = await Client.connect(
+            target_host=temporal_address, data_converter=custom_data_converter
+        )
+    else:
+        runtime = Runtime(
+            telemetry=TelemetryConfig(metrics=OpenTelemetryConfig(url=metrics_url))
+        )
+        client = await Client.connect(
+            target_host=temporal_address,
+            data_converter=custom_data_converter,
+            runtime=runtime,
+        )
+    return client
+
+
+class AgentexWorker:
+    def __init__(
+        self,
+        task_queue,
+        max_workers: int = 10,
+        max_concurrent_activities: int = 10,
+        health_check_port: int = 80,
+    ):
+        self.task_queue = task_queue
+        self.activity_handles = []
+        self.max_workers = max_workers
+        self.max_concurrent_activities = max_concurrent_activities
+        self.health_check_server_running = False
+        self.healthy = False
+        self.health_check_port = health_check_port
+
+    async def run(
+        self,
+        activities: list[Callable],
+        workflow: type,
+    ):
+        await self.start_health_check_server()
+        temporal_client = await get_temporal_client(
+            temporal_address=os.environ.get("TEMPORAL_ADDRESS", "localhost:7233"),
+        )
+        worker = Worker(
+            client=temporal_client,
+            task_queue=self.task_queue,
+            activity_executor=ThreadPoolExecutor(max_workers=self.max_workers),
+            workflows=[workflow],
+            activities=activities,
+            workflow_runner=UnsandboxedWorkflowRunner(),
+            max_concurrent_activities=self.max_concurrent_activities,
+            build_id=str(uuid.uuid4()),
+        )
+
+        logger.info(f"Starting workers for task queue: {self.task_queue}")
+        # Eagerly set the worker status to healthy
+        self.healthy = True
+        logger.info(f"Running workers for task queue: {self.task_queue}")
+        await worker.run()
+
+    async def _health_check(self):
+        return web.json_response(self.healthy)
+
+    async def start_health_check_server(self):
+        if not self.health_check_server_running:
+            app = web.Application()
+            app.router.add_get("/readyz", lambda request: self._health_check())
+
+            # Disable access logging
+            runner = web.AppRunner(app, access_log=None)
+            await runner.setup()
+
+            try:
+                site = web.TCPSite(runner, "0.0.0.0", self.health_check_port)
+                await site.start()
+                logger.info(
+                    f"Health check server running on http://0.0.0.0:{self.health_check_port}/readyz"
+                )
+                self.health_check_server_running = True
+            except OSError as e:
+                logger.error(
+                    f"Failed to start health check server on port {self.health_check_port}: {e}"
+                )
+                # Try alternative port if default fails
+                try:
+                    alt_port = self.health_check_port + 1
+                    site = web.TCPSite(runner, "0.0.0.0", alt_port)
+                    await site.start()
+                    logger.info(
+                        f"Health check server running on alternative port http://0.0.0.0:{alt_port}/readyz"
+                    )
+                    self.health_check_server_running = True
+                except OSError as e:
+                    logger.error(
+                        f"Failed to start health check server on alternative port {alt_port}: {e}"
+                    )
+                    raise

agentex/lib/core/temporal/workflows/workflow.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+
+from temporalio import workflow
+
+from agentex.lib.core.temporal.types.workflow import SignalName
+from agentex.lib.types.acp import CreateTaskParams, SendEventParams
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class BaseWorkflow(ABC):
+    def __init__(
+        self,
+        display_name: str,
+    ):
+        self.display_name = display_name
+
+    @abstractmethod
+    @workflow.signal(name=SignalName.RECEIVE_EVENT)
+    async def on_task_event_send(self, params: SendEventParams) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    async def on_task_create(self, params: CreateTaskParams) -> None:
+        raise NotImplementedError

agentex/lib/core/tracing/__init__.py

@@ -0,0 +1,5 @@
+from agentex.lib.core.tracing.trace import AsyncTrace, Trace
+from agentex.lib.core.tracing.tracer import AsyncTracer, Tracer
+from agentex.types.span import Span
+
+__all__ = ["Trace", "AsyncTrace", "Span", "Tracer", "AsyncTracer"]

agentex/lib/core/tracing/processors/agentex_tracing_processor.py

@@ -0,0 +1,117 @@
+from typing import Any, Dict, override
+
+from agentex import Agentex, AsyncAgentex
+from agentex.lib.core.tracing.processors.tracing_processor_interface import (
+    AsyncTracingProcessor,
+    SyncTracingProcessor,
+)
+from agentex.types.span import Span
+from agentex.lib.types.tracing import AgentexTracingProcessorConfig
+
+
+class AgentexSyncTracingProcessor(SyncTracingProcessor):
+    def __init__(self, config: AgentexTracingProcessorConfig):
+        self.client = Agentex()
+
+    @override
+    def on_span_start(self, span: Span) -> None:
+        self.client.spans.create(
+            name=span.name,
+            start_time=span.start_time,
+            end_time=span.end_time,
+            trace_id=span.trace_id,
+            id=span.id,
+            data=span.data,
+            input=span.input,
+            output=span.output,
+            parent_id=span.parent_id,
+        )
+
+    @override
+    def on_span_end(self, span: Span) -> None:
+        update: Dict[str, Any] = {}
+        if span.trace_id:
+            update["trace_id"] = span.trace_id
+        if span.name:
+            update["name"] = span.name
+        if span.parent_id:
+            update["parent_id"] = span.parent_id
+        if span.start_time:
+            update["start_time"] = span.start_time.isoformat()
+        if span.end_time is not None:
+            update["end_time"] = span.end_time.isoformat()
+        if span.input is not None:
+            update["input"] = span.input
+        if span.output is not None:
+            update["output"] = span.output
+        if span.data is not None:
+            update["data"] = span.data
+
+        self.client.spans.update(
+            span.id,
+            **span.model_dump(
+                mode="json",
+                exclude={"id"},
+                exclude_defaults=True,
+                exclude_none=True,
+                exclude_unset=True,
+            ),
+        )
+
+    @override
+    def shutdown(self) -> None:
+        pass
+
+
+class AgentexAsyncTracingProcessor(AsyncTracingProcessor):
+    def __init__(self, config: AgentexTracingProcessorConfig):
+        self.client = AsyncAgentex()
+
+    @override
+    async def on_span_start(self, span: Span) -> None:
+        await self.client.spans.create(
+            name=span.name,
+            start_time=span.start_time,
+            end_time=span.end_time,
+            id=span.id,
+            trace_id=span.trace_id,
+            parent_id=span.parent_id,
+            input=span.input,
+            output=span.output,
+            data=span.data,
+        )
+
+    @override
+    async def on_span_end(self, span: Span) -> None:
+        update: Dict[str, Any] = {}
+        if span.trace_id:
+            update["trace_id"] = span.trace_id
+        if span.name:
+            update["name"] = span.name
+        if span.parent_id:
+            update["parent_id"] = span.parent_id
+        if span.start_time:
+            update["start_time"] = span.start_time.isoformat()
+        if span.end_time:
+            update["end_time"] = span.end_time.isoformat()
+        if span.input:
+            update["input"] = span.input
+        if span.output:
+            update["output"] = span.output
+        if span.data:
+            update["data"] = span.data
+
+        await self.client.spans.update(
+            span.id,
+            **span.model_dump(
+                mode="json",
+                exclude={"id"},
+                exclude_defaults=True,
+                exclude_none=True,
+                exclude_unset=True,
+            ),
+        )
+
+    @override
+    async def shutdown(self) -> None:
+        pass

agentex/lib/core/tracing/processors/sgp_tracing_processor.py

@@ -0,0 +1,119 @@
+from typing import override
+
+import scale_gp_beta.lib.tracing as tracing
+from scale_gp_beta import AsyncSGPClient, SGPClient
+from scale_gp_beta.lib.tracing import create_span, flush_queue
+from scale_gp_beta.lib.tracing.span import Span as SGPSpan
+
+from agentex.lib.core.tracing.processors.tracing_processor_interface import (
+    AsyncTracingProcessor,
+    SyncTracingProcessor,
+)
+from agentex.types.span import Span
+from agentex.lib.types.tracing import SGPTracingProcessorConfig
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+
+class SGPSyncTracingProcessor(SyncTracingProcessor):
+    def __init__(self, config: SGPTracingProcessorConfig):
+        disabled = config.sgp_api_key == "" or config.sgp_account_id == ""
+        tracing.init(
+            SGPClient(api_key=config.sgp_api_key, account_id=config.sgp_account_id),
+            disabled=disabled,
+        )
+        self._spans: dict[str, SGPSpan] = {}
+
+    @override
+    def on_span_start(self, span: Span) -> None:
+        sgp_span = create_span(
+            name=span.name,
+            span_id=span.id,
+            parent_id=span.parent_id,
+            trace_id=span.trace_id,
+            input=span.input,
+            output=span.output,
+            metadata=span.data,
+        )
+        sgp_span.start_time = span.start_time.isoformat()
+        sgp_span.flush(blocking=False)
+
+        self._spans[span.id] = sgp_span
+
+    @override
+    def on_span_end(self, span: Span) -> None:
+        sgp_span = self._spans.get(span.id)
+        if sgp_span is None:
+            logger.warning(
+                f"Span {span.id} not found in stored spans, skipping span end"
+            )
+            return
+
+        sgp_span.output = span.output
+        sgp_span.metadata = span.data
+        sgp_span.end_time = span.end_time.isoformat()
+        sgp_span.flush(blocking=False)
+
+    @override
+    def shutdown(self) -> None:
+        self._spans.clear()
+        flush_queue()
+
+
+class SGPAsyncTracingProcessor(AsyncTracingProcessor):
+    def __init__(self, config: SGPTracingProcessorConfig):
+        self.disabled = config.sgp_api_key == "" or config.sgp_account_id == ""
+        self._spans: dict[str, SGPSpan] = {}
+        self.sgp_async_client = (
+            AsyncSGPClient(api_key=config.sgp_api_key, account_id=config.sgp_account_id)
+            if not self.disabled
+            else None
+        )
+
+    @override
+    async def on_span_start(self, span: Span) -> None:
+        sgp_span = create_span(
+            name=span.name,
+            span_id=span.id,
+            parent_id=span.parent_id,
+            trace_id=span.trace_id,
+            input=span.input,
+            output=span.output,
+            metadata=span.data,
+        )
+        sgp_span.start_time = span.start_time.isoformat()
+
+        if self.disabled:
+            return
+        await self.sgp_async_client.spans.upsert_batch(
+            items=[sgp_span.to_request_params()]
+        )
+
+        self._spans[span.id] = sgp_span
+
+    @override
+    async def on_span_end(self, span: Span) -> None:
+        sgp_span = self._spans.get(span.id)
+        if sgp_span is None:
+            logger.warning(
+                f"Span {span.id} not found in stored spans, skipping span end"
+            )
+            return
+
+        sgp_span.output = span.output
+        sgp_span.metadata = span.data
+        sgp_span.end_time = span.end_time.isoformat()
+
+        if self.disabled:
+            return
+        await self.sgp_async_client.spans.upsert_batch(
+            items=[sgp_span.to_request_params()]
+        )
+
+    @override
+    async def shutdown(self) -> None:
+        await self.sgp_async_client.spans.upsert_batch(
+            items=[sgp_span.to_request_params() for sgp_span in self._spans.values()]
+        )
+        self._spans.clear()

agentex/lib/core/tracing/processors/tracing_processor_interface.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+
+from agentex.lib.types.tracing import TracingProcessorConfig
+from agentex.types.span import Span
+
+
+class SyncTracingProcessor(ABC):
+    @abstractmethod
+    def __init__(self, config: TracingProcessorConfig):
+        pass
+
+    @abstractmethod
+    def on_span_start(self, span: Span) -> None:
+        pass
+
+    @abstractmethod
+    def on_span_end(self, span: Span) -> None:
+        pass
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        pass
+
+
+class AsyncTracingProcessor(ABC):
+    @abstractmethod
+    def __init__(self, config: TracingProcessorConfig):
+        pass
+
+    @abstractmethod
+    async def on_span_start(self, span: Span) -> None:
+        pass
+
+    @abstractmethod
+    async def on_span_end(self, span: Span) -> None:
+        pass
+
+    @abstractmethod
+    async def shutdown(self) -> None:
+        pass