frogml-core 0.0.1__py3-none-any.whl

frogml_services_mock/mocks/features_online_serving_api.py

@@ -0,0 +1,169 @@
+import json
+from datetime import datetime
+
+import grpc
+import pandas as pd
+
+from _frogml_proto.qwak.feature_store.serving.serving_pb2 import (
+    MultiFeatureValuesResponse,
+)
+from _frogml_proto.qwak.feature_store.serving.serving_pb2_grpc import (
+    ServingServiceServicer,
+)
+from frogml.exceptions import FrogmlException
+
+
+class NpEncoder(json.JSONEncoder):
+    def default(self, obj):
+        import numpy as np
+
+        if isinstance(obj, np.integer):
+            return int(obj)
+        if isinstance(obj, np.floating):
+            return float(obj)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        if isinstance(obj, datetime):
+            return obj.isoformat()
+        return super(NpEncoder, self).default(obj)
+
+
+class ServingServiceMock(ServingServiceServicer):
+    def __init__(self):
+        self.features_by_entity_df = {}
+        self._fail_request_ledger = 0
+
+    def fail_next_request(self):
+        self._fail_request_ledger += 1
+
+    def upsert_features(self, entity_name, features_df):
+        if entity_name not in features_df.columns:
+            raise FrogmlException(
+                "Entity name does not exist in the data frame columns"
+            )
+
+        entity_feature_df = self.features_by_entity_df.get(
+            entity_name, pd.DataFrame(columns=[entity_name])
+        )
+
+        self.features_by_entity_df[entity_name] = (
+            features_df.set_index(entity_name)
+            .combine_first(entity_feature_df.set_index(entity_name))
+            .reset_index()
+            .rename({"index": entity_name})
+        )
+
+    def GetMultiFeatures(self, request, context) -> MultiFeatureValuesResponse:
+        if self._fail_request_ledger > 0:
+            self._fail_request_ledger -= 1
+            raise ValueError("Test failing feature request!")
+
+        if not self._validate_supported_features(request):
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details("Received an unimplemented one-of type for feature")
+            return MultiFeatureValuesResponse()
+
+        indexes, all_values, columns = [], [], []
+        for row_index, entity_values in enumerate(request.entity_values_matrix.rows):
+            values = []
+            for index, entity_name in enumerate(
+                request.entity_values_matrix.header.entity_names
+            ):
+                entity_feature_df = self.features_by_entity_df.get(entity_name)
+                if entity_feature_df is not None:
+                    all_features = self._filter_features(
+                        entity_feature_df, entity_name, entity_values, index
+                    )
+
+                    if len(all_features) > 1:
+                        context.set_code(grpc.StatusCode.INTERNAL)
+                        context.set_details(
+                            "Invalid state, a single entity has multiple feature rows - bug in the store :~("
+                        )
+                        return MultiFeatureValuesResponse()
+
+                    elif len(all_features) == 1:
+                        columns, values = self._populate_data(
+                            all_features, columns, entity_name, request, values
+                        )
+
+            indexes.append(row_index)
+            all_values.append(values)
+        result = {"data": all_values, "columns": columns, "index": indexes}
+
+        return MultiFeatureValuesResponse(
+            pandas_df_as_json=json.dumps(result, cls=NpEncoder)
+        )
+
+    def _populate_data(self, all_features, columns, entity_name, request, values):
+        requested_features = set()
+        self._populate_requested_features(entity_name, request, requested_features)
+        available_features = set(all_features.columns)
+        existing_features = requested_features.intersection(available_features)
+        non_existing_features = requested_features - available_features
+        values = values + (
+            [all_features.loc[0, feature_name] for feature_name in existing_features]
+        )
+
+        columns = self._populate_existing_values(columns, existing_features)
+        columns, values = self._populate_non_existing_values(
+            columns, non_existing_features, values
+        )
+        return columns, values
+
+    def _populate_existing_values(self, columns, existing_features):
+        if not any(x in columns for x in list(existing_features)):
+            columns = columns + (list(existing_features))
+        return columns
+
+    @staticmethod
+    def _populate_non_existing_values(columns, non_existing_features, values):
+        if not any(x in columns for x in list(non_existing_features)):
+            columns = columns + list(non_existing_features)
+            for _ in non_existing_features:
+                values = values + [None]
+        return columns, values
+
+    @staticmethod
+    def _populate_requested_features(entity_name, request, requested_features):
+        for entity_to_features in request.entities_to_features:
+            if entity_to_features.entity_name == entity_name:
+                for feature in entity_to_features.features:
+                    one_of = feature.WhichOneof("type")
+                    if one_of == "batch_feature":
+                        requested_features.add(feature.batch_feature.name)
+                    elif one_of == "batch_v1_feature":
+                        requested_features.add(feature.batch_v1_feature.name)
+                    elif one_of == "streaming_v1_feature":
+                        requested_features.add(feature.streaming_v1_feature.name)
+                    elif one_of == "streaming_feature":
+                        requested_features.add(feature.streaming_feature.name)
+                    elif one_of == "streaming_aggregation":
+                        requested_features.add(feature.streaming_aggregation.name)
+
+    @staticmethod
+    def _validate_supported_features(request):
+        return all(
+            [
+                (
+                    feature.WhichOneof("type") == "batch_v1_feature"
+                    or feature.WhichOneof("type") == "batch_feature"
+                    or feature.WhichOneof("type") == "streaming_feature"
+                    or feature.WhichOneof("type") == "streaming_feature_v1"
+                    or feature.WhichOneof("type") == "streaming_aggregation"
+                )
+                for entities in request.entities_to_features
+                for feature in entities.features
+            ]
+        )
+
+    @staticmethod
+    def _filter_features(entity_feature_df, entity_name, entity_values, index):
+        return (
+            entity_feature_df[
+                entity_feature_df[entity_name].astype(str)
+                == str(entity_values.entity_values[index])
+            ]
+            .drop([entity_name], axis=1)
+            .reset_index(drop=True)
+        )

frogml_services_mock/mocks/features_operator_v3_service.py

@@ -0,0 +1,29 @@
+import uuid
+from typing import Optional
+
+from _frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2 import (
+    ValidationResponse,
+)
+from _frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2_grpc import (
+    FeaturesOperatorAsyncServiceServicer,
+)
+
+
+class FeaturesOperatorV3ServiceMock(FeaturesOperatorAsyncServiceServicer):
+    def __init__(self):
+        self.response: Optional[ValidationResponse] = None
+        super(FeaturesOperatorV3ServiceMock, self).__init__()
+
+    def given_next_response(self, response: ValidationResponse):
+        self.response = response
+
+    def ValidateDataSource(self, request, context) -> str:
+        request_id = str(uuid.uuid4())
+        return ValidationResponse(request_id=request_id)
+
+    def ValidateFeatureSet(self, request, context):
+        request_id = str(uuid.uuid4())
+        return ValidationResponse(request_id=request_id)
+
+    def GetValidationResult(self, request, context):
+        return self.response

frogml_services_mock/mocks/features_set_state_service_api.py

@@ -0,0 +1,57 @@
+import grpc
+
+from _frogml_proto.qwak.feature_store.features.feature_set_state_pb2 import (
+    FeatureSetState,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_state_service_pb2 import (
+    GetFeatureSetStateByIdV1Response,
+    GetFeatureSetStateByNameV1Response,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_state_service_pb2_grpc import (
+    FeatureSetStateServiceServicer,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class FeaturesSetStateServiceApiMock(FeatureSetStateServiceServicer):
+    def __init__(self):
+        self._feature_state_states: dict[(str, str):FeatureSetState] = {}
+        super(FeaturesSetStateServiceApiMock, self).__init__()
+
+    def given_feature_set_state(self, feature_set_state: FeatureSetState):
+        self._feature_state_states[
+            (feature_set_state.feature_set_id, feature_set_state.run_id)
+        ] = feature_set_state
+
+    def GetFeatureSetStateById(self, request, context):
+        """get feature set state by id"""
+        try:
+            return GetFeatureSetStateByIdV1Response(
+                feature_set_state=self._feature_state_states[
+                    (request.feature_set_id, request.run_id)
+                ]
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetFeatureSetStateByName(self, request, context):
+        """get feature set state by name"""
+        try:
+            feature_sets_by_name = [
+                self._feature_state_states[(feature_set_id, run_id)]
+                for feature_set_id, run_id in self._feature_state_states.keys()
+                if self._feature_state_states[(feature_set_id, run_id)].feature_set_name
+                == request.feature_set_name
+            ]
+            if len(feature_sets_by_name) > 0:
+                return GetFeatureSetStateByNameV1Response(
+                    feature_set_state=feature_sets_by_name.pop()
+                )
+            else:
+                context.set_code(grpc.StatusCode.NOT_FOUND)
+                context.set_details("Received a non existing feature set name")
+                return GetFeatureSetStateByNameV1Response()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/feedback_service.py

@@ -0,0 +1,35 @@
+from _frogml_proto.qwak.inference.feedback.feedback_pb2 import (
+    ActualValuesRequest,
+    ActualValuesResponse,
+    ConfigureFeedbackRequest,
+    ConfigureFeedbackResponse,
+)
+from _frogml_proto.qwak.inference.feedback.feedback_pb2_grpc import (
+    FeedbackServiceServicer,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class FeedbackServiceMock(FeedbackServiceServicer):
+    def __init__(self):
+        super(FeedbackServiceMock, self).__init__()
+
+    def PostFeedback(
+        self, request: ActualValuesRequest, context
+    ) -> ActualValuesResponse:
+        try:
+            return ActualValuesResponse(status=ActualValuesResponse.ValueStatus.SUCCESS)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def ConfigureFeedback(
+        self, request: ConfigureFeedbackRequest, context
+    ) -> ConfigureFeedbackResponse:
+        try:
+            return ConfigureFeedbackResponse(
+                status=ConfigureFeedbackResponse.ValueStatus.SUCCESS
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/file_versioning_service.py

@@ -0,0 +1,66 @@
+from collections import defaultdict
+from typing import Dict, List
+
+from _frogml_proto.qwak.file_versioning.file_versioning_pb2 import FileTagSpec
+from _frogml_proto.qwak.file_versioning.file_versioning_service_pb2 import (
+    GetModelFileTagsRequest,
+    GetModelFileTagsResponse,
+    RegisterFileTagRequest,
+    RegisterFileTagResponse,
+)
+from _frogml_proto.qwak.file_versioning.file_versioning_service_pb2_grpc import (
+    FileVersioningManagementServiceServicer,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class FileVersioningServiceMock(FileVersioningManagementServiceServicer):
+    def __init__(self):
+        super(FileVersioningServiceMock, self).__init__()
+        self.tags: Dict[str : List[FileTagSpec]] = defaultdict(list)
+
+    def RegisterFileTag(
+        self, request: RegisterFileTagRequest, context
+    ) -> RegisterFileTagResponse:
+        try:
+            self.tags[request.file_tag_spec.build_id].append(request.file_tag_spec)
+            return RegisterFileTagResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetModelFileTags(
+        self, request: GetModelFileTagsRequest, context
+    ) -> GetModelFileTagsResponse:
+        try:
+            file_tags_list = []
+
+            if not request.build_id:
+                for all_file_tags in self.tags.values():
+                    for file_tag in all_file_tags:
+                        file_tags_list.append(file_tag)
+            else:
+                for file_tag_by_build_id in self.tags[request.build_id]:
+                    if file_tag_by_build_id.model_id == request.model_id:
+                        file_tags_list.append(file_tag_by_build_id)
+
+            if request.filter:
+                file_tags_list_filtered = list(file_tags_list)
+                filter_type: str = request.filter.WhichOneof("filter")
+
+                if filter_type == "tag_contains":
+                    for file_tag in file_tags_list:
+                        if request.filter.tag_contains not in file_tag.tag:
+                            file_tags_list_filtered.remove(file_tag)
+
+                if filter_type == "tag_prefix":
+                    for file_tag in file_tags_list:
+                        if not file_tag.tag.startswith(request.filter.tag_prefix):
+                            file_tags_list_filtered.remove(file_tag)
+
+                file_tags_list = file_tags_list_filtered
+
+            return GetModelFileTagsResponse(file_tags=file_tags_list)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/frogml_mocks.py

@@ -0,0 +1,133 @@
+import dataclasses
+
+from frogml_services_mock.mocks import workspace_manager_service_mock
+from frogml_services_mock.mocks.alert_manager_service_api import (
+    AlertManagerServiceApiMock,
+)
+from frogml_services_mock.mocks.alert_registry_service_api import (
+    AlertsRegistryServiceApiMock,
+)
+from frogml_services_mock.mocks.analytics_api import AnalyticsApiMock
+from frogml_services_mock.mocks.audience_service_api import AudienceServiceApiMock
+from frogml_services_mock.mocks.authentication_service import AuthenticationServiceMock
+from frogml_services_mock.mocks.automation_management_service import (
+    AutomationManagementServiceMock,
+)
+from frogml_services_mock.mocks.autoscaling_service_api import AutoscalingServiceApiMock
+from frogml_services_mock.mocks.batch_job_manager_service import BatchJobManagerService
+from frogml_services_mock.mocks.build_management import BuildsManagementServiceMock
+from frogml_services_mock.mocks.build_orchestrator_build_api import (
+    BuildOrchestratorBuildApiMock,
+)
+from frogml_services_mock.mocks.build_orchestrator_build_settings_api import (
+    BuildOrchestratorBuildSettingsApiMock,
+)
+from frogml_services_mock.mocks.build_orchestrator_service_api import (
+    BuildOrchestratorServiceApiMock,
+)
+from frogml_services_mock.mocks.data_versioning_service import DataVersioningServiceMock
+from frogml_services_mock.mocks.deployment_management_service import (
+    DeploymentManagementServiceMock,
+)
+from frogml_services_mock.mocks.ecosystem_service_api import EcoSystemServiceMock
+from frogml_services_mock.mocks.execution_management_service import (
+    ExecutionManagementServiceMock,
+)
+from frogml_services_mock.mocks.feature_store_data_sources_manager_api import (
+    DataSourceServiceMock,
+)
+from frogml_services_mock.mocks.feature_store_entities_manager_api import (
+    EntityServiceMock,
+)
+from frogml_services_mock.mocks.feature_store_feature_set_manager_api import (
+    FeatureSetServiceMock,
+)
+from frogml_services_mock.mocks.features_operator_v3_service import (
+    FeaturesOperatorV3ServiceMock,
+)
+from frogml_services_mock.mocks.file_versioning_service import FileVersioningServiceMock
+from frogml_services_mock.mocks.fs_offline_serving_service import (
+    FsOfflineServingServiceMock,
+)
+from frogml_services_mock.mocks.instance_template_management_service import (
+    InstanceTemplateManagementServiceMock,
+)
+from frogml_services_mock.mocks.integration_management_service import (
+    IntegrationManagementServiceMock,
+)
+from frogml_services_mock.mocks.internal_build_orchestrator_service import (
+    InternalBuildOrchestratorServiceMock,
+)
+from frogml_services_mock.mocks.job_registry_service_api import (
+    JobRegistryServiceApiMock,
+)
+from frogml_services_mock.mocks.kube_captain_service_api import (
+    KubeCaptainServiceApiMock,
+)
+from frogml_services_mock.mocks.logging_service import LoggingServiceApiMock
+from frogml_services_mock.mocks.model_management_service import (
+    ModelsManagementServiceMock,
+)
+from frogml_services_mock.mocks.model_version_manager_service import (
+    ModelVersionManagerServiceMock,
+)
+from frogml_services_mock.mocks.project_manager_service import ProjectManagerServiceMock
+from frogml_services_mock.mocks.prompt_manager_service import PromptManagerServiceMock
+from frogml_services_mock.mocks.repository_service_mock import RepositoryServiceMock
+from frogml_services_mock.mocks.secret_service import SecretServiceMock
+from frogml_services_mock.mocks.self_service_user_service import (
+    SelfServiceUserServiceMock,
+)
+from frogml_services_mock.mocks.system_secret_service import SystemSecretServiceMock
+from frogml_services_mock.mocks.user_application_instance_service_api import (
+    UserApplicationInstanceServiceApiMock,
+)
+from frogml_services_mock.mocks.vector_serving_api import VectorServingServiceMock
+from frogml_services_mock.mocks.vectors_management_api import (
+    VectorCollectionManagementServiceMock,
+)
+
+
+@dataclasses.dataclass
+class FrogmlMocks:
+    port: int
+    build_orchestrator_build_api: BuildOrchestratorBuildApiMock
+    build_orchestrator_service_api: BuildOrchestratorServiceApiMock
+    build_orchestrator_build_settings_api: BuildOrchestratorBuildSettingsApiMock
+    internal_build_orchestrator_service: InternalBuildOrchestratorServiceMock
+    build_management_service_mock: BuildsManagementServiceMock
+    alert_manager_service_mock: AlertManagerServiceApiMock
+    automation_management_service_mock: AutomationManagementServiceMock
+    job_registry_service_mock: JobRegistryServiceApiMock
+    project_manager_service_mock: ProjectManagerServiceMock
+    kube_captain_service_mock: KubeCaptainServiceApiMock
+    file_versioning_service_mock: FileVersioningServiceMock
+    data_versioning_service_mock: DataVersioningServiceMock
+    model_management_service_mock: ModelsManagementServiceMock
+    logging_service_mock: LoggingServiceApiMock
+    autoscaling_service_mock: AutoscalingServiceApiMock
+    audience_api_mock: AudienceServiceApiMock
+    self_service_user_service_mock: SelfServiceUserServiceMock
+    analytics_api_mock: AnalyticsApiMock
+    deployment_management_service_mock: DeploymentManagementServiceMock
+    batch_job_manager_service: BatchJobManagerService
+    ecosystem_client_mock: EcoSystemServiceMock
+    user_application_instance_service_mock: UserApplicationInstanceServiceApiMock
+    secret_service_mock: SecretServiceMock
+    authentication_service_mock: AuthenticationServiceMock
+    features_operator_service: FeaturesOperatorV3ServiceMock
+    fs_offline_serving_service: FsOfflineServingServiceMock
+    fs_data_sources_service: DataSourceServiceMock
+    fs_entities_service: EntityServiceMock
+    fs_feature_sets_service: FeatureSetServiceMock
+    instance_templates_service: InstanceTemplateManagementServiceMock
+    alerts_registry_service: AlertsRegistryServiceApiMock
+    workspace_manager_service: workspace_manager_service_mock
+    vector_serving_service: VectorServingServiceMock
+    vector_collection_service: VectorCollectionManagementServiceMock
+    execution_management_service: ExecutionManagementServiceMock
+    system_secret_service: SystemSecretServiceMock
+    integration_management_service: IntegrationManagementServiceMock
+    prompt_manager_service: PromptManagerServiceMock
+    model_version_manager_service: ModelVersionManagerServiceMock
+    repository_service: RepositoryServiceMock

frogml_services_mock/mocks/fs_offline_serving_service.py

@@ -0,0 +1,53 @@
+import uuid
+from typing import Optional
+
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFeatureValuesInRangeRequest,
+    GetFeatureValuesInRangeResponse,
+    GetFeatureValuesRequest,
+    GetFeatureValuesResponse,
+    GetFeatureValuesResultResponse,
+    GetFileUploadUrlResponse,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2_grpc import (
+    FeatureStoreOfflineServingAsyncServiceServicer,
+)
+from _frogml_proto.qwak.offline.serving.v1.options_pb2 import (
+    OfflineServingQueryOptions as ProtoOfflineServingQueryOptions,
+)
+
+
+class FsOfflineServingServiceMock(FeatureStoreOfflineServingAsyncServiceServicer):
+    def __init__(self):
+        self.file_upload_url_response: Optional[GetFileUploadUrlResponse] = None
+        self.response: Optional[GetFeatureValuesResultResponse] = None
+        self.latest_query_options: Optional[ProtoOfflineServingQueryOptions] = None
+        super(FsOfflineServingServiceMock, self).__init__()
+
+    def given_next_file_upload_url(
+        self, file_upload_url_response: GetFileUploadUrlResponse
+    ):
+        self.file_upload_url_response = file_upload_url_response
+
+    def given_next_response(self, response: GetFeatureValuesResultResponse):
+        self.response = response
+
+    def GetFileUploadUrl(self, request, context) -> Optional[GetFileUploadUrlResponse]:
+        return self.file_upload_url_response
+
+    def GetFeatureValues(
+        self, request: GetFeatureValuesRequest, context
+    ) -> GetFeatureValuesResponse:
+        request_id = str(uuid.uuid4())
+        self.latest_query_options = request.options
+        return GetFeatureValuesResponse(request_id=request_id)
+
+    def GetFeatureValuesInRange(self, request: GetFeatureValuesInRangeRequest, context):
+        request_id = str(uuid.uuid4())
+        self.latest_query_options = request.options
+        return GetFeatureValuesInRangeResponse(request_id=request_id)
+
+    def GetFeatureValuesResult(
+        self, request, context
+    ) -> Optional[GetFeatureValuesResultResponse]:
+        return self.response