Adding endpoint override feature

Author: kumar-shivam-ranjan

ads/aqua/common/enums.py

@@ -20,6 +20,12 @@ class Resource(ExtendedEnum):
     MODEL_VERSION_SET = "model-version-sets"
 
 
+class PredictEndpoints(ExtendedEnum):
+    CHAT_COMPLETIONS_ENDPOINT = "/v1/chat/completions"
+    TEXT_COMPLETIONS_ENDPOINT = "/v1/completions"
+    EMBEDDING_ENDPOINT = "/v1/embedding"
+
+
 class Tags(ExtendedEnum):
     TASK = "task"
     LICENSE = "license"

ads/aqua/extension/deployment_handler.py

@@ -176,7 +176,7 @@ def list_shapes(self):
 
 class AquaDeploymentStreamingInferenceHandler(AquaAPIhandler):
     @handle_exceptions
-    async def post(self, model_deployment_id):
+    def post(self, model_deployment_id):
         """
         Handles streaming inference request for the Active Model Deployments
         Raises
@@ -201,21 +201,18 @@ async def post(self, model_deployment_id):
             )
         if not input_data.get("model"):
             raise HTTPError(400, Errors.MISSING_REQUIRED_PARAMETER.format("model"))
-
-        if "stream" not in input_data:
-            input_data.update(stream=True)
-
+        route_override_header = self.request.headers.get("route", None)
         self.set_header("Content-Type", "text/event-stream")
         self.set_header("Cache-Control", "no-cache")
         self.set_header("Transfer-Encoding", "chunked")
-        await self.flush()
+        self.flush()
         try:
             response_gen = AquaDeploymentApp().get_model_deployment_response(
-                model_deployment_id, input_data
+                model_deployment_id, input_data, route_override_header
             )
             for chunk in response_gen:
                 self.write(chunk)
-                await self.flush()
+                self.flush()
         except Exception as ex:
             raise HTTPError(500, str(ex)) from ex
         finally:

ads/aqua/modeldeployment/deployment.py

@@ -18,7 +18,12 @@
     ComputeShapeSummary,
     ContainerPath,
 )
-from ads.aqua.common.enums import InferenceContainerTypeFamily, ModelFormat, Tags
+from ads.aqua.common.enums import (
+    InferenceContainerTypeFamily,
+    ModelFormat,
+    PredictEndpoints,
+    Tags,
+)
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     DEFINED_METADATA_TO_FILE_MAP,
@@ -937,7 +942,6 @@ def get(self, model_deployment_id: str, **kwargs) -> "AquaDeploymentDetail":
         model_deployment = self.ds_client.get_model_deployment(
             model_deployment_id=model_deployment_id, **kwargs
         ).data
-
         oci_aqua = (
             (
                 Tags.AQUA_TAG in model_deployment.freeform_tags
@@ -982,7 +986,6 @@ def get(self, model_deployment_id: str, **kwargs) -> "AquaDeploymentDetail":
         aqua_deployment = AquaDeployment.from_oci_model_deployment(
             model_deployment, self.region
         )
-
         if Tags.MULTIMODEL_TYPE_TAG in model_deployment.freeform_tags:
             aqua_model_id = model_deployment.freeform_tags.get(
                 Tags.AQUA_MODEL_ID_TAG, UNKNOWN
@@ -1013,7 +1016,6 @@ def get(self, model_deployment_id: str, **kwargs) -> "AquaDeploymentDetail":
             aqua_deployment.models = [
                 AquaMultiModelRef(**metadata) for metadata in multi_model_metadata
             ]
-
         return AquaDeploymentDetail(
             **vars(aqua_deployment),
             log_group=AquaResourceIdentifier(
@@ -1334,7 +1336,9 @@ def _stream_sanitizer(response):
                 continue
 
     @telemetry(entry_point="plugin=inference&action=get_response", name="aqua")
-    def get_model_deployment_response(self, model_deployment_id: str, payload: dict):
+    def get_model_deployment_response(
+        self, model_deployment_id: str, payload: dict, route_override_header: str
+    ):
         """
         Returns Model deployment inference response in streaming fashion
 
@@ -1365,29 +1369,34 @@ def get_model_deployment_response(self, model_deployment_id: str, payload: dict)
         model_deployment = self.get(model_deployment_id)
         endpoint = model_deployment.endpoint + "/predictWithResponseStream"
         endpoint_type = model_deployment.environment_variables.get(
-            "MODEL_DEPLOY_PREDICT_ENDPOINT", "/v1/completions"
+            "MODEL_DEPLOY_PREDICT_ENDPOINT", PredictEndpoints.TEXT_COMPLETIONS_ENDPOINT
         )
         aqua_client = Client(endpoint=endpoint)
-        if endpoint_type == "/v1/completions":
-            for chunk in aqua_client.generate(
-                prompt=payload.pop("prompt"),
+
+        if PredictEndpoints.CHAT_COMPLETIONS_ENDPOINT in (
+            endpoint_type,
+            route_override_header,
+        ):
+            for chunk in aqua_client.chat(
+                messages=payload.pop("messages"),
                 payload=payload,
                 stream=True,
             ):
                 try:
-                    yield chunk["choices"][0]["text"]
+                    yield chunk["choices"][0]["delta"]["content"]
                 except Exception as e:
                     logger.debug(
                         f"Exception occurred while parsing streaming response: {e}"
                     )
-        else:
-            for chunk in aqua_client.chat(
-                messages=payload.pop("messages"),
+
+        elif endpoint_type == PredictEndpoints.TEXT_COMPLETIONS_ENDPOINT:
+            for chunk in aqua_client.generate(
+                prompt=payload.pop("prompt"),
                 payload=payload,
                 stream=True,
             ):
                 try:
-                    yield chunk["choices"][0]["delta"]["content"]
+                    yield chunk["choices"][0]["text"]
                 except Exception as e:
                     logger.debug(
                         f"Exception occurred while parsing streaming response: {e}"