Add enums for cloud parameters

pinecone/models/__init__.py

@@ -5,6 +5,7 @@
 from .index_list import IndexList
 from .collection_list import CollectionList
 from .index_model import IndexModel
+from .clouds import CloudProvider, AwsRegion, GcpRegion, AzureRegion, PodIndexEnvironment
 
 __all__ = [
     "CollectionDescription",
@@ -15,4 +16,9 @@
     "IndexList",
     "CollectionList",
     "IndexModel",
+    "CloudProvider",
+    "AwsRegion",
+    "GcpRegion",
+    "AzureRegion",
+    "PodIndexEnvironment",
 ]

pinecone/models/clouds.py

@@ -0,0 +1,37 @@
+from enum import Enum
+
+
+class CloudProvider(Enum):
+    AWS = "aws"
+    GCP = "gcp"
+    AZURE = "azure"
+
+
+class AwsRegion(Enum):
+    US_EAST_1 = "us-east-1"
+    US_WEST_2 = "us-west-2"
+    EU_WEST_1 = "eu-west-1"
+
+
+class GcpRegion(Enum):
+    US_CENTRAL1 = "us-central1"
+    EUROPE_WEST4 = "europe-west4"
+
+
+class AzureRegion(Enum):
+    EAST_US = "eastus2"
+
+
+class PodIndexEnvironment(Enum):
+    US_WEST1_GCP = "us-west1-gcp"
+    US_CENTRAL1_GCP = "us-central1-gcp"
+    US_WEST4_GCP = "us-west4-gcp"
+    US_EAST4_GCP = "us-east4-gcp"
+    NORTHAMERICA_NORTHEAST1_GCP = "northamerica-northeast1-gcp"
+    ASIA_NORTHEAST1_GCP = "asia-northeast1-gcp"
+    ASIA_SOUTHEAST1_GCP = "asia-southeast1-gcp"
+    US_EAST1_GCP = "us-east1-gcp"
+    EU_WEST1_GCP = "eu-west1-gcp"
+    EU_WEST4_GCP = "eu-west4-gcp"
+    US_EAST1_AWS = "us-east-1-aws"
+    EASTUS_AZURE = "eastus-azure"

pinecone/models/serverless_spec.py

@@ -1,9 +1,25 @@
-from typing import NamedTuple
+from dataclasses import dataclass
+from typing import Union
+from enum import Enum
 
+from .clouds import CloudProvider, AwsRegion, GcpRegion, AzureRegion
 
-class ServerlessSpec(NamedTuple):
+
+@dataclass(frozen=True)
+class ServerlessSpec:
     cloud: str
     region: str
 
+    def __init__(
+        self,
+        cloud: Union[CloudProvider, str],
+        region: Union[AwsRegion, GcpRegion, AzureRegion, str],
+    ):
+        # Convert Enums to their string values if necessary
+        object.__setattr__(self, "cloud", cloud.value if isinstance(cloud, Enum) else str(cloud))
+        object.__setattr__(
+            self, "region", region.value if isinstance(region, Enum) else str(region)
+        )
+
     def asdict(self):
-        return {"serverless": self._asdict()}
+        return {"serverless": {"cloud": self.cloud, "region": self.region}}

tests/unit/models/test_index_model.py

@@ -6,6 +6,7 @@
     DeletionProtection,
 )
 from pinecone.models import IndexModel
+from pinecone import CloudProvider, AwsRegion
 
 
 class TestIndexModel:
@@ -17,7 +18,11 @@ def test_index_model(self):
             host="https://test-index-1.pinecone.io",
             status=IndexModelStatus(ready=True, state="Ready"),
             deletion_protection=DeletionProtection("enabled"),
-            spec=IndexModelSpec(serverless=ServerlessSpec(cloud="aws", region="us-west-1")),
+            spec=IndexModelSpec(
+                serverless=ServerlessSpec(
+                    cloud=CloudProvider.AWS.value, region=AwsRegion.US_EAST_1.value
+                )
+            ),
         )
 
         wrapped = IndexModel(openapi_model)

tests/unit/test_control.py

@@ -1,7 +1,15 @@
 import pytest
 import re
 from unittest.mock import patch, MagicMock
-from pinecone import ConfigBuilder, Pinecone, PodSpec, ServerlessSpec
+from pinecone import (
+    ConfigBuilder,
+    Pinecone,
+    PodSpec,
+    ServerlessSpec,
+    CloudProvider,
+    AwsRegion,
+    GcpRegion,
+)
 from pinecone.core.openapi.db_control.models import IndexList, IndexModel, DeletionProtection
 from pinecone.core.openapi.db_control.api.manage_indexes_api import ManageIndexesApi
 
@@ -150,8 +158,15 @@ def test_create_index_with_timeout(
     @pytest.mark.parametrize(
         "index_spec",
         [
+            ServerlessSpec(cloud="aws", region="us-west-2"),
+            ServerlessSpec(cloud=CloudProvider.AWS, region=AwsRegion.US_WEST_2),
+            ServerlessSpec(cloud=CloudProvider.AWS, region="us-west-2"),
+            ServerlessSpec(cloud="aws", region="us-west-2"),
+            ServerlessSpec(cloud="aws", region="unknown-region"),
+            ServerlessSpec(cloud=CloudProvider.GCP, region=GcpRegion.US_CENTRAL1),
             {"serverless": {"cloud": "aws", "region": "us-west1"}},
             {"serverless": {"cloud": "aws", "region": "us-west1", "uknown_key": "value"}},
+            PodSpec(environment="us-west1-gcp", pod_type="p1.x1"),
             {"pod": {"environment": "us-west1-gcp", "pod_type": "p1.x1"}},
             {"pod": {"environment": "us-west1-gcp", "pod_type": "p1.x1", "unknown_key": "value"}},
             {