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

qwak/clients/secret_service/client.py

@@ -7,7 +7,6 @@ from _qwak_proto.qwak.secret_service.secret_service_pb2 import (
     SetSecretRequest,
 )
 from _qwak_proto.qwak.secret_service.secret_service_pb2_grpc import SecretServiceStub
-from qwak.clients._inner.edge_communications import get_endpoint_url
 from qwak.clients.administration.eco_system.client import EcosystemClient
 from qwak.exceptions import QwakException
 from qwak.inner.tool.grpc.grpc_tools import create_grpc_channel
@@ -26,7 +25,9 @@ class SecretServiceClient:
         environment_id: Optional[str] = None,
     ):
         if not grpc_channel:
-            edge_services_url = get_endpoint_url(edge_services_url, environment_id)
+            edge_services_url = self.__get_endpoint_url(
+                edge_services_url, environment_id
+            )
             grpc_channel = create_grpc_channel(
                 url=edge_services_url,
                 enable_ssl=enable_ssl,
@@ -44,6 +45,10 @@ class SecretServiceClient:
         if endpoint_url is None:
             user_context = EcosystemClient().get_authenticated_user_context().user
             if environment_id is None:
+                environment_id = user_context.environment_details.id
+
+            # If the environment id is not found, try to use the default environment id
+            if environment_id not in user_context.account_details.environment_by_id:
                 environment_id = user_context.account_details.default_environment_id
 
             if environment_id not in user_context.account_details.environment_by_id:
@@ -63,7 +68,7 @@ class SecretServiceClient:
             return self._secret_service.GetSecret(request).value
         except grpc.RpcError as e:
             if e.code() == grpc.StatusCode.NOT_FOUND:
-                raise QwakException(f"Secret [{name}] not found")
+                raise QwakException(f"Secret [{name}] not found", status_code=e.code())
             else:
                 raise QwakException(
                     f"Failed to get secret, name: [{name}], error code is [{e.code()}], error message is [{e.details()}]"

qwak/clients/system_secret/system_secret_client.py

@@ -0,0 +1,30 @@
+from _qwak_proto.qwak.admiral.secret.v0.secret_pb2 import (
+    EnvironmentSecretIdentifier,
+    SystemSecretValue,
+)
+from _qwak_proto.qwak.admiral.secret.v0.system_secret_service_pb2 import (
+    GetSystemSecretRequest,
+    GetSystemSecretResponse,
+)
+from _qwak_proto.qwak.admiral.secret.v0.system_secret_service_pb2_grpc import (
+    SystemSecretServiceStub,
+)
+from dependency_injector.wiring import Provide, inject
+from qwak.inner.di_configuration import QwakContainer
+
+
+class SystemSecretClient:
+    @inject
+    def __init__(self, grpc_channel=Provide[QwakContainer.core_grpc_channel]):
+        self._grpc_client: SystemSecretServiceStub = SystemSecretServiceStub(
+            grpc_channel
+        )
+
+    def get_system_secret(self, name: str, env_id: str) -> SystemSecretValue:
+        response: GetSystemSecretResponse = self._grpc_client.GetSystemSecret(
+            GetSystemSecretRequest(
+                identifier=EnvironmentSecretIdentifier(name=name, environment_id=env_id)
+            )
+        )
+
+        return response.secret_definition.spec.value

qwak/clients/vector_store/management_client.py

@@ -37,6 +37,7 @@ class VectorManagementClient:
         description: str = None,
         metric: VectorCollectionMetric = VectorCollectionMetric.COLLECTION_METRIC_L2_SQUARED,
         vectorizer: Optional[str] = None,
+        multi_tenant: bool = False,
     ) -> VectorCollection:
         """
         Create a collection
@@ -52,6 +53,7 @@ class VectorManagementClient:
                         ),
                         metric=metric,
                         dimension=dimension,
+                        multi_tenancy_enabled=multi_tenant,
                     )
                 )
             ).vector_collection

qwak/clients/vector_store/serving_client.py

@@ -6,15 +6,17 @@ from _qwak_proto.qwak.vectors.v1.vector_pb2 import (
     DoubleVector,
     SearchResult,
     StoredVector,
+    VectorIdentifier,
 )
 from _qwak_proto.qwak.vectors.v1.vector_service_pb2 import (
     DeleteVectorsRequest,
+    FetchVectorRequest,
     SearchSimilarVectorsRequest,
     UpsertVectorsRequest,
 )
 from _qwak_proto.qwak.vectors.v1.vector_service_pb2_grpc import VectorServiceStub
 from qwak.clients._inner.edge_communications import get_endpoint_url
-from qwak.exceptions import QwakException
+from qwak.exceptions import QwakException, QwakNotFoundException
 from qwak.inner.tool.grpc.grpc_tools import create_grpc_channel
 from typeguard import typechecked
 
@@ -26,6 +28,7 @@ class VectorServingClient:
         environment_id: Optional[str] = None,
     ):
         edge_services_url = get_endpoint_url(edge_services_url, environment_id)
+        self._edge_services_url = edge_services_url
 
         grpc_channel = create_grpc_channel(
             url=edge_services_url,
@@ -51,31 +54,49 @@ class VectorServingClient:
         include_vector: bool = False,
         include_distance: bool = False,
         filters: Optional[ProtoFilter] = None,
+        tenant_id: Optional[str] = None,
     ) -> List[SearchResult]:
         """
         Search for similar vectors
         """
+        request: SearchSimilarVectorsRequest
+        if tenant_id is not None:
+            request = SearchSimilarVectorsRequest(
+                collection_name=collection_name,
+                reference_vector=DoubleVector(element=vector),
+                properties=properties,
+                max_results=top_results,
+                filter=filters,
+                include_id=include_id,
+                include_vector=include_vector,
+                include_distance=include_distance,
+                tenant_id=tenant_id,
+            )
+        else:
+            request = SearchSimilarVectorsRequest(
+                collection_name=collection_name,
+                reference_vector=DoubleVector(element=vector),
+                properties=properties,
+                max_results=top_results,
+                filter=filters,
+                include_id=include_id,
+                include_vector=include_vector,
+                include_distance=include_distance,
+            )
+
         try:
             return list(
                 self._vector_serving_service.SearchSimilarVectors(
-                    SearchSimilarVectorsRequest(
-                        collection_name=collection_name,
-                        reference_vector=DoubleVector(element=vector),
-                        properties=properties,
-                        max_results=top_results,
-                        filter=filters,
-                        include_id=include_id,
-                        include_vector=include_vector,
-                        include_distance=include_distance,
-                    )
+                    request
                 ).search_results
             )
 
         except grpc.RpcError as e:
             raise QwakException(
-                f"Failed to query collection '{collection_name}' for vector '{str(vector)}': {e.details()}"
+                f"Failed to query collection '{collection_name}' for vector '{str(vector)}': {repr(e)}"
             )
 
+    @typechecked
     def upsert_vectors(self, collection_name: str, vectors: List[StoredVector]) -> None:
         """
         Upsert vectors to a collection
@@ -89,22 +110,47 @@ class VectorServingClient:
             )
 
         except grpc.RpcError as e:
-            raise QwakException(
-                f"Failed to upsert vectors with id's '{str([v.id for v in vectors])}': {e.details()}"
-            )
+            raise QwakException(f"Failed to upsert, got: {repr(e)}")
 
-    def delete_vectors(self, collection_name: str, vector_ids: List[str]) -> int:
+    @typechecked
+    def delete_vectors(
+        self, collection_name: str, vector_identifiers: List[VectorIdentifier]
+    ) -> int:
         """
         Delete vectors from a collection
         """
         try:
             return self._vector_serving_service.DeleteVectors(
                 DeleteVectorsRequest(
-                    collection_name=collection_name, vector_id=vector_ids
+                    collection_name=collection_name,
+                    vector_identifiers=vector_identifiers,
                 )
             ).num_vectors_deleted
 
         except grpc.RpcError as e:
             raise QwakException(
-                f"Failed to delete vectors by id '{str(vector_ids)}': {e.details()}"
+                f"Failed to delete vectors by id '{str(vector_identifiers)}': {repr(e)}"
+            )
+
+    @typechecked
+    def fetch_vector(
+        self, collection_name: str, vector_identifier: VectorIdentifier
+    ) -> StoredVector:
+        """
+        Fetch a vector from a collection
+        """
+        try:
+            return self._vector_serving_service.FetchVector(
+                FetchVectorRequest(
+                    collection_name=collection_name, vector_identifier=vector_identifier
+                )
+            ).vector
+
+        except grpc.RpcError as e:
+            if e.args[0].code.name == "NOT_FOUND":
+                raise QwakNotFoundException(
+                    f"Could not find vector {repr(vector_identifier)} in collection {collection_name}"
+                )
+            raise QwakException(
+                f"Failed to get vector by ID '{vector_identifier}': {repr(e)}"
             )

qwak/clients/workspace_manager/client.py

@@ -22,6 +22,7 @@ from _qwak_proto.qwak.workspace.workspace_service_pb2 import (
     UndeployWorkspaceRequest,
     UpdateWorkspaceRequest,
     GetDefaultWorkspaceValuesRequest,
+    MarkBuildAsCopiedRequest,
 )
 from _qwak_proto.qwak.workspace.workspace_service_pb2_grpc import (
     WorkspaceManagementServiceStub,
@@ -206,3 +207,17 @@ class WorkspaceManagerClient:
             raise QwakException(
                 f"Failed to get default workspace values, code: {e.code().name} error is {e.details()}"
             )
+
+    def _mark_build_as_copied(self, workspace_id: str, build_id: str) -> None:
+        """
+        Mark build as copied
+        """
+        try:
+            self._workspace_manager.MarkBuildAsCopied(
+                MarkBuildAsCopiedRequest(
+                    workspace_id=workspace_id,
+                    build_id=build_id,
+                )
+            )
+        except grpc.RpcError as e:
+            raise QwakException(f"Failed to mark build as copied, got: {e.details()}")

qwak/exceptions/__init__.py

@@ -1,7 +1,14 @@
 from .quiet_error import QuietError
 from .qwak_exception import QwakException
+from .qwak_general_build_exception import QwakGeneralBuildException
+from .qwak_grpc_address_exception import QwakGrpcAddressException
 from .qwak_http_exception import QwakHTTPException
+from .qwak_inference_exception import QwakInferenceException
+from .qwak_load_configuration_exception import LoadConfigurationException
+from .qwak_load_model_failed_exception import QwakLoadModelFailedException
 from .qwak_login_exception import QwakLoginException
 from .qwak_mock_http_exception import MockHttpException
 from .qwak_model_initialization_exception import QwakModelInitializationException
 from .qwak_not_found_exception import QwakNotFoundException
+from .qwak_remote_build_failed import QwakRemoteBuildFailedException
+from .qwak_suggestion_exception import QwakSuggestionException

qwak/exceptions/qwak_decode_exception.py

@@ -0,0 +1,7 @@
+class QwakDecodeException(ValueError):
+    """
+    Failed to decode payload
+    """
+
+    def __init__(self, message):
+        self.message = message

qwak/exceptions/qwak_exception.py

@@ -2,5 +2,6 @@ from .quiet_error import QuietError
 
 
 class QwakException(QuietError):
-    def __init__(self, message):
+    def __init__(self, message, status_code=None):
         self.message = message
+        self.status_code = status_code

qwak/exceptions/qwak_external_exception.py

@@ -0,0 +1,11 @@
+from qwak.exceptions import QuietError
+
+
+class QwakExternalException(QuietError):
+    """
+    An external system to Qwak, had an internal error
+    """
+
+    def __init__(self, message):
+        super().__init__(message)
+        self.message = message

qwak/exceptions/qwak_general_build_exception.py

@@ -0,0 +1,13 @@
+class QwakGeneralBuildException(Exception):
+    def __init__(self, message: str, src_exception: Exception = None):
+        self._message = message
+        self._exception_msg = str(src_exception) if src_exception else ""
+
+    @property
+    def message(self) -> str:
+        msg = f"""Message: {self._message}
+Exception message: {self._exception_msg}"""
+        return msg
+
+    def __str__(self):
+        return self.message

qwak/exceptions/qwak_grpc_address_exception.py

@@ -0,0 +1,9 @@
+from typing import Union
+from urllib.parse import ParseResult
+
+from qwak.exceptions import QwakException
+
+
+class QwakGrpcAddressException(QwakException):
+    def __init__(self, details: str, grpc_address: Union[str, ParseResult]):
+        self.message = f"Not a valid gRPC address: '{grpc_address}'. Details: {details}"

qwak/exceptions/qwak_load_configuration_exception.py

@@ -0,0 +1,2 @@
+class LoadConfigurationException(Exception):
+    pass

qwak/exceptions/qwak_remote_build_failed.py

@@ -0,0 +1,5 @@
+from .qwak_suggestion_exception import QwakSuggestionException
+
+
+class QwakRemoteBuildFailedException(QwakSuggestionException):
+    pass

qwak/exceptions/qwak_suggestion_exception.py

@@ -0,0 +1,27 @@
+from typing import Optional
+
+
+class QwakSuggestionException(Exception):
+    def __init__(
+        self,
+        message: str,
+        src_exception: Optional[Exception] = None,
+        suggestion: Optional[str] = "",
+    ):
+        self._message = message
+        self._src_exception = str(src_exception) if src_exception else ""
+        self._suggestion = suggestion
+
+    @property
+    def message(self) -> str:
+        if self._src_exception:
+            msg = f"""Message: {self._message}.
+Exception message: {self._src_exception}.
+Recommendation: {self._suggestion}."""
+        else:
+            msg = f"""Message: {self._message}.
+Recommendation: {self._suggestion}."""
+        return msg
+
+    def __str__(self):
+        return self.message

qwak/feature_store/_common/artifact_utils.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
+
+from qwak.clients.feature_store.management_client import FeatureRegistryClient
+from qwak.feature_store._common.packaging import upload_artifact
+
+if TYPE_CHECKING:
+    from qwak.feature_store.data_sources.streaming.kafka.deserialization import (
+        Deserializer,
+    )
+    from qwak.feature_store.feature_sets.transformations.transformations import (
+        BaseTransformation,
+    )
+
+
+@dataclass
+class ArtifactSpec:
+    """
+    Dataclass for holding artifacts which will be uploaded
+    """
+
+    artifact_name: str
+    root_module_path: Path
+    artifact_object: Union["BaseTransformation", "Deserializer"]
+    callables: List[Callable] = field(default_factory=list)
+    suffix: str = ""
+
+
+class ArtifactsUploader:
+    @staticmethod
+    def upload(artifact: ArtifactSpec) -> str:
+        return upload_artifact(
+            feature_store_object_name=artifact.artifact_name,
+            feature_store_object_name_suffix=artifact.suffix,
+            functions=artifact.callables,
+            feature_module_dir=artifact.root_module_path,
+            features_manager_client=FeatureRegistryClient(),
+            artifact_object=artifact.artifact_object,
+        )
+
+    @staticmethod
+    def get_artifact_spec(
+        transformation: "BaseTransformation",
+        featureset_name: str,
+        __instance_module_path__: str,
+    ) -> Optional[ArtifactSpec]:
+        transformation_functions: Optional[List[Callable[..., Any]]] = (
+            transformation.get_functions()
+        )
+        if not transformation_functions:
+            return None
+
+        return ArtifactSpec(
+            artifact_name=featureset_name,
+            root_module_path=Path(__instance_module_path__).parent,
+            artifact_object=transformation,
+            callables=transformation_functions,
+        )

qwak/feature_store/_common/feature_set_utils.py

@@ -1,21 +1,31 @@
+import time
 from dataclasses import dataclass
-from typing import Optional
+from typing import Dict, Optional
 
-from _qwak_proto.qwak.builds.builds_pb2 import (
+from _qwak_proto.qwak.build.v1.build_pb2 import (
     BatchFeature as ProtoBatchFeature,
     BatchFeatureV1 as ProtoBatchFeatureV1,
     Feature as ProtoFeature,
+    StreamingAggregationFeature as ProtoStreamingAggregationFeature,
+    StreamingFeature as ProtoStreamingFeature,
+    StreamingFeatureV1 as ProtoStreamingFeatureV1,
+)
+from _qwak_proto.qwak.feature_store.entities.entity_pb2 import (
+    EntitySpec as ProtoEntitySpec,
 )
-from _qwak_proto.qwak.feature_store.entities.entity_pb2 import EntitySpec
 from _qwak_proto.qwak.feature_store.features.feature_set_types_pb2 import (
-    BatchFeatureSet,
-    BatchFeatureSetV1,
-    FeatureSetType,
+    BatchFeatureSet as ProtoBatchFeatureSet,
+    BatchFeatureSetV1 as ProtoBatchFeatureSetV1,
+    FeatureSetType as ProtoFeatureSetType,
+    StreamingAggregationFeatureSet as ProtoStreamingAggregationFeatureSet,
+    StreamingFeatureSetV1 as ProtoStreamingFeatureSetV1,
 )
 from qwak.clients.feature_store import FeatureRegistryClient
 from qwak.exceptions import QwakException
 from qwak.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
 
+_INTERNAL_KEY_PREFIX: str = "_qwak_key_"
+
 
 @dataclass(unsafe_hash=True)
 class BatchFeature(BaseFeature):
@@ -41,10 +51,46 @@ class BatchFeatureV1(BaseFeature):
         )
 
 
+@dataclass(unsafe_hash=True)
+class StreamingFeature(BaseFeature):
+    entity: Entity
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_feature=ProtoStreamingFeature(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class StreamingFeatureV1(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_feature_v1=ProtoStreamingFeatureV1(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class StreamingAggregationFeature(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_aggregation_feature=ProtoStreamingAggregationFeature(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
 @dataclass()
 class FeatureSetInfo:
-    entity_spec: EntitySpec
-    feature_set_type: FeatureSetType
+    entity_spec: ProtoEntitySpec
+    feature_set_type: ProtoFeatureSetType
     feature_version: int
 
 
@@ -63,9 +109,9 @@ def get_feature_set_info(
         feature_set_name
     )
     if not feature_set_response:
-        raise QwakException(f"Feature set: {feature_set_name} does not exist")
+        raise QwakException(f"Featureset '{feature_set_name}' does not exist")
 
-    featureset_type: FeatureSetType = get_feature_type(
+    featureset_type: ProtoFeatureSetType = get_feature_type(
         feature_set_response.feature_set.feature_set_definition.feature_set_spec.feature_set_type
     )
     featureset_version = _get_featureset_version(featureset_type)
@@ -77,11 +123,11 @@ def get_feature_set_info(
     )
 
 
-def get_feature_type(feature_set_type: FeatureSetType):
+def get_feature_type(feature_set_type: ProtoFeatureSetType):
     return getattr(feature_set_type, feature_set_type.WhichOneof("set_type"))
 
 
-def _get_featureset_version(feature_set_type: FeatureSetType) -> int:
+def _get_featureset_version(feature_set_type: ProtoFeatureSetType) -> int:
     """
     Get Featureset version. Return 0 if none version is set
     """
@@ -91,7 +137,9 @@ def _get_featureset_version(feature_set_type: FeatureSetType) -> int:
 
 
 def get_typed_feature(
-    feature: FeatureStoreInput, feature_type: FeatureSetType, featureset_version: int
+    feature: FeatureStoreInput,
+    feature_type: ProtoFeatureSetType,
+    featureset_version: int,
 ) -> BaseFeature:
     """
     convert InputFeature to the relevant type
@@ -103,24 +151,99 @@ def get_typed_feature(
     Return:
         BaseFeature with the correct type
     """
-    if isinstance(feature_type, BatchFeatureSet):
+    if isinstance(feature_type, ProtoBatchFeatureSet):
         return BatchFeature(name=feature.name, entity=feature.entity)
-    elif isinstance(feature_type, BatchFeatureSetV1) and featureset_version == 0:
+    elif isinstance(feature_type, ProtoBatchFeatureSetV1) and featureset_version == 0:
         return BatchFeature(name=feature.name, entity=feature.entity)
-    elif isinstance(feature_type, BatchFeatureSetV1) and featureset_version != 0:
+    elif isinstance(feature_type, ProtoBatchFeatureSetV1) and featureset_version != 0:
         return BatchFeatureV1(name=feature.name, entity=feature.entity)
+    elif (
+        isinstance(feature_type, ProtoStreamingFeatureSetV1) and featureset_version == 0
+    ):
+        return StreamingFeature(name=feature.name, entity=feature.entity)
+    elif (
+        isinstance(feature_type, ProtoStreamingFeatureSetV1) and featureset_version != 0
+    ):
+        return StreamingFeatureV1(name=feature.name, entity=feature.entity)
+    elif isinstance(feature_type, ProtoStreamingAggregationFeatureSet):
+        return StreamingAggregationFeature(name=feature.name, entity=feature.entity)
     else:
         raise ValueError(
             f"Feature set type {feature_type} with protocol version {featureset_version} is not supported for extraction"
         )
 
 
-def get_entity_type(value_type: EntitySpec.ValueType):
+def get_env_to_featuresets_mapping(
+    feature_manager_client: FeatureRegistryClient,
+) -> Dict[str, Dict[str, FeatureSetInfo]]:
+    """
+    Get mapping of the account's environment name to feature sets and converts each FeatureSet to FeatureSetInfo
+    :param feature_manager_client:
+    :return: Dict[str, List[FeatureSetInfo]] where the key is the environment name and the value is a list of FeatureSetInfo
+    """
+    env_to_featuresets_mapping = feature_manager_client.get_env_to_featuresets_mapping()
+
+    # Convert the mapping to a dictionary of type Dict[str, List[FeatureSetInfo]]
+    result = {}
+    for (
+        env_name,
+        feature_sets_infos,
+    ) in env_to_featuresets_mapping.env_to_feature_set_mapping.items():
+        env_name = env_name.lower()
+        result[env_name] = {}
+        for feature_set in feature_sets_infos.feature_sets:
+            featureset_type: ProtoFeatureSetType = get_feature_type(
+                feature_set.feature_set_definition.feature_set_spec.feature_set_type
+            )
+            featureset_version = _get_featureset_version(featureset_type)
+
+            feature_set_info = FeatureSetInfo(
+                entity_spec=feature_set.feature_set_definition.feature_set_spec.entity.entity_spec,
+                feature_set_type=featureset_type,
+                feature_version=featureset_version,
+            )
+
+            feature_set_name = (
+                feature_set.feature_set_definition.feature_set_spec.name.lower()
+            )
+            result[env_name][feature_set_name] = feature_set_info
+    return result
+
+
+def get_entity_type(value_type: ProtoEntitySpec.ValueType):
     """
     Normalize entity by the enum
     """
 
-    if value_type == EntitySpec.ValueType.STRING:
+    if value_type == ProtoEntitySpec.ValueType.STRING:
         return str
-    elif value_type == EntitySpec.ValueType.INT:
+    elif value_type == ProtoEntitySpec.ValueType.INT:
         return int
+
+
+def generate_key_unique_name(featureset_name: str):
+    """
+    Generate a name for the feature set key to be registered or treated as a Qwak entity based on its' featuresets'
+    name
+    """
+    return f"{_INTERNAL_KEY_PREFIX}_{featureset_name}_{round(time.time() * 1000)}"
+
+
+def get_batch_source_for_featureset(
+    batch_ds_name: str, feature_registry: FeatureRegistryClient
+):
+    batch_ds = feature_registry.get_data_source_by_name(batch_ds_name)
+    if not batch_ds:
+        raise QwakException(
+            f"Trying to register a feature set with a non registered data source -: {batch_ds_name}"
+        )
+
+    ds_spec = batch_ds.data_source.data_source_definition.data_source_spec
+
+    if ds_spec.WhichOneof("type") != "batch_source":
+        raise ValueError(
+            f"Can only register streaming backfill/batch featuresets with batch sources. "
+            f"Source {batch_ds_name} is of type {ds_spec.WhichOneof('type')}"
+        )
+
+    return ds_spec.batch_source