diff --git a/mindsdb/integrations/handlers/s3vectors_handler/s3vectors_handler.py b/mindsdb/integrations/handlers/s3vectors_handler/s3vectors_handler.py index a60bcd2a727..8ed69cfeddf 100644 --- a/mindsdb/integrations/handlers/s3vectors_handler/s3vectors_handler.py +++ b/mindsdb/integrations/handlers/s3vectors_handler/s3vectors_handler.py @@ -236,6 +236,20 @@ def get_tables(self) -> HandlerResponse: logger.error(f"Error listing indexes: {e}") raise Exception(f"Error listing indexes from bucket '{self.vector_bucket}': {e}") + def get_dimension(self, table_name: str) -> int: + """Get the vector dimension configured for an S3 Vectors index.""" + connection = self.connect() + + try: + response = connection.get_index( + vectorBucketName=self.vector_bucket, + indexName=table_name, + ) + return response["index"]["dimension"] + except Exception as e: + logger.error(f"Error getting dimension for index '{table_name}': {e}") + raise Exception(f"Error getting dimension for index '{table_name}': {e}") + def create_table(self, table_name: str, if_not_exists=True): """Create a vector index with the given name in the S3 Vectors bucket.""" connection = self.connect() @@ -259,7 +273,7 @@ def create_table(self, table_name: str, if_not_exists=True): logger.info(f"Created index '{table_name}' in bucket '{self.vector_bucket}'") except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "") - if error_code == "IndexAlreadyExists" and if_not_exists: + if error_code in {"ConflictException", "IndexAlreadyExists"} and if_not_exists: logger.info(f"Index '{table_name}' already exists") return raise Exception(f"Error creating index '{table_name}': {e}") diff --git a/tests/unit/handlers/test_s3vectors_handler.py b/tests/unit/handlers/test_s3vectors_handler.py new file mode 100644 index 00000000000..3f2b3c412ee --- /dev/null +++ b/tests/unit/handlers/test_s3vectors_handler.py @@ -0,0 +1,63 @@ +import unittest +from unittest.mock import MagicMock + +from botocore.client import ClientError + +from mindsdb.integrations.handlers.s3vectors_handler.s3vectors_handler import S3VectorsHandler +from mindsdb.integrations.libs.vectordatabase_handler import TableField + + +class TestS3VectorsHandler(unittest.TestCase): + def setUp(self): + self.handler = S3VectorsHandler( + "s3vectors", + connection_data={"vector_bucket": "test-bucket"}, + ) + self.connection = MagicMock() + self.handler.connection = self.connection + self.handler.is_connected = True + + def test_get_dimension_returns_index_dimension(self): + self.connection.get_index.return_value = { + "index": { + "dimension": 3072, + }, + } + + dimension = self.handler.get_dimension("test-index") + + self.assertEqual(dimension, 3072) + self.connection.get_index.assert_called_once_with( + vectorBucketName="test-bucket", + indexName="test-index", + ) + + def test_create_table_ignores_existing_index_conflict_when_if_not_exists(self): + self.connection.create_index.side_effect = ClientError( + error_response={ + "Error": { + "Code": "ConflictException", + "Message": "An index with the specified name already exists", + }, + }, + operation_name="CreateIndex", + ) + + self.handler.create_table("test-index") + + self.connection.create_index.assert_called_once_with( + vectorBucketName="test-bucket", + indexName="test-index", + dataType="float32", + dimension=1536, + distanceMetric="cosine", + ) + + def test_select_without_search_vector_or_id_still_fails(self): + with self.assertRaisesRegex(Exception, "search_vector or an ID"): + self.handler.select( + "test-index", + columns=[TableField.EMBEDDINGS.value], + conditions=[], + limit=1, + )