bugfix

mckinsey · Jun 26, 2024 · c756f55 · c756f55
1 parent c385704
commit c756f55
Showing 1 changed file with 11 additions and 9 deletions.
diff --git a/vizro-ai/src/vizro_ai/chains/_llm_models.py b/vizro-ai/src/vizro_ai/chains/_llm_models.py
@@ -9,7 +9,7 @@
     pass
 
 # TODO add new wrappers in if new model support is added
-LLM_MODELS = Union[ChatOpenAI, Optional[ChatAnthropic]]
+LLM_MODELS = Union[ChatOpenAI]
 
 # TODO constant of model inventory, can be converted to yaml and link to docs
 PREDEFINED_MODELS: Dict[str, Dict[str, Union[int, LLM_MODELS]]] = {
@@ -45,16 +45,18 @@
         "max_tokens": 128000,
         "wrapper": ChatOpenAI,
     },
-    "claude-3-haiku-20240307": {
-        "max_tokens": 200000,
-        "wrapper": ChatAnthropic,
-    },
-    "claude-3-sonnet-20240229": {
-        "max_tokens": 200000,
-        "wrapper": ChatAnthropic,
-    },
 }
 
+# TODO add new wrappers in if new model support is added
+if "ChatAnthropic" in globals():
+    LLM_MODELS = Union[ChatOpenAI, ChatAnthropic]
+    PREDEFINED_MODELS = {
+        **PREDEFINED_MODELS,
+        **{"claude-3-haiku-20240307": {"max_tokens": 200000, "wrapper": ChatAnthropic}},
+        **{"claude-3-sonnet-20240229": {"max_tokens": 200000, "wrapper": ChatAnthropic}},
+    }
+
+
 DEFAULT_MODEL = "gpt-3.5-turbo"
 DEFAULT_TEMPERATURE = 0