Handle models like openai/bedrock/claude-3.5-... (#2869)

* Handle models like openai/bedrock/claude-3.5-... * Fix log statement
onyx-dot-app · Oct 22, 2024 · 6e9b6a1 · 6e9b6a1
1 parent e4779c2
commit 6e9b6a1
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/backend/danswer/llm/utils.py b/backend/danswer/llm/utils.py
@@ -342,12 +342,26 @@ def get_llm_max_tokens(
 
     try:
         model_obj = model_map.get(f"{model_provider}/{model_name}")
-        if not model_obj:
-            model_obj = model_map[model_name]
-            logger.debug(f"Using model object for {model_name}")
-        else:
+        if model_obj:
             logger.debug(f"Using model object for {model_provider}/{model_name}")
 
+        if not model_obj:
+            model_obj = model_map.get(model_name)
+            if model_obj:
+                logger.debug(f"Using model object for {model_name}")
+
+        if not model_obj:
+            model_name_split = model_name.split("/")
+            if len(model_name_split) > 1:
+                model_obj = model_map.get(model_name_split[1])
+            if model_obj:
+                logger.debug(f"Using model object for {model_name_split[1]}")
+
+        if not model_obj:
+            raise RuntimeError(
+                f"No litellm entry found for {model_provider}/{model_name}"
+            )
+
         if "max_input_tokens" in model_obj:
             max_tokens = model_obj["max_input_tokens"]
             logger.info(