Skip to content

Commit 0e784a9

Browse files
jhamonjhamon
committed
Implement BYOC
1 parent 5c0e37c commit 0e784a9

12 files changed

+181
-22
lines changed

‎pinecone/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
"IndexList": ("pinecone.db_control.models", "IndexList"),
7272
"IndexModel": ("pinecone.db_control.models", "IndexModel"),
7373
"IndexEmbed": ("pinecone.db_control.models", "IndexEmbed"),
74+
"ByocSpec": ("pinecone.db_control.models", "ByocSpec"),
7475
"ServerlessSpec": ("pinecone.db_control.models", "ServerlessSpec"),
7576
"ServerlessSpecDefinition": ("pinecone.db_control.models", "ServerlessSpecDefinition"),
7677
"PodSpec": ("pinecone.db_control.models", "PodSpec"),

‎pinecone/db_control/models/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from .collection_description import CollectionDescription
33
from .serverless_spec import ServerlessSpec
44
from .pod_spec import PodSpec
5+
from .byoc_spec import ByocSpec
56
from .index_list import IndexList
67
from .collection_list import CollectionList
78
from .index_model import IndexModel
@@ -18,6 +19,7 @@
1819
"PodSpecDefinition",
1920
"ServerlessSpec",
2021
"ServerlessSpecDefinition",
22+
"ByocSpec",
2123
"IndexList",
2224
"CollectionList",
2325
"IndexModel",

‎pinecone/db_control/models/byoc_spec.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from dataclasses import dataclass
2+
3+
4+
@dataclass(frozen=True)
5+
class ByocSpec:
6+
"""
7+
ByocSpec represents the configuration used to deploy a BYOC (Bring Your Own Cloud) index.
8+
9+
To learn more about the options for each configuration, please see [Understanding Indexes](https://docs.pinecone.io/docs/indexes)
10+
"""
11+
12+
environment: str

‎pinecone/db_control/request_factory.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,13 @@
2727
from pinecone.core.openapi.db_control.model.serverless_spec import (
2828
ServerlessSpec as ServerlessSpecModel,
2929
)
30+
from pinecone.core.openapi.db_control.model.byoc_spec import ByocSpec as ByocSpecModel
3031
from pinecone.core.openapi.db_control.model.pod_spec import PodSpec as PodSpecModel
3132
from pinecone.core.openapi.db_control.model.pod_spec_metadata_config import PodSpecMetadataConfig
3233
from pinecone.core.openapi.db_control.model.create_index_from_backup_request import (
3334
CreateIndexFromBackupRequest,
3435
)
35-
from pinecone.db_control.models import ServerlessSpec, PodSpec, IndexModel, IndexEmbed
36+
from pinecone.db_control.models import ServerlessSpec, PodSpec, ByocSpec, IndexModel, IndexEmbed
3637

3738
from pinecone.db_control.enums import (
3839
Metric,
@@ -76,7 +77,7 @@ def __parse_deletion_protection(
7677
raise ValueError("deletion_protection must be either 'enabled' or 'disabled'")
7778

7879
@staticmethod
79-
def __parse_index_spec(spec: Union[Dict, ServerlessSpec, PodSpec]) -> IndexSpec:
80+
def __parse_index_spec(spec: Union[Dict, ServerlessSpec, PodSpec, ByocSpec]) -> IndexSpec:
8081
if isinstance(spec, dict):
8182
if "serverless" in spec:
8283
spec["serverless"]["cloud"] = convert_enum_to_string(spec["serverless"]["cloud"])
@@ -100,8 +101,10 @@ def __parse_index_spec(spec: Union[Dict, ServerlessSpec, PodSpec]) -> IndexSpec:
100101
indexed=args_dict["metadata_config"].get("indexed", None)
101102
)
102103
index_spec = IndexSpec(pod=PodSpecModel(**args_dict))
104+
elif "byoc" in spec:
105+
index_spec = IndexSpec(byoc=ByocSpecModel(**spec["byoc"]))
103106
else:
104-
raise ValueError("spec must contain either 'serverless' or 'pod' key")
107+
raise ValueError("spec must contain either 'serverless', 'pod', or 'byoc' key")
105108
elif isinstance(spec, ServerlessSpec):
106109
index_spec = IndexSpec(
107110
serverless=ServerlessSpecModel(cloud=spec.cloud, region=spec.region)
@@ -123,15 +126,18 @@ def __parse_index_spec(spec: Union[Dict, ServerlessSpec, PodSpec]) -> IndexSpec:
123126
index_spec = IndexSpec(
124127
pod=PodSpecModel(environment=spec.environment, pod_type=spec.pod_type, **args_dict)
125128
)
129+
elif isinstance(spec, ByocSpec):
130+
args_dict = parse_non_empty_args([("environment", spec.environment)])
131+
index_spec = IndexSpec(byoc=ByocSpecModel(**args_dict))
126132
else:
127-
raise TypeError("spec must be of type dict, ServerlessSpec, or PodSpec")
133+
raise TypeError("spec must be of type dict, ServerlessSpec, PodSpec, or ByocSpec")
128134

129135
return index_spec
130136

131137
@staticmethod
132138
def create_index_request(
133139
name: str,
134-
spec: Union[Dict, ServerlessSpec, PodSpec],
140+
spec: Union[Dict, ServerlessSpec, PodSpec, ByocSpec],
135141
dimension: Optional[int] = None,
136142
metric: Optional[Union[Metric, str]] = Metric.COSINE,
137143
deletion_protection: Optional[Union[DeletionProtection, str]] = DeletionProtection.DISABLED,

‎pinecone/db_control/resources/asyncio/index.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,14 @@
33
from typing import Optional, Dict, Union
44

55

6-
from pinecone.db_control.models import ServerlessSpec, PodSpec, IndexModel, IndexList, IndexEmbed
6+
from pinecone.db_control.models import (
7+
ServerlessSpec,
8+
PodSpec,
9+
ByocSpec,
10+
IndexModel,
11+
IndexList,
12+
IndexEmbed,
13+
)
714
from pinecone.utils import docslinks
815

916
from pinecone.db_control.enums import (
@@ -33,7 +40,7 @@ def __init__(self, index_api, config):
3340
async def create(
3441
self,
3542
name: str,
36-
spec: Union[Dict, ServerlessSpec, PodSpec],
43+
spec: Union[Dict, ServerlessSpec, PodSpec, ByocSpec],
3744
dimension: Optional[int] = None,
3845
metric: Optional[Union[Metric, str]] = Metric.COSINE,
3946
timeout: Optional[int] = None,

‎pinecone/db_control/resources/sync/index.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,14 @@
44

55
from pinecone.db_control.index_host_store import IndexHostStore
66

7-
from pinecone.db_control.models import ServerlessSpec, PodSpec, IndexModel, IndexList, IndexEmbed
7+
from pinecone.db_control.models import (
8+
ServerlessSpec,
9+
PodSpec,
10+
ByocSpec,
11+
IndexModel,
12+
IndexList,
13+
IndexEmbed,
14+
)
815
from pinecone.utils import docslinks, require_kwargs
916

1017
from pinecone.db_control.enums import (
@@ -39,7 +46,7 @@ def __init__(self, index_api, config):
3946
def create(
4047
self,
4148
name: str,
42-
spec: Union[Dict, ServerlessSpec, PodSpec],
49+
spec: Union[Dict, ServerlessSpec, PodSpec, ByocSpec],
4350
dimension: Optional[int] = None,
4451
metric: Optional[Union[Metric, str]] = Metric.COSINE,
4552
timeout: Optional[int] = None,

‎pinecone/legacy_pinecone_interface.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from pinecone.db_control.models import (
77
ServerlessSpec,
88
PodSpec,
9+
ByocSpec,
910
IndexList,
1011
CollectionList,
1112
IndexModel,
@@ -194,7 +195,7 @@ def __init__(
194195
def create_index(
195196
self,
196197
name: str,
197-
spec: Union[Dict, "ServerlessSpec", "PodSpec"],
198+
spec: Union[Dict, "ServerlessSpec", "PodSpec", "ByocSpec"],
198199
dimension: Optional[int],
199200
metric: Optional[Union["Metric", str]] = "Metric.COSINE",
200201
timeout: Optional[int] = None,
@@ -214,7 +215,7 @@ def create_index(
214215
:type metric: str, optional
215216
:param spec: A dictionary containing configurations describing how the index should be deployed. For serverless indexes,
216217
specify region and cloud. For pod indexes, specify replicas, shards, pods, pod_type, metadata_config, and source_collection.
217-
Alternatively, use the `ServerlessSpec` or `PodSpec` objects to specify these configurations.
218+
Alternatively, use the `ServerlessSpec`, `PodSpec`, or `ByocSpec` objects to specify these configurations.
218219
:type spec: Dict
219220
: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.
220221
This should match the dimension of the embeddings you will be inserting. For example, if you are using

‎pinecone/pinecone.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
from pinecone.db_control.models import (
3838
ServerlessSpec,
3939
PodSpec,
40+
ByocSpec,
4041
IndexModel,
4142
IndexList,
4243
CollectionList,
@@ -177,7 +178,7 @@ def index_api(self) -> "ManageIndexesApi":
177178
def create_index(
178179
self,
179180
name: str,
180-
spec: Union[Dict, "ServerlessSpec", "PodSpec"],
181+
spec: Union[Dict, "ServerlessSpec", "PodSpec", "ByocSpec"],
181182
dimension: Optional[int] = None,
182183
metric: Optional[Union["Metric", str]] = "cosine",
183184
timeout: Optional[int] = None,

‎pinecone/pinecone_asyncio.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from pinecone.db_control.models import (
2727
ServerlessSpec,
2828
PodSpec,
29+
ByocSpec,
2930
IndexModel,
3031
IndexList,
3132
CollectionList,
@@ -195,7 +196,7 @@ def index_api(self) -> "AsyncioManageIndexesApi":
195196
async def create_index(
196197
self,
197198
name: str,
198-
spec: Union[Dict, "ServerlessSpec", "PodSpec"],
199+
spec: Union[Dict, "ServerlessSpec", "PodSpec", "ByocSpec"],
199200
dimension: Optional[int] = None,
200201
metric: Optional[Union["Metric", str]] = "cosine",
201202
timeout: Optional[int] = None,

‎pinecone/pinecone_interface_asyncio.py

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from pinecone.db_control.models import (
1111
ServerlessSpec,
1212
PodSpec,
13+
ByocSpec,
1314
IndexList,
1415
CollectionList,
1516
IndexModel,
@@ -294,14 +295,12 @@ async def main():
294295
async def create_index(
295296
self,
296297
name: str,
297-
spec: Union[Dict, "ServerlessSpec", "PodSpec"],
298+
spec: Union[Dict, "ServerlessSpec", "PodSpec", "ByocSpec"],
298299
dimension: Optional[int],
299-
metric: Optional[Union["Metric", str]] = "Metric.COSINE",
300+
metric: Optional[Union["Metric", str]] = "cosine",
300301
timeout: Optional[int] = None,
301-
deletion_protection: Optional[
302-
Union["DeletionProtection", str]
303-
] = "DeletionProtection.DISABLED",
304-
vector_type: Optional[Union["VectorType", str]] = "VectorType.DENSE",
302+
deletion_protection: Optional[Union["DeletionProtection", str]] = "disabled",
303+
vector_type: Optional[Union["VectorType", str]] = "dense",
305304
tags: Optional[Dict[str, str]] = None,
306305
):
307306
"""Creates a Pinecone index.
@@ -417,9 +416,7 @@ async def create_index_for_model(
417416
region: Union["AwsRegion", "GcpRegion", "AzureRegion", str],
418417
embed: Union["IndexEmbed", "CreateIndexForModelEmbedTypedDict"],
419418
tags: Optional[Dict[str, str]] = None,
420-
deletion_protection: Optional[
421-
Union["DeletionProtection", str]
422-
] = "DeletionProtection.DISABLED",
419+
deletion_protection: Optional[Union["DeletionProtection", str]] = "disabled",
423420
timeout: Optional[int] = None,
424421
) -> "IndexModel":
425422
"""

‎tests/unit/db_control/test_index.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
import json
2+
3+
from pinecone import Config
4+
5+
from pinecone.db_control.resources.sync.index import IndexResource
6+
from pinecone.openapi_support.api_client import ApiClient
7+
from pinecone.core.openapi.db_control.api.manage_indexes_api import ManageIndexesApi
8+
9+
10+
def build_client_w_faked_response(mocker, body: str, status: int = 200):
11+
response = mocker.Mock()
12+
response.headers = {"content-type": "application/json"}
13+
response.status = status
14+
# Parse the JSON string into a dict
15+
response_data = json.loads(body)
16+
response.data = json.dumps(response_data).encode("utf-8")
17+
18+
api_client = ApiClient()
19+
mock_request = mocker.patch.object(
20+
api_client.rest_client.pool_manager, "request", return_value=response
21+
)
22+
index_api = ManageIndexesApi(api_client=api_client)
23+
return IndexResource(index_api=index_api, config=Config(api_key="test-api-key")), mock_request
24+
25+
26+
class TestIndexResource:
27+
def test_describe_index(self, mocker):
28+
body = """
29+
{
30+
"name": "test-index",
31+
"description": "test-description",
32+
"dimension": 1024,
33+
"metric": "cosine",
34+
"spec": {
35+
"byoc": {
36+
"environment": "test-environment"
37+
}
38+
},
39+
"vector_type": "dense",
40+
"status": {
41+
"ready": true,
42+
"state": "Ready"
43+
},
44+
"host": "test-host.pinecone.io",
45+
"deletion_protection": "disabled",
46+
"tags": {
47+
"test-tag": "test-value"
48+
}
49+
}
50+
"""
51+
index_resource, mock_request = build_client_w_faked_response(mocker, body)
52+
53+
desc = index_resource.describe(name="test-index")
54+
assert desc.name == "test-index"
55+
assert desc.description == "test-description"
56+
assert desc.dimension == 1024
57+
assert desc.metric == "cosine"
58+
assert desc.spec.byoc.environment == "test-environment"
59+
assert desc.vector_type == "dense"
60+
assert desc.status.ready == True
61+
assert desc.deletion_protection == "disabled"
62+
assert desc.tags["test-tag"] == "test-value"

‎tests/unit/db_control/test_index_request_factory.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
from pinecone import ByocSpec, ServerlessSpec
2+
from pinecone.db_control.request_factory import PineconeDBControlRequestFactory
3+
4+
5+
class TestIndexRequestFactory:
6+
def test_create_index_request_with_spec_byoc(self):
7+
req = PineconeDBControlRequestFactory.create_index_request(
8+
name="test-index",
9+
metric="cosine",
10+
dimension=1024,
11+
spec=ByocSpec(environment="test-byoc-spec-id"),
12+
)
13+
assert req.name == "test-index"
14+
assert req.metric == "cosine"
15+
assert req.dimension == 1024
16+
assert req.spec.byoc.environment == "test-byoc-spec-id"
17+
assert req.vector_type == "dense"
18+
assert req.deletion_protection.value == "disabled"
19+
20+
def test_create_index_request_with_spec_serverless(self):
21+
req = PineconeDBControlRequestFactory.create_index_request(
22+
name="test-index",
23+
metric="cosine",
24+
dimension=1024,
25+
spec=ServerlessSpec(cloud="aws", region="us-east-1"),
26+
)
27+
assert req.name == "test-index"
28+
assert req.metric == "cosine"
29+
assert req.dimension == 1024
30+
assert req.spec.serverless.cloud == "aws"
31+
assert req.spec.serverless.region == "us-east-1"
32+
assert req.vector_type == "dense"
33+
assert req.deletion_protection.value == "disabled"
34+
35+
def test_create_index_request_with_spec_serverless_dict(self):
36+
req = PineconeDBControlRequestFactory.create_index_request(
37+
name="test-index",
38+
metric="cosine",
39+
dimension=1024,
40+
spec={"serverless": {"cloud": "aws", "region": "us-east-1"}},
41+
)
42+
assert req.name == "test-index"
43+
assert req.metric == "cosine"
44+
assert req.dimension == 1024
45+
assert req.spec.serverless.cloud == "aws"
46+
assert req.spec.serverless.region == "us-east-1"
47+
assert req.vector_type == "dense"
48+
assert req.deletion_protection.value == "disabled"
49+
50+
def test_create_index_request_with_spec_byoc_dict(self):
51+
req = PineconeDBControlRequestFactory.create_index_request(
52+
name="test-index",
53+
metric="cosine",
54+
dimension=1024,
55+
spec={"byoc": {"environment": "test-byoc-spec-id"}},
56+
)
57+
assert req.name == "test-index"
58+
assert req.metric == "cosine"
59+
assert req.dimension == 1024
60+
assert req.spec.byoc.environment == "test-byoc-spec-id"
61+
assert req.vector_type == "dense"
62+
assert req.deletion_protection.value == "disabled"

0 commit comments

Comments
 (0)