diff --git a/pinecone/__init__.py b/pinecone/__init__.py index d6d73c954..6a66bb226 100644 --- a/pinecone/__init__.py +++ b/pinecone/__init__.py @@ -71,6 +71,7 @@ "IndexList": ("pinecone.db_control.models", "IndexList"), "IndexModel": ("pinecone.db_control.models", "IndexModel"), "IndexEmbed": ("pinecone.db_control.models", "IndexEmbed"), + "ByocSpec": ("pinecone.db_control.models", "ByocSpec"), "ServerlessSpec": ("pinecone.db_control.models", "ServerlessSpec"), "ServerlessSpecDefinition": ("pinecone.db_control.models", "ServerlessSpecDefinition"), "PodSpec": ("pinecone.db_control.models", "PodSpec"), diff --git a/pinecone/db_control/models/__init__.py b/pinecone/db_control/models/__init__.py index 66568de33..cf866f116 100644 --- a/pinecone/db_control/models/__init__.py +++ b/pinecone/db_control/models/__init__.py @@ -2,6 +2,7 @@ from .collection_description import CollectionDescription from .serverless_spec import ServerlessSpec from .pod_spec import PodSpec +from .byoc_spec import ByocSpec from .index_list import IndexList from .collection_list import CollectionList from .index_model import IndexModel @@ -18,6 +19,7 @@ "PodSpecDefinition", "ServerlessSpec", "ServerlessSpecDefinition", + "ByocSpec", "IndexList", "CollectionList", "IndexModel", diff --git a/pinecone/db_control/models/byoc_spec.py b/pinecone/db_control/models/byoc_spec.py new file mode 100644 index 000000000..ccbdff4a0 --- /dev/null +++ b/pinecone/db_control/models/byoc_spec.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ByocSpec: + """ + ByocSpec represents the configuration used to deploy a BYOC (Bring Your Own Cloud) index. + + To learn more about the options for each configuration, please see [Understanding Indexes](https://docs.pinecone.io/docs/indexes) + """ + + environment: str diff --git a/pinecone/db_control/request_factory.py b/pinecone/db_control/request_factory.py index a5d298ca7..76fbd6a0a 100644 --- a/pinecone/db_control/request_factory.py +++ b/pinecone/db_control/request_factory.py @@ -27,12 +27,13 @@ from pinecone.core.openapi.db_control.model.serverless_spec import ( ServerlessSpec as ServerlessSpecModel, ) +from pinecone.core.openapi.db_control.model.byoc_spec import ByocSpec as ByocSpecModel from pinecone.core.openapi.db_control.model.pod_spec import PodSpec as PodSpecModel from pinecone.core.openapi.db_control.model.pod_spec_metadata_config import PodSpecMetadataConfig from pinecone.core.openapi.db_control.model.create_index_from_backup_request import ( CreateIndexFromBackupRequest, ) -from pinecone.db_control.models import ServerlessSpec, PodSpec, IndexModel, IndexEmbed +from pinecone.db_control.models import ServerlessSpec, PodSpec, ByocSpec, IndexModel, IndexEmbed from pinecone.db_control.enums import ( Metric, @@ -76,7 +77,7 @@ def __parse_deletion_protection( raise ValueError("deletion_protection must be either 'enabled' or 'disabled'") @staticmethod - def __parse_index_spec(spec: Union[Dict, ServerlessSpec, PodSpec]) -> IndexSpec: + def __parse_index_spec(spec: Union[Dict, ServerlessSpec, PodSpec, ByocSpec]) -> IndexSpec: if isinstance(spec, dict): if "serverless" in spec: spec["serverless"]["cloud"] = convert_enum_to_string(spec["serverless"]["cloud"]) @@ -100,8 +101,10 @@ def __parse_index_spec(spec: Union[Dict, ServerlessSpec, PodSpec]) -> IndexSpec: indexed=args_dict["metadata_config"].get("indexed", None) ) index_spec = IndexSpec(pod=PodSpecModel(**args_dict)) + elif "byoc" in spec: + index_spec = IndexSpec(byoc=ByocSpecModel(**spec["byoc"])) else: - raise ValueError("spec must contain either 'serverless' or 'pod' key") + raise ValueError("spec must contain either 'serverless', 'pod', or 'byoc' key") elif isinstance(spec, ServerlessSpec): index_spec = IndexSpec( serverless=ServerlessSpecModel(cloud=spec.cloud, region=spec.region) @@ -123,15 +126,18 @@ def __parse_index_spec(spec: Union[Dict, ServerlessSpec, PodSpec]) -> IndexSpec: index_spec = IndexSpec( pod=PodSpecModel(environment=spec.environment, pod_type=spec.pod_type, **args_dict) ) + elif isinstance(spec, ByocSpec): + args_dict = parse_non_empty_args([("environment", spec.environment)]) + index_spec = IndexSpec(byoc=ByocSpecModel(**args_dict)) else: - raise TypeError("spec must be of type dict, ServerlessSpec, or PodSpec") + raise TypeError("spec must be of type dict, ServerlessSpec, PodSpec, or ByocSpec") return index_spec @staticmethod def create_index_request( name: str, - spec: Union[Dict, ServerlessSpec, PodSpec], + spec: Union[Dict, ServerlessSpec, PodSpec, ByocSpec], dimension: Optional[int] = None, metric: Optional[Union[Metric, str]] = Metric.COSINE, deletion_protection: Optional[Union[DeletionProtection, str]] = DeletionProtection.DISABLED, diff --git a/pinecone/db_control/resources/asyncio/index.py b/pinecone/db_control/resources/asyncio/index.py index b48ff99cb..7bb10404e 100644 --- a/pinecone/db_control/resources/asyncio/index.py +++ b/pinecone/db_control/resources/asyncio/index.py @@ -3,7 +3,14 @@ from typing import Optional, Dict, Union -from pinecone.db_control.models import ServerlessSpec, PodSpec, IndexModel, IndexList, IndexEmbed +from pinecone.db_control.models import ( + ServerlessSpec, + PodSpec, + ByocSpec, + IndexModel, + IndexList, + IndexEmbed, +) from pinecone.utils import docslinks from pinecone.db_control.enums import ( @@ -33,7 +40,7 @@ def __init__(self, index_api, config): async def create( self, name: str, - spec: Union[Dict, ServerlessSpec, PodSpec], + spec: Union[Dict, ServerlessSpec, PodSpec, ByocSpec], dimension: Optional[int] = None, metric: Optional[Union[Metric, str]] = Metric.COSINE, timeout: Optional[int] = None, diff --git a/pinecone/db_control/resources/sync/index.py b/pinecone/db_control/resources/sync/index.py index d5e7d6e2d..050683833 100644 --- a/pinecone/db_control/resources/sync/index.py +++ b/pinecone/db_control/resources/sync/index.py @@ -4,7 +4,14 @@ from pinecone.db_control.index_host_store import IndexHostStore -from pinecone.db_control.models import ServerlessSpec, PodSpec, IndexModel, IndexList, IndexEmbed +from pinecone.db_control.models import ( + ServerlessSpec, + PodSpec, + ByocSpec, + IndexModel, + IndexList, + IndexEmbed, +) from pinecone.utils import docslinks, require_kwargs from pinecone.db_control.enums import ( @@ -39,7 +46,7 @@ def __init__(self, index_api, config): def create( self, name: str, - spec: Union[Dict, ServerlessSpec, PodSpec], + spec: Union[Dict, ServerlessSpec, PodSpec, ByocSpec], dimension: Optional[int] = None, metric: Optional[Union[Metric, str]] = Metric.COSINE, timeout: Optional[int] = None, diff --git a/pinecone/legacy_pinecone_interface.py b/pinecone/legacy_pinecone_interface.py index cb896022e..ad315d70f 100644 --- a/pinecone/legacy_pinecone_interface.py +++ b/pinecone/legacy_pinecone_interface.py @@ -6,6 +6,7 @@ from pinecone.db_control.models import ( ServerlessSpec, PodSpec, + ByocSpec, IndexList, CollectionList, IndexModel, @@ -194,7 +195,7 @@ def __init__( def create_index( self, name: str, - spec: Union[Dict, "ServerlessSpec", "PodSpec"], + spec: Union[Dict, "ServerlessSpec", "PodSpec", "ByocSpec"], dimension: Optional[int], metric: Optional[Union["Metric", str]] = "Metric.COSINE", timeout: Optional[int] = None, @@ -214,7 +215,7 @@ def create_index( :type metric: str, optional :param spec: A dictionary containing configurations describing how the index should be deployed. For serverless indexes, specify region and cloud. For pod indexes, specify replicas, shards, pods, pod_type, metadata_config, and source_collection. - Alternatively, use the `ServerlessSpec` or `PodSpec` objects to specify these configurations. + Alternatively, use the `ServerlessSpec`, `PodSpec`, or `ByocSpec` objects to specify these configurations. :type spec: Dict :param dimension: If you are creating an index with `vector_type="dense"` (which is the default), you need to specify `dimension` to indicate the size of your vectors. This should match the dimension of the embeddings you will be inserting. For example, if you are using diff --git a/pinecone/pinecone.py b/pinecone/pinecone.py index ae854129d..25d48cd7a 100644 --- a/pinecone/pinecone.py +++ b/pinecone/pinecone.py @@ -37,6 +37,7 @@ from pinecone.db_control.models import ( ServerlessSpec, PodSpec, + ByocSpec, IndexModel, IndexList, CollectionList, @@ -177,7 +178,7 @@ def index_api(self) -> "ManageIndexesApi": def create_index( self, name: str, - spec: Union[Dict, "ServerlessSpec", "PodSpec"], + spec: Union[Dict, "ServerlessSpec", "PodSpec", "ByocSpec"], dimension: Optional[int] = None, metric: Optional[Union["Metric", str]] = "cosine", timeout: Optional[int] = None, diff --git a/pinecone/pinecone_asyncio.py b/pinecone/pinecone_asyncio.py index 124ac854c..5133f7bdc 100644 --- a/pinecone/pinecone_asyncio.py +++ b/pinecone/pinecone_asyncio.py @@ -26,6 +26,7 @@ from pinecone.db_control.models import ( ServerlessSpec, PodSpec, + ByocSpec, IndexModel, IndexList, CollectionList, @@ -195,7 +196,7 @@ def index_api(self) -> "AsyncioManageIndexesApi": async def create_index( self, name: str, - spec: Union[Dict, "ServerlessSpec", "PodSpec"], + spec: Union[Dict, "ServerlessSpec", "PodSpec", "ByocSpec"], dimension: Optional[int] = None, metric: Optional[Union["Metric", str]] = "cosine", timeout: Optional[int] = None, diff --git a/pinecone/pinecone_interface_asyncio.py b/pinecone/pinecone_interface_asyncio.py index 6dfd953c9..4b8e1cc11 100644 --- a/pinecone/pinecone_interface_asyncio.py +++ b/pinecone/pinecone_interface_asyncio.py @@ -10,6 +10,7 @@ from pinecone.db_control.models import ( ServerlessSpec, PodSpec, + ByocSpec, IndexList, CollectionList, IndexModel, @@ -294,14 +295,12 @@ async def main(): async def create_index( self, name: str, - spec: Union[Dict, "ServerlessSpec", "PodSpec"], + spec: Union[Dict, "ServerlessSpec", "PodSpec", "ByocSpec"], dimension: Optional[int], - metric: Optional[Union["Metric", str]] = "Metric.COSINE", + metric: Optional[Union["Metric", str]] = "cosine", timeout: Optional[int] = None, - deletion_protection: Optional[ - Union["DeletionProtection", str] - ] = "DeletionProtection.DISABLED", - vector_type: Optional[Union["VectorType", str]] = "VectorType.DENSE", + deletion_protection: Optional[Union["DeletionProtection", str]] = "disabled", + vector_type: Optional[Union["VectorType", str]] = "dense", tags: Optional[Dict[str, str]] = None, ): """Creates a Pinecone index. @@ -417,9 +416,7 @@ async def create_index_for_model( region: Union["AwsRegion", "GcpRegion", "AzureRegion", str], embed: Union["IndexEmbed", "CreateIndexForModelEmbedTypedDict"], tags: Optional[Dict[str, str]] = None, - deletion_protection: Optional[ - Union["DeletionProtection", str] - ] = "DeletionProtection.DISABLED", + deletion_protection: Optional[Union["DeletionProtection", str]] = "disabled", timeout: Optional[int] = None, ) -> "IndexModel": """ diff --git a/tests/unit/db_control/test_index.py b/tests/unit/db_control/test_index.py new file mode 100644 index 000000000..36def11c5 --- /dev/null +++ b/tests/unit/db_control/test_index.py @@ -0,0 +1,62 @@ +import json + +from pinecone import Config + +from pinecone.db_control.resources.sync.index import IndexResource +from pinecone.openapi_support.api_client import ApiClient +from pinecone.core.openapi.db_control.api.manage_indexes_api import ManageIndexesApi + + +def build_client_w_faked_response(mocker, body: str, status: int = 200): + response = mocker.Mock() + response.headers = {"content-type": "application/json"} + response.status = status + # Parse the JSON string into a dict + response_data = json.loads(body) + response.data = json.dumps(response_data).encode("utf-8") + + api_client = ApiClient() + mock_request = mocker.patch.object( + api_client.rest_client.pool_manager, "request", return_value=response + ) + index_api = ManageIndexesApi(api_client=api_client) + return IndexResource(index_api=index_api, config=Config(api_key="test-api-key")), mock_request + + +class TestIndexResource: + def test_describe_index(self, mocker): + body = """ + { + "name": "test-index", + "description": "test-description", + "dimension": 1024, + "metric": "cosine", + "spec": { + "byoc": { + "environment": "test-environment" + } + }, + "vector_type": "dense", + "status": { + "ready": true, + "state": "Ready" + }, + "host": "test-host.pinecone.io", + "deletion_protection": "disabled", + "tags": { + "test-tag": "test-value" + } + } + """ + index_resource, mock_request = build_client_w_faked_response(mocker, body) + + desc = index_resource.describe(name="test-index") + assert desc.name == "test-index" + assert desc.description == "test-description" + assert desc.dimension == 1024 + assert desc.metric == "cosine" + assert desc.spec.byoc.environment == "test-environment" + assert desc.vector_type == "dense" + assert desc.status.ready == True + assert desc.deletion_protection == "disabled" + assert desc.tags["test-tag"] == "test-value" diff --git a/tests/unit/db_control/test_index_request_factory.py b/tests/unit/db_control/test_index_request_factory.py new file mode 100644 index 000000000..777486b59 --- /dev/null +++ b/tests/unit/db_control/test_index_request_factory.py @@ -0,0 +1,62 @@ +from pinecone import ByocSpec, ServerlessSpec +from pinecone.db_control.request_factory import PineconeDBControlRequestFactory + + +class TestIndexRequestFactory: + def test_create_index_request_with_spec_byoc(self): + req = PineconeDBControlRequestFactory.create_index_request( + name="test-index", + metric="cosine", + dimension=1024, + spec=ByocSpec(environment="test-byoc-spec-id"), + ) + assert req.name == "test-index" + assert req.metric == "cosine" + assert req.dimension == 1024 + assert req.spec.byoc.environment == "test-byoc-spec-id" + assert req.vector_type == "dense" + assert req.deletion_protection.value == "disabled" + + def test_create_index_request_with_spec_serverless(self): + req = PineconeDBControlRequestFactory.create_index_request( + name="test-index", + metric="cosine", + dimension=1024, + spec=ServerlessSpec(cloud="aws", region="us-east-1"), + ) + assert req.name == "test-index" + assert req.metric == "cosine" + assert req.dimension == 1024 + assert req.spec.serverless.cloud == "aws" + assert req.spec.serverless.region == "us-east-1" + assert req.vector_type == "dense" + assert req.deletion_protection.value == "disabled" + + def test_create_index_request_with_spec_serverless_dict(self): + req = PineconeDBControlRequestFactory.create_index_request( + name="test-index", + metric="cosine", + dimension=1024, + spec={"serverless": {"cloud": "aws", "region": "us-east-1"}}, + ) + assert req.name == "test-index" + assert req.metric == "cosine" + assert req.dimension == 1024 + assert req.spec.serverless.cloud == "aws" + assert req.spec.serverless.region == "us-east-1" + assert req.vector_type == "dense" + assert req.deletion_protection.value == "disabled" + + def test_create_index_request_with_spec_byoc_dict(self): + req = PineconeDBControlRequestFactory.create_index_request( + name="test-index", + metric="cosine", + dimension=1024, + spec={"byoc": {"environment": "test-byoc-spec-id"}}, + ) + assert req.name == "test-index" + assert req.metric == "cosine" + assert req.dimension == 1024 + assert req.spec.byoc.environment == "test-byoc-spec-id" + assert req.vector_type == "dense" + assert req.deletion_protection.value == "disabled"