datarobot-genai 0.2.39__py3-none-any.whl → 0.3.1__py3-none-any.whl

datarobot_genai/drmcp/tools/clients/perplexity.py

@@ -0,0 +1,173 @@
+# 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.
+
+import logging
+from typing import Any
+from typing import Literal
+
+from fastmcp.exceptions import ToolError
+from fastmcp.server.dependencies import get_http_headers
+from perplexity import AsyncPerplexity
+from perplexity.types import search_create_response
+from pydantic import BaseModel
+from pydantic import ConfigDict
+
+logger = logging.getLogger(__name__)
+
+MAX_QUERIES: int = 5
+MAX_RESULTS: int = 20
+MAX_TOKENS_PER_PAGE: int = 8192
+MAX_SEARCH_DOMAIN_FILTER: int = 20
+
+MAX_RESULTS_DEFAULT: int = 10
+MAX_TOKENS_PER_PAGE_DEFAULT: int = 2048
+
+
+async def get_perplexity_access_token() -> str | ToolError:
+    """
+    Get Perplexity API key from HTTP headers.
+
+    At the moment of creating this fn. Perplexity does not support OAuth.
+    It allows only API-KEY authorized flow.
+
+    Returns
+    -------
+        Access token string on success, ToolError on failure
+
+    Example:
+        ```python
+        token = await get_perplexity_access_token()
+        if isinstance(token, ToolError):
+            # Handle error
+            return token
+        # Use token
+        ```
+    """
+    try:
+        headers = get_http_headers()
+
+        if api_key := headers.get("x-perplexity-api-key"):
+            return api_key
+
+        logger.warning("Perplexity API key not found in headers.")
+        return ToolError(
+            "Perplexity API key not found in headers. "
+            "Please provide it via 'x-perplexity-api-key' header."
+        )
+    except Exception as e:
+        logger.error(f"Unexpected error obtaining Perplexity API key: {e}.", exc_info=e)
+        return ToolError("An unexpected error occured while obtaining Perplexity API key.")
+
+
+class PerplexityError(Exception):
+    """Exception for Perplexity API errors."""
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class PerplexitySearchResult(BaseModel):
+    snippet: str
+    title: str
+    url: str
+    date: str | None = None
+    last_updated: str | None = None
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    @classmethod
+    def from_perplexity_sdk(cls, result: search_create_response.Result) -> "PerplexitySearchResult":
+        """Create a PerplexitySearchResult from perplexity sdk response data."""
+        return cls(**result.model_dump())
+
+    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 PerplexityClient:
+    """Client for interacting with Perplexity API.
+    Its simple wrapper around perplexity python sdk.
+    """
+
+    def __init__(self, access_token: str) -> None:
+        self._client = AsyncPerplexity(api_key=access_token)
+
+    async def search(
+        self,
+        query: str | list[str],
+        search_domain_filter: list[str] | None = None,
+        recency: Literal["hour", "day", "week", "month", "year"] | None = None,
+        max_results: int = MAX_RESULTS_DEFAULT,
+        max_tokens_per_page: int = MAX_TOKENS_PER_PAGE_DEFAULT,
+    ) -> list[PerplexitySearchResult]:
+        """
+        Search using Perplexity.
+
+        Args:
+            query: Query to filter results.
+            search_domain_filter: Up to 20 domains/URLs to allowlist or denylist.
+            recency: Filter results by time period.
+            max_results: Number of ranked results to return.
+            max_tokens_per_page: Context extraction cap per page.
+
+        Returns
+        -------
+            List of Perplexity search results.
+        """
+        if not query:
+            raise PerplexityError("Error: query cannot be empty.")
+        if query and isinstance(query, str) and not query.strip():
+            raise PerplexityError("Error: query cannot be empty.")
+        if query and isinstance(query, list) and len(query) > MAX_QUERIES:
+            raise PerplexityError(f"Error: query list cannot be bigger than {MAX_QUERIES}.")
+        if query and isinstance(query, list) and not all(q.strip() for q in query):
+            raise PerplexityError("Error: query cannot contain empty str.")
+        if search_domain_filter and len(search_domain_filter) > MAX_SEARCH_DOMAIN_FILTER:
+            raise PerplexityError("Error: maximum number of search domain filters is 20.")
+        if max_results <= 0:
+            raise PerplexityError("Error: max_results must be greater than 0.")
+        if max_results > MAX_RESULTS:
+            raise PerplexityError("Error: max_results must be smaller than or equal to 20.")
+        if max_tokens_per_page <= 0:
+            raise PerplexityError("Error: max_tokens_per_page must be greater than 0.")
+        if max_tokens_per_page > MAX_TOKENS_PER_PAGE:
+            raise PerplexityError(
+                "Error: max_tokens_per_page must be smaller than or equal to 8192."
+            )
+
+        max_results = min(max_results, MAX_RESULTS)
+        max_tokens_per_page = min(max_tokens_per_page, MAX_TOKENS_PER_PAGE)
+
+        search_result = await self._client.search.create(
+            query=query,
+            search_domain_filter=search_domain_filter,
+            search_recency_filter=recency,
+            max_results=max_results,
+            max_tokens_per_page=max_tokens_per_page,
+        )
+
+        return [
+            PerplexitySearchResult.from_perplexity_sdk(result) for result in search_result.results
+        ]
+
+    async def __aenter__(self) -> "PerplexityClient":
+        """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."""
+        await self._client.close()

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/drmcp/tools/confluence/tools.py

@@ -67,7 +67,6 @@ async def confluence_get_page(
             page_response = await client.get_page_by_title(page_id_or_title, space_key)
 
     return ToolResult(
-        content=f"Successfully retrieved page '{page_response.title}'.",
         structured_content=page_response.as_flat_dict(),
     )
 
@@ -115,7 +114,6 @@ async def confluence_create_page(
         )
 
     return ToolResult(
-        content=f"New page '{title}' created successfully in space '{space_key}'.",
         structured_content={"new_page_id": page_response.page_id, "title": page_response.title},
     )
 
@@ -151,7 +149,6 @@ async def confluence_add_comment(
         )
 
     return ToolResult(
-        content=f"Comment added successfully to page ID {page_id}.",
         structured_content={
             "comment_id": comment_response.comment_id,
             "page_id": page_id,
@@ -212,7 +209,6 @@ async def confluence_search(
 
     n = len(results)
     return ToolResult(
-        content=f"Successfully executed CQL query and retrieved {n} result(s).",
         structured_content={"data": data, "count": n},
     )
 
@@ -267,7 +263,6 @@ async def confluence_update_page(
         )
 
     return ToolResult(
-        content=f"Page ID {page_id} updated successfully to version {page_response.version}.",
         structured_content={
             "updated_page_id": page_response.page_id,
             "new_version": page_response.version,

datarobot_genai/drmcp/tools/gdrive/tools.py

@@ -22,7 +22,6 @@ from fastmcp.exceptions import ToolError
 from fastmcp.tools.tool import ToolResult
 
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
-from datarobot_genai.drmcp.tools.clients.gdrive import GOOGLE_DRIVE_FOLDER_MIME
 from datarobot_genai.drmcp.tools.clients.gdrive import LIMIT
 from datarobot_genai.drmcp.tools.clients.gdrive import MAX_PAGE_SIZE
 from datarobot_genai.drmcp.tools.clients.gdrive import SUPPORTED_FIELDS
@@ -33,9 +32,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[
@@ -66,13 +63,15 @@ async def gdrive_find_contents(
 ) -> ToolResult:
     """
     Search or list files in the user's Google Drive with pagination and filtering support.
-    Use this tool to discover file names and IDs for use with other tools.
+    Use this tool to discover GDrive file names and IDs for use with other tools.
 
     Limit must be bigger than or equal to page size and it must be multiplication of page size.
-    Ex.
-        page size = 10 limit = 50
-        page size = 3 limit = 3
-        page size = 12 limit = 36
+
+    Examples
+    --------
+        - page size = 10 limit = 50
+        - page size = 3 limit = 3
+        - page size = 12 limit = 36
     """
     access_token = await get_gdrive_access_token()
     if isinstance(access_token, ToolError):
@@ -90,13 +89,8 @@ async def gdrive_find_contents(
 
     filtered_fields = set(fields).intersection(SUPPORTED_FIELDS) if fields else SUPPORTED_FIELDS
     number_of_files = len(data.files)
-    next_page_info = (
-        f"Next page token needed to fetch more data: {data.next_page_token}"
-        if data.next_page_token
-        else "There're no more pages."
-    )
+
     return ToolResult(
-        content=f"Successfully listed {number_of_files} files. {next_page_info}",
         structured_content={
             "files": [
                 file.model_dump(by_alias=True, include=filtered_fields) for file in data.files
@@ -119,8 +113,7 @@ async def gdrive_read_content(
     ] = None,
 ) -> ToolResult:
     """
-    Retrieve the content of a specific file by its ID. Google Workspace files are
-    automatically exported to LLM-readable formats (Push-Down).
+    Retrieve the content of a specific Google drive file by its ID.
 
     Usage:
         - Basic: gdrive_read_content(file_id="1ABC123def456")
@@ -150,15 +143,7 @@ async def gdrive_read_content(
     async with GoogleDriveClient(access_token) as client:
         file_content = await client.read_file_content(file_id, target_format)
 
-    export_info = ""
-    if file_content.was_exported:
-        export_info = f" (exported from {file_content.original_mime_type})"
-
     return ToolResult(
-        content=(
-            f"Successfully retrieved content of '{file_content.name}' "
-            f"({file_content.mime_type}){export_info}."
-        ),
         structured_content=file_content.as_flat_dict(),
     )
 
@@ -230,13 +215,7 @@ async def gdrive_create_file(
             initial_content=initial_content,
         )
 
-    file_type = "folder" if mime_type == GOOGLE_DRIVE_FOLDER_MIME else "file"
-    content_info = ""
-    if initial_content and mime_type != GOOGLE_DRIVE_FOLDER_MIME:
-        content_info = " with initial content"
-
     return ToolResult(
-        content=f"Successfully created {file_type} '{created_file.name}'{content_info}.",
         structured_content=created_file.as_flat_dict(),
     )
 
@@ -297,27 +276,12 @@ async def gdrive_update_metadata(
             trashed=trash,
         )
 
-    changes: list[str] = []
-    if new_name is not None:
-        changes.append(f"renamed to '{new_name}'")
-    if starred is True:
-        changes.append("starred")
-    elif starred is False:
-        changes.append("unstarred")
-    if trash is True:
-        changes.append("moved to trash")
-    elif trash is False:
-        changes.append("restored from trash")
-
-    changes_description = ", ".join(changes)
-
     return ToolResult(
-        content=f"Successfully updated file '{updated_file.name}': {changes_description}.",
         structured_content=updated_file.as_flat_dict(),
     )
 
 
-@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."],
@@ -388,17 +352,6 @@ async def gdrive_manage_access(
     # Build response
     structured_content = {"affectedFileId": file_id}
     if action == "add":
-        content = (
-            f"Successfully added role '{role}' for '{email_address}' for gdrive file '{file_id}'. "
-            f"New permission id '{permission_id}'."
-        )
         structured_content["newPermissionId"] = permission_id
-    elif action == "update":
-        content = (
-            f"Successfully updated role '{role}' (permission '{permission_id}') "
-            f"for gdrive file '{file_id}'."
-        )
-    else:  # action == "remove":
-        content = f"Successfully removed permission '{permission_id}' for gdrive file '{file_id}'."
 
-    return ToolResult(content=content, structured_content=structured_content)
+    return ToolResult(structured_content=structured_content)

datarobot_genai/drmcp/tools/jira/tools.py

@@ -53,10 +53,11 @@ async def jira_search_issues(
     async with JiraClient(access_token) as client:
         issues = await client.search_jira_issues(jql_query=jql_query, max_results=max_results)
 
-    n = len(issues)
     return ToolResult(
-        content=f"Successfully executed JQL query and retrieved {n} issue(s).",
-        structured_content={"data": [issue.as_flat_dict() for issue in issues], "count": n},
+        structured_content={
+            "data": [issue.as_flat_dict() for issue in issues],
+            "count": len(issues),
+        },
     )
 
 
@@ -76,7 +77,6 @@ async def jira_get_issue(
         issue = await client.get_jira_issue(issue_key)
 
     return ToolResult(
-        content=f"Successfully retrieved details for issue '{issue_key}'.",
         structured_content=issue.as_flat_dict(),
     )
 
@@ -121,7 +121,6 @@ async def jira_create_issue(
         )
 
     return ToolResult(
-        content=f"Successfully created issue '{issue_key}'.",
         structured_content={"newIssueKey": issue_key, "projectKey": project_key},
     )
 
@@ -174,9 +173,7 @@ async def jira_update_issue(
             issue_key=issue_key, fields=fields_to_update
         )
 
-    updated_fields_str = ",".join(updated_fields)
     return ToolResult(
-        content=f"Successfully updated issue '{issue_key}'. Fields modified: {updated_fields_str}.",
         structured_content={"updatedIssueKey": issue_key, "fields": updated_fields},
     )
 
@@ -216,7 +213,6 @@ async def jira_transition_issue(
         await client.transition_jira_issue(issue_key=issue_key, transition_id=transition_id)
 
     return ToolResult(
-        content=f"Successfully transitioned issue '{issue_key}' to status '{transition_name}'.",
         structured_content={
             "transitionedIssueKey": issue_key,
             "newStatusName": transition_name,