Moving is_o1_model to providers.py

Author: acompa

notdiamond/llms/client.py

@@ -33,6 +33,7 @@
     MissingLLMConfigs,
 )
 from notdiamond.llms.config import LLMConfig
+from notdiamond.llms.providers import is_o1_model
 from notdiamond.llms.request import (
     amodel_select,
     create_preference_id,
@@ -42,7 +43,6 @@
 from notdiamond.metrics.metric import Metric
 from notdiamond.prompts import (
     _curly_escape,
-    _is_o1_model,
     inject_system_prompt,
     o1_system_prompt_translate,
 )
@@ -894,7 +894,7 @@ def invoke(
 
             llm = self._llm_from_config(best_llm, callbacks=self.callbacks)
 
-            if self.tools and not _is_o1_model(best_llm):
+            if self.tools and not is_o1_model(best_llm):
                 llm = llm.bind_tools(self.tools)
 
             if response_model is not None:
@@ -1095,7 +1095,7 @@ async def ainvoke(
 
             llm = self._llm_from_config(best_llm, callbacks=self.callbacks)
 
-            if self.tools and not _is_o1_model(best_llm):
+            if self.tools and not is_o1_model(best_llm):
                 llm = llm.bind_tools(self.tools)
 
             if response_model is not None:
@@ -1524,7 +1524,7 @@ def _llm_from_config(
                     "ChatOpenAI",
                     provider.provider,
                 )
-                if _is_o1_model(provider):
+                if is_o1_model(provider):
                     passed_kwargs["temperature"] = 1.0
 
                 return ChatOpenAI(

notdiamond/llms/providers.py

@@ -175,3 +175,12 @@ class NDLLMProviders(Enum):
 
     def __new__(cls, provider, model):
         return LLMConfig(provider=provider, model=model)
+
+
+def is_o1_model(llm: LLMConfig):
+    return llm in (
+        NDLLMProviders.O1_PREVIEW,
+        NDLLMProviders.O1_PREVIEW_2024_09_12,
+        NDLLMProviders.O1_MINI,
+        NDLLMProviders.O1_MINI_2024_09_12,
+    )

notdiamond/prompts.py

@@ -3,7 +3,7 @@
 from typing import Dict, List
 
 from notdiamond.llms.config import LLMConfig
-from notdiamond.llms.providers import NDLLMProviders
+from notdiamond.llms.providers import is_o1_model
 
 LOGGER = logging.getLogger(__name__)
 LOGGER.setLevel(logging.INFO)
@@ -37,21 +37,10 @@ def _curly_escape(text: str) -> str:
     return re.sub(r"(?<!{){([a-zA-Z])}(?!})", r"{{\1}}", text)
 
 
-def _is_o1_model(llm: LLMConfig):
-    if llm in (
-        NDLLMProviders.O1_PREVIEW,
-        NDLLMProviders.O1_PREVIEW_2024_09_12,
-        NDLLMProviders.O1_MINI,
-        NDLLMProviders.O1_MINI_2024_09_12,
-    ):
-        return True
-    return False
-
-
 def o1_system_prompt_translate(
     messages: List[Dict[str, str]], llm: LLMConfig
 ) -> List[Dict[str, str]]:
-    if _is_o1_model(llm):
+    if is_o1_model(llm):
         translated_messages = []
         for msg in messages:
             if msg["role"] == "system":