From 629b7a5a43d095255a7283c2b42ed2f3d272dd70 Mon Sep 17 00:00:00 2001
From: ccurme <chester.curme@gmail.com>
Date: Fri, 25 Apr 2025 14:38:23 -0400
Subject: [PATCH] openai[patch]: add explicit attribute for service tier
 (#31005)

---
 libs/langchain/tests/unit_tests/chat_models/test_base.py | 1 +
 .../partners/openai/langchain_openai/chat_models/base.py | 5 +++++
 .../tests/integration_tests/chat_models/test_base.py     | 9 ++++++---
 .../openai/tests/unit_tests/chat_models/test_base.py     | 6 ++++++
 4 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/libs/langchain/tests/unit_tests/chat_models/test_base.py b/libs/langchain/tests/unit_tests/chat_models/test_base.py
index 3ed6bc47176..9c24390de48 100644
--- a/libs/langchain/tests/unit_tests/chat_models/test_base.py
+++ b/libs/langchain/tests/unit_tests/chat_models/test_base.py
@@ -119,6 +119,7 @@ def test_configurable() -> None:
             "reasoning_effort": None,
             "frequency_penalty": None,
             "seed": None,
+            "service_tier": None,
             "logprobs": None,
             "top_logprobs": None,
             "logit_bias": None,
diff --git a/libs/partners/openai/langchain_openai/chat_models/base.py b/libs/partners/openai/langchain_openai/chat_models/base.py
index 6a5dc7a599d..9e740f38a95 100644
--- a/libs/partners/openai/langchain_openai/chat_models/base.py
+++ b/libs/partners/openai/langchain_openai/chat_models/base.py
@@ -538,6 +538,10 @@ class BaseChatOpenAI(BaseChatModel):
     However this does not prevent a user from directly passed in the parameter during
     invocation. 
     """
+    service_tier: Optional[str] = None
+    """Latency tier for request. Options are 'auto', 'default', or 'flex'. Relevant
+    for users of OpenAI's scale tier service.
+    """
 
     use_responses_api: Optional[bool] = None
     """Whether to use the Responses API instead of the Chat API.
@@ -655,6 +659,7 @@ class BaseChatOpenAI(BaseChatModel):
             "n": self.n,
             "temperature": self.temperature,
             "reasoning_effort": self.reasoning_effort,
+            "service_tier": self.service_tier,
         }
 
         params = {
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 c36c7390288..c0cbba20f0b 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
@@ -215,12 +215,15 @@ async def test_openai_abatch_tags(use_responses_api: bool) -> None:
         assert isinstance(token.text(), str)
 
 
-@pytest.mark.scheduled
+@pytest.mark.flaky(retries=3, delay=1)
 def test_openai_invoke() -> None:
     """Test invoke tokens from ChatOpenAI."""
-    llm = ChatOpenAI(max_tokens=MAX_TOKEN_COUNT)  # type: ignore[call-arg]
+    llm = ChatOpenAI(
+        model="o4-mini",
+        service_tier="flex",  # Also test service_tier
+    )
 
-    result = llm.invoke("I'm Pickle Rick", config=dict(tags=["foo"]))
+    result = llm.invoke("Hello", config=dict(tags=["foo"]))
     assert isinstance(result.content, str)
 
     # assert no response headers if include_response_headers is not set
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 a212d825405..e114ec06879 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
@@ -1732,3 +1732,9 @@ def test__construct_responses_api_input_multiple_message_types() -> None:
 
     # assert no mutation has occurred
     assert messages_copy == messages
+
+
+def test_service_tier() -> None:
+    llm = ChatOpenAI(model="o4-mini", service_tier="flex")
+    payload = llm._get_request_payload([HumanMessage("Hello")])
+    assert payload["service_tier"] == "flex"