fix ChatOpenAI.agenerate (#1504)

tests/integration_tests/chat_models/test_openai.py

@@ -87,3 +87,44 @@ def test_chat_openai_invalid_streaming_params() -> None:
             temperature=0,
             n=5,
         )
+
+
+@pytest.mark.asyncio
+async def test_async_chat_openai() -> None:
+    """Test async generation."""
+    chat = ChatOpenAI(max_tokens=10, n=2)
+    message = HumanMessage(content="Hello")
+    response = await chat.agenerate([[message], [message]])
+    assert isinstance(response, LLMResult)
+    assert len(response.generations) == 2
+    for generations in response.generations:
+        assert len(generations) == 2
+        for generation in generations:
+            assert isinstance(generation, ChatGeneration)
+            assert isinstance(generation.text, str)
+            assert generation.text == generation.message.content
+
+
+@pytest.mark.asyncio
+async def test_async_chat_openai_streaming() -> None:
+    """Test that streaming correctly invokes on_llm_new_token callback."""
+    callback_handler = FakeCallbackHandler()
+    callback_manager = CallbackManager([callback_handler])
+    chat = ChatOpenAI(
+        max_tokens=10,
+        streaming=True,
+        temperature=0,
+        callback_manager=callback_manager,
+        verbose=True,
+    )
+    message = HumanMessage(content="Hello")
+    response = await chat.agenerate([[message], [message]])
+    assert callback_handler.llm_streams > 0
+    assert isinstance(response, LLMResult)
+    assert len(response.generations) == 2
+    for generations in response.generations:
+        assert len(generations) == 1
+        for generation in generations:
+            assert isinstance(generation, ChatGeneration)
+            assert isinstance(generation.text, str)
+            assert generation.text == generation.message.content

tests/integration_tests/llms/test_openai.py

@@ -54,7 +54,7 @@ def test_openai_stop_error() -> None:
 
 
 def test_saving_loading_llm(tmp_path: Path) -> None:
-    """Test saving/loading an OpenAPI LLM."""
+    """Test saving/loading an OpenAI LLM."""
     llm = OpenAI(max_tokens=10)
     llm.save(file_path=tmp_path / "openai.yaml")
     loaded_llm = load_llm(tmp_path / "openai.yaml")