agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/core/tracing/trace.py

@@ -0,0 +1,311 @@
+import asyncio
+from contextlib import asynccontextmanager, contextmanager
+from datetime import UTC, datetime
+from typing import Any, AsyncGenerator
+import uuid
+
+from pydantic import BaseModel
+
+from agentex import Agentex, AsyncAgentex
+from agentex.lib.core.tracing.processors.tracing_processor_interface import (
+    AsyncTracingProcessor,
+    SyncTracingProcessor,
+)
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.model_utils import recursive_model_dump
+from agentex.types.span import Span
+
+logger = make_logger(__name__)
+
+
+class Trace:
+    """
+    Trace is a wrapper around the Agentex API for tracing.
+    It provides a context manager for spans and a way to start and end spans.
+    It also provides a way to get spans by ID and list all spans in a trace.
+    """
+
+    def __init__(
+        self,
+        processors: list[SyncTracingProcessor],
+        client: Agentex,
+        trace_id: str | None = None,
+    ):
+        """
+        Initialize a new trace with the specified trace ID.
+
+        Args:
+            trace_id: Required trace ID to use for this trace.
+            processors: Optional list of tracing processors to use for this trace.
+        """
+        self.processors = processors
+        self.client = client
+        self.trace_id = trace_id
+
+    def start_span(
+        self,
+        name: str,
+        parent_id: str | None = None,
+        input: dict[str, Any] | list[dict[str, Any]] | BaseModel | None = None,
+        data: dict[str, Any] | list[dict[str, Any]] | BaseModel | None = None,
+    ) -> Span:
+        """
+        Start a new span and register it with the API.
+
+        Args:
+            name: Name of the span.
+            parent_id: Optional parent span ID.
+            input: Optional input data for the span.
+            data: Optional additional data for the span.
+
+        Returns:
+            The newly created span.
+        """
+
+        if not self.trace_id:
+            raise ValueError("Trace ID is required to start a span")
+
+        # Create a span using the client's spans resource
+        start_time = datetime.now(UTC)
+
+        serialized_input = recursive_model_dump(input) if input else None
+        serialized_data = recursive_model_dump(data) if data else None
+        id = str(uuid.uuid4())
+
+        span = Span(
+            id=id,
+            trace_id=self.trace_id,
+            name=name,
+            parent_id=parent_id,
+            start_time=start_time,
+            input=serialized_input,
+            data=serialized_data,
+        )
+
+        for processor in self.processors:
+            processor.on_span_start(span)
+
+        return span
+
+    def end_span(
+        self,
+        span: Span,
+    ) -> Span:
+        """
+        End a span by updating it with any changes made to the span object.
+
+        Args:
+            span: The span object to update.
+
+        Returns:
+            The updated span.
+        """
+        if span.end_time is None:
+            span.end_time = datetime.now(UTC)
+
+        span.input = recursive_model_dump(span.input) if span.input else None
+        span.output = recursive_model_dump(span.output) if span.output else None
+        span.data = recursive_model_dump(span.data) if span.data else None
+
+        for processor in self.processors:
+            processor.on_span_end(span)
+
+        return span
+
+    def get_span(self, span_id: str) -> Span:
+        """
+        Get a span by ID.
+
+        Args:
+            span_id: The ID of the span to get.
+
+        Returns:
+            The requested span.
+        """
+        # Query from Agentex API
+        span = self.client.spans.retrieve(span_id)
+        return span
+
+    def list_spans(self) -> list[Span]:
+        """
+        List all spans in this trace.
+
+        Returns:
+            List of spans in this trace.
+        """
+        # Query from Agentex API
+        spans = self.client.spans.list(trace_id=self.trace_id)
+        return spans
+
+    @contextmanager
+    def span(
+        self,
+        name: str,
+        parent_id: str | None = None,
+        input: dict[str, Any] | list[dict[str, Any]] | BaseModel | None = None,
+        data: dict[str, Any] | list[dict[str, Any]] | BaseModel | None = None,
+    ):
+        """
+        Context manager for spans.
+        If trace_id is falsy, acts as a no-op context manager.
+        """
+        if not self.trace_id:
+            yield None
+            return
+        span = self.start_span(name, parent_id, input, data)
+        try:
+            yield span
+        finally:
+            self.end_span(span)
+
+
+class AsyncTrace:
+    """
+    AsyncTrace is a wrapper around the Agentex API for tracing.
+    It provides a context manager for spans and a way to start and end spans.
+    It also provides a way to get spans by ID and list all spans in a trace.
+    """
+
+    def __init__(
+        self,
+        processors: list[AsyncTracingProcessor],
+        client: AsyncAgentex,
+        trace_id: str | None = None,
+    ):
+        """
+        Initialize a new trace with the specified trace ID.
+
+        Args:
+            trace_id: Required trace ID to use for this trace.
+            processors: Optional list of tracing processors to use for this trace.
+        """
+        self.processors = processors
+        self.client = client
+        self.trace_id = trace_id
+
+    async def start_span(
+        self,
+        name: str,
+        parent_id: str | None = None,
+        input: dict[str, Any] | list[dict[str, Any]] | BaseModel | None = None,
+        data: dict[str, Any] | list[dict[str, Any]] | BaseModel | None = None,
+    ) -> Span:
+        """
+        Start a new span and register it with the API.
+
+        Args:
+            name: Name of the span.
+            parent_id: Optional parent span ID.
+            input: Optional input data for the span.
+            data: Optional additional data for the span.
+
+        Returns:
+            The newly created span.
+        """
+        if not self.trace_id:
+            raise ValueError("Trace ID is required to start a span")
+
+        # Create a span using the client's spans resource
+        start_time = datetime.now(UTC)
+
+        serialized_input = recursive_model_dump(input) if input else None
+        serialized_data = recursive_model_dump(data) if data else None
+        id = str(uuid.uuid4())
+
+        span = Span(
+            id=id,
+            trace_id=self.trace_id,
+            name=name,
+            parent_id=parent_id,
+            start_time=start_time,
+            input=serialized_input,
+            data=serialized_data,
+        )
+
+        if self.processors:
+            await asyncio.gather(
+                *[processor.on_span_start(span) for processor in self.processors]
+            )
+
+        return span
+
+    async def end_span(
+        self,
+        span: Span,
+    ) -> Span:
+        """
+        End a span by updating it with any changes made to the span object.
+
+        Args:
+            span: The span object to update.
+
+        Returns:
+            The updated span.
+        """
+        if span.end_time is None:
+            span.end_time = datetime.now(UTC)
+
+        span.input = recursive_model_dump(span.input) if span.input else None
+        span.output = recursive_model_dump(span.output) if span.output else None
+        span.data = recursive_model_dump(span.data) if span.data else None
+
+        if self.processors:
+            await asyncio.gather(
+                *[processor.on_span_end(span) for processor in self.processors]
+            )
+
+        return span
+
+    async def get_span(self, span_id: str) -> Span:
+        """
+        Get a span by ID.
+
+        Args:
+            span_id: The ID of the span to get.
+
+        Returns:
+            The requested span.
+        """
+        # Query from Agentex API
+        span = await self.client.spans.retrieve(span_id)
+        return span
+
+    async def list_spans(self) -> list[Span]:
+        """
+        List all spans in this trace.
+
+        Returns:
+            List of spans in this trace.
+        """
+        # Query from Agentex API
+        spans = await self.client.spans.list(trace_id=self.trace_id)
+        return spans
+
+    @asynccontextmanager
+    async def span(
+        self,
+        name: str,
+        parent_id: str | None = None,
+        input: dict[str, Any] | list[dict[str, Any]] | BaseModel | None = None,
+        data: dict[str, Any] | list[dict[str, Any]] | BaseModel | None = None,
+    ) -> AsyncGenerator[Span | None, None]:
+        """
+        Context manager for spans.
+
+        Args:
+            name: Name of the span.
+            parent_id: Optional parent span ID.
+            input: Optional input data for the span.
+            data: Optional additional data for the span.
+
+        Yields:
+            The span object.
+        """
+        if not self.trace_id:
+            yield None
+            return
+        span = await self.start_span(name, parent_id, input, data)
+        try:
+            yield span
+        finally:
+            await self.end_span(span)

agentex/lib/core/tracing/tracer.py

@@ -0,0 +1,70 @@
+from agentex import Agentex, AsyncAgentex
+from agentex.lib.core.tracing.trace import AsyncTrace, Trace
+from agentex.lib.core.tracing.tracing_processor_manager import (
+    get_async_tracing_processors,
+    get_sync_tracing_processors,
+)
+
+
+class Tracer:
+    """
+    Tracer is the main entry point for tracing in Agentex.
+    It manages the client connection and creates traces.
+    """
+
+    def __init__(self, client: Agentex):
+        """
+        Initialize a new sync tracer with the provided client.
+
+        Args:
+            client: Agentex client instance used for API communication.
+        """
+        self.client = client
+
+    def trace(self, trace_id: str | None = None) -> Trace:
+        """
+        Create a new trace with the given trace ID.
+
+        Args:
+            trace_id: The trace ID to use.
+
+        Returns:
+            A new Trace instance.
+        """
+        return Trace(
+            processors=get_sync_tracing_processors(),
+            client=self.client,
+            trace_id=trace_id,
+        )
+
+
+class AsyncTracer:
+    """
+    AsyncTracer is the async version of Tracer.
+    It manages the async client connection and creates async traces.
+    """
+
+    def __init__(self, client: AsyncAgentex):
+        """
+        Initialize a new async tracer with the provided client.
+
+        Args:
+            client: AsyncAgentex client instance used for API communication.
+        """
+        self.client = client
+
+    def trace(self, trace_id: str | None = None) -> AsyncTrace:
+        """
+        Create a new trace with the given trace ID.
+
+        Args:
+            trace_id: The trace ID to use.
+
+        Returns:
+            A new AsyncTrace instance.
+        """
+        return AsyncTrace(
+            processors=get_async_tracing_processors(),
+            client=self.client,
+            trace_id=trace_id,
+        )

agentex/lib/core/tracing/tracing_processor_manager.py

@@ -0,0 +1,62 @@
+from threading import Lock
+
+from agentex.lib.core.tracing.processors.agentex_tracing_processor import (
+    AgentexAsyncTracingProcessor,
+    AgentexSyncTracingProcessor,
+)
+from agentex.lib.core.tracing.processors.sgp_tracing_processor import (
+    SGPAsyncTracingProcessor,
+    SGPSyncTracingProcessor,
+)
+from agentex.lib.core.tracing.processors.tracing_processor_interface import (
+    AsyncTracingProcessor,
+    SyncTracingProcessor,
+)
+from agentex.lib.types.tracing import AgentexTracingProcessorConfig, TracingProcessorConfig
+
+
+class TracingProcessorManager:
+    def __init__(self):
+        # Mapping of processor config type to processor class
+        self.sync_config_registry: dict[str, type[SyncTracingProcessor]] = {
+            "agentex": AgentexSyncTracingProcessor,
+            "sgp": SGPSyncTracingProcessor,
+        }
+        self.async_config_registry: dict[str, type[AsyncTracingProcessor]] = {
+            "agentex": AgentexAsyncTracingProcessor,
+            "sgp": SGPAsyncTracingProcessor,
+        }
+        # Cache for processors
+        self.sync_processors: list[SyncTracingProcessor] = []
+        self.async_processors: list[AsyncTracingProcessor] = []
+        self.lock = Lock()
+
+    def add_processor_config(self, processor_config: TracingProcessorConfig) -> None:
+        with self.lock:
+            sync_processor = self.sync_config_registry[processor_config.type]
+            async_processor = self.async_config_registry[processor_config.type]
+            self.sync_processors.append(sync_processor(processor_config))
+            self.async_processors.append(async_processor(processor_config))
+
+    def set_processor_configs(self, processor_configs: list[TracingProcessorConfig]):
+        with self.lock:
+            for processor_config in processor_configs:
+                self.add_processor_config(processor_config)
+
+    def get_sync_processors(self) -> list[SyncTracingProcessor]:
+        return self.sync_processors
+
+    def get_async_processors(self) -> list[AsyncTracingProcessor]:
+        return self.async_processors
+
+
+# Global instance
+GLOBAL_TRACING_PROCESSOR_MANAGER = TracingProcessorManager()
+
+add_tracing_processor_config = GLOBAL_TRACING_PROCESSOR_MANAGER.add_processor_config
+set_tracing_processor_configs = GLOBAL_TRACING_PROCESSOR_MANAGER.set_processor_configs
+get_sync_tracing_processors = GLOBAL_TRACING_PROCESSOR_MANAGER.get_sync_processors
+get_async_tracing_processors = GLOBAL_TRACING_PROCESSOR_MANAGER.get_async_processors
+
+# Add the Agentex tracing processor by default
+add_tracing_processor_config(AgentexTracingProcessorConfig())

agentex/lib/environment_variables.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import os
+from enum import Enum
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+from agentex.lib.utils.model_utils import BaseModel
+
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+
+
+class EnvVarKeys(str, Enum):
+    ENVIRONMENT = "ENVIRONMENT"
+    TEMPORAL_ADDRESS = "TEMPORAL_ADDRESS"
+    REDIS_URL = "REDIS_URL"
+    AGENTEX_BASE_URL = "AGENTEX_BASE_URL"
+    # Agent Identifiers
+    AGENT_NAME = "AGENT_NAME"
+    AGENT_DESCRIPTION = "AGENT_DESCRIPTION"
+    AGENT_ID = "AGENT_ID"
+    # ACP Configuration
+    ACP_URL = "ACP_URL"
+    ACP_PORT = "ACP_PORT"
+    ACP_TYPE = "ACP_TYPE"
+    # Workflow Configuraiton
+    WORKFLOW_NAME = "WORKFLOW_NAME"
+    WORKFLOW_TASK_QUEUE = "WORKFLOW_TASK_QUEUE"
+
+
+class Environment(str, Enum):
+    DEV = "development"
+    STAGING = "staging"
+    PROD = "production"
+
+
+refreshed_environment_variables = None
+
+
+class EnvironmentVariables(BaseModel):
+    ENVIRONMENT: str = Environment.DEV
+    TEMPORAL_ADDRESS: str | None = "localhost:7233"
+    REDIS_URL: str | None = None
+    AGENTEX_BASE_URL: str | None = "http://localhost:5003"
+    # Agent Identifiers
+    AGENT_NAME: str
+    AGENT_DESCRIPTION: str | None = None
+    AGENT_ID: str | None = None
+    ACP_TYPE: str | None = "agentic"
+    # ACP Configuration
+    ACP_URL: str
+    ACP_PORT: int = 8000
+    # Workflow Configuration
+    WORKFLOW_TASK_QUEUE: str | None = None
+    WORKFLOW_NAME: str | None = None
+
+    @classmethod
+    def refresh(cls) -> EnvironmentVariables | None:
+        global refreshed_environment_variables
+        if refreshed_environment_variables is not None:
+            return refreshed_environment_variables
+
+        if os.environ.get(EnvVarKeys.ENVIRONMENT) == Environment.DEV:
+            # Load global .env file first
+            global_env_path = PROJECT_ROOT / ".env"
+            if global_env_path.exists():
+                print(f"Loading global environment variables FROM: {global_env_path}")
+                load_dotenv(dotenv_path=global_env_path, override=False)
+
+            # Load local project .env.local file (takes precedence)
+            local_env_path = Path.cwd().parent / ".env.local"
+            if local_env_path.exists():
+                print(f"Loading local environment variables FROM: {local_env_path}")
+                load_dotenv(dotenv_path=local_env_path, override=True)
+
+        # Create kwargs dict with environment variables, using None for missing values
+        # Pydantic will use the default values when None is passed for optional fields
+        kwargs = {}
+        for key in EnvVarKeys:
+            env_value = os.environ.get(key.value)
+            if env_value is not None:
+                kwargs[key.value] = env_value
+
+        environment_variables = EnvironmentVariables(**kwargs)
+        refreshed_environment_variables = environment_variables
+        return refreshed_environment_variables

agentex/lib/sdk/config/agent_config.py

@@ -0,0 +1,61 @@
+from typing import Any, Literal
+
+from pydantic import Field
+
+from agentex.lib.types.agent_configs import TemporalConfig, TemporalWorkflowConfig
+from agentex.lib.types.credentials import CredentialMapping
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.utils.model_utils import BaseModel
+
+logger = make_logger(__name__)
+
+
+class AgentConfig(BaseModel):
+    name: str = Field(
+        ...,
+        description="The name of the agent.",
+        pattern=r"^[a-z0-9-]+$",
+    )
+    acp_type: Literal["sync", "agentic"] = Field(..., description="The type of agent.")
+    description: str = Field(..., description="The description of the agent.")
+    env: dict[str, str] | None = Field(
+        default=None, description="Environment variables to set directly in the agent deployment"
+    )
+    credentials: list[CredentialMapping | dict[str, Any]] | None = Field(
+        default=None,
+        description="List of credential mappings to mount to the agent deployment. Supports both legacy format and new typed credentials.",
+    )
+    temporal: TemporalConfig | None = Field(
+        default=None, description="Temporal workflow configuration for this agent"
+    )
+
+    def is_temporal_agent(self) -> bool:
+        """Check if this agent uses Temporal workflows"""
+        # Check temporal config with enabled flag
+        if self.temporal and self.temporal.enabled:
+            return True
+        return False
+
+    def get_temporal_workflow_config(self) -> TemporalWorkflowConfig | None:
+        """Get temporal workflow configuration, checking both new and legacy formats"""
+        # Check new workflows list first
+        if self.temporal and self.temporal.enabled and self.temporal.workflows:
+            return self.temporal.workflows[0]  # Return first workflow for backward compatibility
+
+        # Check legacy single workflow
+        if self.temporal and self.temporal.enabled and self.temporal.workflow:
+            return self.temporal.workflow
+
+        return None
+
+    def get_temporal_workflows(self) -> list[TemporalWorkflowConfig]:
+        """Get all temporal workflow configurations"""
+        # Check new workflows list first
+        if self.temporal and self.temporal.enabled and self.temporal.workflows:
+            return self.temporal.workflows
+
+        # Check legacy single workflow
+        if self.temporal and self.temporal.enabled and self.temporal.workflow:
+            return [self.temporal.workflow]
+
+        return []