datarobot-genai 0.1.71__tar.gz → 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (101) hide show
  1. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/PKG-INFO +4 -3
  2. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/pyproject.toml +4 -3
  3. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/agents/base.py +2 -1
  4. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/chat/responses.py +131 -4
  5. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/custom_model.py +0 -32
  6. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/utils/auth.py +16 -1
  7. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/langgraph/agent.py +143 -42
  8. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/.gitignore +0 -0
  9. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/AUTHORS +0 -0
  10. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/LICENSE +0 -0
  11. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/README.md +0 -0
  12. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/__init__.py +0 -0
  13. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/__init__.py +0 -0
  14. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/agents/__init__.py +0 -0
  15. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/chat/__init__.py +0 -0
  16. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/chat/auth.py +0 -0
  17. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/chat/client.py +0 -0
  18. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/cli/__init__.py +0 -0
  19. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/cli/agent_environment.py +0 -0
  20. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/cli/agent_kernel.py +0 -0
  21. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/mcp/__init__.py +0 -0
  22. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/mcp/common.py +0 -0
  23. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/telemetry_agent.py +0 -0
  24. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/utils/__init__.py +0 -0
  25. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/core/utils/urls.py +0 -0
  26. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/crewai/__init__.py +0 -0
  27. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/crewai/agent.py +0 -0
  28. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/crewai/base.py +0 -0
  29. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/crewai/events.py +0 -0
  30. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/crewai/mcp.py +0 -0
  31. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/__init__.py +0 -0
  32. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/__init__.py +0 -0
  33. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/auth.py +0 -0
  34. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/clients.py +0 -0
  35. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/config.py +0 -0
  36. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/config_utils.py +0 -0
  37. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/constants.py +0 -0
  38. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/credentials.py +0 -0
  39. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dr_mcp_server.py +0 -0
  40. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py +0 -0
  41. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/__init__.py +0 -0
  42. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py +0 -0
  43. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py +0 -0
  44. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py +0 -0
  45. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_prompts/utils.py +0 -0
  46. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/__init__.py +0 -0
  47. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/__init__.py +0 -0
  48. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/__init__.py +0 -0
  49. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/base.py +0 -0
  50. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/default.py +0 -0
  51. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/drum.py +0 -0
  52. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/config.py +0 -0
  53. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/controllers.py +0 -0
  54. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/metadata.py +0 -0
  55. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/register.py +0 -0
  56. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_agentic_fallback_schema.json +0 -0
  57. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_prediction_fallback_schema.json +0 -0
  58. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/register.py +0 -0
  59. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/dynamic_tools/schema.py +0 -0
  60. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/exceptions.py +0 -0
  61. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/logging.py +0 -0
  62. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/mcp_instance.py +0 -0
  63. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/mcp_server_tools.py +0 -0
  64. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/memory_management/__init__.py +0 -0
  65. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/memory_management/manager.py +0 -0
  66. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/memory_management/memory_tools.py +0 -0
  67. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/routes.py +0 -0
  68. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/routes_utils.py +0 -0
  69. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/server_life_cycle.py +0 -0
  70. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/telemetry.py +0 -0
  71. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/tool_filter.py +0 -0
  72. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/core/utils.py +0 -0
  73. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/server.py +0 -0
  74. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/test_utils/__init__.py +0 -0
  75. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/test_utils/integration_mcp_server.py +0 -0
  76. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/test_utils/mcp_utils_ete.py +0 -0
  77. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/test_utils/mcp_utils_integration.py +0 -0
  78. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py +0 -0
  79. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/test_utils/tool_base_ete.py +0 -0
  80. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/test_utils/utils.py +0 -0
  81. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/__init__.py +0 -0
  82. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/__init__.py +0 -0
  83. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/data.py +0 -0
  84. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/deployment.py +0 -0
  85. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/deployment_info.py +0 -0
  86. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/model.py +0 -0
  87. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/predict.py +0 -0
  88. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/predict_realtime.py +0 -0
  89. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/project.py +0 -0
  90. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/drmcp/tools/predictive/training.py +0 -0
  91. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/langgraph/__init__.py +0 -0
  92. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/langgraph/mcp.py +0 -0
  93. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/llama_index/__init__.py +0 -0
  94. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/llama_index/agent.py +0 -0
  95. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/llama_index/base.py +0 -0
  96. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/llama_index/mcp.py +0 -0
  97. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/nat/__init__.py +0 -0
  98. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/nat/agent.py +0 -0
  99. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/nat/datarobot_llm_clients.py +0 -0
  100. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/nat/datarobot_llm_providers.py +0 -0
  101. {datarobot_genai-0.1.71 → datarobot_genai-0.2.0}/src/datarobot_genai/py.typed +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: datarobot-genai
3
- Version: 0.1.71
3
+ Version: 0.2.0
4
4
  Summary: Generic helpers for GenAI
5
5
  Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
6
6
  Author: DataRobot, Inc.
@@ -8,6 +8,7 @@ License: Apache-2.0
8
8
  License-File: AUTHORS
9
9
  License-File: LICENSE
10
10
  Requires-Python: <3.13,>=3.10
11
+ Requires-Dist: ag-ui-protocol<0.2.0,>=0.1.9
11
12
  Requires-Dist: datarobot-drum<2.0.0,>=1.17.5
12
13
  Requires-Dist: datarobot-predict<2.0.0,>=1.13.2
13
14
  Requires-Dist: datarobot<4.0.0,>=3.9.1
@@ -23,7 +24,7 @@ Requires-Dist: ragas<0.4.0,>=0.3.8
23
24
  Requires-Dist: requests<3.0.0,>=2.32.4
24
25
  Provides-Extra: crewai
25
26
  Requires-Dist: crewai-tools[mcp]<0.77.0,>=0.69.0; extra == 'crewai'
26
- Requires-Dist: crewai<1.0.0,>=0.193.2; extra == 'crewai'
27
+ Requires-Dist: crewai>=1.1.0; extra == 'crewai'
27
28
  Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'crewai'
28
29
  Requires-Dist: pybase64<2.0.0,>=1.4.2; extra == 'crewai'
29
30
  Provides-Extra: drmcp
@@ -56,8 +57,8 @@ Requires-Dist: llama-index<0.14.0,>=0.13.6; extra == 'llamaindex'
56
57
  Requires-Dist: opentelemetry-instrumentation-llamaindex<1.0.0,>=0.40.5; extra == 'llamaindex'
57
58
  Requires-Dist: pypdf<7.0.0,>=6.0.0; extra == 'llamaindex'
58
59
  Provides-Extra: nat
60
+ Requires-Dist: crewai>=1.1.0; (python_version >= '3.11') and extra == 'nat'
59
61
  Requires-Dist: llama-index-llms-litellm<0.7.0,>=0.4.1; extra == 'nat'
60
- Requires-Dist: nvidia-nat-crewai==1.3.0; (python_version >= '3.11') and extra == 'nat'
61
62
  Requires-Dist: nvidia-nat-langchain==1.3.0; (python_version >= '3.11') and extra == 'nat'
62
63
  Requires-Dist: nvidia-nat-opentelemetry==1.3.0; (python_version >= '3.11') and extra == 'nat'
63
64
  Requires-Dist: nvidia-nat==1.3.0; (python_version >= '3.11') and extra == 'nat'
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "datarobot-genai"
7
- version = "0.1.71"
7
+ version = "0.2.0"
8
8
  description = "Generic helpers for GenAI"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10, <3.13"
@@ -24,6 +24,7 @@ dependencies = [
24
24
  "opentelemetry-instrumentation-aiohttp-client>=0.43b0,<1.0.0",
25
25
  "opentelemetry-instrumentation-httpx>=0.43b0,<1.0.0",
26
26
  "opentelemetry-instrumentation-openai>=0.40.5,<1.0.0",
27
+ "ag-ui-protocol>=0.1.9,<0.2.0",
27
28
  ]
28
29
 
29
30
  [project.urls]
@@ -35,7 +36,7 @@ datarobot_llm_clients = "datarobot_genai.nat.datarobot_llm_clients"
35
36
 
36
37
  [project.optional-dependencies]
37
38
  crewai = [
38
- "crewai>=0.193.2,<1.0.0",
39
+ "crewai>=1.1.0",
39
40
  "crewai-tools[mcp]>=0.69.0,<0.77.0",
40
41
  "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
41
42
  "pybase64>=1.4.2,<2.0.0",
@@ -59,8 +60,8 @@ llamaindex = [
59
60
  nat = [
60
61
  "nvidia-nat==1.3.0; python_version >= '3.11'",
61
62
  "nvidia-nat-opentelemetry==1.3.0; python_version >= '3.11'",
62
- "nvidia-nat-crewai==1.3.0; python_version >= '3.11'",
63
63
  "nvidia-nat-langchain==1.3.0; python_version >= '3.11'",
64
+ "crewai>=1.1.0; python_version >= '3.11'",
64
65
  "llama-index-llms-litellm>=0.4.1,<0.7.0", # Need this to support datarobot-llm plugin
65
66
  "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
66
67
  "opentelemetry-instrumentation-llamaindex>=0.40.5,<1.0.0",
@@ -23,6 +23,7 @@ from typing import TypedDict
23
23
  from typing import TypeVar
24
24
  from typing import cast
25
25
 
26
+ from ag_ui.core import Event
26
27
  from openai.types.chat import CompletionCreateParams
27
28
  from ragas import MultiTurnSample
28
29
 
@@ -167,7 +168,7 @@ class UsageMetrics(TypedDict):
167
168
 
168
169
  # Canonical return type for DRUM-compatible invoke implementations
169
170
  InvokeReturn = (
170
- AsyncGenerator[tuple[str, MultiTurnSample | None, UsageMetrics], None]
171
+ AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]
171
172
  | tuple[str, MultiTurnSample | None, UsageMetrics]
172
173
  )
173
174
 
@@ -14,14 +14,23 @@
14
14
 
15
15
  """OpenAI-compatible response helpers for chat interactions."""
16
16
 
17
+ import asyncio
18
+ import queue
17
19
  import time
18
20
  import traceback as tb
19
21
  import uuid
20
22
  from asyncio import AbstractEventLoop
21
23
  from collections.abc import AsyncGenerator
24
+ from collections.abc import AsyncIterator
22
25
  from collections.abc import Iterator
23
26
  from concurrent.futures import ThreadPoolExecutor
27
+ from typing import Any
28
+ from typing import TypeVar
24
29
 
30
+ from ag_ui.core import BaseEvent
31
+ from ag_ui.core import Event
32
+ from ag_ui.core import TextMessageChunkEvent
33
+ from ag_ui.core import TextMessageContentEvent
25
34
  from openai.types import CompletionUsage
26
35
  from openai.types.chat import ChatCompletion
27
36
  from openai.types.chat import ChatCompletionChunk
@@ -40,6 +49,7 @@ class CustomModelChatResponse(ChatCompletion):
40
49
 
41
50
  class CustomModelStreamingResponse(ChatCompletionChunk):
42
51
  pipeline_interactions: str | None = None
52
+ event: Event | None = None
43
53
 
44
54
 
45
55
  def to_custom_model_chat_response(
@@ -83,7 +93,7 @@ def to_custom_model_streaming_response(
83
93
  thread_pool_executor: ThreadPoolExecutor,
84
94
  event_loop: AbstractEventLoop,
85
95
  streaming_response_generator: AsyncGenerator[
86
- tuple[str, MultiTurnSample | None, dict[str, int]], None
96
+ tuple[str | Event, MultiTurnSample | None, dict[str, int]], None
87
97
  ],
88
98
  model: str | object | None,
89
99
  ) -> Iterator[CustomModelStreamingResponse]:
@@ -105,7 +115,7 @@ def to_custom_model_streaming_response(
105
115
  while True:
106
116
  try:
107
117
  (
108
- response_text,
118
+ response_text_or_event,
109
119
  pipeline_interactions,
110
120
  usage_metrics,
111
121
  ) = thread_pool_executor.submit(
@@ -114,10 +124,10 @@ def to_custom_model_streaming_response(
114
124
  last_pipeline_interactions = pipeline_interactions
115
125
  last_usage_metrics = usage_metrics
116
126
 
117
- if response_text:
127
+ if isinstance(response_text_or_event, str) and response_text_or_event:
118
128
  choice = ChunkChoice(
119
129
  index=0,
120
- delta=ChoiceDelta(role="assistant", content=response_text),
130
+ delta=ChoiceDelta(role="assistant", content=response_text_or_event),
121
131
  finish_reason=None,
122
132
  )
123
133
  yield CustomModelStreamingResponse(
@@ -130,6 +140,29 @@ def to_custom_model_streaming_response(
130
140
  if usage_metrics
131
141
  else None,
132
142
  )
143
+ elif isinstance(response_text_or_event, BaseEvent):
144
+ content = ""
145
+ if isinstance(
146
+ response_text_or_event, (TextMessageContentEvent, TextMessageChunkEvent)
147
+ ):
148
+ content = response_text_or_event.delta or content
149
+ choice = ChunkChoice(
150
+ index=0,
151
+ delta=ChoiceDelta(role="assistant", content=content),
152
+ finish_reason=None,
153
+ )
154
+
155
+ yield CustomModelStreamingResponse(
156
+ id=completion_id,
157
+ object="chat.completion.chunk",
158
+ created=created,
159
+ model=model,
160
+ choices=[choice],
161
+ usage=CompletionUsage.model_validate(required_usage_metrics | usage_metrics)
162
+ if usage_metrics
163
+ else None,
164
+ event=response_text_or_event,
165
+ )
133
166
  except StopAsyncIteration:
134
167
  break
135
168
  event_loop.run_until_complete(streaming_response_generator.aclose())
@@ -168,3 +201,97 @@ def to_custom_model_streaming_response(
168
201
  choices=[choice],
169
202
  usage=None,
170
203
  )
204
+
205
+
206
+ def streaming_iterator_to_custom_model_streaming_response(
207
+ streaming_response_iterator: Iterator[tuple[str, MultiTurnSample | None, dict[str, int]]],
208
+ model: str | object | None,
209
+ ) -> Iterator[CustomModelStreamingResponse]:
210
+ """Convert the OpenAI ChatCompletionChunk response to CustomModelStreamingResponse."""
211
+ completion_id = str(uuid.uuid4())
212
+ created = int(time.time())
213
+
214
+ last_pipeline_interactions = None
215
+ last_usage_metrics = None
216
+
217
+ while True:
218
+ try:
219
+ (
220
+ response_text,
221
+ pipeline_interactions,
222
+ usage_metrics,
223
+ ) = next(streaming_response_iterator)
224
+ last_pipeline_interactions = pipeline_interactions
225
+ last_usage_metrics = usage_metrics
226
+
227
+ if response_text:
228
+ choice = ChunkChoice(
229
+ index=0,
230
+ delta=ChoiceDelta(role="assistant", content=response_text),
231
+ finish_reason=None,
232
+ )
233
+ yield CustomModelStreamingResponse(
234
+ id=completion_id,
235
+ object="chat.completion.chunk",
236
+ created=created,
237
+ model=model,
238
+ choices=[choice],
239
+ usage=CompletionUsage(**usage_metrics) if usage_metrics else None,
240
+ )
241
+ except StopIteration:
242
+ break
243
+ # Yield final chunk indicating end of stream
244
+ choice = ChunkChoice(
245
+ index=0,
246
+ delta=ChoiceDelta(role="assistant"),
247
+ finish_reason="stop",
248
+ )
249
+ yield CustomModelStreamingResponse(
250
+ id=completion_id,
251
+ object="chat.completion.chunk",
252
+ created=created,
253
+ model=model,
254
+ choices=[choice],
255
+ usage=CompletionUsage(**last_usage_metrics) if last_usage_metrics else None,
256
+ pipeline_interactions=last_pipeline_interactions.model_dump_json()
257
+ if last_pipeline_interactions
258
+ else None,
259
+ )
260
+
261
+
262
+ T = TypeVar("T")
263
+
264
+
265
+ def async_gen_to_sync_thread(
266
+ async_iterator: AsyncIterator[T],
267
+ thread_pool_executor: ThreadPoolExecutor,
268
+ event_loop: asyncio.AbstractEventLoop,
269
+ ) -> Iterator[T]:
270
+ """Run an async iterator in a separate thread and provide a sync iterator."""
271
+ # A thread-safe queue for communication
272
+ sync_queue: queue.Queue[Any] = queue.Queue()
273
+ # A sentinel object to signal the end of the async generator
274
+ SENTINEL = object() # noqa: N806
275
+
276
+ async def run_async_to_queue() -> None:
277
+ """Run in the separate thread's event loop."""
278
+ try:
279
+ async for item in async_iterator:
280
+ sync_queue.put(item)
281
+ except Exception as e:
282
+ # Put the exception on the queue to be re-raised in the main thread
283
+ sync_queue.put(e)
284
+ finally:
285
+ # Signal the end of iteration
286
+ sync_queue.put(SENTINEL)
287
+
288
+ thread_pool_executor.submit(event_loop.run_until_complete, run_async_to_queue()).result()
289
+
290
+ # The main thread consumes items synchronously
291
+ while True:
292
+ item = sync_queue.get()
293
+ if item is SENTINEL:
294
+ break
295
+ if isinstance(item, Exception):
296
+ raise item
297
+ yield item
@@ -26,7 +26,6 @@ from concurrent.futures import ThreadPoolExecutor
26
26
  from typing import Any
27
27
  from typing import Literal
28
28
 
29
- from datarobot_drum import RuntimeParameters
30
29
  from openai.types.chat import CompletionCreateParams
31
30
  from openai.types.chat.completion_create_params import CompletionCreateParamsNonStreaming
32
31
  from openai.types.chat.completion_create_params import CompletionCreateParamsStreaming
@@ -41,26 +40,6 @@ from datarobot_genai.core.telemetry_agent import instrument
41
40
  logger = logging.getLogger(__name__)
42
41
 
43
42
 
44
- def maybe_set_env_from_runtime_parameters(key: str) -> None:
45
- """Set an environment variable from a DRUM Runtime Parameter if it exists.
46
-
47
- This is safe to call outside of the DataRobot runtime. If the parameter is not available,
48
- the function does nothing.
49
- """
50
- runtime_parameter_placeholder_value = "SET_VIA_PULUMI_OR_MANUALLY"
51
- try:
52
- runtime_parameter_value = RuntimeParameters.get(key)
53
- if (
54
- runtime_parameter_value
55
- and len(runtime_parameter_value) > 0
56
- and runtime_parameter_value != runtime_parameter_placeholder_value
57
- ):
58
- os.environ[key] = runtime_parameter_value
59
- except ValueError:
60
- # Local dev: runtime parameters may be unavailable
61
- pass
62
-
63
-
64
43
  def load_model() -> tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop]:
65
44
  """Initialize a dedicated event loop within a worker thread.
66
45
 
@@ -83,7 +62,6 @@ def chat_entrypoint(
83
62
  load_model_result: tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop],
84
63
  *,
85
64
  work_dir: str | None = None,
86
- runtime_parameter_keys: list[str] | None = None,
87
65
  framework: Literal["crewai", "langgraph", "llamaindex", "nat"] | None = None,
88
66
  **kwargs: Any,
89
67
  ) -> CustomModelChatResponse | Iterator[CustomModelStreamingResponse]:
@@ -103,10 +81,6 @@ def chat_entrypoint(
103
81
  work_dir : Optional[str]
104
82
  Working directory to ``chdir`` into before invoking the agent. This is useful
105
83
  when relative paths are used in agent templates.
106
- runtime_parameter_keys : Optional[List[str]]
107
- Runtime parameter keys (DataRobot custom model) to propagate into env. When
108
- ``None``, defaults to
109
- ``['EXTERNAL_MCP_URL', 'MCP_DEPLOYMENT_ID']``.
110
84
  framework : Optional[Literal["crewai", "langgraph", "llamaindex", "nat"]]
111
85
  When provided, idempotently instruments HTTP clients, OpenAI SDK, and the
112
86
  given framework. If omitted, general instrumentation is still applied.
@@ -129,12 +103,6 @@ def chat_entrypoint(
129
103
  except Exception as e:
130
104
  logger.warning(f"Failed to change working directory to {work_dir}: {e}")
131
105
 
132
- # Load MCP runtime parameters and session secret if configured
133
- if runtime_parameter_keys is None:
134
- runtime_parameter_keys = ["EXTERNAL_MCP_URL", "MCP_DEPLOYMENT_ID"]
135
- for key in runtime_parameter_keys:
136
- maybe_set_env_from_runtime_parameters(key)
137
-
138
106
  # Retrieve authorization context using all supported methods for downstream agents/tools
139
107
  completion_create_params["authorization_context"] = resolve_authorization_context(
140
108
  completion_create_params, **kwargs
@@ -74,7 +74,22 @@ class AuthContextHeaderHandler:
74
74
  if algorithm is None:
75
75
  raise ValueError("Algorithm None is not allowed. Use a secure algorithm like HS256.")
76
76
 
77
- self.secret_key = secret_key or AuthContextConfig().session_secret_key
77
+ # Get secret key from parameter, config, or environment variable
78
+ # Handle the case where AuthContextConfig() initialization fails due to
79
+ # a bug in the datarobot package when SESSION_SECRET_KEY is not set
80
+ if secret_key:
81
+ self.secret_key = secret_key
82
+ else:
83
+ try:
84
+ config = AuthContextConfig()
85
+ self.secret_key = config.session_secret_key or ""
86
+ except (TypeError, AttributeError, Exception):
87
+ # Fallback to reading environment variable directly if config initialization fails
88
+ # This can happen when SESSION_SECRET_KEY is not set and the datarobot package's
89
+ # getenv function encounters a bug with None values
90
+ # it tries to check if "apiToken" in payload: when payload is None
91
+ self.secret_key = ""
92
+
78
93
  self.algorithm = algorithm
79
94
  self.validate_signature = validate_signature
80
95
 
@@ -17,6 +17,15 @@ from collections.abc import AsyncGenerator
17
17
  from typing import Any
18
18
  from typing import cast
19
19
 
20
+ from ag_ui.core import Event
21
+ from ag_ui.core import EventType
22
+ from ag_ui.core import TextMessageContentEvent
23
+ from ag_ui.core import TextMessageEndEvent
24
+ from ag_ui.core import TextMessageStartEvent
25
+ from ag_ui.core import ToolCallArgsEvent
26
+ from ag_ui.core import ToolCallEndEvent
27
+ from ag_ui.core import ToolCallResultEvent
28
+ from ag_ui.core import ToolCallStartEvent
20
29
  from langchain.tools import BaseTool
21
30
  from langchain_core.messages import AIMessageChunk
22
31
  from langchain_core.messages import ToolMessage
@@ -158,43 +167,7 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
158
167
  # The main difference is returning a generator for streaming or a final response for sync.
159
168
  if is_streaming(completion_create_params):
160
169
  # Streaming response: yield each message as it is generated
161
- async def stream_generator() -> AsyncGenerator[
162
- tuple[str, MultiTurnSample | None, UsageMetrics], None
163
- ]:
164
- # Iterate over the graph stream. For message events, yield the content.
165
- # For update events, accumulate the usage metrics.
166
- events = []
167
- async for _, mode, event in graph_stream:
168
- if mode == "messages":
169
- message_event: tuple[AIMessageChunk, dict[str, Any]] = event # type: ignore[assignment]
170
- llm_token, _ = message_event
171
- yield (
172
- str(llm_token.content),
173
- None,
174
- usage_metrics,
175
- )
176
- elif mode == "updates":
177
- update_event: dict[str, Any] = event # type: ignore[assignment]
178
- events.append(update_event)
179
- current_node = next(iter(update_event))
180
- node_data = update_event[current_node]
181
- current_usage = node_data.get("usage", {}) if node_data is not None else {}
182
- if current_usage:
183
- usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
184
- usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
185
- usage_metrics["completion_tokens"] += current_usage.get(
186
- "completion_tokens", 0
187
- )
188
- else:
189
- raise ValueError(f"Invalid mode: {mode}")
190
-
191
- # Create a list of events from the event listener
192
- pipeline_interactions = self.create_pipeline_interactions_from_events(events)
193
-
194
- # yield the final response indicating completion
195
- yield "", pipeline_interactions, usage_metrics
196
-
197
- return stream_generator()
170
+ return self._stream_generator(graph_stream, usage_metrics)
198
171
  else:
199
172
  # Synchronous response: collect all events and return the final message
200
173
  events: list[dict[str, Any]] = [
@@ -203,6 +176,16 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
203
176
  if mode == "updates"
204
177
  ]
205
178
 
179
+ # Accumulate the usage metrics from the updates
180
+ for update in events:
181
+ current_node = next(iter(update))
182
+ node_data = update[current_node]
183
+ current_usage = node_data.get("usage", {}) if node_data is not None else {}
184
+ if current_usage:
185
+ usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
186
+ usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
187
+ usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
188
+
206
189
  pipeline_interactions = self.create_pipeline_interactions_from_events(events)
207
190
 
208
191
  # Extract the final event from the graph stream as the synchronous response
@@ -214,14 +197,132 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
214
197
  if node_data is not None and "messages" in node_data
215
198
  else ""
216
199
  )
217
- current_usage = node_data.get("usage", {}) if node_data is not None else {}
218
- if current_usage:
219
- usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
220
- usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
221
- usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
222
200
 
223
201
  return response_text, pipeline_interactions, usage_metrics
224
202
 
203
+ async def _stream_generator(
204
+ self, graph_stream: AsyncGenerator[tuple[Any, str, Any], None], usage_metrics: UsageMetrics
205
+ ) -> AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]:
206
+ # Iterate over the graph stream. For message events, yield the content.
207
+ # For update events, accumulate the usage metrics.
208
+ events = []
209
+ current_message_id = None
210
+ tool_call_id = ""
211
+ async for _, mode, event in graph_stream:
212
+ if mode == "messages":
213
+ message_event: tuple[AIMessageChunk | ToolMessage, dict[str, Any]] = event # type: ignore[assignment]
214
+ message = message_event[0]
215
+ if isinstance(message, ToolMessage):
216
+ yield (
217
+ ToolCallEndEvent(
218
+ type=EventType.TOOL_CALL_END, tool_call_id=message.tool_call_id
219
+ ),
220
+ None,
221
+ usage_metrics,
222
+ )
223
+ yield (
224
+ ToolCallResultEvent(
225
+ type=EventType.TOOL_CALL_RESULT,
226
+ message_id=message.id,
227
+ tool_call_id=message.tool_call_id,
228
+ content=message.content,
229
+ role="tool",
230
+ ),
231
+ None,
232
+ usage_metrics,
233
+ )
234
+ tool_call_id = ""
235
+ elif isinstance(message, AIMessageChunk):
236
+ if message.tool_call_chunks:
237
+ # This is a tool call message
238
+ for tool_call_chunk in message.tool_call_chunks:
239
+ if name := tool_call_chunk.get("name"):
240
+ # Its a tool call start message
241
+ tool_call_id = tool_call_chunk["id"]
242
+ yield (
243
+ ToolCallStartEvent(
244
+ type=EventType.TOOL_CALL_START,
245
+ tool_call_id=tool_call_id,
246
+ tool_call_name=name,
247
+ parent_message_id=message.id,
248
+ ),
249
+ None,
250
+ usage_metrics,
251
+ )
252
+ elif args := tool_call_chunk.get("args"):
253
+ # Its a tool call args message
254
+ yield (
255
+ ToolCallArgsEvent(
256
+ type=EventType.TOOL_CALL_ARGS,
257
+ # Its empty when the tool chunk is not a start message
258
+ # So we use the tool call id from a previous start message
259
+ tool_call_id=tool_call_id,
260
+ delta=args,
261
+ ),
262
+ None,
263
+ usage_metrics,
264
+ )
265
+ elif message.content:
266
+ # Its a text message
267
+ # Handle the start and end of the text message
268
+ if message.id != current_message_id:
269
+ if current_message_id:
270
+ yield (
271
+ TextMessageEndEvent(
272
+ type=EventType.TEXT_MESSAGE_END,
273
+ message_id=current_message_id,
274
+ ),
275
+ None,
276
+ usage_metrics,
277
+ )
278
+ current_message_id = message.id
279
+ yield (
280
+ TextMessageStartEvent(
281
+ type=EventType.TEXT_MESSAGE_START,
282
+ message_id=message.id,
283
+ role="assistant",
284
+ ),
285
+ None,
286
+ usage_metrics,
287
+ )
288
+ yield (
289
+ TextMessageContentEvent(
290
+ type=EventType.TEXT_MESSAGE_CONTENT,
291
+ message_id=message.id,
292
+ delta=message.content,
293
+ ),
294
+ None,
295
+ usage_metrics,
296
+ )
297
+ else:
298
+ raise ValueError(f"Invalid message event: {message_event}")
299
+ elif mode == "updates":
300
+ update_event: dict[str, Any] = event # type: ignore[assignment]
301
+ events.append(update_event)
302
+ current_node = next(iter(update_event))
303
+ node_data = update_event[current_node]
304
+ current_usage = node_data.get("usage", {}) if node_data is not None else {}
305
+ if current_usage:
306
+ usage_metrics["total_tokens"] += current_usage.get("total_tokens", 0)
307
+ usage_metrics["prompt_tokens"] += current_usage.get("prompt_tokens", 0)
308
+ usage_metrics["completion_tokens"] += current_usage.get("completion_tokens", 0)
309
+ if current_message_id:
310
+ yield (
311
+ TextMessageEndEvent(
312
+ type=EventType.TEXT_MESSAGE_END,
313
+ message_id=current_message_id,
314
+ ),
315
+ None,
316
+ usage_metrics,
317
+ )
318
+ current_message_id = None
319
+
320
+ # Create a list of events from the event listener
321
+ pipeline_interactions = self.create_pipeline_interactions_from_events(events)
322
+
323
+ # yield the final response indicating completion
324
+ yield "", pipeline_interactions, usage_metrics
325
+
225
326
  @classmethod
226
327
  def create_pipeline_interactions_from_events(
227
328
  cls,