Init update on model_server

model_server/app/tests/test_app.py

@@ -13,6 +13,7 @@ client = TestClient(app)
 logger.info(f"Model will be loaded on device: {glb.DEVICE}")
 
 
+# [TODO] Review: update the following code
 # Unit tests for the health check endpoint
 @pytest.mark.asyncio
 @patch("app.loader.glb.DEVICE", glb.DEVICE)  # Mock the device to 'cpu'
@@ -22,6 +23,7 @@ async def test_healthz():
     assert response.json() == {"status": "ok"}
 
 
+# [TODO] Review: update the following code
 # Unit test for the models endpoint
 @pytest.mark.asyncio
 @patch("app.loader.glb.DEVICE", glb.DEVICE)  # Mock the device to 'cpu'
@@ -32,6 +34,7 @@ async def test_models():
     assert len(response.json()["data"]) > 0
 
 
+# [TODO] Review: update the following code
 # Unit test for embeddings endpoint
 @pytest.mark.asyncio
 @patch("app.loader.glb.DEVICE", glb.DEVICE)  # Mock the device to 'cpu'
@@ -46,6 +49,7 @@ async def test_embedding():
         assert response.status_code == 400
 
 
+# [TODO] Review: update the following code
 # Unit test for the guard endpoint
 @pytest.mark.asyncio
 @patch("app.loader.glb.DEVICE", glb.DEVICE)  # Mock the device to 'cpu'
@@ -56,6 +60,7 @@ async def test_guard():
     assert "jailbreak_verdict" in response.json()
 
 
+# [TODO] Review: update the following code
 # Unit test for the zero-shot endpoint
 @pytest.mark.asyncio
 @patch("app.loader.glb.DEVICE", glb.DEVICE)  # Mock the device to 'cpu'
@@ -73,6 +78,7 @@ async def test_zeroshot():
         assert response.status_code == 400
 
 
+# [TODO] Review: update the following code
 # Unit test for the hallucination endpoint
 @pytest.mark.asyncio
 @patch("app.loader.glb.DEVICE", glb.DEVICE)  # Mock the device to 'cpu'
@@ -90,6 +96,7 @@ async def test_hallucination():
         assert response.status_code == 400
 
 
+# [TODO] Review: update the following code
 # Unit test for the chat completion endpoint
 @pytest.mark.asyncio
 @patch("app.loader.glb.DEVICE", glb.DEVICE)  # Mock the device to 'cpu'

model_server/app/tests/test_cli_stop_server.py

@@ -1,7 +1,6 @@
 import unittest
+
 from unittest.mock import patch, MagicMock
-import subprocess
-import time
 from app.cli import kill_process
 
 

model_server/app/tests/test_function_calling.py

@@ -1,13 +1,12 @@
+import json
 import pytest
-from unittest.mock import AsyncMock, MagicMock, patch
-import app.commons.constants as const
+
 from fastapi import Response
-from app.function_calling.model_utils import (
-    process_messages,
-    chat_completion,
+from unittest.mock import AsyncMock, MagicMock, patch
+from app.commons.globals import handler_map
+from app.model_handler.function_calling import (
     Message,
     ChatMessage,
-    Choice,
     ChatCompletionResponse,
 )
 
@@ -31,14 +30,27 @@ def sample_messages():
 def sample_request(sample_messages):
     return ChatMessage(
         messages=sample_messages,
-        tools=[{"name": "sample_tool", "description": "A sample tool"}],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "sample_tool",
+                    "description": "A sample tool",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {},
+                        "required": [],
+                    },
+                },
+            }
+        ],
     )
 
 
-@patch("app.commons.constants.arch_function_hanlder")
+@patch("app.commons.globals.handler_map")
 def test_process_messages(mock_hanlder):
     messages = sample_messages()
-    processed = process_messages(messages)
+    processed = handler_map["Arch-Function"]._process_messages(messages)
 
     assert len(processed) == 3
     assert processed[0] == {"role": "user", "content": "Hello!"}
@@ -48,10 +60,11 @@ def test_process_messages(mock_hanlder):
     }
     assert processed[2] == {
         "role": "user",
-        "content": "<tool_response>\nResponse from tool\n</tool_response>",
+        "content": f"<tool_response>\n{json.dumps('Response from tool')}\n</tool_response>",
     }
 
 
+# [TODO] Review: Update the following test
 @patch("app.commons.constants.arch_function_client")
 @patch("app.commons.constants.arch_function_hanlder")
 @pytest.mark.asyncio

model_server/app/tests/test_guardrails.py

@@ -0,0 +1,47 @@
+import os
+import pytest
+from unittest.mock import patch, MagicMock
+import app.commons.globals as glb
+
+# Mock constants
+glb.DEVICE = "cpu"  # Adjust as needed for your test case
+arch_guard_model_type = {
+    "cpu": "katanemo/Arch-Guard-cpu",
+    "cuda": "katanemo/Arch-Guard",
+    "mps": "katanemo/Arch-Guard",
+}
+
+
+# [TODO] Review: update the following code to test under `cpu`, `cuda`, and `mps`
+# Test for get_prompt_guard function
+@patch("app.loader.AutoTokenizer.from_pretrained")
+@patch("app.loader.OVModelForSequenceClassification.from_pretrained")
+@patch("app.loader.AutoModelForSequenceClassification.from_pretrained")
+def test_get_prompt_guard(mock_auto_model, mock_ov_model, mock_tokenizer):
+    # Mock model based on device
+    if glb.DEVICE == "cpu":
+        mock_ov_model.return_value = MagicMock()
+    else:
+        mock_auto_model.return_value = MagicMock()
+
+    mock_tokenizer.return_value = MagicMock()
+
+    prompt_guard = get_prompt_guard(arch_guard_model_type[glb.DEVICE])
+
+    # Assertions
+    assert prompt_guard["model_name"] == arch_guard_model_type[glb.DEVICE]
+    mock_tokenizer.assert_called_once_with(
+        arch_guard_model_type[glb.DEVICE], trust_remote_code=True
+    )
+    if glb.DEVICE == "cpu":
+        mock_ov_model.assert_called_once_with(
+            arch_guard_model_type[glb.DEVICE],
+            device_map=glb.DEVICE,
+            low_cpu_mem_usage=True,
+        )
+    else:
+        mock_auto_model.assert_called_once_with(
+            arch_guard_model_type[glb.DEVICE],
+            device_map=glb.DEVICE,
+            low_cpu_mem_usage=True,
+        )

model_server/app/tests/test_hallucination.py

@@ -1,8 +1,11 @@
 import json
-from app.function_calling.hallucination_handler import HallucinationStateHandler
 import pytest
 import os
 
+
+from app.model_handler.hallucination_handler import HallucinationStateHandler
+
+
 # Get the directory of the current file
 current_dir = os.path.dirname(__file__)
 
@@ -45,6 +48,7 @@ if type(function_description) != list:
     function_description = [get_weather_api["function"]]
 
 
+# [TODO] Review: update the following code
 @pytest.mark.parametrize("case", test_cases)
 def test_hallucination(case):
     state = HallucinationStateHandler(
@@ -58,6 +62,7 @@ def test_hallucination(case):
     assert state.hallucination == case["expect"]
 
 
+# [TODO] Review: update the following code
 @pytest.mark.parametrize("is_hallucinate_sample", [True, False])
 def test_hallucination_prompt(is_hallucinate_sample):
     TASK_PROMPT = """

model_server/app/tests/test_loaders_cpu.py

@@ -1,102 +0,0 @@
-import os
-import pytest
-from unittest.mock import patch, MagicMock
-import app.commons.globals as glb
-from app.loader import get_embedding_model, get_zero_shot_model, get_prompt_guard
-
-# Mock constants
-glb.DEVICE = "cpu"  # Adjust as needed for your test case
-arch_guard_model_type = {
-    "cpu": "katanemo/Arch-Guard-cpu",
-    "cuda": "katanemo/Arch-Guard",
-    "mps": "katanemo/Arch-Guard",
-}
-
-
-@pytest.fixture
-def mock_env():
-    # Mock environment variables
-    os.environ["MODELS"] = "katanemo/bge-large-en-v1.5"
-    os.environ["ZERO_SHOT_MODELS"] = "katanemo/bart-large-mnli"
-
-
-# Test for get_embedding_model function
-@patch("app.loader.ORTModelForFeatureExtraction.from_pretrained")
-@patch("app.loader.AutoModel.from_pretrained")
-@patch("app.loader.AutoTokenizer.from_pretrained")
-def test_get_embedding_model(mock_tokenizer, mock_automodel, mock_ort_model, mock_env):
-    mock_automodel.return_value = MagicMock()
-    mock_ort_model.return_value = MagicMock()
-    mock_tokenizer.return_value = MagicMock()
-
-    embedding_model = get_embedding_model()
-
-    # Assertions
-    assert embedding_model["model_name"] == "katanemo/bge-large-en-v1.5"
-    mock_tokenizer.assert_called_once_with(
-        "katanemo/bge-large-en-v1.5", trust_remote_code=True
-    )
-    if glb.DEVICE != "cuda":
-        mock_ort_model.assert_called_once_with(
-            "katanemo/bge-large-en-v1.5", file_name="onnx/model.onnx"
-        )
-    else:
-        mock_automodel.assert_called_once_with(
-            "katanemo/bge-large-en-v1.5", device_map=glb.DEVICE
-        )
-
-
-# Test for get_zero_shot_model function
-@patch("app.loader.ORTModelForSequenceClassification.from_pretrained")
-@patch("app.loader.pipeline")
-@patch("app.loader.AutoTokenizer.from_pretrained")
-def test_get_zero_shot_model(mock_tokenizer, mock_pipeline, mock_ort_model, mock_env):
-    mock_pipeline.return_value = MagicMock()
-    mock_ort_model.return_value = MagicMock()
-    mock_tokenizer.return_value = MagicMock()
-
-    zero_shot_model = get_zero_shot_model()
-
-    # Assertions
-    assert zero_shot_model["model_name"] == "katanemo/bart-large-mnli"
-    mock_tokenizer.assert_called_once_with("katanemo/bart-large-mnli")
-    if glb.DEVICE != "cuda":
-        mock_ort_model.assert_called_once_with(
-            "katanemo/bart-large-mnli", file_name="onnx/model.onnx"
-        )
-    else:
-        assert mock_pipeline.called_once()
-
-
-# Test for get_prompt_guard function
-@patch("app.loader.AutoTokenizer.from_pretrained")
-@patch("app.loader.OVModelForSequenceClassification.from_pretrained")
-@patch("app.loader.AutoModelForSequenceClassification.from_pretrained")
-def test_get_prompt_guard(mock_auto_model, mock_ov_model, mock_tokenizer):
-    # Mock model based on device
-    if glb.DEVICE == "cpu":
-        mock_ov_model.return_value = MagicMock()
-    else:
-        mock_auto_model.return_value = MagicMock()
-
-    mock_tokenizer.return_value = MagicMock()
-
-    prompt_guard = get_prompt_guard(arch_guard_model_type[glb.DEVICE])
-
-    # Assertions
-    assert prompt_guard["model_name"] == arch_guard_model_type[glb.DEVICE]
-    mock_tokenizer.assert_called_once_with(
-        arch_guard_model_type[glb.DEVICE], trust_remote_code=True
-    )
-    if glb.DEVICE == "cpu":
-        mock_ov_model.assert_called_once_with(
-            arch_guard_model_type[glb.DEVICE],
-            device_map=glb.DEVICE,
-            low_cpu_mem_usage=True,
-        )
-    else:
-        mock_auto_model.assert_called_once_with(
-            arch_guard_model_type[glb.DEVICE],
-            device_map=glb.DEVICE,
-            low_cpu_mem_usage=True,
-        )

model_server/app/tests/test_loaders_gpu.py

@@ -1,102 +0,0 @@
-import os
-import pytest
-from unittest.mock import patch, MagicMock
-import app.commons.globals as glb
-from app.loader import get_embedding_model, get_zero_shot_model, get_prompt_guard
-
-# Mock constants
-glb.DEVICE = "cuda"  # Adjust as needed for your test case
-arch_guard_model_type = {
-    "cpu": "katanemo/Arch-Guard-cpu",
-    "cuda": "katanemo/Arch-Guard",
-    "mps": "katanemo/Arch-Guard",
-}
-
-
-@pytest.fixture
-def mock_env():
-    # Mock environment variables
-    os.environ["MODELS"] = "katanemo/bge-large-en-v1.5"
-    os.environ["ZERO_SHOT_MODELS"] = "katanemo/bart-large-mnli"
-
-
-# Test for get_embedding_model function
-@patch("app.loader.ORTModelForFeatureExtraction.from_pretrained")
-@patch("app.loader.AutoModel.from_pretrained")
-@patch("app.loader.AutoTokenizer.from_pretrained")
-def test_get_embedding_model(mock_tokenizer, mock_automodel, mock_ort_model, mock_env):
-    mock_automodel.return_value = MagicMock()
-    mock_ort_model.return_value = MagicMock()
-    mock_tokenizer.return_value = MagicMock()
-
-    embedding_model = get_embedding_model()
-
-    # Assertions
-    assert embedding_model["model_name"] == "katanemo/bge-large-en-v1.5"
-    mock_tokenizer.assert_called_once_with(
-        "katanemo/bge-large-en-v1.5", trust_remote_code=True
-    )
-    if glb.DEVICE != "cuda":
-        mock_ort_model.assert_called_once_with(
-            "katanemo/bge-large-en-v1.5", file_name="onnx/model.onnx"
-        )
-    else:
-        mock_automodel.assert_called_once_with(
-            "katanemo/bge-large-en-v1.5", device_map=glb.DEVICE
-        )
-
-
-# Test for get_zero_shot_model function
-@patch("app.loader.ORTModelForSequenceClassification.from_pretrained")
-@patch("app.loader.pipeline")
-@patch("app.loader.AutoTokenizer.from_pretrained")
-def test_get_zero_shot_model(mock_tokenizer, mock_pipeline, mock_ort_model, mock_env):
-    mock_pipeline.return_value = MagicMock()
-    mock_ort_model.return_value = MagicMock()
-    mock_tokenizer.return_value = MagicMock()
-
-    zero_shot_model = get_zero_shot_model()
-
-    # Assertions
-    assert zero_shot_model["model_name"] == "katanemo/bart-large-mnli"
-    mock_tokenizer.assert_called_once_with("katanemo/bart-large-mnli")
-    if glb.DEVICE != "cuda":
-        mock_ort_model.assert_called_once_with(
-            "katanemo/bart-large-mnli", file_name="onnx/model.onnx"
-        )
-    else:
-        assert mock_pipeline.called_once()
-
-
-# Test for get_prompt_guard function
-@patch("app.loader.AutoTokenizer.from_pretrained")
-@patch("app.loader.OVModelForSequenceClassification.from_pretrained")
-@patch("app.loader.AutoModelForSequenceClassification.from_pretrained")
-def test_get_prompt_guard(mock_auto_model, mock_ov_model, mock_tokenizer):
-    # Mock model based on device
-    if glb.DEVICE == "cpu":
-        mock_ov_model.return_value = MagicMock()
-    else:
-        mock_auto_model.return_value = MagicMock()
-
-    mock_tokenizer.return_value = MagicMock()
-
-    prompt_guard = get_prompt_guard(arch_guard_model_type[glb.DEVICE])
-
-    # Assertions
-    assert prompt_guard["model_name"] == arch_guard_model_type[glb.DEVICE]
-    mock_tokenizer.assert_called_once_with(
-        arch_guard_model_type[glb.DEVICE], trust_remote_code=True
-    )
-    if glb.DEVICE == "cpu":
-        mock_ov_model.assert_called_once_with(
-            arch_guard_model_type[glb.DEVICE],
-            device_map=glb.DEVICE,
-            low_cpu_mem_usage=True,
-        )
-    else:
-        mock_auto_model.assert_called_once_with(
-            arch_guard_model_type[glb.DEVICE],
-            device_map=glb.DEVICE,
-            low_cpu_mem_usage=True,
-        )

model_server/app/tests/test_loaders_mps.py

@@ -1,102 +0,0 @@
-import os
-import pytest
-from unittest.mock import patch, MagicMock
-import app.commons.globals as glb
-from app.loader import get_embedding_model, get_zero_shot_model, get_prompt_guard
-
-# Mock constants
-glb.DEVICE = "mps"  # Adjust as needed for your test case
-arch_guard_model_type = {
-    "cpu": "katanemo/Arch-Guard-cpu",
-    "cuda": "katanemo/Arch-Guard",
-    "mps": "katanemo/Arch-Guard",
-}
-
-
-@pytest.fixture
-def mock_env():
-    # Mock environment variables
-    os.environ["MODELS"] = "katanemo/bge-large-en-v1.5"
-    os.environ["ZERO_SHOT_MODELS"] = "katanemo/bart-large-mnli"
-
-
-# Test for get_embedding_model function
-@patch("app.loader.ORTModelForFeatureExtraction.from_pretrained")
-@patch("app.loader.AutoModel.from_pretrained")
-@patch("app.loader.AutoTokenizer.from_pretrained")
-def test_get_embedding_model(mock_tokenizer, mock_automodel, mock_ort_model, mock_env):
-    mock_automodel.return_value = MagicMock()
-    mock_ort_model.return_value = MagicMock()
-    mock_tokenizer.return_value = MagicMock()
-
-    embedding_model = get_embedding_model()
-
-    # Assertions
-    assert embedding_model["model_name"] == "katanemo/bge-large-en-v1.5"
-    mock_tokenizer.assert_called_once_with(
-        "katanemo/bge-large-en-v1.5", trust_remote_code=True
-    )
-    if glb.DEVICE != "cuda":
-        mock_ort_model.assert_called_once_with(
-            "katanemo/bge-large-en-v1.5", file_name="onnx/model.onnx"
-        )
-    else:
-        mock_automodel.assert_called_once_with(
-            "katanemo/bge-large-en-v1.5", device_map=glb.DEVICE
-        )
-
-
-# Test for get_zero_shot_model function
-@patch("app.loader.ORTModelForSequenceClassification.from_pretrained")
-@patch("app.loader.pipeline")
-@patch("app.loader.AutoTokenizer.from_pretrained")
-def test_get_zero_shot_model(mock_tokenizer, mock_pipeline, mock_ort_model, mock_env):
-    mock_pipeline.return_value = MagicMock()
-    mock_ort_model.return_value = MagicMock()
-    mock_tokenizer.return_value = MagicMock()
-
-    zero_shot_model = get_zero_shot_model()
-
-    # Assertions
-    assert zero_shot_model["model_name"] == "katanemo/bart-large-mnli"
-    mock_tokenizer.assert_called_once_with("katanemo/bart-large-mnli")
-    if glb.DEVICE != "cuda":
-        mock_ort_model.assert_called_once_with(
-            "katanemo/bart-large-mnli", file_name="onnx/model.onnx"
-        )
-    else:
-        assert mock_pipeline.called_once()
-
-
-# Test for get_prompt_guard function
-@patch("app.loader.AutoTokenizer.from_pretrained")
-@patch("app.loader.OVModelForSequenceClassification.from_pretrained")
-@patch("app.loader.AutoModelForSequenceClassification.from_pretrained")
-def test_get_prompt_guard(mock_auto_model, mock_ov_model, mock_tokenizer):
-    # Mock model based on device
-    if glb.DEVICE == "cpu":
-        mock_ov_model.return_value = MagicMock()
-    else:
-        mock_auto_model.return_value = MagicMock()
-
-    mock_tokenizer.return_value = MagicMock()
-
-    prompt_guard = get_prompt_guard(arch_guard_model_type[glb.DEVICE])
-
-    # Assertions
-    assert prompt_guard["model_name"] == arch_guard_model_type[glb.DEVICE]
-    mock_tokenizer.assert_called_once_with(
-        arch_guard_model_type[glb.DEVICE], trust_remote_code=True
-    )
-    if glb.DEVICE == "cpu":
-        mock_ov_model.assert_called_once_with(
-            arch_guard_model_type[glb.DEVICE],
-            device_map=glb.DEVICE,
-            low_cpu_mem_usage=True,
-        )
-    else:
-        mock_auto_model.assert_called_once_with(
-            arch_guard_model_type[glb.DEVICE],
-            device_map=glb.DEVICE,
-            low_cpu_mem_usage=True,
-        )

model_server/app/tests/test_state.py

@@ -1,66 +1,50 @@
-from typing import List
-import pytest
-import json
-from app.function_calling.model_utils import Message, process_messages
+from app.commons.globals import handler_map
+from app.model_handler.function_calling import Message
 
-test_input_history = """
-[
+
+test_input_history = [
+    {"role": "user", "content": "how is the weather in chicago for next 5 days?"},
     {
-      "role": "user",
-      "content": "how is the weather in chicago for next 5 days?"
+        "role": "assistant",
+        "model": "Arch-Function",
+        "tool_calls": [
+            {
+                "id": "call_3394",
+                "type": "function",
+                "function": {
+                    "name": "weather_forecast",
+                    "arguments": {"city": "Chicago", "days": 5},
+                },
+            }
+        ],
     },
+    {"role": "tool", "content": "--", "tool_call_id": "call_3394"},
+    {"role": "assistant", "content": "--", "model": "gpt-3.5-turbo-0125"},
+    {"role": "user", "content": "how is the weather in chicago for next 5 days?"},
     {
-      "role": "assistant",
-      "model": "Arch-Function-1.5B",
-      "tool_calls": [
-        {
-          "id": "call_3394",
-          "type": "function",
-          "function": {
-            "name": "weather_forecast",
-            "arguments": { "city": "Chicago", "days": 5 }
-          }
-        }
-      ]
+        "role": "assistant",
+        "tool_calls": [
+            {
+                "id": "call_5306",
+                "type": "function",
+                "function": {
+                    "name": "weather_forecast",
+                    "arguments": {"city": "Chicago", "days": 5},
+                },
+            }
+        ],
     },
-    {
-      "role": "tool",
-      "content": "--",
-      "tool_call_id": "call_3394"
-    },
-    {
-      "role": "assistant",
-      "content": "--",
-      "model": "gpt-3.5-turbo-0125"
-    },
-    {
-      "role": "user",
-      "content": "how is the weather in chicago for next 5 days?"
-    },
-    {
-      "role": "assistant",
-      "tool_calls": [
-        {
-          "id": "call_5306",
-          "type": "function",
-          "function": {
-            "name": "weather_forecast",
-            "arguments": { "city": "Chicago", "days": 5 }
-          }
-        }
-      ]
-    }
-  ]
-  """
+    {"role": "tool", "content": "--", "tool_call_id": "call_5306"},
+]
 
 
 def test_update_fc_history():
-    history = json.loads(test_input_history)
     message_history = []
-    for h in history:
+
+    for h in test_input_history:
         message_history.append(Message(**h))
 
-    updated_history = process_messages(message_history)
-    assert len(updated_history) == 6
+    updated_history = handler_map["Arch-Function"]._process_messages(message_history)
+    assert len(updated_history) == 7
     # ensure that tool role does not exist anymore
     assert all([h["role"] != "tool" for h in updated_history])