datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/llama_index/base.py

@@ -0,0 +1,299 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Base class for LlamaIndex-based agents.
+
+Provides a standard ``invoke`` that runs an AgentWorkflow, collects events,
+and converts them into pipeline interactions. Subclasses provide the workflow
+and response extraction logic.
+"""
+
+from __future__ import annotations
+
+import abc
+import inspect
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from llama_index.core.tools import BaseTool
+from openai.types.chat import CompletionCreateParams
+from ragas import MultiTurnSample
+
+from datarobot_genai.core.agents.base import BaseAgent
+from datarobot_genai.core.agents.base import InvokeReturn
+from datarobot_genai.core.agents.base import UsageMetrics
+from datarobot_genai.core.agents.base import default_usage_metrics
+from datarobot_genai.core.agents.base import extract_user_prompt_content
+from datarobot_genai.core.agents.base import is_streaming
+
+from .agent import create_pipeline_interactions_from_events
+from .mcp import load_mcp_tools
+
+
+class LlamaIndexAgent(BaseAgent[BaseTool], abc.ABC):
+    """Abstract base agent for LlamaIndex workflows."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._mcp_tools: list[Any] = []
+
+    def set_mcp_tools(self, tools: list[Any]) -> None:
+        """Set MCP tools for this agent."""
+        self._mcp_tools = tools
+
+    @property
+    def mcp_tools(self) -> list[Any]:
+        """Return the list of MCP tools available to this agent.
+
+        Subclasses can use this to wire tools into LlamaIndex agents during
+        workflow construction inside ``build_workflow``.
+        """
+        return self._mcp_tools
+
+    @abc.abstractmethod
+    def build_workflow(self) -> Any:
+        """Return an AgentWorkflow instance ready to run."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def extract_response_text(self, result_state: Any, events: list[Any]) -> str:
+        """Extract final response text from workflow state and/or events."""
+        raise NotImplementedError
+
+    def make_input_message(self, completion_create_params: CompletionCreateParams) -> str:
+        """Create an input string for the workflow from the user prompt."""
+        user_prompt_content = extract_user_prompt_content(completion_create_params)
+        return str(user_prompt_content)
+
+    async def invoke(self, completion_create_params: CompletionCreateParams) -> InvokeReturn:
+        """Run the LlamaIndex workflow with the provided completion parameters."""
+        input_message = self.make_input_message(completion_create_params)
+
+        # Load MCP tools (if configured) asynchronously before building workflow
+        mcp_tools = await load_mcp_tools(
+            authorization_context=self._authorization_context,
+            forwarded_headers=self.forwarded_headers,
+        )
+        self.set_mcp_tools(mcp_tools)
+
+        # Preserve prior template startup print for CLI parity
+        try:
+            print(
+                "Running agent with user prompt:",
+                extract_user_prompt_content(completion_create_params),
+                flush=True,
+            )
+        except Exception:
+            # Printing is best-effort; proceed regardless
+            pass
+
+        workflow = self.build_workflow()
+        handler = workflow.run(user_msg=input_message)
+
+        usage_metrics: UsageMetrics = default_usage_metrics()
+
+        # Streaming parity with LangGraph: yield incremental deltas during event processing
+        if is_streaming(completion_create_params):
+
+            async def _gen() -> AsyncGenerator[tuple[str, MultiTurnSample | None, UsageMetrics]]:
+                events: list[Any] = []
+                current_agent_name: str | None = None
+                async for event in handler.stream_events():
+                    events.append(event)
+                    # Best-effort extraction of incremental text from LlamaIndex events
+                    delta: str | None = None
+                    # Agent switch banner if available on event
+                    try:
+                        if hasattr(event, "current_agent_name"):
+                            new_agent = getattr(event, "current_agent_name")
+                            if (
+                                isinstance(new_agent, str)
+                                and new_agent
+                                and new_agent != current_agent_name
+                            ):
+                                current_agent_name = new_agent
+                                # Print banner for agent switch (do not emit as streamed content)
+                                print("\n" + "=" * 50, flush=True)
+                                print(f"🤖 Agent: {current_agent_name}", flush=True)
+                                print("=" * 50 + "\n", flush=True)
+                    except Exception:
+                        pass
+
+                    try:
+                        if hasattr(event, "delta") and isinstance(getattr(event, "delta"), str):
+                            delta = getattr(event, "delta")
+                        # Some event types may carry incremental text under "text" or similar
+                        elif hasattr(event, "text") and isinstance(getattr(event, "text"), str):
+                            delta = getattr(event, "text")
+                    except Exception:
+                        # Ignore malformed events and continue
+                        delta = None
+
+                    if delta:
+                        # Yield token/content delta with current (accumulated) usage metrics
+                        yield delta, None, usage_metrics
+
+                    # Best-effort debug/event messages printed to CLI (do not stream as content)
+                    try:
+                        event_type = type(event).__name__
+                        if event_type == "AgentInput" and hasattr(event, "input"):
+                            print("📥 Input:", getattr(event, "input"), flush=True)
+                        elif event_type == "AgentOutput":
+                            # Output content
+                            resp = getattr(event, "response", None)
+                            if (
+                                resp is not None
+                                and hasattr(resp, "content")
+                                and getattr(resp, "content")
+                            ):
+                                print("📤 Output:", getattr(resp, "content"), flush=True)
+                            # Planned tool calls
+                            tcalls = getattr(event, "tool_calls", None)
+                            if isinstance(tcalls, list) and tcalls:
+                                names = []
+                                for c in tcalls:
+                                    try:
+                                        nm = getattr(c, "tool_name", None) or (
+                                            c.get("tool_name") if isinstance(c, dict) else None
+                                        )
+                                        if nm:
+                                            names.append(str(nm))
+                                    except Exception:
+                                        pass
+                                if names:
+                                    print("🛠️  Planning to use tools:", names, flush=True)
+                        elif event_type == "ToolCallResult":
+                            tname = getattr(event, "tool_name", None)
+                            tkwargs = getattr(event, "tool_kwargs", None)
+                            tout = getattr(event, "tool_output", None)
+                            print(f"🔧 Tool Result ({tname}):", flush=True)
+                            print(f"  Arguments: {tkwargs}", flush=True)
+                            print(f"  Output: {tout}", flush=True)
+                        elif event_type == "ToolCall":
+                            tname = getattr(event, "tool_name", None)
+                            tkwargs = getattr(event, "tool_kwargs", None)
+                            print(f"🔨 Calling Tool: {tname}", flush=True)
+                            print(f"  With arguments: {tkwargs}", flush=True)
+                    except Exception:
+                        # Ignore best-effort debug rendering errors
+                        pass
+
+                # After streaming completes, build final interactions and finish chunk
+                # Extract state from workflow context (supports sync/async get or attribute)
+                state = None
+                ctx = getattr(handler, "ctx", None)
+                try:
+                    if ctx is not None:
+                        get = getattr(ctx, "get", None)
+                        if callable(get):
+                            result = get("state")
+                            state = await result if inspect.isawaitable(result) else result
+                        elif hasattr(ctx, "state"):
+                            state = getattr(ctx, "state")
+                except (AttributeError, TypeError):
+                    state = None
+
+                # Run subclass-defined response extraction (not streamed) for completeness
+                _ = self.extract_response_text(state, events)
+
+                pipeline_interactions = create_pipeline_interactions_from_events(events)
+                # Final empty chunk indicates end of stream, carrying interactions and usage
+                yield "", pipeline_interactions, usage_metrics
+
+            return _gen()
+
+        # Non-streaming path: run to completion, emit debug prints, then return final response
+        events: list[Any] = []
+        current_agent_name: str | None = None
+        async for event in handler.stream_events():
+            events.append(event)
+
+            # Replicate prior template CLI prints for non-streaming mode
+            try:
+                if hasattr(event, "current_agent_name"):
+                    new_agent = getattr(event, "current_agent_name")
+                    if isinstance(new_agent, str) and new_agent and new_agent != current_agent_name:
+                        current_agent_name = new_agent
+                        print(f"\n{'=' * 50}", flush=True)
+                        print(f"🤖 Agent: {current_agent_name}", flush=True)
+                        print(f"{'=' * 50}\n", flush=True)
+            except Exception:
+                pass
+
+            try:
+                if hasattr(event, "delta") and isinstance(getattr(event, "delta"), str):
+                    print(getattr(event, "delta"), end="", flush=True)
+                elif hasattr(event, "text") and isinstance(getattr(event, "text"), str):
+                    print(getattr(event, "text"), end="", flush=True)
+                else:
+                    event_type = type(event).__name__
+                    if event_type == "AgentInput" and hasattr(event, "input"):
+                        print("📥 Input:", getattr(event, "input"), flush=True)
+                    elif event_type == "AgentOutput":
+                        resp = getattr(event, "response", None)
+                        if (
+                            resp is not None
+                            and hasattr(resp, "content")
+                            and getattr(resp, "content")
+                        ):
+                            print("📤 Output:", getattr(resp, "content"), flush=True)
+                        tcalls = getattr(event, "tool_calls", None)
+                        if isinstance(tcalls, list) and tcalls:
+                            names: list[str] = []
+                            for c in tcalls:
+                                try:
+                                    nm = getattr(c, "tool_name", None) or (
+                                        c.get("tool_name") if isinstance(c, dict) else None
+                                    )
+                                    if nm:
+                                        names.append(str(nm))
+                                except Exception:
+                                    pass
+                            if names:
+                                print("🛠️  Planning to use tools:", names, flush=True)
+                    elif event_type == "ToolCallResult":
+                        tname = getattr(event, "tool_name", None)
+                        tkwargs = getattr(event, "tool_kwargs", None)
+                        tout = getattr(event, "tool_output", None)
+                        print(f"🔧 Tool Result ({tname}):", flush=True)
+                        print(f"  Arguments: {tkwargs}", flush=True)
+                        print(f"  Output: {tout}", flush=True)
+                    elif event_type == "ToolCall":
+                        tname = getattr(event, "tool_name", None)
+                        tkwargs = getattr(event, "tool_kwargs", None)
+                        print(f"🔨 Calling Tool: {tname}", flush=True)
+                        print(f"  With arguments: {tkwargs}", flush=True)
+            except Exception:
+                # Best-effort debug printing; continue on errors
+                pass
+
+        # Extract state from workflow context (supports sync/async get or attribute)
+        state = None
+        ctx = getattr(handler, "ctx", None)
+        try:
+            if ctx is not None:
+                get = getattr(ctx, "get", None)
+                if callable(get):
+                    result = get("state")
+                    state = await result if inspect.isawaitable(result) else result
+                elif hasattr(ctx, "state"):
+                    state = getattr(ctx, "state")
+        except (AttributeError, TypeError):
+            state = None
+        response_text = self.extract_response_text(state, events)
+
+        pipeline_interactions = create_pipeline_interactions_from_events(events)
+
+        return response_text, pipeline_interactions, usage_metrics

datarobot_genai/llama_index/mcp.py

@@ -0,0 +1,79 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+MCP integration for LlamaIndex using llama-index-tools-mcp.
+
+This module provides MCP server connection management for LlamaIndex agents.
+Unlike CrewAI which uses a context manager, LlamaIndex uses async calls to
+fetch tools from MCP servers.
+"""
+
+from typing import Any
+
+from llama_index.tools.mcp import BasicMCPClient
+from llama_index.tools.mcp import aget_tools_from_mcp_url
+
+from datarobot_genai.core.mcp.common import MCPConfig
+
+
+async def load_mcp_tools(
+    authorization_context: dict[str, Any] | None = None,
+    forwarded_headers: dict[str, str] | None = None,
+) -> list[Any]:
+    """
+    Asynchronously load MCP tools for LlamaIndex.
+
+    Args:
+        authorization_context: Optional authorization context for MCP connections
+        forwarded_headers: Optional forwarded headers, e.g. x-datarobot-api-key for MCP auth
+
+    Returns
+    -------
+        List of MCP tools, or empty list if no MCP configuration is present.
+    """
+    config = MCPConfig(
+        authorization_context=authorization_context,
+        forwarded_headers=forwarded_headers,
+    )
+    server_params = config.server_config
+
+    if not server_params:
+        print("No MCP server configured, using empty tools list", flush=True)
+        return []
+
+    url = server_params["url"]
+    headers = server_params.get("headers", {})
+
+    try:
+        print(f"Connecting to MCP server: {url}", flush=True)
+        # Create BasicMCPClient with headers to pass authentication
+        client = BasicMCPClient(command_or_url=url, headers=headers)
+        tools = await aget_tools_from_mcp_url(
+            command_or_url=url,
+            client=client,
+        )
+        # Ensure list
+        tools_list = list(tools) if tools is not None else []
+        print(
+            f"Successfully connected to MCP server, got {len(tools_list)} tools",
+            flush=True,
+        )
+        return tools_list
+    except Exception as e:
+        print(
+            f"Warning: Failed to connect to MCP server {url}: {e}",
+            flush=True,
+        )
+        return []

datarobot_genai/nat/agent.py

@@ -0,0 +1,275 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import asyncio
+import logging
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from nat.builder.context import Context
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import ChatResponse
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.utils.type_utils import StrPath
+from openai.types.chat import CompletionCreateParams
+from ragas import MultiTurnSample
+from ragas.messages import AIMessage
+from ragas.messages import HumanMessage
+from ragas.messages import ToolMessage
+
+from datarobot_genai.core.agents.base import BaseAgent
+from datarobot_genai.core.agents.base import InvokeReturn
+from datarobot_genai.core.agents.base import UsageMetrics
+from datarobot_genai.core.agents.base import extract_user_prompt_content
+from datarobot_genai.core.agents.base import is_streaming
+from datarobot_genai.core.mcp.common import MCPConfig
+from datarobot_genai.nat.helpers import load_workflow
+
+logger = logging.getLogger(__name__)
+
+
+def convert_to_ragas_messages(
+    steps: list[IntermediateStep],
+) -> list[HumanMessage | AIMessage | ToolMessage]:
+    def _to_ragas(step: IntermediateStep) -> HumanMessage | AIMessage | ToolMessage:
+        if step.event_type == IntermediateStepType.LLM_START:
+            return HumanMessage(content=_parse(step.data.input))
+        elif step.event_type == IntermediateStepType.LLM_END:
+            return AIMessage(content=_parse(step.data.output))
+        else:
+            raise ValueError(f"Unknown event type {step.event_type}")
+
+    def _include_step(step: IntermediateStep) -> bool:
+        return step.event_type in {
+            IntermediateStepType.LLM_END,
+            IntermediateStepType.LLM_START,
+        }
+
+    def _parse(messages: Any) -> str:
+        if isinstance(messages, list):
+            last_message = messages[-1]
+        else:
+            last_message = messages
+
+        if isinstance(last_message, dict):
+            content = last_message.get("content") or last_message
+        elif hasattr(last_message, "content"):
+            content = getattr(last_message, "content") or last_message
+        else:
+            content = last_message
+        return str(content)
+
+    return [_to_ragas(step) for step in steps if _include_step(step)]
+
+
+def create_pipeline_interactions_from_steps(
+    steps: list[IntermediateStep],
+) -> MultiTurnSample | None:
+    if not steps:
+        return None
+    ragas_trace = convert_to_ragas_messages(steps)
+    return MultiTurnSample(user_input=ragas_trace)
+
+
+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
+
+
+class NatAgent(BaseAgent[None]):
+    def __init__(
+        self,
+        *,
+        workflow_path: StrPath,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        model: str | None = None,
+        verbose: bool | str | None = True,
+        timeout: int | None = 90,
+        authorization_context: dict[str, Any] | None = None,
+        forwarded_headers: dict[str, str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            api_key=api_key,
+            api_base=api_base,
+            model=model,
+            verbose=verbose,
+            timeout=timeout,
+            authorization_context=authorization_context,
+            forwarded_headers=forwarded_headers,
+            **kwargs,
+        )
+        self.workflow_path = workflow_path
+
+    def make_chat_request(self, completion_create_params: CompletionCreateParams) -> ChatRequest:
+        user_prompt_content = str(extract_user_prompt_content(completion_create_params))
+        return ChatRequest.from_string(user_prompt_content)
+
+    async def invoke(self, completion_create_params: CompletionCreateParams) -> InvokeReturn:
+        """Run the agent with the provided completion parameters.
+
+        [THIS METHOD IS REQUIRED FOR THE AGENT TO WORK WITH DRUM SERVER]
+
+        Args:
+            completion_create_params: The completion request parameters including input topic
+            and settings.
+
+        Returns
+        -------
+            For streaming requests, returns a generator yielding tuples of (response_text,
+            pipeline_interactions, usage_metrics).
+            For non-streaming requests, returns a single tuple of (response_text,
+            pipeline_interactions, usage_metrics).
+
+        """
+        # Retrieve the starting chat request from the CompletionCreateParams
+        chat_request = self.make_chat_request(completion_create_params)
+
+        # Print commands may need flush=True to ensure they are displayed in real-time.
+        print("Running agent with user prompt:", chat_request.messages[0].content, flush=True)
+
+        mcp_config = MCPConfig(
+            authorization_context=self.authorization_context,
+            forwarded_headers=self.forwarded_headers,
+        )
+        server_config = mcp_config.server_config
+        headers = server_config["headers"] if server_config else None
+
+        if is_streaming(completion_create_params):
+
+            async def stream_generator() -> AsyncGenerator[
+                tuple[str, MultiTurnSample | None, UsageMetrics], None
+            ]:
+                default_usage_metrics: UsageMetrics = {
+                    "completion_tokens": 0,
+                    "prompt_tokens": 0,
+                    "total_tokens": 0,
+                }
+                async with load_workflow(self.workflow_path, headers=headers) as workflow:
+                    async with workflow.run(chat_request) as runner:
+                        intermediate_future = pull_intermediate_structured()
+                        async for result in runner.result_stream():
+                            if isinstance(result, ChatResponse):
+                                result_text = result.choices[0].message.content
+                            else:
+                                result_text = str(result)
+
+                            yield (
+                                result_text,
+                                None,
+                                default_usage_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 = create_pipeline_interactions_from_steps(steps)
+                        yield "", pipeline_interactions, usage_metrics
+
+            return stream_generator()
+
+        # Create and invoke the NAT (Nemo Agent Toolkit) Agentic Workflow with the inputs
+        result, steps = await self.run_nat_workflow(self.workflow_path, chat_request, headers)
+
+        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
+
+        if isinstance(result, ChatResponse):
+            result_text = result.choices[0].message.content
+        else:
+            result_text = str(result)
+        pipeline_interactions = create_pipeline_interactions_from_steps(steps)
+
+        return result_text, 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
+            input_str: Input string to process through 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) as workflow:
+            async with workflow.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}"
+
+        print(f"{prefix}{runner_outputs}{suffix}")
+
+        return runner_outputs, steps