Model features: native async

libs/ai-endpoints/tests/integration_tests/test_chat_models.py

@@ -11,6 +11,7 @@
     HumanMessage,
     SystemMessage,
 )
+from langchain_core.outputs import ChatGeneration, LLMResult
 
 from langchain_nvidia_ai_endpoints.chat_models import ChatNVIDIA
 
@@ -441,3 +442,33 @@ def test_stop(
             assert isinstance(token.content, str)
             result += f"{token.content}|"
     assert all(target not in result for target in targets)
+
+
+def test_generate(chat_model: str, mode: dict) -> None:
+    """Test generate method of anthropic."""
+    chat = ChatNVIDIA(model=chat_model, **mode)  # type: ignore[call-arg]
+    chat_messages: List[List[BaseMessage]] = [
+        [HumanMessage(content="How many toes do dogs have?")]
+    ]
+    messages_copy = [messages.copy() for messages in chat_messages]
+    result: LLMResult = chat.generate(chat_messages)
+    assert isinstance(result, LLMResult)
+    for response in result.generations[0]:
+        assert isinstance(response, ChatGeneration)
+        assert isinstance(response.text, str)
+        assert response.text == response.message.content
+    assert chat_messages == messages_copy
+
+
+# @pytest.mark.scheduled
+async def test_async_generate(chat_model: str, mode: dict) -> None:
+    """Test async generation."""
+    llm = ChatNVIDIA(model=chat_model, **mode)
+    message = HumanMessage(content="Hello")
+    response = await llm.agenerate([[message]])
+    assert isinstance(response, LLMResult)
+    for generations in response.generations:
+        for generation in generations:
+            assert isinstance(generation, ChatGeneration)
+            assert isinstance(generation.text, str)
+            assert generation.text == generation.message.content