From dfb6decc1c31c05ba5661822c19b0ce956f6efe0 Mon Sep 17 00:00:00 2001 From: LUIS NOVO Date: Tue, 12 Nov 2024 12:57:06 -0300 Subject: [PATCH] wip - prepare for o1 --- open_notebook/models/llms.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/open_notebook/models/llms.py b/open_notebook/models/llms.py index 52ff94b..c501886 100644 --- a/open_notebook/models/llms.py +++ b/open_notebook/models/llms.py @@ -254,15 +254,24 @@ class OpenAILanguageModel(LanguageModel): """ Convert the language model to a LangChain chat model. """ - kwargs = self.kwargs + kwargs = self.kwargs.copy() # Make a copy to avoid modifying the original if self.json: kwargs["response_format"] = {"type": "json_object"} + # Set the token limit in kwargs with the appropriate key + if self.model_name in ["o1-mini", "o1-preview"]: + kwargs["max_completion_tokens"] = self.max_tokens + top_p = 1 + streaming = False + else: + kwargs["max_tokens"] = self.max_tokens + top_p = self.top_p + streaming = self.streaming + return ChatOpenAI( model=self.model_name, temperature=self.temperature or 0.5, - max_tokens=self.max_tokens, model_kwargs=kwargs, - streaming=self.streaming, - top_p=self.top_p, + streaming=streaming, + top_p=top_p, )