databricks-sdk 0.55.0__py3-none-any.whl → 0.57.0__py3-none-any.whl

databricks/sdk/service/aibuilder.py

@@ -0,0 +1,505 @@
+# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from ._internal import _enum, _from_dict, _repeated_dict
+
+_LOG = logging.getLogger("databricks.sdk")
+
+
+# all definitions in this file are in alphabetical order
+
+
+@dataclass
+class CancelCustomLlmOptimizationRunRequest:
+    id: Optional[str] = None
+
+
+@dataclass
+class CancelOptimizeResponse:
+    def as_dict(self) -> dict:
+        """Serializes the CancelOptimizeResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the CancelOptimizeResponse into a shallow dictionary of its immediate attributes."""
+        body = {}
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> CancelOptimizeResponse:
+        """Deserializes the CancelOptimizeResponse from a dictionary."""
+        return cls()
+
+
+@dataclass
+class CreateCustomLlmRequest:
+    name: str
+    """Name of the custom LLM. Only alphanumeric characters and dashes allowed."""
+
+    instructions: str
+    """Instructions for the custom LLM to follow"""
+
+    agent_artifact_path: Optional[str] = None
+    """Optional: UC path for agent artifacts. If you are using a dataset that you only have read
+    permissions, please provide a destination path where you have write permissions. Please provide
+    this in catalog.schema format."""
+
+    datasets: Optional[List[Dataset]] = None
+    """Datasets used for training and evaluating the model, not for inference. Currently, only 1
+    dataset is accepted."""
+
+    guidelines: Optional[List[str]] = None
+    """Guidelines for the custom LLM to adhere to"""
+
+    def as_dict(self) -> dict:
+        """Serializes the CreateCustomLlmRequest into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.agent_artifact_path is not None:
+            body["agent_artifact_path"] = self.agent_artifact_path
+        if self.datasets:
+            body["datasets"] = [v.as_dict() for v in self.datasets]
+        if self.guidelines:
+            body["guidelines"] = [v for v in self.guidelines]
+        if self.instructions is not None:
+            body["instructions"] = self.instructions
+        if self.name is not None:
+            body["name"] = self.name
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the CreateCustomLlmRequest into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.agent_artifact_path is not None:
+            body["agent_artifact_path"] = self.agent_artifact_path
+        if self.datasets:
+            body["datasets"] = self.datasets
+        if self.guidelines:
+            body["guidelines"] = self.guidelines
+        if self.instructions is not None:
+            body["instructions"] = self.instructions
+        if self.name is not None:
+            body["name"] = self.name
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> CreateCustomLlmRequest:
+        """Deserializes the CreateCustomLlmRequest from a dictionary."""
+        return cls(
+            agent_artifact_path=d.get("agent_artifact_path", None),
+            datasets=_repeated_dict(d, "datasets", Dataset),
+            guidelines=d.get("guidelines", None),
+            instructions=d.get("instructions", None),
+            name=d.get("name", None),
+        )
+
+
+@dataclass
+class CustomLlm:
+    name: str
+    """Name of the custom LLM"""
+
+    instructions: str
+    """Instructions for the custom LLM to follow"""
+
+    optimization_state: State
+    """If optimization is kicked off, tracks the state of the custom LLM"""
+
+    agent_artifact_path: Optional[str] = None
+
+    creation_time: Optional[str] = None
+    """Creation timestamp of the custom LLM"""
+
+    creator: Optional[str] = None
+    """Creator of the custom LLM"""
+
+    datasets: Optional[List[Dataset]] = None
+    """Datasets used for training and evaluating the model, not for inference"""
+
+    endpoint_name: Optional[str] = None
+    """Name of the endpoint that will be used to serve the custom LLM"""
+
+    guidelines: Optional[List[str]] = None
+    """Guidelines for the custom LLM to adhere to"""
+
+    id: Optional[str] = None
+
+    def as_dict(self) -> dict:
+        """Serializes the CustomLlm into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.agent_artifact_path is not None:
+            body["agent_artifact_path"] = self.agent_artifact_path
+        if self.creation_time is not None:
+            body["creation_time"] = self.creation_time
+        if self.creator is not None:
+            body["creator"] = self.creator
+        if self.datasets:
+            body["datasets"] = [v.as_dict() for v in self.datasets]
+        if self.endpoint_name is not None:
+            body["endpoint_name"] = self.endpoint_name
+        if self.guidelines:
+            body["guidelines"] = [v for v in self.guidelines]
+        if self.id is not None:
+            body["id"] = self.id
+        if self.instructions is not None:
+            body["instructions"] = self.instructions
+        if self.name is not None:
+            body["name"] = self.name
+        if self.optimization_state is not None:
+            body["optimization_state"] = self.optimization_state.value
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the CustomLlm into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.agent_artifact_path is not None:
+            body["agent_artifact_path"] = self.agent_artifact_path
+        if self.creation_time is not None:
+            body["creation_time"] = self.creation_time
+        if self.creator is not None:
+            body["creator"] = self.creator
+        if self.datasets:
+            body["datasets"] = self.datasets
+        if self.endpoint_name is not None:
+            body["endpoint_name"] = self.endpoint_name
+        if self.guidelines:
+            body["guidelines"] = self.guidelines
+        if self.id is not None:
+            body["id"] = self.id
+        if self.instructions is not None:
+            body["instructions"] = self.instructions
+        if self.name is not None:
+            body["name"] = self.name
+        if self.optimization_state is not None:
+            body["optimization_state"] = self.optimization_state
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> CustomLlm:
+        """Deserializes the CustomLlm from a dictionary."""
+        return cls(
+            agent_artifact_path=d.get("agent_artifact_path", None),
+            creation_time=d.get("creation_time", None),
+            creator=d.get("creator", None),
+            datasets=_repeated_dict(d, "datasets", Dataset),
+            endpoint_name=d.get("endpoint_name", None),
+            guidelines=d.get("guidelines", None),
+            id=d.get("id", None),
+            instructions=d.get("instructions", None),
+            name=d.get("name", None),
+            optimization_state=_enum(d, "optimization_state", State),
+        )
+
+
+@dataclass
+class Dataset:
+    table: Table
+
+    def as_dict(self) -> dict:
+        """Serializes the Dataset into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.table:
+            body["table"] = self.table.as_dict()
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the Dataset into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.table:
+            body["table"] = self.table
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> Dataset:
+        """Deserializes the Dataset from a dictionary."""
+        return cls(table=_from_dict(d, "table", Table))
+
+
+@dataclass
+class DeleteCustomLlmResponse:
+    def as_dict(self) -> dict:
+        """Serializes the DeleteCustomLlmResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the DeleteCustomLlmResponse into a shallow dictionary of its immediate attributes."""
+        body = {}
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> DeleteCustomLlmResponse:
+        """Deserializes the DeleteCustomLlmResponse from a dictionary."""
+        return cls()
+
+
+@dataclass
+class StartCustomLlmOptimizationRunRequest:
+    id: Optional[str] = None
+    """The Id of the tile."""
+
+
+class State(Enum):
+    """States of Custom LLM optimization lifecycle."""
+
+    CANCELLED = "CANCELLED"
+    COMPLETED = "COMPLETED"
+    CREATED = "CREATED"
+    FAILED = "FAILED"
+    PENDING = "PENDING"
+    RUNNING = "RUNNING"
+
+
+@dataclass
+class Table:
+    table_path: str
+    """Full UC table path in catalog.schema.table_name format"""
+
+    request_col: str
+    """Name of the request column"""
+
+    response_col: Optional[str] = None
+    """Optional: Name of the response column if the data is labeled"""
+
+    def as_dict(self) -> dict:
+        """Serializes the Table into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.request_col is not None:
+            body["request_col"] = self.request_col
+        if self.response_col is not None:
+            body["response_col"] = self.response_col
+        if self.table_path is not None:
+            body["table_path"] = self.table_path
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the Table into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.request_col is not None:
+            body["request_col"] = self.request_col
+        if self.response_col is not None:
+            body["response_col"] = self.response_col
+        if self.table_path is not None:
+            body["table_path"] = self.table_path
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> Table:
+        """Deserializes the Table from a dictionary."""
+        return cls(
+            request_col=d.get("request_col", None),
+            response_col=d.get("response_col", None),
+            table_path=d.get("table_path", None),
+        )
+
+
+@dataclass
+class UpdateCustomLlmRequest:
+    custom_llm: CustomLlm
+    """The CustomLlm containing the fields which should be updated."""
+
+    update_mask: str
+    """The list of the CustomLlm fields to update. These should correspond to the values (or lack
+    thereof) present in `custom_llm`.
+    
+    The field mask must be a single string, with multiple fields separated by commas (no spaces).
+    The field path is relative to the resource object, using a dot (`.`) to navigate sub-fields
+    (e.g., `author.given_name`). Specification of elements in sequence or map fields is not allowed,
+    as only the entire collection field can be specified. Field names must exactly match the
+    resource field names.
+    
+    A field mask of `*` indicates full replacement. It’s recommended to always explicitly list the
+    fields being updated and avoid using `*` wildcards, as it can lead to unintended results if the
+    API changes in the future."""
+
+    id: Optional[str] = None
+    """The id of the custom llm"""
+
+    def as_dict(self) -> dict:
+        """Serializes the UpdateCustomLlmRequest into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.custom_llm:
+            body["custom_llm"] = self.custom_llm.as_dict()
+        if self.id is not None:
+            body["id"] = self.id
+        if self.update_mask is not None:
+            body["update_mask"] = self.update_mask
+        return body
+
+    def as_shallow_dict(self) -> dict:
+        """Serializes the UpdateCustomLlmRequest into a shallow dictionary of its immediate attributes."""
+        body = {}
+        if self.custom_llm:
+            body["custom_llm"] = self.custom_llm
+        if self.id is not None:
+            body["id"] = self.id
+        if self.update_mask is not None:
+            body["update_mask"] = self.update_mask
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> UpdateCustomLlmRequest:
+        """Deserializes the UpdateCustomLlmRequest from a dictionary."""
+        return cls(
+            custom_llm=_from_dict(d, "custom_llm", CustomLlm),
+            id=d.get("id", None),
+            update_mask=d.get("update_mask", None),
+        )
+
+
+class AiBuilderAPI:
+    """The Custom LLMs service manages state and powers the UI for the Custom LLM product."""
+
+    def __init__(self, api_client):
+        self._api = api_client
+
+    def cancel_optimize(self, id: str):
+        """Cancel a Custom LLM Optimization Run.
+
+        :param id: str
+
+
+        """
+
+        headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+        }
+
+        self._api.do("POST", f"/api/2.0/custom-llms/{id}/optimize/cancel", headers=headers)
+
+    def create_custom_llm(
+        self,
+        name: str,
+        instructions: str,
+        *,
+        agent_artifact_path: Optional[str] = None,
+        datasets: Optional[List[Dataset]] = None,
+        guidelines: Optional[List[str]] = None,
+    ) -> CustomLlm:
+        """Create a Custom LLM.
+
+        :param name: str
+          Name of the custom LLM. Only alphanumeric characters and dashes allowed.
+        :param instructions: str
+          Instructions for the custom LLM to follow
+        :param agent_artifact_path: str (optional)
+          Optional: UC path for agent artifacts. If you are using a dataset that you only have read
+          permissions, please provide a destination path where you have write permissions. Please provide this
+          in catalog.schema format.
+        :param datasets: List[:class:`Dataset`] (optional)
+          Datasets used for training and evaluating the model, not for inference. Currently, only 1 dataset is
+          accepted.
+        :param guidelines: List[str] (optional)
+          Guidelines for the custom LLM to adhere to
+
+        :returns: :class:`CustomLlm`
+        """
+        body = {}
+        if agent_artifact_path is not None:
+            body["agent_artifact_path"] = agent_artifact_path
+        if datasets is not None:
+            body["datasets"] = [v.as_dict() for v in datasets]
+        if guidelines is not None:
+            body["guidelines"] = [v for v in guidelines]
+        if instructions is not None:
+            body["instructions"] = instructions
+        if name is not None:
+            body["name"] = name
+        headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+        }
+
+        res = self._api.do("POST", "/api/2.0/custom-llms", body=body, headers=headers)
+        return CustomLlm.from_dict(res)
+
+    def delete_custom_llm(self, id: str):
+        """Delete a Custom LLM.
+
+        :param id: str
+          The id of the custom llm
+
+
+        """
+
+        headers = {
+            "Accept": "application/json",
+        }
+
+        self._api.do("DELETE", f"/api/2.0/custom-lms/{id}", headers=headers)
+
+    def get_custom_llm(self, id: str) -> CustomLlm:
+        """Get a Custom LLM.
+
+        :param id: str
+          The id of the custom llm
+
+        :returns: :class:`CustomLlm`
+        """
+
+        headers = {
+            "Accept": "application/json",
+        }
+
+        res = self._api.do("GET", f"/api/2.0/custom-llms/{id}", headers=headers)
+        return CustomLlm.from_dict(res)
+
+    def start_optimize(self, id: str) -> CustomLlm:
+        """Start a Custom LLM Optimization Run.
+
+        :param id: str
+          The Id of the tile.
+
+        :returns: :class:`CustomLlm`
+        """
+
+        headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+        }
+
+        res = self._api.do("POST", f"/api/2.0/custom-llms/{id}/optimize", headers=headers)
+        return CustomLlm.from_dict(res)
+
+    def update_custom_llm(self, id: str, custom_llm: CustomLlm, update_mask: str) -> CustomLlm:
+        """Update a Custom LLM.
+
+        :param id: str
+          The id of the custom llm
+        :param custom_llm: :class:`CustomLlm`
+          The CustomLlm containing the fields which should be updated.
+        :param update_mask: str
+          The list of the CustomLlm fields to update. These should correspond to the values (or lack thereof)
+          present in `custom_llm`.
+
+          The field mask must be a single string, with multiple fields separated by commas (no spaces). The
+          field path is relative to the resource object, using a dot (`.`) to navigate sub-fields (e.g.,
+          `author.given_name`). Specification of elements in sequence or map fields is not allowed, as only
+          the entire collection field can be specified. Field names must exactly match the resource field
+          names.
+
+          A field mask of `*` indicates full replacement. It’s recommended to always explicitly list the
+          fields being updated and avoid using `*` wildcards, as it can lead to unintended results if the API
+          changes in the future.
+
+        :returns: :class:`CustomLlm`
+        """
+        body = {}
+        if custom_llm is not None:
+            body["custom_llm"] = custom_llm.as_dict()
+        if update_mask is not None:
+            body["update_mask"] = update_mask
+        headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+        }
+
+        res = self._api.do("PATCH", f"/api/2.0/custom-llms/{id}", body=body, headers=headers)
+        return CustomLlm.from_dict(res)

databricks/sdk/service/apps.py

@@ -1232,9 +1232,7 @@ class AppsAPI:
         raise TimeoutError(f"timed out after {timeout}: {status_message}")
 
     def create(self, app: App, *, no_compute: Optional[bool] = None) -> Wait[App]:
-        """Create an app.
-
-        Creates a new app.
+        """Creates a new app.
 
         :param app: :class:`App`
         :param no_compute: bool (optional)
@@ -1260,9 +1258,7 @@ class AppsAPI:
         return self.create(app=app, no_compute=no_compute).result(timeout=timeout)
 
     def delete(self, name: str) -> App:
-        """Delete an app.
-
-        Deletes an app.
+        """Deletes an app.
 
         :param name: str
           The name of the app.
@@ -1278,9 +1274,7 @@ class AppsAPI:
         return App.from_dict(res)
 
     def deploy(self, app_name: str, app_deployment: AppDeployment) -> Wait[AppDeployment]:
-        """Create an app deployment.
-
-        Creates an app deployment for the app with the supplied name.
+        """Creates an app deployment for the app with the supplied name.
 
         :param app_name: str
           The name of the app.
@@ -1310,9 +1304,7 @@ class AppsAPI:
         return self.deploy(app_deployment=app_deployment, app_name=app_name).result(timeout=timeout)
 
     def get(self, name: str) -> App:
-        """Get an app.
-
-        Retrieves information for the app with the supplied name.
+        """Retrieves information for the app with the supplied name.
 
         :param name: str
           The name of the app.
@@ -1328,9 +1320,7 @@ class AppsAPI:
         return App.from_dict(res)
 
     def get_deployment(self, app_name: str, deployment_id: str) -> AppDeployment:
-        """Get an app deployment.
-
-        Retrieves information for the app deployment with the supplied name and deployment id.
+        """Retrieves information for the app deployment with the supplied name and deployment id.
 
         :param app_name: str
           The name of the app.
@@ -1348,9 +1338,7 @@ class AppsAPI:
         return AppDeployment.from_dict(res)
 
     def get_permission_levels(self, app_name: str) -> GetAppPermissionLevelsResponse:
-        """Get app permission levels.
-
-        Gets the permission levels that a user can have on an object.
+        """Gets the permission levels that a user can have on an object.
 
         :param app_name: str
           The app for which to get or manage permissions.
@@ -1366,9 +1354,7 @@ class AppsAPI:
         return GetAppPermissionLevelsResponse.from_dict(res)
 
     def get_permissions(self, app_name: str) -> AppPermissions:
-        """Get app permissions.
-
-        Gets the permissions of an app. Apps can inherit permissions from their root object.
+        """Gets the permissions of an app. Apps can inherit permissions from their root object.
 
         :param app_name: str
           The app for which to get or manage permissions.
@@ -1384,9 +1370,7 @@ class AppsAPI:
         return AppPermissions.from_dict(res)
 
     def list(self, *, page_size: Optional[int] = None, page_token: Optional[str] = None) -> Iterator[App]:
-        """List apps.
-
-        Lists all apps in the workspace.
+        """Lists all apps in the workspace.
 
         :param page_size: int (optional)
           Upper bound for items returned.
@@ -1417,9 +1401,7 @@ class AppsAPI:
     def list_deployments(
         self, app_name: str, *, page_size: Optional[int] = None, page_token: Optional[str] = None
     ) -> Iterator[AppDeployment]:
-        """List app deployments.
-
-        Lists all app deployments for the app with the supplied name.
+        """Lists all app deployments for the app with the supplied name.
 
         :param app_name: str
           The name of the app.
@@ -1452,9 +1434,7 @@ class AppsAPI:
     def set_permissions(
         self, app_name: str, *, access_control_list: Optional[List[AppAccessControlRequest]] = None
     ) -> AppPermissions:
-        """Set app permissions.
-
-        Sets permissions on an object, replacing existing permissions if they exist. Deletes all direct
+        """Sets permissions on an object, replacing existing permissions if they exist. Deletes all direct
         permissions if none are specified. Objects can inherit permissions from their root object.
 
         :param app_name: str
@@ -1475,9 +1455,7 @@ class AppsAPI:
         return AppPermissions.from_dict(res)
 
     def start(self, name: str) -> Wait[App]:
-        """Start an app.
-
-        Start the last active deployment of the app in the workspace.
+        """Start the last active deployment of the app in the workspace.
 
         :param name: str
           The name of the app.
@@ -1499,9 +1477,7 @@ class AppsAPI:
         return self.start(name=name).result(timeout=timeout)
 
     def stop(self, name: str) -> Wait[App]:
-        """Stop an app.
-
-        Stops the active deployment of the app in the workspace.
+        """Stops the active deployment of the app in the workspace.
 
         :param name: str
           The name of the app.
@@ -1523,9 +1499,7 @@ class AppsAPI:
         return self.stop(name=name).result(timeout=timeout)
 
     def update(self, name: str, app: App) -> App:
-        """Update an app.
-
-        Updates the app with the supplied name.
+        """Updates the app with the supplied name.
 
         :param name: str
           The name of the app. The name must contain only lowercase alphanumeric characters and hyphens. It
@@ -1546,9 +1520,7 @@ class AppsAPI:
     def update_permissions(
         self, app_name: str, *, access_control_list: Optional[List[AppAccessControlRequest]] = None
     ) -> AppPermissions:
-        """Update app permissions.
-
-        Updates the permissions on an app. Apps can inherit permissions from their root object.
+        """Updates the permissions on an app. Apps can inherit permissions from their root object.
 
         :param app_name: str
           The app for which to get or manage permissions.