370 lines
12 KiB
Python
370 lines
12 KiB
Python
# 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 databricks.sdk.client_types import HostType
|
||
from databricks.sdk.service._internal import _enum, _from_dict, _repeated_dict
|
||
|
||
_LOG = logging.getLogger("databricks.sdk")
|
||
|
||
|
||
# all definitions in this file are in alphabetical order
|
||
|
||
|
||
@dataclass
|
||
class CustomLlm:
|
||
name: str
|
||
"""Name of the custom LLM"""
|
||
|
||
instructions: str
|
||
"""Instructions for the custom LLM to follow"""
|
||
|
||
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
|
||
|
||
optimization_state: Optional[State] = None
|
||
"""If optimization is kicked off, tracks the state of the custom LLM"""
|
||
|
||
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))
|
||
|
||
|
||
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),
|
||
)
|
||
|
||
|
||
class AgentBricksAPI:
|
||
"""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",
|
||
}
|
||
|
||
cfg = self._api._cfg
|
||
if cfg.host_type == HostType.UNIFIED and cfg.workspace_id:
|
||
headers["X-Databricks-Org-Id"] = cfg.workspace_id
|
||
|
||
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)
|
||
This will soon be deprecated!! 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",
|
||
}
|
||
|
||
cfg = self._api._cfg
|
||
if cfg.host_type == HostType.UNIFIED and cfg.workspace_id:
|
||
headers["X-Databricks-Org-Id"] = cfg.workspace_id
|
||
|
||
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",
|
||
}
|
||
|
||
cfg = self._api._cfg
|
||
if cfg.host_type == HostType.UNIFIED and cfg.workspace_id:
|
||
headers["X-Databricks-Org-Id"] = cfg.workspace_id
|
||
|
||
self._api.do("DELETE", f"/api/2.0/custom-llms/{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",
|
||
}
|
||
|
||
cfg = self._api._cfg
|
||
if cfg.host_type == HostType.UNIFIED and cfg.workspace_id:
|
||
headers["X-Databricks-Org-Id"] = cfg.workspace_id
|
||
|
||
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",
|
||
}
|
||
|
||
cfg = self._api._cfg
|
||
if cfg.host_type == HostType.UNIFIED and cfg.workspace_id:
|
||
headers["X-Databricks-Org-Id"] = cfg.workspace_id
|
||
|
||
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",
|
||
}
|
||
|
||
cfg = self._api._cfg
|
||
if cfg.host_type == HostType.UNIFIED and cfg.workspace_id:
|
||
headers["X-Databricks-Org-Id"] = cfg.workspace_id
|
||
|
||
res = self._api.do("PATCH", f"/api/2.0/custom-llms/{id}", body=body, headers=headers)
|
||
return CustomLlm.from_dict(res)
|