agentex-sdk 0.4.18__py3-none-any.whl → 0.4.20__py3-none-any.whl

agentex/lib/sdk/state_machine/state_machine.py

@@ -2,9 +2,9 @@ from abc import ABC, abstractmethod
 from typing import Any, Generic, TypeVar
 
 from agentex.lib import adk
+from agentex.lib.utils.model_utils import BaseModel
 from agentex.lib.sdk.state_machine.state import State
 from agentex.lib.sdk.state_machine.state_workflow import StateWorkflow
-from agentex.lib.utils.model_utils import BaseModel
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -55,7 +55,13 @@ class StateMachine(ABC, Generic[T]):
             raise ValueError(f"State {target_state_name} not found")
         self._current_state = self._state_map[target_state_name]
 
-    def get_state_machine_data(self) -> T:
+    def get_state_machine_data(self) -> T | None:
+        return self.state_machine_data
+
+    def require_state_machine_data(self) -> T:
+        """Get state machine data, raising an error if not set."""
+        if self.state_machine_data is None:
+            raise ValueError("State machine data not initialized - ensure data is provided")
         return self.state_machine_data
 
     @abstractmethod
@@ -70,7 +76,10 @@ class StateMachine(ABC, Generic[T]):
     async def step(self) -> str:
         current_state_name = self.get_current_state()
         current_state = self._state_map.get(current_state_name)
+        if current_state is None:
+            raise ValueError(f"Current state '{current_state_name}' not found in state map")
 
+        span = None
         if self._trace_transitions:
             if self._task_id is None:
                 raise ValueError(
@@ -79,7 +88,7 @@ class StateMachine(ABC, Generic[T]):
             span = await adk.tracing.start_span(
                 trace_id=self._task_id,
                 name="state_transition",
-                input=self.state_machine_data.model_dump(),
+                input=self.require_state_machine_data().model_dump(),
                 data={"input_state": current_state_name},
             )
 
@@ -87,13 +96,10 @@ class StateMachine(ABC, Generic[T]):
             state_machine=self, state_machine_data=self.state_machine_data
         )
 
-        if self._trace_transitions:
-            if self._task_id is None:
-                raise ValueError(
-                    "Task ID is must be set before tracing can be enabled"
-                )
-            span.output = self.state_machine_data.model_dump()
-            span.data["output_state"] = next_state_name
+        if self._trace_transitions and span is not None:
+            span.output = self.require_state_machine_data().model_dump()  # type: ignore[assignment]
+            if span.data is not None:
+                span.data["output_state"] = next_state_name  # type: ignore[index]
             await adk.tracing.end_span(trace_id=self._task_id, span=span)
 
         await self.transition(next_state_name)
@@ -118,7 +124,7 @@ class StateMachine(ABC, Generic[T]):
         await self.transition(self._initial_state)
 
         if self._trace_transitions:
-            span.output = {"output_state": self._initial_state}
+            span.output = {"output_state": self._initial_state}  # type: ignore[assignment,union-attr]
             await adk.tracing.end_span(trace_id=self._task_id, span=span)
 
     def dump(self) -> dict[str, Any]:
@@ -168,7 +174,7 @@ class StateMachine(ABC, Generic[T]):
             state_machine_data = None
             if state_machine_data_dict is not None:
                 # Get the actual model type from the class's type parameters
-                model_type = cls.__orig_bases__[0].__args__[0]
+                model_type = cls.__orig_bases__[0].__args__[0]  # type: ignore[attr-defined]
                 state_machine_data = model_type.model_validate(state_machine_data_dict)
 
             # Create a new instance

agentex/lib/sdk/utils/messages.py

@@ -2,17 +2,17 @@ import json
 from abc import ABC, abstractmethod
 from typing import Any, Literal, override
 
+from agentex.types.data_content import DataContent
+from agentex.types.task_message import TaskMessage
+from agentex.types.text_content import TextContent
 from agentex.lib.types.llm_messages import (
-    AssistantMessage,
     Message,
     ToolCall,
-    ToolCallRequest,
     ToolMessage,
     UserMessage,
+    ToolCallRequest,
+    AssistantMessage,
 )
-from agentex.types.data_content import DataContent
-from agentex.types.task_message import TaskMessage
-from agentex.types.text_content import TextContent
 from agentex.types.tool_request_content import ToolRequestContent
 from agentex.types.tool_response_content import ToolResponseContent
 

agentex/lib/types/acp.py

@@ -1,12 +1,12 @@
 from enum import Enum
 from typing import Any
 
-from pydantic import BaseModel, Field
+from pydantic import Field, BaseModel
 
+from agentex.types.task import Task
 from agentex.types.agent import Agent
 from agentex.types.event import Event
 from agentex.types.task_message_content import TaskMessageContent
-from agentex.types.task import Task
 
 
 class RPCMethod(str, Enum):

agentex/lib/types/agent_configs.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, Field, model_validator, validator
+from pydantic import Field, BaseModel, validator, model_validator
 
 
 class TemporalWorkflowConfig(BaseModel):

agentex/lib/types/converters.py

@@ -1,9 +1,11 @@
+import json
+
+from agents import TResponseInputItem
+
 from agentex.types.task_message import TaskMessage
 from agentex.types.text_content import TextContent
 from agentex.types.tool_request_content import ToolRequestContent
 from agentex.types.tool_response_content import ToolResponseContent
-import json
-from agents import TResponseInputItem
 
 
 def convert_task_messages_to_oai_agents_inputs(

agentex/lib/types/credentials.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, Field
+from pydantic import Field, BaseModel
 
 
 class CredentialMapping(BaseModel):

agentex/lib/types/fastacp.py

@@ -1,6 +1,7 @@
 from typing import Any, Literal
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import Field, BaseModel, field_validator
+
 from agentex.lib.core.clients.temporal.utils import validate_client_plugins
 
 
@@ -48,11 +49,9 @@ class TemporalACPConfig(AgenticACPConfig):
     """
 
     type: Literal["temporal"] = Field(default="temporal", frozen=True)
-    temporal_address: str = Field(
-        default="temporal-frontend.temporal.svc.cluster.local:7233", frozen=True
-    )
+    temporal_address: str = Field(default="temporal-frontend.temporal.svc.cluster.local:7233", frozen=True)
     plugins: list[Any] = Field(default=[], frozen=True)
-    
+
     @field_validator("plugins")
     @classmethod
     def validate_plugins(cls, v: list[Any]) -> list[Any]:

agentex/lib/types/tracing.py

@@ -1,4 +1,4 @@
-from typing import Annotated, Literal
+from typing import Literal, Annotated
 
 from pydantic import Field
 

agentex/lib/utils/completions.py

@@ -1,18 +1,18 @@
 from copy import deepcopy
-from functools import reduce, singledispatch
 from typing import Any
+from functools import reduce, singledispatch
 
 from agentex.lib.types.llm_messages import (
-    Choice,
-    Completion,
     Delta,
+    Choice,
     ToolCall,
+    Completion,
     ToolCallRequest,
 )
 
 
 @singledispatch
-def _concat_chunks(a: None, b: Any):
+def _concat_chunks(_a: None, b: Any):
     return b
 
 
@@ -116,7 +116,7 @@ def concat_completion_chunks(chunks: list[Completion]) -> Completion:
     single `CompletionChunk`. Finally we convert the type to the appropriate non-streaming type `Completion` and return it.
     """
     if not chunks:
-        return None
+        raise ValueError("Cannot concatenate empty chunks list")
 
     chunks_copy = chunks.copy()
     chunks_copy[0] = deepcopy(chunks_copy[0])  # _concat_chunks mutates first argument

agentex/lib/utils/console.py

@@ -1,6 +1,6 @@
 from rich import box
-from rich.console import Console
 from rich.table import Table
+from rich.console import Console
 
 console = Console()
 

agentex/lib/utils/debug.py

@@ -5,8 +5,13 @@ Provides debugging setup functionality that can be used across different compone
 """
 
 import os
+
 import debugpy  # type: ignore
 
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
 
 def setup_debug_if_enabled() -> None:
     """
@@ -33,14 +38,14 @@ def setup_debug_if_enabled() -> None:
         debugpy.configure(subProcess=False)
         debugpy.listen(debug_port)
         
-        print(f"🐛 [{debug_type.upper()}] Debug server listening on port {debug_port}")
-        
+        logger.info(f"🐛 [{debug_type.upper()}] Debug server listening on port {debug_port}")
+
         if wait_for_attach:
-            print(f"⏳ [{debug_type.upper()}] Waiting for debugger to attach...")
+            logger.info(f"⏳ [{debug_type.upper()}] Waiting for debugger to attach...")
             debugpy.wait_for_client()
-            print(f"✅ [{debug_type.upper()}] Debugger attached!")
+            logger.info(f"✅ [{debug_type.upper()}] Debugger attached!")
         else:
-            print(f"📡 [{debug_type.upper()}] Ready for debugger attachment")
+            logger.info(f"📡 [{debug_type.upper()}] Ready for debugger attachment")
 
 
 def is_debug_enabled() -> bool:

agentex/lib/utils/dev_tools/async_messages.py

@@ -6,26 +6,25 @@ to new streaming messages, handling mid-stream connections gracefully.
 """
 
 import json
-from datetime import datetime, timezone
 from typing import List, Optional
+from datetime import datetime, timezone
 
-from yaspin.core import Yaspin
+from yaspin import yaspin  # type: ignore[import-untyped]
+from rich.panel import Panel
+from yaspin.core import Yaspin  # type: ignore[import-untyped]
+from rich.console import Console
+from rich.markdown import Markdown
 
 from agentex import Agentex
-from agentex.types import Task, TaskMessage, TextContent, ToolRequestContent, ToolResponseContent, ReasoningContent
+from agentex.types import Task, TaskMessage, TextContent, ReasoningContent, ToolRequestContent, ToolResponseContent
+from agentex.types.text_delta import TextDelta
 from agentex.types.task_message_update import (
     TaskMessageUpdate,
-    StreamTaskMessageStart,
-    StreamTaskMessageDelta,
+    StreamTaskMessageDone,
     StreamTaskMessageFull,
-    StreamTaskMessageDone
+    StreamTaskMessageDelta,
+    StreamTaskMessageStart,
 )
-from agentex.types.text_delta import TextDelta
-
-from rich.console import Console
-from rich.panel import Panel
-from rich.markdown import Markdown
-from yaspin import yaspin
 
 
 def print_task_message(

agentex/lib/utils/iterables.py

@@ -1,5 +1,5 @@
-from collections.abc import AsyncGenerator
 from typing import Any
+from collections.abc import AsyncGenerator
 
 
 async def async_enumerate(

agentex/lib/utils/json_schema.py

@@ -10,10 +10,10 @@ def resolve_refs(schema: dict) -> dict:
     """
     resolved = jsonref.replace_refs(schema, proxies=False, lazy_load=False)
     serializable = {
-        "type": resolved.get("type"),
-        "properties": resolved.get("properties"),
-        "required": list(resolved.get("required", [])),
-        "additionalProperties": resolved.get("additionalProperties", False),
+        "type": resolved.get("type"),  # type: ignore[union-attr]
+        "properties": resolved.get("properties"),  # type: ignore[union-attr]
+        "required": list(resolved.get("required", [])),  # type: ignore[union-attr]
+        "additionalProperties": resolved.get("additionalProperties", False),  # type: ignore[union-attr]
     }
     return serializable
 

agentex/lib/utils/logging.py

@@ -1,31 +1,79 @@
+import os
 import logging
+import contextvars
 
+import ddtrace
+import json_log_formatter
 from rich.console import Console
 from rich.logging import RichHandler
 
+_is_datadog_configured = bool(os.environ.get("DD_AGENT_HOST"))
 
-def make_logger(name: str):
+ctx_var_request_id = contextvars.ContextVar[str]("request_id")
+
+
+class CustomJSONFormatter(json_log_formatter.JSONFormatter):
+    def json_record(self, message: str, extra: dict, record: logging.LogRecord) -> dict:  # type: ignore[override]
+        extra = super().json_record(message, extra, record)
+        extra["level"] = record.levelname
+        extra["name"] = record.name
+        extra["lineno"] = record.lineno
+        extra["pathname"] = record.pathname
+        extra["request_id"] = ctx_var_request_id.get(None)
+        if _is_datadog_configured:
+            extra["dd.trace_id"] = ddtrace.tracer.get_log_correlation_context().get("dd.trace_id", None) or getattr(  # type: ignore[attr-defined]
+                record, "dd.trace_id", 0
+            )
+            extra["dd.span_id"] = ddtrace.tracer.get_log_correlation_context().get("dd.span_id", None) or getattr(  # type: ignore[attr-defined]
+                record, "dd.span_id", 0
+            )
+        # add the env, service, and version configured for the tracer
+        # If tracing is not set up, then this should pull values from DD_ENV, DD_SERVICE, and DD_VERSION.
+        service_override = ddtrace.config.service or os.getenv("DD_SERVICE")
+        if service_override:
+            extra["dd.service"] = service_override
+
+        env_override = ddtrace.config.env or os.getenv("DD_ENV")
+        if env_override:
+            extra["dd.env"] = env_override
+
+        version_override = ddtrace.config.version or os.getenv("DD_VERSION")
+        if version_override:
+            extra["dd.version"] = version_override
+
+        return extra
+
+def make_logger(name: str) -> logging.Logger:
     """
     Creates a logger object with a RichHandler to print colored text.
     :param name: The name of the module to create the logger for.
     :return: A logger object.
     """
     # Create a console object to print colored text
-    console = Console()
-
-    # Create a logger object with the name of the current module
     logger = logging.getLogger(name)
-
-    # Set the global log level to INFO
     logger.setLevel(logging.INFO)
 
-    # Add the RichHandler to the logger to print colored text
-    handler = RichHandler(
-        console=console,
-        show_level=False,
-        show_path=False,
-        show_time=False,
-    )
-    logger.addHandler(handler)
+    environment = os.getenv("ENVIRONMENT")
+    if environment == "local":
+        console = Console()
+        # Add the RichHandler to the logger to print colored text
+        handler = RichHandler(
+            console=console,
+            show_level=False,
+            show_path=False,
+            show_time=False,
+        )
+        logger.addHandler(handler)
+        return logger
 
+    stream_handler = logging.StreamHandler()
+    if _is_datadog_configured:
+        stream_handler.setFormatter(CustomJSONFormatter())
+    else:
+        stream_handler.setFormatter(
+            logging.Formatter("%(asctime)s %(levelname)s [%(name)s] [%(filename)s:%(lineno)d] - %(message)s")
+        )
+
+    logger.addHandler(stream_handler)
+    # Create a logger object with the name of the current module
     return logger

agentex/lib/utils/mcp.py

@@ -1,4 +1,5 @@
 from typing import Any
+
 from mcp import StdioServerParameters
 
 

agentex/lib/utils/model_utils.py

@@ -1,9 +1,8 @@
-from collections.abc import Iterable, Mapping
-from datetime import datetime
 from typing import Any, TypeVar
+from datetime import datetime
+from collections.abc import Mapping, Iterable
 
-from pydantic import BaseModel as PydanticBaseModel
-from pydantic import ConfigDict
+from pydantic import BaseModel as PydanticBaseModel, ConfigDict
 
 from agentex.lib.utils.io import load_yaml_file
 
@@ -28,7 +27,7 @@ class BaseModel(PydanticBaseModel):
     def to_json(self, *args, **kwargs) -> str:
         return self.model_dump_json(*args, **kwargs)
 
-    def to_dict(self, *args, **kwargs) -> dict[str, Any]:
+    def to_dict(self, *_args, **_kwargs) -> dict[str, Any]:
         return recursive_model_dump(self)
 
 

agentex/lib/utils/registration.py

@@ -1,11 +1,12 @@
-import base64
-import json
 import os
-import httpx
+import json
+import base64
 import asyncio
 
-from agentex.lib.environment_variables import EnvironmentVariables, refreshed_environment_variables
+import httpx
+
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.environment_variables import EnvironmentVariables
 
 logger = make_logger(__name__)
 

agentex/resources/agents.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import json
-from typing import AsyncGenerator, Generator, Union, Optional
+from typing import Union, Optional, Generator, AsyncGenerator
 from typing_extensions import Literal
 
 import httpx
@@ -22,7 +22,14 @@ from .._response import (
 )
 from ..types.agent import Agent
 from .._base_client import make_request_options
-from ..types.agent_rpc_response import AgentRpcResponse, CancelTaskResponse, CreateTaskResponse, SendEventResponse, SendMessageResponse, SendMessageStreamResponse
+from ..types.agent_rpc_response import (
+    AgentRpcResponse,
+    SendEventResponse,
+    CancelTaskResponse,
+    CreateTaskResponse,
+    SendMessageResponse,
+    SendMessageStreamResponse,
+)
 from ..types.agent_list_response import AgentListResponse
 from ..types.shared.delete_response import DeleteResponse
 
@@ -1074,8 +1081,8 @@ class AsyncAgentsResource(AsyncAPIResource):
           except json.JSONDecodeError:
             # Skip invalid JSON lines
             continue
-          except ValidationError:
-            raise ValueError(f"Invalid SendMessageStreamResponse returned: {line}")
+          except ValidationError as e:
+            raise ValueError(f"Invalid SendMessageStreamResponse returned: {line}") from e
     
     async def send_event(
       self,

agentex/types/agent_rpc_response.py

@@ -3,11 +3,11 @@
 from typing import Union, Optional
 from typing_extensions import Literal
 
-from .._models import BaseModel
-from .agent_rpc_result import AgentRpcResult
-from .event import Event
 from .task import Task
+from .event import Event
+from .._models import BaseModel
 from .task_message import TaskMessage
+from .agent_rpc_result import AgentRpcResult
 from .task_message_update import TaskMessageUpdate
 
 __all__ = ["AgentRpcResponse"]

agentex/types/data_content.py

@@ -1,6 +1,6 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing import Dict, Optional
+from typing import Dict
 from typing_extensions import Literal
 
 from .._models import BaseModel

agentex/types/tool_request_content.py

@@ -1,6 +1,6 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing import Dict, Optional
+from typing import Dict
 from typing_extensions import Literal
 
 from .._models import BaseModel

agentex/types/tool_response_content.py

@@ -1,6 +1,5 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing import Optional
 from typing_extensions import Literal
 
 from .._models import BaseModel

{agentex_sdk-0.4.18.dist-info → agentex_sdk-0.4.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agentex-sdk
-Version: 0.4.18
+Version: 0.4.20
 Summary: The official Python library for the agentex API
 Project-URL: Homepage, https://github.com/scaleapi/agentex-python
 Project-URL: Repository, https://github.com/scaleapi/agentex-python
@@ -22,11 +22,14 @@ Requires-Python: <4,>=3.12
 Requires-Dist: aiohttp<4,>=3.10.10
 Requires-Dist: anyio<5,>=3.5.0
 Requires-Dist: cloudpickle>=3.1.1
+Requires-Dist: datadog>=0.52.1
+Requires-Dist: ddtrace>=3.13.0
 Requires-Dist: distro<2,>=1.7.0
 Requires-Dist: fastapi<0.116,>=0.115.0
 Requires-Dist: httpx<0.28,>=0.27.2
 Requires-Dist: ipykernel>=6.29.5
 Requires-Dist: jinja2<4,>=3.1.3
+Requires-Dist: json-log-formatter>=1.1.1
 Requires-Dist: jsonref<2,>=1.1.0
 Requires-Dist: jsonschema<5,>=4.23.0
 Requires-Dist: kubernetes<29.0.0,>=25.0.0