databricks-sdk 0.19.1__py3-none-any.whl → 0.21.0__py3-none-any.whl

databricks/sdk/service/vectorsearch.py

@@ -61,6 +61,9 @@ class CreateVectorIndexRequest:
     name: str
     """Name of the index"""
 
+    endpoint_name: str
+    """Name of the endpoint to be used for serving the index"""
+
     primary_key: str
     """Primary key of the index"""
 
@@ -78,9 +81,6 @@ class CreateVectorIndexRequest:
     direct_access_index_spec: Optional[DirectAccessVectorIndexSpec] = None
     """Specification for Direct Vector Access Index. Required if `index_type` is `DIRECT_ACCESS`."""
 
-    endpoint_name: Optional[str] = None
-    """Name of the endpoint to be used for serving the index"""
-
     def as_dict(self) -> dict:
         """Serializes the CreateVectorIndexRequest into a dictionary suitable for use as a JSON request body."""
         body = {}
@@ -161,20 +161,20 @@ class DeleteDataVectorIndexRequest:
     primary_keys: List[str]
     """List of primary keys for the data to be deleted."""
 
-    name: Optional[str] = None
+    index_name: Optional[str] = None
     """Name of the vector index where data is to be deleted. Must be a Direct Vector Access Index."""
 
     def as_dict(self) -> dict:
         """Serializes the DeleteDataVectorIndexRequest into a dictionary suitable for use as a JSON request body."""
         body = {}
-        if self.name is not None: body['name'] = self.name
+        if self.index_name is not None: body['index_name'] = self.index_name
         if self.primary_keys: body['primary_keys'] = [v for v in self.primary_keys]
         return body
 
     @classmethod
     def from_dict(cls, d: Dict[str, any]) -> DeleteDataVectorIndexRequest:
         """Deserializes the DeleteDataVectorIndexRequest from a dictionary."""
-        return cls(name=d.get('name', None), primary_keys=d.get('primary_keys', None))
+        return cls(index_name=d.get('index_name', None), primary_keys=d.get('primary_keys', None))
 
 
 @dataclass
@@ -201,6 +201,34 @@ class DeleteDataVectorIndexResponse:
                    status=_enum(d, 'status', DeleteDataStatus))
 
 
+@dataclass
+class DeleteEndpointResponse:
+
+    def as_dict(self) -> dict:
+        """Serializes the DeleteEndpointResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> DeleteEndpointResponse:
+        """Deserializes the DeleteEndpointResponse from a dictionary."""
+        return cls()
+
+
+@dataclass
+class DeleteIndexResponse:
+
+    def as_dict(self) -> dict:
+        """Serializes the DeleteIndexResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> DeleteIndexResponse:
+        """Deserializes the DeleteIndexResponse from a dictionary."""
+        return cls()
+
+
 @dataclass
 class DeltaSyncVectorIndexSpecRequest:
     embedding_source_columns: Optional[List[EmbeddingSourceColumn]] = None
@@ -319,41 +347,25 @@ class DirectAccessVectorIndexSpec:
 
 
 @dataclass
-class EmbeddingConfig:
+class EmbeddingSourceColumn:
     embedding_model_endpoint_name: Optional[str] = None
     """Name of the embedding model endpoint"""
 
-    def as_dict(self) -> dict:
-        """Serializes the EmbeddingConfig into a dictionary suitable for use as a JSON request body."""
-        body = {}
-        if self.embedding_model_endpoint_name is not None:
-            body['embedding_model_endpoint_name'] = self.embedding_model_endpoint_name
-        return body
-
-    @classmethod
-    def from_dict(cls, d: Dict[str, any]) -> EmbeddingConfig:
-        """Deserializes the EmbeddingConfig from a dictionary."""
-        return cls(embedding_model_endpoint_name=d.get('embedding_model_endpoint_name', None))
-
-
-@dataclass
-class EmbeddingSourceColumn:
-    embedding_config: Optional[EmbeddingConfig] = None
-
     name: Optional[str] = None
     """Name of the column"""
 
     def as_dict(self) -> dict:
         """Serializes the EmbeddingSourceColumn into a dictionary suitable for use as a JSON request body."""
         body = {}
-        if self.embedding_config: body['embedding_config'] = self.embedding_config.as_dict()
+        if self.embedding_model_endpoint_name is not None:
+            body['embedding_model_endpoint_name'] = self.embedding_model_endpoint_name
         if self.name is not None: body['name'] = self.name
         return body
 
     @classmethod
     def from_dict(cls, d: Dict[str, any]) -> EmbeddingSourceColumn:
         """Deserializes the EmbeddingSourceColumn from a dictionary."""
-        return cls(embedding_config=_from_dict(d, 'embedding_config', EmbeddingConfig),
+        return cls(embedding_model_endpoint_name=d.get('embedding_model_endpoint_name', None),
                    name=d.get('name', None))
 
 
@@ -598,6 +610,9 @@ class QueryVectorIndexRequest:
     """Query vector. Required for Direct Vector Access Index and Delta Sync Index using self-managed
     vectors."""
 
+    score_threshold: Optional[float] = None
+    """Threshold for the approximate nearest neighbor search. Defaults to 0.0."""
+
     def as_dict(self) -> dict:
         """Serializes the QueryVectorIndexRequest into a dictionary suitable for use as a JSON request body."""
         body = {}
@@ -607,6 +622,7 @@ class QueryVectorIndexRequest:
         if self.num_results is not None: body['num_results'] = self.num_results
         if self.query_text is not None: body['query_text'] = self.query_text
         if self.query_vector: body['query_vector'] = [v for v in self.query_vector]
+        if self.score_threshold is not None: body['score_threshold'] = self.score_threshold
         return body
 
     @classmethod
@@ -617,7 +633,8 @@ class QueryVectorIndexRequest:
                    index_name=d.get('index_name', None),
                    num_results=d.get('num_results', None),
                    query_text=d.get('query_text', None),
-                   query_vector=d.get('query_vector', None))
+                   query_vector=d.get('query_vector', None),
+                   score_threshold=d.get('score_threshold', None))
 
 
 @dataclass
@@ -688,6 +705,20 @@ class ResultManifest:
         return cls(column_count=d.get('column_count', None), columns=_repeated_dict(d, 'columns', ColumnInfo))
 
 
+@dataclass
+class SyncIndexResponse:
+
+    def as_dict(self) -> dict:
+        """Serializes the SyncIndexResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> SyncIndexResponse:
+        """Deserializes the SyncIndexResponse from a dictionary."""
+        return cls()
+
+
 @dataclass
 class UpsertDataResult:
     """Result of the upsert or delete operation."""
@@ -727,20 +758,20 @@ class UpsertDataVectorIndexRequest:
     inputs_json: str
     """JSON string representing the data to be upserted."""
 
-    name: Optional[str] = None
+    index_name: Optional[str] = None
     """Name of the vector index where data is to be upserted. Must be a Direct Vector Access Index."""
 
     def as_dict(self) -> dict:
         """Serializes the UpsertDataVectorIndexRequest into a dictionary suitable for use as a JSON request body."""
         body = {}
+        if self.index_name is not None: body['index_name'] = self.index_name
         if self.inputs_json is not None: body['inputs_json'] = self.inputs_json
-        if self.name is not None: body['name'] = self.name
         return body
 
     @classmethod
     def from_dict(cls, d: Dict[str, any]) -> UpsertDataVectorIndexRequest:
         """Deserializes the UpsertDataVectorIndexRequest from a dictionary."""
-        return cls(inputs_json=d.get('inputs_json', None), name=d.get('name', None))
+        return cls(index_name=d.get('index_name', None), inputs_json=d.get('inputs_json', None))
 
 
 @dataclass
@@ -772,9 +803,9 @@ class VectorIndex:
     creator: Optional[str] = None
     """The user who created the index."""
 
-    delta_sync_vector_index_spec: Optional[DeltaSyncVectorIndexSpecResponse] = None
+    delta_sync_index_spec: Optional[DeltaSyncVectorIndexSpecResponse] = None
 
-    direct_access_vector_index_spec: Optional[DirectAccessVectorIndexSpec] = None
+    direct_access_index_spec: Optional[DirectAccessVectorIndexSpec] = None
 
     endpoint_name: Optional[str] = None
     """Name of the endpoint associated with the index"""
@@ -799,10 +830,9 @@ class VectorIndex:
         """Serializes the VectorIndex into a dictionary suitable for use as a JSON request body."""
         body = {}
         if self.creator is not None: body['creator'] = self.creator
-        if self.delta_sync_vector_index_spec:
-            body['delta_sync_vector_index_spec'] = self.delta_sync_vector_index_spec.as_dict()
-        if self.direct_access_vector_index_spec:
-            body['direct_access_vector_index_spec'] = self.direct_access_vector_index_spec.as_dict()
+        if self.delta_sync_index_spec: body['delta_sync_index_spec'] = self.delta_sync_index_spec.as_dict()
+        if self.direct_access_index_spec:
+            body['direct_access_index_spec'] = self.direct_access_index_spec.as_dict()
         if self.endpoint_name is not None: body['endpoint_name'] = self.endpoint_name
         if self.index_type is not None: body['index_type'] = self.index_type.value
         if self.name is not None: body['name'] = self.name
@@ -814,10 +844,10 @@ class VectorIndex:
     def from_dict(cls, d: Dict[str, any]) -> VectorIndex:
         """Deserializes the VectorIndex from a dictionary."""
         return cls(creator=d.get('creator', None),
-                   delta_sync_vector_index_spec=_from_dict(d, 'delta_sync_vector_index_spec',
-                                                           DeltaSyncVectorIndexSpecResponse),
-                   direct_access_vector_index_spec=_from_dict(d, 'direct_access_vector_index_spec',
-                                                              DirectAccessVectorIndexSpec),
+                   delta_sync_index_spec=_from_dict(d, 'delta_sync_index_spec',
+                                                    DeltaSyncVectorIndexSpecResponse),
+                   direct_access_index_spec=_from_dict(d, 'direct_access_index_spec',
+                                                       DirectAccessVectorIndexSpec),
                    endpoint_name=d.get('endpoint_name', None),
                    index_type=_enum(d, 'index_type', VectorIndexType),
                    name=d.get('name', None),
@@ -926,6 +956,7 @@ class VectorSearchEndpointsAPI:
         if endpoint_type is not None: body['endpoint_type'] = endpoint_type.value
         if name is not None: body['name'] = name
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         op_response = self._api.do('POST', '/api/2.0/vector-search/endpoints', body=body, headers=headers)
         return Wait(self.wait_get_endpoint_vector_search_endpoint_online,
                     response=EndpointInfo.from_dict(op_response),
@@ -935,18 +966,17 @@ class VectorSearchEndpointsAPI:
                                  timeout=timedelta(minutes=20)) -> EndpointInfo:
         return self.create_endpoint(endpoint_type=endpoint_type, name=name).result(timeout=timeout)
 
-    def delete_endpoint(self, endpoint_name: str, name: str):
+    def delete_endpoint(self, endpoint_name: str):
         """Delete an endpoint.
         
         :param endpoint_name: str
           Name of the endpoint
-        :param name: str
-          Name of the endpoint to delete
         
         
         """
 
         headers = {}
+
         self._api.do('DELETE', f'/api/2.0/vector-search/endpoints/{endpoint_name}', headers=headers)
 
     def get_endpoint(self, endpoint_name: str) -> EndpointInfo:
@@ -959,6 +989,7 @@ class VectorSearchEndpointsAPI:
         """
 
         headers = {'Accept': 'application/json', }
+
         res = self._api.do('GET', f'/api/2.0/vector-search/endpoints/{endpoint_name}', headers=headers)
         return EndpointInfo.from_dict(res)
 
@@ -997,20 +1028,24 @@ class VectorSearchIndexesAPI:
     def __init__(self, api_client):
         self._api = api_client
 
-    def create_index(self,
-                     name: str,
-                     primary_key: str,
-                     index_type: VectorIndexType,
-                     *,
-                     delta_sync_index_spec: Optional[DeltaSyncVectorIndexSpecRequest] = None,
-                     direct_access_index_spec: Optional[DirectAccessVectorIndexSpec] = None,
-                     endpoint_name: Optional[str] = None) -> CreateVectorIndexResponse:
+    def create_index(
+            self,
+            name: str,
+            endpoint_name: str,
+            primary_key: str,
+            index_type: VectorIndexType,
+            *,
+            delta_sync_index_spec: Optional[DeltaSyncVectorIndexSpecRequest] = None,
+            direct_access_index_spec: Optional[DirectAccessVectorIndexSpec] = None
+    ) -> CreateVectorIndexResponse:
         """Create an index.
         
         Create a new index.
         
         :param name: str
           Name of the index
+        :param endpoint_name: str
+          Name of the endpoint to be used for serving the index
         :param primary_key: str
           Primary key of the index
         :param index_type: :class:`VectorIndexType`
@@ -1024,8 +1059,6 @@ class VectorSearchIndexesAPI:
           Specification for Delta Sync Index. Required if `index_type` is `DELTA_SYNC`.
         :param direct_access_index_spec: :class:`DirectAccessVectorIndexSpec` (optional)
           Specification for Direct Vector Access Index. Required if `index_type` is `DIRECT_ACCESS`.
-        :param endpoint_name: str (optional)
-          Name of the endpoint to be used for serving the index
         
         :returns: :class:`CreateVectorIndexResponse`
         """
@@ -1038,15 +1071,17 @@ class VectorSearchIndexesAPI:
         if name is not None: body['name'] = name
         if primary_key is not None: body['primary_key'] = primary_key
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         res = self._api.do('POST', '/api/2.0/vector-search/indexes', body=body, headers=headers)
         return CreateVectorIndexResponse.from_dict(res)
 
-    def delete_data_vector_index(self, name: str, primary_keys: List[str]) -> DeleteDataVectorIndexResponse:
+    def delete_data_vector_index(self, index_name: str,
+                                 primary_keys: List[str]) -> DeleteDataVectorIndexResponse:
         """Delete data from index.
         
         Handles the deletion of data from a specified vector index.
         
-        :param name: str
+        :param index_name: str
           Name of the vector index where data is to be deleted. Must be a Direct Vector Access Index.
         :param primary_keys: List[str]
           List of primary keys for the data to be deleted.
@@ -1056,8 +1091,9 @@ class VectorSearchIndexesAPI:
         body = {}
         if primary_keys is not None: body['primary_keys'] = [v for v in primary_keys]
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         res = self._api.do('POST',
-                           f'/api/2.0/vector-search/indexes/{name}/delete-data',
+                           f'/api/2.0/vector-search/indexes/{index_name}/delete-data',
                            body=body,
                            headers=headers)
         return DeleteDataVectorIndexResponse.from_dict(res)
@@ -1074,6 +1110,7 @@ class VectorSearchIndexesAPI:
         """
 
         headers = {}
+
         self._api.do('DELETE', f'/api/2.0/vector-search/indexes/{index_name}', headers=headers)
 
     def get_index(self, index_name: str) -> VectorIndex:
@@ -1088,6 +1125,7 @@ class VectorSearchIndexesAPI:
         """
 
         headers = {'Accept': 'application/json', }
+
         res = self._api.do('GET', f'/api/2.0/vector-search/indexes/{index_name}', headers=headers)
         return VectorIndex.from_dict(res)
 
@@ -1128,7 +1166,8 @@ class VectorSearchIndexesAPI:
                     filters_json: Optional[str] = None,
                     num_results: Optional[int] = None,
                     query_text: Optional[str] = None,
-                    query_vector: Optional[List[float]] = None) -> QueryVectorIndexResponse:
+                    query_vector: Optional[List[float]] = None,
+                    score_threshold: Optional[float] = None) -> QueryVectorIndexResponse:
         """Query an index.
         
         Query the specified vector index.
@@ -1150,6 +1189,8 @@ class VectorSearchIndexesAPI:
         :param query_vector: List[float] (optional)
           Query vector. Required for Direct Vector Access Index and Delta Sync Index using self-managed
           vectors.
+        :param score_threshold: float (optional)
+          Threshold for the approximate nearest neighbor search. Defaults to 0.0.
         
         :returns: :class:`QueryVectorIndexResponse`
         """
@@ -1159,7 +1200,9 @@ class VectorSearchIndexesAPI:
         if num_results is not None: body['num_results'] = num_results
         if query_text is not None: body['query_text'] = query_text
         if query_vector is not None: body['query_vector'] = [v for v in query_vector]
+        if score_threshold is not None: body['score_threshold'] = score_threshold
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         res = self._api.do('POST',
                            f'/api/2.0/vector-search/indexes/{index_name}/query',
                            body=body,
@@ -1178,14 +1221,15 @@ class VectorSearchIndexesAPI:
         """
 
         headers = {}
+
         self._api.do('POST', f'/api/2.0/vector-search/indexes/{index_name}/sync', headers=headers)
 
-    def upsert_data_vector_index(self, name: str, inputs_json: str) -> UpsertDataVectorIndexResponse:
+    def upsert_data_vector_index(self, index_name: str, inputs_json: str) -> UpsertDataVectorIndexResponse:
         """Upsert data into an index.
         
         Handles the upserting of data into a specified vector index.
         
-        :param name: str
+        :param index_name: str
           Name of the vector index where data is to be upserted. Must be a Direct Vector Access Index.
         :param inputs_json: str
           JSON string representing the data to be upserted.
@@ -1195,8 +1239,9 @@ class VectorSearchIndexesAPI:
         body = {}
         if inputs_json is not None: body['inputs_json'] = inputs_json
         headers = {'Accept': 'application/json', 'Content-Type': 'application/json', }
+
         res = self._api.do('POST',
-                           f'/api/2.0/vector-search/indexes/{name}/upsert-data',
+                           f'/api/2.0/vector-search/indexes/{index_name}/upsert-data',
                            body=body,
                            headers=headers)
         return UpsertDataVectorIndexResponse.from_dict(res)