datarobot-genai 0.2.27__tar.gz → 0.2.28__tar.gz

{datarobot_genai-0.2.27 → datarobot_genai-0.2.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.27
+Version: 0.2.28
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.

{datarobot_genai-0.2.27 → datarobot_genai-0.2.28}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.27"
+version = "0.2.28"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"

{datarobot_genai-0.2.27 → datarobot_genai-0.2.28}/src/datarobot_genai/drmcp/core/config.py

@@ -31,6 +31,124 @@ from .constants import DEFAULT_DATAROBOT_ENDPOINT
 from .constants import RUNTIME_PARAM_ENV_VAR_NAME_PREFIX
 
 
+class MCPToolConfig(BaseSettings):
+    """Tool configuration for MCP server."""
+
+    enable_predictive_tools: bool = Field(
+        default=True,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
+            "ENABLE_PREDICTIVE_TOOLS",
+        ),
+        description="Enable/disable predictive tools",
+    )
+
+    enable_jira_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
+            "ENABLE_JIRA_TOOLS",
+        ),
+        description="Enable/disable Jira tools",
+    )
+
+    enable_confluence_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
+            "ENABLE_CONFLUENCE_TOOLS",
+        ),
+        description="Enable/disable Confluence tools",
+    )
+
+    enable_gdrive_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
+            "ENABLE_GDRIVE_TOOLS",
+        ),
+        description="Enable/disable GDrive tools",
+    )
+
+    enable_microsoft_graph_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_MICROSOFT_GRAPH_TOOLS",
+            "ENABLE_MICROSOFT_GRAPH_TOOLS",
+        ),
+        description="Enable/disable Sharepoint tools",
+    )
+
+    is_atlassian_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
+            "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Atlassian OAuth provider is configured for Atlassian integration",
+    )
+
+    @property
+    def is_atlassian_oauth_configured(self) -> bool:
+        """Check if Atlassian OAuth is configured via provider flag or environment variables."""
+        return self.is_atlassian_oauth_provider_configured or bool(
+            os.getenv("ATLASSIAN_CLIENT_ID") and os.getenv("ATLASSIAN_CLIENT_SECRET")
+        )
+
+    is_google_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Google OAuth provider is configured for Google integration",
+    )
+
+    @property
+    def is_google_oauth_configured(self) -> bool:
+        return self.is_google_oauth_provider_configured or bool(
+            os.getenv("GOOGLE_CLIENT_ID") and os.getenv("GOOGLE_CLIENT_SECRET")
+        )
+
+    is_microsoft_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
+            "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Microsoft OAuth provider is configured for Microsoft integration",
+    )
+
+    @property
+    def is_microsoft_oauth_configured(self) -> bool:
+        return self.is_microsoft_oauth_provider_configured or bool(
+            os.getenv("MICROSOFT_CLIENT_ID") and os.getenv("MICROSOFT_CLIENT_SECRET")
+        )
+
+    @field_validator(
+        "enable_predictive_tools",
+        "enable_jira_tools",
+        "enable_confluence_tools",
+        "enable_gdrive_tools",
+        "enable_microsoft_graph_tools",
+        "is_atlassian_oauth_provider_configured",
+        "is_google_oauth_provider_configured",
+        "is_microsoft_oauth_provider_configured",
+        mode="before",
+    )
+    @classmethod
+    def validate_runtime_params(cls, v: Any) -> Any:
+        """Validate runtime parameters."""
+        return extract_datarobot_runtime_param_payload(v)
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        case_sensitive=False,
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+
 class MCPServerConfig(BaseSettings):
     """MCP Server configuration using pydantic settings."""
 
@@ -188,86 +306,11 @@ class MCPServerConfig(BaseSettings):
         ),
         description="Enable/disable memory management",
     )
-    enable_predictive_tools: bool = Field(
-        default=True,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
-            "ENABLE_PREDICTIVE_TOOLS",
-        ),
-        description="Enable/disable predictive tools",
-    )
 
-    # Jira tools
-    enable_jira_tools: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
-            "ENABLE_JIRA_TOOLS",
-        ),
-        description="Enable/disable Jira tools",
+    tool_config: MCPToolConfig = Field(
+        default_factory=MCPToolConfig,
+        description="Tool configuration",
     )
-    is_jira_oauth_provider_configured: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
-            "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
-        ),
-        description="Whether Jira OAuth provider is configured for Jira integration",
-    )
-
-    @property
-    def is_jira_oauth_configured(self) -> bool:
-        return self.is_jira_oauth_provider_configured or bool(
-            os.getenv("JIRA_CLIENT_ID") and os.getenv("JIRA_CLIENT_SECRET")
-        )
-
-    # Confluence tools
-    enable_confluence_tools: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
-            "ENABLE_CONFLUENCE_TOOLS",
-        ),
-        description="Enable/disable Confluence tools",
-    )
-    is_confluence_oauth_provider_configured: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
-            "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
-        ),
-        description="Whether Confluence OAuth provider is configured for Confluence integration",
-    )
-
-    @property
-    def is_confluence_oauth_configured(self) -> bool:
-        return self.is_confluence_oauth_provider_configured or bool(
-            os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
-        )
-
-    # Gdrive tools
-    enable_gdrive_tools: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
-            "ENABLE_GDRIVE_TOOLS",
-        ),
-        description="Enable/disable GDrive tools",
-    )
-    is_gdrive_oauth_provider_configured: bool = Field(
-        default=False,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
-            "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
-        ),
-        description="Whether GDrive OAuth provider is configured for GDrive integration",
-    )
-
-    @property
-    def is_gdrive_oauth_configured(self) -> bool:
-        return self.is_gdrive_oauth_provider_configured or bool(
-            os.getenv("GDRIVE_CLIENT_ID") and os.getenv("GDRIVE_CLIENT_SECRET")
-        )
 
     @field_validator(
         "otel_attributes",
@@ -291,11 +334,6 @@ class MCPServerConfig(BaseSettings):
         "mcp_server_register_dynamic_tools_on_startup",
         "tool_registration_duplicate_behavior",
         "mcp_server_register_dynamic_prompts_on_startup",
-        "enable_predictive_tools",
-        "enable_jira_tools",
-        "is_jira_oauth_provider_configured",
-        "enable_confluence_tools",
-        "is_confluence_oauth_provider_configured",
         mode="before",
     )
     @classmethod

{datarobot_genai-0.2.27 → datarobot_genai-0.2.28}/src/datarobot_genai/drmcp/core/tool_config.py

@@ -30,6 +30,7 @@ class ToolType(str, Enum):
     JIRA = "jira"
     CONFLUENCE = "confluence"
     GDRIVE = "gdrive"
+    MICROSOFT_GRAPH = "microsoft_graph"
 
 
 class ToolConfig(TypedDict):
@@ -39,7 +40,7 @@ class ToolConfig(TypedDict):
     oauth_check: Callable[["MCPServerConfig"], bool] | None
     directory: str
     package_prefix: str
-    config_field_name: str  # Name of the config field (e.g., "enable_predictive_tools")
+    config_field_name: str
 
 
 # Tool configuration registry
@@ -53,25 +54,32 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
     ),
     ToolType.JIRA: ToolConfig(
         name="jira",
-        oauth_check=lambda config: config.is_jira_oauth_configured,
+        oauth_check=lambda config: config.tool_config.is_atlassian_oauth_configured,
         directory="jira",
         package_prefix="datarobot_genai.drmcp.tools.jira",
         config_field_name="enable_jira_tools",
     ),
     ToolType.CONFLUENCE: ToolConfig(
         name="confluence",
-        oauth_check=lambda config: config.is_confluence_oauth_configured,
+        oauth_check=lambda config: config.tool_config.is_atlassian_oauth_configured,
         directory="confluence",
         package_prefix="datarobot_genai.drmcp.tools.confluence",
         config_field_name="enable_confluence_tools",
     ),
     ToolType.GDRIVE: ToolConfig(
         name="gdrive",
-        oauth_check=lambda config: config.is_gdrive_oauth_configured,
+        oauth_check=lambda config: config.tool_config.is_google_oauth_configured,
         directory="gdrive",
         package_prefix="datarobot_genai.drmcp.tools.gdrive",
         config_field_name="enable_gdrive_tools",
     ),
+    ToolType.MICROSOFT_GRAPH: ToolConfig(
+        name="microsoft_graph",
+        oauth_check=lambda config: config.tool_config.is_microsoft_oauth_configured,
+        directory="microsoft_graph",
+        package_prefix="datarobot_genai.drmcp.tools.microsoft_graph",
+        config_field_name="enable_microsoft_graph_tools",
+    ),
 }
 
 
@@ -92,12 +100,12 @@ def is_tool_enabled(tool_type: ToolType, config: "MCPServerConfig") -> bool:
     -------
         True if the tool is enabled, False otherwise
     """
-    tool_config = TOOL_CONFIGS[tool_type]
-    enable_config_name = tool_config["config_field_name"]
-    is_enabled = getattr(config, enable_config_name)
+    tool_config_registry = TOOL_CONFIGS[tool_type]
+    enable_config_name = tool_config_registry["config_field_name"]
+    is_enabled = getattr(config.tool_config, enable_config_name)
 
     # If tool is enabled, check OAuth requirements if needed
-    if is_enabled and tool_config["oauth_check"] is not None:
-        return tool_config["oauth_check"](config)
+    if is_enabled and tool_config_registry["oauth_check"] is not None:
+        return tool_config_registry["oauth_check"](config)
 
     return is_enabled

{datarobot_genai-0.2.27 → datarobot_genai-0.2.28}/src/datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -39,6 +39,54 @@ class ETETestExpectations(BaseModel):
 SHOULD_NOT_BE_EMPTY = "SHOULD_NOT_BE_EMPTY"
 
 
+def _extract_structured_content(tool_result: str) -> Any:
+    r"""
+    Extract and parse structured content from tool result string.
+
+    Tool results are formatted as:
+    "Content: {content}\nStructured content: {structured_content}"
+
+    Structured content can be:
+    1. A JSON object with a "result" key: {"result": "..."} or {"result": "{...}"}
+    2. A direct JSON object: {"key": "value", ...}
+    3. Empty or missing
+
+    Args:
+        tool_result: The tool result string
+
+    Returns
+    -------
+        Parsed structured content, or None if not available
+    """
+    # Early returns for invalid inputs
+    if not tool_result or "Structured content: " not in tool_result:
+        return None
+
+    structured_part = tool_result.split("Structured content: ", 1)[1].strip()
+    # Parse JSON, return None on failure or empty structured_part
+    if not structured_part:
+        return None
+    try:
+        structured_data = json.loads(structured_part)
+    except json.JSONDecodeError:
+        return None
+
+    # If structured data has a "result" key, extract and parse that
+    if isinstance(structured_data, dict) and "result" in structured_data:
+        result_value = structured_data["result"]
+        # If result is a JSON string (starts with { or [), try to parse it
+        if isinstance(result_value, str) and result_value.strip().startswith(("{", "[")):
+            try:
+                parsed_result = json.loads(result_value)
+            except json.JSONDecodeError:
+                parsed_result = result_value  # Return string as-is if parsing fails
+            return parsed_result
+        return result_value  # Return result value directly
+
+    # If it's a direct JSON object (not wrapped in {"result": ...}), return it as-is
+    return structured_data
+
+
 def _check_dict_has_keys(
     expected: dict[str, Any],
     actual: dict[str, Any] | list[dict[str, Any]],
@@ -130,7 +178,26 @@ class ToolBaseE2E:
                             f"result, but got: {response.tool_results[i]}"
                         )
                     else:
-                        actual_result = json.loads(response.tool_results[i])
+                        actual_result = _extract_structured_content(response.tool_results[i])
+                        if actual_result is None:
+                            # Fallback: try to parse the entire tool result as JSON
+                            try:
+                                actual_result = json.loads(response.tool_results[i])
+                            except json.JSONDecodeError:
+                                # If that fails, try to extract content part
+                                if "Content: " in response.tool_results[i]:
+                                    content_part = response.tool_results[i].split("Content: ", 1)[1]
+                                    if "\nStructured content: " in content_part:
+                                        content_part = content_part.split(
+                                            "\nStructured content: ", 1
+                                        )[0]
+                                    try:
+                                        actual_result = json.loads(content_part.strip())
+                                    except json.JSONDecodeError:
+                                        raise AssertionError(
+                                            f"Could not parse tool result for "
+                                            f"{tool_call.tool_name}: {response.tool_results[i]}"
+                                        )
                         assert _check_dict_has_keys(expected_result, actual_result), (
                             f"Should have called {tool_call.tool_name} tool with the correct "
                             f"result structure, but got: {response.tool_results[i]}"