qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak/vector_store/collection.py

@@ -1,16 +1,22 @@
+import re
 from dataclasses import dataclass
 from typing import Dict, List, Optional, TypeVar, Union
 
 from _qwak_proto.qwak.vectors.v1.vector_pb2 import (
-    DoubleVector,
     Property,
     SearchResult as ProtoSearchResult,
-    StoredVector,
+    StoredVector as ProtoStoredVector,
+    VectorIdentifier as ProtoVectorIdentifier,
 )
 from qwak.clients.vector_store.serving_client import VectorServingClient
 from qwak.exceptions import QwakException
 from qwak.vector_store.filters import Filter
 from qwak.vector_store.inference_client import VectorStoreInferenceClient
+from qwak.vector_store.utils.upsert_utils import (
+    _divide_chunks,
+    _upsert_natural_input,
+    _upsert_vectors,
+)
 from typeguard import typechecked
 
 NaturalInput = TypeVar("T")
@@ -18,6 +24,8 @@ NaturalInputs = List[NaturalInput]
 Vector = List[float]
 Properties = Dict[str, Union[str, int, bool, float]]
 
+_TENANT_ID_PATTERN = r"^[a-zA-Z0-9_-]{4,64}$"
+
 
 @dataclass
 class SearchResult:
@@ -57,6 +65,7 @@ class Collection:
     dimension: int
     description: Optional[str]
     vectorizer: Optional[str]
+    multi_tenant: bool
 
     _vector_serving_client: VectorServingClient
     _type_to_proto_property_mapping: Dict[str, TypeVar] = {
@@ -79,6 +88,7 @@ class Collection:
         vector_serving_client: VectorServingClient,
         description: Optional[str] = None,
         vectorizer: Optional[str] = None,
+        muli_tenant: bool = False,
     ):
         """
         Initializes a `Collection` client object to interact with Qwak's vector serving service. Should not be created
@@ -92,6 +102,7 @@ class Collection:
         self.vectorizer = vectorizer
         self._vector_serving_client = vector_serving_client
         self._realtime_inference_client = None
+        self.multi_tenant = muli_tenant
 
         if vectorizer:
             self._realtime_inference_client = VectorStoreInferenceClient(
@@ -109,6 +120,7 @@ class Collection:
         include_vector: bool = False,
         include_distance: bool = False,
         filter: Optional[Filter] = None,
+        tenant_id: Optional[str] = None,
     ) -> List[SearchResult]:
         """
         Searches for vectors in the collection that are most similar to a given query vector.
@@ -126,6 +138,7 @@ class Collection:
             include_vector (list): Whether to include the vector values themselves in the result set
             include_distance (list): Whether to include the distance calculations to the result set
             filter (Filter): Pre-filtering search results
+            tenant_id (str): tenant ID, passed if and only if the collection has multi tenancy enabled
 
         Returns:
             List[SearchResult]: A list of SearchResult object, which is used as a container for the search results
@@ -133,10 +146,11 @@ class Collection:
         Raises:
             QwakException: If you don't provide either vectors or natural_inputs
             QwakException: If you provide both vectors and natural_inputs
+            QwakException: If the tenant provided mismatches the configuration
         """
-        if bool(vector) and bool(natural_input):
+        if not (bool(vector) ^ bool(natural_input)):
             raise QwakException(
-                "Only one of `vectors` or `natural` inputs should be defined and not empty."
+                "Exactly one of {'vectors', 'natural_input'} should be passed"
             )
 
         if natural_input:
@@ -144,6 +158,7 @@ class Collection:
                 natural_input=natural_input
             )
         proto_filter = filter._to_proto() if filter else None
+        self._validate_tenant(tenant_id)
 
         return [
             self._to_search_result(
@@ -161,6 +176,7 @@ class Collection:
                 include_vector=include_vector,
                 include_distance=include_distance,
                 filters=proto_filter,
+                tenant_id=tenant_id,
             )
         ]
 
@@ -171,6 +187,11 @@ class Collection:
         properties: List[Properties],
         vectors: Optional[List[Vector]] = None,
         natural_inputs: Optional[NaturalInputs] = None,
+        batch_size: int = 1000,
+        multiproc: bool = False,
+        max_processes: Optional[int] = None,
+        *,
+        tenant_ids: Optional[List[str]] = None,
     ) -> None:
         """
         Inserts new vectors into the collection or updates existing vectors. Notice that this method will overwrite
@@ -182,76 +203,135 @@ class Collection:
             natural_inputs (list): Natural inputs (text, image) which should be embedded by the collection and added
               to the store. This attribute or `vectors` must be set
             properties (dict): A dictionary of properties to attach with the vectors
+            batch_size(int): maximum batch size when upserting against the backend Vector Store, defaults to 1000
+            multiproc (bool): whether to use multiprocessing, defaults to False
+            max_processes (Optional[int]): max number of processes if multiproc is selected, multiproc must be enabled
+            tenant_ids (List[str]): tenant ids, should be specified if and only if the collection has multi tenancy enabled.
 
         Raises:
             QwakException: If you don't provide either vectors or natural_inputs
             QwakException: If you provide both vectors and natural_inputs
         """
 
-        if bool(vectors) and bool(natural_inputs):
+        if not (bool(vectors) ^ bool(natural_inputs)):
             raise QwakException(
                 "`vectors` or `natural` inputs should be defined and not empty. But not both"
             )
 
-        if natural_inputs:
-            vectors = self._transform_natural_input_list_to_vectors(
-                natural_inputs=natural_inputs
-            )
+        if max_processes is not None and not multiproc:
+            raise QwakException("Can not set max_processes if multiproc is not enabled")
 
-        if (len(vectors) != len(ids)) or (len(properties) != len(ids)):
+        self._validate_tenant_ids(
+            vector_ids=vectors, tenant_ids=tenant_ids, verb="upserting"
+        )
+        id_tpls = zip(ids, tenant_ids) if self.multi_tenant else ids
+
+        if (len(vectors or natural_inputs) != len(ids)) or (
+            len(properties) != len(ids)
+        ):
             raise QwakException(
-                "Non matching lengths for input list (vectors / natural inputs), ID's, and properties list. "
+                "Non matching lengths for input list (vectors / natural inputs), IDs, and properties list. "
                 "Make sure all 3 fields are aligned in length"
             )
-
-        for zipped_vectors_chunks in self._divide_chunks(
-            list(zip(ids, vectors, properties)), 100
-        ):
-            self._vector_serving_client.upsert_vectors(
+        if bool(natural_inputs):
+            _upsert_natural_input(
+                vector_tuples=list(zip(id_tpls, natural_inputs, properties)),
+                chunk_size=batch_size,
+                vectorizer_name=self.vectorizer,
+                collection_name=self.name,
+                edge_services_url=self._vector_serving_client._edge_services_url,
+                multiproc=multiproc,
+                max_processes=max_processes,
+            )
+        else:
+            _upsert_vectors(
+                vector_tuples=list(zip(id_tpls, vectors, properties)),
+                chunk_size=batch_size,
                 collection_name=self.name,
-                vectors=[
-                    StoredVector(
-                        id=vector_object[0],
-                        vector=DoubleVector(element=vector_object[1]),
-                        properties=[
-                            self._build_property(key, value)
-                            for (key, value) in vector_object[2].items()
-                        ],
-                    )
-                    for vector_object in zipped_vectors_chunks
-                ],
+                edge_services_url=self._vector_serving_client._edge_services_url,
+                multiproc=multiproc,
+                max_processes=max_processes,
             )
 
     @typechecked
-    def delete(self, vector_ids: List[str]) -> int:
+    def delete(
+        self,
+        vector_ids: List[str],
+        *,
+        tenant_ids: Optional[List[str]] = None,
+        batch_size: int = 10000,
+    ) -> int:
         """
         Deletes vectors from the collection based on their IDs.
 
         Parameters:
             vector_ids (list): A list of vector IDs to delete.
+            batch_size (int): optional batch size, defaults to 10000
+            tenant_ids (list): tenant IDs (same length as vector_ids, used only when multi tenancy is enabled)
 
         Returns:
             int: Number of actual vectors deleted from the collection
         """
+        self._validate_tenant_ids(
+            vector_ids=vector_ids, tenant_ids=tenant_ids, verb="deleting"
+        )
+        vector_identifiers: List[ProtoVectorIdentifier] = self._extract_tenant_ids(
+            vector_ids, tenant_ids
+        )
+
         return sum(
             self._vector_serving_client.delete_vectors(
-                collection_name=self.name, vector_ids=ids_chunk
+                collection_name=self.name, vector_identifiers=ids_chunk
             )
-            for ids_chunk in self._divide_chunks(vector_ids, 100)
+            for ids_chunk in _divide_chunks(vector_identifiers, batch_size)
+        )
+
+    @typechecked
+    def fetch(self, vector_id: str, *, tenant_id: Optional[str] = None) -> SearchResult:
+        """
+        Fetches a vector from the collection based on its ID.
+
+        Parameters:
+            vector_id (str): The ID of the vector to fetch.
+            tenant_id (str, optional): Tenant id, passed if and only if multi tenancy is enabled
+
+        Returns:
+            SearchResult: A SearchResult object, which is used as a container for the search results
+        """
+        self._validate_tenant(tenant_id)
+        vector_identifier: ProtoVectorIdentifier
+        if tenant_id is not None:
+            vector_identifier = ProtoVectorIdentifier(
+                vector_id=vector_id, tenant_id=tenant_id
+            )
+        else:
+            vector_identifier = ProtoVectorIdentifier(vector_id=vector_id)
+
+        result = self._vector_serving_client.fetch_vector(
+            collection_name=self.name, vector_identifier=vector_identifier
+        )
+
+        return self._to_search_result(
+            result, include_id=True, include_distance=False, include_vector=True
         )
 
     def _to_search_result(
         self,
-        search_result: ProtoSearchResult,
+        search_result: Union[ProtoSearchResult, ProtoStoredVector],
         include_id: bool,
         include_vector: bool,
         include_distance: bool,
     ) -> SearchResult:
+        id = (
+            search_result.vector_identifier.vector_id
+            if type(search_result) is ProtoStoredVector
+            else search_result.id
+        )
         return SearchResult(
-            id=search_result.id if include_id else None,
-            vector=[e for e in search_result.vector.element]
-            if include_vector
-            else None,
+            id=id if include_id else None,
+            vector=(
+                [e for e in search_result.vector.element] if include_vector else None
+            ),
             distance=search_result.distance if include_distance else None,
             properties={
                 prop.name: self._extract_value_with_type(prop)
@@ -259,22 +339,6 @@ class Collection:
             },
         )
 
-    @staticmethod
-    def _divide_chunks(list_to_break, chunk_size):
-        for i in range(0, len(list_to_break), chunk_size):
-            yield list_to_break[i : i + chunk_size]
-
-    def _build_property(self, key: str, value: Union[str, int, bool, float]):
-        type_val = self._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(self._type_to_proto_property_mapping.keys())}"
-            )
-
-        property_args = {"name": key, type_val: value}
-        return Property(**property_args)
-
     def _extract_value_with_type(self, prop: Property):
         type_caster = self._proto_property_to_type_mapping.get(
             prop.WhichOneof("value_type"), None
@@ -294,28 +358,7 @@ class Collection:
             raise QwakException(
                 "Unable to search by natural input because the collection does not have a Vectorizer defined."
             )
-        feature_vector = [
-            {
-                "input": natural_input,
-            }
-        ]
-        try:
-            result_list = self._realtime_inference_client.predict(feature_vector)
-        except Exception as e:
-            raise QwakException(
-                f"Vectorizer {self.vectorizer} failed to transform input {feature_vector} to vectors. Error is: {str(e)}"
-            )
-        try:
-            vector = result_list[0]["embeddings"]
-        except Exception:
-            raise QwakException(
-                f"Vectorizer {self.vectorizer} must return a dataframe containing an `embeddings` column"
-            )
-        if not vector:
-            raise QwakException(
-                f"Vectorizer {self.vectorizer} did not return embeddings for the given natural input. Unable to continue with the query."
-            )
-        return vector
+        return self._realtime_inference_client.get_embedding(natural_input)
 
     def _transform_natural_input_list_to_vectors(
         self, natural_inputs: NaturalInputs
@@ -324,3 +367,60 @@ class Collection:
             self._transform_natural_input_to_vectors(natural_input=natural_input)
             for natural_input in natural_inputs
         ]
+
+    def _validate_tenant(self, tenant_id: Optional[str] = None):
+        if self.multi_tenant:
+            # we are multi tenant, assert a valid tenant is passed
+            if tenant_id is None:
+                raise QwakException(
+                    "Tenant ID must be passed when multi tenancy is enabled"
+                )
+
+            self._validate_tenant_id(tenant_id)
+        else:
+            if tenant_id is not None:
+                raise QwakException(
+                    f"Collection {self.name} is not multi tenant, can not specify tenant"
+                )
+
+    def _validate_tenant_id(self, tenant_id: str):
+        if not (bool(re.match(_TENANT_ID_PATTERN, tenant_id))):
+            raise QwakException(
+                f"Tenant ID {tenant_id} does not conform to {_TENANT_ID_PATTERN}"
+            )
+
+    def _validate_tenant_ids(
+        self, vector_ids: List[str], tenant_ids: Optional[List[str]], verb: str
+    ) -> None:
+        if self.multi_tenant:
+            if tenant_ids is None:
+                raise QwakException(
+                    f"Tenant IDs must be provided when {verb} against multitenant collections"
+                )
+            if len(tenant_ids) != len(vector_ids):
+                raise QwakException(
+                    f"Got different number of vector ids {len(vector_ids)} and tenant ids {len(tenant_ids)}"
+                )
+            for tenant_id in tenant_ids:
+                self._validate_tenant_id(tenant_id=tenant_id)
+        else:
+            if tenant_ids is not None:
+                raise QwakException(
+                    f"Collection {self.name} does not have multi tenancy enabled, do not pass tenant ids"
+                )
+
+    def _extract_tenant_ids(
+        self, vector_ids: List[str], tenant_ids: Optional[List[str]]
+    ) -> List[ProtoVectorIdentifier]:
+        vector_identifiers: List[ProtoVectorIdentifier]
+        if self.multi_tenant:
+            vector_identifiers = [
+                ProtoVectorIdentifier(vector_id=vector_id, tenant_id=tenant_id)
+                for vector_id, tenant_id in zip(vector_ids, tenant_ids)
+            ]
+        else:
+            vector_identifiers = [
+                ProtoVectorIdentifier(vector_id=vector_id) for vector_id in vector_ids
+            ]
+
+        return vector_identifiers

qwak/vector_store/inference_client.py

@@ -1,7 +1,7 @@
 import json
 
 from qwak.clients.administration.eco_system.client import EcosystemClient
-from qwak.exceptions import QwakHTTPException
+from qwak.exceptions import QwakException, QwakHTTPException
 from qwak.vector_store.rest_helpers import RestSession
 
 
@@ -70,6 +70,34 @@ class VectorStoreInferenceClient:
         except Exception as e:
             raise RuntimeError(f"Failed to make a prediction request. Error is: {e}")
 
+    def get_embedding(self, natural_input: str):
+        feature_vector = [
+            {
+                "input": natural_input,
+            }
+        ]
+        try:
+            result_list = self.predict(feature_vector)
+        except Exception as e:
+            raise QwakException(
+                f"Vectorizer {self.model_id} failed to transform input {feature_vector} to vectors. Error is: {str(e)}"
+            )
+        try:
+            vector = result_list[0]["embeddings"]
+        except Exception:
+            raise QwakException(
+                f"Vectorizer {self.model_id} must return a dataframe containing an 'embeddings' column"
+            )
+        if not vector:
+            raise QwakException(
+                f"Vectorizer {self.model_id} did not return embeddings for the given natural input. Unable to continue with the query."
+            )
+        return vector
+
 
 def _get_model_url(model_id: str, model_url_prefix: str) -> str:
-    return f"https://{model_url_prefix}/v1/{model_id}/predict"
+    scheme = "http" if model_url_prefix.startswith("localhost") else "https"
+    effective_model_id = model_id.replace("-", "_")
+    return (
+        f"{scheme}://{model_url_prefix}/v1/{effective_model_id}/predict"  # noqa: E231
+    )

qwak/vector_store/rest_helpers.py

@@ -4,21 +4,15 @@ import socket
 from datetime import datetime
 
 import requests
-from qwak.clients.administration import AuthenticationClient
-from qwak.inner.di_configuration.account import UserAccountConfiguration
+from qwak.inner.tool.auth import Auth0ClientBase
 
 
 def _get_authorization():
-    user_account_configuration = UserAccountConfiguration()
-    token = (
-        AuthenticationClient()
-        .authenticate(user_account_configuration.get_user_apikey())
-        .access_token
-    )
+    auth_client = Auth0ClientBase()
+    token = auth_client.get_token()
     token_split = token.split(".")
     decoded_token = json.loads(_base64url_decode(token_split[1]).decode("utf-8"))
     token_expiration = datetime.fromtimestamp(decoded_token["exp"])
-
     return f"Bearer {token}", token_expiration
 
 
@@ -30,15 +24,15 @@ def _base64url_decode(input):
     return base64.urlsafe_b64decode(input)
 
 
-class HTTPAdapterWithSocketOptions(requests.adapters.HTTPAdapter):
-    def __init__(self, *args, **kwargs):
-        self.socket_options = kwargs.pop("socket_options", None)
-        super(HTTPAdapterWithSocketOptions, self).__init__(*args, **kwargs)
+class SocketAdapter(requests.adapters.HTTPAdapter):
+    def __init__(self, *args, socket_params=None, **kwargs):
+        self.socket_params = socket_params
+        super().__init__(*args, **kwargs)
 
     def init_poolmanager(self, *args, **kwargs):
-        if self.socket_options is not None:
-            kwargs["socket_options"] = self.socket_options
-        super(HTTPAdapterWithSocketOptions, self).init_poolmanager(*args, **kwargs)
+        if self.socket_params:
+            kwargs["socket_options"] = self.socket_params
+        super().init_poolmanager(*args, **kwargs)
 
 
 # Configure only TCP attributes that are available in the OS
@@ -61,7 +55,7 @@ class RestSession(requests.Session):
             ("SOL_TCP", "TCP_KEEPCNT", 9),
         ]
         socket_options = validate_socket_config(socket_options)
-        adapter = HTTPAdapterWithSocketOptions(socket_options=socket_options)
+        adapter = SocketAdapter(socket_params=socket_options)
         self.mount("https://", adapter)
 
     def prepare_request(self, request):
@@ -74,7 +68,5 @@ class RestSession(requests.Session):
         return super().prepare_request(request)
 
     def prepare_request_token(self):
-        (
-            self.headers["Authorization"],
-            self.jwt_expiration,
-        ) = _get_authorization()
+        auth_token, self.jwt_expiration = _get_authorization()
+        self.headers["Authorization"] = auth_token

qwak/vector_store/utils/filter_utils.py

@@ -1,8 +1,8 @@
-from datetime import datetime, timezone
-from typing import Any, Optional
+from datetime import datetime
+from typing import Any
 
 from _qwak_proto.qwak.vectors.v1.filters_pb2 import AtomicLiteral as ProtoAtomicLiteral
-from google.protobuf.timestamp_pb2 import Timestamp as ProtoTimestamp
+from qwak.utils.datetime_utils import datetime_to_pts
 
 
 def transform(value: Any) -> ProtoAtomicLiteral:
@@ -16,19 +16,6 @@ def transform(value: Any) -> ProtoAtomicLiteral:
         return ProtoAtomicLiteral(double_literal=value)
     elif isinstance(value, datetime):
         # Assuming that timestamp is a datetime
-        return ProtoAtomicLiteral(timestamp_literal=_datetime_to_pts(value))
+        return ProtoAtomicLiteral(timestamp_literal=datetime_to_pts(value))
     else:
         raise ValueError(f"Unsupported data type: {type(value)}")
-
-
-def _datetime_to_pts(dtime: Optional[datetime]) -> Optional[ProtoTimestamp]:
-    """
-    converts a python datetime to Protobuf Timestamp
-    @param dtime: python datetime.datetime
-    @return: if the input is None, returns None. else converts to proto timestamp in utc timezone
-    """
-    if not dtime:
-        return None
-    res: ProtoTimestamp = ProtoTimestamp()
-    res.FromDatetime(dtime.replace(tzinfo=timezone.utc))
-    return res