truefoundry 0.5.9__py3-none-any.whl → 0.5.10__py3-none-any.whl

truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -4,7 +4,7 @@ import json
 import os
 import time
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, cast
 from urllib.parse import urljoin
 
 import requests
@@ -24,6 +24,7 @@ from truefoundry.common.servicefoundry_client import (
     session_with_retries,
 )
 from truefoundry.common.session import Session
+from truefoundry.common.utils import get_user_agent
 from truefoundry.deploy.auto_gen import models as auto_gen_models
 from truefoundry.deploy.io.output_callback import OutputCallBack
 from truefoundry.deploy.lib.model.entity import (
@@ -32,6 +33,8 @@ from truefoundry.deploy.lib.model.entity import (
     Deployment,
     DockerRegistryCredentials,
     JobRun,
+    LogBody,
+    SocketEvent,
     TriggerJobResult,
     Workspace,
     WorkspaceResources,
@@ -51,7 +54,7 @@ BUILD_LOGS_SUBSCRIBE_MESSAGE = "BUILD_LOGS"
 MAX_RETRIES_WORKFLOW_TRIGGER = 3
 
 
-def _upload_packaged_code(metadata, package_file):
+def _upload_packaged_code(metadata: Dict[str, Any], package_file: str) -> None:
     file_size = os.stat(package_file).st_size
     with open(package_file, "rb") as file_to_upload:
         with tqdm(
@@ -74,7 +77,9 @@ def _upload_packaged_code(metadata, package_file):
 
 
 class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
-    def __init__(self, init_session: bool = True, tfy_host: Optional[str] = None):
+    def __init__(
+        self, init_session: bool = True, tfy_host: Optional[str] = None
+    ) -> None:
         self._session: Optional[Session] = None
         if init_session:
             if tfy_host:
@@ -85,25 +90,23 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             raise Exception("Neither session, not tfy_host provided")
         super().__init__(tfy_host=tfy_host)
 
-    def _get_header(self):
+    def _get_headers(self) -> Dict[str, str]:
+        headers = {"User-Agent": get_user_agent()}
         if not self._session:
-            return {}
-        return {"Authorization": f"Bearer {self._session.access_token}"}
+            return headers
+        return {
+            **headers,
+            "Authorization": f"Bearer {self._session.access_token}",
+        }
 
     @check_min_cli_version
-    def get_id_from_fqn(self, fqn_type: str, fqn: str):
+    def get_id_from_fqn(self, fqn_type: str, fqn: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/fqn/{fqn_type}"
         response = session_with_retries().get(
-            url, headers=self._get_header(), params={"fqn": fqn}
+            url, headers=self._get_headers(), params={"fqn": fqn}
         )
         return request_handling(response)
 
-    @check_min_cli_version
-    def list_workspace(self):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace"
-        response = session_with_retries().get(url, headers=self._get_header())
-        return request_handling(response)
-
     @check_min_cli_version
     def list_workspaces(
         self,
@@ -120,7 +123,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if workspace_fqn:
             params["workspaceFqn"] = workspace_fqn
         response = session_with_retries().get(
-            url, params=params, headers=self._get_header()
+            url, params=params, headers=self._get_headers()
         )
         response = request_handling(response)
         return parse_obj_as(List[Workspace], response)
@@ -142,75 +145,44 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                     "resources": resources.dict(exclude_none=True),
                 }
             },
-            headers=self._get_header(),
+            headers=self._get_headers(),
         )
         response_data = request_handling(response)
         return Workspace.parse_obj(response_data)
 
     @check_min_cli_version
-    def remove_workspace(self, workspace_id, force=False) -> Workspace:
+    def remove_workspace(self, workspace_id: str, force: bool = False) -> Workspace:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace/{workspace_id}"
-        force = json.dumps(
+        force_str = json.dumps(
             force
         )  # this dumb conversion is required because `params` just casts as str
         response = session_with_retries().delete(
-            url, headers=self._get_header(), params={"force": force}
+            url, headers=self._get_headers(), params={"force": force_str}
         )
-        response = request_handling(response)
+        response = cast(Dict[str, Any], request_handling(response))
         return Workspace.parse_obj(response["workspace"])
 
-    @check_min_cli_version
-    def get_workspace_by_name(self, workspace_name, cluster_id):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace"
-        response = session_with_retries().get(
-            url,
-            headers=self._get_header(),
-            params={"name": workspace_name, "clusterId": cluster_id},
-        )
-        return request_handling(response)
-
-    @check_min_cli_version
-    def get_workspace(self, workspace_id):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace/{workspace_id}"
-        response = session_with_retries().get(url, headers=self._get_header())
-        return request_handling(response)
-
     @check_min_cli_version
     def get_workspace_by_fqn(self, workspace_fqn: str) -> List[Workspace]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace"
         response = session_with_retries().get(
             url,
-            headers=self._get_header(),
+            headers=self._get_headers(),
             params={"fqn": workspace_fqn},
         )
         response = request_handling(response)
         return parse_obj_as(List[Workspace], response)
 
     @check_min_cli_version
-    def list_deployments(self, workspace_id: str = None):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/deployment"
-        params = {}
-        if workspace_id:
-            params["workspaceId"] = workspace_id
-        response = session_with_retries().get(
-            url=url, params=params, headers=self._get_header()
-        )
-        return request_handling(response)
-
-    @check_min_cli_version
-    def list_cluster(self):
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/cluster"
-        response = session_with_retries().get(url, headers=self._get_header())
-        return request_handling(response)
-
-    @check_min_cli_version
-    def get_cluster(self, cluster_id):
+    def get_cluster(self, cluster_id: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/cluster/{cluster_id}"
-        response = session_with_retries().get(url, headers=self._get_header())
-        return request_handling(response)
+        response = session_with_retries().get(url, headers=self._get_headers())
+        return cast(Dict[str, Any], request_handling(response))
 
     @check_min_cli_version
-    def get_presigned_url(self, space_name, service_name, env):
+    def get_presigned_url(
+        self, space_name: str, service_name: str, env: str
+    ) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/deployment/code-upload-url"
         response = session_with_retries().post(
             url,
@@ -219,9 +191,9 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 "serviceName": service_name,
                 "stage": env,
             },
-            headers=self._get_header(),
+            headers=self._get_headers(),
         )
-        return request_handling(response)
+        return cast(Dict[str, Any], request_handling(response))
 
     @check_min_cli_version
     def upload_code_package(
@@ -231,7 +203,6 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             space_name=workspace_fqn, service_name=component_name, env="default"
         )
         _upload_packaged_code(metadata=http_response, package_file=package_local_path)
-
         return http_response["uri"]
 
     @check_min_cli_version
@@ -250,30 +221,30 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         logger.debug(json.dumps(data))
         url = f"{self._api_server_url}/{VERSION_PREFIX}/deployment"
         response = session_with_retries().post(
-            url, json=data, headers=self._get_header()
+            url, json=data, headers=self._get_headers()
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return Deployment.parse_obj(response_data["deployment"])
 
-    def _get_log_print_line(self, log: dict):
-        timestamp = int(log["time"]) / 1e6
+    def _get_log_print_line(self, log_data: LogBody) -> str:
+        timestamp = int(log_data.time) / 1e6
 
         time_obj = datetime.fromtimestamp(timestamp / 1000.0, tz=timezone.utc)
         time_obj.replace(tzinfo=timezone.utc)
         local_time = time_obj.astimezone(tzlocal())
         local_time_str = local_time.isoformat()
-        return f"[{local_time_str}] {log['log'].strip()}"
+        return f"[{local_time_str}] {log_data.log.strip()}"
 
     def _tail_logs(
         self,
         tail_logs_url: str,
-        query_dict: dict,
+        query_dict: Dict[str, Any],
         # NOTE: Rather making this printer callback an argument,
         # we should have global printer callback
         # which will be initialized based on the running env (cli, lib, notebook)
         subscribe_message: str,
         socketio_path: str = "socket.io",
-        callback=None,
+        callback: Optional[OutputCallBack] = None,
     ) -> socketio.Client:
         callback = callback or OutputCallBack()
         sio = socketio.Client(request_timeout=60)
@@ -281,18 +252,20 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         next_log_start_timestamp = query_dict.get("startTs")
 
         @sio.on(subscribe_message)
-        def logs(data):
+        def logs(data: Any) -> None:
             try:
-                _log = json.loads(data)
-                body = _log["body"]
-                callback.print_line(self._get_log_print_line(body))
+                event = SocketEvent.parse_obj(json.loads(data))
+                if not isinstance(event.body, LogBody):
+                    logger.debug(f"Skipped log for {data!r}")
+                    return
+                callback.print_line(self._get_log_print_line(event.body))
                 nonlocal next_log_start_timestamp
-                next_log_start_timestamp = body["time"]
+                next_log_start_timestamp = event.body.time
             except Exception:
-                logger.exception(f"Error while parsing log line, {data!r}")
+                logger.debug(f"Error while parsing log line, {data!r}")
 
         @sio.on("connect")
-        def on_connect():
+        def on_connect() -> None:
             # TODO: We should have have a timeout here. `emit` does
             #   not support timeout. Explore `sio.call`.
             query_dict["startTs"] = next_log_start_timestamp
@@ -301,7 +274,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 json.dumps(query_dict),
             )
 
-        def sio_disconnect_no_exception():
+        def sio_disconnect_no_exception() -> None:
             try:
                 sio.disconnect()
             except Exception:
@@ -311,7 +284,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             sio.connect(
                 tail_logs_url,
                 transports="websocket",
-                headers=self._get_header(),
+                headers=self._get_headers(),
                 socketio_path=socketio_path,
             )
             return sio
@@ -319,7 +292,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     @check_min_cli_version
     def get_deployment(self, application_id: str, deployment_id: str) -> Deployment:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}/deployments/{deployment_id}"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response_data = request_handling(response)
         return Deployment.parse_obj(response_data)
 
@@ -328,7 +301,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         self, application_id: str, deployment_id: str
     ) -> List[AppDeploymentStatusResponse]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}/deployments/{deployment_id}/statuses"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response_data = request_handling(response)
         return parse_obj_as(List[AppDeploymentStatusResponse], response_data)
 
@@ -337,7 +310,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         self, application_id: str, deployment_id: str
     ) -> List[BuildResponse]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}/deployments/{deployment_id}/builds"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response_data = request_handling(response)
         return parse_obj_as(List[BuildResponse], response_data)
 
@@ -351,11 +324,11 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         end_ts_nano: Optional[int] = None,
         limit: Optional[int] = None,
         num_logs_to_ignore: Optional[int] = None,
-    ) -> List:
+    ) -> List[LogBody]:
         get_logs_query = {"applicationId": application_id}
         if deployment_id:
             get_logs_query["deploymentId"] = deployment_id
-        data = {"getLogsQuery": json.dumps(get_logs_query)}
+        data: Dict[str, Any] = {"getLogsQuery": json.dumps(get_logs_query)}
         if start_ts_nano:
             data["startTs"] = str(start_ts_nano)
         if end_ts_nano:
@@ -369,22 +342,23 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
 
         url = f"{self._api_server_url}/{VERSION_PREFIX}/logs/{workspace_id}"
         response = session_with_retries().get(
-            url=url, params=data, headers=self._get_header()
+            url=url, params=data, headers=self._get_headers()
         )
-        response_data = request_handling(response)
-        return list(response_data["logs"])
+        response_data = cast(Dict[str, Any], request_handling(response))
+        return parse_obj_as(List[LogBody], response_data["data"])
 
     @check_min_cli_version
     def tail_build_logs(
         self,
         build_response: BuildResponse,
-        callback=None,
+        callback: Optional[OutputCallBack] = None,
     ) -> socketio.Client:
         callback = callback or OutputCallBack()
         tail_logs_obj = json.loads(build_response.tailLogsUrl)
         socket = self._tail_logs(
             tail_logs_url=urljoin(
-                tail_logs_obj["uri"], f"/?type={BUILD_LOGS_SUBSCRIBE_MESSAGE}"
+                tail_logs_obj["uri"],
+                f"/?type={BUILD_LOGS_SUBSCRIBE_MESSAGE}",
             ),
             socketio_path=tail_logs_obj["path"],
             query_dict={
@@ -404,9 +378,9 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         deployment_id: str,
         start_ts: int,
         limit: int,
-        callback=None,
+        callback: Optional[OutputCallBack] = None,
         wait: bool = True,
-    ):
+    ) -> None:
         callback = callback or OutputCallBack()
         self._tail_logs(
             tail_logs_url=urljoin(
@@ -435,8 +409,8 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         start_ts: int,
         limit: int,
         poll_interval_seconds: int,
-        callback=None,
-    ):
+        callback: Optional[OutputCallBack] = None,
+    ) -> None:
         callback = callback or OutputCallBack()
         start_ts_nano = int(start_ts * 1e6)
 
@@ -444,7 +418,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             num_logs_to_ignore = 0
 
             while True:
-                logs_list = self._get_deployment_logs(
+                logs = self._get_deployment_logs(
                     workspace_id=workspace_id,
                     application_id=application_id,
                     deployment_id=deployment_id,
@@ -454,18 +428,18 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                     num_logs_to_ignore=num_logs_to_ignore,
                 )
 
-                if len(logs_list) == 0:
+                if not logs:
                     logger.warning("Did not receive any logs")
                     time.sleep(poll_interval_seconds)
                     continue
 
-                for log in logs_list:
+                for log in logs:
                     callback.print_line(self._get_log_print_line(log))
 
-                last_log_time = logs_list[-1]["time"]
+                last_log_time = logs[-1].time
                 num_logs_to_ignore = 0
-                for log in reversed(logs_list):
-                    if log["time"] != last_log_time:
+                for log in reversed(logs):
+                    if log.time != last_log_time:
                         break
                     num_logs_to_ignore += 1
 
@@ -483,40 +457,43 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         start_ts: Optional[int],
         end_ts: Optional[int],
         limit: Optional[int],
-        callback=None,
-    ):
+        callback: Optional[OutputCallBack] = None,
+    ) -> None:
         callback = callback or OutputCallBack()
-        logs_list = self._get_deployment_logs(
+        logs = self._get_deployment_logs(
             workspace_id=workspace_id,
             application_id=application_id,
             deployment_id=deployment_id,
             job_run_name=job_run_name,
-            start_ts_nano=int(start_ts * 1e6),
-            end_ts_nano=int(end_ts * 1e6),
+            start_ts_nano=int(start_ts * 1e6) if start_ts else None,
+            end_ts_nano=int(end_ts * 1e6) if end_ts else None,
             limit=limit,
         )
-        for log in logs_list:
+        for log in logs:
             callback.print_line(self._get_log_print_line(log))
 
     @check_min_cli_version
     def fetch_build_logs(
         self,
         build_response: BuildResponse,
-        callback=None,
+        callback: Optional[OutputCallBack] = None,
     ) -> None:
         callback = callback or OutputCallBack()
         url = build_response.getLogsUrl
-        response = session_with_retries().get(url=url, headers=self._get_header())
-        logs_list = request_handling(response)
-        for log in logs_list["logs"]:
-            # TODO: Have to establish a log line format that includes timestamp, level, message
-            callback.print_line(self._get_log_print_line(log))
+        response = session_with_retries().get(url=url, headers=self._get_headers())
+        logs_response = cast(Dict[str, Any], request_handling(response))
+        for _log_body in logs_response["logs"]:
+            try:
+                log_body = LogBody.parse_obj(_log_body)
+                callback.print_line(self._get_log_print_line(log_body))
+            except Exception:
+                logger.debug(f"Failed to parse log body {_log_body}")
 
     @check_min_cli_version
     def get_deployment_info_by_fqn(self, deployment_fqn: str) -> DeploymentFqnResponse:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/fqn/deployment"
         response = session_with_retries().get(
-            url, headers=self._get_header(), params={"fqn": deployment_fqn}
+            url, headers=self._get_headers(), params={"fqn": deployment_fqn}
         )
         response_data = request_handling(response)
         return DeploymentFqnResponse.parse_obj(response_data)
@@ -527,23 +504,23 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     ) -> ApplicationFqnResponse:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/fqn/app"
         response = session_with_retries().get(
-            url, headers=self._get_header(), params={"fqn": application_fqn}
+            url, headers=self._get_headers(), params={"fqn": application_fqn}
         )
         response_data = request_handling(response)
         return ApplicationFqnResponse.parse_obj(response_data)
 
     @check_min_cli_version
-    def remove_application(self, application_id: str):
+    def remove_application(self, application_id: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}"
-        response = session_with_retries().delete(url, headers=self._get_header())
-        response = request_handling(response)
+        response = session_with_retries().delete(url, headers=self._get_headers())
+        response = cast(Dict[str, Any], request_handling(response))
         # TODO: Add pydantic here.
         return response
 
     @check_min_cli_version
     def get_application_info(self, application_id: str) -> Application:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response = request_handling(response)
         return Application.parse_obj(response)
 
@@ -563,18 +540,18 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if search_prefix:
             params["searchPrefix"] = search_prefix
         response = session_with_retries().get(
-            url, headers=self._get_header(), params=params
+            url, headers=self._get_headers(), params=params
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return parse_obj_as(List[JobRun], response_data["data"])
 
     def get_job_run(
         self,
         application_id: str,
         job_run_name: str,
-    ):
+    ) -> JobRun:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/jobs/{application_id}/runs/{job_run_name}"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response_data = request_handling(response)
         return parse_obj_as(JobRun, response_data)
 
@@ -594,18 +571,20 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if params:
             body["input"]["params"] = params
         response = session_with_retries().post(
-            url, json=body, headers=self._get_header()
+            url, json=body, headers=self._get_headers()
         )
         response = request_handling(response)
         return TriggerJobResult.parse_obj(response)
 
-    def trigger_workflow(self, application_id: str, inputs: Dict[str, Any]):
+    def trigger_workflow(
+        self, application_id: str, inputs: Dict[str, Any]
+    ) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/workflow/{application_id}/executions"
         body = {"inputs": inputs}
         response = session_with_retries(retries=MAX_RETRIES_WORKFLOW_TRIGGER).post(
-            url, json=body, headers=self._get_header()
+            url, json=body, headers=self._get_headers()
         )
-        response = request_handling(response)
+        response = cast(Dict[str, Any], request_handling(response))
         return response
 
     @check_min_cli_version
@@ -615,7 +594,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/docker-registry/creds"
         response = session_with_retries().get(
             url,
-            headers=self._get_header(),
+            headers=self._get_headers(),
             params={
                 "fqn": docker_registry_fqn,
                 "clusterId": cluster_id,
@@ -631,7 +610,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/docker-registry/create-repo"
         response = session_with_retries().post(
             url,
-            headers=self._get_header(),
+            headers=self._get_headers(),
             data={
                 "fqn": docker_registry_fqn,
                 "workspaceFqn": workspace_fqn,
@@ -657,7 +636,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if application_name:
             params["applicationName"] = application_name
         response = session_with_retries().get(
-            url, params=params, headers=self._get_header()
+            url, params=params, headers=self._get_headers()
         )
         response = request_handling(response)
         return parse_obj_as(List[Application], response)
@@ -678,36 +657,36 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if deployment_id:
             params["deploymentId"] = deployment_id
         response = session_with_retries().get(
-            url, params=params, headers=self._get_header()
+            url, params=params, headers=self._get_headers()
         )
         response = request_handling(response)
         return parse_obj_as(List[Deployment], response)
 
     @check_min_cli_version
-    def apply(self, manifest: Dict[str, Any], dry_run: bool = False):
+    def apply(self, manifest: Dict[str, Any], dry_run: bool = False) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/apply"
         body = {"manifest": manifest, "dryRun": dry_run}
         response = session_with_retries().put(
-            url, headers=self._get_header(), json=body
+            url, headers=self._get_headers(), json=body
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return response_data
 
     @check_min_cli_version
-    def delete(self, manifest: Dict[str, Any]):
+    def delete(self, manifest: Dict[str, Any]) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/delete"
         body = {"manifest": manifest}
         response = session_with_retries().post(
-            url, headers=self._get_header(), json=body
+            url, headers=self._get_headers(), json=body
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return response_data
 
     def terminate_job_run(
         self,
         deployment_id: str,
         job_run_name: str,
-    ):
+    ) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/jobs/terminate"
         body = {
             "deploymentId": deployment_id,
@@ -718,7 +697,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             # TODO (chiragjn): Check if this is supposed to be params or json
             params=body,
             json=body,
-            headers=self._get_header(),
+            headers=self._get_headers(),
         )
-        response_data = request_handling(response)
+        response_data = cast(Dict[str, Any], request_handling(response))
         return response_data

truefoundry/deploy/lib/dao/apply.py

@@ -6,7 +6,7 @@ import yaml
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
-from truefoundry.deploy.lib.model.entity import ApplyResult, Manifest
+from truefoundry.deploy.lib.model.entity import ApplyResult, ManifestLike
 from truefoundry.pydantic_v1 import ValidationError
 
 
@@ -26,7 +26,7 @@ def _apply_manifest(
         file_metadata += f" from file {filename}"
 
     try:
-        manifest = Manifest.parse_obj(manifest)
+        manifest = ManifestLike.parse_obj(manifest)
     except ValidationError as ex:
         return ApplyResult(
             success=False,

truefoundry/deploy/lib/dao/delete.py

@@ -6,7 +6,7 @@ import yaml
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
-from truefoundry.deploy.lib.model.entity import DeleteResult, Manifest
+from truefoundry.deploy.lib.model.entity import DeleteResult, ManifestLike
 from truefoundry.pydantic_v1 import ValidationError
 
 
@@ -25,7 +25,7 @@ def _delete_manifest(
         file_metadata += f" from file {filename}"
 
     try:
-        manifest = Manifest.parse_obj(manifest)
+        manifest = ManifestLike.parse_obj(manifest)
     except ValidationError as ex:
         return DeleteResult(
             success=False,