From 51cb7de97acdbd5517b335303db373fa600f84e9 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Tue, 14 Jul 2026 07:35:52 +0530 Subject: [PATCH] Add AgentDraftSession as private class --- .../agents/__init__.py | 4 +- .../agents/agent_session.py | 132 ++++++++++----- .../agents/prepared_turn.py | 18 +- .../agents/private/__init__.py | 3 + .../agents/private/agent_draft_session.py | 156 ++++++++++++++++++ src/truefoundry_gateway_sdk/agents/turn.py | 18 +- 6 files changed, 272 insertions(+), 59 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 diff --git a/src/truefoundry_gateway_sdk/agents/__init__.py b/src/truefoundry_gateway_sdk/agents/__init__.py index 4e53269..fc02f9d 100644 --- a/src/truefoundry_gateway_sdk/agents/__init__.py +++ b/src/truefoundry_gateway_sdk/agents/__init__.py @@ -28,7 +28,7 @@ from ..types.turn_streaming_event import TurnStreamingEvent from ..types.user_tool_approval_event import UserToolApprovalEvent from ..types.user_tool_response_event import UserToolResponseEvent -from .agent_session import AgentSession, AsyncAgentSession +from .agent_session import AgentSession, AsyncAgentSession, AsyncBaseAgentSession, BaseAgentSession from .agent_session_client import AgentSessionClient, AsyncAgentSessionClient from .event_delta import DeltaEvents, is_event_delta, merge_event_delta from .prepared_turn import AsyncPreparedTurn, PreparedTurn @@ -42,6 +42,8 @@ "AgentSessionClient", "AsyncAgentSession", "AsyncAgentSessionClient", + "AsyncBaseAgentSession", + "BaseAgentSession", "AsyncPreparedTurn", "AsyncTurn", "DeltaEvents", diff --git a/src/truefoundry_gateway_sdk/agents/agent_session.py b/src/truefoundry_gateway_sdk/agents/agent_session.py index 792967d..e0ed16b 100644 --- a/src/truefoundry_gateway_sdk/agents/agent_session.py +++ b/src/truefoundry_gateway_sdk/agents/agent_session.py @@ -21,7 +21,7 @@ def _wrap_turns_pager( raw_pager: SyncPager[RawTurn, ListTurnsResponse], - session: "AgentSession", + session: "BaseAgentSession", client: TrueFoundryGateway, ) -> SyncPager[Turn, ListTurnsResponse]: wrapped_items = [Turn(t, session, client) for t in (raw_pager.items or [])] @@ -44,7 +44,7 @@ def get_next() -> typing.Optional[SyncPager[Turn, ListTurnsResponse]]: async def _async_wrap_turns_pager( raw_pager: AsyncPager[RawTurn, ListTurnsResponse], - session: "AsyncAgentSession", + session: "AsyncBaseAgentSession", client: AsyncTrueFoundryGateway, ) -> AsyncPager[AsyncTurn, ListTurnsResponse]: wrapped_items = [AsyncTurn(t, session, client) for t in (raw_pager.items or [])] @@ -65,23 +65,29 @@ async def get_next() -> typing.Optional[AsyncPager[AsyncTurn, ListTurnsResponse] ) -class AgentSession: +class BaseAgentSession: """ - A session enriched with convenience methods: prepare_turn, list_turns, get_turn, list_events, cancel. + Shared base for AgentSession and AgentDraftSession. Holds common identity fields + and provides turn/event convenience methods. """ - def __init__(self, session: RawSession, client: TrueFoundryGateway) -> None: - self._id: str = session.id - self._agent_name: 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 + def __init__( + self, + *, + id: str, + title: typing.Optional[str], + created_by_subject: Subject, + created_at: str, + updated_at: str, + client: TrueFoundryGateway, + ) -> None: + self._id: str = id + self._title: typing.Optional[str] = title + self._created_by_subject: Subject = created_by_subject + self._created_at: str = created_at + self._updated_at: str = updated_at self._client = client - def __repr__(self) -> str: - return f"AgentSession(id={self._id!r}, agent_name={self._agent_name!r})" - @property def id(self) -> str: """ @@ -92,16 +98,6 @@ def id(self) -> str: """ return self._id - @property - def agent_name(self) -> str: - """ - Returns - ------- - str - Name of the agent for this session. - """ - return self._agent_name - @property def title(self) -> typing.Optional[str]: """ @@ -274,42 +270,68 @@ def list_events( ) -class AsyncAgentSession: +class AgentSession(BaseAgentSession): """ - Async version of AgentSession. + A session enriched with convenience methods: prepare_turn, list_turns, get_turn, list_events, cancel. """ - def __init__(self, session: RawSession, client: AsyncTrueFoundryGateway) -> None: - self._id: str = session.id + def __init__(self, session: RawSession, client: TrueFoundryGateway) -> None: + super().__init__( + id=session.id, + title=session.title, + created_by_subject=session.created_by_subject, + created_at=session.created_at, + updated_at=session.updated_at, + client=client, + ) self._agent_name: 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._client = client def __repr__(self) -> str: - return f"AsyncAgentSession(id={self._id!r}, agent_name={self._agent_name!r})" + return f"AgentSession(id={self._id!r}, agent_name={self._agent_name!r})" @property - def id(self) -> str: + def agent_name(self) -> str: """ Returns ------- str - Unique identifier of this session. + Name of the agent for this session. """ - return self._id + return self._agent_name + + +class AsyncBaseAgentSession: + """ + Async shared base for AsyncAgentSession and AsyncAgentDraftSession. Holds common identity + fields and provides turn/event convenience methods. + """ + + def __init__( + self, + *, + id: str, + title: typing.Optional[str], + created_by_subject: Subject, + created_at: str, + updated_at: str, + client: AsyncTrueFoundryGateway, + ) -> None: + self._id: str = id + self._title: typing.Optional[str] = title + self._created_by_subject: Subject = created_by_subject + self._created_at: str = created_at + self._updated_at: str = updated_at + self._client = client @property - def agent_name(self) -> str: + def id(self) -> str: """ Returns ------- str - Name of the agent for this session. + Unique identifier of this session. """ - return self._agent_name + return self._id @property def title(self) -> typing.Optional[str]: @@ -481,3 +503,33 @@ async def list_events( limit=limit, request_options=request_options, ) + + +class AsyncAgentSession(AsyncBaseAgentSession): + """ + Async version of AgentSession. + """ + + def __init__(self, session: RawSession, client: AsyncTrueFoundryGateway) -> None: + super().__init__( + id=session.id, + title=session.title, + created_by_subject=session.created_by_subject, + created_at=session.created_at, + updated_at=session.updated_at, + client=client, + ) + self._agent_name: str = session.agent_name + + def __repr__(self) -> str: + return f"AsyncAgentSession(id={self._id!r}, agent_name={self._agent_name!r})" + + @property + def agent_name(self) -> str: + """ + Returns + ------- + str + Name of the agent for this session. + """ + return self._agent_name diff --git a/src/truefoundry_gateway_sdk/agents/prepared_turn.py b/src/truefoundry_gateway_sdk/agents/prepared_turn.py index ec720ca..bd5ad14 100644 --- a/src/truefoundry_gateway_sdk/agents/prepared_turn.py +++ b/src/truefoundry_gateway_sdk/agents/prepared_turn.py @@ -19,7 +19,7 @@ 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 + from .agent_session import AsyncBaseAgentSession, BaseAgentSession class PreparedTurn: @@ -33,12 +33,12 @@ def __init__( *, input: typing.Optional[typing.Sequence[TurnInputItem]], previous_turn_id: typing.Optional[PreviousTurnIdInput], - session: "AgentSession", + session: "BaseAgentSession", 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: "BaseAgentSession" = session self._session_id: str = session.id self._client = client self._started: bool = False @@ -50,11 +50,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) -> "BaseAgentSession": """ Returns ------- - AgentSession + BaseAgentSession Parent session this prepared turn belongs to. """ return self._session @@ -385,12 +385,12 @@ def __init__( *, input: typing.Optional[typing.Sequence[TurnInputItem]], previous_turn_id: typing.Optional[PreviousTurnIdInput], - session: "AsyncAgentSession", + session: "AsyncBaseAgentSession", 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: "AsyncBaseAgentSession" = session self._session_id: str = session.id self._client = client self._started: bool = False @@ -402,11 +402,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) -> "AsyncBaseAgentSession": """ Returns ------- - AsyncAgentSession + AsyncBaseAgentSession 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..5a2ff38 --- /dev/null +++ b/src/truefoundry_gateway_sdk/agents/private/__init__.py @@ -0,0 +1,3 @@ +from .agent_draft_session import AgentDraftSession, AsyncAgentDraftSession + +__all__ = ["AgentDraftSession", "AsyncAgentDraftSession"] 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..946acd3 --- /dev/null +++ b/src/truefoundry_gateway_sdk/agents/private/agent_draft_session.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import typing + +from ..agent_session import AsyncBaseAgentSession, BaseAgentSession + +if typing.TYPE_CHECKING: + from ...client import AsyncTrueFoundryGateway, TrueFoundryGateway + from ...core.request_options import RequestOptions + from ...types.agent_spec import AgentSpec + from ...types.draft_session import DraftSession + + +class AgentDraftSession(BaseAgentSession): + """ + A draft session enriched with convenience methods inherited from BaseAgentSession, plus + an ``update()`` method to mutate the inline agent spec. + """ + + def __init__(self, session: DraftSession, client: TrueFoundryGateway) -> None: + super().__init__( + id=session.id, + title=session.title, + created_by_subject=session.created_by_subject, + created_at=session.created_at, + updated_at=session.updated_at, + client=client, + ) + self._agent_name: typing.Optional[str] = session.agent_name + self._agent_spec: AgentSpec = session.agent_spec + + def __repr__(self) -> str: + return f"AgentDraftSession(id={self._id!r}, agent_name={self._agent_name!r})" + + @property + def agent_name(self) -> typing.Optional[str]: + """ + Returns + ------- + typing.Optional[str] + Name of the agent this draft session is linked to, if any. + """ + return self._agent_name + + @property + def agent_spec(self) -> AgentSpec: + """ + Returns + ------- + AgentSpec + The inline agent specification for this draft session. + """ + return self._agent_spec + + def update( + self, + agent_spec: AgentSpec, + *, + request_options: typing.Optional[RequestOptions] = None, + ) -> None: + """ + Update the inline agent spec for this draft session. + + Calls ``PATCH v1/agents/draft-sessions/{id}`` and syncs ``agent_spec`` + and ``updated_at`` on this object in place. + + Parameters + ---------- + agent_spec : AgentSpec + The new agent specification to store. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + None + """ + response = self._client.agents.private.draft_sessions.update( + self._id, + agent_spec=agent_spec, + request_options=request_options, + ) + self._agent_spec = response.data.agent_spec + self._updated_at = response.data.updated_at + + +class AsyncAgentDraftSession(AsyncBaseAgentSession): + """ + Async version of AgentDraftSession. + """ + + def __init__(self, session: DraftSession, client: AsyncTrueFoundryGateway) -> None: + super().__init__( + id=session.id, + title=session.title, + created_by_subject=session.created_by_subject, + created_at=session.created_at, + updated_at=session.updated_at, + client=client, + ) + self._agent_name: typing.Optional[str] = session.agent_name + self._agent_spec: AgentSpec = session.agent_spec + + def __repr__(self) -> str: + return f"AsyncAgentDraftSession(id={self._id!r}, agent_name={self._agent_name!r})" + + @property + def agent_name(self) -> typing.Optional[str]: + """ + Returns + ------- + typing.Optional[str] + Name of the agent this draft session is linked to, if any. + """ + return self._agent_name + + @property + def agent_spec(self) -> AgentSpec: + """ + Returns + ------- + AgentSpec + The inline agent specification for this draft session. + """ + return self._agent_spec + + async def update( + self, + agent_spec: AgentSpec, + *, + request_options: typing.Optional[RequestOptions] = None, + ) -> None: + """ + Update the inline agent spec for this draft session. + + Calls ``PATCH v1/agents/draft-sessions/{id}`` and syncs ``agent_spec`` + and ``updated_at`` on this object in place. + + Parameters + ---------- + agent_spec : AgentSpec + The new agent specification to store. + request_options : typing.Optional[RequestOptions] + Overrides client timeout, retries, headers, and stream reconnect. + + Returns + ------- + None + """ + response = await self._client.agents.private.draft_sessions.update( + self._id, + agent_spec=agent_spec, + request_options=request_options, + ) + self._agent_spec = response.data.agent_spec + self._updated_at = response.data.updated_at diff --git a/src/truefoundry_gateway_sdk/agents/turn.py b/src/truefoundry_gateway_sdk/agents/turn.py index 15a097b..1f7af89 100644 --- a/src/truefoundry_gateway_sdk/agents/turn.py +++ b/src/truefoundry_gateway_sdk/agents/turn.py @@ -22,7 +22,7 @@ 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 + from .agent_session import AsyncBaseAgentSession, BaseAgentSession # waitForCompletion poll interval in milliseconds: default and enforced minimum. _DEFAULT_POLL_INTERVAL_MS = 3000 @@ -38,7 +38,7 @@ class Turn: def __init__( self, turn: RawTurn, - session: "AgentSession", + session: "BaseAgentSession", client: TrueFoundryGateway, ) -> None: self._id: str = turn.id @@ -48,7 +48,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: "BaseAgentSession" = session self._client = client def __repr__(self) -> str: @@ -125,11 +125,11 @@ def state(self) -> TurnState: return self._state @property - def session(self) -> "AgentSession": + def session(self) -> "BaseAgentSession": """ Returns ------- - AgentSession + BaseAgentSession Parent session this turn belongs to. """ return self._session @@ -280,7 +280,7 @@ class AsyncTurn: def __init__( self, turn: RawTurn, - session: "AsyncAgentSession", + session: "AsyncBaseAgentSession", client: AsyncTrueFoundryGateway, ) -> None: self._id: str = turn.id @@ -290,7 +290,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: "AsyncBaseAgentSession" = session self._client: AsyncTrueFoundryGateway = client def __repr__(self) -> str: @@ -367,11 +367,11 @@ def state(self) -> TurnState: return self._state @property - def session(self) -> "AsyncAgentSession": + def session(self) -> "AsyncBaseAgentSession": """ Returns ------- - AsyncAgentSession + AsyncBaseAgentSession Parent session this turn belongs to. """ return self._session