Adding aqua client

Author: kumar-shivam-ranjan

ads/aqua/extension/deployment_handler.py

@@ -5,7 +5,6 @@
 from typing import List, Union
 from urllib.parse import urlparse
 
-from tornado.iostream import StreamClosedError
 from tornado.web import HTTPError
 
 from ads.aqua.common.decorator import handle_exceptions
@@ -215,11 +214,9 @@ async def post(self, model_deployment_id):
                 model_deployment_id, input_data
             )
             for chunk in response_gen:
-                if not chunk:
-                    continue
-                self.write(f"data: {chunk}\n\n")
+                self.write(chunk)
                 await self.flush()
-        except StreamClosedError as ex:
+        except Exception as ex:
             raise HTTPError(500, str(ex)) from ex
         finally:
             self.finish()

ads/aqua/modeldeployment/deployment.py

@@ -11,6 +11,7 @@
 from oci.data_science.models import ModelDeploymentShapeSummary
 from pydantic import ValidationError
 
+from ads.aqua import Client
 from ads.aqua.app import AquaApp, logger
 from ads.aqua.common.entities import (
     AquaMultiModelRef,
@@ -1361,7 +1362,33 @@ def get_model_deployment_response(self, model_deployment_id: str, payload: dict)
 
         """
 
-        response = self.model_deployment_client.predict_with_response_stream(
-            model_deployment_id=model_deployment_id, request_body=payload
+        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"
         )
-        yield from self._stream_sanitizer(response)
+        aqua_client = Client(endpoint=endpoint)
+        if endpoint_type == "/v1/completions":
+            for chunk in aqua_client.generate(
+                prompt=payload.pop("prompt"),
+                payload=payload,
+                stream=True,
+            ):
+                try:
+                    yield chunk["choices"][0]["text"]
+                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"),
+                payload=payload,
+                stream=True,
+            ):
+                try:
+                    yield chunk["choices"][0]["delta"]["content"]
+                except Exception as e:
+                    logger.debug(
+                        f"Exception occurred while parsing streaming response: {e}"
+                    )