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

qwak/qwak_client/client.py

@@ -1,22 +1,36 @@
 from datetime import timedelta
 from functools import lru_cache
+from pathlib import Path
+from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from _qwak_proto.qwak.build.v1.build_pb2 import DESCRIPTOR
 from qwak.clients.analytics.client import AnalyticsEngineClient
 from qwak.clients.automation_management.client import AutomationsManagementClient
-from qwak.clients.build_management.client import BuildsManagementClient
+from qwak.clients.batch_job_management.client import BatchJobManagerClient
 from qwak.clients.build_orchestrator.client import BuildFilter, BuildOrchestratorClient
+from qwak.clients.data_versioning.client import DataVersioningManagementClient
+from qwak.clients.data_versioning.data_tag_filter import DataTagFilter
 from qwak.clients.deployment.client import DeploymentManagementClient
 from qwak.clients.feature_store import FeatureRegistryClient
+from qwak.clients.file_versioning.client import FileVersioningManagementClient
+from qwak.clients.file_versioning.file_tag_filter import FileTagFilter
+from qwak.clients.instance_template.client import InstanceTemplateManagementClient
 from qwak.clients.model_management.client import ModelsManagementClient
 from qwak.clients.project.client import ProjectsManagementClient
 from qwak.exceptions import QwakException
-from qwak.feature_store.feature_sets.batch import BatchFeatureSetV1
+from qwak.feature_store.feature_sets.batch import BatchFeatureSet
+from qwak.inner.build_config.build_config_v1 import BuildConfigV1
+from qwak.inner.tool import Auth0ClientBase
+from qwak.model.base import QwakModel
+from qwak.qwak_client.batch_jobs.execution import Execution
+from qwak.qwak_client.batch_jobs.task import Task
 from qwak.qwak_client.builds.build import Build
 from qwak.qwak_client.builds.filters.metric_filter import MetricFilter
 from qwak.qwak_client.builds.filters.parameter_filter import ParameterFilter
+from qwak.qwak_client.data_versioning.data_tag import DataTag
 from qwak.qwak_client.deployments.deployment import Deployment, EnvironmentAudienceRoute
+from qwak.qwak_client.file_versioning.file_tag import FileTag
 from qwak.qwak_client.models.model import Model
 from qwak.qwak_client.models.model_metadata import ModelMetadata
 from qwak.qwak_client.projects.project import Project
@@ -43,14 +57,22 @@ class QwakClient:
     def _get_model_management(self):
         return ModelsManagementClient()
 
-    @lru_cache(maxsize=1)
-    def _get_build_management(self):
-        return BuildsManagementClient()
-
     @lru_cache(maxsize=1)
     def _get_build_orchestrator(self):
         return BuildOrchestratorClient()
 
+    @lru_cache(maxsize=1)
+    def _get_batch_job_manager(self):
+        return BatchJobManagerClient()
+
+    @lru_cache(maxsize=1)
+    def _get_file_versioning_management(self):
+        return FileVersioningManagementClient()
+
+    @lru_cache(maxsize=1)
+    def _get_data_versioning_management(self):
+        return DataVersioningManagementClient()
+
     @lru_cache(maxsize=1)
     def _get_deployment_management(self):
         return DeploymentManagementClient()
@@ -67,6 +89,29 @@ class QwakClient:
     def _get_analytics_engine(self):
         return AnalyticsEngineClient()
 
+    @lru_cache(maxsize=1)
+    def _get_instance_template_client(self):
+        return InstanceTemplateManagementClient()
+
+    @lru_cache(maxsize=1)
+    def _get_auth_client(self, api_key):
+        return Auth0ClientBase(api_key)
+
+    def get_token(self, api_key: Optional[str] = None) -> str:
+        """
+        Retrieve a token by apy key.
+        Use Cases:
+        Default - from already configured (qwak configured) - conf file
+        Support QWAK_API_KEY environment variable
+        Get token for a specified API key
+        Args:
+            api_key (str): a key for token exchange
+        Returns:
+            str: a client token
+        """
+        client = self._get_auth_client(api_key)
+        return client.get_token()
+
     def get_latest_build(
         self, model_id: str, build_status: str = "SUCCESSFUL"
     ) -> Optional[str]:
@@ -193,8 +238,8 @@ class QwakClient:
         model = self._get_model_management().get_model(model_id=model_id)
         model_uuid = model.uuid
 
-        metric_filters = list()
-        parameter_filters = list()
+        metric_filters = []
+        parameter_filters = []
 
         for build_filter in filters:
             if isinstance(build_filter, MetricFilter):
@@ -217,6 +262,75 @@ class QwakClient:
 
         return [Build.from_proto(build) for build in builds_proto]
 
+    def list_file_tags(
+        self,
+        model_id: str,
+        build_id: str = "",
+        filter: Optional[FileTagFilter] = None,
+    ) -> List[FileTag]:
+        """
+        List file tags by its model ID
+
+        Args:
+            model_id (str): The model ID
+            build_id (str): The build ID - optional.
+                If not specified, returns all model file tags.
+            filter (FileTagFilter): Filter list by - optional.
+                value (str): The filter value.
+                type (enum): The filter type - FILE_TAG_FILTER_TYPE_CONTAINS/FILE_TAG_FILTER_TYPE_PREFIX.
+                If not specified, returns all model file tags.
+
+
+        Returns:
+             List[FileTag]: List of file tags with their specifications.
+
+        """
+        model_file_tags_proto = (
+            self._get_file_versioning_management().get_model_file_tags(
+                model_id=model_id,
+                build_id=build_id,
+                file_tag_filter=None if filter is None else filter.to_proto(),
+            )
+        )
+
+        return [
+            FileTag.from_proto(file_tag) for file_tag in model_file_tags_proto.file_tags
+        ]
+
+    def list_data_tags(
+        self,
+        model_id: str,
+        build_id: str = "",
+        filter: Optional[DataTagFilter] = None,
+    ) -> List[DataTag]:
+        """
+        List data tags by its model ID
+
+        Args:
+            model_id (str): The model ID
+            build_id (str): The build ID - optional.
+                If not specified, returns all model data tags.
+            filter (DataTagFilter): Filter list by - optional.
+                value (str): The filter value.
+                type (enum): The filter type - DATA_TAG_FILTER_TYPE_CONTAINS/DATA_TAG_FILTER_TYPE_PREFIX.
+                If not specified, returns all model file tags.
+
+        Returns:
+             List[DataTag]: List of data tags with their specifications.
+
+        """
+        model_data_tags_proto = (
+            self._get_data_versioning_management().get_model_data_tags(
+                model_id=model_id,
+                build_id=build_id,
+                data_tag_filter=None if filter is None else filter.to_proto(),
+            )
+        )
+
+        return [
+            DataTag.from_proto(data_tag) for data_tag in model_data_tags_proto.data_tags
+        ]
+
     def set_tag(self, build_id: str, tag: str) -> None:
         """
         Assign a tag to an existing build
@@ -238,7 +352,11 @@ class QwakClient:
         """
         self._get_build_orchestrator().register_tags(build_id=build_id, tags=tags)
 
-    def create_project(self, project_name: str, project_description: str) -> str:
+    def create_project(
+        self,
+        project_name: str,
+        project_description: str,
+    ) -> str:
         """
         Create project
 
@@ -251,7 +369,8 @@ class QwakClient:
 
         """
         project = self._get_project_management().create_project(
-            project_name=project_name, project_description=project_description
+            project_name=project_name,
+            project_description=project_description,
         )
 
         return project.project.project_id
@@ -268,7 +387,7 @@ class QwakClient:
 
         """
         return Project.from_proto(
-            self._get_project_management().get_project(project_id)
+            self._get_project_management().get_project(project_id).project.spec
         )
 
     def list_projects(self) -> List[Project]:
@@ -293,7 +412,10 @@ class QwakClient:
         return self._get_project_management().delete_project(project_id=project_id)
 
     def create_model(
-        self, project_id: str, model_name: str, model_description: str
+        self,
+        project_id: str,
+        model_name: str,
+        model_description: str,
     ) -> str:
         """
         Create model
@@ -458,7 +580,7 @@ class QwakClient:
 
     def get_deployed_build_id_per_environment(
         self, model_id: str
-    ) -> Optional[Dict[str, str]]:
+    ) -> Optional[Dict[str, Dict[str, str]]]:
         """
         Get deployed build ID per environment by its model ID
 
@@ -482,16 +604,14 @@ class QwakClient:
         environment_to_deployment_details = (
             deployment_details.environment_to_deployment_details
         )
-        env_to_deployed_build_id = {}
-        for (
-            env_id,
-            environment_deployment_details_message,
-        ) in environment_to_deployment_details.items():
-            env_to_deployed_build_id[
-                env_id
-            ] = environment_deployment_details_message.deployments_details[0].build_id
 
-        return env_to_deployed_build_id
+        return {
+            env_id: {
+                deployment.variation.name: deployment.build_id
+                for deployment in v.deployments_details
+            }
+            for env_id, v in environment_to_deployment_details.items()
+        }
 
     def run_analytics_query(
         self, query: str, timeout: timedelta = None
@@ -517,7 +637,7 @@ class QwakClient:
         url = client.get_analytics_data(query=query, timeout=timeout)
         return pd.read_csv(url)
 
-    def list_feature_sets(self) -> List[BatchFeatureSetV1]:
+    def list_feature_sets(self) -> List[BatchFeatureSet]:
         """
         List all feature sets
 
@@ -528,7 +648,7 @@ class QwakClient:
         feature_set_protos = client.list_feature_sets().feature_families
 
         return [
-            BatchFeatureSetV1.from_proto(
+            BatchFeatureSet._from_proto(
                 feature_set.feature_set_definition.feature_set_spec
             )
             for feature_set in feature_set_protos
@@ -552,6 +672,34 @@ class QwakClient:
         else:
             raise QwakException(f"Feature set named '{feature_set_name}' not found")
 
+    def delete_featureset_version(self, featureset_name: str, version_number: int):
+        """
+        Delete a feature set version
+
+        Args:
+            featureset_name (str): The feature set name
+            version_number (int): version number
+        """
+        client = self._get_feature_registry_client()
+
+        client.delete_featureset_version(
+            featureset_name=featureset_name, version_number=version_number
+        )
+
+    def set_featureset_active_version(self, featureset_name: str, version_number: int):
+        """
+        Set active feature set version
+
+        Args:
+            featureset_name (str): The feature set name
+            version_number (int): version number
+        """
+        client = self._get_feature_registry_client()
+
+        client.set_active_featureset_version(
+            featureset_name=featureset_name, version_number=version_number
+        )
+
     def delete_data_source(self, data_source_name: str):
         """
         Delete a data source
@@ -608,3 +756,107 @@ class QwakClient:
                 f"Could not find automation with given name '{automation_name}'"
             )
         client.run_automation(automation_id=automation.id)
+
+    def build_model(
+        self,
+        model_id: str,
+        main_module_path: str = ".",
+        dependencies_file: Optional[str] = None,
+        dependencies_list: Optional[List[str]] = None,
+        tags: Optional[List[str]] = None,
+        instance: str = "small",
+        gpu_compatible: bool = False,
+        run_tests: bool = True,
+        validate_build_artifact: bool = True,
+        validate_build_artifact_timeout: Optional[int] = None,
+        prebuilt_qwak_model: Optional[QwakModel] = None,
+        base_image: Optional[str] = None,
+        build_name: str = "",
+    ):
+        from qwak.qwak_client.build_api_helpers.trigger_build_api import (
+            trigger_build_api,
+        )
+
+        config = BuildConfigV1()
+        config.build_properties.model_id = model_id
+        config.build_properties.build_name = build_name
+        config.build_properties.model_uri.uri = main_module_path
+        config.build_env.python_env.dependency_file_path = dependencies_file
+        config.build_properties.tags = tags if tags else []
+        config.build_env.remote.resources.instance = instance
+        config.build_properties.gpu_compatible = gpu_compatible
+        config.step.validate_build_artifact = validate_build_artifact
+        config.step.tests = run_tests
+        config.fetch_base_docker_image_name(
+            self._get_build_orchestrator(), self._get_instance_template_client()
+        )
+        config.pre_built_model = prebuilt_qwak_model
+        if validate_build_artifact_timeout:
+            config.step.validate_build_artifact_timeout = (
+                validate_build_artifact_timeout
+            )
+        if base_image:
+            config.build_env.docker.base_image = base_image
+        if dependencies_list:
+            with TemporaryDirectory() as temporary_directory:
+                requirements_filename = Path(temporary_directory) / "requirements.txt"
+                with open(requirements_filename, "w") as temp_dependencies_file:
+                    temp_dependencies_file.write("\n".join(dependencies_list))
+                config.build_env.python_env.dependency_file_path = requirements_filename
+                return trigger_build_api(config)
+        else:
+            return trigger_build_api(config)
+
+    def list_executions(
+        self,
+        model_id: str,
+        build_id: str = "",
+    ) -> List[Execution]:
+        """
+        List batch executions by its model ID
+
+        Args:
+            model_id (str): The model ID
+            build_id (str): The build ID - optional.
+                If not specified, returns all model executions.
+
+
+        Returns:
+             List[Execution]: List of executions with their specifications.
+
+        """
+        model_executions_proto = self._get_batch_job_manager().list_batch_jobs(
+            model_id=model_id,
+            build_id=build_id,
+        )
+
+        return [
+            Execution.from_proto(execution)
+            for execution in model_executions_proto.batch_jobs
+        ]
+
+    def list_execution_tasks(
+        self,
+        execution_id: str,
+    ) -> List[Task]:
+        """
+        List batch executions tasks by its job ID
+
+        Args:
+            execution_id (str): The execution ID
+
+
+        Returns:
+             List[Task]: List of execution tasks with their specifications.
+
+        """
+        model_execution_tasks_proto = (
+            self._get_batch_job_manager().get_batch_job_details(
+                job_id=execution_id,
+            )
+        )
+
+        return [
+            Task.from_proto(task)
+            for task in model_execution_tasks_proto.batch_job.task_executions
+        ]

qwak/qwak_client/data_versioning/data_tag.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass, field
+
+from _qwak_proto.qwak.data_versioning.data_versioning_pb2 import (
+    DataTagSpec as DataTagSpecProto,
+)
+
+
+@dataclass
+class DataTag:
+    model_id: str = field(default=None)
+    build_id: str = field(default=None)
+    tag: str = field(default=None)
+    extension_type: str = field(default=None)
+    environment_id: str = field(default=None)
+
+    @staticmethod
+    def from_proto(data_versioning_spec_proto: DataTagSpecProto):
+        return DataTag(
+            model_id=data_versioning_spec_proto.model_id,
+            build_id=data_versioning_spec_proto.build_id,
+            tag=data_versioning_spec_proto.tag,
+            extension_type=data_versioning_spec_proto.extension_type,
+            environment_id=data_versioning_spec_proto.environment_id,
+        )

qwak/qwak_client/deployments/deployment.py

@@ -282,7 +282,7 @@ class Deployment:
                         proto_trigger.prometheus_trigger.query_spec.aggregation_type
                     ),
                     time_period=proto_trigger.prometheus_trigger.query_spec.time_period,
-                    error_code=proto_trigger.prometheus_trigger.query_spec.error_code,
+                    error_code=proto_trigger.prometheus_trigger.query_spec.status_code,
                 ),
                 threshold=proto_trigger.prometheus_trigger.threshold,
             )
@@ -337,20 +337,22 @@ class Deployment:
                     topic=proto_config.kafka.consumer.topic,
                     group=proto_config.kafka.consumer.group,
                     timeout=proto_config.kafka.consumer.timeout,
-                    auto_offset=AutoOffSet(proto_config.kafka.consumer.auto_offset_type)
-                    if proto_config.kafka.consumer.auto_offset_type
-                    else None,
+                    auto_offset=(
+                        AutoOffSet(proto_config.kafka.consumer.auto_offset_type)
+                        if proto_config.kafka.consumer.auto_offset_type
+                        else None
+                    ),
                     max_batch_size=proto_config.kafka.consumer.max_batch_size,
                     max_poll_latency=proto_config.kafka.consumer.max_poll_latency,
                 ),
                 producer=Producer(
                     bootstrap_server=list(proto_config.kafka.producer.bootstrap_server),
                     topic=proto_config.kafka.producer.topic,
-                    compression=Compression(
-                        proto_config.kafka.producer.compression_type
-                    )
-                    if proto_config.kafka.producer.compression_type
-                    else None,
+                    compression=(
+                        Compression(proto_config.kafka.producer.compression_type)
+                        if proto_config.kafka.producer.compression_type
+                        else None
+                    ),
                 ),
                 workers=proto_config.kafka.workers,
             )

qwak/qwak_client/file_versioning/file_tag.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass, field
+
+from _qwak_proto.qwak.file_versioning.file_versioning_pb2 import (
+    FileTagSpec as FileTagSpecProto,
+)
+
+
+@dataclass
+class FileTag:
+    model_id: str = field(default=None)
+    build_id: str = field(default=None)
+    tag: str = field(default=None)
+    extension_type: str = field(default=None)
+    environment_id: str = field(default=None)
+
+    @staticmethod
+    def from_proto(file_versioning_spec_proto: FileTagSpecProto):
+        return FileTag(
+            model_id=file_versioning_spec_proto.model_id,
+            build_id=file_versioning_spec_proto.build_id,
+            tag=file_versioning_spec_proto.tag,
+            extension_type=file_versioning_spec_proto.extension_type,
+            environment_id=file_versioning_spec_proto.environment_id,
+        )

qwak/qwak_client/projects/project.py

@@ -1,8 +1,10 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 
-from _qwak_proto.qwak.projects.projects_pb2 import Project as ProjectProto
-from _qwak_proto.qwak.projects.projects_pb2 import ProjectSpec as ProjectSpecProto
+from _qwak_proto.qwak.projects.projects_pb2 import (
+    Project as ProjectProto,
+    ProjectSpec as ProjectSpecProto,
+)
 from google.protobuf.timestamp_pb2 import Timestamp
 
 

qwak/tools/logger/logger.py

@@ -275,7 +275,7 @@ def set_file_handler_log_file(
     logger: logging.Logger, handler_name: str, log_file: Path
 ):
     existing_handler = get_handler_from_logger(logger, handler_name)
-    if type(existing_handler) != RotatingFileHandler:
+    if not isinstance(existing_handler, RotatingFileHandler):
         raise QwakException(
             f"Error in setting log file. Error message: handler of name {handler_name} is not a file logger handler"
         )

qwak/utils/datetime_utils.py

@@ -0,0 +1,18 @@
+from datetime import datetime, timezone
+from typing import Optional
+
+from google.protobuf.timestamp_pb2 import Timestamp as ProtoTimestamp
+
+
+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

qwak/utils/dict_utils.py

@@ -0,0 +1,12 @@
+from copy import deepcopy
+from typing import Any, Dict
+
+
+def remove_none_value_keys(dict: Dict[Any, Any]):
+    copy = deepcopy(dict)
+    keys_to_remove = [k for k, v in dict.items() if v is None]
+
+    for k in keys_to_remove:
+        del copy[k]
+
+    return copy

qwak/vector_store/client.py

@@ -26,6 +26,9 @@ class VectorStoreClient:
     _metric_dict: Dict[str, VectorCollectionMetric] = {
         "l2_squared": VectorCollectionMetric.COLLECTION_METRIC_L2_SQUARED,
         "cosine": VectorCollectionMetric.COLLECTION_METRIC_COSINE,
+        "dot_product": VectorCollectionMetric.COLLECTION_METRIC_DOT_PRODUCT,
+        "l1": VectorCollectionMetric.COLLECTION_METRIC_L1,
+        "hamming": VectorCollectionMetric.COLLECTION_METRIC_HAMMING,
     }
 
     def __init__(self, edge_services_url: Optional[str] = None):
@@ -45,6 +48,7 @@ class VectorStoreClient:
         description: str = None,
         metric: str = "l2_squared",
         vectorizer: Optional[str] = None,
+        multi_tenant: bool = False,
     ) -> Collection:
         """
         Creates a new collection with the given name and dimension.
@@ -58,6 +62,8 @@ class VectorStoreClient:
             description (str, optional): A human-readable description of the collection.
             metric (int): The distance metric used by the collection when executing similarity search
             vectorizer (str): an optional Qwak model used for vector embedding in case natural input is provided
+            multi_tenant (bool): Whether this collection has multitenancy. defaults to False and cannot be changed
+                                after the collection is created
 
         Returns:
             Collection: The Collection object which is used to interact with the vector store.
@@ -71,6 +77,7 @@ class VectorStoreClient:
             dimension=dimension,
             metric=self._metrics_mapper(metric),
             vectorizer=vectorizer,
+            multi_tenant=multi_tenant,
         ).definition
 
         return self._collection_from_definition(proto_definition)
@@ -131,6 +138,7 @@ class VectorStoreClient:
             description=collection_definition.collection_spec.description,
             vectorizer=collection_definition.collection_spec.vectorizer.qwak_model_name,
             vector_serving_client=self._vector_serving_client,
+            muli_tenant=collection_definition.collection_spec.multi_tenancy_enabled,
         )
 
     def _metrics_mapper(self, metric: str) -> VectorCollectionMetric: