Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 64 additions & 30 deletions src/truefoundry_gateway_sdk/agents/prepared_turn.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@

class PreparedTurn:
"""
Output of prepare_turn: not yet started (no HTTP). execute() fires the create_turn POST
and drives the SSE stream. The inner Turn is adopted once turn.created is received.
Output of prepare_turn: not yet started (no HTTP). execute() starts the turn via
create_turn_stream (SSE) or create_turn (JSON), then drives the inner Turn.
"""

def __init__(
Expand Down Expand Up @@ -165,12 +165,13 @@ def execute(
request_options: typing.Optional[RequestOptions] = None,
) -> typing.Union[typing.Iterator[TurnStreamData], TurnState]:
"""
Start the turn via ``create_turn``.
Start the turn via ``create_turn_stream`` / ``create_turn``.

Parameters
----------
stream : bool
Stream ``create_turn`` SSE when true. Default true.
Stream ``create_turn_stream`` SSE when true. Default true.
When false, uses ``create_turn`` JSON then polls to terminal.
poll_interval_ms : int
Poll interval ms when ``stream=False``. Minimum 3000.
request_options : typing.Optional[RequestOptions]
Expand All @@ -184,7 +185,7 @@ def execute(
Yields
------
TurnStreamData
SSE stream from create_turn.
SSE stream from create_turn_stream.
"""
if self._started:
raise RuntimeError("Turn already started; use stream() / wait_for_completion().")
Expand Down Expand Up @@ -313,19 +314,25 @@ def list_events(
# --- Private helpers ---

def _run_streaming(self, request_options: typing.Optional[RequestOptions]) -> typing.Iterator[TurnStreamData]:
"""execute(stream=True) path: open create_turn SSE, adopt inner Turn on turn.created."""
"""execute(stream=True) path: open create_turn_stream SSE, adopt inner Turn on turn.created."""
yield from self._consume_stream(request_options)

def _start_and_wait(self, poll_interval_ms: int, request_options: typing.Optional[RequestOptions]) -> TurnState:
"""execute(stream=False) path: drive SSE until turn.created mints the inner Turn, then poll."""
self._create_turn_if_not_exist(request_options)
"""execute(stream=False) path: create_turn JSON returns the running turn; then poll."""
response = self._client.agents.sessions.create_turn(
self._session_id,
input=self._input_param,
previous_turn_id=self._previous_turn_id,
request_options=request_options,
)
self._adopt_turn_from_api(response.data)
return self._must_get_turn().wait_for_completion(
poll_interval_ms=poll_interval_ms, request_options=request_options
)

def _consume_stream(self, request_options: typing.Optional[RequestOptions]) -> typing.Iterator[TurnStreamData]:
"""Consume the create_turn SSE, adopting the inner Turn from the first turn.created."""
with self._client.agents.sessions.create_turn(
"""Consume create_turn_stream SSE, adopting the inner Turn from the first turn.created."""
with self._client.agents.sessions.create_turn_stream(
self._session_id,
input=self._input_param,
previous_turn_id=self._previous_turn_id,
Expand All @@ -334,7 +341,7 @@ def _consume_stream(self, request_options: typing.Optional[RequestOptions]) -> t
for event in sse.with_metadata():
sequence_number = parse_sequence_number(event.id)
if isinstance(event.data, TurnCreatedEvent) and self._turn is None:
self._adopt_turn(event.data)
self._adopt_turn_from_created_event(event.data)
elif self._turn is not None and isinstance(event.data, TurnDoneEvent):
self._replace_turn_state(event.data.state)
yield TurnStreamData(sequence_number=sequence_number, event=event.data)
Expand All @@ -344,14 +351,23 @@ def _must_get_turn(self) -> Turn:
raise RuntimeError("Turn not started yet; call execute() first.")
return self._turn

def _create_turn_if_not_exist(self, request_options: typing.Optional[RequestOptions]) -> None:
"""Drive the create_turn SSE only until the first turn.created builds the inner Turn, then stop."""
if self._turn is None:
for _ in self._consume_stream(request_options):
if self._turn is not None:
break
def _adopt_turn_from_api(self, turn: RawTurn) -> None:
"""Build the inner Turn from create_turn / get_turn response data."""
self._turn = Turn(
RawTurn(
id=turn.id,
session_id=turn.session_id,
previous_turn_id=turn.previous_turn_id,
input=turn.input if turn.input is not None else self._input, # type: ignore[arg-type]
state=turn.state,
created_by_subject=turn.created_by_subject,
created_at=turn.created_at,
),
self._session,
self._client,
Comment thread
bhaveshpatel640 marked this conversation as resolved.
)

def _adopt_turn(self, event: TurnCreatedEvent) -> None:
def _adopt_turn_from_created_event(self, event: TurnCreatedEvent) -> None:
"""Build the inner Turn directly from the turn.created event."""
self._turn = Turn(
RawTurn(
Expand Down Expand Up @@ -523,12 +539,13 @@ def execute(
request_options: typing.Optional[RequestOptions] = None,
) -> typing.Union[typing.AsyncIterator[TurnStreamData], "typing.Coroutine[typing.Any, typing.Any, TurnState]"]:
"""
Start the turn via ``create_turn``.
Start the turn via ``create_turn_stream`` / ``create_turn``.

Parameters
----------
stream : bool
Stream ``create_turn`` SSE when true. Default true.
Stream ``create_turn_stream`` SSE when true. Default true.
When false, uses ``create_turn`` JSON then polls to terminal.
poll_interval_ms : int
Poll interval ms when ``stream=False``. Minimum 3000.
request_options : typing.Optional[RequestOptions]
Expand All @@ -542,7 +559,7 @@ def execute(
Yields
------
TurnStreamData
SSE stream from create_turn.
SSE stream from create_turn_stream.
"""
if self._started:
raise RuntimeError("Turn already started; use stream() / wait_for_completion().")
Expand Down Expand Up @@ -680,15 +697,23 @@ async def _run_streaming(
async def _start_and_wait(
self, poll_interval_ms: int, request_options: typing.Optional[RequestOptions]
) -> TurnState:
await self._create_turn_if_not_exist(request_options)
"""execute(stream=False) path: create_turn JSON returns the running turn; then poll."""
response = await self._client.agents.sessions.create_turn(
self._session_id,
input=self._input_param,
previous_turn_id=self._previous_turn_id,
request_options=request_options,
)
self._adopt_turn_from_api(response.data)
return await self._must_get_turn().wait_for_completion(
poll_interval_ms=poll_interval_ms, request_options=request_options
)

async def _consume_stream(
self, request_options: typing.Optional[RequestOptions]
) -> typing.AsyncIterator[TurnStreamData]:
async with self._client.agents.sessions.create_turn(
"""Consume create_turn_stream SSE, adopting the inner Turn from the first turn.created."""
async with self._client.agents.sessions.create_turn_stream(
self._session_id,
input=self._input_param,
previous_turn_id=self._previous_turn_id,
Expand All @@ -697,7 +722,7 @@ async def _consume_stream(
async for event in sse.with_metadata():
sequence_number = parse_sequence_number(event.id)
if isinstance(event.data, TurnCreatedEvent) and self._turn is None:
self._adopt_turn(event.data)
self._adopt_turn_from_created_event(event.data)
elif self._turn is not None and isinstance(event.data, TurnDoneEvent):
self._replace_turn_state(event.data.state)
yield TurnStreamData(sequence_number=sequence_number, event=event.data)
Expand All @@ -707,13 +732,22 @@ def _must_get_turn(self) -> AsyncTurn:
raise RuntimeError("Turn not started yet; call execute() first.")
return self._turn

async def _create_turn_if_not_exist(self, request_options: typing.Optional[RequestOptions]) -> None:
if self._turn is None:
async for _ in self._consume_stream(request_options):
if self._turn is not None:
break
def _adopt_turn_from_api(self, turn: RawTurn) -> None:
self._turn = AsyncTurn(
RawTurn(
id=turn.id,
session_id=turn.session_id,
previous_turn_id=turn.previous_turn_id,
input=turn.input if turn.input is not None else self._input, # type: ignore[arg-type]
Comment thread
thesujai marked this conversation as resolved.
state=turn.state,
created_by_subject=turn.created_by_subject,
created_at=turn.created_at,
),
Comment thread
bhaveshpatel640 marked this conversation as resolved.
self._session,
self._client,
)

def _adopt_turn(self, event: TurnCreatedEvent) -> None:
def _adopt_turn_from_created_event(self, event: TurnCreatedEvent) -> None:
self._turn = AsyncTurn(
RawTurn(
id=event.turn_id,
Expand Down
Loading