Feat/openai llm trial paid config (langgenius#2545)

to-matsu · Feb 23, 2024 · 3ad943a · 3ad943a
1 parent 3082093
commit 3ad943a
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 24 deletions.
diff --git a/api/.env.example b/api/.env.example
@@ -111,8 +111,10 @@ HOSTED_OPENAI_API_KEY=
 HOSTED_OPENAI_API_BASE=
 HOSTED_OPENAI_API_ORGANIZATION=
 HOSTED_OPENAI_TRIAL_ENABLED=false
+HOSTED_OPENAI_TRIAL_MODELS=gpt-3.5-turbo,gpt-3.5-turbo-1106,gpt-3.5-turbo-instruct,gpt-3.5-turbo-16k,gpt-3.5-turbo-16k-0613,gpt-3.5-turbo-0613,gpt-3.5-turbo-0125,text-davinci-003
 HOSTED_OPENAI_QUOTA_LIMIT=200
 HOSTED_OPENAI_PAID_ENABLED=false
+HOSTED_OPENAI_PAID_MODELS=gpt-4,gpt-4-turbo-preview,gpt-4-32k,gpt-4-1106-preview,gpt-4-0125-preview,gpt-3.5-turbo,gpt-3.5-turbo-16k,gpt-3.5-turbo-16k-0613,gpt-3.5-turbo-1106,gpt-3.5-turbo-0613,gpt-3.5-turbo-0125,gpt-3.5-turbo-instruct,text-davinci-003
 
 HOSTED_AZURE_OPENAI_ENABLED=false
 HOSTED_AZURE_OPENAI_API_KEY=

diff --git a/api/config.py b/api/config.py
@@ -38,7 +38,9 @@
     'LOG_LEVEL': 'INFO',
     'HOSTED_OPENAI_QUOTA_LIMIT': 200,
     'HOSTED_OPENAI_TRIAL_ENABLED': 'False',
+    'HOSTED_OPENAI_TRIAL_MODELS': '',
     'HOSTED_OPENAI_PAID_ENABLED': 'False',
+    'HOSTED_OPENAI_PAID_MODELS': '',
     'HOSTED_AZURE_OPENAI_ENABLED': 'False',
     'HOSTED_AZURE_OPENAI_QUOTA_LIMIT': 200,
     'HOSTED_ANTHROPIC_QUOTA_LIMIT': 600000,
@@ -261,8 +263,10 @@ def __init__(self):
         self.HOSTED_OPENAI_API_BASE = get_env('HOSTED_OPENAI_API_BASE')
         self.HOSTED_OPENAI_API_ORGANIZATION = get_env('HOSTED_OPENAI_API_ORGANIZATION')
         self.HOSTED_OPENAI_TRIAL_ENABLED = get_bool_env('HOSTED_OPENAI_TRIAL_ENABLED')
+        self.HOSTED_OPENAI_TRIAL_MODELS = get_env('HOSTED_OPENAI_TRIAL_MODELS')
         self.HOSTED_OPENAI_QUOTA_LIMIT = int(get_env('HOSTED_OPENAI_QUOTA_LIMIT'))
         self.HOSTED_OPENAI_PAID_ENABLED = get_bool_env('HOSTED_OPENAI_PAID_ENABLED')
+        self.HOSTED_OPENAI_PAID_MODELS = get_env('HOSTED_OPENAI_PAID_MODELS')
 
         self.HOSTED_AZURE_OPENAI_ENABLED = get_bool_env('HOSTED_AZURE_OPENAI_ENABLED')
         self.HOSTED_AZURE_OPENAI_API_KEY = get_env('HOSTED_AZURE_OPENAI_API_KEY')

diff --git a/api/core/hosting_configuration.py b/api/core/hosting_configuration.py
@@ -104,37 +104,17 @@ def init_openai(self, app_config: Config) -> HostingProvider:
 
         if app_config.get("HOSTED_OPENAI_TRIAL_ENABLED"):
             hosted_quota_limit = int(app_config.get("HOSTED_OPENAI_QUOTA_LIMIT", "200"))
+            trial_models = self.parse_restrict_models_from_env(app_config, "HOSTED_OPENAI_TRIAL_MODELS")
             trial_quota = TrialHostingQuota(
                 quota_limit=hosted_quota_limit,
-                restrict_models=[
-                    RestrictModel(model="gpt-3.5-turbo", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-1106", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-instruct", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-16k", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-16k-0613", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-0613", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-0125", model_type=ModelType.LLM),
-                    RestrictModel(model="text-davinci-003", model_type=ModelType.LLM),
-                ]
+                restrict_models=trial_models
             )
             quotas.append(trial_quota)
 
         if app_config.get("HOSTED_OPENAI_PAID_ENABLED"):
+            paid_models = self.parse_restrict_models_from_env(app_config, "HOSTED_OPENAI_PAID_MODELS")
             paid_quota = PaidHostingQuota(
-                restrict_models=[
-                    RestrictModel(model="gpt-4", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-4-turbo-preview", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-4-1106-preview", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-4-0125-preview", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-16k", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-16k-0613", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-1106", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-0613", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-0125", model_type=ModelType.LLM),
-                    RestrictModel(model="gpt-3.5-turbo-instruct", model_type=ModelType.LLM),
-                    RestrictModel(model="text-davinci-003", model_type=ModelType.LLM),
-                ]
+                restrict_models=paid_models
             )
             quotas.append(paid_quota)
 
@@ -258,3 +238,11 @@ def init_moderation_config(self, app_config: Config) -> HostedModerationConfig:
         return HostedModerationConfig(
             enabled=False
         )
+
+    @staticmethod
+    def parse_restrict_models_from_env(app_config: Config, env_var: str) -> list[RestrictModel]:
+        models_str = app_config.get(env_var)
+        models_list = models_str.split(",") if models_str else []
+        return [RestrictModel(model=model_name.strip(), model_type=ModelType.LLM) for model_name in models_list if
+                model_name.strip()]
+