From ac8c5314c9144568cf354a1957daf2de22d410f2 Mon Sep 17 00:00:00 2001 From: Srajan Asthana Date: Fri, 17 Jul 2026 20:21:54 +0530 Subject: [PATCH 1/2] Add internal list sessions to SDK --- .../agents/__init__.py | 13 + .../agents/agent_session.py | 119 ++-- .../agents/prepared_turn.py | 17 +- .../agents/private/__init__.py | 9 + .../agents/private/agent_draft_session.py | 462 ++++++++++++++++ .../private/private_agent_session_client.py | 506 ++++++++++++++++++ .../agents/session_mixin.py | 378 +++++++++++++ src/truefoundry_gateway_sdk/agents/turn.py | 17 +- .../test_private_agent_session_client.py | 128 +++++ 9 files changed, 1551 insertions(+), 98 deletions(-) create mode 100644 src/truefoundry_gateway_sdk/agents/private/__init__.py create mode 100644 src/truefoundry_gateway_sdk/agents/private/agent_draft_session.py create mode 100644 src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py create mode 100644 src/truefoundry_gateway_sdk/agents/session_mixin.py create mode 100644 tests/custom/test_private_agent_session_client.py diff --git a/src/truefoundry_gateway_sdk/agents/__init__.py b/src/truefoundry_gateway_sdk/agents/__init__.py index 4e53269..01394ce 100644 --- a/src/truefoundry_gateway_sdk/agents/__init__.py +++ b/src/truefoundry_gateway_sdk/agents/__init__.py @@ -32,17 +32,28 @@ from .agent_session_client import AgentSessionClient, AsyncAgentSessionClient from .event_delta import DeltaEvents, is_event_delta, merge_event_delta from .prepared_turn import AsyncPreparedTurn, PreparedTurn +from .private import ( + AgentDraftSession, + AsyncAgentDraftSession, + AsyncPrivateAgentSessionClient, + PrivateAgentSessionClient, +) +from .session_mixin import AsyncSessionMixin, SessionMixin from .turn import AsyncTurn, Turn from .turn_stream_data import TurnStreamData __all__ = [ # type: ignore[reportUnsupportedDunderAll] *_private_agents.__all__, "ActionRequiredEvent", + "AgentDraftSession", "AgentSession", "AgentSessionClient", + "AsyncAgentDraftSession", "AsyncAgentSession", "AsyncAgentSessionClient", "AsyncPreparedTurn", + "AsyncPrivateAgentSessionClient", + "AsyncSessionMixin", "AsyncTurn", "DeltaEvents", "McpAuthRequiredEvent", @@ -52,7 +63,9 @@ "ModelMessageUsage", "ModelMessageUsageInputTokensBreakdown", "PreparedTurn", + "PrivateAgentSessionClient", "SandboxCreatedEvent", + "SessionMixin", "ThreadCreatedEvent", "ThreadDoneEvent", "ToolApprovalRequiredEvent", diff --git a/src/truefoundry_gateway_sdk/agents/agent_session.py b/src/truefoundry_gateway_sdk/agents/agent_session.py index 792967d..67c17a6 100644 --- a/src/truefoundry_gateway_sdk/agents/agent_session.py +++ b/src/truefoundry_gateway_sdk/agents/agent_session.py @@ -2,12 +2,11 @@ import typing -from ..core.pagination import AsyncPager, SyncPager -from .prepared_turn import AsyncPreparedTurn, PreparedTurn -from .turn import AsyncTurn, Turn +from .session_mixin import AsyncSessionMixin, SessionMixin if typing.TYPE_CHECKING: from ..client import AsyncTrueFoundryGateway, TrueFoundryGateway + 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 @@ -15,59 +14,15 @@ 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: TrueFoundryGateway, -) -> SyncPager[Turn, ListTurnsResponse]: - wrapped_items = [Turn(t, session, client) for t in (raw_pager.items or [])] - - def get_next() -> typing.Optional[SyncPager[Turn, ListTurnsResponse]]: - if raw_pager.get_next is None: - return None - next_raw = raw_pager.get_next() - if next_raw is None: - return None - return _wrap_turns_pager(next_raw, session, client) - - return SyncPager( - get_next=get_next if raw_pager.has_next else None, - has_next=raw_pager.has_next, - items=wrapped_items, - response=raw_pager.response, - ) - - -async def _async_wrap_turns_pager( - raw_pager: AsyncPager[RawTurn, ListTurnsResponse], - session: "AsyncAgentSession", - client: AsyncTrueFoundryGateway, -) -> AsyncPager[AsyncTurn, ListTurnsResponse]: - wrapped_items = [AsyncTurn(t, session, client) for t in (raw_pager.items or [])] - - async def get_next() -> typing.Optional[AsyncPager[AsyncTurn, ListTurnsResponse]]: - if raw_pager.get_next is None: - return None - next_raw = await raw_pager.get_next() - if next_raw is None: - return None - return await _async_wrap_turns_pager(next_raw, session, client) - - return AsyncPager( - get_next=get_next if raw_pager.has_next else None, - has_next=raw_pager.has_next, - items=wrapped_items, - response=raw_pager.response, - ) + from .prepared_turn import AsyncPreparedTurn, PreparedTurn + from .turn import AsyncTurn, Turn class AgentSession: """ A session enriched with convenience methods: prepare_turn, list_turns, get_turn, list_events, cancel. + Turn operations are delegated to a shared :class:`SessionMixin`. """ def __init__(self, session: RawSession, client: TrueFoundryGateway) -> None: @@ -77,11 +32,21 @@ def __init__(self, session: RawSession, client: TrueFoundryGateway) -> None: self._created_by_subject: Subject = session.created_by_subject self._created_at: str = session.created_at self._updated_at: str = session.updated_at - self._client = client + self._mixin = SessionMixin(session.id, client) def __repr__(self) -> str: return f"AgentSession(id={self._id!r}, agent_name={self._agent_name!r})" + @property + def type(self) -> typing.Literal["session"]: + """ + Returns + ------- + typing.Literal["session"] + Discriminant distinguishing a saved session from a draft session. + """ + return "session" + @property def id(self) -> str: """ @@ -163,12 +128,7 @@ def prepare_turn( PreparedTurn Staged turn. """ - return PreparedTurn( - input=input, - previous_turn_id=previous_turn_id, - session=self, - client=self._client, - ) + return self._mixin.prepare_turn(self, input=input, previous_turn_id=previous_turn_id) def list_turns( self, @@ -194,10 +154,7 @@ def list_turns( SyncPager[Turn, ListTurnsResponse] Paginated turns. """ - raw_pager = self._client.agents.sessions.list_turns( - self._id, page_token=page_token, limit=limit, request_options=request_options - ) - return _wrap_turns_pager(raw_pager, self, self._client) + return self._mixin.list_turns(self, page_token=page_token, limit=limit, request_options=request_options) def get_turn( self, @@ -220,8 +177,7 @@ def get_turn( Turn Turn data. """ - response = self._client.agents.sessions.get_turn(self._id, turn_id, request_options=request_options) - return Turn(response.data, self, self._client) + return self._mixin.get_turn(self, turn_id, request_options=request_options) def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None: """ @@ -236,7 +192,7 @@ def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> ------- None """ - self._client.agents.sessions.cancel(self._id, request_options=request_options) + self._mixin.cancel(request_options=request_options) def list_events( self, @@ -265,8 +221,7 @@ def list_events( SyncPager[SessionEventItem, ListSessionEventsResponse] Paginated session events. """ - return self._client.agents.sessions.list_events( - self._id, + return self._mixin.list_events( page_token=page_token, last_turn_id=last_turn_id, limit=limit, @@ -286,11 +241,21 @@ def __init__(self, session: RawSession, client: AsyncTrueFoundryGateway) -> None self._created_by_subject: Subject = session.created_by_subject self._created_at: str = session.created_at self._updated_at: str = session.updated_at - self._client = client + self._mixin = AsyncSessionMixin(session.id, client) def __repr__(self) -> str: return f"AsyncAgentSession(id={self._id!r}, agent_name={self._agent_name!r})" + @property + def type(self) -> typing.Literal["session"]: + """ + Returns + ------- + typing.Literal["session"] + Discriminant distinguishing a saved session from a draft session. + """ + return "session" + @property def id(self) -> str: """ @@ -372,12 +337,7 @@ def prepare_turn( AsyncPreparedTurn Staged turn. """ - return AsyncPreparedTurn( - input=input, - previous_turn_id=previous_turn_id, - session=self, - client=self._client, - ) + return self._mixin.prepare_turn(self, input=input, previous_turn_id=previous_turn_id) async def list_turns( self, @@ -403,10 +363,7 @@ async def list_turns( AsyncPager[AsyncTurn, ListTurnsResponse] Paginated turns. """ - raw_pager = await self._client.agents.sessions.list_turns( - self._id, page_token=page_token, limit=limit, request_options=request_options - ) - return await _async_wrap_turns_pager(raw_pager, self, self._client) + return await self._mixin.list_turns(self, page_token=page_token, limit=limit, request_options=request_options) async def get_turn( self, @@ -429,8 +386,7 @@ async def get_turn( AsyncTurn Turn data. """ - response = await self._client.agents.sessions.get_turn(self._id, turn_id, request_options=request_options) - return AsyncTurn(response.data, self, self._client) + return await self._mixin.get_turn(self, turn_id, request_options=request_options) async def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None: """ @@ -445,7 +401,7 @@ async def cancel(self, *, request_options: typing.Optional[RequestOptions] = Non ------- None """ - await self._client.agents.sessions.cancel(self._id, request_options=request_options) + await self._mixin.cancel(request_options=request_options) async def list_events( self, @@ -474,8 +430,7 @@ async def list_events( AsyncPager[SessionEventItem, ListSessionEventsResponse] Paginated session events. """ - return await self._client.agents.sessions.list_events( - self._id, + return await self._mixin.list_events( page_token=page_token, last_turn_id=last_turn_id, limit=limit, diff --git a/src/truefoundry_gateway_sdk/agents/prepared_turn.py b/src/truefoundry_gateway_sdk/agents/prepared_turn.py index ec720ca..5b8c029 100644 --- a/src/truefoundry_gateway_sdk/agents/prepared_turn.py +++ b/src/truefoundry_gateway_sdk/agents/prepared_turn.py @@ -20,6 +20,7 @@ from ..types.turn_input_item import TurnInputItem from ..types.turn_state import TurnState from .agent_session import AgentSession, AsyncAgentSession + from .private.agent_draft_session import AgentDraftSession, AsyncAgentDraftSession class PreparedTurn: @@ -33,12 +34,12 @@ def __init__( *, input: typing.Optional[typing.Sequence[TurnInputItem]], previous_turn_id: typing.Optional[PreviousTurnIdInput], - session: "AgentSession", + session: typing.Union["AgentSession", "AgentDraftSession"], client: TrueFoundryGateway, ) -> None: self._input = list(input) if input is not None else None self._previous_turn_id = previous_turn_id - self._session: "AgentSession" = session + self._session: typing.Union["AgentSession", "AgentDraftSession"] = session self._session_id: str = session.id self._client = client self._started: bool = False @@ -50,11 +51,11 @@ def __repr__(self) -> str: # --- Properties that delegate to the inner Turn (None until execute is called) --- @property - def session(self) -> "AgentSession": + def session(self) -> typing.Union["AgentSession", "AgentDraftSession"]: """ Returns ------- - AgentSession + typing.Union[AgentSession, AgentDraftSession] Parent session this prepared turn belongs to. """ return self._session @@ -385,12 +386,12 @@ def __init__( *, input: typing.Optional[typing.Sequence[TurnInputItem]], previous_turn_id: typing.Optional[PreviousTurnIdInput], - session: "AsyncAgentSession", + session: typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"], client: AsyncTrueFoundryGateway, ) -> None: self._input = list(input) if input is not None else None self._previous_turn_id = previous_turn_id - self._session: "AsyncAgentSession" = session + self._session: typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"] = session self._session_id: str = session.id self._client = client self._started: bool = False @@ -402,11 +403,11 @@ def __repr__(self) -> str: # --- Properties that delegate to the inner AsyncTurn (None until execute is called) --- @property - def session(self) -> "AsyncAgentSession": + def session(self) -> typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"]: """ Returns ------- - AsyncAgentSession + typing.Union[AsyncAgentSession, AsyncAgentDraftSession] Parent session this prepared turn belongs to. """ return self._session diff --git a/src/truefoundry_gateway_sdk/agents/private/__init__.py b/src/truefoundry_gateway_sdk/agents/private/__init__.py new file mode 100644 index 0000000..a00cea6 --- /dev/null +++ b/src/truefoundry_gateway_sdk/agents/private/__init__.py @@ -0,0 +1,9 @@ +from .agent_draft_session import AgentDraftSession, AsyncAgentDraftSession +from .private_agent_session_client import AsyncPrivateAgentSessionClient, PrivateAgentSessionClient + +__all__ = [ + "AgentDraftSession", + "AsyncAgentDraftSession", + "AsyncPrivateAgentSessionClient", + "PrivateAgentSessionClient", +] diff --git a/src/truefoundry_gateway_sdk/agents/private/agent_draft_session.py b/src/truefoundry_gateway_sdk/agents/private/agent_draft_session.py new file mode 100644 index 0000000..401fe7d --- /dev/null +++ b/src/truefoundry_gateway_sdk/agents/private/agent_draft_session.py @@ -0,0 +1,462 @@ +from __future__ import annotations + +import typing + +from ..session_mixin import AsyncSessionMixin, SessionMixin + +if typing.TYPE_CHECKING: + from ...client import AsyncTrueFoundryGateway, TrueFoundryGateway + from ...core.pagination import AsyncPager, SyncPager + from ...core.request_options import RequestOptions + from ...types.agent_spec import AgentSpec + 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_event_item import SessionEventItem + from ...types.subject import Subject + from ...types.turn_input_item import TurnInputItem + from ..prepared_turn import AsyncPreparedTurn, PreparedTurn + from ..turn import AsyncTurn, Turn + + +class AgentDraftSession: + """ + A draft session enriched with the same convenience methods as :class:`AgentSession`: + prepare_turn, list_turns, get_turn, list_events, cancel. Turn operations are delegated to a + shared :class:`SessionMixin`, so drafts and saved sessions expose an identical turn API. + """ + + def __init__(self, session: RawDraftSession, client: TrueFoundryGateway) -> None: + self._id: str = session.id + self._agent_spec: AgentSpec = session.agent_spec + 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 + self._updated_at: str = session.updated_at + self._mixin = SessionMixin(session.id, client) + + def __repr__(self) -> str: + return f"AgentDraftSession(id={self._id!r}, agent_name={self._agent_name!r})" + + @property + def type(self) -> typing.Literal["session/draft"]: + """ + Returns + ------- + typing.Literal["session/draft"] + Discriminant distinguishing a draft session from a saved session. + """ + return "session/draft" + + @property + def id(self) -> str: + """ + Returns + ------- + str + Unique identifier of this draft session. + """ + return self._id + + @property + def agent_spec(self) -> AgentSpec: + """ + Returns + ------- + AgentSpec + Inline agent spec held by this draft. + """ + return self._agent_spec + + @property + def agent_name(self) -> typing.Optional[str]: + """ + Returns + ------- + typing.Optional[str] + Optional saved agent this draft is linked to. + """ + return self._agent_name + + @property + def title(self) -> typing.Optional[str]: + """ + Returns + ------- + typing.Optional[str] + Optional user-visible title for the draft session. + """ + return self._title + + @property + def created_by_subject(self) -> Subject: + """ + Returns + ------- + Subject + Subject that created this draft session. + """ + return self._created_by_subject + + @property + def created_at(self) -> str: + """ + Returns + ------- + str + ISO-8601 timestamp when the draft session was created. + """ + return self._created_at + + @property + def updated_at(self) -> str: + """ + Returns + ------- + str + ISO-8601 timestamp when the draft session was last updated. + """ + return self._updated_at + + def prepare_turn( + self, + *, + input: typing.Optional[typing.Sequence[TurnInputItem]] = None, + previous_turn_id: typing.Optional[PreviousTurnIdInput] = None, + ) -> PreparedTurn: + """ + Stage a turn locally; call ``execute()`` to start ``create_turn``. + + Parameters + ---------- + input : typing.Optional[typing.Sequence[TurnInputItem]] + Turn input items passed to create turn. + previous_turn_id : typing.Optional[PreviousTurnIdInput] + Previous turn to chain from. Defaults to ``auto``. + + Returns + ------- + PreparedTurn + Staged turn. + """ + return self._mixin.prepare_turn(self, input=input, previous_turn_id=previous_turn_id) + + def list_turns( + self, + *, + page_token: typing.Optional[str] = None, + limit: typing.Optional[int] = 10, + request_options: typing.Optional[RequestOptions] = None, + ) -> SyncPager[Turn, ListTurnsResponse]: + """ + List turns in this session. + + Parameters + ---------- + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + limit : typing.Optional[int] + Page size. Default 10. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + SyncPager[Turn, ListTurnsResponse] + Paginated turns. + """ + return self._mixin.list_turns(self, page_token=page_token, limit=limit, request_options=request_options) + + def get_turn( + self, + turn_id: str, + *, + request_options: typing.Optional[RequestOptions] = None, + ) -> Turn: + """ + Fetch a turn by ID. + + Parameters + ---------- + turn_id : str + Unique identifier of the turn to fetch. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + Turn + Turn data. + """ + return self._mixin.get_turn(self, turn_id, request_options=request_options) + + def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None: + """ + Cancel the running last turn for the session. + + Parameters + ---------- + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + None + """ + self._mixin.cancel(request_options=request_options) + + def list_events( + self, + *, + page_token: typing.Optional[str] = None, + last_turn_id: typing.Optional[str] = None, + limit: typing.Optional[int] = 100, + request_options: typing.Optional[RequestOptions] = None, + ) -> SyncPager[SessionEventItem, ListSessionEventsResponse]: + """ + Paginated session events across turns (newest first); subscribe to a running turn for live events. + + Parameters + ---------- + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + last_turn_id : typing.Optional[str] + Newest turn in the listing window (initial load only). Omit to use the session last turn. + limit : typing.Optional[int] + Page size. Default 100. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + SyncPager[SessionEventItem, ListSessionEventsResponse] + Paginated session events. + """ + return self._mixin.list_events( + page_token=page_token, + last_turn_id=last_turn_id, + limit=limit, + request_options=request_options, + ) + + +class AsyncAgentDraftSession: + """ + Async version of AgentDraftSession. + """ + + def __init__(self, session: RawDraftSession, client: AsyncTrueFoundryGateway) -> None: + self._id: str = session.id + self._agent_spec: AgentSpec = session.agent_spec + 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 + self._updated_at: str = session.updated_at + self._mixin = AsyncSessionMixin(session.id, client) + + def __repr__(self) -> str: + return f"AsyncAgentDraftSession(id={self._id!r}, agent_name={self._agent_name!r})" + + @property + def type(self) -> typing.Literal["session/draft"]: + """ + Returns + ------- + typing.Literal["session/draft"] + Discriminant distinguishing a draft session from a saved session. + """ + return "session/draft" + + @property + def id(self) -> str: + """ + Returns + ------- + str + Unique identifier of this draft session. + """ + return self._id + + @property + def agent_spec(self) -> AgentSpec: + """ + Returns + ------- + AgentSpec + Inline agent spec held by this draft. + """ + return self._agent_spec + + @property + def agent_name(self) -> typing.Optional[str]: + """ + Returns + ------- + typing.Optional[str] + Optional saved agent this draft is linked to. + """ + return self._agent_name + + @property + def title(self) -> typing.Optional[str]: + """ + Returns + ------- + typing.Optional[str] + Optional user-visible title for the draft session. + """ + return self._title + + @property + def created_by_subject(self) -> Subject: + """ + Returns + ------- + Subject + Subject that created this draft session. + """ + return self._created_by_subject + + @property + def created_at(self) -> str: + """ + Returns + ------- + str + ISO-8601 timestamp when the draft session was created. + """ + return self._created_at + + @property + def updated_at(self) -> str: + """ + Returns + ------- + str + ISO-8601 timestamp when the draft session was last updated. + """ + return self._updated_at + + def prepare_turn( + self, + *, + input: typing.Optional[typing.Sequence[TurnInputItem]] = None, + previous_turn_id: typing.Optional[PreviousTurnIdInput] = None, + ) -> AsyncPreparedTurn: + """ + Stage a turn locally; call ``execute()`` to start ``create_turn``. + + Parameters + ---------- + input : typing.Optional[typing.Sequence[TurnInputItem]] + Turn input items passed to create turn. + previous_turn_id : typing.Optional[PreviousTurnIdInput] + Previous turn to chain from. Defaults to ``auto``. + + Returns + ------- + AsyncPreparedTurn + Staged turn. + """ + return self._mixin.prepare_turn(self, input=input, previous_turn_id=previous_turn_id) + + async def list_turns( + self, + *, + page_token: typing.Optional[str] = None, + limit: typing.Optional[int] = 10, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncPager[AsyncTurn, ListTurnsResponse]: + """ + List turns in this session. + + Parameters + ---------- + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + limit : typing.Optional[int] + Page size. Default 10. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncPager[AsyncTurn, ListTurnsResponse] + Paginated turns. + """ + return await self._mixin.list_turns(self, page_token=page_token, limit=limit, request_options=request_options) + + async def get_turn( + self, + turn_id: str, + *, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncTurn: + """ + Fetch a turn by ID. + + Parameters + ---------- + turn_id : str + Unique identifier of the turn to fetch. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncTurn + Turn data. + """ + return await self._mixin.get_turn(self, turn_id, request_options=request_options) + + async def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None: + """ + Cancel the running last turn for the session. + + Parameters + ---------- + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + None + """ + await self._mixin.cancel(request_options=request_options) + + async def list_events( + self, + *, + page_token: typing.Optional[str] = None, + last_turn_id: typing.Optional[str] = None, + limit: typing.Optional[int] = 100, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncPager[SessionEventItem, ListSessionEventsResponse]: + """ + Paginated session events across turns (newest first); subscribe to a running turn for live events. + + Parameters + ---------- + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + last_turn_id : typing.Optional[str] + Newest turn in the listing window (initial load only). Omit to use the session last turn. + limit : typing.Optional[int] + Page size. Default 100. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncPager[SessionEventItem, ListSessionEventsResponse] + Paginated session events. + """ + return await self._mixin.list_events( + page_token=page_token, + last_turn_id=last_turn_id, + limit=limit, + request_options=request_options, + ) diff --git a/src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py b/src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py new file mode 100644 index 0000000..6522b1e --- /dev/null +++ b/src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +import os +import typing + +from ...core.pagination import AsyncPager, SyncPager +from ...types.draft_session import DraftSession as RawDraftSession +from ..agent_session import AgentSession, AsyncAgentSession +from .agent_draft_session import AgentDraftSession, AsyncAgentDraftSession + +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.agent_spec import AgentSpec + from ...types.list_draft_sessions_order import ListDraftSessionsOrder + from ...types.list_draft_sessions_response import ListDraftSessionsResponse + from ...types.list_owned_sessions_order import ListOwnedSessionsOrder + from ...types.list_owned_sessions_response import ListOwnedSessionsResponse + from ...types.list_owned_sessions_response_data_item import ListOwnedSessionsResponseDataItem + + +def _wrap_draft_sessions_pager( + raw_pager: SyncPager[RawDraftSession, ListDraftSessionsResponse], + client: TrueFoundryGateway, +) -> SyncPager[AgentDraftSession, ListDraftSessionsResponse]: + wrapped_items = [AgentDraftSession(d, client) for d in (raw_pager.items or [])] + + def get_next() -> typing.Optional[SyncPager[AgentDraftSession, ListDraftSessionsResponse]]: + if raw_pager.get_next is None: + return None + next_raw = raw_pager.get_next() + if next_raw is None: + return None + return _wrap_draft_sessions_pager(next_raw, client) + + return SyncPager( + get_next=get_next if raw_pager.has_next else None, + has_next=raw_pager.has_next, + items=wrapped_items, + response=raw_pager.response, + ) + + +def _wrap_owned_session( + raw: ListOwnedSessionsResponseDataItem, + client: TrueFoundryGateway, +) -> typing.Union[AgentSession, AgentDraftSession]: + # Dispatch a raw owned-session union member into its enriched wrapper, keyed off the `type` discriminant. + if isinstance(raw, RawDraftSession): + return AgentDraftSession(raw, client) + return AgentSession(raw, client) + + +def _wrap_owned_sessions_pager( + raw_pager: SyncPager[ListOwnedSessionsResponseDataItem, ListOwnedSessionsResponse], + client: TrueFoundryGateway, +) -> SyncPager[typing.Union[AgentSession, AgentDraftSession], ListOwnedSessionsResponse]: + wrapped_items = [_wrap_owned_session(s, client) for s in (raw_pager.items or [])] + + def get_next() -> typing.Optional[ + SyncPager[typing.Union[AgentSession, AgentDraftSession], ListOwnedSessionsResponse] + ]: + if raw_pager.get_next is None: + return None + next_raw = raw_pager.get_next() + if next_raw is None: + return None + return _wrap_owned_sessions_pager(next_raw, client) + + return SyncPager( + get_next=get_next if raw_pager.has_next else None, + has_next=raw_pager.has_next, + items=wrapped_items, + response=raw_pager.response, + ) + + +async def _async_wrap_draft_sessions_pager( + raw_pager: AsyncPager[RawDraftSession, ListDraftSessionsResponse], + client: AsyncTrueFoundryGateway, +) -> AsyncPager[AsyncAgentDraftSession, ListDraftSessionsResponse]: + wrapped_items = [AsyncAgentDraftSession(d, client) for d in (raw_pager.items or [])] + + async def get_next() -> typing.Optional[AsyncPager[AsyncAgentDraftSession, ListDraftSessionsResponse]]: + if raw_pager.get_next is None: + return None + next_raw = await raw_pager.get_next() + if next_raw is None: + return None + return await _async_wrap_draft_sessions_pager(next_raw, client) + + return AsyncPager( + get_next=get_next if raw_pager.has_next else None, + has_next=raw_pager.has_next, + items=wrapped_items, + response=raw_pager.response, + ) + + +def _async_wrap_owned_session( + raw: ListOwnedSessionsResponseDataItem, + client: AsyncTrueFoundryGateway, +) -> typing.Union[AsyncAgentSession, AsyncAgentDraftSession]: + if isinstance(raw, RawDraftSession): + return AsyncAgentDraftSession(raw, client) + return AsyncAgentSession(raw, client) + + +async def _async_wrap_owned_sessions_pager( + raw_pager: AsyncPager[ListOwnedSessionsResponseDataItem, ListOwnedSessionsResponse], + client: AsyncTrueFoundryGateway, +) -> AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], ListOwnedSessionsResponse]: + wrapped_items = [_async_wrap_owned_session(s, client) for s in (raw_pager.items or [])] + + async def get_next() -> typing.Optional[ + AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], ListOwnedSessionsResponse] + ]: + if raw_pager.get_next is None: + return None + next_raw = await raw_pager.get_next() + if next_raw is None: + return None + return await _async_wrap_owned_sessions_pager(next_raw, client) + + return AsyncPager( + get_next=get_next if raw_pager.has_next else None, + has_next=raw_pager.has_next, + items=wrapped_items, + response=raw_pager.response, + ) + + +class PrivateAgentSessionClient: + """ + High-level client for caller-scoped agent session internals: draft session create/get and + listings that span a subject's saved sessions and drafts, returning enriched + :class:`AgentSession` / :class:`AgentDraftSession` objects instead of raw response types. + """ + + def __init__( + self, + *, + base_url: str, + api_key: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = os.getenv("TFY_API_KEY"), + headers: typing.Optional[typing.Dict[str, str]] = None, + timeout: typing.Optional[float] = None, + max_retries: typing.Optional[int] = None, + follow_redirects: typing.Optional[bool] = True, + httpx_client: typing.Optional[httpx.Client] = None, + logging: typing.Optional[typing.Union[LogConfig, Logger]] = None, + ) -> None: + from ...client import TrueFoundryGateway + + self._client = TrueFoundryGateway( + base_url=base_url, + api_key=api_key, + headers=headers, + timeout=timeout, + max_retries=max_retries, + follow_redirects=follow_redirects, + httpx_client=httpx_client, + logging=logging, + ) + + def create_draft_session( + self, + *, + agent_spec: AgentSpec, + agent_name: typing.Optional[str] = None, + tfy_metadata: typing.Optional[str] = None, + request_options: typing.Optional[RequestOptions] = None, + ) -> AgentDraftSession: + """ + Create a draft session holding an inline agent spec, optionally linked to a saved agent. + + Parameters + ---------- + agent_spec : AgentSpec + Inline agent spec held by the draft. + agent_name : typing.Optional[str] + Optionally link the draft to an existing saved agent. Omit for a standalone draft. + tfy_metadata : typing.Optional[str] + Optional request metadata (x-tfy-metadata) persisted at creation. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AgentDraftSession + The created draft session. + """ + response = self._client.agents.private.draft_sessions.create( + agent_spec=agent_spec, + agent_name=agent_name, + tfy_metadata=tfy_metadata, + request_options=request_options, + ) + return AgentDraftSession(response.data, self._client) + + def get_draft_session( + self, + draft_session_id: str, + *, + request_options: typing.Optional[RequestOptions] = None, + ) -> AgentDraftSession: + """ + Fetch a draft session by ID (owner-only). + + Parameters + ---------- + draft_session_id : str + Unique identifier of the draft session to fetch. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AgentDraftSession + Draft session data. + """ + response = self._client.agents.private.draft_sessions.get(draft_session_id, request_options=request_options) + return AgentDraftSession(response.data, self._client) + + def list_draft_sessions( + self, + *, + agent_name: typing.Optional[str] = None, + limit: typing.Optional[int] = 10, + order: typing.Optional[ListDraftSessionsOrder] = None, + page_token: typing.Optional[str] = None, + start_timestamp: typing.Optional[str] = None, + end_timestamp: typing.Optional[str] = None, + request_options: typing.Optional[RequestOptions] = None, + ) -> SyncPager[AgentDraftSession, ListDraftSessionsResponse]: + """ + List the caller-owned draft sessions (newest first by default). + + Parameters + ---------- + agent_name : typing.Optional[str] + Filter to drafts linked to this saved agent. Omit for all owned drafts. + limit : typing.Optional[int] + Page size. Default 10. + order : typing.Optional[ListDraftSessionsOrder] + Sort by creation time. Default ``desc``. + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + start_timestamp : typing.Optional[str] + Inclusive lower bound on ``created_at`` (ISO-8601). + end_timestamp : typing.Optional[str] + Inclusive upper bound on ``created_at`` (ISO-8601). + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + SyncPager[AgentDraftSession, ListDraftSessionsResponse] + Paginated draft sessions. + """ + raw_pager = self._client.agents.private.draft_sessions.list( + agent_name=agent_name, + limit=limit, + order=order, + page_token=page_token, + start_timestamp=start_timestamp, + end_timestamp=end_timestamp, + request_options=request_options, + ) + return _wrap_draft_sessions_pager(raw_pager, self._client) + + def list_owned_sessions( + self, + *, + agent_name: typing.Optional[str] = None, + limit: typing.Optional[int] = 10, + order: typing.Optional[ListOwnedSessionsOrder] = None, + page_token: typing.Optional[str] = None, + start_timestamp: typing.Optional[str] = None, + end_timestamp: typing.Optional[str] = None, + request_options: typing.Optional[RequestOptions] = None, + ) -> SyncPager[typing.Union[AgentSession, AgentDraftSession], ListOwnedSessionsResponse]: + """ + List all sessions owned by the caller, spanning both saved sessions and drafts (newest first by default). + + Parameters + ---------- + agent_name : typing.Optional[str] + Filter to sessions linked to this saved agent. Omit for all owned sessions. + limit : typing.Optional[int] + Page size. Default 10. + order : typing.Optional[ListOwnedSessionsOrder] + Sort by creation time. Default ``desc``. + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + start_timestamp : typing.Optional[str] + Inclusive lower bound on ``created_at`` (ISO-8601). + end_timestamp : typing.Optional[str] + Inclusive upper bound on ``created_at`` (ISO-8601). + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + SyncPager[typing.Union[AgentSession, AgentDraftSession], ListOwnedSessionsResponse] + Paginated owned sessions. + """ + raw_pager = self._client.agents.private.sessions.list_owned_sessions( + agent_name=agent_name, + limit=limit, + order=order, + page_token=page_token, + start_timestamp=start_timestamp, + end_timestamp=end_timestamp, + request_options=request_options, + ) + return _wrap_owned_sessions_pager(raw_pager, self._client) + + +class AsyncPrivateAgentSessionClient: + """ + Async version of PrivateAgentSessionClient. + """ + + def __init__( + self, + *, + base_url: str, + api_key: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = os.getenv("TFY_API_KEY"), + headers: typing.Optional[typing.Dict[str, str]] = None, + async_token: typing.Optional[typing.Callable[[], typing.Awaitable[str]]] = None, + timeout: typing.Optional[float] = None, + max_retries: typing.Optional[int] = None, + follow_redirects: typing.Optional[bool] = True, + httpx_client: typing.Optional[httpx.AsyncClient] = None, + logging: typing.Optional[typing.Union[LogConfig, Logger]] = None, + ) -> None: + from ...client import AsyncTrueFoundryGateway + + self._client = AsyncTrueFoundryGateway( + base_url=base_url, + api_key=api_key, + headers=headers, + async_token=async_token, + timeout=timeout, + max_retries=max_retries, + follow_redirects=follow_redirects, + httpx_client=httpx_client, + logging=logging, + ) + + async def create_draft_session( + self, + *, + agent_spec: AgentSpec, + agent_name: typing.Optional[str] = None, + tfy_metadata: typing.Optional[str] = None, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncAgentDraftSession: + """ + Create a draft session holding an inline agent spec, optionally linked to a saved agent. + + Parameters + ---------- + agent_spec : AgentSpec + Inline agent spec held by the draft. + agent_name : typing.Optional[str] + Optionally link the draft to an existing saved agent. Omit for a standalone draft. + tfy_metadata : typing.Optional[str] + Optional request metadata (x-tfy-metadata) persisted at creation. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncAgentDraftSession + The created draft session. + """ + response = await self._client.agents.private.draft_sessions.create( + agent_spec=agent_spec, + agent_name=agent_name, + tfy_metadata=tfy_metadata, + request_options=request_options, + ) + return AsyncAgentDraftSession(response.data, self._client) + + async def get_draft_session( + self, + draft_session_id: str, + *, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncAgentDraftSession: + """ + Fetch a draft session by ID (owner-only). + + Parameters + ---------- + draft_session_id : str + Unique identifier of the draft session to fetch. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncAgentDraftSession + Draft session data. + """ + response = await self._client.agents.private.draft_sessions.get( + draft_session_id, request_options=request_options + ) + return AsyncAgentDraftSession(response.data, self._client) + + async def list_draft_sessions( + self, + *, + agent_name: typing.Optional[str] = None, + limit: typing.Optional[int] = 10, + order: typing.Optional[ListDraftSessionsOrder] = None, + page_token: typing.Optional[str] = None, + start_timestamp: typing.Optional[str] = None, + end_timestamp: typing.Optional[str] = None, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncPager[AsyncAgentDraftSession, ListDraftSessionsResponse]: + """ + List the caller-owned draft sessions (newest first by default). + + Parameters + ---------- + agent_name : typing.Optional[str] + Filter to drafts linked to this saved agent. Omit for all owned drafts. + limit : typing.Optional[int] + Page size. Default 10. + order : typing.Optional[ListDraftSessionsOrder] + Sort by creation time. Default ``desc``. + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + start_timestamp : typing.Optional[str] + Inclusive lower bound on ``created_at`` (ISO-8601). + end_timestamp : typing.Optional[str] + Inclusive upper bound on ``created_at`` (ISO-8601). + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncPager[AsyncAgentDraftSession, ListDraftSessionsResponse] + Paginated draft sessions. + """ + raw_pager = await self._client.agents.private.draft_sessions.list( + agent_name=agent_name, + limit=limit, + order=order, + page_token=page_token, + start_timestamp=start_timestamp, + end_timestamp=end_timestamp, + request_options=request_options, + ) + return await _async_wrap_draft_sessions_pager(raw_pager, self._client) + + async def list_owned_sessions( + self, + *, + agent_name: typing.Optional[str] = None, + limit: typing.Optional[int] = 10, + order: typing.Optional[ListOwnedSessionsOrder] = None, + page_token: typing.Optional[str] = None, + start_timestamp: typing.Optional[str] = None, + end_timestamp: typing.Optional[str] = None, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], ListOwnedSessionsResponse]: + """ + List all sessions owned by the caller, spanning both saved sessions and drafts (newest first by default). + + Parameters + ---------- + agent_name : typing.Optional[str] + Filter to sessions linked to this saved agent. Omit for all owned sessions. + limit : typing.Optional[int] + Page size. Default 10. + order : typing.Optional[ListOwnedSessionsOrder] + Sort by creation time. Default ``desc``. + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + start_timestamp : typing.Optional[str] + Inclusive lower bound on ``created_at`` (ISO-8601). + end_timestamp : typing.Optional[str] + Inclusive upper bound on ``created_at`` (ISO-8601). + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncPager[typing.Union[AsyncAgentSession, AsyncAgentDraftSession], ListOwnedSessionsResponse] + Paginated owned sessions. + """ + raw_pager = await self._client.agents.private.sessions.list_owned_sessions( + agent_name=agent_name, + limit=limit, + order=order, + page_token=page_token, + start_timestamp=start_timestamp, + end_timestamp=end_timestamp, + request_options=request_options, + ) + return await _async_wrap_owned_sessions_pager(raw_pager, self._client) diff --git a/src/truefoundry_gateway_sdk/agents/session_mixin.py b/src/truefoundry_gateway_sdk/agents/session_mixin.py new file mode 100644 index 0000000..a225f90 --- /dev/null +++ b/src/truefoundry_gateway_sdk/agents/session_mixin.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +import typing + +from ..core.pagination import AsyncPager, SyncPager +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.list_session_events_response import ListSessionEventsResponse + from ..types.list_turns_response import ListTurnsResponse + from ..types.previous_turn_id_input import PreviousTurnIdInput + from ..types.session_event_item import SessionEventItem + from ..types.turn import Turn as RawTurn + from ..types.turn_input_item import TurnInputItem + from .agent_session import AgentSession, AsyncAgentSession + from .private.agent_draft_session import AgentDraftSession, AsyncAgentDraftSession + + SyncSessionOwner = typing.Union[AgentSession, AgentDraftSession] + AsyncSessionOwner = typing.Union[AsyncAgentSession, AsyncAgentDraftSession] + + +def _wrap_turns_pager( + raw_pager: SyncPager[RawTurn, ListTurnsResponse], + owner: SyncSessionOwner, + client: TrueFoundryGateway, +) -> SyncPager[Turn, ListTurnsResponse]: + wrapped_items = [Turn(t, owner, client) for t in (raw_pager.items or [])] + + def get_next() -> typing.Optional[SyncPager[Turn, ListTurnsResponse]]: + if raw_pager.get_next is None: + return None + next_raw = raw_pager.get_next() + if next_raw is None: + return None + return _wrap_turns_pager(next_raw, owner, client) + + return SyncPager( + get_next=get_next if raw_pager.has_next else None, + has_next=raw_pager.has_next, + items=wrapped_items, + response=raw_pager.response, + ) + + +async def _async_wrap_turns_pager( + raw_pager: AsyncPager[RawTurn, ListTurnsResponse], + owner: AsyncSessionOwner, + client: AsyncTrueFoundryGateway, +) -> AsyncPager[AsyncTurn, ListTurnsResponse]: + wrapped_items = [AsyncTurn(t, owner, client) for t in (raw_pager.items or [])] + + async def get_next() -> typing.Optional[AsyncPager[AsyncTurn, ListTurnsResponse]]: + if raw_pager.get_next is None: + return None + next_raw = await raw_pager.get_next() + if next_raw is None: + return None + return await _async_wrap_turns_pager(next_raw, owner, client) + + return AsyncPager( + get_next=get_next if raw_pager.has_next else None, + has_next=raw_pager.has_next, + items=wrapped_items, + response=raw_pager.response, + ) + + +class SessionMixin: + """ + Shared turn behavior keyed by a session id. Both :class:`AgentSession` and + :class:`AgentDraftSession` hold a SessionMixin and delegate prepare_turn / list_turns / + get_turn / cancel / list_events to it, so the two session wrappers expose identical turn + operations without duplicating logic. + + The owning wrapper is NOT stored on the mixin (that back-reference would form an + ``AgentSession`` <-> ``SessionMixin`` cycle that outlives both). Instead the wrapper passes + itself as the first argument to each turn-producing method, so turns still expose the + enriched wrapper (with ``agent_name``, ``title``, etc.) as ``turn.session``. + """ + + def __init__(self, id: str, client: TrueFoundryGateway) -> None: + self._id = id + self._client = client + + def prepare_turn( + self, + owner: SyncSessionOwner, + *, + input: typing.Optional[typing.Sequence[TurnInputItem]] = None, + previous_turn_id: typing.Optional[PreviousTurnIdInput] = None, + ) -> PreparedTurn: + """ + Stage a turn locally; call ``execute()`` to start ``create_turn``. + + Parameters + ---------- + owner : SyncSessionOwner + Enriched wrapper surfaced as ``turn.session`` on the resulting turn. + input : typing.Optional[typing.Sequence[TurnInputItem]] + Turn input items passed to create turn. + previous_turn_id : typing.Optional[PreviousTurnIdInput] + Previous turn to chain from. Defaults to ``auto``. + + Returns + ------- + PreparedTurn + Staged turn. + """ + return PreparedTurn( + input=input, + previous_turn_id=previous_turn_id, + session=owner, + client=self._client, + ) + + def list_turns( + self, + owner: SyncSessionOwner, + *, + page_token: typing.Optional[str] = None, + limit: typing.Optional[int] = 10, + request_options: typing.Optional[RequestOptions] = None, + ) -> SyncPager[Turn, ListTurnsResponse]: + """ + List turns in this session. + + Parameters + ---------- + owner : SyncSessionOwner + Enriched wrapper surfaced as ``turn.session`` on each listed turn. + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + limit : typing.Optional[int] + Page size. Default 10. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + SyncPager[Turn, ListTurnsResponse] + Paginated turns. + """ + raw_pager = self._client.agents.sessions.list_turns( + self._id, page_token=page_token, limit=limit, request_options=request_options + ) + return _wrap_turns_pager(raw_pager, owner, self._client) + + def get_turn( + self, + owner: SyncSessionOwner, + turn_id: str, + *, + request_options: typing.Optional[RequestOptions] = None, + ) -> Turn: + """ + Fetch a turn by ID. + + Parameters + ---------- + owner : SyncSessionOwner + Enriched wrapper surfaced as ``turn.session`` on the resulting turn. + turn_id : str + Unique identifier of the turn to fetch. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + Turn + Turn data. + """ + response = self._client.agents.sessions.get_turn(self._id, turn_id, request_options=request_options) + return Turn(response.data, owner, self._client) + + def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None: + """ + Cancel the running last turn for the session. + + Parameters + ---------- + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + None + """ + self._client.agents.sessions.cancel(self._id, request_options=request_options) + + def list_events( + self, + *, + page_token: typing.Optional[str] = None, + last_turn_id: typing.Optional[str] = None, + limit: typing.Optional[int] = 100, + request_options: typing.Optional[RequestOptions] = None, + ) -> SyncPager[SessionEventItem, ListSessionEventsResponse]: + """ + Paginated session events across turns (newest first); subscribe to a running turn for live events. + + Parameters + ---------- + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + last_turn_id : typing.Optional[str] + Newest turn in the listing window (initial load only). Omit to use the session last turn. + limit : typing.Optional[int] + Page size. Default 100. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + SyncPager[SessionEventItem, ListSessionEventsResponse] + Paginated session events. + """ + return self._client.agents.sessions.list_events( + self._id, + page_token=page_token, + last_turn_id=last_turn_id, + limit=limit, + request_options=request_options, + ) + + +class AsyncSessionMixin: + """ + Async version of :class:`SessionMixin`. Same delegation contract; the owning wrapper passes + itself as the first argument to each turn-producing method so turns expose the enriched + wrapper as ``turn.session``. + """ + + def __init__(self, id: str, client: AsyncTrueFoundryGateway) -> None: + self._id = id + self._client = client + + def prepare_turn( + self, + owner: AsyncSessionOwner, + *, + input: typing.Optional[typing.Sequence[TurnInputItem]] = None, + previous_turn_id: typing.Optional[PreviousTurnIdInput] = None, + ) -> AsyncPreparedTurn: + """ + Stage a turn locally; call ``execute()`` to start ``create_turn``. + + Parameters + ---------- + owner : AsyncSessionOwner + Enriched wrapper surfaced as ``turn.session`` on the resulting turn. + input : typing.Optional[typing.Sequence[TurnInputItem]] + Turn input items passed to create turn. + previous_turn_id : typing.Optional[PreviousTurnIdInput] + Previous turn to chain from. Defaults to ``auto``. + + Returns + ------- + AsyncPreparedTurn + Staged turn. + """ + return AsyncPreparedTurn( + input=input, + previous_turn_id=previous_turn_id, + session=owner, + client=self._client, + ) + + async def list_turns( + self, + owner: AsyncSessionOwner, + *, + page_token: typing.Optional[str] = None, + limit: typing.Optional[int] = 10, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncPager[AsyncTurn, ListTurnsResponse]: + """ + List turns in this session. + + Parameters + ---------- + owner : AsyncSessionOwner + Enriched wrapper surfaced as ``turn.session`` on each listed turn. + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + limit : typing.Optional[int] + Page size. Default 10. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncPager[AsyncTurn, ListTurnsResponse] + Paginated turns. + """ + raw_pager = await self._client.agents.sessions.list_turns( + self._id, page_token=page_token, limit=limit, request_options=request_options + ) + return await _async_wrap_turns_pager(raw_pager, owner, self._client) + + async def get_turn( + self, + owner: AsyncSessionOwner, + turn_id: str, + *, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncTurn: + """ + Fetch a turn by ID. + + Parameters + ---------- + owner : AsyncSessionOwner + Enriched wrapper surfaced as ``turn.session`` on the resulting turn. + turn_id : str + Unique identifier of the turn to fetch. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncTurn + Turn data. + """ + response = await self._client.agents.sessions.get_turn(self._id, turn_id, request_options=request_options) + return AsyncTurn(response.data, owner, self._client) + + async def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None: + """ + Cancel the running last turn for the session. + + Parameters + ---------- + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + None + """ + await self._client.agents.sessions.cancel(self._id, request_options=request_options) + + async def list_events( + self, + *, + page_token: typing.Optional[str] = None, + last_turn_id: typing.Optional[str] = None, + limit: typing.Optional[int] = 100, + request_options: typing.Optional[RequestOptions] = None, + ) -> AsyncPager[SessionEventItem, ListSessionEventsResponse]: + """ + Paginated session events across turns (newest first); subscribe to a running turn for live events. + + Parameters + ---------- + page_token : typing.Optional[str] + Token from the previous response ``next_page_token``. + last_turn_id : typing.Optional[str] + Newest turn in the listing window (initial load only). Omit to use the session last turn. + limit : typing.Optional[int] + Page size. Default 100. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + AsyncPager[SessionEventItem, ListSessionEventsResponse] + Paginated session events. + """ + return await self._client.agents.sessions.list_events( + self._id, + page_token=page_token, + last_turn_id=last_turn_id, + limit=limit, + request_options=request_options, + ) diff --git a/src/truefoundry_gateway_sdk/agents/turn.py b/src/truefoundry_gateway_sdk/agents/turn.py index 15a097b..cc78277 100644 --- a/src/truefoundry_gateway_sdk/agents/turn.py +++ b/src/truefoundry_gateway_sdk/agents/turn.py @@ -23,6 +23,7 @@ from ..types.turn_state import TurnState from ..types.turn_streaming_event import TurnStreamingEvent from .agent_session import AgentSession, AsyncAgentSession + from .private.agent_draft_session import AgentDraftSession, AsyncAgentDraftSession # waitForCompletion poll interval in milliseconds: default and enforced minimum. _DEFAULT_POLL_INTERVAL_MS = 3000 @@ -38,7 +39,7 @@ class Turn: def __init__( self, turn: RawTurn, - session: "AgentSession", + session: typing.Union["AgentSession", "AgentDraftSession"], client: TrueFoundryGateway, ) -> None: self._id: str = turn.id @@ -48,7 +49,7 @@ def __init__( self._created_by_subject: Subject = turn.created_by_subject self._created_at: str = turn.created_at self._state: TurnState = turn.state - self._session: "AgentSession" = session + self._session: typing.Union["AgentSession", "AgentDraftSession"] = session self._client = client def __repr__(self) -> str: @@ -125,11 +126,11 @@ def state(self) -> TurnState: return self._state @property - def session(self) -> "AgentSession": + def session(self) -> typing.Union["AgentSession", "AgentDraftSession"]: """ Returns ------- - AgentSession + typing.Union[AgentSession, AgentDraftSession] Parent session this turn belongs to. """ return self._session @@ -280,7 +281,7 @@ class AsyncTurn: def __init__( self, turn: RawTurn, - session: "AsyncAgentSession", + session: typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"], client: AsyncTrueFoundryGateway, ) -> None: self._id: str = turn.id @@ -290,7 +291,7 @@ def __init__( self._created_by_subject: Subject = turn.created_by_subject self._created_at: str = turn.created_at self._state: TurnState = turn.state - self._session: "AsyncAgentSession" = session + self._session: typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"] = session self._client: AsyncTrueFoundryGateway = client def __repr__(self) -> str: @@ -367,11 +368,11 @@ def state(self) -> TurnState: return self._state @property - def session(self) -> "AsyncAgentSession": + def session(self) -> typing.Union["AsyncAgentSession", "AsyncAgentDraftSession"]: """ Returns ------- - AsyncAgentSession + typing.Union[AsyncAgentSession, AsyncAgentDraftSession] Parent session this turn belongs to. """ return self._session diff --git a/tests/custom/test_private_agent_session_client.py b/tests/custom/test_private_agent_session_client.py new file mode 100644 index 0000000..16853fc --- /dev/null +++ b/tests/custom/test_private_agent_session_client.py @@ -0,0 +1,128 @@ +import types + +from truefoundry_gateway_sdk.agents.agent_session import AgentSession +from truefoundry_gateway_sdk.agents.prepared_turn import PreparedTurn +from truefoundry_gateway_sdk.agents.private import AgentDraftSession, PrivateAgentSessionClient +from truefoundry_gateway_sdk.agents.turn import Turn +from truefoundry_gateway_sdk.core.pagination import SyncPager +from truefoundry_gateway_sdk.types.agent_spec import AgentSpec +from truefoundry_gateway_sdk.types.draft_session import DraftSession as RawDraftSession +from truefoundry_gateway_sdk.types.model import Model +from truefoundry_gateway_sdk.types.session import Session as RawSession +from truefoundry_gateway_sdk.types.subject import Subject +from truefoundry_gateway_sdk.types.turn import Turn as RawTurn +from truefoundry_gateway_sdk.types.turn_state_running import TurnStateRunning + +SUBJECT = Subject(subject_id="subject_id", subject_type="subject_type", subject_slug="subject_slug") + +RAW_SESSION = RawSession( + id="session-1", + agent_name="agent_name", + title="a session", + created_by_subject=SUBJECT, + created_at="2026-01-01T00:00:00Z", + updated_at="2026-01-01T00:00:00Z", +) + +RAW_DRAFT = RawDraftSession( + id="draft-1", + agent_spec=AgentSpec(model=Model(name="gpt-4o")), + agent_name="agent_name", + title="a draft", + created_by_subject=SUBJECT, + created_at="2026-01-01T00:00:00Z", + updated_at="2026-01-01T00:00:00Z", +) + + +def _new_client() -> PrivateAgentSessionClient: + return PrivateAgentSessionClient(base_url="https://example.com", api_key="test", max_retries=0) + + +def _pager(items, *, next_pager=None): + return SyncPager( + get_next=(lambda: next_pager) if next_pager is not None else None, + has_next=next_pager is not None, + items=items, + response=None, + ) + + +def test_list_owned_sessions_wraps_each_member_by_discriminant(monkeypatch): + client = _new_client() + raw_sessions = client._client.agents.private.sessions + monkeypatch.setattr(raw_sessions, "list_owned_sessions", lambda **kw: _pager([RAW_SESSION, RAW_DRAFT])) + + items = list(client.list_owned_sessions()) + assert len(items) == 2 + + session, draft = items + assert isinstance(session, AgentSession) + assert session.type == "session" + assert session.id == "session-1" + + assert isinstance(draft, AgentDraftSession) + assert draft.type == "session/draft" + assert draft.id == "draft-1" + assert draft.agent_spec.model.name == "gpt-4o" + + +def test_list_owned_sessions_paginates_via_next_page_token(monkeypatch): + client = _new_client() + raw_sessions = client._client.agents.private.sessions + page2 = _pager([RAW_SESSION]) + monkeypatch.setattr(raw_sessions, "list_owned_sessions", lambda **kw: _pager([RAW_SESSION], next_pager=page2)) + + page = client.list_owned_sessions(limit=1) + assert page.has_next is True + next_page = page.get_next() + assert next_page is not None + assert isinstance(next_page.items[0], AgentSession) + + +def test_list_draft_sessions_wraps_each_raw_draft(monkeypatch): + client = _new_client() + raw_drafts = client._client.agents.private.draft_sessions + monkeypatch.setattr(raw_drafts, "list", lambda **kw: _pager([RAW_DRAFT])) + + items = list(client.list_draft_sessions()) + assert len(items) == 1 + assert isinstance(items[0], AgentDraftSession) + assert items[0].id == "draft-1" + + +def test_prepare_turn_returns_prepared_turn_without_request(monkeypatch): + client = _new_client() + raw_sessions = client._client.agents.private.sessions + monkeypatch.setattr(raw_sessions, "list_owned_sessions", lambda **kw: _pager([RAW_DRAFT])) + + (draft,) = list(client.list_owned_sessions()) + prepared = draft.prepare_turn() + assert isinstance(prepared, PreparedTurn) + assert prepared.session is draft + assert isinstance(prepared.session, AgentDraftSession) + assert prepared.session.agent_name == "agent_name" + + +def test_get_turn_session_is_enriched_owner(monkeypatch): + client = _new_client() + raw_owned = client._client.agents.private.sessions + monkeypatch.setattr(raw_owned, "list_owned_sessions", lambda **kw: _pager([RAW_SESSION])) + + raw_turn = RawTurn( + id="turn-1", + session_id="session-1", + created_by_subject=SUBJECT, + created_at="2026-01-01T00:00:00Z", + state=TurnStateRunning(), + ) + raw_turns = client._client.agents.sessions + monkeypatch.setattr(raw_turns, "get_turn", lambda session_id, turn_id, **kw: types.SimpleNamespace(data=raw_turn)) + + (session,) = list(client.list_owned_sessions()) + turn = session.get_turn("turn-1") + assert isinstance(turn, Turn) + assert turn.session is session + assert isinstance(turn.session, AgentSession) + assert turn.session.agent_name == "agent_name" + assert turn.session.title == "a session" From c6d02b6fd50aef8422f606904fd6b48613634112 Mon Sep 17 00:00:00 2001 From: Srajan Asthana Date: Fri, 17 Jul 2026 20:48:20 +0530 Subject: [PATCH 2/2] nit --- .../private/private_agent_session_client.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py b/src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py index 6522b1e..9a78ee3 100644 --- a/src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py +++ b/src/truefoundry_gateway_sdk/agents/private/private_agent_session_client.py @@ -4,7 +4,6 @@ import typing from ...core.pagination import AsyncPager, SyncPager -from ...types.draft_session import DraftSession as RawDraftSession from ..agent_session import AgentSession, AsyncAgentSession from .agent_draft_session import AgentDraftSession, AsyncAgentDraftSession @@ -14,6 +13,7 @@ from ...core.logging import LogConfig, Logger from ...core.request_options import RequestOptions from ...types.agent_spec import AgentSpec + from ...types.draft_session import DraftSession as RawDraftSession from ...types.list_draft_sessions_order import ListDraftSessionsOrder from ...types.list_draft_sessions_response import ListDraftSessionsResponse from ...types.list_owned_sessions_order import ListOwnedSessionsOrder @@ -48,9 +48,11 @@ def _wrap_owned_session( client: TrueFoundryGateway, ) -> typing.Union[AgentSession, AgentDraftSession]: # Dispatch a raw owned-session union member into its enriched wrapper, keyed off the `type` discriminant. - if isinstance(raw, RawDraftSession): + if raw.type == "session/draft": return AgentDraftSession(raw, client) - return AgentSession(raw, client) + if raw.type == "session": + return AgentSession(raw, client) + raise ValueError(f"Unknown owned session type: {raw.type!r}") def _wrap_owned_sessions_pager( @@ -103,9 +105,12 @@ def _async_wrap_owned_session( raw: ListOwnedSessionsResponseDataItem, client: AsyncTrueFoundryGateway, ) -> typing.Union[AsyncAgentSession, AsyncAgentDraftSession]: - if isinstance(raw, RawDraftSession): + # Dispatch a raw owned-session union member into its enriched wrapper, keyed off the `type` discriminant. + if raw.type == "session/draft": return AsyncAgentDraftSession(raw, client) - return AsyncAgentSession(raw, client) + if raw.type == "session": + return AsyncAgentSession(raw, client) + raise ValueError(f"Unknown owned session type: {raw.type!r}") async def _async_wrap_owned_sessions_pager(