codemie-sdk-python 0.1.226__py3-none-any.whl → 0.1.256__py3-none-any.whl

codemie_sdk/services/datasource.py

@@ -15,7 +15,11 @@ from ..models.datasource import (
     UpdateCodeDataSourceRequest,
     BaseUpdateDataSourceRequest,
     FileDataSourceRequest,
+    CodeAnalysisDataSourceRequest,
+    CodeExplorationDataSourceRequest,
+    ElasticsearchStatsResponse,
 )
+from ..models.assistant import AssistantListResponse
 from ..utils import ApiRequestHandler
 
 
@@ -206,3 +210,66 @@ class DatasourceService:
             Deletion confirmation
         """
         return self._api.delete(f"/v1/index/{datasource_id}", dict)
+
+    def get_assistants_using_datasource(
+        self, datasource_id: str
+    ) -> List[AssistantListResponse]:
+        """Get list of assistants that are using this datasource.
+
+        Args:
+            datasource_id: ID of the datasource
+
+        Returns:
+            List of AssistantListResponse objects containing assistants using this datasource
+
+        Raises:
+            ApiError: If the datasource is not found or other API errors occur.
+        """
+        return self._api.get(
+            f"/v1/index/{datasource_id}/assistants", List[AssistantListResponse]
+        )
+
+    def create_provider_datasource(
+        self,
+        toolkit_id: str,
+        provider_name: str,
+        request: Union[CodeAnalysisDataSourceRequest, CodeExplorationDataSourceRequest],
+    ) -> dict:
+        """Create a provider-based datasource.
+
+        Args:
+            toolkit_id: ID of the toolkit to use
+            provider_name: Name of the provider
+            request: Provider datasource creation request (CodeAnalysisDataSourceRequest or CodeExplorationDataSourceRequest)
+
+        Returns:
+            dict: Response from the server containing operation status
+        """
+        endpoint = (
+            f"/v1/index/provider?toolkit_id={toolkit_id}&provider_name={provider_name}"
+        )
+
+        return self._api.post(
+            endpoint,
+            dict,
+            json_data=request.model_dump(by_alias=True, exclude_none=True),
+        )
+
+    def get_elasticsearch_stats(self, datasource_id: str) -> ElasticsearchStatsResponse:
+        """Get Elasticsearch statistics for a specific datasource index.
+
+        Args:
+            datasource_id: ID of the datasource
+
+        Returns:
+            ElasticsearchStatsResponse with Elasticsearch statistics including:
+            - index_name: Name of the index in Elasticsearch
+            - size_in_bytes: Size of the index in bytes
+
+        Raises:
+            ApiError: If the datasource is not found, platform datasources are not supported,
+                     or Elasticsearch statistics are not available.
+        """
+        return self._api.get(
+            f"/v1/index/{datasource_id}/elasticsearch", ElasticsearchStatsResponse
+        )

codemie_sdk/services/files.py

@@ -59,3 +59,24 @@ class FileOperationService:
         )
 
         return response
+
+    def get_file(self, file_id: str) -> bytes:
+        """Get a file by its ID.
+
+        Args:
+            file_id: The file identifier (base64 encoded ID from file_url)
+
+        Returns:
+            bytes: The file content as binary data
+
+        Raises:
+            ApiError: If the file doesn't exist or there's an API error
+        """
+        import requests
+
+        response = self._api.get(
+            f"/v1/files/{file_id}",
+            response_model=requests.Response,
+            wrap_response=False,
+        )
+        return response.content

codemie_sdk/services/vendor_guardrail.py

@@ -0,0 +1,375 @@
+"""Vendor guardrail service implementation for managing cloud vendor guardrail settings."""
+
+from typing import Union, Optional, List
+
+from ..models.vendor_assistant import VendorType
+from ..models.vendor_guardrail import (
+    VendorGuardrailSettingsResponse,
+    VendorGuardrailsResponse,
+    VendorGuardrail,
+    VendorGuardrailVersion,
+    VendorGuardrailVersionsResponse,
+    VendorGuardrailInstallRequest,
+    VendorGuardrailInstallResponse,
+    VendorGuardrailUninstallResponse,
+)
+from ..utils import ApiRequestHandler
+
+
+class VendorGuardrailService:
+    """Service for managing cloud vendor guardrail settings (AWS, Azure, GCP)."""
+
+    def __init__(self, api_domain: str, token: str, verify_ssl: bool = True):
+        """Initialize the vendor guardrail service.
+
+        Args:
+            api_domain: Base URL for the CodeMie API
+            token: Authentication token
+            verify_ssl: Whether to verify SSL certificates
+        """
+        self._api = ApiRequestHandler(api_domain, token, verify_ssl)
+
+    def get_guardrail_settings(
+        self,
+        vendor: Union[VendorType, str],
+        page: int = 0,
+        per_page: int = 10,
+    ) -> VendorGuardrailSettingsResponse:
+        """Get guardrail settings for a specific cloud vendor.
+
+        Args:
+            vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
+            page: Page number for pagination (0-based)
+            per_page: Number of items per page
+
+        Returns:
+            VendorGuardrailSettingsResponse containing list of settings and pagination info
+
+        Example:
+            >>> # Using enum
+            >>> settings = client.vendor_guardrails.get_guardrail_settings(VendorType.AWS, page=0, per_page=10)
+            >>> # Using string
+            >>> settings = client.vendor_guardrails.get_guardrail_settings("aws", page=0, per_page=10)
+            >>> # Access settings data
+            >>> for setting in settings.data:
+            ...     print(f"Setting: {setting.setting_name}, Project: {setting.project}")
+            ...     if setting.invalid:
+            ...         print(f"Error: {setting.error}")
+        """
+        # Convert enum to string value if needed
+        vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
+
+        params = {
+            "page": page,
+            "per_page": per_page,
+        }
+
+        return self._api.get(
+            f"/v1/vendors/{vendor_str}/guardrails/settings",
+            VendorGuardrailSettingsResponse,
+            params=params,
+            wrap_response=False,
+        )
+
+    def get_guardrails(
+        self,
+        vendor: Union[VendorType, str],
+        setting_id: str,
+        per_page: int = 10,
+        next_token: Optional[str] = None,
+    ) -> VendorGuardrailsResponse:
+        """Get guardrails for a specific vendor setting.
+
+        Args:
+            vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
+            setting_id: ID of the vendor setting to retrieve guardrails for
+            per_page: Number of items per page
+            next_token: Token for pagination (optional, for retrieving next page)
+
+        Returns:
+            VendorGuardrailsResponse containing list of guardrails and pagination token
+
+        Example:
+            >>> # Get first page
+            >>> guardrails = client.vendor_guardrails.get_guardrails(
+            ...     vendor=VendorType.AWS,
+            ...     setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
+            ...     per_page=8
+            ... )
+            >>> # Access guardrail data
+            >>> for guardrail in guardrails.data:
+            ...     print(f"Name: {guardrail.name}, Status: {guardrail.status}")
+            ...     print(f"Version: {guardrail.version}, Description: {guardrail.description}")
+            >>> # Get next page if available
+            >>> if guardrails.pagination.next_token:
+            ...     next_page = client.vendor_guardrails.get_guardrails(
+            ...         vendor=VendorType.AWS,
+            ...         setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
+            ...         per_page=8,
+            ...         next_token=guardrails.pagination.next_token
+            ...     )
+        """
+        # Convert enum to string value if needed
+        vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
+
+        params = {
+            "setting_id": setting_id,
+            "per_page": per_page,
+        }
+
+        if next_token:
+            params["next_token"] = next_token
+
+        return self._api.get(
+            f"/v1/vendors/{vendor_str}/guardrails",
+            VendorGuardrailsResponse,
+            params=params,
+            wrap_response=False,
+        )
+
+    def get_guardrail(
+        self,
+        vendor: Union[VendorType, str],
+        guardrail_id: str,
+        setting_id: str,
+    ) -> VendorGuardrail:
+        """Get a specific guardrail by ID for a vendor setting.
+
+        Args:
+            vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
+            guardrail_id: ID of the guardrail to retrieve
+            setting_id: ID of the vendor setting
+
+        Returns:
+            VendorGuardrail containing guardrail details
+
+        Example:
+            >>> guardrail = client.vendor_guardrails.get_guardrail(
+            ...     vendor=VendorType.AWS,
+            ...     guardrail_id="lss9vxro9oxg",
+            ...     setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
+            ... )
+            >>> print(f"Name: {guardrail.name}, Status: {guardrail.status}")
+            >>> print(f"Version: {guardrail.version}, Description: {guardrail.description}")
+        """
+        # Convert enum to string value if needed
+        vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
+
+        params = {
+            "setting_id": setting_id,
+        }
+
+        return self._api.get(
+            f"/v1/vendors/{vendor_str}/guardrails/{guardrail_id}",
+            VendorGuardrail,
+            params=params,
+            wrap_response=False,
+        )
+
+    def get_guardrail_version(
+        self,
+        vendor: Union[VendorType, str],
+        guardrail_id: str,
+        version: str,
+        setting_id: str,
+    ) -> VendorGuardrailVersion:
+        """Get a specific version of a vendor guardrail with detailed information.
+
+        Args:
+            vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
+            guardrail_id: ID of the guardrail to retrieve
+            version: Version number to retrieve (e.g., "1", "2", "DRAFT")
+            setting_id: ID of the vendor setting
+
+        Returns:
+            VendorGuardrailVersion containing detailed version information including
+            blocked messaging settings and timestamps
+
+        Example:
+            >>> version_details = client.vendor_guardrails.get_guardrail_version(
+            ...     vendor=VendorType.AWS,
+            ...     guardrail_id="lss9vxro9oxg",
+            ...     version="1",
+            ...     setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
+            ... )
+            >>> print(f"Name: {version_details.name}")
+            >>> print(f"Version: {version_details.version}")
+            >>> print(f"Status: {version_details.status}")
+            >>> print(f"Blocked Input Message: {version_details.blockedInputMessaging}")
+            >>> print(f"Blocked Output Message: {version_details.blockedOutputsMessaging}")
+        """
+        # Convert enum to string value if needed
+        vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
+
+        params = {
+            "setting_id": setting_id,
+        }
+
+        return self._api.get(
+            f"/v1/vendors/{vendor_str}/guardrails/{guardrail_id}/{version}",
+            VendorGuardrailVersion,
+            params=params,
+            wrap_response=False,
+        )
+
+    def get_guardrail_versions(
+        self,
+        vendor: Union[VendorType, str],
+        guardrail_id: str,
+        setting_id: str,
+        per_page: int = 10,
+        next_token: Optional[str] = None,
+    ) -> VendorGuardrailVersionsResponse:
+        """Get versions for a specific vendor guardrail.
+
+        Args:
+            vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
+            guardrail_id: ID of the guardrail to retrieve versions for
+            setting_id: ID of the vendor setting
+            per_page: Number of items per page
+            next_token: Token for pagination (optional, for retrieving next page)
+
+        Returns:
+            VendorGuardrailVersionsResponse containing list of versions and pagination token
+
+        Example:
+            >>> # Get first page of versions
+            >>> versions = client.vendor_guardrails.get_guardrail_versions(
+            ...     vendor=VendorType.AWS,
+            ...     guardrail_id="lss9vxro9oxg",
+            ...     setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
+            ...     per_page=5
+            ... )
+            >>> for version in versions.data:
+            ...     print(f"{version.name} (v{version.version}): {version.status}")
+            >>> # Get next page if available
+            >>> if versions.pagination.next_token:
+            ...     next_page = client.vendor_guardrails.get_guardrail_versions(
+            ...         vendor=VendorType.AWS,
+            ...         guardrail_id="lss9vxro9oxg",
+            ...         setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c",
+            ...         per_page=5,
+            ...         next_token=versions.pagination.next_token
+            ...     )
+        """
+        # Convert enum to string value if needed
+        vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
+
+        params = {
+            "setting_id": setting_id,
+            "per_page": per_page,
+        }
+
+        if next_token:
+            params["next_token"] = next_token
+
+        return self._api.get(
+            f"/v1/vendors/{vendor_str}/guardrails/{guardrail_id}/versions",
+            VendorGuardrailVersionsResponse,
+            params=params,
+            wrap_response=False,
+        )
+
+    def install_guardrails(
+        self,
+        vendor: Union[VendorType, str],
+        guardrails: List[VendorGuardrailInstallRequest],
+    ) -> VendorGuardrailInstallResponse:
+        """Install/activate vendor guardrails.
+
+        Args:
+            vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
+            guardrails: List of guardrail installation requests with guardrail ID, version, and setting ID
+
+        Returns:
+            VendorGuardrailInstallResponse containing installation summary with AI run IDs
+
+        Example:
+            >>> from codemie_sdk import VendorGuardrailInstallRequest
+            >>> # Install single guardrail
+            >>> install_request = VendorGuardrailInstallRequest(
+            ...     id="lss9vxro9oxg",
+            ...     version="1",
+            ...     setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
+            ... )
+            >>> response = client.vendor_guardrails.install_guardrails(
+            ...     vendor=VendorType.AWS,
+            ...     guardrails=[install_request]
+            ... )
+            >>> for item in response.summary:
+            ...     print(f"Installed {item.guardrailId} version {item.version} with run ID: {item.aiRunId}")
+            >>>
+            >>> # Install multiple guardrails
+            >>> requests = [
+            ...     VendorGuardrailInstallRequest(
+            ...         id="GUARDRAIL_ID_1",
+            ...         version="1",
+            ...         setting_id="SETTING_ID"
+            ...     ),
+            ...     VendorGuardrailInstallRequest(
+            ...         id="GUARDRAIL_ID_2",
+            ...         version="2",
+            ...         setting_id="SETTING_ID"
+            ...     )
+            ... ]
+            >>> response = client.vendor_guardrails.install_guardrails(
+            ...     vendor=VendorType.AWS,
+            ...     guardrails=requests
+            ... )
+        """
+        # Convert enum to string value if needed
+        vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
+
+        # Convert list of Pydantic models to list of dicts
+        payload = [guardrail.model_dump(by_alias=True) for guardrail in guardrails]
+
+        return self._api.post(
+            f"/v1/vendors/{vendor_str}/guardrails",
+            VendorGuardrailInstallResponse,
+            json_data=payload,
+            wrap_response=False,
+        )
+
+    def uninstall_guardrail(
+        self,
+        vendor: Union[VendorType, str],
+        ai_run_id: str,
+    ) -> VendorGuardrailUninstallResponse:
+        """Uninstall/deactivate a vendor guardrail.
+
+        Args:
+            vendor: Cloud vendor type (aws, azure, gcp). Can be VendorType enum or string.
+            ai_run_id: AI run ID returned from the install operation
+
+        Returns:
+            VendorGuardrailUninstallResponse with success status
+
+        Example:
+            >>> # First, install a guardrail
+            >>> install_request = VendorGuardrailInstallRequest(
+            ...     id="lss9vxro9oxg",
+            ...     version="1",
+            ...     setting_id="cac90788-39b7-4ffe-8b57-e8b047fa1f6c"
+            ... )
+            >>> install_response = client.vendor_guardrails.install_guardrails(
+            ...     vendor=VendorType.AWS,
+            ...     guardrails=[install_request]
+            ... )
+            >>> ai_run_id = install_response.summary[0].aiRunId
+            >>>
+            >>> # Later, uninstall the guardrail using the AI run ID
+            >>> response = client.vendor_guardrails.uninstall_guardrail(
+            ...     vendor=VendorType.AWS,
+            ...     ai_run_id=ai_run_id
+            ... )
+            >>> if response.success:
+            ...     print("Guardrail successfully uninstalled!")
+        """
+        # Convert enum to string value if needed
+        vendor_str = vendor.value if isinstance(vendor, VendorType) else vendor
+
+        return self._api.delete(
+            f"/v1/vendors/{vendor_str}/guardrails/{ai_run_id}",
+            VendorGuardrailUninstallResponse,
+            wrap_response=False,
+        )

codemie_sdk/services/workflow.py

@@ -135,6 +135,8 @@ class WorkflowService:
         workflow_id: str,
         user_input: Optional[str] = None,
         file_name: Optional[str] = None,
+        propagate_headers: bool = False,
+        headers: Optional[dict[str, str]] = None,
     ) -> dict:
         """Run a workflow with optional input parameters.
 
@@ -142,11 +144,18 @@ class WorkflowService:
             workflow_id: ID of the workflow to run
             user_input: Optional user input for the workflow execution
             file_name: Optional file name for the workflow execution
+            propagate_headers: Enable propagation of X-* HTTP headers to MCP servers
+            headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
 
         Returns:
             dict: Created workflow execution details
         """
-        return self.executions(workflow_id).create(user_input, file_name)
+        return self.executions(workflow_id).create(
+            user_input=user_input,
+            file_name=file_name,
+            propagate_headers=propagate_headers,
+            headers=headers,
+        )
 
     def executions(self, workflow_id: str) -> WorkflowExecutionService:
         """Get workflow execution service for the specified workflow.

codemie_sdk/services/workflow_execution.py

@@ -46,24 +46,33 @@ class WorkflowExecutionService:
         )
 
     def create(
-        self, user_input: Optional[str] = None, file_name: Optional[str] = None
+        self,
+        user_input: Optional[str] = None,
+        file_name: Optional[str] = None,
+        propagate_headers: bool = False,
+        headers: Optional[dict[str, str]] = None,
     ) -> dict:
         """Create a new workflow execution.
 
         Args:
             user_input: Optional input data for the workflow execution.
             file_name: Optional file name associated with the workflow execution.
+            propagate_headers: Enable propagation of X-* HTTP headers to MCP servers.
+            headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
 
         Returns:
             dict: Created workflow execution details
         """
         payload = WorkflowExecutionCreateRequest(
-            user_input=user_input, file_name=file_name
+            user_input=user_input,
+            file_name=file_name,
+            propagate_headers=propagate_headers,
         )
         return self._api.post(
             f"/v1/workflows/{self._workflow_id}/executions",
             dict,
             json_data=payload.model_dump(),
+            extra_headers=headers,
         )
 
     def get(self, execution_id: str) -> WorkflowExecution:
@@ -108,17 +117,30 @@ class WorkflowExecutionService:
             f"/v1/workflows/{self._workflow_id}/executions/{execution_id}/abort", dict
         )
 
-    def resume(self, execution_id: str) -> dict:
+    def resume(
+        self,
+        execution_id: str,
+        propagate_headers: bool = False,
+        headers: Optional[dict[str, str]] = None,
+    ) -> dict:
         """Resume an interrupted workflow execution.
 
         Args:
             execution_id: ID of the execution to resume
+            propagate_headers: Enable propagation of X-* HTTP headers to MCP servers.
+            headers: Optional additional HTTP headers (e.g., X-* for MCP propagation)
 
         Returns:
             dict: Updated workflow execution details
         """
+        params = {"propagate_headers": propagate_headers}
+        # Empty body per API; passing empty dict to satisfy typing
         return self._api.put(
-            f"/v1/workflows/{self._workflow_id}/executions/{execution_id}/resume", dict
+            f"/v1/workflows/{self._workflow_id}/executions/{execution_id}/resume",
+            dict,
+            json_data={},
+            params=params,
+            extra_headers=headers,
         )
 
     def get_thoughts(

codemie_sdk/utils/http.py

@@ -1,12 +1,11 @@
 """HTTP utilities for CodeMie SDK."""
 
-from typing import TypeVar, Type, Optional, Any, Union, Dict, List, get_origin, get_args
-from pydantic import BaseModel
+from typing import TypeVar, Type, Optional, Any, Dict, List, get_origin, get_args
 import requests
 import logging
 from functools import wraps
 
-T = TypeVar("T", bound=Union[BaseModel, List[BaseModel], dict])
+T = TypeVar("T")
 
 logger = logging.getLogger(__name__)
 
@@ -55,6 +54,7 @@ class ApiRequestHandler:
             "localhost",
             "127.0.0.1",
             "0.0.0.0",
+            "192.168",
         ]
         return any(pattern in domain_lower for pattern in localhost_patterns)
 
@@ -138,7 +138,7 @@ class ApiRequestHandler:
             wrap_response: Whether response is wrapped in 'data' field
 
         Returns:
-            Parsed response object or list of objects
+            Parsed response object or list of objects, or raw Response if response_model is requests.Response
         """
         if params:
             logger.debug(f"Request params: {params}")
@@ -150,6 +150,9 @@ class ApiRequestHandler:
         )
         response.raise_for_status()
 
+        if response_model is None or response_model is requests.Response:
+            return response
+
         return self._parse_response(response, response_model, wrap_response)
 
     @log_request
@@ -161,7 +164,8 @@ class ApiRequestHandler:
         stream: bool = False,
         wrap_response: bool = True,
         raise_on_error: bool = True,
-    ) -> Union[T, requests.Response]:
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> T:
         """Makes a POST request and parses the response.
 
         Args:
@@ -171,6 +175,7 @@ class ApiRequestHandler:
             stream: Whether to return streaming response
             wrap_response: Whether response is wrapped in 'data' field
             raise_on_error: Whether to raise exception on HTTP error status codes
+            extra_headers: Optional additional HTTP headers to include (e.g., X-* for MCP propagation)
 
         Returns:
             Parsed response object/list or streaming response
@@ -178,9 +183,13 @@ class ApiRequestHandler:
         if json_data:
             logger.debug(f"Request body: {json_data}")
 
+        headers = self._get_headers()
+        if extra_headers:
+            headers.update(extra_headers)
+
         response = requests.post(
             url=f"{self._base_url}{endpoint}",
-            headers=self._get_headers(),
+            headers=headers,
             json=json_data,
             verify=self._verify_ssl,
             stream=stream,
@@ -242,6 +251,7 @@ class ApiRequestHandler:
         json_data: Dict[str, Any],
         params: Optional[Dict[str, Any]] = None,
         wrap_response: bool = True,
+        extra_headers: Optional[Dict[str, str]] = None,
     ) -> T:
         """Makes a PUT request and parses the response.
 
@@ -251,14 +261,18 @@ class ApiRequestHandler:
             json_data: JSON request body
             params: Query parameters
             wrap_response: Whether response is wrapped in 'data' field
+            extra_headers: Optional additional HTTP headers to include (e.g., X-* for MCP propagation)
 
         Returns:
             Parsed response object or list of objects
         """
         logger.debug(f"Request body: {json_data}")
+        headers = self._get_headers()
+        if extra_headers:
+            headers.update(extra_headers)
         response = requests.put(
             url=f"{self._base_url}{endpoint}",
-            headers=self._get_headers(),
+            headers=headers,
             json=json_data,
             params=params,
             verify=self._verify_ssl,