diff --git a/src/agentex/resources/agents/agents.py b/src/agentex/resources/agents/agents.py index ab2949fc7..3dd0c770d 100644 --- a/src/agentex/resources/agents/agents.py +++ b/src/agentex/resources/agents/agents.py @@ -10,7 +10,7 @@ from pydantic import ValidationError from ...types import agent_rpc_params, agent_list_params, agent_rpc_by_name_params -from ..._types import Body, Omit, Query, Headers, NotGiven, omit, not_given +from ..._types import NOT_GIVEN, Body, Omit, Query, Headers, NotGiven, omit, not_given from ..._utils import path_template, maybe_transform, async_maybe_transform from ..._compat import cached_property from .schedules import ( @@ -38,7 +38,14 @@ ) from ...types.agent import Agent from ..._base_client import make_request_options -from ...types.agent_rpc_response import AgentRpcResponse +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 diff --git a/src/agentex/types/agent_rpc_response.py b/src/agentex/types/agent_rpc_response.py index 8b464b8d7..97f0f9c2f 100644 --- a/src/agentex/types/agent_rpc_response.py +++ b/src/agentex/types/agent_rpc_response.py @@ -11,7 +11,14 @@ from .agent_rpc_result import AgentRpcResult from .task_message_update import TaskMessageUpdate -__all__ = ["AgentRpcResponse"] +__all__ = [ + "AgentRpcResponse", + "CancelTaskResponse", + "CreateTaskResponse", + "SendEventResponse", + "SendMessageResponse", + "SendMessageStreamResponse", +] class BaseAgentRpcResponse(BaseModel): diff --git a/tests/api_resources/agents/test_deployments.py b/tests/api_resources/agents/test_deployments.py index 3a429c0f1..7bddf7c49 100644 --- a/tests/api_resources/agents/test_deployments.py +++ b/tests/api_resources/agents/test_deployments.py @@ -8,7 +8,6 @@ import pytest from agentex import Agentex, AsyncAgentex -from tests.utils import assert_matches_type from agentex.types import AgentRpcResponse from agentex.types.agents import ( DeploymentListResponse, @@ -18,6 +17,8 @@ ) from agentex.types.shared import DeleteResponse +from ...utils import assert_matches_type + base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") diff --git a/tests/api_resources/agents/test_schedules.py b/tests/api_resources/agents/test_schedules.py index 46ef24aa1..0431de18e 100644 --- a/tests/api_resources/agents/test_schedules.py +++ b/tests/api_resources/agents/test_schedules.py @@ -8,7 +8,6 @@ import pytest from agentex import Agentex, AsyncAgentex -from tests.utils import assert_matches_type from agentex._utils import parse_datetime from agentex.types.agents import ( ScheduleListResponse, @@ -20,6 +19,8 @@ ) from agentex.types.shared import DeleteResponse +from ...utils import assert_matches_type + base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") diff --git a/tests/api_resources/test_checkpoints.py b/tests/api_resources/test_checkpoints.py index 68d297cab..3c13fd43a 100644 --- a/tests/api_resources/test_checkpoints.py +++ b/tests/api_resources/test_checkpoints.py @@ -8,13 +8,14 @@ import pytest from agentex import Agentex, AsyncAgentex -from tests.utils import assert_matches_type from agentex.types import ( CheckpointPutResponse, CheckpointListResponse, CheckpointGetTupleResponse, ) +from ..utils import assert_matches_type + base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")