pinecone.grpc.utils

  1import uuid
  2
  3from google.protobuf.struct_pb2 import Struct
  4
  5
  6def _generate_request_id() -> str:
  7    return str(uuid.uuid4())
  8
  9
 10from pinecone.core.openapi.data.models import (
 11    Vector as _Vector,
 12    Usage,
 13    ScoredVector,
 14    SparseValues,
 15    FetchResponse,
 16    QueryResponse,
 17    DescribeIndexStatsResponse,
 18    NamespaceSummary,
 19)
 20
 21from typing import Optional
 22
 23
 24def dict_to_proto_struct(d: Optional[dict]) -> "Struct":
 25    if not d:
 26        d = {}
 27    s = Struct()
 28    s.update(d)
 29    return s
 30
 31
 32def parse_sparse_values(sparse_values: dict):
 33    return (
 34        SparseValues(indices=sparse_values["indices"], values=sparse_values["values"])
 35        if sparse_values
 36        else SparseValues(indices=[], values=[])
 37    )
 38
 39
 40def parse_fetch_response(response: dict):
 41    vd = {}
 42    vectors = response.get("vectors", {})
 43    namespace = response.get("namespace", "")
 44
 45    for id, vec in vectors.items():
 46        vd[id] = _Vector(
 47            id=vec["id"],
 48            values=vec["values"],
 49            sparse_values=parse_sparse_values(vec.get("sparseValues")),
 50            metadata=vec.get("metadata", None),
 51            _check_type=False,
 52        )
 53
 54    return FetchResponse(
 55        vectors=vd,
 56        namespace=namespace,
 57        usage=parse_usage(response.get("usage", {})),
 58        _check_type=False,
 59    )
 60
 61
 62def parse_usage(usage: dict):
 63    return Usage(read_units=int(usage.get("readUnits", 0)))
 64
 65
 66def parse_query_response(response: dict, _check_type: bool = False):
 67    matches = []
 68    for item in response.get("matches", []):
 69        sc = ScoredVector(
 70            id=item["id"],
 71            score=item.get("score", 0.0),
 72            values=item.get("values", []),
 73            sparse_values=parse_sparse_values(item.get("sparseValues")),
 74            metadata=item.get("metadata", None),
 75            _check_type=_check_type,
 76        )
 77        matches.append(sc)
 78
 79    # Due to OpenAPI model classes / actual parsing cost, we want to avoid
 80    # creating empty `Usage` objects and then passing them into QueryResponse
 81    # when they are not actually present in the response from the server.
 82    args = {
 83        "namespace": response.get("namespace", ""),
 84        "matches": matches,
 85        "_check_type": _check_type,
 86    }
 87    usage = response.get("usage")
 88    if usage:
 89        args["usage"] = parse_usage(usage)
 90    return QueryResponse(**args)
 91
 92
 93def parse_stats_response(response: dict):
 94    fullness = response.get("indexFullness", 0.0)
 95    total_vector_count = response.get("totalVectorCount", 0)
 96    dimension = response.get("dimension", 0)
 97    summaries = response.get("namespaces", {})
 98    namespace_summaries = {}
 99    for key in summaries:
100        vc = summaries[key].get("vectorCount", 0)
101        namespace_summaries[key] = NamespaceSummary(vector_count=vc)
102    return DescribeIndexStatsResponse(
103        namespaces=namespace_summaries,
104        dimension=dimension,
105        index_fullness=fullness,
106        total_vector_count=total_vector_count,
107        _check_type=False,
108    )