datarobot_genai.core.chat.completions

completions

is_streaming

is_streaming(completion_create_params: CompletionCreateParams | Mapping[str, Any]) -> bool

Return True when the request asks for streaming, False otherwise.

Accepts both pydantic types and plain dictionaries.

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

def is_streaming(completion_create_params: CompletionCreateParams | Mapping[str, Any]) -> bool:
    """Return True when the request asks for streaming, False otherwise.

    Accepts both pydantic types and plain dictionaries.
    """
    params = cast(Mapping[str, Any], completion_create_params)
    value = params.get("stream", False)
    # Handle non-bool truthy values defensively (e.g., "true")
    if isinstance(value, str):
        return value.lower() == "true"
    return bool(value)

backfill_model

backfill_model(current: str | None, fallback: str | None) -> str | None

Replace NAT's "unknown-model" placeholder (or None) with fallback.

NAT defaults ChatResponse.model / ChatResponseChunk.model to the literal "unknown-model" whenever the workflow output didn't carry one. Callers pass the agent's configured model (:func:default_response_model) as fallback so the response reports the model the agent actually ran. A real model the workflow produced — or a deliberately-set one such as moderation's MODERATION_MODEL_NAME — is preserved.

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

def backfill_model(current: str | None, fallback: str | None) -> str | None:
    """Replace NAT's ``"unknown-model"`` placeholder (or ``None``) with ``fallback``.

    NAT defaults ``ChatResponse.model`` / ``ChatResponseChunk.model`` to the literal
    ``"unknown-model"`` whenever the workflow output didn't carry one. Callers pass the
    agent's configured model (:func:`default_response_model`) as ``fallback`` so the
    response reports the model the agent actually ran. A real model the workflow
    produced — or a deliberately-set one such as moderation's ``MODERATION_MODEL_NAME``
    — is preserved.
    """
    if fallback and current in (None, "unknown-model"):
        return fallback
    return current

convert_chat_completion_params_to_run_agent_input

convert_chat_completion_params_to_run_agent_input(chat_completion_params: CompletionCreateParams | Mapping[str, Any]) -> RunAgentInput

Convert a chat completion parameters to a run agent input.

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

def convert_chat_completion_params_to_run_agent_input(
    chat_completion_params: CompletionCreateParams | Mapping[str, Any],
) -> RunAgentInput:
    """Convert a chat completion parameters to a run agent input."""
    tools = [
        Tool(
            name=tool.get("function").get("name"),
            description=tool.get("function").get("description"),
            parameters=tool.get("function").get("parameters"),
        )
        for tool in chat_completion_params.get("tools", []) or []
        if tool.get("type") == "function"  # type: ignore[union-attr]
    ]
    messages: list[Message] = []
    for i, message in enumerate(chat_completion_params.get("messages", [])):  # type: ignore[arg-type]
        id = f"message_{i}"
        if message.get("role") == "user":
            messages.append(UserMessage(id=id, content=message.get("content")))
        elif message.get("role") == "assistant":
            tool_calls = []
            for tool_call in message.get("tool_calls", []) or []:
                function = tool_call.get("function") or {}
                tool_calls.append(
                    ToolCall(
                        id=tool_call.get("id"),
                        type=tool_call.get("type", "function"),
                        function=FunctionCall(
                            name=function.get("name"),
                            arguments=function.get("arguments", "{}"),
                        ),
                    )
                )
            messages.append(
                AssistantMessage(
                    id=id,
                    content=message.get("content"),
                    tool_calls=tool_calls or None,
                )
            )
        elif message.get("role") == "tool":
            messages.append(
                ToolMessage(
                    id=id,
                    content=message.get("content"),
                    tool_call_id=message.get("tool_call_id"),
                    error=message.get("error"),
                )
            )
        elif message.get("role") == "system":
            messages.append(SystemMessage(id=id, content=message.get("content")))

    forwarded_props: dict[str, Any] = {
        "model": chat_completion_params.get("model"),
        "authorization_context": chat_completion_params.get("authorization_context"),
        "forwarded_headers": chat_completion_params.get("forwarded_headers"),
    }

    thread_id, run_id = _resolve_thread_and_run_ids(chat_completion_params)

    return RunAgentInput(
        messages=messages,
        tools=tools,
        forwarded_props=forwarded_props,
        thread_id=thread_id,
        run_id=run_id,
        state={},
        context=[],
    )

agent_chat_completion_wrapper async

agent_chat_completion_wrapper(agent: BaseAgent, chat_completion_params: CompletionCreateParams | Mapping[str, Any], mcp_tools_factory: Callable[[], Any]) -> InvokeReturn | tuple[str, MultiTurnSample | None, UsageMetrics]

Wrap the agent's invoke method in a chat completion wrapper.

MCP tools from mcp_tools_factory are combined with any tools already on agent (MCP first, then existing agent.tools).

Returns

InvokeReturn When streaming is requested - the raw async event generator tuple[str, MultiTurnSample | None, UsageMetrics] When non-streaming - the reassembled final text, pipeline interactions, and accumulated usage metrics.

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

async def agent_chat_completion_wrapper(
    agent: BaseAgent,
    chat_completion_params: CompletionCreateParams | Mapping[str, Any],
    mcp_tools_factory: Callable[[], Any],
) -> InvokeReturn | tuple[str, MultiTurnSample | None, UsageMetrics]:
    """Wrap the agent's invoke method in a chat completion wrapper.

    MCP tools from ``mcp_tools_factory`` are combined with any tools already on
    ``agent`` (MCP first, then existing ``agent.tools``).

    Returns
    -------
    InvokeReturn
        When streaming is requested - the raw async event generator
    tuple[str, MultiTurnSample | None, UsageMetrics]
        When non-streaming - the reassembled final text, pipeline
        interactions, and accumulated usage metrics.
    """
    run_agent_input = convert_chat_completion_params_to_run_agent_input(chat_completion_params)

    if is_streaming(chat_completion_params):

        async def _stream_with_mcp() -> InvokeReturn:
            async with mcp_tools_factory() as mcp_tools:
                agent.set_tools(_merge_mcp_tools_with_agent_tools(mcp_tools, agent))
                async for item in agent.invoke(run_agent_input):
                    yield item

        return _stream_with_mcp()
    else:
        async with mcp_tools_factory() as mcp_tools:
            agent.set_tools(_merge_mcp_tools_with_agent_tools(mcp_tools, agent))
            final_response = ""
            pipeline_interactions = None
            usage_metrics = default_usage_metrics()
            received_run_finished = False
            async for event, iter_interactions, iter_metrics in agent.invoke(run_agent_input):
                # When we work in non-streaming mode, we only send back the final message
                # It is because of limitation of completions interface we can not send back the
                # intermediate messages
                if isinstance(event, TextMessageStartEvent):
                    final_response = ""
                elif isinstance(event, (TextMessageContentEvent, TextMessageChunkEvent)):
                    final_response += event.delta
                elif isinstance(event, RunFinishedEvent):
                    received_run_finished = True
                    pipeline_interactions = iter_interactions
                    usage_metrics = iter_metrics

            if not received_run_finished:
                logger.warning("Agent stream ended without RunFinishedEvent")

            return final_response, pipeline_interactions, usage_metrics