:recycle: [Refactor] Move AVAILABLE_MODELS(_DICTS) to constants

apis/chat_api.py

@@ -5,16 +5,19 @@ import sys
 import uvicorn
 
 from pathlib import Path
+from typing import Union
+
 from fastapi import FastAPI, Depends
 from fastapi.responses import HTMLResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, Field
-from typing import Union
 from sse_starlette.sse import EventSourceResponse, ServerSentEvent
-from utils.logger import logger
-from networks.message_streamer import MessageStreamer
+
 from messagers.message_composer import MessageComposer
 from mocks.stream_chat_mocker import stream_chat_mock
+from networks.message_streamer import MessageStreamer
+from utils.logger import logger
+from constants.models import AVAILABLE_MODELS_DICTS
 
 
 class ChatAPIApp:
@@ -28,49 +31,7 @@ class ChatAPIApp:
         self.setup_routes()
 
     def get_available_models(self):
-        # https://platform.openai.com/docs/api-reference/models/list
-        # ANCHOR[id=available-models]: Available models
-        self.available_models = {
-            "object": "list",
-            "data": [
-                {
-                    "id": "mixtral-8x7b",
-                    "description": "[mistralai/Mixtral-8x7B-Instruct-v0.1]: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
-                    "object": "model",
-                    "created": 1700000000,
-                    "owned_by": "mistralai",
-                },
-                {
-                    "id": "nous-mixtral-8x7b",
-                    "description": "[NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO]: https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
-                    "object": "model",
-                    "created": 1700000000,
-                    "owned_by": "NousResearch",
-                },
-                {
-                    "id": "mistral-7b",
-                    "description": "[mistralai/Mistral-7B-Instruct-v0.2]: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
-                    "object": "model",
-                    "created": 1700000000,
-                    "owned_by": "mistralai",
-                },
-                {
-                    "id": "openchat-3.5",
-                    "description": "[openchat/openchat-3.5-0106]: https://huggingface.co/openchat/openchat-3.5-0106",
-                    "object": "model",
-                    "created": 1700000000,
-                    "owned_by": "openchat",
-                },
-                {
-                    "id": "gemma-7b",
-                    "description": "[google/gemma-7b-it]: https://huggingface.co/google/gemma-7b-it",
-                    "object": "model",
-                    "created": 1700000000,
-                    "owned_by": "Google",
-                },
-            ],
-        }
-        return self.available_models
+        return {"object": "list", "data": AVAILABLE_MODELS_DICTS}
 
     def extract_api_key(
         credentials: HTTPAuthorizationCredentials = Depends(

constants/models.py

@@ -25,3 +25,51 @@ TOKEN_LIMIT_MAP = {
 }
 
 TOKEN_RESERVED = 20
+
+
+AVAILABLE_MODELS = [
+    "mixtral-8x7b",
+    "nous-mixtral-8x7b",
+    "mistral-7b",
+    "openchat-3.5",
+    "gemma-7b",
+]
+
+# https://platform.openai.com/docs/api-reference/models/list
+AVAILABLE_MODELS_DICTS = [
+    {
+        "id": "mixtral-8x7b",
+        "description": "[mistralai/Mixtral-8x7B-Instruct-v0.1]: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
+        "object": "model",
+        "created": 1700000000,
+        "owned_by": "mistralai",
+    },
+    {
+        "id": "nous-mixtral-8x7b",
+        "description": "[NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO]: https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
+        "object": "model",
+        "created": 1700000000,
+        "owned_by": "NousResearch",
+    },
+    {
+        "id": "mistral-7b",
+        "description": "[mistralai/Mistral-7B-Instruct-v0.2]: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
+        "object": "model",
+        "created": 1700000000,
+        "owned_by": "mistralai",
+    },
+    {
+        "id": "openchat-3.5",
+        "description": "[openchat/openchat-3.5-0106]: https://huggingface.co/openchat/openchat-3.5-0106",
+        "object": "model",
+        "created": 1700000000,
+        "owned_by": "openchat",
+    },
+    {
+        "id": "gemma-7b",
+        "description": "[google/gemma-7b-it]: https://huggingface.co/google/gemma-7b-it",
+        "object": "model",
+        "created": 1700000000,
+        "owned_by": "Google",
+    },
+]

messagers/message_composer.py

@@ -1,21 +1,15 @@
 import re
 from pprint import pprint
-from utils.logger import logger
+
 from transformers import AutoTokenizer
 
+from constants.models import AVAILABLE_MODELS
+from utils.logger import logger
+
 
 class MessageComposer:
-    # LINK - apis/chat_api.py#available-models
-    AVALAIBLE_MODELS = [
-        "mixtral-8x7b",
-        "mistral-7b",
-        "nous-mixtral-8x7b",
-        "openchat-3.5",
-        "gemma-7b",
-    ]
-
     def __init__(self, model: str = None):
-        if model in self.AVALAIBLE_MODELS:
+        if model in AVAILABLE_MODELS:
             self.model = model
         else:
             self.model = "mixtral-8x7b"