refactor: (models) use regex for model ID filtering in OpenAIChatBot

yufeikang · Dec 26, 2024 · 006fac7 · 006fac7
1 parent 1e6aec9
commit 006fac7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/app/models.py b/app/models.py
@@ -2,6 +2,7 @@
 import json
 import logging
 import os
+import re
 from functools import cache
 
 import anthropic
@@ -347,7 +348,7 @@ async def get_models(self):
         openai_models = (await self.openai_client.models.list()).data
         models = []
         for model in openai_models:
-            if not model.id.startswith("gpt-4") and not model.id.startswith("o"):
+            if not re.match(r"gpt-\d", model.id) and not re.match(r"o\d", model.id):
                 # skip other models
                 logger.debug(f"Skipping model: {model.id}")
                 continue