added test cases and PR review comments

Author: agrimk

ads/aqua/extension/deployment_handler.py

@@ -15,6 +15,7 @@
 from ads.aqua.extension.errors import Errors
 from ads.aqua.modeldeployment import AquaDeploymentApp
 from ads.config import COMPARTMENT_OCID
+from ads.aqua import logger
 
 
 class AquaDeploymentHandler(AquaAPIhandler):
@@ -222,7 +223,36 @@ def list_shapes(self):
 
 class AquaDeploymentStreamingInferenceHandler(AquaAPIhandler):
 
-    def _extract_text_from_choice(self, choice):
+    def _extract_text_from_choice(self, choice: dict) -> str:
+        """
+        Extract text content from a single choice structure.
+
+        Handles both dictionary-based API responses and object-based SDK responses.
+        For dict choices, it checks delta-based streaming fields, message-based
+        non-streaming fields, and finally top-level text/content keys.
+        For object choices, it inspects `.delta`, `.message`, and top-level
+        `.text` or `.content` attributes.
+
+        Parameters
+        ----------
+        choice : dict
+            A choice entry from a model response. It may be:
+                - A dict originating from a JSON API response (streaming or non-streaming).
+                - An SDK-style object with attributes such as `delta`, `message`,
+                `text`, or `content`.
+
+                For dicts, the method checks:
+                    • delta → content/text  
+                    • message → content/text  
+                    • top-level → text/content  
+
+                For objects, the method checks the same fields via attributes.
+
+        Returns
+        -------
+        str | None:
+            The extracted text if present; otherwise None.
+        """
         # choice may be a dict or an object
         if isinstance(choice, dict):
             # streaming chunk: {"delta": {"content": "..."}}
@@ -246,7 +276,31 @@ def _extract_text_from_choice(self, choice):
             return getattr(msg, "content", None) or getattr(msg, "text", None)
         return getattr(choice, "text", None) or getattr(choice, "content", None)
 
-    def _extract_text_from_chunk(self, chunk):
+    def _extract_text_from_chunk(self, chunk: dict) -> str :
+        """
+        Extract text content from a model response chunk.
+
+        Supports both dict-form chunks (streaming or non-streaming) and SDK-style
+        object chunks. When choices are present, extraction is delegated to
+        `_extract_text_from_choice`. If no choices exist, top-level text/content
+        fields or attributes are used.
+
+        Parameters
+        ----------
+        chunk : dict
+            A chunk returned from a model stream or full response. It may be:
+            - A dict containing a `choices` list or top-level text/content fields.
+            - An SDK-style object with a `choices` attribute or top-level
+              `text`/`content` attributes.
+
+            If `choices` is present, the method extracts text from the first
+            choice using `_extract_text_from_choice`. Otherwise, it falls back
+            to top-level text/content.
+        Returns
+        -------
+        str
+            The extracted text if present; otherwise None.
+        """
         if chunk : 
             if isinstance(chunk, dict):
                 choices = chunk.get("choices") or []
@@ -311,6 +365,13 @@ def _get_model_deployment_response(
 
         model_deployment = AquaDeploymentApp().get(model_deployment_id)
         endpoint = model_deployment.endpoint + "/predictWithResponseStream/v1"
+
+        required_keys = ["endpoint_type", "prompt", "model"]
+        missing = [k for k in required_keys if k not in payload]
+
+        if missing:
+            raise HTTPError(400, f"Missing required payload keys: {', '.join(missing)}")
+        
         endpoint_type = payload["endpoint_type"]
         aqua_client = OpenAI(base_url=endpoint)
 
@@ -381,7 +442,7 @@ def _get_model_deployment_response(
                                     {"type": "text", "text": payload["prompt"]},
                                     {
                                         "type": "image_url",
-                                        "image_url": {"url": f"{self.encoded_image}"},
+                                        "image_url": {"url": f"{encoded_image}"},
                                     },
                                 ],
                             }
@@ -397,7 +458,7 @@ def _get_model_deployment_response(
 
                     response = aqua_client.chat.completions.create(**api_kwargs)
 
-                elif self.file_type.startswith("audio"):
+                elif file_type.startswith("audio"):
                     api_kwargs = {
                         "model": model,
                         "messages": [
@@ -407,7 +468,7 @@ def _get_model_deployment_response(
                                     {"type": "text", "text": payload["prompt"]},
                                     {
                                         "type": "audio_url",
-                                        "audio_url": {"url": f"{self.encoded_image}"},
+                                        "audio_url": {"url": f"{encoded_image}"},
                                     },
                                 ],
                             }
@@ -426,7 +487,7 @@ def _get_model_deployment_response(
                     for chunk in response:
                         piece = self._extract_text_from_chunk(chunk)
                         if piece:
-                            print(piece, end="", flush=True)
+                            yield piece
                 except ExtendedRequestError as ex:
                     raise HTTPError(400, str(ex))
                 except Exception as ex:
@@ -468,6 +529,8 @@ def _get_model_deployment_response(
                 raise HTTPError(400, str(ex))
             except Exception as ex:
                 raise HTTPError(500, str(ex))
+        else:
+            raise HTTPError(400, f"Unsupported endpoint_type: {endpoint_type}")
 
     @handle_exceptions
     def post(self, model_deployment_id):
@@ -502,12 +565,11 @@ def post(self, model_deployment_id):
         )
         try:
             for chunk in response_gen:
-                print(chunk)
                 self.write(chunk)
                 self.flush()
             self.finish()
         except Exception as ex:
-            self.set_status(ex.status_code)
+            self.set_status(getattr(ex, "status_code", 500))
             self.write({"message": "Error occurred", "reason": str(ex)})
             self.finish()
 

tests/unitary/with_extras/aqua/test_deployment_handler.py

@@ -8,7 +8,9 @@
 import unittest
 from importlib import reload
 from unittest.mock import MagicMock, patch
+from urllib.error import HTTPError
 
+from ads.aqua.common.enums import PredictEndpoints
 from notebook.base.handlers import IPythonHandler
 from parameterized import parameterized
 
@@ -280,6 +282,220 @@ def test_post(self, mock_get_model_deployment_response):
         self.handler.write.assert_any_call("chunk2")
         self.handler.finish.assert_called_once()
 
+    def test_extract_text_from_choice_dict_delta_content(self):
+        """Test dict choice with delta.content."""
+        choice = {"delta": {"content": "hello"}}
+        result = self.handler._extract_text_from_choice(choice)
+        self.assertEqual(result, "hello")
+
+    def test_extract_text_from_choice_dict_delta_text(self):
+        """Test dict choice with delta.text fallback."""
+        choice = {"delta": {"text": "world"}}
+        result = self.handler._extract_text_from_choice(choice)
+        self.assertEqual(result, "world")
+
+    def test_extract_text_from_choice_dict_message_content(self):
+        """Test dict choice with message.content."""
+        choice = {"message": {"content": "foo"}}
+        result = self.handler._extract_text_from_choice(choice)
+        self.assertEqual(result, "foo")
+
+    def test_extract_text_from_choice_dict_top_level_text(self):
+        """Test dict choice with top-level text."""
+        choice = {"text": "bar"}
+        result = self.handler._extract_text_from_choice(choice)
+        self.assertEqual(result, "bar")
+
+    def test_extract_text_from_choice_object_delta_content(self):
+        """Test object choice with delta.content attribute."""
+        choice = MagicMock()
+        choice.delta = MagicMock(content="obj-content", text=None)
+        result = self.handler._extract_text_from_choice(choice)
+        self.assertEqual(result, "obj-content")
+
+    def test_extract_text_from_choice_object_message_str(self):
+        """Test object choice with message as string."""
+        choice = MagicMock(message="direct-string")
+        result = self.handler._extract_text_from_choice(choice)
+        self.assertEqual(result, "direct-string")
+
+    def test_extract_text_from_choice_none_return(self):
+        """Test choice with no text content returns None."""
+        choice = {}
+        result = self.handler._extract_text_from_choice(choice)
+        self.assertIsNone(result)
+
+    def test_extract_text_from_chunk_dict_with_choices(self):
+        """Test chunk dict with choices list."""
+        chunk = {"choices": [{"delta": {"content": "chunk-text"}}]}
+        result = self.handler._extract_text_from_chunk(chunk)
+        self.assertEqual(result, "chunk-text")
+
+    def test_extract_text_from_chunk_dict_top_level_content(self):
+        """Test chunk dict with top-level content (no choices)."""
+        chunk = {"content": "direct-content"}
+        result = self.handler._extract_text_from_chunk(chunk)
+        self.assertEqual(result, "direct-content")
+
+    def test_extract_text_from_chunk_object_choices(self):
+        """Test object chunk with choices attribute."""
+        chunk = MagicMock()
+        chunk.choices = [{"message": {"content": "obj-chunk"}}]
+        result = self.handler._extract_text_from_chunk(chunk)
+        self.assertEqual(result, "obj-chunk")
+
+    def test_extract_text_from_chunk_empty(self):
+        """Test empty/None chunk returns None."""
+        result = self.handler._extract_text_from_chunk({})
+        self.assertIsNone(result)
+        result = self.handler._extract_text_from_chunk(None)
+        self.assertIsNone(result)
+    
+    @patch('ads.aqua.modeldeployment.AquaDeploymentApp')
+    def test_missing_required_keys_raises_http_error(self, mock_aqua_app):
+        """Test missing required payload keys raises HTTPError."""
+        payload = {"prompt": "test"}
+        with self.assertRaises(HTTPError) as cm:
+            list(self.handler._get_model_deployment_response("test-id", payload))
+        self.assertEqual(cm.exception.status_code, 400)
+        self.assertIn("model", str(cm.exception))
+
+    @patch('ads.aqua.modeldeployment.AquaDeploymentApp')
+    @patch.object(AquaDeploymentStreamingInferenceHandler, '_extract_text_from_chunk')
+    def test_chat_completions_no_image_yields_chunks(self, mock_extract, mock_aqua_app):
+        """Test chat completions without image streams correctly."""
+        mock_deployment = MagicMock()
+        mock_deployment.endpoint = "https://test-endpoint"
+        mock_aqua_app.return_value.get.return_value = mock_deployment
+        
+        mock_stream = iter([MagicMock(choices=[{"delta": {"content": "hello"}}])])
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.return_value = mock_stream
+        with patch.object(self.handler, 'OpenAI', return_value=mock_client):
+            payload = {
+                "endpoint_type": PredictEndpoints.CHAT_COMPLETIONS_ENDPOINT,
+                "prompt": "test prompt",
+                "model": "test-model"
+            }
+            result = list(self.handler._get_model_deployment_response("test-id", payload))
+        
+        mock_extract.assert_called()
+        self.assertEqual(result, ["hello"])
+
+    @patch('ads.aqua.modeldeployment.AquaDeploymentApp')
+    @patch.object(AquaDeploymentStreamingInferenceHandler, '_extract_text_from_chunk')
+    def test_text_completions_endpoint(self, mock_extract, mock_aqua_app):
+        """Test text completions endpoint path."""
+        mock_deployment = MagicMock()
+        mock_deployment.endpoint = "https://test-endpoint"
+        mock_aqua_app.return_value.get.return_value = mock_deployment
+        
+        mock_stream = iter([MagicMock(choices=[{"delta": {"content": "text"}}])])
+        mock_client = MagicMock()
+        mock_client.completions.create.return_value = mock_stream
+        with patch.object(self.handler, 'OpenAI', return_value=mock_client):
+            payload = {
+                "endpoint_type": PredictEndpoints.TEXT_COMPLETIONS_ENDPOINT,
+                "prompt": "test",
+                "model": "test-model"
+            }
+            result = list(self.handler._get_model_deployment_response("test-id", payload))
+        
+        self.assertEqual(result, ["text"])
+
+    @patch('ads.aqua.modeldeployment.AquaDeploymentApp')
+    @patch.object(AquaDeploymentStreamingInferenceHandler, '_extract_text_from_chunk')
+    def test_image_chat_completions(self, mock_extract, mock_aqua_app):
+        """Test chat completions with image input."""
+        mock_deployment = MagicMock()
+        mock_deployment.endpoint = "https://test-endpoint"
+        mock_aqua_app.return_value.get.return_value = mock_deployment
+        
+        mock_stream = iter([MagicMock()])
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.return_value = mock_stream
+        with patch.object(self.handler, 'OpenAI', return_value=mock_client):
+            payload = {
+                "endpoint_type": PredictEndpoints.CHAT_COMPLETIONS_ENDPOINT,
+                "prompt": "describe image",
+                "model": "test-model",
+                "encoded_image": "data:image/jpeg;base64,...",
+                "file_type": "image/jpeg"
+            }
+            list(self.handler._get_model_deployment_response("test-id", payload))
+        
+        expected_call = call(
+            model="test-model",
+            messages=[{
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "describe image"},
+                    {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}  # Note: f-string expands
+                ]
+            }],
+            stream=True
+        )
+        mock_client.chat.completions.create.assert_has_calls([expected_call])
+
+    @patch('ads.aqua.modeldeployment.AquaDeploymentApp')
+    def test_unsupported_endpoint_type_raises_error(self, mock_aqua_app):
+        """Test unsupported endpoint_type raises HTTPError."""
+        mock_aqua_app.return_value.get.return_value = MagicMock(endpoint="test")
+        payload = {
+            "endpoint_type": "invalid-type",
+            "prompt": "test",
+            "model": "test-model"
+        }
+        with self.assertRaises(HTTPError) as cm:
+            list(self.handler._get_model_deployment_response("test-id", payload))
+        self.assertEqual(cm.exception.status_code, 400)
+
+    @patch('ads.aqua.modeldeployment.AquaDeploymentApp')
+    @patch.object(AquaDeploymentStreamingInferenceHandler, '_extract_text_from_chunk')
+    def test_responses_endpoint_with_params(self, mock_extract, mock_aqua_app):
+        """Test responses endpoint with temperature/top_p filtering."""
+        mock_deployment = MagicMock()
+        mock_deployment.endpoint = "https://test-endpoint"
+        mock_aqua_app.return_value.get.return_value = mock_deployment
+        
+        mock_stream = iter([MagicMock()])
+        mock_client = MagicMock()
+        mock_client.responses.create.return_value = mock_stream
+        with patch.object(self.handler, 'OpenAI', return_value=mock_client):
+            payload = {
+                "endpoint_type": PredictEndpoints.RESPONSES,
+                "prompt": "test",
+                "model": "test-model",
+                "temperature": 0.7,
+                "top_p": 0.9
+            }
+            list(self.handler._get_model_deployment_response("test-id", payload))
+        
+        mock_client.responses.create.assert_called_once_with(
+            model="test-model",
+            input="test",
+            stream=True,
+            temperature=0.7,
+            top_p=0.9
+        )
+
+    @patch('ads.aqua.modeldeployment.AquaDeploymentApp')
+    def test_stop_param_normalization(self, mock_aqua_app):
+        """Test stop=[] gets normalized to None."""
+        mock_aqua_app.return_value.get.return_value = MagicMock(endpoint="test")
+        payload = {
+            "endpoint_type": PredictEndpoints.CHAT_COMPLETIONS_ENDPOINT,
+            "prompt": "test",
+            "model": "test-model",
+            "stop": []
+        }
+        # Just verify it doesn't crash - normalization happens before API calls
+        try:
+            next(self.handler._get_model_deployment_response("test-id", payload))
+        except HTTPError:
+            pass  # Expected due to missing client mocks, but normalization should work
+
+
 
 class AquaModelListHandlerTestCase(unittest.TestCase):
     default_params = {