Merge pull request #35 from Not-Diamond/t7-fix-system-prompt-o1

o1 fixes

Author: r0ymanesco

notdiamond/llms/client.py

@@ -33,14 +33,19 @@
     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,
     model_select,
     report_latency,
 )
 from notdiamond.metrics.metric import Metric
-from notdiamond.prompts import _curly_escape, inject_system_prompt
+from notdiamond.prompts import (
+    _curly_escape,
+    inject_system_prompt,
+    o1_system_prompt_translate,
+)
 from notdiamond.types import NDApiKeyValidator
 
 LOGGER = logging.getLogger(__name__)
@@ -883,11 +888,13 @@ def invoke(
                     messages, best_llm.system_prompt
                 )
 
+            messages = o1_system_prompt_translate(messages, best_llm)
+
             self.call_callbacks("on_model_select", best_llm, best_llm.model)
 
             llm = self._llm_from_config(best_llm, callbacks=self.callbacks)
 
-            if self.tools:
+            if self.tools and not is_o1_model(best_llm):
                 llm = llm.bind_tools(self.tools)
 
             if response_model is not None:
@@ -1082,11 +1089,13 @@ async def ainvoke(
                     messages, best_llm.system_prompt
                 )
 
+            messages = o1_system_prompt_translate(messages, best_llm)
+
             self.call_callbacks("on_model_select", best_llm, best_llm.model)
 
             llm = self._llm_from_config(best_llm, callbacks=self.callbacks)
 
-            if self.tools:
+            if self.tools and not is_o1_model(best_llm):
                 llm = llm.bind_tools(self.tools)
 
             if response_model is not None:
@@ -1515,6 +1524,9 @@ def _llm_from_config(
                     "ChatOpenAI",
                     provider.provider,
                 )
+                if is_o1_model(provider):
+                    passed_kwargs["temperature"] = 1.0
+
                 return ChatOpenAI(
                     openai_api_key=provider.api_key,
                     model_name=provider.model,

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

@@ -1,7 +1,10 @@
-import re
 import logging
+import re
 from typing import Dict, List
 
+from notdiamond.llms.config import LLMConfig
+from notdiamond.llms.providers import is_o1_model
+
 LOGGER = logging.getLogger(__name__)
 LOGGER.setLevel(logging.INFO)
 
@@ -32,3 +35,19 @@ def _curly_escape(text: str) -> str:
     This function will not escape double curly braces or non-alphabetic characters.
     """
     return re.sub(r"(?<!{){([a-zA-Z])}(?!})", r"{{\1}}", text)
+
+
+def o1_system_prompt_translate(
+    messages: List[Dict[str, str]], llm: LLMConfig
+) -> List[Dict[str, str]]:
+    if is_o1_model(llm):
+        translated_messages = []
+        for msg in messages:
+            if msg["role"] == "system":
+                translated_messages.append(
+                    {"role": "user", "content": msg["content"]}
+                )
+            else:
+                translated_messages.append(msg)
+        return translated_messages
+    return messages

tests/test_llm_calls/test_openai.py

@@ -74,3 +74,19 @@ def test_response_model(self, response_model, provider):
         assert isinstance(result, response_model)
         assert result.setup
         assert result.punchline
+
+
+def test_o1_with_system_prompt():
+    provider = NDLLMProviders.O1_MINI
+    nd_llm = NotDiamond(
+        llm_configs=[provider], latency_tracking=False, hash_content=True
+    )
+    result, session_id, _ = nd_llm.invoke(
+        [
+            {"role": "system", "content": "You are a funny AI"},
+            {"role": "user", "content": "Tell me a joke"},
+        ],
+    )
+
+    assert session_id != "NO-SESSION-ID"
+    assert len(result.content) > 0