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

qwak_services_mock/mocks/location_discovery_service_api.py

@@ -0,0 +1,104 @@
+from typing import Dict, Optional
+
+import grpc
+from _qwak_proto.qwak.service_discovery.service_discovery_location_pb2 import (
+    ServiceLocationDescriptor,
+)
+from _qwak_proto.qwak.service_discovery.service_discovery_location_service_pb2 import (
+    GetServingUrlRequestResponse,
+)
+from _qwak_proto.qwak.service_discovery.service_discovery_location_service_pb2_grpc import (
+    LocationDiscoveryServiceServicer,
+)
+
+
+class LocationDiscoveryServiceApiMock(LocationDiscoveryServiceServicer):
+    """
+    Mock implementation of the LocationDiscoveryService for testing SDK behavior.
+    Allows setting mock responses and optional error codes for each endpoint.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._responses: Dict[str, Optional[ServiceLocationDescriptor]] = {}
+        self._error_codes: Dict[str, grpc.StatusCode] = {}
+
+    def _set_mock(
+        self,
+        key: str,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._responses[key] = response
+        self._error_codes[key] = error_code
+
+    def _handle(
+        self, key: str, context: grpc.ServicerContext
+    ) -> GetServingUrlRequestResponse:
+        response = self._responses.get(key)
+        if response:
+            return GetServingUrlRequestResponse(location=response)
+        context.set_code(self._error_codes.get(key, grpc.StatusCode.NOT_FOUND))
+        context.set_details(f"No mock response set for {key}")
+        return GetServingUrlRequestResponse()
+
+    # Setters
+    def set_get_offline_serving_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("offline", response, error_code)
+
+    def set_get_distribution_manager_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("distribution", response, error_code)
+
+    def set_get_analytics_engine_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("analytics", response, error_code)
+
+    def set_get_metrics_gateway_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("metrics", response, error_code)
+
+    def set_get_features_operator_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("features", response, error_code)
+
+    def set_get_hosting_gateway_url_response(
+        self,
+        response: Optional[ServiceLocationDescriptor],
+        error_code: grpc.StatusCode = grpc.StatusCode.NOT_FOUND,
+    ):
+        self._set_mock("hosting", response, error_code)
+
+    def GetOfflineServingUrl(self, request, context):
+        return self._handle("offline", context)
+
+    def GetDistributionManagerUrl(self, request, context):
+        return self._handle("distribution", context)
+
+    def GetAnalyticsEngineUrl(self, request, context):
+        return self._handle("analytics", context)
+
+    def GetMetricsGatewayUrl(self, request, context):
+        return self._handle("metrics", context)
+
+    def GetFeaturesOperatorUrl(self, request, context):
+        return self._handle("features", context)
+
+    def GetHostingGatewayUrl(self, request, context):
+        return self._handle("hosting", context)

qwak_services_mock/mocks/model_management_service.py

@@ -10,9 +10,8 @@ from _qwak_proto.qwak.models.models_pb2 import (
     GetModelMetadataResponse,
     GetModelRequest,
     GetModelResponse,
-    ListModelsResponse,
     ListModelsMetadataResponse,
-    Branch,
+    ListModelsResponse,
     Model,
     ModelMetadata,
     ModelSpec,
@@ -85,13 +84,20 @@ class ModelsManagementServiceMock(ModelsManagementServiceServicer):
         except Exception as e:
             raise_internal_grpc_error(context, e)
 
-
     def ListModels(self, request, context):
-        models = [model for _, model in self.models.items() if model.project_id == request.project_id]
+        models = [
+            model
+            for _, model in self.models.items()
+            if model.project_id == request.project_id
+        ]
         return ListModelsResponse(models=models)
 
     def ListModelsMetadata(self, request, context):
-        models = [ModelMetadata(model=model) for _, model in self.models.items() if model.project_id == request.project_id]
+        models = [
+            ModelMetadata(model=model)
+            for _, model in self.models.items()
+            if model.project_id == request.project_id
+        ]
         return ListModelsMetadataResponse(model_metadata=models)
 
     @staticmethod

qwak_services_mock/mocks/prompt_manager_service.py

@@ -0,0 +1,244 @@
+from typing import Dict, List, Tuple
+
+import grpc
+from _qwak_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
+    CreatePromptRequest as ProtoCreatePromptRequest,
+    CreatePromptResponse as ProtoCreatePromptResponse,
+    CreatePromptVersionRequest as ProtoCreatePromptVersionRequest,
+    CreatePromptVersionResponse as ProtoCreatePromptVersionResponse,
+    DeletePromptRequest as ProtoDeletePromptRequest,
+    DeletePromptResponse as ProtoDeletePromptResponse,
+    DeletePromptVersionRequest as ProtoDeletePromptVersionRequest,
+    DeletePromptVersionResponse as ProtoDeletePromptVersionResponse,
+    GetPromptByNameRequest as ProtoGetPromptByNameRequest,
+    GetPromptByNameResponse as ProtoGetPromptByNameResponse,
+    GetPromptVersionByPromptNameRequest as ProtoGetPromptVersionByPromptNameRequest,
+    GetPromptVersionByPromptNameResponse as ProtoGetPromptVersionByPromptNameResponse,
+    SetDefaultPromptVersionRequest as ProtoSetDefaultPromptVersionRequest,
+    SetDefaultPromptVersionResponse as ProtoSetDefaultPromptVersionResponse,
+    UpdatePromptRequest as ProtoUpdatePromptRequest,
+    UpdatePromptResponse as ProtoUpdatePromptResponse,
+)
+from _qwak_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2_grpc import (
+    PromptManagerServiceServicer,
+)
+from _qwak_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
+    Prompt as ProtoPrompt,
+    PromptMetadata as ProtoPromptMetadata,
+    PromptSpec as ProtoPromptSpec,
+    PromptVersion as ProtoPromptVersion,
+    PromptVersionDefinition as ProtoPromptVersionDefinition,
+    PromptVersionMetadata as ProtoPromptVersionMetadata,
+)
+from google.protobuf.timestamp_pb2 import Timestamp
+
+
+class PromptManagerServiceMock(PromptManagerServiceServicer):
+    prompts: Dict[str, ProtoPrompt] = {}
+    prompts_versions: Dict[Tuple[str, int], ProtoPromptVersionDefinition] = {}
+    prompts_next_version: Dict[str, int] = {}
+
+    def clear(self):
+        self.prompts = {}
+        self.prompts_versions = {}
+        self.prompts_next_version = {}
+
+    def _is_prompt_name_valid(self, test_str: str) -> bool:
+        import re
+
+        regex = r"^[a-z0-9](?:[-_]?[a-z0-9]+)+$"
+        return re.match(pattern=regex, string=test_str) is not None
+
+    def _raise_on_missing_prompt(self, prompt_name: str, context):
+        if prompt_name not in self.prompts:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Prompt name does not exist")
+            raise ValueError("Prompt name does not exist")
+
+    def _raise_on_missing_prompt_version(self, prompt_name: str, version: int, context):
+        if (
+            prompt_name not in self.prompts
+            or (prompt_name, version) not in self.prompts_versions
+        ):
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Prompt version does not exist")
+            raise ValueError("Prompt version does not exist")
+
+    def CreatePrompt(self, request: ProtoCreatePromptRequest, context):
+        """Create a new prompt"""
+
+        if request.prompt_name in self.prompts:
+            context.set_code(grpc.StatusCode.ALREADY_EXISTS)
+            context.set_details("Prompt name already exists")
+            raise ValueError("Prompt name already exists")
+
+        if not self._is_prompt_name_valid(request.prompt_name):
+            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
+            context.set_details("Bad Prompt name")
+            raise ValueError("Bad Prompt name")
+
+        prompt_version_definition = ProtoPromptVersionDefinition(
+            version_number=1,
+            version_spec=request.prompt_version_spec,
+            version_metadata=ProtoPromptVersionMetadata(
+                created_by="nirh@qwak.com", created_at=Timestamp().GetCurrentTime()
+            ),
+        )
+
+        prompt = ProtoPrompt(
+            name=request.prompt_name,
+            default_version_definition=prompt_version_definition,
+            metadata=ProtoPromptMetadata(
+                created_at=Timestamp().GetCurrentTime(), created_by="nirh@qwak.com"
+            ),
+            prompt_spec=ProtoPromptSpec(description=request.prompt_spec.description),
+        )
+
+        self.prompts[prompt.name] = prompt
+        self.prompts_next_version[prompt.name] = 2
+        self.prompts_versions[
+            (prompt.name, prompt.default_version_definition.version_number)
+        ] = prompt_version_definition
+
+        return ProtoCreatePromptResponse(prompt=prompt)
+
+    def CreatePromptVersion(self, request: ProtoCreatePromptVersionRequest, context):
+        """
+        Create new prompt version.
+        Set new version in versions dict.
+        Bump prompt next version.
+        Set default version if needed
+        """
+        self._raise_on_missing_prompt(prompt_name=request.prompt_name, context=context)
+
+        existing_prompt: ProtoPrompt = self.prompts[request.prompt_name]
+        new_prompt_version_number: int = self.prompts_next_version[request.prompt_name]
+        prompt_version_definition = ProtoPromptVersionDefinition(
+            version_number=new_prompt_version_number,
+            version_spec=request.prompt_version_spec,
+            version_metadata=ProtoPromptVersionMetadata(
+                created_by="nirh@qwak.com", created_at=Timestamp().GetCurrentTime()
+            ),
+        )
+
+        self.prompts_next_version[request.prompt_name] = new_prompt_version_number + 1
+
+        if request.set_default:
+            existing_prompt.default_version_definition.CopyFrom(
+                prompt_version_definition
+            )
+
+        self.prompts_versions[
+            (request.prompt_name, new_prompt_version_number)
+        ] = prompt_version_definition
+        return ProtoCreatePromptVersionResponse(
+            prompt_version=ProtoPromptVersion(
+                prompt_name=request.prompt_name,
+                prompt_version_definition=prompt_version_definition,
+            )
+        )
+
+    def UpdatePrompt(self, request: ProtoUpdatePromptRequest, context):
+        """Update Prompt"""
+        self._raise_on_missing_prompt(prompt_name=request.prompt_name, context=context)
+
+        existing_prompt: ProtoPrompt = self.prompts[request.prompt_name]
+        existing_prompt.prompt_spec.CopyFrom(request.prompt_spec)
+
+        return ProtoUpdatePromptResponse()
+
+    def GetPromptByName(self, request: ProtoGetPromptByNameRequest, context):
+        """Get prompt with its default version by name"""
+        self._raise_on_missing_prompt(prompt_name=request.prompt_name, context=context)
+
+        existing_prompt: ProtoPrompt = self.prompts[request.prompt_name]
+        return ProtoGetPromptByNameResponse(prompt=existing_prompt)
+
+    def GetPromptVersionByPromptName(
+        self, request: ProtoGetPromptVersionByPromptNameRequest, context
+    ):
+        """Get prompt version by name and optional version number"""
+
+        self._raise_on_missing_prompt(
+            prompt_name=request.prompt_name,
+            context=context,
+        )
+
+        self._raise_on_missing_prompt_version(
+            prompt_name=request.prompt_name,
+            version=request.version_number,
+            context=context,
+        )
+
+        return ProtoGetPromptVersionByPromptNameResponse(
+            prompt_version=ProtoPromptVersion(
+                prompt_name=request.prompt_name,
+                prompt_version_definition=self.prompts_versions[
+                    (request.prompt_name, request.version_number)
+                ],
+            )
+        )
+
+    def DeletePrompt(self, request: ProtoDeletePromptRequest, context):
+        """Delete prompt"""
+        self._raise_on_missing_prompt(prompt_name=request.prompt_name, context=context)
+
+        del self.prompts[request.prompt_name]
+        del self.prompts_next_version[request.prompt_name]
+
+        versions_to_delete: List[int] = [
+            version
+            for (name, version) in self.prompts_versions.keys()
+            if name == request.prompt_name
+        ]
+
+        for v in versions_to_delete:
+            del self.prompts_versions[(request.prompt_name, v)]
+
+        return ProtoDeletePromptResponse()
+
+    def DeletePromptVersion(self, request: ProtoDeletePromptVersionRequest, context):
+        """Delete prompt version"""
+        self._raise_on_missing_prompt(
+            prompt_name=request.prompt_name,
+            context=context,
+        )
+
+        self._raise_on_missing_prompt_version(
+            prompt_name=request.prompt_name,
+            version=request.version_number,
+            context=context,
+        )
+
+        prompt_default_version: ProtoPrompt = self.prompts[request.prompt_name]
+        if (
+            prompt_default_version.default_version_definition.version_number
+            == request.version_number
+        ):
+            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
+            context.set_details("Can't delete prompt default version")
+            raise ValueError("Can't delete prompt default version")
+
+        del self.prompts_versions[(request.prompt_name, request.version_number)]
+        return ProtoDeletePromptVersionResponse()
+
+    def SetDefaultPromptVersion(
+        self, request: ProtoSetDefaultPromptVersionRequest, context
+    ):
+        """Set default prompt version"""
+        self._raise_on_missing_prompt(
+            prompt_name=request.prompt_name,
+            context=context,
+        )
+
+        self._raise_on_missing_prompt_version(
+            prompt_name=request.prompt_name,
+            version=request.version_number,
+            context=context,
+        )
+        existing_prompt: ProtoPrompt = self.prompts[request.prompt_name]
+        prompt_version_definition: ProtoPromptVersionDefinition = self.prompts_versions[
+            (request.prompt_name, request.version_number)
+        ]
+        existing_prompt.default_version_definition.CopyFrom(prompt_version_definition)
+        return ProtoSetDefaultPromptVersionResponse()

qwak_services_mock/mocks/qwak_mocks.py

@@ -15,10 +15,12 @@ from qwak_services_mock.mocks.automation_management_service import (
 )
 from qwak_services_mock.mocks.autoscaling_service_api import AutoscalingServiceApiMock
 from qwak_services_mock.mocks.batch_job_manager_service import BatchJobManagerService
-from qwak_services_mock.mocks.build_management import BuildsManagementServiceMock
 from qwak_services_mock.mocks.build_orchestrator_build_api import (
     BuildOrchestratorBuildApiMock,
 )
+from qwak_services_mock.mocks.build_orchestrator_build_settings_api import (
+    BuildOrchestratorBuildSettingsApiMock,
+)
 from qwak_services_mock.mocks.build_orchestrator_service_api import (
     BuildOrchestratorServiceApiMock,
 )
@@ -43,23 +45,34 @@ from qwak_services_mock.mocks.features_operator_v3_service import (
     FeaturesOperatorV3ServiceMock,
 )
 from qwak_services_mock.mocks.file_versioning_service import FileVersioningServiceMock
+from qwak_services_mock.mocks.fs_offline_serving_service import (
+    FsOfflineServingServiceMock,
+)
 from qwak_services_mock.mocks.instance_template_management_service import (
     InstanceTemplateManagementServiceMock,
 )
+from qwak_services_mock.mocks.integration_management_service import (
+    IntegrationManagementServiceMock,
+)
 from qwak_services_mock.mocks.internal_build_orchestrator_service import (
     InternalBuildOrchestratorServiceMock,
 )
 from qwak_services_mock.mocks.job_registry_service_api import JobRegistryServiceApiMock
 from qwak_services_mock.mocks.kube_captain_service_api import KubeCaptainServiceApiMock
+from qwak_services_mock.mocks.location_discovery_service_api import (
+    LocationDiscoveryServiceApiMock,
+)
 from qwak_services_mock.mocks.logging_service import LoggingServiceApiMock
 from qwak_services_mock.mocks.model_management_service import (
     ModelsManagementServiceMock,
 )
 from qwak_services_mock.mocks.project_manager_service import ProjectManagerServiceMock
+from qwak_services_mock.mocks.prompt_manager_service import PromptManagerServiceMock
 from qwak_services_mock.mocks.secret_service import SecretServiceMock
 from qwak_services_mock.mocks.self_service_user_service import (
     SelfServiceUserServiceMock,
 )
+from qwak_services_mock.mocks.system_secret_service import SystemSecretServiceMock
 from qwak_services_mock.mocks.user_application_instance_service_api import (
     UserApplicationInstanceServiceApiMock,
 )
@@ -74,8 +87,8 @@ class QwakMocks:
     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
@@ -96,6 +109,7 @@ class QwakMocks:
     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
@@ -105,3 +119,7 @@ class QwakMocks:
     vector_serving_service: VectorServingServiceMock
     vector_collection_service: VectorCollectionManagementServiceMock
     execution_management_service: ExecutionManagementServiceMock
+    system_secret_service: SystemSecretServiceMock
+    integration_management_service: IntegrationManagementServiceMock
+    prompt_manager_service: PromptManagerServiceMock
+    location_discovery_service: LocationDiscoveryServiceApiMock

qwak_services_mock/mocks/system_secret_service.py

@@ -0,0 +1,54 @@
+from typing import Dict, Tuple
+
+import grpc
+from _qwak_proto.qwak.admiral.secret.v0.secret_pb2 import (
+    EnvironmentSecretIdentifier,
+    SetSystemEnvironmentSecretOptions,
+    SystemSecretDefinition,
+)
+from _qwak_proto.qwak.admiral.secret.v0.system_secret_service_pb2 import (
+    DeleteSystemSecretRequest,
+    DeleteSystemSecretResponse,
+    GetSystemSecretRequest,
+    GetSystemSecretResponse,
+    SetSystemSecretRequest,
+    SetSystemSecretResponse,
+)
+from _qwak_proto.qwak.admiral.secret.v0.system_secret_service_pb2_grpc import (
+    SystemSecretServiceServicer,
+)
+
+
+class SystemSecretServiceMock(SystemSecretServiceServicer):
+    _secrets: Dict[Tuple[str, str], SystemSecretDefinition] = dict()
+
+    def clear(self):
+        self._secrets.clear()
+
+    def SetSystemSecret(self, request: SetSystemSecretRequest, context):
+        """Set a value for a secret"""
+
+        options: SetSystemEnvironmentSecretOptions = request.options
+        identifier: EnvironmentSecretIdentifier = options.identifier
+        definition: SystemSecretDefinition = SystemSecretDefinition(
+            identifier=identifier, spec=options.spec
+        )
+
+        self._secrets[(identifier.name, identifier.environment_id)] = definition
+        return SetSystemSecretResponse()
+
+    def DeleteSystemSecret(self, request: DeleteSystemSecretRequest, context):
+        """Delete secret"""
+        name: str = request.identifier.name
+        env_id: str = request.identifier.environment_id
+
+        if (name, env_id) in self._secrets:
+            del self._secrets[(name, env_id)]
+
+        return DeleteSystemSecretResponse()
+
+    def GetSystemSecret(self, request: GetSystemSecretRequest, context):
+        """Get secret"""
+        name: str = request.identifier.name
+        env_id: str = request.identifier.environment_id
+        return GetSystemSecretResponse(secret_definition=self._secrets[(name, env_id)])

qwak_services_mock/mocks/vector_serving_api.py

@@ -2,10 +2,20 @@ from dataclasses import dataclass
 from typing import Dict, List
 
 import grpc
-from _qwak_proto.qwak.vectors.v1.vector_pb2 import DoubleVector, Property, SearchResult
+from _qwak_proto.qwak.vectors.v1.vector_pb2 import (
+    DoubleVector,
+    Property,
+    SearchResult,
+    StoredVector,
+    VectorIdentifier,
+)
 from _qwak_proto.qwak.vectors.v1.vector_service_pb2 import (
     DeleteVectorsResponse,
+    FetchVectorRequest,
+    FetchVectorResponse,
+    SearchSimilarVectorsRequest,
     SearchSimilarVectorsResponse,
+    UpsertVectorsRequest,
     UpsertVectorsResponse,
 )
 from _qwak_proto.qwak.vectors.v1.vector_service_pb2_grpc import VectorServiceServicer
@@ -45,7 +55,7 @@ class VectorServingServiceMock(VectorServiceServicer):
         collection = self._vector_collections[collection_name]
         return [collection.get(_id) for _id in vector_ids if _id in collection]
 
-    def SearchSimilarVectors(self, request, context):
+    def SearchSimilarVectors(self, request: SearchSimilarVectorsRequest, context):
         if request.collection_name not in self._vector_collections:
             context.set_details(f"Collection named {request.collection} doesn't exist'")
             context.set_code(grpc.StatusCode.NOT_FOUND)
@@ -81,11 +91,12 @@ class VectorServingServiceMock(VectorServiceServicer):
             ]
         )
 
-    def UpsertVectors(self, request, context):
+    def UpsertVectors(self, request: UpsertVectorsRequest, context):
         collection_store = self._vector_collections.get(request.collection_name, dict())
         for stored_vector in request.vector:
-            collection_store[stored_vector.id] = {
-                "id": stored_vector.id,
+            id = stored_vector.vector_identifier.vector_id
+            collection_store[id] = {
+                "id": id,
                 "vector": list(stored_vector.vector.element),
                 "properties": stored_vector.properties,
             }
@@ -103,9 +114,9 @@ class VectorServingServiceMock(VectorServiceServicer):
 
         collection_store = self._vector_collections[request.collection_name]
         ids_in_collection = [
-            vector_id
-            for vector_id in request.vector_id
-            if vector_id in collection_store
+            vector_identifier.vector_id
+            for vector_identifier in request.vector_identifiers
+            if vector_identifier.vector_id in collection_store
         ]
         for id in ids_in_collection:
             collection_store.pop(id)
@@ -113,6 +124,31 @@ class VectorServingServiceMock(VectorServiceServicer):
         self._vector_collections[request.collection_name] = collection_store
         return DeleteVectorsResponse(num_vectors_deleted=len(ids_in_collection))
 
+    def FetchVector(self, request, context):
+        if request.collection_name not in self._vector_collections:
+            context.set_details(
+                f"Collection named {request.collection_name} doesn't exist'"
+            )
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return
+
+        collection_store = self._vector_collections[request.collection_name]
+        vector_identifier = request.vector_identifier
+        vector_id = vector_identifier.vector_id
+        if vector_id not in collection_store:
+            context.set_details(f"Vector with id {vector_id} doesn't exist'")
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return
+
+        stored_vector_dict = collection_store[vector_id]
+        return FetchVectorResponse(
+            vector=StoredVector(
+                vector_identifier=request.vector_identifier,
+                vector=DoubleVector(element=stored_vector_dict["vector"]),
+                properties=stored_vector_dict["properties"],
+            )
+        )
+
 
 def _cos_sim(a, b):
     return dot(a, b) / (norm(a) * norm(b))

qwak_services_mock/mocks/workspace_manager_service_mock.py

@@ -3,20 +3,27 @@ import uuid
 from dataclasses import dataclass
 
 import grpc
-
 from _qwak_proto.qwak.user_application.common.v0.resources_pb2 import (
     ClientPodComputeResources,
     PodComputeResourceTemplateSpec,
 )
-from _qwak_proto.qwak.workspace.workspace_pb2 import Workspace, WorkspaceImage, DefaultWorkspaceDetails
+from _qwak_proto.qwak.workspace.workspace_pb2 import (
+    DefaultWorkspaceDetails,
+    Workspace,
+    WorkspaceImage,
+)
 from _qwak_proto.qwak.workspace.workspace_service_pb2 import (
+    CountWorkspacesRequest,
+    CountWorkspacesResponse,
     CreateWorkspaceResponse,
     DeleteWorkspaceResponse,
+    GetDefaultWorkspaceValuesRequest,
+    GetDefaultWorkspaceValuesResponse,
     GetWorkspaceByIdResponse,
     ListWorkspaceImagesResponse,
     ListWorkspacesResponse,
-    UpdateWorkspaceResponse, CountWorkspacesRequest, CountWorkspacesResponse, GetDefaultWorkspaceValuesRequest,
-    GetDefaultWorkspaceValuesResponse, )
+    UpdateWorkspaceResponse,
+)
 from _qwak_proto.qwak.workspace.workspace_service_pb2_grpc import (
     WorkspaceManagementServiceServicer,
 )
@@ -28,6 +35,7 @@ DEFAULT_GPU_IMAGE = "3"
 DEFAULT_CPU_INSTANCE = "small"
 DEFAULT_GPU_INSTANCE = "g5_xlarge"
 
+
 @dataclass
 class WorkspaceData:
     id: str
@@ -163,19 +171,27 @@ class WorkspaceManagerServiceMock(WorkspaceManagementServiceServicer):
             workspace_images.append(WorkspaceImage(name=v, id=k))
         return ListWorkspaceImagesResponse(workspace_images=workspace_images)
 
-    def CountWorkspaces(self, request: CountWorkspacesRequest, context) -> CountWorkspacesResponse:
+    def CountWorkspaces(
+        self, request: CountWorkspacesRequest, context
+    ) -> CountWorkspacesResponse:
         return CountWorkspacesResponse(count=len(self.workspaces))
 
-    def GetDefaultWorkspaceValues(self, request: GetDefaultWorkspaceValuesRequest, context) -> GetDefaultWorkspaceValuesResponse:
+    def GetDefaultWorkspaceValues(
+        self, request: GetDefaultWorkspaceValuesRequest, context
+    ) -> GetDefaultWorkspaceValuesResponse:
         return GetDefaultWorkspaceValuesResponse(
             default_workspace_details=DefaultWorkspaceDetails(
                 cpu_image_id=DEFAULT_CPU_IMAGE,
                 gpu_image_id=DEFAULT_GPU_IMAGE,
                 cpu_compute_resources=ClientPodComputeResources(
-                    template_spec=PodComputeResourceTemplateSpec(template_id=DEFAULT_CPU_INSTANCE)
+                    template_spec=PodComputeResourceTemplateSpec(
+                        template_id=DEFAULT_CPU_INSTANCE
+                    )
                 ),
                 gpu_compute_resources=ClientPodComputeResources(
-                    template_spec=PodComputeResourceTemplateSpec(template_id=DEFAULT_GPU_INSTANCE)
+                    template_spec=PodComputeResourceTemplateSpec(
+                        template_id=DEFAULT_GPU_INSTANCE
+                    )
                 ),
             )
         )