Add support for gpt-4o-mini

open-cogsci · Jul 19, 2024 · d4fa3b7 · d4fa3b7
1 parent 9b2255c
commit d4fa3b7
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/sigmund/config.py b/sigmund/config.py
@@ -89,10 +89,12 @@
 mistral_api_key = os.environ.get('MISTRAL_API_KEY', None)
 # Supported models are currently:
 # - gpt-3.5
-# - gpt-4
+# - gpt-4o
+# - gpt-4o-mini
 # - claude-2.1
 # - claude-3-sonnet
 # - claude-3-opus
+# - claude-3.5-sonnet
 # - mistral-tiny
 # - mistral-small
 # - mistral-medium
@@ -106,9 +108,9 @@
 #   model
 model_config = {
     'openai': {
-        'search_model': 'gpt-3.5',
-        'condense_model': 'gpt-3.5',
-        'public_model': 'claude-3-haiku',
+        'search_model': 'gpt-4o-mini',
+        'condense_model': 'gpt-4o-mini',
+        'public_model': 'gpt-4o-mini',
         'answer_model': 'gpt-4'
     },
     'anthropic': {
@@ -120,7 +122,7 @@
     'mistral': {
         'search_model': 'mistral-large',
         'condense_model': 'mistral-medium',
-        'public_model': 'claude-3-haiku',
+        'public_model': 'gpt-4o-mini',
         'answer_model': 'mistral-large'
     },
     'dummy': {

diff --git a/sigmund/model/__init__.py b/sigmund/model/__init__.py
@@ -3,9 +3,12 @@
 
 def model(sigmund, model, **kwargs):
     """A factory function that returns a Model instance."""
-    if model == 'gpt-4':
+    if model in ('gpt-4', 'gpt-4o'):
         from ._openai_model import OpenAIModel
         return OpenAIModel(sigmund, 'gpt-4o', **kwargs)
+    if model == 'gpt-4o-mini':
+        from ._openai_model import OpenAIModel
+        return OpenAIModel(sigmund, 'gpt-4o-mini', **kwargs)
     if model == 'gpt-3.5':
         from ._openai_model import OpenAIModel
         return OpenAIModel(sigmund, 'gpt-3.5-turbo', **kwargs)