From 98ba7a8a5cc581c3096757092ff8e3edb8ec9b0e Mon Sep 17 00:00:00 2001 From: Chester Curme Date: Sat, 18 Jan 2025 17:25:56 -0500 Subject: [PATCH] add test for merging --- .../tests/unit_tests/chat_models/test_perplexity.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/libs/community/tests/unit_tests/chat_models/test_perplexity.py b/libs/community/tests/unit_tests/chat_models/test_perplexity.py index 225cc44d5d4b5..d4991de422a36 100644 --- a/libs/community/tests/unit_tests/chat_models/test_perplexity.py +++ b/libs/community/tests/unit_tests/chat_models/test_perplexity.py @@ -1,10 +1,11 @@ """Test Perplexity Chat API wrapper.""" import os -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from unittest.mock import MagicMock import pytest +from langchain_core.messages import AIMessageChunk, BaseMessageChunk from pytest_mock import MockerFixture from langchain_community.chat_models import ChatPerplexity @@ -82,7 +83,9 @@ def test_perplexity_stream_includes_citations(mocker: MockerFixture) -> None: llm.client.chat.completions, "create", return_value=mock_stream ) stream = llm.stream("Hello langchain") + full: Optional[BaseMessageChunk] = None for i, chunk in enumerate(stream): + full = chunk if full is None else full + chunk assert chunk.content == mock_chunks[i]["choices"][0]["delta"]["content"] if i == 0: assert chunk.additional_kwargs["citations"] == [ @@ -91,5 +94,8 @@ def test_perplexity_stream_includes_citations(mocker: MockerFixture) -> None: ] else: assert "citations" not in chunk.additional_kwargs + assert isinstance(full, AIMessageChunk) + assert full.content == "Hello Perplexity" + assert full.additional_kwargs == {"citations": ["example.com", "example2.com"]} patcher.assert_called_once()