Repository navigation
Expand file tree
/
Copy pathvector_storage.py
More file actions
701 lines (575 loc) · 29.8 KB
/
Copy pathvector_storage.py
File metadata and controls
701 lines (575 loc) · 29.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import logging
import sys
from abc import abstractmethod
from collections.abc import Awaitable, Callable, Mapping, Sequence
from contextlib import suppress
from typing import Any, ClassVar, Generic, TypeVar, overload
from pydantic import BaseModel, Field, model_validator
from semantic_kernel.data.record_definition import (
SerializeMethodProtocol,
ToDictMethodProtocol,
VectorStoreRecordDefinition,
VectorStoreRecordVectorField,
)
from semantic_kernel.exceptions import (
VectorStoreModelDeserializationException,
VectorStoreModelSerializationException,
VectorStoreModelValidationError,
VectorStoreOperationException,
)
from semantic_kernel.kernel_pydantic import KernelBaseModel
from semantic_kernel.kernel_types import OneOrMany
from semantic_kernel.utils.feature_stage_decorator import experimental
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
if sys.version_info >= (3, 13):
from warnings import deprecated # pragma: no cover
else:
from typing_extensions import deprecated # pragma: no cover
TModel = TypeVar("TModel", bound=object)
TKey = TypeVar("TKey")
_T = TypeVar("_T", bound="VectorStoreRecordHandler")
logger = logging.getLogger(__name__)
class VectorStoreRecordHandler(KernelBaseModel, Generic[TKey, TModel]):
"""Vector Store Record Handler class.
This class is used to serialize and deserialize records to and from a vector store.
As well as validating the data model against the vector store.
It is subclassed by VectorStoreRecordCollection and VectorSearchBase.
"""
data_model_type: type[TModel]
data_model_definition: VectorStoreRecordDefinition
supported_key_types: ClassVar[list[str] | None] = None
supported_vector_types: ClassVar[list[str] | None] = None
@property
def _key_field_name(self) -> str:
return self.data_model_definition.key_field_name
@property
def _container_mode(self) -> bool:
return self.data_model_definition.container_mode
@model_validator(mode="before")
@classmethod
def _ensure_data_model_definition(cls: type[_T], data: Any) -> dict[str, Any]:
"""Ensure there is a data model definition, if it isn't passed, try to get it from the data model type."""
if isinstance(data, dict) and not data.get("data_model_definition"):
data["data_model_definition"] = getattr(
data["data_model_type"], "__kernel_vectorstoremodel_definition__", None
)
return data
def model_post_init(self, __context: object | None = None):
"""Post init function that sets the key field and container mode values, and validates the datamodel."""
self._validate_data_model()
def _validate_data_model(self):
"""Internal function that can be overloaded by child classes to validate datatypes, etc.
This should take the VectorStoreRecordDefinition from the item_type and validate it against the store.
Checks can include, allowed naming of parameters, allowed data types, allowed vector dimensions.
Default checks are that the key field is in the allowed key types and the vector fields
are in the allowed vector types.
Raises:
VectorStoreModelValidationError: If the key field is not in the allowed key types.
VectorStoreModelValidationError: If the vector fields are not in the allowed vector types.
"""
if (
self.supported_key_types
and self.data_model_definition.key_field.property_type
and self.data_model_definition.key_field.property_type not in self.supported_key_types
):
raise VectorStoreModelValidationError(
f"Key field must be one of {self.supported_key_types}, "
f"got {self.data_model_definition.key_field.property_type}"
)
if not self.supported_vector_types:
return
for field in self.data_model_definition.vector_fields:
if field.property_type and field.property_type not in self.supported_vector_types:
raise VectorStoreModelValidationError(
f"Vector field {field.name} must be one of {self.supported_vector_types}, got {field.property_type}"
)
@abstractmethod
def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]:
"""Serialize a list of dicts of the data to the store model.
This method should be overridden by the child class to convert the dict to the store model.
"""
... # pragma: no cover
@abstractmethod
def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]:
"""Deserialize the store models to a list of dicts.
This method should be overridden by the child class to convert the store model to a list of dicts.
"""
... # pragma: no cover
# region Serialization methods
def serialize(self, records: OneOrMany[TModel], **kwargs: Any) -> OneOrMany[Any]:
"""Serialize the data model to the store model.
This method follows the following steps:
1. Check if the data model has a serialize method.
Use that method to serialize and return the result.
2. Serialize the records into a dict, using the data model specific method.
3. Convert the dict to the store model, using the store specific method.
If overriding this method, make sure to first try to serialize the data model to the store model,
before doing the store specific version,
the user supplied version should have precedence.
Raises:
VectorStoreModelSerializationException: If an error occurs during serialization.
"""
try:
if serialized := self._serialize_data_model_to_store_model(records):
return serialized
if isinstance(records, Sequence):
dict_records = [self._serialize_data_model_to_dict(rec) for rec in records]
return self._serialize_dicts_to_store_models(dict_records, **kwargs) # type: ignore
dict_records = self._serialize_data_model_to_dict(records) # type: ignore
if isinstance(dict_records, Sequence):
# most likely this is a container, so we return all records as a list
# can also be a single record, but the to_dict returns a list
# hence we will treat it as a container.
return self._serialize_dicts_to_store_models(dict_records, **kwargs) # type: ignore
# this case is single record in, single record out
return self._serialize_dicts_to_store_models([dict_records], **kwargs)[0]
except VectorStoreModelSerializationException:
raise # pragma: no cover
except Exception as exc:
raise VectorStoreModelSerializationException(f"Error serializing records: {exc}") from exc
def _serialize_data_model_to_store_model(self, record: OneOrMany[TModel], **kwargs: Any) -> OneOrMany[Any] | None:
"""Serialize the data model to the store model.
This works when the data model has supplied a serialize method, specific to a data source.
This is a method called 'serialize()' on the data model or part of the vector store record definition.
The developer is responsible for correctly serializing for the specific data source.
"""
if isinstance(record, Sequence):
result = [self._serialize_data_model_to_store_model(rec, **kwargs) for rec in record]
if not all(result):
return None
return result
if self.data_model_definition.serialize:
return self.data_model_definition.serialize(record, **kwargs)
if isinstance(record, SerializeMethodProtocol):
return record.serialize(**kwargs)
return None
def _serialize_data_model_to_dict(self, record: TModel, **kwargs: Any) -> OneOrMany[dict[str, Any]]:
"""This function is used if no serialize method is found on the data model.
This will generally serialize the data model to a dict, should not be overridden by child classes.
The output of this should be passed to the serialize_dict_to_store_model method.
"""
if self.data_model_definition.to_dict:
return self.data_model_definition.to_dict(record, **kwargs)
if isinstance(record, ToDictMethodProtocol):
return self._serialize_vectors(record.to_dict())
if isinstance(record, BaseModel):
return self._serialize_vectors(record.model_dump())
store_model = {}
for field_name, field in self.data_model_definition.fields.items():
if isinstance(field, VectorStoreRecordVectorField) and not field.local_embedding:
logger.info(f"Vector field {field_name} is not local, skipping serialization.")
continue
value = record.get(field_name, None) if isinstance(record, Mapping) else getattr(record, field_name)
if isinstance(field, VectorStoreRecordVectorField):
if (func := getattr(field, "serialize_function", None)) and value is not None:
value = func(value)
elif value is None:
# if the field is not a vector field, then it should have a value.
raise VectorStoreModelSerializationException(f"Field {field_name} is None, cannot serialize.")
store_model[field_name] = value
return store_model
def _serialize_vectors(self, record: dict[str, Any]) -> dict[str, Any]:
for field in self.data_model_definition.vector_fields:
if field.serialize_function:
record[field.name or ""] = field.serialize_function(record[field.name or ""])
return record
# region Deserialization methods
def deserialize(self, records: OneOrMany[Any | dict[str, Any]], **kwargs: Any) -> OneOrMany[TModel] | None:
"""Deserialize the store model to the data model.
This method follows the following steps:
1. Check if the data model has a deserialize method.
Use that method to deserialize and return the result.
2. Deserialize the store model to a dict, using the store specific method.
3. Convert the dict to the data model, using the data model specific method.
Raises:
VectorStoreModelDeserializationException: If an error occurs during deserialization.
"""
try:
if not records:
return None
if deserialized := self._deserialize_store_model_to_data_model(records, **kwargs):
return deserialized
if isinstance(records, Sequence):
dict_records = self._deserialize_store_models_to_dicts(records, **kwargs)
return (
self._deserialize_dict_to_data_model(dict_records, **kwargs)
if self._container_mode
else [self._deserialize_dict_to_data_model(rec, **kwargs) for rec in dict_records]
)
dict_record = self._deserialize_store_models_to_dicts([records], **kwargs)[0]
# regardless of mode, only 1 object is returned.
return self._deserialize_dict_to_data_model(dict_record, **kwargs)
except VectorStoreModelDeserializationException:
raise # pragma: no cover
except Exception as exc:
raise VectorStoreModelDeserializationException(f"Error deserializing records: {exc}") from exc
def _deserialize_store_model_to_data_model(self, record: OneOrMany[Any], **kwargs: Any) -> OneOrMany[TModel] | None:
"""Deserialize the store model to the data model.
This works when the data model has supplied a deserialize method, specific to a data source.
This uses a method called 'deserialize()' on the data model or part of the vector store record definition.
The developer is responsible for correctly deserializing for the specific data source.
"""
if self.data_model_definition.deserialize:
if isinstance(record, Sequence):
return self.data_model_definition.deserialize(record, **kwargs)
return self.data_model_definition.deserialize([record], **kwargs)
if func := getattr(self.data_model_type, "deserialize", None):
if isinstance(record, Sequence):
return [func(rec, **kwargs) for rec in record]
return func(record, **kwargs)
return None
def _deserialize_dict_to_data_model(self, record: OneOrMany[dict[str, Any]], **kwargs: Any) -> TModel:
"""This function is used if no deserialize method is found on the data model.
This method is the second step and will deserialize a dict to the data model,
should not be overridden by child classes.
The input of this should come from the _deserialized_store_model_to_dict function.
"""
include_vectors = kwargs.get("include_vectors", True)
if self.data_model_definition.from_dict:
if isinstance(record, Sequence):
return self.data_model_definition.from_dict(record, **kwargs)
ret = self.data_model_definition.from_dict([record], **kwargs)
return ret if self._container_mode else ret[0]
if isinstance(record, Sequence):
if len(record) > 1:
raise VectorStoreModelDeserializationException(
"Cannot deserialize multiple records to a single record unless you are using a container."
)
record = record[0]
if func := getattr(self.data_model_type, "from_dict", None):
if include_vectors:
record = self._deserialize_vector(record)
return func(record)
if issubclass(self.data_model_type, BaseModel):
if include_vectors:
record = self._deserialize_vector(record)
return self.data_model_type.model_validate(record) # type: ignore
data_model_dict: dict[str, Any] = {}
for field_name, field in self.data_model_definition.fields.items():
value = record.get(field_name, None)
if isinstance(field, VectorStoreRecordVectorField):
if not include_vectors or not field.local_embedding:
continue
if field.deserialize_function and value is not None:
value = field.deserialize_function(value)
elif value is None:
# if the field is not a vector field, then it should have a value.
raise VectorStoreModelDeserializationException(f"Field {field_name} is None, cannot deserialize.")
data_model_dict[field_name] = value
if self.data_model_type is dict:
return data_model_dict # type: ignore
return self.data_model_type(**data_model_dict)
def _deserialize_vector(self, record: dict[str, Any]) -> dict[str, Any]:
for field in self.data_model_definition.vector_fields:
if field.deserialize_function:
if not field.local_embedding:
logger.info(f"Vector field {field.name} is not local, skipping deserialization.")
continue
record[field.name] = field.deserialize_function(record[field.name])
return record
@experimental
class VectorStoreRecordCollection(VectorStoreRecordHandler, Generic[TKey, TModel]):
"""Base class for a vector store record collection."""
collection_name: str
managed_client: bool = True
async def __aenter__(self) -> Self:
"""Enter the context manager."""
return self
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
"""Exit the context manager.
Should be overridden by subclasses, if necessary.
If the client is passed in the constructor, it should not be closed,
in that case the managed_client should be set to False.
If the store supplied the managed client, it is responsible for closing it,
and it should not be closed here and so managed_client should be False.
Some services use two clients, one for the store and one for the collection,
in that case, the collection client should be closed here,
but the store client should only be closed when it is created in the collection.
A additional flag might be needed for that.
"""
pass
@abstractmethod
async def _inner_upsert(
self,
records: Sequence[Any],
**kwargs: Any,
) -> Sequence[TKey]:
"""Upsert the records, this should be overridden by the child class.
Args:
records: The records, the format is specific to the store.
**kwargs (Any): Additional arguments, to be passed to the store.
Returns:
The keys of the upserted records.
Raises:
Exception: If an error occurs during the upsert.
There is no need to catch and parse exceptions in the inner functions,
they are handled by the public methods.
The only exception is raises exceptions yourself, such as a ValueError.
This is then caught and turned into the relevant exception by the public method.
This setup promotes a limited depth of the stack trace.
"""
... # pragma: no cover
@abstractmethod
async def _inner_get(self, keys: Sequence[TKey], **kwargs: Any) -> OneOrMany[Any] | None:
"""Get the records, this should be overridden by the child class.
Args:
keys: The keys to get.
**kwargs: Additional arguments.
Returns:
The records from the store, not deserialized.
Raises:
Exception: If an error occurs during the upsert.
There is no need to catch and parse exceptions in the inner functions,
they are handled by the public methods.
The only exception is raises exceptions yourself, such as a ValueError.
This is then caught and turned into the relevant exception by the public method.
This setup promotes a limited depth of the stack trace.
"""
... # pragma: no cover
@abstractmethod
async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None:
"""Delete the records, this should be overridden by the child class.
Args:
keys: The keys.
**kwargs: Additional arguments.
Raises:
Exception: If an error occurs during the upsert.
There is no need to catch and parse exceptions in the inner functions,
they are handled by the public methods.
The only exception is raises exceptions yourself, such as a ValueError.
This is then caught and turned into the relevant exception by the public method.
This setup promotes a limited depth of the stack trace.
"""
... # pragma: no cover
async def create_collection_if_not_exists(self, **kwargs: Any) -> bool:
"""Create the collection in the service if it does not exists.
First uses does_collection_exist to check if it exists, if it does returns False.
Otherwise, creates the collection and returns True.
"""
if await self.does_collection_exist(**kwargs):
return False
await self.create_collection(**kwargs)
return True
@abstractmethod
async def create_collection(self, **kwargs: Any) -> None:
"""Create the collection in the service.
This should be overridden by the child class.
Raises:
Make sure the implementation of this function raises relevant exceptions with good descriptions.
This is different then the `_inner_x` methods, as this is a public method.
"""
... # pragma: no cover
@abstractmethod
async def does_collection_exist(self, **kwargs: Any) -> bool:
"""Check if the collection exists.
This should be overridden by the child class.
Raises:
Make sure the implementation of this function raises relevant exceptions with good descriptions.
This is different then the `_inner_x` methods, as this is a public method.
"""
... # pragma: no cover
@abstractmethod
async def delete_collection(self, **kwargs: Any) -> None:
"""Delete the collection.
This should be overridden by the child class.
Raises:
Make sure the implementation of this function raises relevant exceptions with good descriptions.
This is different then the `_inner_x` methods, as this is a public method.
"""
... # pragma: no cover
# region Public Methods
@deprecated("upsert_batch is deprecated, use upsert instead.")
async def upsert_batch(self, *args: Any, **kwargs: Any) -> Sequence[TKey]:
"""Upsert a batch of records, this method is deprecated, use upsert instead."""
return await self.upsert(*args, **kwargs)
@overload
async def upsert(
self,
record: TModel = ...,
embedding_generation_function: Callable[
[TModel, type[TModel] | None, VectorStoreRecordDefinition | None], Awaitable[TModel]
]
| None = None,
**kwargs: Any,
) -> OneOrMany[TKey] | None: ...
@overload
async def upsert(
self,
records: OneOrMany[TModel] = ...,
embedding_generation_function: Callable[
[OneOrMany[TModel], type[TModel] | None, VectorStoreRecordDefinition | None], Awaitable[OneOrMany[TModel]]
]
| None = None,
**kwargs: Any,
) -> Sequence[TKey]: ...
async def upsert(self, record=None, records=None, embedding_generation_function=None, **kwargs):
"""Upsert a one or more records.
If the key of the record already exists, the existing record will be updated.
If the key does not exist, a new record will be created.
Args:
record: The record to upsert, can be a list of records, or a single container.
records: The records to upsert, can be a list of records, or a single container,
if supplied, record is ignored.
embedding_generation_function: Supply this function to generate embeddings.
This will be called with the data model definition and the records,
should return the records with vectors.
This can be supplied by using the add_vector_to_records method.
**kwargs: Additional arguments.
Returns:
Sequence[TKey]: The keys of the upserted records, this is always a list,
corresponds to the input or the items in the container.
Raises:
VectorStoreModelSerializationException: If an error occurs during serialization.
VectorStoreOperationException: If an error occurs during upserting.
"""
batch = True
if records is None and record is not None:
if not isinstance(record, list) and not self._container_mode:
records = [record]
batch = False
else:
records = record
if records is None:
raise VectorStoreOperationException("Either record or records must be provided.")
if embedding_generation_function:
records = await embedding_generation_function(records, self.data_model_type, self.data_model_definition)
try:
data = self.serialize(records)
# the serialize method will parse any exception into a VectorStoreModelSerializationException
except VectorStoreModelSerializationException:
raise
try:
results = await self._inner_upsert(data, **kwargs) # type: ignore
except Exception as exc:
raise VectorStoreOperationException(f"Error upserting record(s): {exc}") from exc
if batch or self._container_mode:
return results
return results[0]
@deprecated("get_batch is deprecated, use get instead.")
async def get_batch(self, *args: Any, **kwargs: Any) -> OneOrMany[TModel] | None:
"""Get a batch of records, this method is deprecated, use get instead."""
return await self.get(*args, **kwargs)
@overload
async def get(self, key: TKey = ..., include_vectors: bool = True, **kwargs: Any) -> TModel | None: ...
@overload
async def get(
self, keys: Sequence[TKey] = ..., include_vectors: bool = True, **kwargs: Any
) -> OneOrMany[TModel] | None: ...
async def get(self, key=None, keys=None, include_vectors=True, **kwargs):
"""Get a batch of records whose keys exist in the collection, i.e. keys that do not exist are ignored.
Args:
key: The key to get.
keys: The keys to get, if keys are provided, key is ignored.
include_vectors: Include the vectors in the response. Default is True.
Some vector stores do not support retrieving without vectors, even when set to false.
Some vector stores have specific parameters to control that behavior, when
that parameter is set, include_vectors is ignored.
**kwargs: Additional arguments.
Returns:
The records, either a list of TModel or the container type.
Raises:
VectorStoreOperationException: If an error occurs during the get.
VectorStoreModelDeserializationException: If an error occurs during deserialization.
"""
batch = True
if not keys and key:
if not isinstance(key, list):
keys = [key]
batch = False
else:
keys = key
if not keys:
raise VectorStoreOperationException("Either key or keys must be provided.")
try:
records = await self._inner_get(keys, include_vectors=include_vectors, **kwargs)
except Exception as exc:
raise VectorStoreOperationException(f"Error getting record(s): {exc}") from exc
if not records:
return None
try:
model_records = self.deserialize(records if batch else records[0], **kwargs)
# the deserialize method will parse any exception into a VectorStoreModelDeserializationException
except VectorStoreModelDeserializationException:
raise
# there are many code paths within the deserialize method, some supplied by the developer,
# and so depending on what is used,
# it might return a sequence, so we just return the first element,
# there should never be multiple elements (this is not a batch get),
# hence a raise if there are.
if batch:
return model_records
if not isinstance(model_records, Sequence):
return model_records
if len(model_records) == 1:
return model_records[0]
raise VectorStoreModelDeserializationException(
f"Error deserializing record, multiple records returned: {model_records}"
)
@deprecated("delete_batch is deprecated, use delete instead.")
async def delete_batch(self, *args: Any, **kwargs: Any) -> None:
"""Delete a batch of records, this method is deprecated, use delete instead."""
return await self.delete(*args, **kwargs)
@overload
async def delete(self, key: TKey = ..., **kwargs: Any) -> None: ...
@overload
async def delete(self, keys: Sequence[TKey] = ..., **kwargs: Any) -> None: ...
async def delete(self, key=None, keys=None, **kwargs):
"""Delete a one or more records.
An exception will be raised at the end if any record does not exist.
Args:
key: The key to delete.
keys: The keys, if keys are provided, key is ignored.
**kwargs: Additional arguments.
Exceptions:
VectorStoreOperationException: If an error occurs during deletion or a record does not exist.
"""
if not keys:
if key and not isinstance(key, list):
keys = [key]
elif isinstance(key, list):
keys = key
else:
raise VectorStoreOperationException("Either key or keys must be provided.")
try:
await self._inner_delete(keys, **kwargs)
except Exception as exc:
raise VectorStoreOperationException(f"Error deleting record(s): {exc}") from exc
# region Internal Functions
def __del__(self):
"""Delete the instance."""
with suppress(Exception):
asyncio.get_running_loop().create_task(self.__aexit__(None, None, None))
@experimental
class VectorStore(KernelBaseModel):
"""Base class for vector stores."""
vector_record_collections: dict[str, VectorStoreRecordCollection] = Field(default_factory=dict)
managed_client: bool = True
@abstractmethod
def get_collection(
self,
collection_name: str,
data_model_type: type[object],
data_model_definition: VectorStoreRecordDefinition | None = None,
**kwargs: Any,
) -> "VectorStoreRecordCollection":
"""Get a vector record store."""
... # pragma: no cover
@abstractmethod
async def list_collection_names(self, **kwargs) -> Sequence[str]:
"""Get the names of all collections."""
... # pragma: no cover
async def __aenter__(self) -> Self:
"""Enter the context manager."""
return self
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
"""Exit the context manager.
Should be overridden by subclasses, if necessary.
If the client is passed in the constructor, it should not be closed,
in that case the managed_client should be set to False.
"""
pass # pragma: no cover