datarobot_genai.core.agents.base

base

BaseAgent

Bases: Generic[TTool], ABC

BaseAgent centralizes common initialization for agent templates.

Fields

• api_key: DataRobot API token
• api_base: Endpoint for DataRobot, normalized for LLM Gateway usage
• llm: Framework-specific LLM client (constructed outside the agent and passed in)
• model: Optional model identifier string (e.g. LiteLLM / deployment model name)
• timeout: Request timeout
• verbose: Verbosity flag
• forwarded_headers: Forwarded headers for the agent
• max_history_messages: Maximum number of prior messages to include in chat history

Source code in datarobot_genai/core/agents/base.py

class BaseAgent(Generic[TTool], abc.ABC):
    """BaseAgent centralizes common initialization for agent templates.

    Fields:
      - api_key: DataRobot API token
      - api_base: Endpoint for DataRobot, normalized for LLM Gateway usage
      - llm: Framework-specific LLM client (constructed outside the agent and passed in)
      - model: Optional model identifier string (e.g. LiteLLM / deployment model name)
      - timeout: Request timeout
      - verbose: Verbosity flag
      - forwarded_headers: Forwarded headers for the agent
      - max_history_messages: Maximum number of prior messages to include in chat history
    """

    def __init__(
        self,
        api_key: str | None = None,
        api_base: str | None = None,
        llm: Any | None = None,
        tools: list[TTool] | None = None,
        verbose: bool = True,
        timeout: int = 90,
        forwarded_headers: dict[str, str] | None = None,
        max_history_messages: int | None = None,
        model: str | None = None,
    ) -> None:
        self.api_key = api_key or os.environ.get("DATAROBOT_API_TOKEN")
        self.api_base = (
            api_base or os.environ.get("DATAROBOT_ENDPOINT") or "https://app.datarobot.com"
        )
        self.set_llm(llm)
        self.set_tools(tools or [])
        self.set_timeout(timeout)
        self.set_verbose(verbose)
        self.set_model(model)
        self._forwarded_headers: dict[str, str] = forwarded_headers or {}
        self._identity_header: dict[str, str] = prepare_identity_header(self._forwarded_headers)
        self._max_history_messages = max_history_messages

    def set_llm(self, llm: Any | None) -> None:
        self._llm = llm

    @property
    def llm(self) -> Any | None:
        return self._llm

    def set_model(self, model: str | None) -> None:
        self._model = model

    @property
    def model(self) -> str | None:
        return self._model

    def set_timeout(self, timeout: int) -> None:
        self._timeout = timeout

    @property
    def timeout(self) -> int:
        return self._timeout

    def set_verbose(self, verbose: bool) -> None:
        self._verbose = verbose

    @property
    def verbose(self) -> bool:
        return self._verbose

    def set_tools(self, tools: list[TTool]) -> None:
        self._tools = tools

    @property
    def tools(self) -> list[TTool]:
        """Return the list of tools available to this agent.

        Subclasses can use this to wire tools into the agent.
        """
        return self._tools

    @property
    def max_history_messages(self) -> int:
        """Maximum number of prior messages to include in chat history.

        Defaults to ``DATAROBOT_GENAI_MAX_HISTORY_MESSAGES`` env var (read at
        call time). Subclasses can override via the constructor parameter or
        by overriding this property.
        """
        if self._max_history_messages is not None:
            return self._max_history_messages
        return get_max_history_messages_default()

    @property
    def forwarded_headers(self) -> dict[str, str]:
        """Return the forwarded headers for this agent."""
        return self._forwarded_headers

    def litellm_api_base(self, deployment_id: str | None) -> str:
        return get_api_base(self.api_base, deployment_id)

    @abc.abstractmethod
    def invoke(self, run_agent_input: RunAgentInput) -> InvokeReturn:
        raise NotImplementedError("Not implemented")

    async def invoke_single_message(self, user_message: str) -> InvokeReturn:
        """
        Invoke the agent without chat history with a single user message.

        Parameters
        ----------
        user_message: str
            The user message to invoke the agent with.

        Returns
        -------
        InvokeReturn
            Same async stream as :meth:`invoke` for a fresh run with a single user turn.
        """
        # Generate a new thread and run ID
        run_input = RunAgentInput(
            thread_id=str(uuid.uuid4()),
            run_id=str(uuid.uuid4()),
            messages=[UserMessage(id=str(uuid.uuid4()), role="user", content=user_message)],
            state=[],
            tools=[],
            context=[],
            forwardedProps=None,
        )
        async for output in self.invoke(run_input):
            yield output

    def build_history_summary(
        self,
        run_agent_input: RunAgentInput,
    ) -> str:
        """Instance helper to summarize prior turns as plain-text transcript.

        Subclasses can override ``max_history_messages`` to control how many
        prior messages are included. This is primarily intended for exposing a
        ``chat_history`` variable in prompts across different agent types.
        """
        return build_history_summary_from_messages(run_agent_input, self.max_history_messages)

    def history_messages(
        self,
        run_agent_input: RunAgentInput,
    ) -> list[NormalizedHistoryMessage]:
        """Prior turns as structured dicts (role/content + tool_calls/tool_call_id).

        Same history selection and ``max_history_messages`` truncation as
        :meth:`build_history_summary`, but preserves structured tool-call metadata so
        framework adapters can reconstruct native message history instead of
        flattening to text. See ``datarobot_genai.langgraph.history`` and
        ``datarobot_genai.llama_index.history``.

        Tool-call/result turns left unpaired by ``max_history`` truncation are dropped
        so the reconstructed message sequence stays valid for the model.
        """
        history = extract_history_messages(run_agent_input, self.max_history_messages)
        return drop_unpaired_boundary_tool_turns(history)

    @classmethod
    def create_pipeline_interactions_from_events(
        cls,
        events: list[Any] | None,
    ) -> MultiTurnSample | None:
        """Create a simple MultiTurnSample from a list of generic events/messages."""
        if not events:
            return None
        # Lazy import to reduce memory overhead when ragas is not used
        from ragas import MultiTurnSample

        return MultiTurnSample(user_input=events)

tools property

tools: list[TTool]

Return the list of tools available to this agent.

Subclasses can use this to wire tools into the agent.

max_history_messages property

max_history_messages: int

Maximum number of prior messages to include in chat history.

Defaults to DATAROBOT_GENAI_MAX_HISTORY_MESSAGES env var (read at call time). Subclasses can override via the constructor parameter or by overriding this property.

forwarded_headers property

forwarded_headers: dict[str, str]

Return the forwarded headers for this agent.

invoke_single_message async

invoke_single_message(user_message: str) -> InvokeReturn

Invoke the agent without chat history with a single user message.

Parameters

user_message: str The user message to invoke the agent with.

Returns

InvokeReturn Same async stream as :meth:invoke for a fresh run with a single user turn.

Source code in datarobot_genai/core/agents/base.py

async def invoke_single_message(self, user_message: str) -> InvokeReturn:
    """
    Invoke the agent without chat history with a single user message.

    Parameters
    ----------
    user_message: str
        The user message to invoke the agent with.

    Returns
    -------
    InvokeReturn
        Same async stream as :meth:`invoke` for a fresh run with a single user turn.
    """
    # Generate a new thread and run ID
    run_input = RunAgentInput(
        thread_id=str(uuid.uuid4()),
        run_id=str(uuid.uuid4()),
        messages=[UserMessage(id=str(uuid.uuid4()), role="user", content=user_message)],
        state=[],
        tools=[],
        context=[],
        forwardedProps=None,
    )
    async for output in self.invoke(run_input):
        yield output

build_history_summary

build_history_summary(run_agent_input: RunAgentInput) -> str

Instance helper to summarize prior turns as plain-text transcript.

Subclasses can override max_history_messages to control how many prior messages are included. This is primarily intended for exposing a chat_history variable in prompts across different agent types.

Source code in datarobot_genai/core/agents/base.py

def build_history_summary(
    self,
    run_agent_input: RunAgentInput,
) -> str:
    """Instance helper to summarize prior turns as plain-text transcript.

    Subclasses can override ``max_history_messages`` to control how many
    prior messages are included. This is primarily intended for exposing a
    ``chat_history`` variable in prompts across different agent types.
    """
    return build_history_summary_from_messages(run_agent_input, self.max_history_messages)

history_messages

history_messages(run_agent_input: RunAgentInput) -> list[NormalizedHistoryMessage]

Prior turns as structured dicts (role/content + tool_calls/tool_call_id).

Same history selection and max_history_messages truncation as :meth:build_history_summary, but preserves structured tool-call metadata so framework adapters can reconstruct native message history instead of flattening to text. See datarobot_genai.langgraph.history and datarobot_genai.llama_index.history.

Tool-call/result turns left unpaired by max_history truncation are dropped so the reconstructed message sequence stays valid for the model.

Source code in datarobot_genai/core/agents/base.py

def history_messages(
    self,
    run_agent_input: RunAgentInput,
) -> list[NormalizedHistoryMessage]:
    """Prior turns as structured dicts (role/content + tool_calls/tool_call_id).

    Same history selection and ``max_history_messages`` truncation as
    :meth:`build_history_summary`, but preserves structured tool-call metadata so
    framework adapters can reconstruct native message history instead of
    flattening to text. See ``datarobot_genai.langgraph.history`` and
    ``datarobot_genai.llama_index.history``.

    Tool-call/result turns left unpaired by ``max_history`` truncation are dropped
    so the reconstructed message sequence stays valid for the model.
    """
    history = extract_history_messages(run_agent_input, self.max_history_messages)
    return drop_unpaired_boundary_tool_turns(history)

create_pipeline_interactions_from_events classmethod

create_pipeline_interactions_from_events(events: list[Any] | None) -> MultiTurnSample | None

Create a simple MultiTurnSample from a list of generic events/messages.

Source code in datarobot_genai/core/agents/base.py

@classmethod
def create_pipeline_interactions_from_events(
    cls,
    events: list[Any] | None,
) -> MultiTurnSample | None:
    """Create a simple MultiTurnSample from a list of generic events/messages."""
    if not events:
        return None
    # Lazy import to reduce memory overhead when ragas is not used
    from ragas import MultiTurnSample

    return MultiTurnSample(user_input=events)

extract_streaming_memory_context

extract_streaming_memory_context(run_agent_input: RunAgentInput) -> str | None

Return streaming_memory_agent context to prepend to a user prompt, or None.

Source code in datarobot_genai/core/agents/base.py

def extract_streaming_memory_context(run_agent_input: RunAgentInput) -> str | None:
    """Return ``streaming_memory_agent`` context to prepend to a user prompt, or None."""
    messages = list(run_agent_input.messages)
    last_user_idx = _last_user_message_index(messages)
    if last_user_idx is None or last_user_idx == 0:
        return None

    memory_message = messages[last_user_idx - 1]
    if getattr(memory_message, "role", None) != "system":
        return None

    memory_text = _message_content_as_str(getattr(memory_message, "content", None))
    if not memory_text.startswith(STREAMING_MEMORY_CONTEXT_PREFIX):
        return None

    return memory_text

prepend_streaming_memory_to_prompt

prepend_streaming_memory_to_prompt(prompt: str, run_agent_input: RunAgentInput) -> str

Prepend streaming_memory_agent context to a processed user prompt.

Source code in datarobot_genai/core/agents/base.py

def prepend_streaming_memory_to_prompt(prompt: str, run_agent_input: RunAgentInput) -> str:
    """Prepend ``streaming_memory_agent`` context to a processed user prompt."""
    streaming_context = extract_streaming_memory_context(run_agent_input)
    if not streaming_context:
        return prompt
    return streaming_context + "\n\n" + prompt

extract_user_prompt_content

extract_user_prompt_content(run_agent_input: RunAgentInput) -> Any

Extract the last user message content from input.

Source code in datarobot_genai/core/agents/base.py

def extract_user_prompt_content(run_agent_input: RunAgentInput) -> Any:
    """Extract the last user message content from input."""
    user_messages = [msg for msg in run_agent_input.messages if msg.role == "user"]
    # Get the last user message
    content: str = user_messages[-1].content if user_messages else ""
    # Try converting prompt from json to a dict
    if isinstance(content, str):
        try:
            content = json.loads(content)
        except json.JSONDecodeError:
            pass

    return content

make_system_prompt

make_system_prompt(suffix: str = '', *, prefix: str | None = None) -> str

Build a system prompt with optional prefix and suffix.

Parameters

suffix : str, default "" Text appended after the prefix. If non-empty, it is placed on a new line. prefix : str | None, keyword-only, default None Custom prefix text. When None, a default collaborative assistant instruction is used.

Returns

str The composed system prompt string.

Source code in datarobot_genai/core/agents/base.py

def make_system_prompt(suffix: str = "", *, prefix: str | None = None) -> str:
    """Build a system prompt with optional prefix and suffix.

    Parameters
    ----------
    suffix : str, default ""
        Text appended after the prefix. If non-empty, it is placed on a new line.
    prefix : str | None, keyword-only, default None
        Custom prefix text. When ``None``, a default collaborative assistant
        instruction is used.

    Returns
    -------
    str
        The composed system prompt string.
    """
    default_prefix = (
        "You are a helpful AI assistant, collaborating with other assistants."
        " Use the provided tools to progress towards answering the question."
        " If you are unable to fully answer, that's OK, another assistant with different tools "
        " will help where you left off. Execute what you can to make progress."
    )
    head = prefix if prefix is not None else default_prefix
    if suffix:
        return head + "\n" + suffix
    return head

default_usage_metrics

default_usage_metrics() -> UsageMetrics

Return a metrics dict with required keys for OpenAI-compatible responses.

Source code in datarobot_genai/core/agents/base.py

def default_usage_metrics() -> UsageMetrics:
    """Return a metrics dict with required keys for OpenAI-compatible responses."""
    return {
        "completion_tokens": 0,
        "prompt_tokens": 0,
        "total_tokens": 0,
    }