datarobot-genai 0.1.70__tar.gz → 0.1.72__tar.gz

{datarobot_genai-0.1.70 → datarobot_genai-0.1.72}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.1.70
+Version: 0.1.72
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -23,7 +23,7 @@ Requires-Dist: ragas<0.4.0,>=0.3.8
 Requires-Dist: requests<3.0.0,>=2.32.4
 Provides-Extra: crewai
 Requires-Dist: crewai-tools[mcp]<0.77.0,>=0.69.0; extra == 'crewai'
-Requires-Dist: crewai<1.0.0,>=0.193.2; extra == 'crewai'
+Requires-Dist: crewai>=1.1.0; extra == 'crewai'
 Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'crewai'
 Requires-Dist: pybase64<2.0.0,>=1.4.2; extra == 'crewai'
 Provides-Extra: drmcp
@@ -56,8 +56,8 @@ Requires-Dist: llama-index<0.14.0,>=0.13.6; extra == 'llamaindex'
 Requires-Dist: opentelemetry-instrumentation-llamaindex<1.0.0,>=0.40.5; extra == 'llamaindex'
 Requires-Dist: pypdf<7.0.0,>=6.0.0; extra == 'llamaindex'
 Provides-Extra: nat
+Requires-Dist: crewai>=1.1.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: llama-index-llms-litellm<0.7.0,>=0.4.1; extra == 'nat'
-Requires-Dist: nvidia-nat-crewai==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat-langchain==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat-opentelemetry==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat==1.3.0; (python_version >= '3.11') and extra == 'nat'

{datarobot_genai-0.1.70 → datarobot_genai-0.1.72}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.1.70"
+version = "0.1.72"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -35,7 +35,7 @@ datarobot_llm_clients = "datarobot_genai.nat.datarobot_llm_clients"
 
 [project.optional-dependencies]
 crewai = [
-  "crewai>=0.193.2,<1.0.0",
+  "crewai>=1.1.0",
   "crewai-tools[mcp]>=0.69.0,<0.77.0",
   "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
   "pybase64>=1.4.2,<2.0.0",
@@ -59,8 +59,8 @@ llamaindex = [
 nat = [
   "nvidia-nat==1.3.0; python_version >= '3.11'",
   "nvidia-nat-opentelemetry==1.3.0; python_version >= '3.11'",
-  "nvidia-nat-crewai==1.3.0; python_version >= '3.11'",
   "nvidia-nat-langchain==1.3.0; python_version >= '3.11'",
+  "crewai>=1.1.0; python_version >= '3.11'",
   "llama-index-llms-litellm>=0.4.1,<0.7.0",  # Need this to support datarobot-llm plugin
   "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",
   "opentelemetry-instrumentation-llamaindex>=0.40.5,<1.0.0",

{datarobot_genai-0.1.70 → datarobot_genai-0.1.72}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py

@@ -38,7 +38,7 @@ def _apply_green(text: str) -> str:
     return "\n".join(colored_lines)
 
 
-DR_LOGO_ASCII = _apply_green("""\
+DR_LOGO_ASCII = _apply_green(r"""
  ____        _        ____       _           _   
 |  _ \  __ _| |_ __ _|  _ \ ___ | |__   ___ | |_ 
 | | | |/ _` | __/ _` | |_) / _ \| '_ \ / _ \| __|

{datarobot_genai-0.1.70 → datarobot_genai-0.1.72}/src/datarobot_genai/langgraph/agent.py

@@ -84,21 +84,38 @@ class LangGraphAgent(BaseAgent[BaseTool], abc.ABC):
             async def wrapped_generator() -> AsyncGenerator[
                 tuple[str, Any | None, UsageMetrics], None
             ]:
-                async with mcp_tools_context(
-                    authorization_context=self._authorization_context,
-                    forwarded_headers=self.forwarded_headers,
-                ) as mcp_tools:
-                    self.set_mcp_tools(mcp_tools)
-                    result = await self._invoke(completion_create_params)
-
-                    # Yield all items from the result generator
-                    # The context will be closed when this generator is exhausted
-                    # Cast to async generator since we know stream=True means it's a generator
-                    result_generator = cast(
-                        AsyncGenerator[tuple[str, Any | None, UsageMetrics], None], result
-                    )
-                    async for item in result_generator:
-                        yield item
+                try:
+                    async with mcp_tools_context(
+                        authorization_context=self._authorization_context,
+                        forwarded_headers=self.forwarded_headers,
+                    ) as mcp_tools:
+                        self.set_mcp_tools(mcp_tools)
+                        result = await self._invoke(completion_create_params)
+
+                        # Yield all items from the result generator
+                        # The context will be closed when this generator is exhausted
+                        # Cast to async generator since we know stream=True means it's a generator
+                        result_generator = cast(
+                            AsyncGenerator[tuple[str, Any | None, UsageMetrics], None], result
+                        )
+                        async for item in result_generator:
+                            yield item
+                except RuntimeError as e:
+                    error_message = str(e).lower()
+                    if "different task" in error_message and "cancel scope" in error_message:
+                        # Due to anyio task group constraints when consuming async generators
+                        # across task boundaries, we cannot always clean up properly.
+                        # The underlying HTTP client/connection pool should handle resource cleanup
+                        # via timeouts and connection pooling, but this
+                        # may lead to delayed resource release.
+                        logger.debug(
+                            "MCP context cleanup attempted in different task. "
+                            "This is a limitation when consuming async generators "
+                            "across task boundaries."
+                        )
+                    else:
+                        # Re-raise if it's a different RuntimeError
+                        raise
 
             return wrapped_generator()
         else:

{datarobot_genai-0.1.70 → datarobot_genai-0.1.72}/src/datarobot_genai/nat/datarobot_llm_clients.py

@@ -23,6 +23,7 @@ from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.register_workflow import register_llm_client
 
+from ..nat.datarobot_llm_providers import DataRobotLLMComponentModelConfig
 from ..nat.datarobot_llm_providers import DataRobotLLMDeploymentModelConfig
 from ..nat.datarobot_llm_providers import DataRobotLLMGatewayModelConfig
 from ..nat.datarobot_llm_providers import DataRobotNIMModelConfig
@@ -75,6 +76,7 @@ async def datarobot_llm_gateway_langchain(
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
     config["base_url"] = config["base_url"] + "/genai/llmgw"
     config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
     yield DataRobotChatOpenAI(**config)
 
 
@@ -85,7 +87,8 @@ async def datarobot_llm_gateway_crewai(
     llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
 ) -> AsyncGenerator[LLM]:
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["base_url"] = config["base_url"].removesuffix("/api/v2")
     yield LLM(**config)
 
@@ -97,7 +100,8 @@ async def datarobot_llm_gateway_llamaindex(
     llm_config: DataRobotLLMGatewayModelConfig, builder: Builder
 ) -> AsyncGenerator[LLM]:
     config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url").removesuffix("/api/v2")
     yield DataRobotLiteLLM(**config)
 
@@ -109,11 +113,12 @@ async def datarobot_llm_deployment_langchain(
     llm_config: DataRobotLLMDeploymentModelConfig, builder: Builder
 ) -> AsyncGenerator[ChatOpenAI]:
     config = llm_config.model_dump(
-        exclude={"type", "thinking", "datarobot_endpoint", "llm_deployment_id"},
+        exclude={"type", "thinking"},
         by_alias=True,
         exclude_none=True,
     )
     config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
     yield DataRobotChatOpenAI(**config)
 
 
@@ -128,7 +133,8 @@ async def datarobot_llm_deployment_crewai(
         by_alias=True,
         exclude_none=True,
     )
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
     yield LLM(**config)
 
@@ -144,7 +150,8 @@ async def datarobot_llm_deployment_llamaindex(
         by_alias=True,
         exclude_none=True,
     )
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
     yield DataRobotLiteLLM(**config)
 
@@ -159,6 +166,7 @@ async def datarobot_nim_langchain(
         exclude_none=True,
     )
     config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
     yield DataRobotChatOpenAI(**config)
 
 
@@ -171,7 +179,8 @@ async def datarobot_nim_crewai(
         by_alias=True,
         exclude_none=True,
     )
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
     yield LLM(**config)
 
@@ -185,6 +194,56 @@ async def datarobot_nim_llamaindex(
         by_alias=True,
         exclude_none=True,
     )
-    config["model"] = "datarobot/" + config["model"]
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
     config["api_base"] = config.pop("base_url") + "/chat/completions"
     yield DataRobotLiteLLM(**config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN
+)
+async def datarobot_llm_component_langchain(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[ChatOpenAI]:
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if config["use_datarobot_llm_gateway"]:
+        config["base_url"] = config["base_url"] + "/genai/llmgw"
+    config["stream_options"] = {"include_usage": True}
+    config["model"] = config["model"].removeprefix("datarobot/")
+    config.pop("use_datarobot_llm_gateway")
+    yield DataRobotChatOpenAI(**config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.CREWAI
+)
+async def datarobot_llm_component_crewai(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[LLM]:
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    if config["use_datarobot_llm_gateway"]:
+        config["base_url"] = config["base_url"].removesuffix("/api/v2")
+    else:
+        config["api_base"] = config.pop("base_url") + "/chat/completions"
+    config.pop("use_datarobot_llm_gateway")
+    yield LLM(**config)
+
+
+@register_llm_client(
+    config_type=DataRobotLLMComponentModelConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX
+)
+async def datarobot_llm_component_llamaindex(
+    llm_config: DataRobotLLMComponentModelConfig, builder: Builder
+) -> AsyncGenerator[LLM]:
+    config = llm_config.model_dump(exclude={"type", "thinking"}, by_alias=True, exclude_none=True)
+    if not config["model"].startswith("datarobot/"):
+        config["model"] = "datarobot/" + config["model"]
+    if config["use_datarobot_llm_gateway"]:
+        config["api_base"] = config.pop("base_url").removesuffix("/api/v2")
+    else:
+        config["api_base"] = config.pop("base_url") + "/chat/completions"
+    config.pop("use_datarobot_llm_gateway")
+    yield DataRobotLiteLLM(**config)

{datarobot_genai-0.1.70 → datarobot_genai-0.1.72}/src/datarobot_genai/nat/datarobot_llm_providers.py

@@ -32,11 +32,43 @@ class Config(DataRobotAppFrameworkBaseSettings):
     datarobot_api_token: str | None = None
     llm_deployment_id: str | None = None
     nim_deployment_id: str | None = None
+    use_datarobot_llm_gateway: bool = False
+    llm_default_model: str | None = None
 
 
 config = Config()
 
 
+class DataRobotLLMComponentModelConfig(OpenAIModelConfig, name="datarobot-llm-component"):  # type: ignore[call-arg]
+    """A DataRobot LLM provider to be used with an LLM client."""
+
+    api_key: str | None = Field(
+        default=config.datarobot_api_token, description="DataRobot API key."
+    )
+    base_url: str | None = Field(
+        default=config.datarobot_endpoint.rstrip("/")
+        if config.use_datarobot_llm_gateway
+        else config.datarobot_endpoint + f"/deployments/{config.llm_deployment_id}",
+        description="DataRobot LLM URL.",
+    )
+    model_name: str = Field(
+        validation_alias=AliasChoices("model_name", "model"),
+        serialization_alias="model",
+        description="The model name.",
+        default=config.llm_default_model or "datarobot-deployed-llm",
+    )
+    use_datarobot_llm_gateway: bool = config.use_datarobot_llm_gateway
+
+
+@register_llm_provider(config_type=DataRobotLLMComponentModelConfig)
+async def datarobot_llm_component(
+    config: DataRobotLLMComponentModelConfig, _builder: Builder
+) -> LLMProviderInfo:
+    yield LLMProviderInfo(
+        config=config, description="DataRobot LLM Component for use with an LLM client."
+    )
+
+
 class DataRobotLLMGatewayModelConfig(OpenAIModelConfig, name="datarobot-llm-gateway"):  # type: ignore[call-arg]
     """A DataRobot LLM provider to be used with an LLM client."""