From a599176bbf0459bc067213c9f55e4815e22854dc Mon Sep 17 00:00:00 2001
From: Chris Weaver <25087905+Weves@users.noreply.github.com>
Date: Tue, 3 Jun 2025 13:01:12 -0700
Subject: [PATCH] Improve reasoning detection (#4817)

* Improve reasoning detection

* Address greptile comments

* Fix mypy
---
 backend/onyx/llm/chat_llm.py                  |  2 +-
 backend/onyx/llm/utils.py                     | 40 ++++++++++++++-----
 .../unit/onyx/llm/test_model_is_reasoning.py  | 33 +++++++++++++++
 3 files changed, 65 insertions(+), 10 deletions(-)
 create mode 100644 backend/tests/unit/onyx/llm/test_model_is_reasoning.py

diff --git a/backend/onyx/llm/chat_llm.py b/backend/onyx/llm/chat_llm.py
index e365d782cb..556a37f1c2 100644
--- a/backend/onyx/llm/chat_llm.py
+++ b/backend/onyx/llm/chat_llm.py
@@ -264,7 +264,7 @@ class DefaultMultiLLM(LLM):
     ):
         self._timeout = timeout
         if timeout is None:
-            if model_is_reasoning_model(model_name):
+            if model_is_reasoning_model(model_name, model_provider):
                 self._timeout = QA_TIMEOUT * 10  # Reasoning models are slow
             else:
                 self._timeout = QA_TIMEOUT
diff --git a/backend/onyx/llm/utils.py b/backend/onyx/llm/utils.py
index 9674eb4a4a..730abbfe2b 100644
--- a/backend/onyx/llm/utils.py
+++ b/backend/onyx/llm/utils.py
@@ -663,12 +663,34 @@ def model_supports_image_input(model_name: str, model_provider: str) -> bool:
         return False
 
 
-def model_is_reasoning_model(model_name: str) -> bool:
-    _REASONING_MODEL_NAMES = [
-        "o1",
-        "o1-mini",
-        "o3-mini",
-        "deepseek-reasoner",
-        "deepseek-r1",
-    ]
-    return model_name.lower() in _REASONING_MODEL_NAMES
+def model_is_reasoning_model(model_name: str, model_provider: str) -> bool:
+    model_map = get_model_map()
+    try:
+        model_obj = find_model_obj(
+            model_map,
+            model_provider,
+            model_name,
+        )
+        if model_obj and "supports_reasoning" in model_obj:
+            return model_obj["supports_reasoning"]
+
+        # Fallback: try using litellm.supports_reasoning() for newer models
+        try:
+            logger.debug("Falling back to `litellm.supports_reasoning`")
+            full_model_name = (
+                f"{model_provider}/{model_name}"
+                if model_provider not in model_name
+                else model_name
+            )
+            return litellm.supports_reasoning(model=full_model_name)
+        except Exception:
+            logger.exception(
+                f"Failed to check if {model_provider}/{model_name} supports reasoning"
+            )
+            return False
+
+    except Exception:
+        logger.exception(
+            f"Failed to get model object for {model_provider}/{model_name}"
+        )
+        return False
diff --git a/backend/tests/unit/onyx/llm/test_model_is_reasoning.py b/backend/tests/unit/onyx/llm/test_model_is_reasoning.py
new file mode 100644
index 0000000000..5d7fbb048f
--- /dev/null
+++ b/backend/tests/unit/onyx/llm/test_model_is_reasoning.py
@@ -0,0 +1,33 @@
+from onyx.llm.utils import model_is_reasoning_model
+
+
+def test_model_is_reasoning_model() -> None:
+    """Test that reasoning models are correctly identified and non-reasoning models are not"""
+
+    # Models that should be identified as reasoning models
+    reasoning_models = [
+        ("o3", "openai"),
+        ("o3-mini", "openai"),
+        ("o4-mini", "openai"),
+        ("deepseek-reasoner", "deepseek"),
+        ("deepseek-r1", "openrouter/deepseek"),
+        ("claude-sonnet-4-20250514", "anthropic"),
+    ]
+
+    # Models that should NOT be identified as reasoning models
+    non_reasoning_models = [
+        ("gpt-4o", "openai"),
+        ("claude-3-5-sonnet-20240620", "anthropic"),
+    ]
+
+    # Test reasoning models
+    for model_name, provider in reasoning_models:
+        assert (
+            model_is_reasoning_model(model_name, provider) is True
+        ), f"Expected {provider}/{model_name} to be identified as a reasoning model"
+
+    # Test non-reasoning models
+    for model_name, provider in non_reasoning_models:
+        assert (
+            model_is_reasoning_model(model_name, provider) is False
+        ), f"Expected {provider}/{model_name} to NOT be identified as a reasoning model"