From b4f4695f16c53f5f8ba0923123a0819733d1f45f Mon Sep 17 00:00:00 2001
From: Shuguang Chen <54548843+nehcgs@users.noreply.github.com>
Date: Sun, 8 Dec 2024 09:21:53 -0800
Subject: [PATCH] Reorganize model_server

---
 model_server/.vscode/launch.json                          | 2 +-
 model_server/{app => src}/__init__.py                     | 0
 model_server/{app => src}/cli.py                          | 0
 model_server/{app => src}/commons/__init__.py             | 0
 model_server/{app => src}/commons/constants.py            | 0
 model_server/{app => src}/commons/globals.py              | 8 ++++----
 model_server/{app => src}/commons/utilities.py            | 0
 model_server/{app/model_handler => src/core}/__init__.py  | 0
 .../{app/model_handler => src/core}/base_handler.py       | 0
 .../{app/model_handler => src/core}/function_calling.py   | 3 +--
 .../{app/model_handler => src/core}/guardrails.py         | 2 +-
 .../model_handler => src/core}/hallucination_handler.py   | 3 ++-
 model_server/{app => src}/main.py                         | 6 +++---
 model_server/{app => }/tests/__init__.py                  | 0
 .../{app/tests => tests/core}/test_function_calling.py    | 4 ++--
 model_server/{app/tests => tests/core}/test_guardrails.py | 2 +-
 .../{app/tests => tests/core}/test_hallucination.py       | 2 +-
 model_server/{app/tests => tests/core}/test_state.py      | 4 ++--
 model_server/{app => }/tests/test_app.py                  | 2 +-
 model_server/{app => }/tests/test_cli_stop_server.py      | 2 +-
 20 files changed, 20 insertions(+), 20 deletions(-)
 rename model_server/{app => src}/__init__.py (100%)
 rename model_server/{app => src}/cli.py (100%)
 rename model_server/{app => src}/commons/__init__.py (100%)
 rename model_server/{app => src}/commons/constants.py (100%)
 rename model_server/{app => src}/commons/globals.py (79%)
 rename model_server/{app => src}/commons/utilities.py (100%)
 rename model_server/{app/model_handler => src/core}/__init__.py (100%)
 rename model_server/{app/model_handler => src/core}/base_handler.py (100%)
 rename model_server/{app/model_handler => src/core}/function_calling.py (99%)
 rename model_server/{app/model_handler => src/core}/guardrails.py (99%)
 rename model_server/{app/model_handler => src/core}/hallucination_handler.py (99%)
 rename model_server/{app => src}/main.py (94%)
 rename model_server/{app => }/tests/__init__.py (100%)
 rename model_server/{app/tests => tests/core}/test_function_calling.py (97%)
 rename model_server/{app/tests => tests/core}/test_guardrails.py (97%)
 rename model_server/{app/tests => tests/core}/test_hallucination.py (98%)
 rename model_server/{app/tests => tests/core}/test_state.py (93%)
 rename model_server/{app => }/tests/test_app.py (98%)
 rename model_server/{app => }/tests/test_cli_stop_server.py (98%)

diff --git a/model_server/.vscode/launch.json b/model_server/.vscode/launch.json
index dd045c6f..ca83be87 100644
--- a/model_server/.vscode/launch.json
+++ b/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"]
     }
   ]
 }
diff --git a/model_server/app/__init__.py b/model_server/src/__init__.py
similarity index 100%
rename from model_server/app/__init__.py
rename to model_server/src/__init__.py
diff --git a/model_server/app/cli.py b/model_server/src/cli.py
similarity index 100%
rename from model_server/app/cli.py
rename to model_server/src/cli.py
diff --git a/model_server/app/commons/__init__.py b/model_server/src/commons/__init__.py
similarity index 100%
rename from model_server/app/commons/__init__.py
rename to model_server/src/commons/__init__.py
diff --git a/model_server/app/commons/constants.py b/model_server/src/commons/constants.py
similarity index 100%
rename from model_server/app/commons/constants.py
rename to model_server/src/commons/constants.py
diff --git a/model_server/app/commons/globals.py b/model_server/src/commons/globals.py
similarity index 79%
rename from model_server/app/commons/globals.py
rename to model_server/src/commons/globals.py
index e62286b3..c29bc35b 100644
--- a/model_server/app/commons/globals.py
+++ b/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()
diff --git a/model_server/app/commons/utilities.py b/model_server/src/commons/utilities.py
similarity index 100%
rename from model_server/app/commons/utilities.py
rename to model_server/src/commons/utilities.py
diff --git a/model_server/app/model_handler/__init__.py b/model_server/src/core/__init__.py
similarity index 100%
rename from model_server/app/model_handler/__init__.py
rename to model_server/src/core/__init__.py
diff --git a/model_server/app/model_handler/base_handler.py b/model_server/src/core/base_handler.py
similarity index 100%
rename from model_server/app/model_handler/base_handler.py
rename to model_server/src/core/base_handler.py
diff --git a/model_server/app/model_handler/function_calling.py b/model_server/src/core/function_calling.py
similarity index 99%
rename from model_server/app/model_handler/function_calling.py
rename to model_server/src/core/function_calling.py
index 167f52ff..9eb8ccd6 100644
--- a/model_server/app/model_handler/function_calling.py
+++ b/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"]
diff --git a/model_server/app/model_handler/guardrails.py b/model_server/src/core/guardrails.py
similarity index 99%
rename from model_server/app/model_handler/guardrails.py
rename to model_server/src/core/guardrails.py
index 4f6eaf0e..0d02d0e5 100644
--- a/model_server/app/model_handler/guardrails.py
+++ b/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
diff --git a/model_server/app/model_handler/hallucination_handler.py b/model_server/src/core/hallucination_handler.py
similarity index 99%
rename from model_server/app/model_handler/hallucination_handler.py
rename to model_server/src/core/hallucination_handler.py
index 7353312a..4d923ce1 100644
--- a/model_server/app/model_handler/hallucination_handler.py
+++ b/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
diff --git a/model_server/app/main.py b/model_server/src/main.py
similarity index 94%
rename from model_server/app/main.py
rename to model_server/src/main.py
index c798bd8d..19fef239 100644
--- a/model_server/app/main.py
+++ b/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
diff --git a/model_server/app/tests/__init__.py b/model_server/tests/__init__.py
similarity index 100%
rename from model_server/app/tests/__init__.py
rename to model_server/tests/__init__.py
diff --git a/model_server/app/tests/test_function_calling.py b/model_server/tests/core/test_function_calling.py
similarity index 97%
rename from model_server/app/tests/test_function_calling.py
rename to model_server/tests/core/test_function_calling.py
index 49b1ea07..c59a28dc 100644
--- a/model_server/app/tests/test_function_calling.py
+++ b/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,
diff --git a/model_server/app/tests/test_guardrails.py b/model_server/tests/core/test_guardrails.py
similarity index 97%
rename from model_server/app/tests/test_guardrails.py
rename to model_server/tests/core/test_guardrails.py
index de9176b3..5ba7ad11 100644
--- a/model_server/app/tests/test_guardrails.py
+++ b/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 = {
diff --git a/model_server/app/tests/test_hallucination.py b/model_server/tests/core/test_hallucination.py
similarity index 98%
rename from model_server/app/tests/test_hallucination.py
rename to model_server/tests/core/test_hallucination.py
index 2afb85c9..3de80c37 100644
--- a/model_server/app/tests/test_hallucination.py
+++ b/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
diff --git a/model_server/app/tests/test_state.py b/model_server/tests/core/test_state.py
similarity index 93%
rename from model_server/app/tests/test_state.py
rename to model_server/tests/core/test_state.py
index dcb3a808..b5a0332a 100644
--- a/model_server/app/tests/test_state.py
+++ b/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 = [
diff --git a/model_server/app/tests/test_app.py b/model_server/tests/test_app.py
similarity index 98%
rename from model_server/app/tests/test_app.py
rename to model_server/tests/test_app.py
index 5784ca55..12a77772 100644
--- a/model_server/app/tests/test_app.py
+++ b/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)
diff --git a/model_server/app/tests/test_cli_stop_server.py b/model_server/tests/test_cli_stop_server.py
similarity index 98%
rename from model_server/app/tests/test_cli_stop_server.py
rename to model_server/tests/test_cli_stop_server.py
index d5bad43f..5c16475a 100644
--- a/model_server/app/tests/test_cli_stop_server.py
+++ b/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):