openai: raw response headers (#24150)

libs/partners/openai/tests/integration_tests/chat_models/test_base.py

@@ -319,6 +319,9 @@ def test_openai_invoke() -> None:
     result = llm.invoke("I'm Pickle Rick", config=dict(tags=["foo"]))
     assert isinstance(result.content, str)
 
+    # assert no response headers if include_response_headers is not set
+    assert "headers" not in result.response_metadata
+
 
 def test_stream() -> None:
     """Test streaming tokens from OpenAI."""
@@ -671,3 +674,13 @@ def test_openai_proxy() -> None:
         assert proxy.scheme == b"http"
         assert proxy.host == b"localhost"
         assert proxy.port == 8080
+
+
+def test_openai_response_headers_invoke() -> None:
+    """Test ChatOpenAI response headers."""
+    chat_openai = ChatOpenAI(include_response_headers=True)
+    result = chat_openai.invoke("I'm Pickle Rick")
+    headers = result.response_metadata["headers"]
+    assert headers
+    assert isinstance(headers, dict)
+    assert "content-type" in headers

libs/partners/openai/tests/unit_tests/chat_models/test_base.py

@@ -189,38 +189,58 @@ def mock_completion() -> dict:
     }
 
 
-def test_openai_invoke(mock_completion: dict) -> None:
+@pytest.fixture
+def mock_client(mock_completion: dict) -> MagicMock:
+    rtn = MagicMock()
+
+    mock_create = MagicMock()
+
+    mock_resp = MagicMock()
+    mock_resp.headers = {"content-type": "application/json"}
+    mock_resp.parse.return_value = mock_completion
+    mock_create.return_value = mock_resp
+
+    rtn.with_raw_response.create = mock_create
+    rtn.create.return_value = mock_completion
+    return rtn
+
+
+@pytest.fixture
+def mock_async_client(mock_completion: dict) -> AsyncMock:
+    rtn = AsyncMock()
+
+    mock_create = AsyncMock()
+    mock_resp = MagicMock()
+    mock_resp.parse.return_value = mock_completion
+    mock_create.return_value = mock_resp
+
+    rtn.with_raw_response.create = mock_create
+    rtn.create.return_value = mock_completion
+    return rtn
+
+
+def test_openai_invoke(mock_client: MagicMock) -> None:
     llm = ChatOpenAI()
-    mock_client = MagicMock()
-    completed = False
 
-    def mock_create(*args: Any, **kwargs: Any) -> Any:
-        nonlocal completed
-        completed = True
-        return mock_completion
-
-    mock_client.create = mock_create
     with patch.object(llm, "client", mock_client):
         res = llm.invoke("bar")
         assert res.content == "Bar Baz"
-    assert completed
+
+        # headers are not in response_metadata if include_response_headers not set
+        assert "headers" not in res.response_metadata
+    assert mock_client.create.called
 
 
-async def test_openai_ainvoke(mock_completion: dict) -> None:
+async def test_openai_ainvoke(mock_async_client: AsyncMock) -> None:
     llm = ChatOpenAI()
-    mock_client = AsyncMock()
-    completed = False
 
-    async def mock_create(*args: Any, **kwargs: Any) -> Any:
-        nonlocal completed
-        completed = True
-        return mock_completion
-
-    mock_client.create = mock_create
-    with patch.object(llm, "async_client", mock_client):
+    with patch.object(llm, "async_client", mock_async_client):
         res = await llm.ainvoke("bar")
         assert res.content == "Bar Baz"
-    assert completed
+
+        # headers are not in response_metadata if include_response_headers not set
+        assert "headers" not in res.response_metadata
+    assert mock_async_client.create.called
 
 
 @pytest.mark.parametrize(
@@ -239,12 +259,9 @@ def test__get_encoding_model(model: str) -> None:
     return
 
 
-def test_openai_invoke_name(mock_completion: dict) -> None:
+def test_openai_invoke_name(mock_client: MagicMock) -> None:
     llm = ChatOpenAI()
 
-    mock_client = MagicMock()
-    mock_client.create.return_value = mock_completion
-
     with patch.object(llm, "client", mock_client):
         messages = [HumanMessage(content="Foo", name="Katie")]
         res = llm.invoke(messages)