Accept DraftSession in AgentSession and Fix client types

src/truefoundry_gateway_sdk/agents/agent_session.py

@@ -1,23 +1,29 @@
+from __future__ import annotations
+
 import typing
 
 from ..core.pagination import AsyncPager, SyncPager
-from ..core.request_options import RequestOptions
-from ..types.list_session_events_response import ListSessionEventsResponse
-from ..types.list_turns_response import ListTurnsResponse
-from ..types.session_event_item import SessionEventItem
-from ..types.previous_turn_id_input import PreviousTurnIdInput
-from ..types.session import Session as RawSession
-from ..types.subject import Subject
-from ..types.turn import Turn as RawTurn
-from ..types.turn_input_item import TurnInputItem
 from .prepared_turn import AsyncPreparedTurn, PreparedTurn
 from .turn import AsyncTurn, Turn
 
+if typing.TYPE_CHECKING:
+    from ..client import AsyncTrueFoundryGateway, TrueFoundryGateway
+    from ..core.request_options import RequestOptions
+    from ..types.draft_session import DraftSession as RawDraftSession
+    from ..types.list_session_events_response import ListSessionEventsResponse
+    from ..types.list_turns_response import ListTurnsResponse
+    from ..types.previous_turn_id_input import PreviousTurnIdInput
+    from ..types.session import Session as RawSession
+    from ..types.session_event_item import SessionEventItem
+    from ..types.subject import Subject
+    from ..types.turn import Turn as RawTurn
+    from ..types.turn_input_item import TurnInputItem
+
 
 def _wrap_turns_pager(
     raw_pager: SyncPager[RawTurn, ListTurnsResponse],
     session: "AgentSession",
-    client: typing.Any,
+    client: TrueFoundryGateway,
 ) -> SyncPager[Turn, ListTurnsResponse]:
     wrapped_items = [Turn(t, session, client) for t in (raw_pager.items or [])]
 
@@ -40,7 +46,7 @@ def get_next() -> typing.Optional[SyncPager[Turn, ListTurnsResponse]]:
 async def _async_wrap_turns_pager(
     raw_pager: AsyncPager[RawTurn, ListTurnsResponse],
     session: "AsyncAgentSession",
-    client: typing.Any,
+    client: AsyncTrueFoundryGateway,
 ) -> AsyncPager[AsyncTurn, ListTurnsResponse]:
     wrapped_items = [AsyncTurn(t, session, client) for t in (raw_pager.items or [])]
 
@@ -65,9 +71,9 @@ class AgentSession:
     A session enriched with convenience methods: prepare_turn, list_turns, get_turn, list_events, cancel.
     """
 
-    def __init__(self, session: RawSession, client: typing.Any) -> None:
+    def __init__(self, session: typing.Union[RawSession, RawDraftSession], client: TrueFoundryGateway) -> None:
         self._id: str = session.id
-        self._agent_name: str = session.agent_name
+        self._agent_name: typing.Optional[str] = session.agent_name
         self._title: typing.Optional[str] = session.title
         self._created_by_subject: Subject = session.created_by_subject
         self._created_at: str = session.created_at
@@ -88,7 +94,7 @@ def id(self) -> str:
         return self._id
 
     @property
-    def agent_name(self) -> str:
+    def agent_name(self) -> typing.Optional[str]:
         """
         Returns
         -------
@@ -274,9 +280,9 @@ class AsyncAgentSession:
     Async version of AgentSession.
     """
 
-    def __init__(self, session: RawSession, client: typing.Any) -> None:
+    def __init__(self, session: typing.Union[RawSession, RawDraftSession], client: AsyncTrueFoundryGateway) -> None:
         self._id: str = session.id
-        self._agent_name: str = session.agent_name
+        self._agent_name: typing.Optional[str] = session.agent_name
         self._title: typing.Optional[str] = session.title
         self._created_by_subject: Subject = session.created_by_subject
         self._created_at: str = session.created_at
@@ -297,7 +303,7 @@ def id(self) -> str:
         return self._id
 
     @property
-    def agent_name(self) -> str:
+    def agent_name(self) -> typing.Optional[str]:
         """
         Returns
         -------

src/truefoundry_gateway_sdk/agents/agent_session_client.py

@@ -1,19 +1,24 @@
+from __future__ import annotations
+
 import os
 import typing
 
-import httpx
-from ..core.logging import LogConfig, Logger
 from ..core.pagination import AsyncPager, SyncPager
-from ..core.request_options import RequestOptions
-from ..types.list_sessions_order import ListSessionsOrder
-from ..types.list_sessions_response import ListSessionsResponse
-from ..types.session import Session as RawSession
 from .agent_session import AgentSession, AsyncAgentSession
 
+if typing.TYPE_CHECKING:
+    import httpx
+    from ..client import AsyncTrueFoundryGateway, TrueFoundryGateway
+    from ..core.logging import LogConfig, Logger
+    from ..core.request_options import RequestOptions
+    from ..types.list_sessions_order import ListSessionsOrder
+    from ..types.list_sessions_response import ListSessionsResponse
+    from ..types.session import Session as RawSession
+
 
 def _wrap_sessions_pager(
     raw_pager: SyncPager[RawSession, ListSessionsResponse],
-    client: typing.Any,
+    client: TrueFoundryGateway,
 ) -> SyncPager[AgentSession, ListSessionsResponse]:
     wrapped_items = [AgentSession(s, client) for s in (raw_pager.items or [])]
 
@@ -35,7 +40,7 @@ def get_next() -> typing.Optional[SyncPager[AgentSession, ListSessionsResponse]]
 
 async def _async_wrap_sessions_pager(
     raw_pager: AsyncPager[RawSession, ListSessionsResponse],
-    client: typing.Any,
+    client: AsyncTrueFoundryGateway,
 ) -> AsyncPager[AsyncAgentSession, ListSessionsResponse]:
     wrapped_items = [AsyncAgentSession(s, client) for s in (raw_pager.items or [])]
 

src/truefoundry_gateway_sdk/agents/event_delta.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import typing
 
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
@@ -6,8 +8,10 @@
 from ..types.model_message_delta_event import ModelMessageDeltaEvent
 from ..types.model_message_event import ModelMessageEvent
 from ..types.tool_call import ToolCall
-from ..types.turn_event import TurnEvent
-from ..types.turn_streaming_event import TurnStreamingEvent
+
+if typing.TYPE_CHECKING:
+    from ..types.turn_event import TurnEvent
+    from ..types.turn_streaming_event import TurnStreamingEvent
 
 # Union of all streaming delta events. Expand as more `.delta` events are added.
 DeltaEvents = ModelMessageDeltaEvent

src/truefoundry_gateway_sdk/agents/prepared_turn.py

@@ -1,21 +1,24 @@
+from __future__ import annotations
+
 import typing
 
-from ..core.pagination import AsyncPager, SyncPager
-from ..core.request_options import RequestOptions
-from ..types.list_events_order import ListEventsOrder
-from ..types.list_events_response import ListEventsResponse
-from ..types.previous_turn_id_input import PreviousTurnIdInput
-from ..types.subject import Subject
 from ..types.turn import Turn as RawTurn
 from ..types.turn_created_event import TurnCreatedEvent
 from ..types.turn_done_event import TurnDoneEvent
-from ..types.turn_event import TurnEvent
-from ..types.turn_input_item import TurnInputItem
-from ..types.turn_state import TurnState
 from .turn import AsyncTurn, Turn
 from .turn_stream_data import TurnStreamData
 
 if typing.TYPE_CHECKING:
+    from ..client import AsyncTrueFoundryGateway, TrueFoundryGateway
+    from ..core.pagination import AsyncPager, SyncPager
+    from ..core.request_options import RequestOptions
+    from ..types.list_events_order import ListEventsOrder
+    from ..types.list_events_response import ListEventsResponse
+    from ..types.previous_turn_id_input import PreviousTurnIdInput
+    from ..types.subject import Subject
+    from ..types.turn_event import TurnEvent
+    from ..types.turn_input_item import TurnInputItem
+    from ..types.turn_state import TurnState
     from .agent_session import AgentSession, AsyncAgentSession
 
 
@@ -31,7 +34,7 @@ def __init__(
         input: typing.Optional[typing.Sequence[TurnInputItem]],
         previous_turn_id: typing.Optional[PreviousTurnIdInput],
         session: "AgentSession",
-        client: typing.Any,
+        client: TrueFoundryGateway,
     ) -> None:
         self._input = list(input) if input is not None else None
         self._previous_turn_id = previous_turn_id
@@ -383,7 +386,7 @@ def __init__(
         input: typing.Optional[typing.Sequence[TurnInputItem]],
         previous_turn_id: typing.Optional[PreviousTurnIdInput],
         session: "AsyncAgentSession",
-        client: typing.Any,
+        client: AsyncTrueFoundryGateway,
     ) -> None:
         self._input = list(input) if input is not None else None
         self._previous_turn_id = previous_turn_id

src/truefoundry_gateway_sdk/agents/turn.py

@@ -1,24 +1,27 @@
+from __future__ import annotations
+
 import time
 import typing
 
-from ..core.pagination import AsyncPager, SyncPager
-from ..core.request_options import RequestOptions
-from ..types.list_events_order import ListEventsOrder
-from ..types.list_events_response import ListEventsResponse
-from ..types.subject import Subject
-from ..types.turn import Turn as RawTurn
 from ..types.turn_created_event import TurnCreatedEvent
 from ..types.turn_done_event import TurnDoneEvent
-from ..types.turn_event import TurnEvent
-from ..types.turn_input_item import TurnInputItem
-from ..types.turn_state import TurnState
 from ..types.turn_state_cancelled import TurnStateCancelled
 from ..types.turn_state_done import TurnStateDone
 from ..types.turn_state_error import TurnStateError
-from ..types.turn_streaming_event import TurnStreamingEvent
 from .turn_stream_data import TurnStreamData
 
 if typing.TYPE_CHECKING:
+    from ..client import AsyncTrueFoundryGateway, TrueFoundryGateway
+    from ..core.pagination import AsyncPager, SyncPager
+    from ..core.request_options import RequestOptions
+    from ..types.list_events_order import ListEventsOrder
+    from ..types.list_events_response import ListEventsResponse
+    from ..types.subject import Subject
+    from ..types.turn import Turn as RawTurn
+    from ..types.turn_event import TurnEvent
+    from ..types.turn_input_item import TurnInputItem
+    from ..types.turn_state import TurnState
+    from ..types.turn_streaming_event import TurnStreamingEvent
     from .agent_session import AgentSession, AsyncAgentSession
 
 # waitForCompletion poll interval in milliseconds: default and enforced minimum.
@@ -36,10 +39,8 @@ def __init__(
         self,
         turn: RawTurn,
         session: "AgentSession",
-        client: typing.Any,
+        client: TrueFoundryGateway,
     ) -> None:
-        from ..client import TrueFoundryGateway
-
         self._id: str = turn.id
         self._session_id: str = turn.session_id
         self._previous_turn_id: typing.Optional[str] = turn.previous_turn_id
@@ -48,7 +49,7 @@ def __init__(
         self._created_at: str = turn.created_at
         self._state: TurnState = turn.state
         self._session: "AgentSession" = session
-        self._client: TrueFoundryGateway = client
+        self._client = client
 
     def __repr__(self) -> str:
         return f"Turn(id={self._id!r}, session_id={self._session_id!r}, status={self._state.status!r})"
@@ -280,10 +281,8 @@ def __init__(
         self,
         turn: RawTurn,
         session: "AsyncAgentSession",
-        client: typing.Any,
+        client: AsyncTrueFoundryGateway,
     ) -> None:
-        from ..client import AsyncTrueFoundryGateway
-
         self._id: str = turn.id
         self._session_id: str = turn.session_id
         self._previous_turn_id: typing.Optional[str] = turn.previous_turn_id

src/truefoundry_gateway_sdk/agents/turn_stream_data.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 import dataclasses
 import typing
 
-from ..types.turn_streaming_event import TurnStreamingEvent
+if typing.TYPE_CHECKING:
+    from ..types.turn_streaming_event import TurnStreamingEvent
 
 
 @dataclasses.dataclass