datarobot-genai 0.2.41__tar.gz → 0.2.44__tar.gz

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.41
+Version: 0.2.44
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -46,6 +46,7 @@ Requires-Dist: opentelemetry-sdk<2.0.0,>=1.22.0; extra == 'drmcp'
 Requires-Dist: pydantic-settings<3.0.0,>=2.1.0; extra == 'drmcp'
 Requires-Dist: pydantic<3.0.0,>=2.6.1; extra == 'drmcp'
 Requires-Dist: python-dotenv<2.0.0,>=1.1.0; extra == 'drmcp'
+Requires-Dist: tavily-python<1.0.0,>=0.7.20; extra == 'drmcp'
 Provides-Extra: langgraph
 Requires-Dist: langchain-mcp-adapters<0.2.0,>=0.1.12; extra == 'langgraph'
 Requires-Dist: langgraph-prebuilt<0.7.0,>=0.2.3; extra == 'langgraph'

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.41"
+version = "0.2.44"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -84,6 +84,7 @@ drmcp = [
   "python-dotenv>=1.1.0,<2.0.0",
   "boto3>=1.34.0,<2.0.0",
   "httpx>=0.28.1,<1.0.0",
+  "tavily-python>=0.7.20,<1.0.0",
   "pydantic>=2.6.1,<3.0.0",
   "pydantic-settings>=2.1.0,<3.0.0",
   "opentelemetry-api>=1.22.0,<2.0.0",
@@ -187,6 +188,9 @@ fixable = [
     "I001",   # isort import order within functions
     "PLC0415",# pylint import-outside-toplevel for controlled local imports
 ]
+"src/datarobot_genai/**/*.py" = [
+    "PLC0415",# pylint import-outside-toplevel: allow lazy imports for memory optimization
+]
 
 [tool.ruff.lint.isort]
 force-single-line = true

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/core/agents/base.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import abc
 import json
 import os
@@ -19,6 +21,7 @@ from collections.abc import AsyncGenerator
 from collections.abc import Mapping
 from typing import Any
 from typing import Generic
+from typing import TypeAlias
 from typing import TypedDict
 from typing import TypeVar
 from typing import cast
@@ -98,7 +101,7 @@ class BaseAgent(Generic[TTool], abc.ABC):
         return get_api_base(self.api_base, deployment_id)
 
     @abc.abstractmethod
-    async def invoke(self, completion_create_params: CompletionCreateParams) -> "InvokeReturn":
+    async def invoke(self, completion_create_params: CompletionCreateParams) -> InvokeReturn:
         raise NotImplementedError("Not implemented")
 
     @classmethod
@@ -167,7 +170,7 @@ class UsageMetrics(TypedDict):
 
 
 # Canonical return type for DRUM-compatible invoke implementations
-InvokeReturn = (
+InvokeReturn: TypeAlias = (
     AsyncGenerator[tuple[str | Event, MultiTurnSample | None, UsageMetrics], None]
     | tuple[str, MultiTurnSample | None, UsageMetrics]
 )

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/core/chat/responses.py

@@ -14,6 +14,8 @@
 
 """OpenAI-compatible response helpers for chat interactions."""
 
+from __future__ import annotations
+
 import asyncio
 import queue
 import time
@@ -24,6 +26,7 @@ from collections.abc import AsyncGenerator
 from collections.abc import AsyncIterator
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import TypeVar
 
@@ -38,10 +41,12 @@ from openai.types.chat import ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
 from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
 from openai.types.chat.chat_completion_chunk import ChoiceDelta
-from ragas import MultiTurnSample
 
 from datarobot_genai.core.agents import default_usage_metrics
 
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+
 
 class CustomModelChatResponse(ChatCompletion):
     pipeline_interactions: str | None = None

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/crewai/agent.py

@@ -11,15 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from crewai import LLM
-from ragas import MultiTurnSample
-from ragas.messages import AIMessage
-from ragas.messages import HumanMessage
-from ragas.messages import ToolMessage
 
 from datarobot_genai.core.utils.urls import get_api_base
 
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+    from ragas.messages import AIMessage
+    from ragas.messages import HumanMessage
+    from ragas.messages import ToolMessage
+
 
 def build_llm(
     *,
@@ -39,4 +44,7 @@ def create_pipeline_interactions_from_messages(
 ) -> MultiTurnSample | None:
     if not messages:
         return None
+    # Lazy import to reduce memory overhead when ragas is not used
+    from ragas import MultiTurnSample
+
     return MultiTurnSample(user_input=messages)

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/crewai/base.py

@@ -26,13 +26,13 @@ from __future__ import annotations
 import abc
 import asyncio
 from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
 from typing import Any
 
 from crewai import Crew
 from crewai.events.event_bus import CrewAIEventsBus
 from crewai.tools import BaseTool
 from openai.types.chat import CompletionCreateParams
-from ragas import MultiTurnSample
 
 from datarobot_genai.core.agents.base import BaseAgent
 from datarobot_genai.core.agents.base import InvokeReturn
@@ -44,6 +44,9 @@ from datarobot_genai.core.agents.base import is_streaming
 from .agent import create_pipeline_interactions_from_messages
 from .mcp import mcp_tools_context
 
+if TYPE_CHECKING:
+    from ragas import MultiTurnSample
+
 
 class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
     """Abstract base agent for CrewAI workflows.

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/crewai/events.py

@@ -16,12 +16,13 @@ from __future__ import annotations
 import importlib
 import json
 import logging
+from typing import TYPE_CHECKING
 from typing import Any
 
-from ragas.messages import AIMessage
-from ragas.messages import HumanMessage
-from ragas.messages import ToolCall
-from ragas.messages import ToolMessage
+if TYPE_CHECKING:
+    from ragas.messages import AIMessage
+    from ragas.messages import HumanMessage
+    from ragas.messages import ToolMessage
 
 # Resolve crewai symbols at runtime to avoid mypy issues with untyped packages
 try:
@@ -65,6 +66,12 @@ class CrewAIEventListener:
         self.messages: list[HumanMessage | AIMessage | ToolMessage] = []
 
     def setup_listeners(self, crewai_event_bus: Any) -> None:
+        # Lazy import to reduce memory overhead when ragas is not used
+        from ragas.messages import AIMessage
+        from ragas.messages import HumanMessage
+        from ragas.messages import ToolCall
+        from ragas.messages import ToolMessage
+
         @crewai_event_bus.on(CrewKickoffStartedEvent)
         def on_crew_execution_started(_: Any, event: Any) -> None:
             self.messages.append(

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/drmcp/core/config.py

@@ -79,6 +79,15 @@ class MCPToolConfig(BaseSettings):
         description="Enable/disable Sharepoint tools",
     )
 
+    enable_tavily_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_TAVILY_TOOLS",
+            "ENABLE_TAVILY_TOOLS",
+        ),
+        description="Enable/disable Tavily search tools",
+    )
+
     is_atlassian_oauth_provider_configured: bool = Field(
         default=False,
         validation_alias=AliasChoices(
@@ -131,6 +140,7 @@ class MCPToolConfig(BaseSettings):
         "enable_confluence_tools",
         "enable_gdrive_tools",
         "enable_microsoft_graph_tools",
+        "enable_tavily_tools",
         "is_atlassian_oauth_provider_configured",
         "is_google_oauth_provider_configured",
         "is_microsoft_oauth_provider_configured",

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/drmcp/core/tool_config.py

@@ -31,6 +31,7 @@ class ToolType(str, Enum):
     CONFLUENCE = "confluence"
     GDRIVE = "gdrive"
     MICROSOFT_GRAPH = "microsoft_graph"
+    TAVILY = "tavily"
 
 
 class ToolConfig(TypedDict):
@@ -80,6 +81,13 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
         package_prefix="datarobot_genai.drmcp.tools.microsoft_graph",
         config_field_name="enable_microsoft_graph_tools",
     ),
+    ToolType.TAVILY: ToolConfig(
+        name="tavily",
+        oauth_check=None,
+        directory="tavily",
+        package_prefix="datarobot_genai.drmcp.tools.tavily",
+        config_field_name="enable_tavily_tools",
+    ),
 }
 
 

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/drmcp/tools/clients/microsoft_graph.py

@@ -586,6 +586,102 @@ class MicrosoftGraphClient:
                 f"Microsoft Graph API error {response.status_code}: {response.text}"
             )
 
+    async def update_item_metadata(
+        self,
+        item_id: str,
+        fields_to_update: dict[str, Any],
+        site_id: str | None = None,
+        list_id: str | None = None,
+        drive_id: str | None = None,
+    ) -> dict[str, Any]:
+        """Update metadata on a SharePoint list item or OneDrive/SharePoint drive item.
+
+        For SharePoint list items: Updates custom column values via the fields endpoint.
+        For drive items: Updates properties like name and description.
+
+        Args:
+            item_id: The ID of the item to update.
+            fields_to_update: Key-value pairs of metadata fields to modify.
+            site_id: For SharePoint list items - the site ID.
+            list_id: For SharePoint list items - the list ID.
+            drive_id: For OneDrive/drive items - the drive ID.
+
+        Returns
+        -------
+            The API response containing updated item data.
+
+        Raises
+        ------
+            MicrosoftGraphError: If validation fails or the API request fails.
+        """
+        if not item_id or not item_id.strip():
+            raise MicrosoftGraphError("item_id cannot be empty")
+        if not fields_to_update:
+            raise MicrosoftGraphError("fields_to_update cannot be empty")
+
+        # Determine the endpoint based on provided parameters
+        has_sharepoint_context = site_id is not None and list_id is not None
+        has_drive_context = drive_id is not None
+
+        if has_sharepoint_context and has_drive_context:
+            raise MicrosoftGraphError(
+                "Cannot specify both SharePoint (site_id + list_id) and OneDrive "
+                "(document_library_id) context. Choose one."
+            )
+
+        if not has_sharepoint_context and not has_drive_context:
+            raise MicrosoftGraphError(
+                "Must specify either SharePoint context (site_id + list_id) or "
+                "OneDrive context (document_library_id)."
+            )
+
+        try:
+            if has_sharepoint_context:
+                # PATCH /sites/{site-id}/lists/{list-id}/items/{item-id}/fields
+                url = f"{GRAPH_API_BASE}/sites/{site_id}/lists/{list_id}/items/{item_id}/fields"
+                response = await self._client.patch(url, json=fields_to_update)
+            else:
+                # Drive item: PATCH /drives/{drive-id}/items/{item-id}
+                url = f"{GRAPH_API_BASE}/drives/{drive_id}/items/{item_id}"
+                response = await self._client.patch(url, json=fields_to_update)
+
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPStatusError as e:
+            raise self._handle_update_metadata_error(e, item_id) from e
+
+    def _handle_update_metadata_error(
+        self,
+        error: httpx.HTTPStatusError,
+        item_id: str,
+    ) -> MicrosoftGraphError:
+        """Handle HTTP errors for metadata updates and return appropriate MicrosoftGraphError."""
+        status_code = error.response.status_code
+        error_msg = f"Failed to update metadata: HTTP {status_code}"
+
+        if status_code == 400:
+            try:
+                error_data = error.response.json()
+                api_message = error_data.get("error", {}).get("message", "Invalid request")
+                error_msg = f"Bad request updating metadata: {api_message}"
+            except Exception:
+                error_msg = "Bad request: invalid field names or values."
+        elif status_code == 401:
+            error_msg = "Authentication failed. Access token may be expired or invalid."
+        elif status_code == 403:
+            error_msg = (
+                f"Permission denied: you don't have permission to update item '{item_id}'. "
+                "Requires Sites.ReadWrite.All or Files.ReadWrite.All permission."
+            )
+        elif status_code == 404:
+            error_msg = f"Item '{item_id}' not found."
+        elif status_code == 409:
+            error_msg = f"Conflict: item '{item_id}' may have been modified concurrently."
+        elif status_code == 429:
+            error_msg = "Rate limit exceeded. Please try again later."
+
+        return MicrosoftGraphError(error_msg)
+
     async def __aenter__(self) -> "MicrosoftGraphClient":
         """Async context manager entry."""
         return self

datarobot_genai-0.2.44/src/datarobot_genai/drmcp/tools/clients/tavily.py

@@ -0,0 +1,199 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tavily API Client and utilities for API key authentication."""
+
+import logging
+from typing import Any
+from typing import Literal
+
+from fastmcp.exceptions import ToolError
+from fastmcp.server.dependencies import get_http_headers
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from tavily import AsyncTavilyClient
+
+logger = logging.getLogger(__name__)
+
+MAX_RESULTS: int = 20
+MAX_CHUNKS_PER_SOURCE: int = 3
+
+MAX_RESULTS_DEFAULT: int = 5
+CHUNKS_PER_SOURCE_DEFAULT: int = 1
+
+
+async def get_tavily_access_token() -> str:
+    """
+    Get Tavily API key from HTTP headers.
+
+    Returns
+    -------
+        API key string
+
+    Raises
+    ------
+        ToolError: If API key is not found in headers
+    """
+    headers = get_http_headers()
+
+    api_key = headers.get("x-tavily-api-key")
+    if api_key:
+        return api_key
+
+    logger.warning("Tavily API key not found in headers")
+    raise ToolError(
+        "Tavily API key not found in headers. Please provide it via 'x-tavily-api-key' header."
+    )
+
+
+class TavilySearchResult(BaseModel):
+    """A single search result from Tavily API."""
+
+    title: str
+    url: str
+    content: str
+    score: float
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    @classmethod
+    def from_tavily_sdk(cls, result: dict[str, Any]) -> "TavilySearchResult":
+        """Create a TavilySearchResult from Tavily SDK response data."""
+        return cls(
+            title=result.get("title", ""),
+            url=result.get("url", ""),
+            content=result.get("content", ""),
+            score=result.get("score", 0.0),
+        )
+
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the search result."""
+        return self.model_dump(by_alias=True)
+
+
+class TavilyImage(BaseModel):
+    """An image result from Tavily API."""
+
+    url: str
+    description: str | None = None
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    @classmethod
+    def from_tavily_sdk(cls, image: dict[str, Any] | str) -> "TavilyImage":
+        """Create a TavilyImage from Tavily SDK response data."""
+        if isinstance(image, str):
+            return cls(url=image)
+        return cls(
+            url=image.get("url", ""),
+            description=image.get("description"),
+        )
+
+
+class TavilyClient:
+    """Client for interacting with Tavily Search API.
+
+    This is a wrapper around the official tavily-python SDK.
+    """
+
+    def __init__(self, api_key: str) -> None:
+        self._client = AsyncTavilyClient(api_key=api_key)
+
+    async def search(
+        self,
+        query: str,
+        *,
+        topic: Literal["general", "news", "finance"] = "general",
+        search_depth: Literal["basic", "advanced"] = "basic",
+        max_results: int = MAX_RESULTS_DEFAULT,
+        time_range: Literal["day", "week", "month", "year"] | None = None,
+        include_images: bool = False,
+        include_image_descriptions: bool = False,
+        chunks_per_source: int = CHUNKS_PER_SOURCE_DEFAULT,
+        include_answer: bool = False,
+    ) -> dict[str, Any]:
+        """
+        Perform a web search using Tavily API.
+
+        Args:
+            query: The search query to execute.
+            topic: The category of search ("general", "news", or "finance").
+            search_depth: The depth of search ("basic" or "advanced").
+            max_results: Maximum number of results to return (1-20).
+            time_range: Time range filter ("day", "week", "month", "year").
+            include_images: Whether to include images in results.
+            include_image_descriptions: Whether to include image descriptions.
+            chunks_per_source: Maximum content snippets per URL (1-3).
+            include_answer: Whether to include an AI-generated answer.
+
+        Returns
+        -------
+            Dict with search results from Tavily API.
+
+        Raises
+        ------
+            ValueError: If validation fails.
+            TavilyInvalidAPIKeyError: If the API key is invalid.
+            TavilyUsageLimitExceededError: If usage limit is exceeded.
+            TavilyForbiddenError: If access is forbidden.
+            TavilyBadRequestError: If the request is malformed.
+        """
+        # Validate inputs
+        if not query:
+            raise ValueError("query cannot be empty.")
+        if isinstance(query, str) and not query.strip():
+            raise ValueError("query cannot be empty.")
+        if max_results <= 0:
+            raise ValueError("max_results must be greater than 0.")
+        if max_results > MAX_RESULTS:
+            raise ValueError(f"max_results must be smaller than or equal to {MAX_RESULTS}.")
+        if chunks_per_source <= 0:
+            raise ValueError("chunks_per_source must be greater than 0.")
+        if chunks_per_source > MAX_CHUNKS_PER_SOURCE:
+            raise ValueError(
+                f"chunks_per_source must be smaller than or equal to {MAX_CHUNKS_PER_SOURCE}."
+            )
+
+        # Clamp values to valid ranges
+        max_results = min(max_results, MAX_RESULTS)
+        chunks_per_source = min(chunks_per_source, MAX_CHUNKS_PER_SOURCE)
+
+        # Build search parameters
+        search_kwargs: dict[str, Any] = {
+            "query": query,
+            "topic": topic,
+            "search_depth": search_depth,
+            "max_results": max_results,
+            "include_images": include_images,
+            "include_image_descriptions": include_image_descriptions,
+            "chunks_per_source": chunks_per_source,
+            "include_answer": include_answer,
+        }
+
+        if time_range:
+            search_kwargs["time_range"] = time_range
+
+        return await self._client.search(**search_kwargs)
+
+    async def __aenter__(self) -> "TavilyClient":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
+    ) -> None:
+        """Async context manager exit."""
+        # AsyncTavilyClient doesn't have a close method, but we keep the context manager
+        # pattern for consistency with other clients
+        pass

{datarobot_genai-0.2.41 → datarobot_genai-0.2.44}/src/datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -33,9 +33,7 @@ from datarobot_genai.drmcp.tools.clients.gdrive import get_gdrive_access_token
 logger = logging.getLogger(__name__)
 
 
-@dr_mcp_tool(
-    tags={"google", "gdrive", "list", "search", "files", "find", "contents"}, enabled=False
-)
+@dr_mcp_tool(tags={"google", "gdrive", "list", "search", "files", "find", "contents"})
 async def gdrive_find_contents(
     *,
     page_size: Annotated[
@@ -317,7 +315,7 @@ async def gdrive_update_metadata(
     )
 
 
-@dr_mcp_tool(tags={"google", "gdrive", "manage", "access", "acl"})
+@dr_mcp_tool(tags={"google", "gdrive", "manage", "access", "acl"}, enabled=False)
 async def gdrive_manage_access(
     *,
     file_id: Annotated[str, "The ID of the file or folder."],