datarobot-genai 0.1.75__py3-none-any.whl → 0.2.3__py3-none-any.whl

datarobot_genai/core/agents/base.py

@@ -23,6 +23,7 @@ from typing import TypedDict
 from typing import TypeVar
 from typing import cast
 
+from ag_ui.core import Event
 from openai.types.chat import CompletionCreateParams
 from ragas import MultiTurnSample
 
@@ -167,7 +168,7 @@ class UsageMetrics(TypedDict):
 
 # Canonical return type for DRUM-compatible invoke implementations
 InvokeReturn = (
-    AsyncGenerator[tuple[str, MultiTurnSample | None, UsageMetrics], None]
+    AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]
     | tuple[str, MultiTurnSample | None, UsageMetrics]
 )
 

datarobot_genai/core/chat/responses.py

@@ -27,6 +27,10 @@ from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 from typing import TypeVar
 
+from ag_ui.core import BaseEvent
+from ag_ui.core import Event
+from ag_ui.core import TextMessageChunkEvent
+from ag_ui.core import TextMessageContentEvent
 from openai.types import CompletionUsage
 from openai.types.chat import ChatCompletion
 from openai.types.chat import ChatCompletionChunk
@@ -45,6 +49,7 @@ class CustomModelChatResponse(ChatCompletion):
 
 class CustomModelStreamingResponse(ChatCompletionChunk):
     pipeline_interactions: str | None = None
+    event: Event | None = None
 
 
 def to_custom_model_chat_response(
@@ -88,7 +93,7 @@ def to_custom_model_streaming_response(
     thread_pool_executor: ThreadPoolExecutor,
     event_loop: AbstractEventLoop,
     streaming_response_generator: AsyncGenerator[
-        tuple[str, MultiTurnSample | None, dict[str, int]], None
+        tuple[str | Event, MultiTurnSample | None, dict[str, int]], None
     ],
     model: str | object | None,
 ) -> Iterator[CustomModelStreamingResponse]:
@@ -110,7 +115,7 @@ def to_custom_model_streaming_response(
         while True:
             try:
                 (
-                    response_text,
+                    response_text_or_event,
                     pipeline_interactions,
                     usage_metrics,
                 ) = thread_pool_executor.submit(
@@ -119,10 +124,10 @@ def to_custom_model_streaming_response(
                 last_pipeline_interactions = pipeline_interactions
                 last_usage_metrics = usage_metrics
 
-                if response_text:
+                if isinstance(response_text_or_event, str) and response_text_or_event:
                     choice = ChunkChoice(
                         index=0,
-                        delta=ChoiceDelta(role="assistant", content=response_text),
+                        delta=ChoiceDelta(role="assistant", content=response_text_or_event),
                         finish_reason=None,
                     )
                     yield CustomModelStreamingResponse(
@@ -135,6 +140,29 @@ def to_custom_model_streaming_response(
                         if usage_metrics
                         else None,
                     )
+                elif isinstance(response_text_or_event, BaseEvent):
+                    content = ""
+                    if isinstance(
+                        response_text_or_event, (TextMessageContentEvent, TextMessageChunkEvent)
+                    ):
+                        content = response_text_or_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=response_text_or_event,
+                    )
             except StopAsyncIteration:
                 break
         event_loop.run_until_complete(streaming_response_generator.aclose())

datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -115,6 +115,9 @@ class DataRobotMCPServer:
         self._mcp = mcp
         self._mcp_transport = transport
 
+        # Configure MCP server capabilities
+        self._configure_mcp_capabilities()
+
         # Initialize telemetry
         initialize_telemetry(mcp)
 
@@ -163,6 +166,37 @@ class DataRobotMCPServer:
         if transport == "streamable-http":
             register_routes(self._mcp)
 
+    def _configure_mcp_capabilities(self) -> None:
+        """Configure MCP capabilities that FastMCP doesn't expose directly.
+
+        See: https://github.com/modelcontextprotocol/python-sdk/issues/1126
+        """
+        server = self._mcp._mcp_server
+
+        # Declare prompts_changed capability  (capabilities.prompts.listChanged: true)
+        server.notification_options.prompts_changed = True
+
+        # Declare experimental capabilities ( experimental.dynamic_prompts: true)
+        server.experimental_capabilities = {"dynamic_prompts": {"enabled": True}}
+
+        # Patch to include experimental_capabilities (FastMCP doesn't expose this)
+        original = server.create_initialization_options
+
+        def patched(
+            notification_options: Any = None,
+            experimental_capabilities: dict[str, dict[str, Any]] | None = None,
+            **kwargs: Any,
+        ) -> Any:
+            if experimental_capabilities is None:
+                experimental_capabilities = getattr(server, "experimental_capabilities", None)
+            return original(
+                notification_options=notification_options,
+                experimental_capabilities=experimental_capabilities,
+                **kwargs,
+            )
+
+        server.create_initialization_options = patched
+
     def run(self, show_banner: bool = False) -> None:
         """Run the DataRobot MCP server synchronously."""
         try:

datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py

@@ -12,78 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections import defaultdict
-from dataclasses import dataclass
 
 import datarobot as dr
 
 from datarobot_genai.drmcp.core.clients import get_api_client
 
-# Needed SDK version (3.10.0) is not published yet. We'll reimplement simplified version of it.
-# get_datarobot_prompt_templates = dr.genai.PromptTemplate.list()
-# DrPrompt = dr.genai.PromptTemplate
-# DrPromptVersion = dr.genai.PromptTemplateVersion
-# DrVariable = dr.genai.Variable
 
-
-@dataclass
-class DrVariable:
-    name: str
-    description: str
-
-
-@dataclass
-class DrPromptVersion:
-    id: str
-    prompt_template_id: str
-    version: int
-    prompt_text: str
-    variables: list[DrVariable]
-
-    @classmethod
-    def from_dict(cls, d: dict) -> "DrPromptVersion":
-        variables = [
-            DrVariable(name=v["name"], description=v["description"]) for v in d["variables"]
-        ]
-        return cls(
-            id=d["id"],
-            prompt_template_id=d["promptTemplateId"],
-            version=d["version"],
-            prompt_text=d["promptText"],
-            variables=variables,
-        )
-
-
-@dataclass
-class DrPrompt:
-    id: str
-    name: str
-    description: str
-
-    def get_latest_version(self) -> DrPromptVersion | None:
-        all_prompt_template_versions = get_datarobot_prompt_template_versions([self.id])
-        prompt_template_versions = all_prompt_template_versions.get(self.id)
-
-        if not prompt_template_versions:
-            return None
-        latest_version = max(prompt_template_versions, key=lambda v: v.version)
-        return latest_version
-
-    @classmethod
-    def from_dict(cls, d: dict) -> "DrPrompt":
-        return cls(id=d["id"], name=d["name"], description=d["description"])
-
-
-def get_datarobot_prompt_templates() -> list[DrPrompt]:
-    prompt_templates_data = dr.utils.pagination.unpaginate(
-        initial_url="genai/promptTemplates/", initial_params={}, client=get_api_client()
-    )
-
-    return [DrPrompt.from_dict(prompt_template) for prompt_template in prompt_templates_data]
+def get_datarobot_prompt_templates() -> list[dr.genai.PromptTemplate]:
+    try:
+        return dr.genai.PromptTemplate.list()
+    except Exception:
+        return []
 
 
 def get_datarobot_prompt_template_versions(
     prompt_template_ids: list[str],
-) -> dict[str, list[DrPromptVersion]]:
+) -> dict[str, list[dr.genai.PromptTemplateVersion]]:
+    # Still missing in SDK
     prompt_template_versions_data = dr.utils.pagination.unpaginate(
         initial_url="genai/promptTemplates/versions/",
         initial_params={
@@ -94,35 +39,32 @@ def get_datarobot_prompt_template_versions(
     prompt_template_versions = defaultdict(list)
     for prompt_template_version in prompt_template_versions_data:
         prompt_template_versions[prompt_template_version["promptTemplateId"]].append(
-            DrPromptVersion.from_dict(prompt_template_version)
+            dr.genai.PromptTemplateVersion(
+                id=prompt_template_version["id"],
+                prompt_template_id=prompt_template_version["promptTemplateId"],
+                prompt_text=prompt_template_version["promptText"],
+                commit_comment=prompt_template_version["commitComment"],
+                version=prompt_template_version["version"],
+                variables=prompt_template_version["variables"],
+                creation_date=prompt_template_version["creationDate"],
+                creation_user_id=prompt_template_version["creationUserId"],
+                user_name=prompt_template_version["userName"],
+            )
         )
     return prompt_template_versions
 
 
-def get_datarobot_prompt_template(prompt_template_id: str) -> DrPrompt | None:
-    api_client = get_api_client()
+def get_datarobot_prompt_template(prompt_template_id: str) -> dr.genai.PromptTemplate | None:
     try:
-        prompt_template_response = api_client.get(
-            f"genai/promptTemplates/{prompt_template_id}/", join_endpoint=True
-        )
-        prompt_template_json = prompt_template_response.json()
+        return dr.genai.PromptTemplate.get(prompt_template_id)
     except Exception:
         return None
 
-    return DrPrompt.from_dict(prompt_template_json)
-
 
 def get_datarobot_prompt_template_version(
     prompt_template_id: str, prompt_template_version_id: str
-) -> DrPromptVersion | None:
-    api_client = get_api_client()
+) -> dr.genai.PromptTemplateVersion | None:
     try:
-        prompt_template_version_response = api_client.get(
-            f"genai/promptTemplates/{prompt_template_id}/versions/{prompt_template_version_id}/",
-            join_endpoint=True,
-        )
-        prompt_template_version_json = prompt_template_version_response.json()
+        return dr.genai.PromptTemplateVersion.get(prompt_template_id, prompt_template_version_id)
     except Exception:
         return None
-
-    return DrPromptVersion.from_dict(prompt_template_version_json)

datarobot_genai/drmcp/core/dynamic_prompts/register.py

@@ -18,15 +18,13 @@ from collections.abc import Callable
 from inspect import Parameter
 from inspect import Signature
 
+import datarobot as dr
 from fastmcp.prompts.prompt import Prompt
 from pydantic import Field
 
 from datarobot_genai.drmcp.core.exceptions import DynamicPromptRegistrationError
 from datarobot_genai.drmcp.core.mcp_instance import register_prompt
 
-from .dr_lib import DrPrompt
-from .dr_lib import DrPromptVersion
-from .dr_lib import DrVariable
 from .dr_lib import get_datarobot_prompt_template_versions
 from .dr_lib import get_datarobot_prompt_templates
 
@@ -57,7 +55,8 @@ async def register_prompts_from_datarobot_prompt_management() -> None:
 
 
 async def register_prompt_from_datarobot_prompt_management(
-    prompt_template: DrPrompt, prompt_template_version: DrPromptVersion | None = None
+    prompt_template: dr.genai.PromptTemplate,
+    prompt_template_version: dr.genai.PromptTemplateVersion | None = None,
 ) -> Prompt:
     """Register a single prompt.
 
@@ -173,7 +172,7 @@ def to_valid_mcp_prompt_name(s: str) -> str:
 
 
 def make_prompt_function(
-    name: str, description: str, prompt_text: str, variables: list[DrVariable]
+    name: str, description: str, prompt_text: str, variables: list[dr.genai.Variable]
 ) -> Callable:
     params = []
     for v in variables:

datarobot_genai/drmcp/core/mcp_instance.py

@@ -22,6 +22,7 @@ from fastmcp import Context
 from fastmcp import FastMCP
 from fastmcp.exceptions import NotFoundError
 from fastmcp.prompts.prompt import Prompt
+from fastmcp.server.dependencies import get_context
 from fastmcp.tools import FunctionTool
 from fastmcp.tools import Tool
 from fastmcp.utilities.types import NotSet
@@ -91,6 +92,34 @@ class TaggedFastMCP(FastMCP):
         self._deployments_map: dict[str, str] = {}
         self._prompts_map: dict[str, tuple[str, str]] = {}
 
+    async def notify_prompts_changed(self) -> None:
+        """
+        Notify connected clients that the prompt list has changed.
+
+        This method attempts to send a prompts/list_changed notification to inform
+        clients that they should refresh their prompt list.
+
+        Note: In stateless HTTP mode (default for this server), notifications may not
+        reach clients since each request is independent. This method still logs the
+        change for auditing purposes and will work if the server is configured for
+        stateful connections.
+
+        See: https://github.com/modelcontextprotocol/python-sdk/issues/710
+        """
+        logger.info("Prompt list changed - attempting to notify connected clients")
+
+        # Try to use FastMCP's built-in notification mechanism if in an MCP context
+        try:
+            context = get_context()
+            context._queue_prompt_list_changed()
+            logger.debug("Queued prompts_changed notification via MCP context")
+        except RuntimeError:
+            # No active MCP context - this is expected when called from REST API
+            logger.debug(
+                "No active MCP context for notification. "
+                "In stateless mode, clients will see changes on next request."
+            )
+
     @overload
     def tool(
         self,
@@ -286,6 +315,9 @@ class TaggedFastMCP(FastMCP):
                 f"already mapped to {existing_prompt_template_version_id}. "
                 f"Updating to version id = {prompt_template_version_id} and name = {prompt_name}"
             )
+            await self.remove_prompt_mapping(
+                prompt_template_id, existing_prompt_template_version_id
+            )
 
         self._prompts_map[prompt_template_id] = (prompt_template_version_id, prompt_name)
 
@@ -308,7 +340,7 @@ class TaggedFastMCP(FastMCP):
                     f"skipping removal."
                 )
             else:
-                prompts_d = await mcp.get_prompts()
+                prompts_d = await self.get_prompts()
                 for prompt in prompts_d.values():
                     if (
                         prompt.meta is not None
@@ -319,6 +351,9 @@ class TaggedFastMCP(FastMCP):
                         prompt.disable()
 
                 self._prompts_map.pop(prompt_template_id, None)
+
+                # Notify clients that the prompt list has changed
+                await self.notify_prompts_changed()
         else:
             logger.debug(
                 f"Do not found prompt template with id = {prompt_template_id} in registry, "
@@ -526,17 +561,21 @@ async def register_prompt(
     )
 
     # Register the prompt
-    registered_prompt = mcp.add_prompt(prompt)
     if prompt_template:
         prompt_template_id, prompt_template_version_id = prompt_template
         await mcp.set_prompt_mapping(
             prompt_template_id, prompt_template_version_id, prompt_name_no_duplicate
         )
 
+    registered_prompt = mcp.add_prompt(prompt)
+
     # Verify prompt is registered
     prompts = await mcp.get_prompts()
     if not any(prompt.name == prompt_name_no_duplicate for prompt in prompts.values()):
         raise RuntimeError(f"Prompt {prompt_name_no_duplicate} was not registered successfully")
     logger.info(f"Registered prompts: {len(prompts)}")
 
+    # Notify clients that the prompt list has changed
+    await mcp.notify_prompts_changed()
+
     return registered_prompt

datarobot_genai/drmcp/core/routes.py

@@ -428,7 +428,10 @@ def register_routes(mcp: TaggedFastMCP) -> None:
         """Refresh prompt templates."""
         try:
             await refresh_registered_prompt_template()
-            return JSONResponse(status_code=HTTPStatus.NO_CONTENT, content=None)
+            return JSONResponse(
+                status_code=HTTPStatus.OK,
+                content={"message": "Prompts refreshed successfully"},
+            )
         except Exception as e:
             return JSONResponse(
                 status_code=HTTPStatus.INTERNAL_SERVER_ERROR,

datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@@ -16,6 +16,8 @@ import os
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 
+import aiohttp
+from aiohttp import ClientSession as HttpClientSession
 from mcp import ClientSession
 from mcp.client.streamable_http import streamablehttp_client
 
@@ -29,6 +31,11 @@ def get_dr_mcp_server_url() -> str | None:
     return os.environ.get("DR_MCP_SERVER_URL")
 
 
+def get_dr_mcp_server_http_url() -> str | None:
+    """Get DataRobot MCP server http URL."""
+    return os.environ.get("DR_MCP_SERVER_HTTP_URL")
+
+
 def get_openai_llm_client_config() -> dict[str, str]:
     """Get OpenAI LLM client configuration."""
     openai_api_key = os.environ.get("OPENAI_API_KEY")
@@ -94,3 +101,25 @@ async def ete_test_mcp_session(
                 yield session
     except asyncio.TimeoutError:
         raise TimeoutError(f"Check if the MCP server is running at {get_dr_mcp_server_url()}")
+
+
+@asynccontextmanager
+async def ete_test_http_session(
+    additional_headers: dict[str, str] | None = None,
+) -> AsyncGenerator[HttpClientSession, None]:
+    """Create an HTTP session for each test that can connect to MCP custom http routes.
+
+    Parameters
+    ----------
+    additional_headers : dict[str, str], optional
+        Additional headers to include in the HTTP session (e.g., auth headers for testing).
+    """
+    headers = get_headers()
+    if additional_headers:
+        headers.update(additional_headers)
+
+    async with ete_test_mcp_session(additional_headers=additional_headers):
+        async with aiohttp.ClientSession(
+            base_url=get_dr_mcp_server_http_url(), headers=headers
+        ) as client:
+            yield client

datarobot_genai/langgraph/agent.py

@@ -17,6 +17,15 @@ from collections.abc import AsyncGenerator
 from typing import Any
 from typing import cast
 
+from ag_ui.core import Event
+from ag_ui.core import EventType
+from ag_ui.core import TextMessageContentEvent
+from ag_ui.core import TextMessageEndEvent
+from ag_ui.core import TextMessageStartEvent
+from ag_ui.core import ToolCallArgsEvent
+from ag_ui.core import ToolCallEndEvent
+from ag_ui.core import ToolCallResultEvent
+from ag_ui.core import ToolCallStartEvent
 from langchain.tools import BaseTool
 from langchain_core.messages import AIMessageChunk
 from langchain_core.messages import ToolMessage
@@ -158,43 +167,7 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
         # The main difference is returning a generator for streaming or a final response for sync.
         if is_streaming(completion_create_params):
             # Streaming response: yield each message as it is generated
-            async def stream_generator() -> AsyncGenerator[
-                tuple[str, MultiTurnSample | None, UsageMetrics], None
-            ]:
-                # Iterate over the graph stream. For message events, yield the content.
-                # For update events, accumulate the usage metrics.
-                events = []
-                async for _, mode, event in graph_stream:
-                    if mode == "messages":
-                        message_event: tuple[AIMessageChunk, dict[str, Any]] = event  # type: ignore[assignment]
-                        llm_token, _ = message_event
-                        yield (
-                            str(llm_token.content),
-                            None,
-                            usage_metrics,
-                        )
-                    elif mode == "updates":
-                        update_event: dict[str, Any] = event  # type: ignore[assignment]
-                        events.append(update_event)
-                        current_node = next(iter(update_event))
-                        node_data = update_event[current_node]
-                        current_usage = node_data.get("usage", {}) if node_data is not None else {}
-                        if current_usage:
-                            usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
-                            usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
-                            usage_metrics["completion_tokens"] += current_usage.get(
-                                "completion_tokens", 0
-                            )
-                    else:
-                        raise ValueError(f"Invalid mode: {mode}")
-
-                # Create a list of events from the event listener
-                pipeline_interactions = self.create_pipeline_interactions_from_events(events)
-
-                # yield the final response indicating completion
-                yield "", pipeline_interactions, usage_metrics
-
-            return stream_generator()
+            return self._stream_generator(graph_stream, usage_metrics)
         else:
             # Synchronous response: collect all events and return the final message
             events: list[dict[str, Any]] = [
@@ -203,6 +176,16 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
                 if mode == "updates"
             ]
 
+            # Accumulate the usage metrics from the updates
+            for update in events:
+                current_node = next(iter(update))
+                node_data = update[current_node]
+                current_usage = node_data.get("usage", {}) if node_data is not None else {}
+                if current_usage:
+                    usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
+                    usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
+                    usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
+
             pipeline_interactions = self.create_pipeline_interactions_from_events(events)
 
             # Extract the final event from the graph stream as the synchronous response
@@ -214,14 +197,132 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
                 if node_data is not None and "messages" in node_data
                 else ""
             )
-            current_usage = node_data.get("usage", {}) if node_data is not None else {}
-            if current_usage:
-                usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
-                usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
-                usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
 
             return response_text, pipeline_interactions, usage_metrics
 
+    async def _stream_generator(
+        self, graph_stream: AsyncGenerator[tuple[Any, str, Any], None], usage_metrics: UsageMetrics
+    ) -> AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]:
+        # Iterate over the graph stream. For message events, yield the content.
+        # For update events, accumulate the usage metrics.
+        events = []
+        current_message_id = None
+        tool_call_id = ""
+        async for _, mode, event in graph_stream:
+            if mode == "messages":
+                message_event: tuple[AIMessageChunk | ToolMessage, dict[str, Any]] = event  # type: ignore[assignment]
+                message = message_event[0]
+                if isinstance(message, ToolMessage):
+                    yield (
+                        ToolCallEndEvent(
+                            type=EventType.TOOL_CALL_END, tool_call_id=message.tool_call_id
+                        ),
+                        None,
+                        usage_metrics,
+                    )
+                    yield (
+                        ToolCallResultEvent(
+                            type=EventType.TOOL_CALL_RESULT,
+                            message_id=message.id,
+                            tool_call_id=message.tool_call_id,
+                            content=message.content,
+                            role="tool",
+                        ),
+                        None,
+                        usage_metrics,
+                    )
+                    tool_call_id = ""
+                elif isinstance(message, AIMessageChunk):
+                    if message.tool_call_chunks:
+                        # This is a tool call message
+                        for tool_call_chunk in message.tool_call_chunks:
+                            if name := tool_call_chunk.get("name"):
+                                # Its a tool call start message
+                                tool_call_id = tool_call_chunk["id"]
+                                yield (
+                                    ToolCallStartEvent(
+                                        type=EventType.TOOL_CALL_START,
+                                        tool_call_id=tool_call_id,
+                                        tool_call_name=name,
+                                        parent_message_id=message.id,
+                                    ),
+                                    None,
+                                    usage_metrics,
+                                )
+                            elif args := tool_call_chunk.get("args"):
+                                # Its a tool call args message
+                                yield (
+                                    ToolCallArgsEvent(
+                                        type=EventType.TOOL_CALL_ARGS,
+                                        # Its empty when the tool chunk is not a start message
+                                        # So we use the tool call id from a previous start message
+                                        tool_call_id=tool_call_id,
+                                        delta=args,
+                                    ),
+                                    None,
+                                    usage_metrics,
+                                )
+                    elif message.content:
+                        # Its a text message
+                        # Handle the start and end of the text message
+                        if message.id != current_message_id:
+                            if current_message_id:
+                                yield (
+                                    TextMessageEndEvent(
+                                        type=EventType.TEXT_MESSAGE_END,
+                                        message_id=current_message_id,
+                                    ),
+                                    None,
+                                    usage_metrics,
+                                )
+                            current_message_id = message.id
+                            yield (
+                                TextMessageStartEvent(
+                                    type=EventType.TEXT_MESSAGE_START,
+                                    message_id=message.id,
+                                    role="assistant",
+                                ),
+                                None,
+                                usage_metrics,
+                            )
+                        yield (
+                            TextMessageContentEvent(
+                                type=EventType.TEXT_MESSAGE_CONTENT,
+                                message_id=message.id,
+                                delta=message.content,
+                            ),
+                            None,
+                            usage_metrics,
+                        )
+                else:
+                    raise ValueError(f"Invalid message event: {message_event}")
+            elif mode == "updates":
+                update_event: dict[str, Any] = event  # type: ignore[assignment]
+                events.append(update_event)
+                current_node = next(iter(update_event))
+                node_data = update_event[current_node]
+                current_usage = node_data.get("usage", {}) if node_data is not None else {}
+                if current_usage:
+                    usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
+                    usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
+                    usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
+                if current_message_id:
+                    yield (
+                        TextMessageEndEvent(
+                            type=EventType.TEXT_MESSAGE_END,
+                            message_id=current_message_id,
+                        ),
+                        None,
+                        usage_metrics,
+                    )
+                    current_message_id = None
+
+        # Create a list of events from the event listener
+        pipeline_interactions = self.create_pipeline_interactions_from_events(events)
+
+        # yield the final response indicating completion
+        yield "", pipeline_interactions, usage_metrics
+
     @classmethod
     def create_pipeline_interactions_from_events(
         cls,

{datarobot_genai-0.1.75.dist-info → datarobot_genai-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.75
+Version: 0.2.3
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -8,14 +8,16 @@ License: Apache-2.0
 License-File: AUTHORS
 License-File: LICENSE
 Requires-Python: <3.13,>=3.10
+Requires-Dist: ag-ui-protocol<0.2.0,>=0.1.9
 Requires-Dist: datarobot-drum<2.0.0,>=1.17.5
 Requires-Dist: datarobot-predict<2.0.0,>=1.13.2
-Requires-Dist: datarobot<4.0.0,>=3.9.1
+Requires-Dist: datarobot<4.0.0,>=3.10.0
 Requires-Dist: openai<2.0.0,>=1.76.2
 Requires-Dist: opentelemetry-instrumentation-aiohttp-client<1.0.0,>=0.43b0
 Requires-Dist: opentelemetry-instrumentation-httpx<1.0.0,>=0.43b0
 Requires-Dist: opentelemetry-instrumentation-openai<1.0.0,>=0.40.5
 Requires-Dist: opentelemetry-instrumentation-requests<1.0.0,>=0.43b0
+Requires-Dist: opentelemetry-instrumentation-threading<1.0.0,>=0.43b0
 Requires-Dist: pandas<3.0.0,>=2.2.3
 Requires-Dist: pyjwt<3.0.0,>=2.10.1
 Requires-Dist: pypdf<7.0.0,>=6.1.3

{datarobot_genai-0.1.75.dist-info → datarobot_genai-0.2.3.dist-info}/RECORD

@@ -4,11 +4,11 @@ datarobot_genai/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hS
 datarobot_genai/core/custom_model.py,sha256=9rSJhKZ1Fo8BpVVLnRMUj3QV2AYUr5ADHEs0aW6s2tY,5813
 datarobot_genai/core/telemetry_agent.py,sha256=CxvoyResG3jXQ7ucU26NXCzWjWQyua-5qSYvVxpZJQg,5343
 datarobot_genai/core/agents/__init__.py,sha256=mTG_QVV5aoOWOgVA3KEq7KQLJllyxtG2ZQoq9wiUNYo,1542
-datarobot_genai/core/agents/base.py,sha256=kyMk9m1qGFT_XretYIEbP0F-vstPIWh_OZlECqWXvdw,6685
+datarobot_genai/core/agents/base.py,sha256=63JRi0r7L94p6QHpor4JBFu_g5nxaQkqb_I4dsfrDc0,6722
 datarobot_genai/core/chat/__init__.py,sha256=kAxp4Dc-6HIM_cdBl-3IxwzJQr13UYYQ2Zc-hMwz2F8,638
 datarobot_genai/core/chat/auth.py,sha256=6qITKTHFtESsBc2NsA6cvJf78pPUrcA5XV3Vxlhb5us,5457
 datarobot_genai/core/chat/client.py,sha256=fk8MebXa8_R33VK0_DrXCS0Fgw3wFvPEvsuubC27c3s,6639
-datarobot_genai/core/chat/responses.py,sha256=CCkeRpoG7RDI0-JBFhfSUVNaUOspTeH93q3YTUHpZ4A,9242
+datarobot_genai/core/chat/responses.py,sha256=vGxTA433f2AxGVlijV6O4EghyNPJCDmEqpAK2oWnsIs,10583
 datarobot_genai/core/cli/__init__.py,sha256=B93Yb6VavoZpatrh8ltCL6YglIfR5FHgytXbO9UuxBw,733
 datarobot_genai/core/cli/agent_environment.py,sha256=BJzQoiDvZF5gW4mFE71U0yeg-l72C--kxiE-fv6W194,1662
 datarobot_genai/core/cli/agent_kernel.py,sha256=3XX58DQ6XPpWB_tn5m3iGb3XTfhZf5X3W9tc6ADieU4,7790
@@ -31,13 +31,13 @@ datarobot_genai/drmcp/core/config.py,sha256=D7bSi40Yc5J71_JxmpfppG83snbIJW9iz1J7
 datarobot_genai/drmcp/core/config_utils.py,sha256=U-aieWw7MyP03cGDFIp97JH99ZUfr3vD9uuTzBzxn7w,6428
 datarobot_genai/drmcp/core/constants.py,sha256=lUwoW_PTrbaBGqRJifKqCn3EoFacoEgdO-CpoFVrUoU,739
 datarobot_genai/drmcp/core/credentials.py,sha256=PYEUDNMVw1BoMzZKLkPVTypNkVevEPtmk3scKnE-zYg,6706
-datarobot_genai/drmcp/core/dr_mcp_server.py,sha256=7mu5UXHQmKNbIpNoQE0lPJaUI7AZa03avfHZRRtpjNI,12841
+datarobot_genai/drmcp/core/dr_mcp_server.py,sha256=peAq0TL4ZL0P6XJjwCEWNfi0OVPPSQMSRQyKFH6yjtY,14228
 datarobot_genai/drmcp/core/dr_mcp_server_logo.py,sha256=hib-nfR1SNTW6CnpFsFCkL9H_OMwa4YYyinV7VNOuLk,4708
 datarobot_genai/drmcp/core/exceptions.py,sha256=eqsGI-lxybgvWL5w4BFhbm3XzH1eU5tetwjnhJxelpc,905
 datarobot_genai/drmcp/core/logging.py,sha256=Y_hig4eBWiXGaVV7B_3wBcaYVRNH4ydptbEQhrP9-mY,3414
-datarobot_genai/drmcp/core/mcp_instance.py,sha256=wMsP39xqTmNBYqd49olEQb5UHTSsxj6BOIoIElorRB0,19235
+datarobot_genai/drmcp/core/mcp_instance.py,sha256=hArS-BIdsIdRyRA21a4_ILgqqzmuRxZts-Ewgtf1H60,20917
 datarobot_genai/drmcp/core/mcp_server_tools.py,sha256=odNZKozfx0VV38SLZHw9lY0C0JM_JnRI06W3BBXnyE4,4278
-datarobot_genai/drmcp/core/routes.py,sha256=nrgzYkiWWcwSy0BUJ-k5AuXdUHNrnGRudCo4AhtYvlY,17856
+datarobot_genai/drmcp/core/routes.py,sha256=dqE2M0UzAyyN9vQjlyTjYW4rpju3LT039po5weuO__I,17936
 datarobot_genai/drmcp/core/routes_utils.py,sha256=vSseXWlplMSnRgoJgtP_rHxWSAVYcx_tpTv4lyTpQoc,944
 datarobot_genai/drmcp/core/server_life_cycle.py,sha256=WKGJWGxalvqxupzJ2y67Kklc_9PgpZT0uyjlv_sr5wc,3419
 datarobot_genai/drmcp/core/telemetry.py,sha256=NEkSTC1w6uQgtukLHI-sWvR4EMgInysgATcvfQ5CplM,15378
@@ -45,8 +45,8 @@ datarobot_genai/drmcp/core/tool_filter.py,sha256=tLOcG50QBvS48cOVHM6OqoODYiiS6Ke
 datarobot_genai/drmcp/core/utils.py,sha256=dSjrayWVcnC5GxQcvOIOSHaoEymPIVtG_s2ZBMlmSOw,4336
 datarobot_genai/drmcp/core/dynamic_prompts/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
 datarobot_genai/drmcp/core/dynamic_prompts/controllers.py,sha256=AGJlKqgHRO0Kd7Gl-Ulw9KYBgzjTTFXWBvOUF-SuKUI,5454
-datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py,sha256=IEdD2Gqm4SfUdiXJB99RiWxkN6frGaxJ2SfATetMM3c,4243
-datarobot_genai/drmcp/core/dynamic_prompts/register.py,sha256=5AEh1m8GX-gPZHUdiE1VATt7IKJQk-eThcxh01sWn0I,7204
+datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py,sha256=4j33AKmq7kQX_EE2_RWAbP8-K5KPVEvpUginTWn_MHs,2701
+datarobot_genai/drmcp/core/dynamic_prompts/register.py,sha256=2c-vBaTfu3mq_8tSFfDAzG5hG06uS9CghIC1sJxHRNw,7173
 datarobot_genai/drmcp/core/dynamic_prompts/utils.py,sha256=BZ3792AgfvYlwL0_J0MzQfGecyEA5_OKUMynEZYzCds,1136
 datarobot_genai/drmcp/core/dynamic_tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/core/dynamic_tools/register.py,sha256=3M5-F0mhUYTZJWmFDmqzsj3QAd7ut7b0kPv-JZyaTzg,9204
@@ -67,7 +67,7 @@ datarobot_genai/drmcp/core/memory_management/manager.py,sha256=gmc_SQs12YQFMWl2U
 datarobot_genai/drmcp/core/memory_management/memory_tools.py,sha256=AxzpwOlldmhhDfKZcAxaGs7Xih2SCe0XbQuXX5nQczI,6397
 datarobot_genai/drmcp/test_utils/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
 datarobot_genai/drmcp/test_utils/integration_mcp_server.py,sha256=MdoR7r3m9uT7crodyhY69yhkrM7Thpe__BBD9lB_2oA,3328
-datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=rvLOePXF9epIEksJzTOqdL9-fGn21bRLbK3yzpvi-6E,3430
+datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=rgZkPF26YCHX2FGppWE4v22l_NQ3kLSPSUimO0tD4nM,4402
 datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=0sU29Khal0CelnHBDInyTRiuPKrFFbTbIomOoUbyMhs,3271
 datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=Va3_5c2ToZyfIsEjK2ef5d3z-FA5SE51voikvjKPt8Q,8837
 datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=-mKHBkGkyOKQCVS2LHFhSnRofIqJBbeAPRkwizBDtTg,6104
@@ -83,7 +83,7 @@ datarobot_genai/drmcp/tools/predictive/predict_realtime.py,sha256=t7f28y_ealZoA6
 datarobot_genai/drmcp/tools/predictive/project.py,sha256=KaMDAvJY4s12j_4ybA7-KcCS1yMOj-KPIKNBgCSE2iM,2536
 datarobot_genai/drmcp/tools/predictive/training.py,sha256=kxeDVLqUh9ajDk8wK7CZRRydDK8UNuTVZCB3huUihF8,23660
 datarobot_genai/langgraph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/langgraph/agent.py,sha256=P_vqNO-gFOLREHh8NdUXwc9cRlz6dWUY3B_z0r2vpQQ,10945
+datarobot_genai/langgraph/agent.py,sha256=DRnywmS9KDywyChtuIZZwNKbJs8BpC259EG_kxYbiQ8,15828
 datarobot_genai/langgraph/mcp.py,sha256=iA2_j46mZAaNaL7ntXT-LW6C-NMJkzr3VfKDDfe7mh8,2851
 datarobot_genai/llama_index/__init__.py,sha256=JEMkLQLuP8n14kNE3bZ2j08NdajnkJMfYjDQYqj7C0c,407
 datarobot_genai/llama_index/agent.py,sha256=V6ZsD9GcBDJS-RJo1tJtIHhyW69_78gM6_fOHFV-Piw,1829
@@ -93,9 +93,9 @@ datarobot_genai/nat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 datarobot_genai/nat/agent.py,sha256=siBLDWAff2-JwZ8Q3iNpM_e4_IoSwG9IvY0hyEjNenw,10292
 datarobot_genai/nat/datarobot_llm_clients.py,sha256=STzAZ4OF8U-Y_cUTywxmKBGVotwsnbGP6vTojnu6q0g,9921
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
-datarobot_genai-0.1.75.dist-info/METADATA,sha256=7tWV-nzYpfF4oJgeUO8oINJ7sf5pkC2fvmIfFQuc3x8,5898
-datarobot_genai-0.1.75.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.1.75.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
-datarobot_genai-0.1.75.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.1.75.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.1.75.dist-info/RECORD,,
+datarobot_genai-0.2.3.dist-info/METADATA,sha256=TVfYbmNgyNUZ0hv9lg-p3lrcJAGJUgmaGU6AVK6qxSE,6012
+datarobot_genai-0.2.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.3.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
+datarobot_genai-0.2.3.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.3.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.3.dist-info/RECORD,,