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

datarobot_genai/core/chat/responses.py

@@ -14,13 +14,18 @@
 
 """OpenAI-compatible response helpers for chat interactions."""
 
+import asyncio
+import queue
 import time
 import traceback as tb
 import uuid
 from asyncio import AbstractEventLoop
 from collections.abc import AsyncGenerator
+from collections.abc import AsyncIterator
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+from typing import TypeVar
 
 from openai.types import CompletionUsage
 from openai.types.chat import ChatCompletion
@@ -168,3 +173,97 @@ def to_custom_model_streaming_response(
             choices=[choice],
             usage=None,
         )
+
+
+def streaming_iterator_to_custom_model_streaming_response(
+    streaming_response_iterator: Iterator[tuple[str, MultiTurnSample | None, dict[str, int]]],
+    model: str | object | None,
+) -> Iterator[CustomModelStreamingResponse]:
+    """Convert the OpenAI ChatCompletionChunk response to CustomModelStreamingResponse."""
+    completion_id = str(uuid.uuid4())
+    created = int(time.time())
+
+    last_pipeline_interactions = None
+    last_usage_metrics = None
+
+    while True:
+        try:
+            (
+                response_text,
+                pipeline_interactions,
+                usage_metrics,
+            ) = next(streaming_response_iterator)
+            last_pipeline_interactions = pipeline_interactions
+            last_usage_metrics = usage_metrics
+
+            if response_text:
+                choice = ChunkChoice(
+                    index=0,
+                    delta=ChoiceDelta(role="assistant", content=response_text),
+                    finish_reason=None,
+                )
+                yield CustomModelStreamingResponse(
+                    id=completion_id,
+                    object="chat.completion.chunk",
+                    created=created,
+                    model=model,
+                    choices=[choice],
+                    usage=CompletionUsage(**usage_metrics) if usage_metrics else None,
+                )
+        except StopIteration:
+            break
+    # Yield final chunk indicating end of stream
+    choice = ChunkChoice(
+        index=0,
+        delta=ChoiceDelta(role="assistant"),
+        finish_reason="stop",
+    )
+    yield CustomModelStreamingResponse(
+        id=completion_id,
+        object="chat.completion.chunk",
+        created=created,
+        model=model,
+        choices=[choice],
+        usage=CompletionUsage(**last_usage_metrics) if last_usage_metrics else None,
+        pipeline_interactions=last_pipeline_interactions.model_dump_json()
+        if last_pipeline_interactions
+        else None,
+    )
+
+
+T = TypeVar("T")
+
+
+def async_gen_to_sync_thread(
+    async_iterator: AsyncIterator[T],
+    thread_pool_executor: ThreadPoolExecutor,
+    event_loop: asyncio.AbstractEventLoop,
+) -> Iterator[T]:
+    """Run an async iterator in a separate thread and provide a sync iterator."""
+    # A thread-safe queue for communication
+    sync_queue: queue.Queue[Any] = queue.Queue()
+    # A sentinel object to signal the end of the async generator
+    SENTINEL = object()  # noqa: N806
+
+    async def run_async_to_queue() -> None:
+        """Run in the separate thread's event loop."""
+        try:
+            async for item in async_iterator:
+                sync_queue.put(item)
+        except Exception as e:
+            # Put the exception on the queue to be re-raised in the main thread
+            sync_queue.put(e)
+        finally:
+            # Signal the end of iteration
+            sync_queue.put(SENTINEL)
+
+    thread_pool_executor.submit(event_loop.run_until_complete, run_async_to_queue()).result()
+
+    # The main thread consumes items synchronously
+    while True:
+        item = sync_queue.get()
+        if item is SENTINEL:
+            break
+        if isinstance(item, Exception):
+            raise item
+        yield item

datarobot_genai/core/custom_model.py

@@ -26,7 +26,6 @@ from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 from typing import Literal
 
-from datarobot_drum import RuntimeParameters
 from openai.types.chat import CompletionCreateParams
 from openai.types.chat.completion_create_params import CompletionCreateParamsNonStreaming
 from openai.types.chat.completion_create_params import CompletionCreateParamsStreaming
@@ -41,26 +40,6 @@ from datarobot_genai.core.telemetry_agent import instrument
 logger = logging.getLogger(__name__)
 
 
-def maybe_set_env_from_runtime_parameters(key: str) -> None:
-    """Set an environment variable from a DRUM Runtime Parameter if it exists.
-
-    This is safe to call outside of the DataRobot runtime. If the parameter is not available,
-    the function does nothing.
-    """
-    runtime_parameter_placeholder_value = "SET_VIA_PULUMI_OR_MANUALLY"
-    try:
-        runtime_parameter_value = RuntimeParameters.get(key)
-        if (
-            runtime_parameter_value
-            and len(runtime_parameter_value) > 0
-            and runtime_parameter_value != runtime_parameter_placeholder_value
-        ):
-            os.environ[key] = runtime_parameter_value
-    except ValueError:
-        # Local dev: runtime parameters may be unavailable
-        pass
-
-
 def load_model() -> tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop]:
     """Initialize a dedicated event loop within a worker thread.
 
@@ -83,7 +62,6 @@ def chat_entrypoint(
     load_model_result: tuple[ThreadPoolExecutor, asyncio.AbstractEventLoop],
     *,
     work_dir: str | None = None,
-    runtime_parameter_keys: list[str] | None = None,
     framework: Literal["crewai", "langgraph", "llamaindex", "nat"] | None = None,
     **kwargs: Any,
 ) -> CustomModelChatResponse | Iterator[CustomModelStreamingResponse]:
@@ -103,10 +81,6 @@ def chat_entrypoint(
     work_dir : Optional[str]
         Working directory to ``chdir`` into before invoking the agent. This is useful
         when relative paths are used in agent templates.
-    runtime_parameter_keys : Optional[List[str]]
-        Runtime parameter keys (DataRobot custom model) to propagate into env. When
-        ``None``, defaults to
-        ``['EXTERNAL_MCP_URL', 'MCP_DEPLOYMENT_ID']``.
     framework : Optional[Literal["crewai", "langgraph", "llamaindex", "nat"]]
         When provided, idempotently instruments HTTP clients, OpenAI SDK, and the
         given framework. If omitted, general instrumentation is still applied.
@@ -129,12 +103,6 @@ def chat_entrypoint(
         except Exception as e:
             logger.warning(f"Failed to change working directory to {work_dir}: {e}")
 
-    # Load MCP runtime parameters and session secret if configured
-    if runtime_parameter_keys is None:
-        runtime_parameter_keys = ["EXTERNAL_MCP_URL", "MCP_DEPLOYMENT_ID"]
-    for key in runtime_parameter_keys:
-        maybe_set_env_from_runtime_parameters(key)
-
     # Retrieve authorization context using all supported methods for downstream agents/tools
     completion_create_params["authorization_context"] = resolve_authorization_context(
         completion_create_params, **kwargs

datarobot_genai/core/utils/auth.py

@@ -74,7 +74,22 @@ class AuthContextHeaderHandler:
         if algorithm is None:
             raise ValueError("Algorithm None is not allowed. Use a secure algorithm like HS256.")
 
-        self.secret_key = secret_key or AuthContextConfig().session_secret_key
+        # Get secret key from parameter, config, or environment variable
+        # Handle the case where AuthContextConfig() initialization fails due to
+        # a bug in the datarobot package when SESSION_SECRET_KEY is not set
+        if secret_key:
+            self.secret_key = secret_key
+        else:
+            try:
+                config = AuthContextConfig()
+                self.secret_key = config.session_secret_key or ""
+            except (TypeError, AttributeError, Exception):
+                # Fallback to reading environment variable directly if config initialization fails
+                # This can happen when SESSION_SECRET_KEY is not set and the datarobot package's
+                # getenv function encounters a bug with None values
+                # it tries to check if "apiToken" in payload: when payload is None
+                self.secret_key = ""
+
         self.algorithm = algorithm
         self.validate_signature = validate_signature
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.72
+Version: 0.1.75
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.

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

@@ -1,21 +1,21 @@
 datarobot_genai/__init__.py,sha256=QwWCADIZigzGvew8tGT9rAjgQjchSMjaZtFPbi3Bt2s,527
 datarobot_genai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datarobot_genai/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/core/custom_model.py,sha256=8SMW3BAToI44331pF6fATQlS-FulDXop1OmmrVxuB1U,7155
+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/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=DEMkSLyDKeCWl_RhWxIfncWWCDrKIZr6HwRkVWC6n3c,5911
+datarobot_genai/core/chat/responses.py,sha256=CCkeRpoG7RDI0-JBFhfSUVNaUOspTeH93q3YTUHpZ4A,9242
 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
 datarobot_genai/core/mcp/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datarobot_genai/core/mcp/common.py,sha256=Y8SjuquUODKEfI7T9X-QuTMKdIlpCWFI1b3xs6tmHFA,7812
 datarobot_genai/core/utils/__init__.py,sha256=VxtRUz6iwb04eFQQy0zqTNXLAkYpPXcJxVoKV0nOdXk,59
-datarobot_genai/core/utils/auth.py,sha256=Xo1PxVr6oMgtMHkmHdS02klDKK1cyDpjGvIMF4Tx0Lo,7874
+datarobot_genai/core/utils/auth.py,sha256=LpSoHdPD2siskYwG8q4f9cike4VQdgFWJpuJrpiszXU,8674
 datarobot_genai/core/utils/urls.py,sha256=tk0t13duDEPcmwz2OnS4vwEdatruiuX8lnxMMhSaJik,2289
 datarobot_genai/crewai/__init__.py,sha256=MtFnHA3EtmgiK_GjwUGPgQQ6G1MCEzz1SDBwQi9lE8M,706
 datarobot_genai/crewai/agent.py,sha256=vp8_2LExpeLls7Fpzo0R6ud5I6Ryfu3n3oVTN4Yyi6A,1417
@@ -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.72.dist-info/METADATA,sha256=NKR8UuyddfZt7nhQslbYM-2G7pnebcE47j3_kI7YjBQ,5898
-datarobot_genai-0.1.72.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.1.72.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
-datarobot_genai-0.1.72.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.1.72.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.1.72.dist-info/RECORD,,
+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,,