chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,633 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import ast
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import MutableSequence, Sequence
|
||||
from importlib import metadata
|
||||
from typing import Any, ClassVar, Final, Generic, TypeVar
|
||||
|
||||
from pydantic import SecretStr, ValidationError
|
||||
from pymongo import AsyncMongoClient, ReplaceOne
|
||||
from pymongo.asynchronous.collection import AsyncCollection
|
||||
from pymongo.asynchronous.database import AsyncDatabase
|
||||
from pymongo.driver_info import DriverInfo
|
||||
from pymongo.operations import SearchIndexModel
|
||||
|
||||
from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase
|
||||
from semantic_kernel.data.vector import (
|
||||
DistanceFunction,
|
||||
GetFilteredRecordOptions,
|
||||
KernelSearchResults,
|
||||
SearchType,
|
||||
TModel,
|
||||
VectorSearch,
|
||||
VectorSearchOptions,
|
||||
VectorSearchResult,
|
||||
VectorStore,
|
||||
VectorStoreCollection,
|
||||
VectorStoreCollectionDefinition,
|
||||
VectorStoreField,
|
||||
_get_collection_name_from_model,
|
||||
)
|
||||
from semantic_kernel.exceptions import (
|
||||
VectorSearchExecutionException,
|
||||
VectorStoreInitializationException,
|
||||
VectorStoreOperationException,
|
||||
)
|
||||
from semantic_kernel.exceptions.vector_store_exceptions import VectorStoreModelException
|
||||
from semantic_kernel.kernel_pydantic import KernelBaseSettings
|
||||
from semantic_kernel.utils.feature_stage_decorator import release_candidate
|
||||
from semantic_kernel.utils.telemetry.user_agent import SEMANTIC_KERNEL_USER_AGENT
|
||||
|
||||
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, 12):
|
||||
from typing import override # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
TKey = TypeVar("TKey", bound=str)
|
||||
|
||||
DEFAULT_DB_NAME: Final[str] = "default"
|
||||
DEFAULT_SEARCH_INDEX_NAME: Final[str] = "default"
|
||||
MONGODB_ID_FIELD: Final[str] = "_id"
|
||||
MONGODB_SCORE_FIELD: Final[str] = "score"
|
||||
NUM_CANDIDATES_SCALAR: Final[int] = 10
|
||||
DISTANCE_FUNCTION_MAP: Final[dict[DistanceFunction, str]] = {
|
||||
DistanceFunction.EUCLIDEAN_DISTANCE: "euclidean",
|
||||
DistanceFunction.COSINE_SIMILARITY: "cosine",
|
||||
DistanceFunction.DOT_PROD: "dotProduct",
|
||||
DistanceFunction.DEFAULT: "euclidean",
|
||||
}
|
||||
DRIVER_METADATA = DriverInfo(SEMANTIC_KERNEL_USER_AGENT, metadata.version("semantic-kernel"))
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@release_candidate
|
||||
class MongoDBAtlasSettings(KernelBaseSettings):
|
||||
"""MongoDB Atlas model settings.
|
||||
|
||||
Args:
|
||||
- connection_string: str - MongoDB Atlas connection string
|
||||
(Env var MONGODB_ATLAS_CONNECTION_STRING)
|
||||
- database_name: str - MongoDB Atlas database name, defaults to 'default'
|
||||
(Env var MONGODB_ATLAS_DATABASE_NAME)
|
||||
- index_name: str - MongoDB Atlas search index name, defaults to 'default'
|
||||
(Env var MONGODB_ATLAS_INDEX_NAME)
|
||||
"""
|
||||
|
||||
env_prefix: ClassVar[str] = "MONGODB_ATLAS_"
|
||||
|
||||
connection_string: SecretStr
|
||||
database_name: str = DEFAULT_DB_NAME
|
||||
index_name: str = DEFAULT_SEARCH_INDEX_NAME
|
||||
|
||||
|
||||
def _create_vector_field(field: VectorStoreField) -> dict:
|
||||
"""Create a vector field.
|
||||
|
||||
Args:
|
||||
field (VectorStoreRecordVectorField): The vector field.
|
||||
|
||||
Returns:
|
||||
dict: The vector field.
|
||||
"""
|
||||
if field.distance_function not in DISTANCE_FUNCTION_MAP:
|
||||
raise VectorStoreInitializationException(
|
||||
f"Distance function {field.distance_function} is not supported. "
|
||||
f"Supported distance functions are: {list(DISTANCE_FUNCTION_MAP.keys())}"
|
||||
)
|
||||
return {
|
||||
"type": "vector",
|
||||
"numDimensions": field.dimensions,
|
||||
"path": field.storage_name or field.name,
|
||||
"similarity": DISTANCE_FUNCTION_MAP[field.distance_function],
|
||||
}
|
||||
|
||||
|
||||
def _create_index_definitions(
|
||||
record_definition: VectorStoreCollectionDefinition, index_name: str
|
||||
) -> list[SearchIndexModel]:
|
||||
"""Create the index definitions."""
|
||||
indexes = []
|
||||
if record_definition.vector_fields:
|
||||
vector_fields = [_create_vector_field(field) for field in record_definition.vector_fields]
|
||||
filterable_fields = [
|
||||
{"path": field.storage_name or field.name, "type": "filter"}
|
||||
for field in record_definition.data_fields
|
||||
if field.is_indexed
|
||||
]
|
||||
filterable_fields.append({"path": record_definition.key_field.name, "type": "filter"})
|
||||
indexes.append(
|
||||
SearchIndexModel(
|
||||
type="vectorSearch",
|
||||
name=index_name,
|
||||
definition={"fields": vector_fields + filterable_fields},
|
||||
)
|
||||
)
|
||||
if record_definition.data_fields:
|
||||
ft_indexed_fields = [
|
||||
{field.storage_name or field.name: {"type": "string"}}
|
||||
for field in record_definition.data_fields
|
||||
if field.is_full_text_indexed
|
||||
]
|
||||
if ft_indexed_fields:
|
||||
indexes.append(
|
||||
SearchIndexModel(
|
||||
type="search",
|
||||
name=f"{index_name}_ft",
|
||||
definition={
|
||||
"mapping": {"dynamic": True, "fields": ft_indexed_fields},
|
||||
},
|
||||
)
|
||||
)
|
||||
return indexes
|
||||
|
||||
|
||||
@release_candidate
|
||||
class MongoDBAtlasCollection(
|
||||
VectorStoreCollection[TKey, TModel],
|
||||
VectorSearch[TKey, TModel],
|
||||
Generic[TKey, TModel],
|
||||
):
|
||||
"""MongoDB Atlas collection implementation."""
|
||||
|
||||
mongo_client: AsyncMongoClient
|
||||
database_name: str
|
||||
index_name: str
|
||||
supported_key_types: ClassVar[set[str] | None] = {"str"}
|
||||
supported_vector_types: ClassVar[set[str] | None] = {"float", "int"}
|
||||
supported_search_types: ClassVar[set[SearchType]] = {SearchType.VECTOR, SearchType.KEYWORD_HYBRID}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
record_type: type[TModel],
|
||||
definition: VectorStoreCollectionDefinition | None = None,
|
||||
collection_name: str | None = None,
|
||||
embedding_generator: EmbeddingGeneratorBase | None = None,
|
||||
index_name: str | None = None,
|
||||
mongo_client: AsyncMongoClient | None = None,
|
||||
connection_string: str | None = None,
|
||||
database_name: str | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initializes a new instance of the MongoDBAtlasCollection class.
|
||||
|
||||
Args:
|
||||
record_type: The type of the data model.
|
||||
definition: The model definition, optional.
|
||||
collection_name: The name of the collection, optional.
|
||||
embedding_generator: The embedding generator, optional.
|
||||
index_name: The name of the index to use for searching, when not passed, will use <collection_name>_idx.
|
||||
mongo_client: The MongoDB client for interacting with MongoDB Atlas,
|
||||
used for creating and deleting collections.
|
||||
connection_string: The connection string for MongoDB Atlas, optional.
|
||||
Can be read from environment variables.
|
||||
database_name: The name of the database, will be filled from the env when this is not set.
|
||||
connection_string: str | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None
|
||||
**kwargs: Additional keyword arguments
|
||||
"""
|
||||
if not collection_name:
|
||||
collection_name = _get_collection_name_from_model(record_type, definition)
|
||||
managed_client = kwargs.get("managed_client", not mongo_client)
|
||||
if mongo_client:
|
||||
super().__init__(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
mongo_client=mongo_client,
|
||||
collection_name=collection_name,
|
||||
database_name=database_name or DEFAULT_DB_NAME,
|
||||
index_name=index_name or DEFAULT_SEARCH_INDEX_NAME,
|
||||
managed_client=managed_client,
|
||||
embedding_generator=embedding_generator,
|
||||
)
|
||||
if callable(mongo_client.append_metadata):
|
||||
mongo_client.append_metadata(DRIVER_METADATA)
|
||||
return
|
||||
|
||||
try:
|
||||
mongodb_atlas_settings = MongoDBAtlasSettings(
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
connection_string=connection_string,
|
||||
database_name=database_name,
|
||||
index_name=index_name,
|
||||
)
|
||||
except ValidationError as exc:
|
||||
raise VectorStoreInitializationException("Failed to create MongoDB Atlas settings.") from exc
|
||||
|
||||
mongo_client = AsyncMongoClient(
|
||||
mongodb_atlas_settings.connection_string.get_secret_value(),
|
||||
driver=DRIVER_METADATA,
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
mongo_client=mongo_client,
|
||||
managed_client=managed_client,
|
||||
database_name=mongodb_atlas_settings.database_name,
|
||||
index_name=mongodb_atlas_settings.index_name,
|
||||
embedding_generator=embedding_generator,
|
||||
)
|
||||
|
||||
def _get_database(self) -> AsyncDatabase:
|
||||
"""Get the database.
|
||||
|
||||
If you need control over things like read preference, you can override this method.
|
||||
"""
|
||||
return self.mongo_client.get_database(self.database_name)
|
||||
|
||||
def _get_collection(self) -> AsyncCollection:
|
||||
"""Get the collection.
|
||||
|
||||
If you need control over things like read preference, you can override this method.
|
||||
"""
|
||||
return self.mongo_client.get_database(self.database_name).get_collection(self.collection_name)
|
||||
|
||||
@override
|
||||
async def _inner_upsert(
|
||||
self,
|
||||
records: Sequence[Any],
|
||||
**kwargs: Any,
|
||||
) -> Sequence[TKey]:
|
||||
operations: MutableSequence[ReplaceOne] = []
|
||||
for record in records:
|
||||
operations.append(
|
||||
ReplaceOne(
|
||||
filter={MONGODB_ID_FIELD: record[MONGODB_ID_FIELD]},
|
||||
replacement=record,
|
||||
upsert=True,
|
||||
)
|
||||
)
|
||||
result = await self._get_collection().bulk_write(operations, ordered=False)
|
||||
return [str(value) for _, value in result.upserted_ids.items()] # type: ignore
|
||||
|
||||
@override
|
||||
async def _inner_get(
|
||||
self,
|
||||
keys: Sequence[TKey] | None = None,
|
||||
options: GetFilteredRecordOptions | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Sequence[dict[str, Any]] | None:
|
||||
if not keys:
|
||||
if options is not None:
|
||||
raise NotImplementedError("Get without keys is not yet implemented.")
|
||||
return None
|
||||
result = self._get_collection().find({MONGODB_ID_FIELD: {"$in": keys}})
|
||||
return await result.to_list(length=len(keys))
|
||||
|
||||
@override
|
||||
async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None:
|
||||
collection = self._get_collection()
|
||||
await collection.delete_many({MONGODB_ID_FIELD: {"$in": keys}})
|
||||
|
||||
def _replace_key_field(self, record: dict[str, Any]) -> dict[str, Any]:
|
||||
if self._key_field_name == MONGODB_ID_FIELD:
|
||||
return record
|
||||
return {
|
||||
MONGODB_ID_FIELD: record.pop(self._key_field_name, None),
|
||||
**record,
|
||||
}
|
||||
|
||||
def _reset_key_field(self, record: dict[str, Any]) -> dict[str, Any]:
|
||||
if self._key_field_name == MONGODB_ID_FIELD:
|
||||
return record
|
||||
return {
|
||||
self._key_field_name: record.pop(MONGODB_ID_FIELD, None),
|
||||
**record,
|
||||
}
|
||||
|
||||
@override
|
||||
def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]:
|
||||
return [self._replace_key_field(record) for record in records]
|
||||
|
||||
@override
|
||||
def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]:
|
||||
return [self._reset_key_field(record) for record in records]
|
||||
|
||||
@override
|
||||
async def ensure_collection_exists(self, **kwargs) -> None:
|
||||
"""Create a new collection in MongoDB.
|
||||
|
||||
This first creates a collection, with the kwargs.
|
||||
Then creates a search index based on the data model definition.
|
||||
|
||||
Args:
|
||||
**kwargs: Additional keyword arguments.
|
||||
"""
|
||||
collection = await self._get_database().create_collection(self.collection_name, **kwargs)
|
||||
await collection.create_search_indexes(models=_create_index_definitions(self.definition, self.index_name))
|
||||
|
||||
@override
|
||||
async def collection_exists(self, **kwargs) -> bool:
|
||||
return bool(await self._get_database().list_collection_names(filter={"name": self.collection_name}))
|
||||
|
||||
@override
|
||||
async def ensure_collection_deleted(self, **kwargs) -> None:
|
||||
await self._get_database().drop_collection(self.collection_name, **kwargs)
|
||||
|
||||
@override
|
||||
async def _inner_search(
|
||||
self,
|
||||
search_type: SearchType,
|
||||
options: VectorSearchOptions,
|
||||
values: Any | None = None,
|
||||
vector: Sequence[float | int] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> KernelSearchResults[VectorSearchResult[TModel]]:
|
||||
if search_type == SearchType.VECTOR:
|
||||
return await self._inner_vector_search(options, values, vector, **kwargs)
|
||||
if search_type == SearchType.KEYWORD_HYBRID:
|
||||
return await self._inner_keyword_hybrid_search(options, values, vector, **kwargs)
|
||||
raise VectorStoreOperationException("Vector is required for search.")
|
||||
|
||||
async def _inner_vector_search(
|
||||
self,
|
||||
options: VectorSearchOptions,
|
||||
values: Any | None = None,
|
||||
vector: Sequence[float | int] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> KernelSearchResults[VectorSearchResult[TModel]]:
|
||||
collection = self._get_collection()
|
||||
vector_field = self.definition.try_get_vector_field(options.vector_property_name)
|
||||
if not vector_field:
|
||||
raise VectorStoreModelException(
|
||||
f"Vector field '{options.vector_property_name}' not found in the data model definition."
|
||||
)
|
||||
if not vector:
|
||||
vector = await self._generate_vector_from_values(values, options)
|
||||
vector_search_query: dict[str, Any] = {
|
||||
"limit": options.top + options.skip,
|
||||
"index": self.index_name,
|
||||
"queryVector": vector,
|
||||
"path": vector_field.storage_name or vector_field.name,
|
||||
}
|
||||
if filter := self._build_filter(options.filter):
|
||||
vector_search_query["filter"] = filter if isinstance(filter, dict) else {"$and": filter}
|
||||
|
||||
projection_query: dict[str, int | dict] = {
|
||||
field: 1
|
||||
for field in self.definition.get_names(
|
||||
include_vector_fields=options.include_vectors,
|
||||
include_key_field=False, # _id is always included
|
||||
)
|
||||
}
|
||||
projection_query[MONGODB_SCORE_FIELD] = {"$meta": "vectorSearchScore"}
|
||||
try:
|
||||
raw_results = await collection.aggregate([
|
||||
{"$vectorSearch": vector_search_query},
|
||||
{"$project": projection_query},
|
||||
])
|
||||
except Exception as exc:
|
||||
raise VectorSearchExecutionException("Failed to search the collection.") from exc
|
||||
return KernelSearchResults(
|
||||
results=self._get_vector_search_results_from_results(raw_results, options),
|
||||
total_count=None, # no way to get a count before looping through the result cursor
|
||||
)
|
||||
|
||||
async def _inner_keyword_hybrid_search(
|
||||
self,
|
||||
options: VectorSearchOptions,
|
||||
values: Any | None = None,
|
||||
vector: Sequence[float | int] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> KernelSearchResults[VectorSearchResult[TModel]]:
|
||||
collection = self._get_collection()
|
||||
vector_field = self.definition.try_get_vector_field(options.vector_property_name)
|
||||
if not vector_field:
|
||||
raise VectorStoreModelException(
|
||||
f"Vector field '{options.vector_property_name}' not found in the data model definition."
|
||||
)
|
||||
if not vector:
|
||||
vector = await self._generate_vector_from_values(values, options)
|
||||
vector_search_query: dict[str, Any] = {
|
||||
"limit": options.top + options.skip,
|
||||
"index": self.index_name,
|
||||
"queryVector": vector,
|
||||
"path": vector_field.storage_name or vector_field.name,
|
||||
}
|
||||
if filter := self._build_filter(options.filter):
|
||||
vector_search_query["filter"] = filter if isinstance(filter, dict) else {"$and": filter}
|
||||
|
||||
projection_query: dict[str, int | dict] = {
|
||||
field: 1
|
||||
for field in self.definition.get_names(
|
||||
include_vector_fields=options.include_vectors,
|
||||
include_key_field=False, # _id is always included
|
||||
)
|
||||
}
|
||||
projection_query[MONGODB_SCORE_FIELD] = {"$meta": "vectorSearchScore"}
|
||||
try:
|
||||
raw_results = await collection.aggregate([
|
||||
{"$vectorSearch": vector_search_query},
|
||||
{"$project": projection_query},
|
||||
])
|
||||
except Exception as exc:
|
||||
raise VectorSearchExecutionException("Failed to search the collection.") from exc
|
||||
return KernelSearchResults(
|
||||
results=self._get_vector_search_results_from_results(raw_results, options),
|
||||
total_count=None, # no way to get a count before looping through the result cursor
|
||||
)
|
||||
|
||||
@override
|
||||
def _lambda_parser(self, node: ast.AST) -> Any:
|
||||
# Comparison operations
|
||||
match node:
|
||||
case ast.Compare():
|
||||
if len(node.ops) > 1:
|
||||
# Chain comparisons (e.g., 1 < x < 3) become $and of each comparison
|
||||
values = []
|
||||
for idx in range(len(node.ops)):
|
||||
left = node.left if idx == 0 else node.comparators[idx - 1]
|
||||
right = node.comparators[idx]
|
||||
op = node.ops[idx]
|
||||
values.append(self._lambda_parser(ast.Compare(left=left, ops=[op], comparators=[right])))
|
||||
return {"$and": values}
|
||||
left = self._lambda_parser(node.left)
|
||||
right = self._lambda_parser(node.comparators[0])
|
||||
op = node.ops[0]
|
||||
match op:
|
||||
case ast.In():
|
||||
return {left: {"$in": right}}
|
||||
case ast.NotIn():
|
||||
return {left: {"$nin": right}}
|
||||
case ast.Eq():
|
||||
# MongoDB allows short form: {field: value}
|
||||
return {left: right}
|
||||
case ast.NotEq():
|
||||
return {left: {"$ne": right}}
|
||||
case ast.Gt():
|
||||
return {left: {"$gt": right}}
|
||||
case ast.GtE():
|
||||
return {left: {"$gte": right}}
|
||||
case ast.Lt():
|
||||
return {left: {"$lt": right}}
|
||||
case ast.LtE():
|
||||
return {left: {"$lte": right}}
|
||||
raise NotImplementedError(f"Unsupported operator: {type(op)}")
|
||||
case ast.BoolOp():
|
||||
op = node.op # type: ignore
|
||||
values = [self._lambda_parser(v) for v in node.values]
|
||||
if isinstance(op, ast.And):
|
||||
return {"$and": values}
|
||||
if isinstance(op, ast.Or):
|
||||
return {"$or": values}
|
||||
raise NotImplementedError(f"Unsupported BoolOp: {type(op)}")
|
||||
case ast.UnaryOp():
|
||||
match node.op:
|
||||
case ast.Not():
|
||||
operand = self._lambda_parser(node.operand)
|
||||
return {"$not": operand}
|
||||
case ast.UAdd() | ast.USub() | ast.Invert():
|
||||
raise NotImplementedError("Unary +, -, ~ are not supported in MongoDB filters.")
|
||||
case ast.Attribute():
|
||||
# Only allow attributes that are in the data model
|
||||
if node.attr not in self.definition.storage_names:
|
||||
raise VectorStoreOperationException(
|
||||
f"Field '{node.attr}' not in data model (storage property names are used)."
|
||||
)
|
||||
return node.attr
|
||||
case ast.Name():
|
||||
# Only allow names that are in the data model
|
||||
if node.id not in self.definition.storage_names:
|
||||
raise VectorStoreOperationException(
|
||||
f"Field '{node.id}' not in data model (storage property names are used)."
|
||||
)
|
||||
return node.id
|
||||
case ast.Constant():
|
||||
return node.value
|
||||
raise NotImplementedError(f"Unsupported AST node: {type(node)}")
|
||||
|
||||
@override
|
||||
def _get_record_from_result(self, result: dict[str, Any]) -> dict[str, Any]:
|
||||
return result
|
||||
|
||||
@override
|
||||
def _get_score_from_result(self, result: dict[str, Any]) -> float | None:
|
||||
return result.get(MONGODB_SCORE_FIELD)
|
||||
|
||||
@override
|
||||
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
|
||||
"""Exit the context manager."""
|
||||
if self.managed_client:
|
||||
await self.mongo_client.close()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
"""Enter the context manager."""
|
||||
await self.mongo_client.aconnect()
|
||||
return self
|
||||
|
||||
|
||||
@release_candidate
|
||||
class MongoDBAtlasStore(VectorStore):
|
||||
"""MongoDB Atlas store implementation."""
|
||||
|
||||
mongo_client: AsyncMongoClient
|
||||
database_name: str
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
connection_string: str | None = None,
|
||||
database_name: str | None = None,
|
||||
embedding_generator: EmbeddingGeneratorBase | None = None,
|
||||
mongo_client: AsyncMongoClient | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initializes a new instance of the MongoDBAtlasStore client.
|
||||
|
||||
Args:
|
||||
connection_string: The connection string for MongoDB Atlas, optional.
|
||||
Can be read from environment variables.
|
||||
database_name: The name of the database, optional. Can be read from environment variables.
|
||||
embedding_generator: The embedding generator, optional.
|
||||
mongo_client: The MongoDB client, optional.
|
||||
env_file_path: Use the environment settings file as a fallback
|
||||
to environment variables.
|
||||
env_file_encoding: The encoding of the environment settings file.
|
||||
kwargs: Additional keyword arguments.
|
||||
"""
|
||||
managed_client = kwargs.get("managed_client", not mongo_client)
|
||||
if mongo_client:
|
||||
super().__init__(
|
||||
mongo_client=mongo_client,
|
||||
managed_client=managed_client,
|
||||
database_name=database_name or DEFAULT_DB_NAME,
|
||||
embedding_generator=embedding_generator,
|
||||
)
|
||||
if callable(mongo_client.append_metadata):
|
||||
mongo_client.append_metadata(DRIVER_METADATA)
|
||||
return
|
||||
|
||||
try:
|
||||
mongodb_atlas_settings = MongoDBAtlasSettings(
|
||||
env_file_path=env_file_path,
|
||||
connection_string=connection_string,
|
||||
database_name=database_name,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
except ValidationError as exc:
|
||||
raise VectorStoreInitializationException("Failed to create MongoDB Atlas settings.") from exc
|
||||
if not mongodb_atlas_settings.connection_string:
|
||||
raise VectorStoreInitializationException("The connection string is missing.")
|
||||
|
||||
mongo_client = AsyncMongoClient(
|
||||
mongodb_atlas_settings.connection_string.get_secret_value(),
|
||||
driver=DRIVER_METADATA,
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
mongo_client=mongo_client,
|
||||
managed_client=managed_client,
|
||||
database_name=mongodb_atlas_settings.database_name,
|
||||
embedding_generator=embedding_generator,
|
||||
)
|
||||
|
||||
@override
|
||||
def get_collection(
|
||||
self,
|
||||
record_type: type[TModel],
|
||||
*,
|
||||
definition: VectorStoreCollectionDefinition | None = None,
|
||||
collection_name: str | None = None,
|
||||
embedding_generator: EmbeddingGeneratorBase | None = None,
|
||||
**kwargs: Any,
|
||||
) -> MongoDBAtlasCollection:
|
||||
return MongoDBAtlasCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
mongo_client=self.mongo_client,
|
||||
managed_client=False,
|
||||
database_name=self.database_name,
|
||||
embedding_generator=embedding_generator or self.embedding_generator,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@override
|
||||
async def list_collection_names(self, **kwargs: Any) -> list[str]:
|
||||
database: AsyncDatabase = self.mongo_client.get_database(self.database_name)
|
||||
return await database.list_collection_names()
|
||||
|
||||
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
|
||||
"""Exit the context manager."""
|
||||
if self.managed_client:
|
||||
await self.mongo_client.close()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
"""Enter the context manager."""
|
||||
await self.mongo_client.aconnect()
|
||||
return self
|
||||
Reference in New Issue
Block a user