fix: resolve lint and test failures from new endpoints (#360)

Author: declan-scale

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
 

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):

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")
 
 

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")
 
 

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")