datarobot-genai 0.2.37__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

@@ -55,29 +55,18 @@ async def confluence_get_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            if page_id_or_title.isdigit():
-                page_response = await client.get_page_by_id(page_id_or_title)
-            else:
-                if not space_key:
-                    raise ToolError(
-                        "Argument validation error: "
-                        "'space_key' is required when identifying a page by title."
-                    )
-                page_response = await client.get_page_by_title(page_id_or_title, space_key)
-    except ConfluenceError as e:
-        logger.error(f"Confluence error getting page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error getting Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while getting Confluence page "
-            f"'{page_id_or_title}': {str(e)}"
-        )
+    async with ConfluenceClient(access_token) as client:
+        if page_id_or_title.isdigit():
+            page_response = await client.get_page_by_id(page_id_or_title)
+        else:
+            if not space_key:
+                raise ToolError(
+                    "Argument validation error: "
+                    "'space_key' is required when identifying a page by title."
+                )
+            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(),
     )
 
@@ -116,26 +105,15 @@ async def confluence_create_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            page_response = await client.create_page(
-                space_key=space_key,
-                title=title,
-                body_content=body_content,
-                parent_id=parent_id,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error creating page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error creating Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while creating Confluence page "
-            f"'{title}' in space '{space_key}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        page_response = await client.create_page(
+            space_key=space_key,
+            title=title,
+            body_content=body_content,
+            parent_id=parent_id,
         )
 
     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},
     )
 
@@ -164,23 +142,13 @@ async def confluence_add_comment(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            comment_response = await client.add_comment(
-                page_id=page_id,
-                comment_body=comment_body,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error adding comment: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error adding comment to Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while adding comment to page '{page_id}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        comment_response = await client.add_comment(
+            page_id=page_id,
+            comment_body=comment_body,
         )
 
     return ToolResult(
-        content=f"Comment added successfully to page ID {page_id}.",
         structured_content={
             "comment_id": comment_response.comment_id,
             "page_id": page_id,
@@ -220,36 +188,27 @@ async def confluence_search(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            results = await client.search_confluence_content(
-                cql_query=cql_query, max_results=max_results
-            )
-
-            # If include_body is True, fetch full content for each page
-            if include_body and results:
-                data = []
-                for result in results:
-                    flat = result.as_flat_dict()
-                    try:
-                        page = await client.get_page_by_id(result.id)
-                        flat["body"] = page.body
-                    except ConfluenceError:
-                        flat["body"] = None  # Keep excerpt if page fetch fails
-                    data.append(flat)
-            else:
-                data = [result.as_flat_dict() for result in results]
-
-    except ConfluenceError as e:
-        logger.error(f"Confluence error searching content: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error searching Confluence content: {e}")
-        raise ToolError(f"An unexpected error occurred while searching Confluence: {str(e)}")
+    async with ConfluenceClient(access_token) as client:
+        results = await client.search_confluence_content(
+            cql_query=cql_query, max_results=max_results
+        )
+
+        # If include_body is True, fetch full content for each page
+        if include_body and results:
+            data = []
+            for result in results:
+                flat = result.as_flat_dict()
+                try:
+                    page = await client.get_page_by_id(result.id)
+                    flat["body"] = page.body
+                except ConfluenceError:
+                    flat["body"] = None  # Keep excerpt if page fetch fails
+                data.append(flat)
+        else:
+            data = [result.as_flat_dict() for result in results]
 
     n = len(results)
     return ToolResult(
-        content=f"Successfully executed CQL query and retrieved {n} result(s).",
         structured_content={"data": data, "count": n},
     )
 
@@ -296,24 +255,14 @@ async def confluence_update_page(
     if isinstance(access_token, ToolError):
         raise access_token
 
-    try:
-        async with ConfluenceClient(access_token) as client:
-            page_response = await client.update_page(
-                page_id=page_id,
-                new_body_content=new_body_content,
-                version_number=version_number,
-            )
-    except ConfluenceError as e:
-        logger.error(f"Confluence error updating page: {e}")
-        raise ToolError(str(e))
-    except Exception as e:
-        logger.error(f"Unexpected error updating Confluence page: {e}")
-        raise ToolError(
-            f"An unexpected error occurred while updating Confluence page '{page_id}': {str(e)}"
+    async with ConfluenceClient(access_token) as client:
+        page_response = await client.update_page(
+            page_id=page_id,
+            new_body_content=new_body_content,
+            version_number=version_number,
         )
 
     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,