Optional setting MODELS_WHITELIST added. (#25)

README.md

@@ -125,6 +125,7 @@ You can configure Chibi using the following environment variables:
 | MAX_TOKENS                    | Maximum tokens for the OpenAI Client                                                                                                                              | No       | `1000`                                                                           |
 | MESSAGE_FOR_DISALLOWED_USERS  | Message to show disallowed users                                                                                                                                  | No       | `You're not allowed to interact with me, sorry. Contact my owner first, please.` |
 | MODEL_DEFAULT                 | Default OpenAI model to use                                                                                                                                       | No       | `gpt-3.5-turbo`                                                                  |
+| MODELS_WHITELIST              | Comma-separated list of allowed models, i.e.: "gpt-4,gpt-3.5-turbo"                                                                                               | No       |                                                                                  |
 | OPENAI_FREQUENCY_PENALTY      | OpenAI frequency penalty                                                                                                                                          | No       | `0`                                                                              |
 | OPENAI_IMAGE_N_CHOICES        | Number of choices for image generation in DALL-E                                                                                                                  | No       | `4`                                                                              |
 | OPENAI_PRESENCE_PENALTY       | OpenAI presence penalty                                                                                                                                           | No       | `0`                                                                              |

changelog.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.6.2] - 2024-02-18
+
+### Added
+- Optional `MODELS_WHITELIST` setting to limit the number of available GPT models.
+
 ## [0.6.1] - 2024-01-12
 
 ### Added
@@ -159,7 +164,12 @@ applied.
 - Flake8 and Mypy setups.
 - GitHub Action for linters.
 
-[Unreleased]: https://github.com/s-nagaev/chibi/compare/v0.5.2...HEAD
+[Unreleased]: https://github.com/s-nagaev/chibi/compare/v0.6.3...HEAD
+[0.6.3]: https://github.com/s-nagaev/chibi/compare/v0.6.2...v0.6.3
+[0.6.2]: https://github.com/s-nagaev/chibi/compare/v0.6.1...v0.6.2
+[0.6.1]: https://github.com/s-nagaev/chibi/compare/v0.6.0...v0.6.1
+[0.6.0]: https://github.com/s-nagaev/chibi/compare/v0.5.3...v0.6.0
+[0.5.3]: https://github.com/s-nagaev/chibi/compare/v0.5.2...v0.5.3
 [0.5.2]: https://github.com/s-nagaev/chibi/compare/v0.5.1...v0.5.2
 [0.5.1]: https://github.com/s-nagaev/chibi/compare/v0.5.0...v0.5.1
 [0.5.0]: https://github.com/s-nagaev/chibi/compare/v0.4.0...v0.5.0

chibi/config/gpt.py

@@ -28,6 +28,7 @@ class GPTSettings(BaseSettings):
     max_history_tokens: int = Field(env="MAX_HISTORY_TOKENS", default=1800)
     max_tokens: int = Field(env="MAX_TOKENS", default=1000)
     model_default: str = Field(env="MODEL_DEFAULT", default="gpt-3.5-turbo")
+    models_whitelist: list[str] = Field(env="MODELS_WHITELIST", default_factory=list)
     presence_penalty: float = Field(env="OPENAI_PRESENCE_PENALTY", default=0)
     proxy: str | None = Field(env="PROXY", default=None)
     temperature: float = Field(env="OPENAI_TEMPERATURE", default=1)
@@ -38,7 +39,7 @@ class Config:
 
         @classmethod
         def parse_env_var(cls, field_name: str, raw_val: str) -> Any:
-            if field_name == "image_generations_whitelist":
+            if field_name in ("image_generations_whitelist", "models_whitelist"):
                 return raw_val.split(",")
             if field_name == "gpt4_whitelist":
                 return [str(username).strip().strip("@") for username in raw_val.split(",")]

chibi/models.py

@@ -27,7 +27,25 @@ class User(BaseModel):
 
     @property
     def model(self) -> str:
-        return self.gpt_model or gpt_settings.model_default
+        """Get user's preferred model.
+
+        If user never set the preferred model, returns the default one.
+        If user's preferred model is not compatible with the models whitelist, returns the default one.
+        If none of the previous conditions are met, returns user's preferred model.
+
+        Returns:
+            The GPT model name.
+        """
+        if not self.gpt_model:
+            return gpt_settings.model_default
+
+        if not gpt_settings.models_whitelist:
+            return self.gpt_model
+
+        if self.gpt_model in gpt_settings.models_whitelist:
+            return self.gpt_model
+
+        return gpt_settings.model_default
 
     @property
     def has_reached_image_limits(self) -> bool:

chibi/services/gpt.py

@@ -61,8 +61,13 @@ async def api_key_is_valid(api_key: str) -> bool:
 
 async def retrieve_available_models(api_key: str, include_gpt4: bool) -> list[str]:
     client = AsyncOpenAI(api_key=api_key)
-    all_models = await client.models.list()
+    available_models = await client.models.list()
+
+    if gpt_settings.models_whitelist:
+        allowed_model_names = [model.id for model in available_models.data if model.id in gpt_settings.models_whitelist]
+    else:
+        allowed_model_names = [model.id for model in available_models.data]
 
     if include_gpt4:
-        return sorted([model.id for model in all_models.data if "gpt" in model.id])
-    return sorted([model.id for model in all_models.data if "gpt-3" in model.id])
+        return sorted([model for model in allowed_model_names if "gpt" in model])
+    return sorted([model for model in allowed_model_names if "gpt-3" in model])

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "chibi"
-version = "0.6.0"
+version = "0.6.2"
 description = "Another one ChatGPT Telegram bot"
 authors = ["Sergei Nagaev <nagaev.sv@gmail.com>"]
 license = "MIT"