From 1378ddfa5fa901179a26e0e2a91cbea52de625cf Mon Sep 17 00:00:00 2001
From: Bagatur <22008038+baskaryan@users.noreply.github.com>
Date: Thu, 19 Dec 2024 11:36:49 -0800
Subject: [PATCH] openai[patch]: type reasoning_effort (#28825)

---
 .../langchain_openai/chat_models/base.py      | 11 ++++++++++
 libs/partners/openai/poetry.lock              |  8 ++++----
 libs/partners/openai/pyproject.toml           |  2 +-
 .../chat_models/test_base.py                  | 20 +++++++------------
 .../tests/unit_tests/chat_models/test_base.py |  6 ++++++
 5 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/libs/partners/openai/langchain_openai/chat_models/base.py b/libs/partners/openai/langchain_openai/chat_models/base.py
index d866e30fe36e1..142e7eca1a84b 100644
--- a/libs/partners/openai/langchain_openai/chat_models/base.py
+++ b/libs/partners/openai/langchain_openai/chat_models/base.py
@@ -454,6 +454,16 @@ class BaseChatOpenAI(BaseChatModel):
     """Total probability mass of tokens to consider at each step."""
     max_tokens: Optional[int] = Field(default=None)
     """Maximum number of tokens to generate."""
+    reasoning_effort: Optional[str] = None
+    """Constrains effort on reasoning for reasoning models. 
+    
+    o1 models only.
+
+    Currently supported values are low, medium, and high. Reducing reasoning effort 
+    can result in faster responses and fewer tokens used on reasoning in a response.
+    
+    .. versionadded:: 0.2.14
+    """
     tiktoken_model_name: Optional[str] = None
     """The model name to pass to tiktoken when using this class. 
     Tiktoken is used to count the number of tokens in documents to constrain 
@@ -599,6 +609,7 @@ def _default_params(self) -> Dict[str, Any]:
             "stop": self.stop or None,  # also exclude empty list for this
             "max_tokens": self.max_tokens,
             "extra_body": self.extra_body,
+            "reasoning_effort": self.reasoning_effort,
         }
 
         params = {
diff --git a/libs/partners/openai/poetry.lock b/libs/partners/openai/poetry.lock
index ee8fba9970162..d7610271ad57f 100644
--- a/libs/partners/openai/poetry.lock
+++ b/libs/partners/openai/poetry.lock
@@ -546,13 +546,13 @@ url = "../../standard-tests"
 
 [[package]]
 name = "langsmith"
-version = "0.2.3"
+version = "0.2.4"
 description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
 optional = false
 python-versions = "<4.0,>=3.9"
 files = [
-    {file = "langsmith-0.2.3-py3-none-any.whl", hash = "sha256:4958b6e918f57fedba6ddc55b8534d1e06478bb44c779aa73713ce898ca6ae87"},
-    {file = "langsmith-0.2.3.tar.gz", hash = "sha256:54c231b07fdeb0f8472925074a0ec0ed2cb654a0437d63c6ccf76a9da635900d"},
+    {file = "langsmith-0.2.4-py3-none-any.whl", hash = "sha256:fa797e4ecba968b76bccf351053e48bd6c6de7455515588cb46b74765e8a4127"},
+    {file = "langsmith-0.2.4.tar.gz", hash = "sha256:386fedc815b45f94fa17571860e561c0d79facc0a2979a532eb8b893a4d98fa9"},
 ]
 
 [package.dependencies]
@@ -1647,4 +1647,4 @@ watchmedo = ["PyYAML (>=3.10)"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9,<4.0"
-content-hash = "1bc76c2d222284109d4dd1bb79d309c4157cb9fe9d426870183cb32fa16465be"
+content-hash = "71de53990a6cfb9cd6a25249b40eeef52e089840a9a06b54ac556fe7fa60504c"
diff --git a/libs/partners/openai/pyproject.toml b/libs/partners/openai/pyproject.toml
index 084cbd0e739bd..3c93a2335e768 100644
--- a/libs/partners/openai/pyproject.toml
+++ b/libs/partners/openai/pyproject.toml
@@ -24,7 +24,7 @@ ignore_missing_imports = true
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
 langchain-core = "^0.3.27"
-openai = "^1.55.3"
+openai = "^1.58.1"
 tiktoken = ">=0.7,<1"
 
 [tool.ruff.lint]
diff --git a/libs/partners/openai/tests/integration_tests/chat_models/test_base.py b/libs/partners/openai/tests/integration_tests/chat_models/test_base.py
index f997ce44d6852..f2490f89501f8 100644
--- a/libs/partners/openai/tests/integration_tests/chat_models/test_base.py
+++ b/libs/partners/openai/tests/integration_tests/chat_models/test_base.py
@@ -1089,19 +1089,13 @@ async def test_astream_response_format() -> None:
         pass
 
 
-def test_o1_max_tokens() -> None:
-    response = ChatOpenAI(model="o1-mini", max_tokens=10).invoke("how are you")  # type: ignore[call-arg]
-    assert isinstance(response, AIMessage)
-
-    response = ChatOpenAI(model="gpt-4o", max_completion_tokens=10).invoke(
-        "how are you"
-    )
-    assert isinstance(response, AIMessage)
-
-
-def test_developer_message() -> None:
-    llm = ChatOpenAI(model="o1", max_tokens=10)  # type: ignore[call-arg]
-    response = llm.invoke(
+@pytest.mark.parametrize("use_max_completion_tokens", [True, False])
+def test_o1(use_max_completion_tokens: bool) -> None:
+    if use_max_completion_tokens:
+        kwargs: dict = {"max_completion_tokens": 10}
+    else:
+        kwargs = {"max_tokens": 10}
+    response = ChatOpenAI(model="o1", reasoning_effort="low", **kwargs).invoke(
         [
             {"role": "developer", "content": "respond in all caps"},
             {"role": "user", "content": "HOW ARE YOU"},
diff --git a/libs/partners/openai/tests/unit_tests/chat_models/test_base.py b/libs/partners/openai/tests/unit_tests/chat_models/test_base.py
index 6e889b80a562c..2e6cca0cd2d96 100644
--- a/libs/partners/openai/tests/unit_tests/chat_models/test_base.py
+++ b/libs/partners/openai/tests/unit_tests/chat_models/test_base.py
@@ -882,3 +882,9 @@ def test__get_request_payload() -> None:
     }
     payload = llm._get_request_payload(messages)
     assert payload == expected
+
+
+def test_init_o1() -> None:
+    with pytest.warns(None) as record:  # type: ignore[call-overload]
+        ChatOpenAI(model="o1", reasoning_effort="medium")
+    assert len(record) == 0