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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.42
+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.42 → datarobot_genai-0.2.44}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.42"
+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",

{datarobot_genai-0.2.42 → 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.42 → 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.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.42 → 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."],

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

@@ -0,0 +1,13 @@
+# 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.

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

@@ -0,0 +1,148 @@
+# 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 MCP tools for web search."""
+
+import logging
+from typing import Annotated
+from typing import Literal
+
+from fastmcp.tools.tool import ToolResult
+
+from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.tools.clients.tavily import CHUNKS_PER_SOURCE_DEFAULT
+from datarobot_genai.drmcp.tools.clients.tavily import MAX_CHUNKS_PER_SOURCE
+from datarobot_genai.drmcp.tools.clients.tavily import MAX_RESULTS
+from datarobot_genai.drmcp.tools.clients.tavily import MAX_RESULTS_DEFAULT
+from datarobot_genai.drmcp.tools.clients.tavily import TavilyClient
+from datarobot_genai.drmcp.tools.clients.tavily import TavilyImage
+from datarobot_genai.drmcp.tools.clients.tavily import TavilySearchResult
+from datarobot_genai.drmcp.tools.clients.tavily import get_tavily_access_token
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"search", "tavily", "web", "websearch"})
+async def tavily_search(
+    *,
+    query: Annotated[str, "The search query to execute."],
+    topic: Annotated[
+        Literal["general", "news", "finance"],
+        "The category of search. Use 'general' for broad web search, "
+        "'news' for recent news articles, or 'finance' for financial information.",
+    ] = "general",
+    search_depth: Annotated[
+        Literal["basic", "advanced"],
+        "The depth of search. 'basic' is faster and cheaper, "
+        "'advanced' provides more comprehensive results.",
+    ] = "basic",
+    max_results: Annotated[
+        int,
+        f"Maximum number of search results to return (1-{MAX_RESULTS}).",
+    ] = MAX_RESULTS_DEFAULT,
+    time_range: Annotated[
+        Literal["day", "week", "month", "year"] | None,
+        "Filter results by time range. Use 'day' for last 24 hours, "
+        "'week' for last 7 days, 'month' for last 30 days, or 'year' for last year.",
+    ] = None,
+    include_images: Annotated[
+        bool,
+        "Whether to include related images in the search results.",
+    ] = False,
+    include_image_descriptions: Annotated[
+        bool,
+        "Whether to include AI-generated descriptions for images. "
+        "Only applicable when include_images is True.",
+    ] = False,
+    chunks_per_source: Annotated[
+        int,
+        f"Maximum number of content snippets to return per source URL (1-{MAX_CHUNKS_PER_SOURCE}).",
+    ] = CHUNKS_PER_SOURCE_DEFAULT,
+    include_answer: Annotated[
+        bool,
+        "Whether to include an AI-generated answer summarizing the search results.",
+    ] = False,
+) -> ToolResult:
+    """
+    Perform a real-time web search using Tavily API.
+
+    Tavily is optimized for AI agents and provides clean, relevant search results
+    suitable for LLM consumption. Use this tool to search the web for current
+    information, news, financial data, or general knowledge.
+
+    Usage:
+        - Basic search: tavily_search(query="Python best practices 2026")
+        - News search: tavily_search(query="AI regulations", topic="news", time_range="week")
+        - Financial search: tavily_search(query="AAPL stock analysis", topic="finance")
+        - Comprehensive search: tavily_search(
+            query="climate change solutions",
+            search_depth="advanced",
+            max_results=10,
+            include_answer=True
+          )
+
+    Note:
+        - Advanced search depth consumes more API credits but provides better results
+        - Time range filtering is useful for finding recent information
+    """
+    api_key = await get_tavily_access_token()
+
+    async with TavilyClient(api_key) as client:
+        response = await client.search(
+            query=query,
+            topic=topic,
+            search_depth=search_depth,
+            max_results=max_results,
+            time_range=time_range,
+            include_images=include_images,
+            include_image_descriptions=include_image_descriptions,
+            chunks_per_source=chunks_per_source,
+            include_answer=include_answer,
+        )
+
+    results = [TavilySearchResult.from_tavily_sdk(r) for r in response.get("results", [])]
+
+    images: list[TavilyImage] | None = None
+    if include_images and response.get("images"):
+        images = [TavilyImage.from_tavily_sdk(img) for img in response.get("images", [])]
+
+    result_count = len(results)
+    answer = response.get("answer")
+    response_time = response.get("response_time", 0.0)
+
+    answer_info = " with AI-generated answer" if answer else ""
+    image_info = f" and {len(images)} images" if images else ""
+
+    structured_content: dict = {
+        "query": response.get("query", query),
+        "results": [r.as_flat_dict() for r in results],
+        "resultCount": result_count,
+        "responseTime": response_time,
+    }
+
+    if answer:
+        structured_content["answer"] = answer
+
+    if images:
+        structured_content["images"] = [
+            {"url": img.url, "description": img.description} for img in images
+        ]
+
+    return ToolResult(
+        content=(
+            f"Successfully searched for '{query}'. "
+            f"Found {result_count} results{answer_info}{image_info}."
+        ),
+        structured_content=structured_content,
+    )