Separating out setting model type into a separate method in the job service

.pre-commit-config.yaml

@@ -17,3 +17,13 @@ repos:
       - id: end-of-file-fixer
       - id: requirements-txt-fixer
         exclude: requirements_lock.txt
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.7.4
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args: [ --fix ]
+      # Run the formatter.
+      - id: ruff-format

lumigator/python/mzai/backend/backend/api/routes/health.py

@@ -33,7 +33,7 @@ def get_job_metadata(job_id: UUID) -> JobSubmissionResponse:
         )
     elif resp.status_code != HTTPStatus.OK:
         loguru.logger.error(
-            f"Unexpected status code getting job metadata text: "
+            "Unexpected status code getting job metadata text: "
              "{resp.status_code}, error: {resp.text or ''}"
         )
         raise HTTPException(

lumigator/python/mzai/backend/backend/services/jobs.py

@@ -104,17 +104,23 @@ def _get_config_template(self, job_type: str, model_name: str) -> str:
 
         return config_template
 
-    def _get_job_params(self, job_type: str, record, request: BaseModel) -> dict:
-        # get dataset S3 path from UUID
-        dataset_s3_path = self.data_service.get_dataset_s3_path(request.dataset)
-
+    def _set_model_type(self, request: BaseModel) -> str:
+        """Sets model URL based on protocol address"""
         if request.model.startswith("oai://"):
             model_url = settings.OAI_API_URL
         elif request.model.startswith("mistral://"):
             model_url = settings.MISTRAL_API_URL
         else:
             model_url = request.model_url
 
+        return model_url
+
+    def _get_job_params(self, job_type: str, record, request: BaseModel) -> dict:
+        # get dataset S3 path from UUID
+        dataset_s3_path = self.data_service.get_dataset_s3_path(request.dataset)
+
+        model_url = self._set_model_type(request)
+
         # provide a reasonable system prompt for services where none was specified
         if request.system_prompt is None and not request.model.startswith("hf://"):
             request.system_prompt = settings.DEFAULT_SUMMARIZER_PROMPT

lumigator/python/mzai/backend/backend/tests/services/test_job_service.py

@@ -9,7 +9,7 @@
 def test_set_null_inference_job_params(job_record,job_service ):
     request = JobInferenceCreate(name="test_run_hugging_face",
                                  description="Test run for Huggingface model",
-                                    model="Test run for Huggingface model",
+                                 model="hf://facebook/bart-large-cnn",
                                  dataset="cced289c-f869-4af1-9195-1d58e32d1cc1")
     params = job_service._get_job_params("INFERENCE", job_record,request)
     assert params["max_samples"] == -1
@@ -18,7 +18,38 @@ def test_set_explicit_inference_job_params(job_record,job_service ):
     request = JobInferenceCreate(name="test_run_hugging_face",
                                  description="Test run for Huggingface model",
                                  max_samples=10,
-                                    model="Test run for Huggingface model",
+                                 model="hf://facebook/bart-large-cnn",
                                  dataset="cced289c-f869-4af1-9195-1d58e32d1cc1")
     params = job_service._get_job_params("INFERENCE", job_record,request)
     assert params["max_samples"] == 10
+
+
+def test_set_model_hf(job_record,job_service):
+    request = JobInferenceCreate(name="test_run_hugging_face",
+                                 description="Test run for Huggingface model",
+                                 max_samples=10,
+                                 model="hf://facebook/bart-large-cnn",
+                                 model_url="hf://facebook/bart-large-cnn",
+                                 dataset="cced289c-f869-4af1-9195-1d58e32d1cc1")
+    model_url = job_service._set_model_type(request)
+    assert model_url == "hf://facebook/bart-large-cnn"
+
+def test_set_model_openai(job_record,job_service):
+    request = JobInferenceCreate(name="test_run_openai",
+                                 description="Test run for OpenAI server",
+                                 max_samples=10,
+                                 model="oai://gpt-4-turbo",
+                                 dataset="cced289c-f869-4af1-9195-1d58e32d1cc1")
+
+    model_url = job_service._set_model_type(request)
+    assert model_url == "https://api.openai.com/v1"
+
+def test_set_model_mistral(job_record,job_service):
+    request = JobInferenceCreate(name="test_run_mistral",
+                                 description="Test run for Mistral server",
+                                 max_samples=10,
+                                 model="mistral://open-mistral-7b",
+                                 dataset="cced289c-f869-4af1-9195-1d58e32d1cc1")
+
+    model_url = job_service._set_model_type(request)
+    assert model_url == "https://api.mistral.ai/v1"

lumigator/python/mzai/sdk/tests/data/job-all.json

@@ -4,8 +4,8 @@
     "model": "some_model",
     "dataset": "6f6487ac-7170-4a11-af7a-0f6db1ec9a74",
     "max_samples": 5,
-    "model_url": "s3://some_url",
+    "model_url": "hf://facebook/bart-large-cnn",
     "system_prompt": "some prompt",
     "config_infer_template": "some_infer_template",
     "config_eval_template": "some_eval_template"
-  }
+  }