pinecone.manage

  1import time
  2from typing import NamedTuple, Optional
  3import copy
  4
  5import pinecone
  6from pinecone.config import Config
  7from pinecone.core.client.api.index_operations_api import IndexOperationsApi
  8from pinecone.core.client.api_client import ApiClient
  9from pinecone.core.client.model.create_request import CreateRequest
 10from pinecone.core.client.model.patch_request import PatchRequest
 11from pinecone.core.client.model.create_collection_request import CreateCollectionRequest
 12from pinecone.core.utils import get_user_agent
 13
 14__all__ = [
 15    "create_index",
 16    "delete_index",
 17    "describe_index",
 18    "list_indexes",
 19    "scale_index",
 20    "create_collection",
 21    "describe_collection",
 22    "list_collections",
 23    "delete_collection",
 24    "configure_index",
 25    "CollectionDescription",
 26    "IndexDescription",
 27]
 28
 29
 30class IndexDescription(NamedTuple):
 31    name: str
 32    metric: str
 33    replicas: int
 34    dimension: int
 35    shards: int
 36    pods: int
 37    pod_type: str
 38    status: None
 39    metadata_config: None
 40    source_collection: None
 41
 42
 43class CollectionDescription(object):
 44    def __init__(self, keys, values):
 45        for k, v in zip(keys, values):
 46            self.__dict__[k] = v
 47
 48    def __str__(self):
 49        return str(self.__dict__)
 50
 51
 52def _get_api_instance():
 53    client_config = copy.deepcopy(Config.OPENAPI_CONFIG)
 54    client_config.api_key = client_config.api_key or {}
 55    client_config.api_key["ApiKeyAuth"] = client_config.api_key.get("ApiKeyAuth", Config.API_KEY)
 56    client_config.server_variables = {**{"environment": Config.ENVIRONMENT}, **client_config.server_variables}
 57
 58    # If a custom host has been passed with initialization pass it to the client_config
 59    if (Config.CONTROLLER_HOST):
 60        client_config.host = Config.CONTROLLER_HOST
 61
 62    api_client = ApiClient(configuration=client_config)
 63    api_client.user_agent = get_user_agent()
 64    api_instance = IndexOperationsApi(api_client)
 65    return api_instance
 66
 67
 68def _get_status(name: str):
 69    api_instance = _get_api_instance()
 70    response = api_instance.describe_index(name)
 71    return response["status"]
 72
 73
 74def create_index(
 75    name: str,
 76    dimension: int,
 77    timeout: int = None,
 78    index_type: str = "approximated",
 79    metric: str = "cosine",
 80    replicas: int = 1,
 81    shards: int = 1,
 82    pods: int = 1,
 83    pod_type: str = "p1",
 84    index_config: dict = None,
 85    metadata_config: dict = None,
 86    source_collection: str = "",
 87):
 88    """Creates a Pinecone index.
 89
 90    :param name: the name of the index.
 91    :type name: str
 92    :param dimension: the dimension of vectors that would be inserted in the index
 93    :param index_type: type of index, one of `{"approximated", "exact"}`, defaults to "approximated".
 94        The "approximated" index uses fast approximate search algorithms developed by Pinecone.
 95        The "exact" index uses accurate exact search algorithms.
 96        It performs exhaustive searches and thus it is usually slower than the "approximated" index.
 97    :type index_type: str, optional
 98    :param metric: type of metric used in the vector index, one of `{"cosine", "dotproduct", "euclidean"}`, defaults to "cosine".
 99        Use "cosine" for cosine similarity,
100        "dotproduct" for dot-product,
101        and "euclidean" for euclidean distance.
102    :type metric: str, optional
103    :param replicas: the number of replicas, defaults to 1.
104        Use at least 2 replicas if you need high availability (99.99% uptime) for querying.
105        For additional throughput (QPS) your index needs to support, provision additional replicas.
106    :type replicas: int, optional
107    :param shards: the number of shards per index, defaults to 1.
108        Use 1 shard per 1GB of vectors
109    :type shards: int,optional
110    :param pods: Total number of pods to be used by the index. pods = shard*replicas
111    :type pods: int,optional
112    :param pod_type: the pod type to be used for the index. can be one of p1 or s1.
113    :type pod_type: str,optional
114    :param index_config: Advanced configuration options for the index
115    :param metadata_config: Configuration related to the metadata index
116    :type metadata_config: dict, optional
117    :param source_collection: Collection name to create the index from
118    :type metadata_config: str, optional
119    :type timeout: int, optional
120    :param timeout: Timeout for wait until index gets ready. If None, wait indefinitely; if >=0, time out after this many seconds;
121        if -1, return immediately and do not wait. Default: None
122    """
123    api_instance = _get_api_instance()
124
125    api_instance.create_index(
126        create_request=CreateRequest(
127            name=name,
128            dimension=dimension,
129            index_type=index_type,
130            metric=metric,
131            replicas=replicas,
132            shards=shards,
133            pods=pods,
134            pod_type=pod_type,
135            index_config=index_config or {},
136            metadata_config=metadata_config,
137            source_collection=source_collection,
138        )
139    )
140
141    def is_ready():
142        status = _get_status(name)
143        ready = status["ready"]
144        return ready
145
146    if timeout == -1:
147        return
148    if timeout is None:
149        while not is_ready():
150            time.sleep(5)
151    else:
152        while (not is_ready()) and timeout >= 0:
153            time.sleep(5)
154            timeout -= 5
155    if timeout and timeout < 0:
156        raise (
157            TimeoutError(
158                "Please call the describe_index API ({}) to confirm index status.".format(
159                    "https://www.pinecone.io/docs/api/operation/describe_index/"
160                )
161            )
162        )
163
164
165def delete_index(name: str, timeout: int = None):
166    """Deletes a Pinecone index.
167
168    :param name: the name of the index.
169    :type name: str
170    :param timeout: Timeout for wait until index gets ready. If None, wait indefinitely; if >=0, time out after this many seconds;
171        if -1, return immediately and do not wait. Default: None
172    :type timeout: int, optional
173    """
174    api_instance = _get_api_instance()
175    api_instance.delete_index(name)
176
177    def get_remaining():
178        return name in api_instance.list_indexes()
179
180    if timeout == -1:
181        return
182
183    if timeout is None:
184        while get_remaining():
185            time.sleep(5)
186    else:
187        while get_remaining() and timeout >= 0:
188            time.sleep(5)
189            timeout -= 5
190    if timeout and timeout < 0:
191        raise (
192            TimeoutError(
193                "Please call the list_indexes API ({}) to confirm if index is deleted".format(
194                    "https://www.pinecone.io/docs/api/operation/list_indexes/"
195                )
196            )
197        )
198
199
200def list_indexes():
201    """Lists all indexes."""
202    api_instance = _get_api_instance()
203    response = api_instance.list_indexes()
204    return response
205
206
207def describe_index(name: str):
208    """Describes a Pinecone index.
209
210    :param name: the name of the index to describe.
211    :return: Returns an `IndexDescription` object
212    """
213    api_instance = _get_api_instance()
214    response = api_instance.describe_index(name)
215    db = response["database"]
216    ready = response["status"]["ready"]
217    state = response["status"]["state"]
218    return IndexDescription(
219        name=db["name"],
220        metric=db["metric"],
221        replicas=db["replicas"],
222        dimension=db["dimension"],
223        shards=db["shards"],
224        pods=db.get("pods", db["shards"] * db["replicas"]),
225        pod_type=db.get("pod_type", "p1"),
226        status={"ready": ready, "state": state},
227        metadata_config=db.get("metadata_config"),
228        source_collection=db.get("source_collection", ""),
229    )
230
231
232def scale_index(name: str, replicas: int):
233    """Increases number of replicas for the index.
234
235    :param name: the name of the Index
236    :type name: str
237    :param replicas: the number of replicas in the index now, lowest value is 0.
238    :type replicas: int
239    """
240    api_instance = _get_api_instance()
241    api_instance.configure_index(name, patch_request=PatchRequest(replicas=replicas, pod_type=""))
242
243
244def create_collection(name: str, source: str):
245    """Create a collection
246    :param name: Name of the collection
247    :param source: Name of the source index
248    """
249    api_instance = _get_api_instance()
250    api_instance.create_collection(create_collection_request=CreateCollectionRequest(name=name, source=source))
251
252
253def list_collections():
254    """List all collections"""
255    api_instance = _get_api_instance()
256    response = api_instance.list_collections()
257    return response
258
259
260def delete_collection(name: str):
261    """Deletes a collection.
262    :param: name: The name of the collection
263    """
264    api_instance = _get_api_instance()
265    api_instance.delete_collection(name)
266
267
268def describe_collection(name: str):
269    """Describes a collection.
270    :param: The name of the collection
271    :return: Description of the collection
272    """
273    api_instance = _get_api_instance()
274    response = api_instance.describe_collection(name).to_dict()
275    response_object = CollectionDescription(response.keys(), response.values())
276    return response_object
277
278
279def configure_index(name: str, replicas: Optional[int] = None, pod_type: Optional[str] = ""):
280    """Changes current configuration of the index.
281    :param: name: the name of the Index
282    :param: replicas: the desired number of replicas, lowest value is 0.
283    :param: pod_type: the new pod_type for the index.
284    """
285    api_instance = _get_api_instance()
286    config_args = {}
287    if pod_type != "":
288        config_args.update(pod_type=pod_type)
289    if replicas:
290        config_args.update(replicas=replicas)
291    patch_request = PatchRequest(**config_args)
292    api_instance.configure_index(name, patch_request=patch_request)