datarobot_genai.nat.agent

agent

NatAgent

Bases: BaseAgent[None]

Source code in datarobot_genai/nat/agent.py

class NatAgent(BaseAgent[None]):
    def __init__(
        self,
        workflow_path: StrPath,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.workflow_path = workflow_path

    def make_user_prompt(self, run_agent_input: RunAgentInput) -> str:
        """Create the user prompt text. Override to customize formatting.

        Chat history is automatically appended by `invoke` when
        max_history_messages > 0 (controlled via DATAROBOT_GENAI_MAX_HISTORY_MESSAGES env var).

        Default implementation returns the raw user message content.
        """
        user_prompt_content = extract_user_prompt_content(run_agent_input)
        return str(user_prompt_content)

    def make_chat_request(self, user_prompt: str) -> ChatRequest:
        """Create a NAT ChatRequest from the processed user prompt."""
        return ChatRequest.from_string(user_prompt)

    async def invoke(self, run_agent_input: RunAgentInput) -> InvokeReturn:
        """Run the agent with the provided input.

        Args:
            run_agent_input: The agent run input including messages, tools, and context.

        Returns
        -------
            Returns a generator yielding tuples of (event, pipeline_interactions, usage_metrics).

        """
        # Build the user prompt from the template
        user_prompt = self.make_user_prompt(run_agent_input)

        # Automatically inject chat history when enabled (max_history_messages > 0)
        history_summary = self.build_history_summary(run_agent_input)
        if history_summary:
            user_prompt = f"{user_prompt}\n\nPrior conversation:\n{history_summary}"

        # Create the chat request from the processed prompt
        chat_request = self.make_chat_request(user_prompt)

        # Print commands may need flush=True to ensure they are displayed in real-time.
        logger.info(
            f"Running agent with user prompt: {chat_request.messages[0].content}",
        )

        thread_id = run_agent_input.thread_id
        run_id = run_agent_input.run_id
        zero_metrics: UsageMetrics = {
            "completion_tokens": 0,
            "prompt_tokens": 0,
            "total_tokens": 0,
        }

        # Partial AG-UI: workflow lifecycle + text message events
        yield (
            RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=run_id),
            None,
            zero_metrics,
        )

        message_id = str(uuid.uuid4())
        text_started = False

        async with load_workflow(
            self.workflow_path,
            headers=self.forwarded_headers,
            disable_datarobot_moderation=True,
        ) as workflow:
            async with workflow.session(user_id=thread_id) as session:
                async with session.run(chat_request) as runner:
                    intermediate_future = pull_intermediate_structured()
                    async for result in runner.result_stream():
                        result_text = _extract_text(result)
                        if result_text:
                            if not text_started:
                                yield (
                                    TextMessageStartEvent(
                                        type=EventType.TEXT_MESSAGE_START,
                                        message_id=message_id,
                                    ),
                                    None,
                                    zero_metrics,
                                )
                                text_started = True
                            yield (
                                TextMessageContentEvent(
                                    type=EventType.TEXT_MESSAGE_CONTENT,
                                    message_id=message_id,
                                    delta=result_text,
                                ),
                                None,
                                zero_metrics,
                            )

                    if text_started:
                        yield (
                            TextMessageEndEvent(
                                type=EventType.TEXT_MESSAGE_END, message_id=message_id
                            ),
                            None,
                            zero_metrics,
                        )

                    steps = await intermediate_future
                    llm_end_steps = [
                        step for step in steps if step.event_type == IntermediateStepType.LLM_END
                    ]
                    usage_metrics: UsageMetrics = {
                        "completion_tokens": 0,
                        "prompt_tokens": 0,
                        "total_tokens": 0,
                    }
                    for step in llm_end_steps:
                        if step.usage_info:
                            token_usage = step.usage_info.token_usage
                            usage_metrics["total_tokens"] += token_usage.total_tokens
                            usage_metrics["prompt_tokens"] += token_usage.prompt_tokens
                            usage_metrics["completion_tokens"] += token_usage.completion_tokens

                    pipeline_interactions = self.create_pipeline_interactions_from_steps(steps)
                    yield (
                        RunFinishedEvent(
                            type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=run_id
                        ),
                        pipeline_interactions,
                        usage_metrics,
                    )

    async def run_nat_workflow(
        self, workflow_path: StrPath, chat_request: ChatRequest, headers: dict[str, str] | None
    ) -> tuple[ChatResponse | str, list[IntermediateStep]]:
        """Run the NAT workflow with the provided config file and input string.

        Args:
            workflow_path: Path to the NAT workflow configuration file
            chat_request: The chat request to process through the workflow
            headers: Optional HTTP headers to forward to the workflow

        Returns
        -------
            ChatResponse | str: The result from the NAT workflow
            list[IntermediateStep]: The list of intermediate steps
        """
        async with load_workflow(
            workflow_path,
            headers=headers,
            disable_datarobot_moderation=True,
        ) as workflow:
            async with workflow.session(user_id=str(uuid.uuid4())) as session:
                async with session.run(chat_request) as runner:
                    intermediate_future = pull_intermediate_structured()
                    runner_outputs = await runner.result()
                    steps = await intermediate_future

        line = f"{'-' * 50}"
        prefix = f"{line}\nWorkflow Result:\n"
        suffix = f"\n{line}"

        logger.info(f"{prefix}{runner_outputs}{suffix}")

        return runner_outputs, steps

    @classmethod
    def create_pipeline_interactions_from_steps(
        cls,
        steps: list[IntermediateStep],
    ) -> MultiTurnSample | None:
        if not steps:
            return None
        # Lazy import to reduce memory overhead when ragas is not used
        from ragas import MultiTurnSample

        ragas_trace = convert_to_ragas_messages(steps)
        return MultiTurnSample(user_input=ragas_trace)

make_user_prompt

make_user_prompt(run_agent_input: RunAgentInput) -> str

Create the user prompt text. Override to customize formatting.

Chat history is automatically appended by invoke when max_history_messages > 0 (controlled via DATAROBOT_GENAI_MAX_HISTORY_MESSAGES env var).

Default implementation returns the raw user message content.

Source code in datarobot_genai/nat/agent.py

def make_user_prompt(self, run_agent_input: RunAgentInput) -> str:
    """Create the user prompt text. Override to customize formatting.

    Chat history is automatically appended by `invoke` when
    max_history_messages > 0 (controlled via DATAROBOT_GENAI_MAX_HISTORY_MESSAGES env var).

    Default implementation returns the raw user message content.
    """
    user_prompt_content = extract_user_prompt_content(run_agent_input)
    return str(user_prompt_content)

make_chat_request

make_chat_request(user_prompt: str) -> ChatRequest

Create a NAT ChatRequest from the processed user prompt.

Source code in datarobot_genai/nat/agent.py

def make_chat_request(self, user_prompt: str) -> ChatRequest:
    """Create a NAT ChatRequest from the processed user prompt."""
    return ChatRequest.from_string(user_prompt)

invoke async

invoke(run_agent_input: RunAgentInput) -> InvokeReturn

Run the agent with the provided input.

Parameters:

Name	Type	Description	Default
run_agent_input	RunAgentInput	The agent run input including messages, tools, and context.	required

Returns

Returns a generator yielding tuples of (event, pipeline_interactions, usage_metrics).

Source code in datarobot_genai/nat/agent.py

async def invoke(self, run_agent_input: RunAgentInput) -> InvokeReturn:
    """Run the agent with the provided input.

    Args:
        run_agent_input: The agent run input including messages, tools, and context.

    Returns
    -------
        Returns a generator yielding tuples of (event, pipeline_interactions, usage_metrics).

    """
    # Build the user prompt from the template
    user_prompt = self.make_user_prompt(run_agent_input)

    # Automatically inject chat history when enabled (max_history_messages > 0)
    history_summary = self.build_history_summary(run_agent_input)
    if history_summary:
        user_prompt = f"{user_prompt}\n\nPrior conversation:\n{history_summary}"

    # Create the chat request from the processed prompt
    chat_request = self.make_chat_request(user_prompt)

    # Print commands may need flush=True to ensure they are displayed in real-time.
    logger.info(
        f"Running agent with user prompt: {chat_request.messages[0].content}",
    )

    thread_id = run_agent_input.thread_id
    run_id = run_agent_input.run_id
    zero_metrics: UsageMetrics = {
        "completion_tokens": 0,
        "prompt_tokens": 0,
        "total_tokens": 0,
    }

    # Partial AG-UI: workflow lifecycle + text message events
    yield (
        RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=run_id),
        None,
        zero_metrics,
    )

    message_id = str(uuid.uuid4())
    text_started = False

    async with load_workflow(
        self.workflow_path,
        headers=self.forwarded_headers,
        disable_datarobot_moderation=True,
    ) as workflow:
        async with workflow.session(user_id=thread_id) as session:
            async with session.run(chat_request) as runner:
                intermediate_future = pull_intermediate_structured()
                async for result in runner.result_stream():
                    result_text = _extract_text(result)
                    if result_text:
                        if not text_started:
                            yield (
                                TextMessageStartEvent(
                                    type=EventType.TEXT_MESSAGE_START,
                                    message_id=message_id,
                                ),
                                None,
                                zero_metrics,
                            )
                            text_started = True
                        yield (
                            TextMessageContentEvent(
                                type=EventType.TEXT_MESSAGE_CONTENT,
                                message_id=message_id,
                                delta=result_text,
                            ),
                            None,
                            zero_metrics,
                        )

                if text_started:
                    yield (
                        TextMessageEndEvent(
                            type=EventType.TEXT_MESSAGE_END, message_id=message_id
                        ),
                        None,
                        zero_metrics,
                    )

                steps = await intermediate_future
                llm_end_steps = [
                    step for step in steps if step.event_type == IntermediateStepType.LLM_END
                ]
                usage_metrics: UsageMetrics = {
                    "completion_tokens": 0,
                    "prompt_tokens": 0,
                    "total_tokens": 0,
                }
                for step in llm_end_steps:
                    if step.usage_info:
                        token_usage = step.usage_info.token_usage
                        usage_metrics["total_tokens"] += token_usage.total_tokens
                        usage_metrics["prompt_tokens"] += token_usage.prompt_tokens
                        usage_metrics["completion_tokens"] += token_usage.completion_tokens

                pipeline_interactions = self.create_pipeline_interactions_from_steps(steps)
                yield (
                    RunFinishedEvent(
                        type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=run_id
                    ),
                    pipeline_interactions,
                    usage_metrics,
                )

run_nat_workflow async

run_nat_workflow(workflow_path: StrPath, chat_request: ChatRequest, headers: dict[str, str] | None) -> tuple[ChatResponse | str, list[IntermediateStep]]

Run the NAT workflow with the provided config file and input string.

Parameters:

Name	Type	Description	Default
workflow_path	StrPath	Path to the NAT workflow configuration file	required
chat_request	ChatRequest	The chat request to process through the workflow	required
headers	dict[str, str] | None	Optional HTTP headers to forward to the workflow	required

Returns

ChatResponse | str: The result from the NAT workflow
list[IntermediateStep]: The list of intermediate steps

Source code in datarobot_genai/nat/agent.py

async def run_nat_workflow(
    self, workflow_path: StrPath, chat_request: ChatRequest, headers: dict[str, str] | None
) -> tuple[ChatResponse | str, list[IntermediateStep]]:
    """Run the NAT workflow with the provided config file and input string.

    Args:
        workflow_path: Path to the NAT workflow configuration file
        chat_request: The chat request to process through the workflow
        headers: Optional HTTP headers to forward to the workflow

    Returns
    -------
        ChatResponse | str: The result from the NAT workflow
        list[IntermediateStep]: The list of intermediate steps
    """
    async with load_workflow(
        workflow_path,
        headers=headers,
        disable_datarobot_moderation=True,
    ) as workflow:
        async with workflow.session(user_id=str(uuid.uuid4())) as session:
            async with session.run(chat_request) as runner:
                intermediate_future = pull_intermediate_structured()
                runner_outputs = await runner.result()
                steps = await intermediate_future

    line = f"{'-' * 50}"
    prefix = f"{line}\nWorkflow Result:\n"
    suffix = f"\n{line}"

    logger.info(f"{prefix}{runner_outputs}{suffix}")

    return runner_outputs, steps

pull_intermediate_structured

pull_intermediate_structured() -> asyncio.Future[list[IntermediateStep]]

Subscribe to the runner's event stream using callbacks. Intermediate steps are collected and, when complete, the future is set with the list of dumped intermediate steps.

Source code in datarobot_genai/nat/agent.py

def pull_intermediate_structured() -> asyncio.Future[list[IntermediateStep]]:
    """
    Subscribe to the runner's event stream using callbacks.
    Intermediate steps are collected and, when complete, the future is set
    with the list of dumped intermediate steps.
    """
    future: asyncio.Future[list[IntermediateStep]] = asyncio.Future()
    intermediate_steps = []  # We'll store the dumped steps here.
    context = Context.get()

    def on_next_cb(item: IntermediateStep) -> None:
        # Append each new intermediate step to the list.
        intermediate_steps.append(item)

    def on_error_cb(exc: Exception) -> None:
        logger.error("Hit on_error: %s", exc)
        if not future.done():
            future.set_exception(exc)

    def on_complete_cb() -> None:
        logger.debug("Completed reading intermediate steps")
        if not future.done():
            future.set_result(intermediate_steps)

    # Subscribe with our callbacks.
    context.intermediate_step_manager.subscribe(
        on_next=on_next_cb, on_error=on_error_cb, on_complete=on_complete_cb
    )

    return future