qwak-core 0.6.7__py3-none-any.whl → 0.7.0__py3-none-any.whl

qwak/vector_store/utils/upsert_utils.py

@@ -1,217 +0,0 @@
-import os
-from functools import partial
-from math import ceil
-from multiprocessing import Pool, set_start_method
-from typing import Dict, Iterable, List, Optional, Tuple, TypeVar, Union
-
-from _qwak_proto.qwak.vectors.v1.vector_pb2 import (
-    DoubleVector,
-    Property,
-    StoredVector,
-    VectorIdentifier,
-)
-from dependency_injector.wiring import Provide, inject
-from qwak.clients.vector_store.serving_client import VectorServingClient
-from qwak.exceptions import QwakException
-from qwak.inner.di_configuration import QwakContainer
-from qwak.vector_store.inference_client import VectorStoreInferenceClient
-from typeguard import typechecked
-
-_NaturalInput = TypeVar("T")
-_NaturalInputs = List[_NaturalInput]
-_Vector = List[float]
-_Properties = Dict[str, Union[str, int, bool, float]]
-
-_type_to_proto_property_mapping: Dict = {
-    str: "string_val",
-    bool: "bool_val",
-    int: "int_val",
-    float: "double_val",
-}
-
-
-def _build_property(key: str, value: Union[str, int, bool, float]):
-    type_val = _type_to_proto_property_mapping.get(type(value), None)
-    if not type_val:
-        raise QwakException(
-            f"Cannot upsert vector with property value type {type(value)}. "
-            f"Supported types are: {list(_type_to_proto_property_mapping.keys())}"
-        )
-
-    property_args = {"name": key, type_val: value}
-    return Property(**property_args)
-
-
-def _rewire_qwak_container(config):
-    # re-creating the container using the config from the original container.
-    # note that this runs in a fresh interpreter - at that point there's a running
-    # container because of the imports, but it won't necessarily have the same config
-    # as the container in the parent process that spawned this one.
-    # rewiring only the vector store + ecosystem and authentication - if using stuff from
-    # other modules pls feel free to add it here.
-    new_container = QwakContainer(config=config)
-    from qwak.clients import vector_store
-    from qwak.clients.administration import authentication, eco_system
-
-    new_container.wire(
-        packages=[
-            authentication,
-            eco_system,
-            vector_store,
-        ]
-    )
-
-
-@typechecked
-def _divide_chunks(lst: List, chunk_size: int):
-    if chunk_size <= 0:
-        raise QwakException("Chunk size must be a positive integer")
-
-    num_items: int = len(lst)
-    num_chunks: int = ceil(num_items / chunk_size)
-    for i in range(num_chunks):
-        yield lst[i * chunk_size : (i + 1) * chunk_size]
-
-
-def _get_vector_identifier(t: Union[str, Tuple[str, str]]):
-    if type(t) is str:
-        return VectorIdentifier(vector_id=t)
-    return VectorIdentifier(vector_id=t[0], tenant_id=t[1])
-
-
-def _upsert_vector_block(
-    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _Vector, _Properties]],
-    chunk_size: int,
-    collection_name: str,
-    edge_services_url: str,
-) -> None:
-    vector_serving_client: VectorServingClient = VectorServingClient(
-        edge_services_url=edge_services_url
-    )
-    for chunk in _divide_chunks(vector_tuples, chunk_size):
-        # chunk is a list of (id, vector, properties) tuples
-        vector_serving_client.upsert_vectors(
-            collection_name=collection_name,
-            vectors=[
-                StoredVector(
-                    vector_identifier=_get_vector_identifier(tpl[0]),
-                    vector=DoubleVector(element=tpl[1]),
-                    properties=[
-                        _build_property(key, value) for (key, value) in tpl[2].items()
-                    ],
-                )
-                for tpl in chunk
-            ],
-        )
-
-
-def _upsert_natural_input_block(
-    vector_tuples: List[Tuple[VectorIdentifier, _NaturalInput, _Properties]],
-    chunk_size: int,
-    vectorizer_name: str,
-    collection_name: str,
-    edge_services_url: str,
-) -> None:
-    vector_serving_client: VectorServingClient = VectorServingClient(
-        edge_services_url=edge_services_url
-    )
-    inference_client: VectorStoreInferenceClient = VectorStoreInferenceClient(
-        model_id=vectorizer_name
-    )
-    for chunk in _divide_chunks(vector_tuples, chunk_size):
-        # chunk is a list of (id, _NaturalInput, properties) tuples
-        vector_serving_client.upsert_vectors(
-            collection_name=collection_name,
-            vectors=[
-                StoredVector(
-                    vector_identifier=_get_vector_identifier(tpl[0]),
-                    vector=DoubleVector(
-                        element=inference_client.get_embedding(natural_input=tpl[1])
-                    ),
-                    properties=[
-                        _build_property(key, value) for (key, value) in tpl[2].items()
-                    ],
-                )
-                for tpl in chunk
-            ],
-        )
-
-
-@inject
-def _upsert_natural_input(
-    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _NaturalInput, _Properties]],
-    chunk_size: int,
-    vectorizer_name: str,
-    collection_name: str,
-    edge_services_url: str,
-    multiproc: bool = False,
-    max_processes: Optional[int] = None,
-    config=Provide[QwakContainer.config],
-):
-    if not multiproc:
-        _upsert_natural_input_block(
-            vector_tuples=vector_tuples,
-            chunk_size=chunk_size,
-            vectorizer_name=vectorizer_name,
-            collection_name=collection_name,
-            edge_services_url=edge_services_url,
-        )
-    else:
-        if max_processes is None:
-            max_processes = os.cpu_count()
-        effective_block_size: int = ceil(len(vector_tuples) / (max_processes * 4))
-
-        f = partial(
-            _upsert_natural_input_block,
-            chunk_size=chunk_size,
-            vectorizer_name=vectorizer_name,
-            collection_name=collection_name,
-            edge_services_url=edge_services_url,
-        )
-
-        blocks: Iterable[List[Tuple[str, _NaturalInput, _Properties]]] = _divide_chunks(
-            vector_tuples, effective_block_size
-        )
-        initializer = partial(_rewire_qwak_container, config=config)
-        set_start_method("spawn", force=True)
-
-        with Pool(processes=max_processes, initializer=initializer) as p:
-            p.map(f, blocks)
-
-
-@inject
-def _upsert_vectors(
-    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _Vector, _Properties]],
-    chunk_size: int,
-    collection_name: str,
-    edge_services_url: str,
-    multiproc: bool = False,
-    max_processes: Optional[int] = None,
-    config=Provide[QwakContainer.config],
-):
-    if not multiproc:
-        _upsert_vector_block(
-            vector_tuples=vector_tuples,
-            chunk_size=chunk_size,
-            collection_name=collection_name,
-            edge_services_url=edge_services_url,
-        )
-    else:
-        if max_processes is None:
-            max_processes = os.cpu_count()
-        effective_block_size: int = ceil(len(vector_tuples) / (max_processes * 4))
-
-        f = partial(
-            _upsert_vector_block,
-            chunk_size=chunk_size,
-            collection_name=collection_name,
-            edge_services_url=edge_services_url,
-        )
-        blocks: Iterable[List[Tuple[str, _Vector, _Properties]]] = list(
-            _divide_chunks(vector_tuples, effective_block_size)
-        )
-
-        set_start_method("spawn", force=True)
-        initializer = partial(_rewire_qwak_container, config=config)
-        with Pool(processes=max_processes, initializer=initializer) as p:
-            p.map(f, blocks)

qwak_services_mock/mocks/vector_serving_api.py

@@ -1,154 +0,0 @@
-from dataclasses import dataclass
-from typing import Dict, List
-
-import grpc
-from _qwak_proto.qwak.vectors.v1.vector_pb2 import (
-    DoubleVector,
-    Property,
-    SearchResult,
-    StoredVector,
-    VectorIdentifier,
-)
-from _qwak_proto.qwak.vectors.v1.vector_service_pb2 import (
-    DeleteVectorsResponse,
-    FetchVectorRequest,
-    FetchVectorResponse,
-    SearchSimilarVectorsRequest,
-    SearchSimilarVectorsResponse,
-    UpsertVectorsRequest,
-    UpsertVectorsResponse,
-)
-from _qwak_proto.qwak.vectors.v1.vector_service_pb2_grpc import VectorServiceServicer
-from numpy import dot
-from numpy.linalg import norm
-
-
-@dataclass
-class VectorObject:
-    id: str
-    vector: List[float]
-    property: List[Property]
-
-
-class VectorServingServiceMock(VectorServiceServicer):
-    def __init__(self):
-        self._vector_collections: Dict[str, Dict] = dict()
-
-    def reset_vector_store(self) -> None:
-        self._vector_collections.clear()
-
-    def get_num_of_vectors(self, collection_name: str) -> int:
-        if collection_name not in self._vector_collections:
-            return 0
-
-        return len(self._vector_collections[collection_name].values())
-
-    def create_collection(self, collection_name: str) -> None:
-        self._vector_collections[collection_name] = dict()
-
-    def get_vector_by_ids(
-        self, collection_name: str, vector_ids: List[str]
-    ) -> List[dict]:
-        if collection_name not in self._vector_collections:
-            raise ValueError(f"Collection named {collection_name} doesn't exist")
-
-        collection = self._vector_collections[collection_name]
-        return [collection.get(_id) for _id in vector_ids if _id in collection]
-
-    def SearchSimilarVectors(self, request: SearchSimilarVectorsRequest, context):
-        if request.collection_name not in self._vector_collections:
-            context.set_details(f"Collection named {request.collection} doesn't exist'")
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            return
-
-        reference_vector = list(request.reference_vector.element)
-        collection_store = self._vector_collections.get(request.collection_name, dict())
-
-        # a naive impl of a "vector similarity" search - compute pairwise cosine distance on
-        # the entire set and return top results
-        result_set = sorted(
-            collection_store.values(),
-            key=lambda b: _cos_sim(reference_vector, b["vector"]),
-        )[-request.max_results :]
-
-        return SearchSimilarVectorsResponse(
-            search_results=[
-                SearchResult(
-                    id=result["id"] if request.include_id else None,
-                    properties=[
-                        p
-                        for p in result_set[0]["properties"]
-                        if p.name in request.properties
-                    ],
-                    vector=DoubleVector(element=result["vector"])
-                    if request.include_vector
-                    else None,
-                    distance=_cos_sim(reference_vector, result["vector"])
-                    if request.include_distance
-                    else None,
-                )
-                for result in result_set
-            ]
-        )
-
-    def UpsertVectors(self, request: UpsertVectorsRequest, context):
-        collection_store = self._vector_collections.get(request.collection_name, dict())
-        for stored_vector in request.vector:
-            id = stored_vector.vector_identifier.vector_id
-            collection_store[id] = {
-                "id": id,
-                "vector": list(stored_vector.vector.element),
-                "properties": stored_vector.properties,
-            }
-
-        self._vector_collections[request.collection_name] = collection_store
-        return UpsertVectorsResponse()
-
-    def DeleteVectors(self, request, context):
-        if request.collection_name not in self._vector_collections:
-            context.set_details(
-                f"Collection named {request.collection_name} doesn't exist'"
-            )
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            return
-
-        collection_store = self._vector_collections[request.collection_name]
-        ids_in_collection = [
-            vector_identifier.vector_id
-            for vector_identifier in request.vector_identifiers
-            if vector_identifier.vector_id in collection_store
-        ]
-        for id in ids_in_collection:
-            collection_store.pop(id)
-
-        self._vector_collections[request.collection_name] = collection_store
-        return DeleteVectorsResponse(num_vectors_deleted=len(ids_in_collection))
-
-    def FetchVector(self, request, context):
-        if request.collection_name not in self._vector_collections:
-            context.set_details(
-                f"Collection named {request.collection_name} doesn't exist'"
-            )
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            return
-
-        collection_store = self._vector_collections[request.collection_name]
-        vector_identifier = request.vector_identifier
-        vector_id = vector_identifier.vector_id
-        if vector_id not in collection_store:
-            context.set_details(f"Vector with id {vector_id} doesn't exist'")
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            return
-
-        stored_vector_dict = collection_store[vector_id]
-        return FetchVectorResponse(
-            vector=StoredVector(
-                vector_identifier=request.vector_identifier,
-                vector=DoubleVector(element=stored_vector_dict["vector"]),
-                properties=stored_vector_dict["properties"],
-            )
-        )
-
-
-def _cos_sim(a, b):
-    return dot(a, b) / (norm(a) * norm(b))

qwak_services_mock/mocks/vectors_management_api.py

@@ -1,96 +0,0 @@
-import uuid
-from datetime import datetime
-
-import grpc
-from _qwak_proto.qwak.vectors.v1.collection.collection_pb2 import (
-    COLLECTION_STATUS_CREATE_REQUESTED,
-    QwakMetadata,
-    VectorCollection,
-    VectorCollectionDefinition,
-)
-from _qwak_proto.qwak.vectors.v1.collection.collection_service_pb2 import (
-    CreateCollectionResponse,
-    DeleteCollectionByIdResponse,
-    DeleteCollectionByNameResponse,
-    GetCollectionByIdResponse,
-    GetCollectionByNameResponse,
-    ListCollectionsResponse,
-)
-from _qwak_proto.qwak.vectors.v1.collection.collection_service_pb2_grpc import (
-    VectorCollectionServiceServicer,
-)
-from google.protobuf.timestamp_pb2 import Timestamp
-
-
-class VectorCollectionManagementServiceMock(VectorCollectionServiceServicer):
-    def __init__(self):
-        self._collections_spec_by_ids = {}
-        self._collections_spec_by_name = {}
-
-    def reset_collections(self):
-        self._collections_spec_by_ids.clear()
-        self._collections_spec_by_name.clear()
-
-    def CreateCollection(self, request, context):
-        timestamp = Timestamp()
-        timestamp.FromDatetime(datetime.now())
-
-        collection_id = str(uuid.uuid4())
-        vector_collection = VectorCollection(
-            metadata=QwakMetadata(
-                created_at=timestamp,
-                created_by="it@qwak.com",
-                last_modified_at=timestamp,
-                last_modified_by="it@qwak.com",
-            ),
-            definition=VectorCollectionDefinition(
-                id=collection_id,
-                collection_spec=request.collection_spec,
-            ),
-            status=COLLECTION_STATUS_CREATE_REQUESTED,
-        )
-
-        self._collections_spec_by_ids[collection_id] = vector_collection
-        self._collections_spec_by_name[request.collection_spec.name] = vector_collection
-        return CreateCollectionResponse(vector_collection=vector_collection)
-
-    def GetCollectionById(self, request, context):
-        if request.id in self._collections_spec_by_ids:
-            return GetCollectionByIdResponse(
-                vector_collection=self._collections_spec_by_ids[request.id]
-            )
-
-        context.set_details(f"Collection ID {request.id} doesn't exist'")
-        context.set_code(grpc.StatusCode.NOT_FOUND)
-
-    def GetCollectionByName(self, request, context):
-        if request.name in self._collections_spec_by_name:
-            return GetCollectionByNameResponse(
-                vector_collection=self._collections_spec_by_name[request.name]
-            )
-
-        context.set_details(f"Collection name {request.name} doesn't exist'")
-        context.set_code(grpc.StatusCode.NOT_FOUND)
-
-    def DeleteCollectionById(self, request, context):
-        if request.id in self._collections_spec_by_ids:
-            self._collections_spec_by_ids.pop(request.id)
-            return DeleteCollectionByIdResponse()
-
-        context.set_details(f"Collection ID {request.id} doesn't exist'")
-        context.set_code(grpc.StatusCode.NOT_FOUND)
-
-    def DeleteCollectionByName(self, request, context):
-        if request.name in self._collections_spec_by_name:
-            self._collections_spec_by_name.pop(request.name)
-            return DeleteCollectionByNameResponse()
-
-        context.set_details(f"Collection name {request.name} doesn't exist'")
-        context.set_code(grpc.StatusCode.NOT_FOUND)
-
-    def ListCollections(self, request, context):
-        return ListCollectionsResponse(
-            vector_collections=[
-                collection for collection in self._collections_spec_by_ids.values()
-            ]
-        )