datarobot_genai.core.chat.responses

responses

OpenAI-compatible response helpers for chat interactions.

CustomModelStreamingResponse

Bases: ChatCompletionChunk

Source code in datarobot_genai/core/chat/responses.py

class CustomModelStreamingResponse(ChatCompletionChunk):
    pipeline_interactions: str | None = None
    event: Event | None = None

    @field_serializer("event")
    def _serialize_event(self, value: Event | None) -> dict[str, Any] | None:
        """Serialize event with discriminator always included.

        When the parent is dumped with exclude_unset=True, nested Event models
        lose their default 'type' field, breaking tagged-union validation on
        round-trip. We always dump the full event (including type) here.
        """
        if value is None:
            return None
        return value.model_dump(mode="json")

to_custom_model_chat_response

to_custom_model_chat_response(response_text: str, pipeline_interactions: MultiTurnSample | None, usage_metrics: dict[str, int], model: str | object | None) -> CustomModelChatResponse

Convert the OpenAI ChatCompletion response to CustomModelChatResponse.

Source code in datarobot_genai/core/chat/responses.py

def to_custom_model_chat_response(
    response_text: str,
    pipeline_interactions: MultiTurnSample | None,
    usage_metrics: dict[str, int],
    model: str | object | None,
) -> CustomModelChatResponse:
    """Convert the OpenAI ChatCompletion response to CustomModelChatResponse."""
    choice = Choice(
        index=0,
        message=ChatCompletionMessage(role="assistant", content=response_text),
        finish_reason="stop",
    )

    if model is None:
        model = "unspecified-model"
    else:
        model = str(model)

    required_usage_metrics: dict[str, int] = {
        "completion_tokens": 0,
        "prompt_tokens": 0,
        "total_tokens": 0,
    }

    return CustomModelChatResponse(
        id=str(uuid.uuid4()),
        object="chat.completion",
        choices=[choice],
        created=int(time.time()),
        model=model,
        usage=CompletionUsage.model_validate(required_usage_metrics | usage_metrics),
        pipeline_interactions=pipeline_interactions.model_dump_json()
        if pipeline_interactions
        else None,
    )

to_custom_model_streaming_response

to_custom_model_streaming_response(thread_pool_executor: ThreadPoolExecutor, event_loop: AbstractEventLoop, streaming_response_generator: AsyncGenerator[tuple[Event, MultiTurnSample | None, dict[str, int]]], model: str | object | None) -> Iterator[CustomModelStreamingResponse]

Convert the OpenAI ChatCompletionChunk response to CustomModelStreamingResponse.

Source code in datarobot_genai/core/chat/responses.py

def to_custom_model_streaming_response(
    thread_pool_executor: ThreadPoolExecutor,
    event_loop: AbstractEventLoop,
    streaming_response_generator: AsyncGenerator[
        tuple[Event, MultiTurnSample | None, dict[str, int]]
    ],
    model: str | object | None,
) -> Iterator[CustomModelStreamingResponse]:
    """Convert the OpenAI ChatCompletionChunk response to CustomModelStreamingResponse."""
    completion_id = str(uuid.uuid4())
    created = int(time.time())

    last_pipeline_interactions = None
    last_usage_metrics = None

    if model is None:
        model = "unspecified-model"
    else:
        model = str(model)

    required_usage_metrics = default_usage_metrics()
    try:
        agent_response = aiter(streaming_response_generator)
        while True:
            try:
                (
                    event,
                    pipeline_interactions,
                    usage_metrics,
                ) = thread_pool_executor.submit(
                    event_loop.run_until_complete, anext(agent_response)
                ).result()
                last_pipeline_interactions = pipeline_interactions
                last_usage_metrics = usage_metrics

                # Skip only RunStarted: the OpenAI client stream does not replace the
                # API gateway's own RunStarted (e.g. Fastapi AG-UI layer may emit one first).
                # Do **not** skip RunFinished: HttpAgent / useAgUiChat rely on RUN_FINISHED
                # arriving with the model stream; otherwise the UI can stay in "running"
                # and never process follow-up turns.
                if isinstance(event, RunStartedEvent):
                    continue

                if isinstance(event, BaseEvent):
                    content = ""
                    if isinstance(event, (TextMessageContentEvent, TextMessageChunkEvent)):
                        content = event.delta or content
                    choice = ChunkChoice(
                        index=0,
                        delta=ChoiceDelta(role="assistant", content=content),
                        finish_reason=None,
                    )

                    yield CustomModelStreamingResponse(
                        id=completion_id,
                        object="chat.completion.chunk",
                        created=created,
                        model=model,
                        choices=[choice],
                        usage=CompletionUsage.model_validate(required_usage_metrics | usage_metrics)
                        if usage_metrics
                        else None,
                        event=event,
                    )
            except StopAsyncIteration:
                break
            except RuntimeError as e:
                if "Attempted to exit cancel scope" in str(e):
                    # Known issue: MCP/AnyIO cancel scope may be torn down in a different
                    # asyncio Task than the one that entered it when streaming is driven via
                    # repeated run_until_complete(anext(...)). Full async execution avoids this;
                    # log and continue until a sync-free execution path is available.
                    logger.warning(
                        "Ignoring MCP context teardown RuntimeError during streaming: %s",
                        e,
                    )
                    continue
                raise
        event_loop.run_until_complete(streaming_response_generator.aclose())
        # Yield final chunk indicating end of stream
        choice = ChunkChoice(
            index=0,
            delta=ChoiceDelta(role="assistant"),
            finish_reason="stop",
        )
        yield CustomModelStreamingResponse(
            id=completion_id,
            object="chat.completion.chunk",
            created=created,
            model=model,
            choices=[choice],
            usage=CompletionUsage.model_validate(required_usage_metrics | last_usage_metrics)
            if last_usage_metrics
            else None,
            pipeline_interactions=last_pipeline_interactions.model_dump_json()
            if last_pipeline_interactions
            else None,
        )
    except Exception as e:
        tb.print_exc()
        created = int(time.time())
        choice = ChunkChoice(
            index=0,
            delta=ChoiceDelta(role="assistant", content=str(e), refusal="error"),
            finish_reason="stop",
        )
        yield CustomModelStreamingResponse(
            id=completion_id,
            object="chat.completion.chunk",
            created=created,
            model=model,
            choices=[choice],
            usage=None,
        )

streaming_iterator_to_custom_model_streaming_response

streaming_iterator_to_custom_model_streaming_response(streaming_response_iterator: Iterator[tuple[Event, MultiTurnSample | None, dict[str, int]]], model: str | object | None) -> Iterator[CustomModelStreamingResponse]

Convert the OpenAI ChatCompletionChunk response to CustomModelStreamingResponse.

Source code in datarobot_genai/core/chat/responses.py

def streaming_iterator_to_custom_model_streaming_response(
    streaming_response_iterator: Iterator[tuple[Event, MultiTurnSample | None, dict[str, int]]],
    model: str | object | None,
) -> Iterator[CustomModelStreamingResponse]:
    """Convert the OpenAI ChatCompletionChunk response to CustomModelStreamingResponse."""
    completion_id = str(uuid.uuid4())
    created = int(time.time())

    last_pipeline_interactions = None
    last_usage_metrics = None

    if model is None:
        model = "unspecified-model"
    else:
        model = str(model)

    required_usage_metrics = default_usage_metrics()

    try:
        while True:
            try:
                (
                    event,
                    pipeline_interactions,
                    usage_metrics,
                ) = next(streaming_response_iterator)
                last_pipeline_interactions = pipeline_interactions
                last_usage_metrics = usage_metrics

                # Skip only RunStarted: the OpenAI client stream does not replace the
                # API gateway's own RunStarted (e.g. Fastapi AG-UI layer may emit one first).
                # Do **not** skip RunFinished: HttpAgent / useAgUiChat rely on RUN_FINISHED
                # arriving with the model stream; otherwise the UI can stay in "running"
                # and never process follow-up turns.
                if isinstance(event, RunStartedEvent):
                    continue

                if isinstance(event, BaseEvent):
                    content = ""
                    if isinstance(event, (TextMessageContentEvent, TextMessageChunkEvent)):
                        content = event.delta or content
                    choice = ChunkChoice(
                        index=0,
                        delta=ChoiceDelta(role="assistant", content=content),
                        finish_reason=None,
                    )
                    yield CustomModelStreamingResponse(
                        id=completion_id,
                        object="chat.completion.chunk",
                        created=created,
                        model=model,
                        choices=[choice],
                        usage=CompletionUsage.model_validate(required_usage_metrics | usage_metrics)
                        if usage_metrics
                        else None,
                        event=event,
                    )
            except StopIteration:
                break
        # Yield final chunk indicating end of stream
        choice = ChunkChoice(
            index=0,
            delta=ChoiceDelta(role="assistant"),
            finish_reason="stop",
        )
        yield CustomModelStreamingResponse(
            id=completion_id,
            object="chat.completion.chunk",
            created=created,
            model=model,
            choices=[choice],
            usage=CompletionUsage.model_validate(required_usage_metrics | last_usage_metrics)
            if last_usage_metrics
            else None,
            pipeline_interactions=last_pipeline_interactions.model_dump_json()
            if last_pipeline_interactions
            else None,
        )
    except Exception as e:
        tb.print_exc()
        created = int(time.time())
        choice = ChunkChoice(
            index=0,
            delta=ChoiceDelta(role="assistant", content=str(e), refusal="error"),
            finish_reason="stop",
        )
        yield CustomModelStreamingResponse(
            id=completion_id,
            object="chat.completion.chunk",
            created=created,
            model=model,
            choices=[choice],
            usage=None,
        )

async_gen_to_sync_thread

async_gen_to_sync_thread(async_iterator: AsyncIterator[T], thread_pool_executor: ThreadPoolExecutor, event_loop: AbstractEventLoop) -> Iterator[T]

Run an async iterator in a separate thread and provide a sync iterator.

Source code in datarobot_genai/core/chat/responses.py

def async_gen_to_sync_thread(
    async_iterator: AsyncIterator[T],
    thread_pool_executor: ThreadPoolExecutor,
    event_loop: asyncio.AbstractEventLoop,
) -> Iterator[T]:
    """Run an async iterator in a separate thread and provide a sync iterator."""
    # A thread-safe queue for communication
    sync_queue: queue.Queue[Any] = queue.Queue()
    # A sentinel object to signal the end of the async generator
    SENTINEL = object()  # noqa: N806

    async def run_async_to_queue() -> None:
        """Run in the separate thread's event loop."""
        try:
            async for item in async_iterator:
                sync_queue.put(item)
        except Exception as e:
            # Put the exception on the queue to be re-raised in the main thread
            sync_queue.put(e)
        finally:
            # Signal the end of iteration
            sync_queue.put(SENTINEL)

    thread_pool_executor.submit(event_loop.run_until_complete, run_async_to_queue()).result()

    # The main thread consumes items synchronously
    while True:
        item = sync_queue.get()
        if item is SENTINEL:
            break
        if isinstance(item, Exception):
            raise item
        yield item