Reorganize model_server

model_server/.vscode/launch.json

@@ -9,7 +9,7 @@
       "type": "debugpy",
       "request": "launch",
       "module": "uvicorn",
-      "args": ["app.main:app","--reload", "--port", "51000"]
+      "args": ["src.main:app","--reload", "--port", "51000"]
     }
   ]
 }

model_server/src/commons/globals.py

@@ -1,9 +1,9 @@
-import app.commons.utilities as utils
+import src.commons.utilities as utils
 
 from openai import OpenAI
-from app.commons.constants import *
-from app.model_handler.function_calling import ArchIntentHandler, ArchFunctionHandler
-from app.model_handler.guardrails import get_guardrail_handler
+from src.commons.constants import *
+from src.core.function_calling import ArchIntentHandler, ArchFunctionHandler
+from src.core.guardrails import get_guardrail_handler
 
 
 logger = utils.get_model_server_logger()

model_server/src/core/function_calling.py

@@ -5,14 +5,13 @@ import builtins
 from openai import OpenAI
 from typing import Any, Dict, List, Tuple, Union
 from overrides import override
-from app.model_handler.base_handler import (
+from src.core.base_handler import (
     Message,
     ChatMessage,
     Choice,
     ChatCompletionResponse,
     ArchBaseHandler,
 )
-from app.function_calling.hallucination_handler import HallucinationStateHandler
 
 
 SUPPORT_DATA_TYPES = ["int", "float", "bool", "str", "list", "tuple", "set", "dict"]

model_server/src/core/guardrails.py

@@ -1,7 +1,7 @@
 import time
 import torch
 import numpy as np
-import app.commons.utilities as utils
+import src.commons.utilities as utils
 
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_server/src/core/hallucination_handler.py

@@ -1,7 +1,8 @@
 import math
 import torch
-from typing import Dict, List, Tuple
 import itertools
+
+from typing import Dict, List, Tuple
 from enum import Enum
 
 # constants

model_server/src/main.py

@@ -1,8 +1,8 @@
 import os
 
-from app.commons.globals import handler_map
-from app.model_handler.base_handler import ChatMessage
-from app.model_handler.guardrails import GuardRequest
+from src.commons.globals import handler_map
+from src.core.base_handler import ChatMessage
+from src.core.guardrails import GuardRequest
 
 from fastapi import FastAPI, Response
 from opentelemetry import trace

model_server/tests/core/test_function_calling.py

@@ -3,8 +3,8 @@ import pytest
 
 from fastapi import Response
 from unittest.mock import AsyncMock, MagicMock, patch
-from app.commons.globals import handler_map
-from app.model_handler.base_handler import (
+from src.commons.globals import handler_map
+from src.core.base_handler import (
     Message,
     ChatMessage,
     ChatCompletionResponse,

model_server/tests/core/test_guardrails.py

@@ -1,5 +1,5 @@
 from unittest.mock import patch, MagicMock
-from app.model_handler.guardrails import get_guardrail_handler
+from src.core.guardrails import get_guardrail_handler
 
 # Mock constants
 arch_guard_model_type = {

model_server/tests/core/test_hallucination.py

@@ -3,7 +3,7 @@ import pytest
 import os
 
 
-from app.model_handler.hallucination_handler import HallucinationStateHandler
+from src.core.hallucination_handler import HallucinationStateHandler
 
 
 # Get the directory of the current file

model_server/tests/core/test_state.py

@@ -1,5 +1,5 @@
-from app.commons.globals import handler_map
-from app.model_handler.function_calling import Message
+from src.commons.globals import handler_map
+from src.core.function_calling import Message
 
 
 test_input_history = [

model_server/tests/test_app.py

@@ -2,7 +2,7 @@ import pytest
 import httpx
 
 from fastapi.testclient import TestClient
-from app.main import app
+from src.main import app
 
 
 client = TestClient(app)

model_server/tests/test_cli_stop_server.py

@@ -1,7 +1,7 @@
 import unittest
 
 from unittest.mock import patch, MagicMock
-from app.cli import kill_process
+from src.cli import kill_process
 
 
 class TestStopServer(unittest.TestCase):