[AQUA] Add API endpoints to manage chat templates via model custom metadata (#1213)

Author: sidhantkmund

ads/aqua/constants.py

@@ -55,6 +55,8 @@
 SUPPORTED_FILE_FORMATS = ["jsonl"]
 MODEL_BY_REFERENCE_OSS_PATH_KEY = "artifact_location"
 
+AQUA_CHAT_TEMPLATE_METADATA_KEY = "chat_template"
+
 CONSOLE_LINK_RESOURCE_TYPE_MAPPING = {
  "datasciencemodel": "models",
  "datasciencemodeldeployment": "model-deployments",

ads/aqua/extension/model_handler.py

@@ -11,12 +11,15 @@
 from ads.aqua.common.enums import CustomInferenceContainerTypeFamily
 from ads.aqua.common.errors import AquaRuntimeError
 from ads.aqua.common.utils import get_hf_model_info, is_valid_ocid, list_hf_models
+from ads.aqua.constants import AQUA_CHAT_TEMPLATE_METADATA_KEY
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModelSummary, HFModelSummary
 from ads.config import SERVICE
+from ads.model import DataScienceModel
 from ads.model.common.utils import MetadataArtifactPathType
+from ads.model.service.oci_datascience_model import OCIDataScienceModel
 
 
 class AquaModelHandler(AquaAPIhandler):
@@ -320,26 +323,65 @@ def post(self, *args, **kwargs): # noqa: ARG002
  )
 
 
-class AquaModelTokenizerConfigHandler(AquaAPIhandler):
+class AquaModelChatTemplateHandler(AquaAPIhandler):
  def get(self, model_id):
  """
- Handles requests for retrieving the Hugging Face tokenizer configuration of a specified model.
- Expected request format: GET /aqua/models/<model-ocid>/tokenizer
+ Handles requests for retrieving the chat template from custom metadata of a specified model.
+ Expected request format: GET /aqua/models/<model-ocid>/chat-template
 
  """
 
  path_list = urlparse(self.request.path).path.strip("/").split("/")
- # Path should be /aqua/models/ocid1.iad.ahdxxx/tokenizer
- # path_list=['aqua','models','<model-ocid>','tokenizer']
+ # Path should be /aqua/models/ocid1.iad.ahdxxx/chat-template
+ # path_list=['aqua','models','<model-ocid>','chat-template']
  if (
  len(path_list) == 4
  and is_valid_ocid(path_list[2])
- and path_list[3] == "tokenizer"
+ and path_list[3] == "chat-template"
  ):
- return self.finish(AquaModelApp().get_hf_tokenizer_config(model_id))
+ try:
+ oci_data_science_model = OCIDataScienceModel.from_id(model_id)
+ except Exception as e:
+ raise HTTPError(404, f"Model not found for id: {model_id}. Details: {str(e)}")
+ return self.finish(oci_data_science_model.get_custom_metadata_artifact("chat_template"))
 
  raise HTTPError(400, f"The request {self.request.path} is invalid.")
 
+ @handle_exceptions
+ def post(self, model_id: str):
+ """
+ Handles POST requests to add a custom chat_template metadata artifact to a model.
+
+ Expected request format:
+ POST /aqua/models/<model-ocid>/chat-template
+ Body: { "chat_template": "<your_template_string>" }
+
+ """
+ try:
+ input_body = self.get_json_body()
+ except Exception as e:
+ raise HTTPError(400, f"Invalid JSON body: {str(e)}")
+
+ chat_template = input_body.get("chat_template")
+ if not chat_template:
+ raise HTTPError(400, "Missing required field: 'chat_template'")
+
+ try:
+ data_science_model = DataScienceModel.from_id(model_id)
+ except Exception as e:
+ raise HTTPError(404, f"Model not found for id: {model_id}. Details: {str(e)}")
+
+ try:
+ result = data_science_model.create_custom_metadata_artifact(
+ metadata_key_name=AQUA_CHAT_TEMPLATE_METADATA_KEY,
+ path_type=MetadataArtifactPathType.CONTENT,
+ artifact_path_or_content=chat_template.encode()
+ )
+ except Exception as e:
+ raise HTTPError(500, f"Failed to create metadata artifact: {str(e)}")
+
+ return self.finish(result)
+
 
 class AquaModelDefinedMetadataArtifactHandler(AquaAPIhandler):
  """
@@ -381,7 +423,7 @@ def post(self, model_id: str, metadata_key: str):
  ("model/?([^/]*)", AquaModelHandler),
  ("model/?([^/]*)/license", AquaModelLicenseHandler),
  ("model/?([^/]*)/readme", AquaModelReadmeHandler),
- ("model/?([^/]*)/tokenizer", AquaModelTokenizerConfigHandler),
+ ("model/?([^/]*)/chat-template", AquaModelChatTemplateHandler),
  ("model/hf/search/?([^/]*)", AquaHuggingFaceHandler),
  (
  "model/?([^/]*)/definedMetadata/?([^/]*)",

tests/unitary/with_extras/aqua/test_model_handler.py

@@ -4,7 +4,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 from unicodedata import category
 from unittest import TestCase
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, patch, ANY
 
 import pytest
 from huggingface_hub.hf_api import HfApi, ModelInfo
@@ -14,13 +14,13 @@
 
 from ads.aqua.common.errors import AquaRuntimeError
 from ads.aqua.common.utils import get_hf_model_info
-from ads.aqua.constants import AQUA_TROUBLESHOOTING_LINK, STATUS_CODE_MESSAGES
+from ads.aqua.constants import AQUA_TROUBLESHOOTING_LINK, STATUS_CODE_MESSAGES, AQUA_CHAT_TEMPLATE_METADATA_KEY
 from ads.aqua.extension.errors import ReplyDetails
 from ads.aqua.extension.model_handler import (
  AquaHuggingFaceHandler,
  AquaModelHandler,
  AquaModelLicenseHandler,
- AquaModelTokenizerConfigHandler,
+ AquaModelChatTemplateHandler
 )
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModel, AquaModelSummary, HFModelSummary
@@ -254,39 +254,114 @@ def test_get(self, mock_load_license):
  mock_load_license.assert_called_with("test_model_id")
 
 
-class ModelTokenizerConfigHandlerTestCase(TestCase):
+class AquaModelChatTemplateHandlerTestCase(TestCase):
  @patch.object(IPythonHandler, "__init__")
  def setUp(self, ipython_init_mock) -> None:
  ipython_init_mock.return_value = None
- self.model_tokenizer_config_handler = AquaModelTokenizerConfigHandler(
+ self.model_chat_template_handler = AquaModelChatTemplateHandler(
  MagicMock(), MagicMock()
  )
- self.model_tokenizer_config_handler.finish = MagicMock()
- self.model_tokenizer_config_handler.request = MagicMock()
+ self.model_chat_template_handler.finish = MagicMock()
+ self.model_chat_template_handler.request = MagicMock()
+ self.model_chat_template_handler._headers = {}
 
- @patch.object(AquaModelApp, "get_hf_tokenizer_config")
+ @patch("ads.aqua.extension.model_handler.OCIDataScienceModel.from_id")
  @patch("ads.aqua.extension.model_handler.urlparse")
- def test_get(self, mock_urlparse, mock_get_hf_tokenizer_config):
- request_path = MagicMock(path="aqua/model/ocid1.xx./tokenizer")
+ def test_get_valid_path(self, mock_urlparse, mock_from_id):
+ request_path = MagicMock(path="/aqua/models/ocid1.xx./chat-template")
  mock_urlparse.return_value = request_path
- self.model_tokenizer_config_handler.get(model_id="test_model_id")
- self.model_tokenizer_config_handler.finish.assert_called_with(
- mock_get_hf_tokenizer_config.return_value
- )
- mock_get_hf_tokenizer_config.assert_called_with("test_model_id")
 
- @patch.object(AquaModelApp, "get_hf_tokenizer_config")
+ model_mock = MagicMock()
+ model_mock.get_custom_metadata_artifact.return_value = "chat_template_string"
+ mock_from_id.return_value = model_mock
+
+ self.model_chat_template_handler.get(model_id="test_model_id")
+ self.model_chat_template_handler.finish.assert_called_with("chat_template_string")
+ model_mock.get_custom_metadata_artifact.assert_called_with("chat_template")
+
  @patch("ads.aqua.extension.model_handler.urlparse")
- def test_get_invalid_path(self, mock_urlparse, mock_get_hf_tokenizer_config):
- """Test invalid request path should raise HTTPError(400)"""
- request_path = MagicMock(path="/invalid/path")
+ def test_get_invalid_path(self, mock_urlparse):
+ request_path = MagicMock(path="/wrong/path")
  mock_urlparse.return_value = request_path
 
  with self.assertRaises(HTTPError) as context:
- self.model_tokenizer_config_handler.get(model_id="test_model_id")
+ self.model_chat_template_handler.get("ocid1.test.chat")
  self.assertEqual(context.exception.status_code, 400)
- self.model_tokenizer_config_handler.finish.assert_not_called()
- mock_get_hf_tokenizer_config.assert_not_called()
+
+ @patch("ads.aqua.extension.model_handler.OCIDataScienceModel.from_id", side_effect=Exception("Not found"))
+ @patch("ads.aqua.extension.model_handler.urlparse")
+ def test_get_model_not_found(self, mock_urlparse, mock_from_id):
+ request_path = MagicMock(path="/aqua/models/ocid1.invalid/chat-template")
+ mock_urlparse.return_value = request_path
+
+ with self.assertRaises(HTTPError) as context:
+ self.model_chat_template_handler.get("ocid1.invalid")
+ self.assertEqual(context.exception.status_code, 404)
+
+ @patch("ads.aqua.extension.model_handler.DataScienceModel.from_id")
+ def test_post_valid(self, mock_from_id):
+ model_mock = MagicMock()
+ model_mock.create_custom_metadata_artifact.return_value = {"result": "success"}
+ mock_from_id.return_value = model_mock
+
+ self.model_chat_template_handler.get_json_body = MagicMock(return_value={"chat_template": "Hello <|user|>"})
+ result = self.model_chat_template_handler.post("ocid1.valid")
+ self.model_chat_template_handler.finish.assert_called_with({"result": "success"})
+
+ model_mock.create_custom_metadata_artifact.assert_called_with(
+ metadata_key_name=AQUA_CHAT_TEMPLATE_METADATA_KEY,
+ path_type=ANY,
+ artifact_path_or_content=b"Hello <|user|>"
+ )
+
+ @patch.object(AquaModelChatTemplateHandler, "write_error")
+ def test_post_invalid_json(self, mock_write_error):
+ self.model_chat_template_handler.get_json_body = MagicMock(side_effect=Exception("Invalid JSON"))
+ self.model_chat_template_handler._headers = {}
+ self.model_chat_template_handler.post("ocid1.test.invalidjson")
+
+ mock_write_error.assert_called_once()
+
+ kwargs = mock_write_error.call_args.kwargs
+ exc_info = kwargs.get("exc_info")
+
+ assert exc_info is not None
+ exc_type, exc_instance, _ = exc_info
+
+ assert isinstance(exc_instance, HTTPError)
+ assert exc_instance.status_code == 400
+ assert "Invalid JSON body" in str(exc_instance)
+
+ @patch.object(AquaModelChatTemplateHandler, "write_error")
+ def test_post_missing_chat_template(self, mock_write_error):
+ self.model_chat_template_handler.get_json_body = MagicMock(return_value={})
+ self.model_chat_template_handler._headers = {}
+
+ self.model_chat_template_handler.post("ocid1.test.model")
+
+ mock_write_error.assert_called_once()
+ exc_info = mock_write_error.call_args.kwargs.get("exc_info")
+ assert exc_info is not None
+ _, exc_instance, _ = exc_info
+ assert isinstance(exc_instance, HTTPError)
+ assert exc_instance.status_code == 400
+ assert "Missing required field: 'chat_template'" in str(exc_instance)
+
+ @patch("ads.aqua.extension.model_handler.DataScienceModel.from_id", side_effect=Exception("Not found"))
+ @patch.object(AquaModelChatTemplateHandler, "write_error")
+ def test_post_model_not_found(self, mock_write_error, mock_from_id):
+ self.model_chat_template_handler.get_json_body = MagicMock(return_value={"chat_template": "test template"})
+ self.model_chat_template_handler._headers = {}
+
+ self.model_chat_template_handler.post("ocid1.invalid.model")
+
+ mock_write_error.assert_called_once()
+ exc_info = mock_write_error.call_args.kwargs.get("exc_info")
+ assert exc_info is not None
+ _, exc_instance, _ = exc_info
+ assert isinstance(exc_instance, HTTPError)
+ assert exc_instance.status_code == 404
+ assert "Model not found" in str(exc_instance)
 
 
 class TestAquaHuggingFaceHandler: