diff --git a/libs/partners/anthropic/Makefile b/libs/partners/anthropic/Makefile index 136c4debc85..43404f5c284 100644 --- a/libs/partners/anthropic/Makefile +++ b/libs/partners/anthropic/Makefile @@ -14,7 +14,7 @@ test tests: uv run --group test pytest -vvv --disable-socket --allow-unix-socket $(TEST_FILE) integration_test integration_tests: - uv run --group test --group test_integration pytest -vvv --timeout 30 $(TEST_FILE) + uv run --group test --group test_integration pytest -vvv --timeout 30 tests/integration_tests/test_chat_models.py::test_async_stream_twice test_watch: uv run --group test ptw --snapshot-update --now . -- -vv $(TEST_FILE) diff --git a/libs/partners/anthropic/tests/integration_tests/test_chat_models.py b/libs/partners/anthropic/tests/integration_tests/test_chat_models.py index a2563cc480d..5d3669a6acb 100644 --- a/libs/partners/anthropic/tests/integration_tests/test_chat_models.py +++ b/libs/partners/anthropic/tests/integration_tests/test_chat_models.py @@ -145,6 +145,15 @@ async def test_stream_usage_override() -> None: assert token.usage_metadata is None +async def test_async_stream_twice() -> None: + model = ChatAnthropic(model_name=MODEL_NAME) + async for token in model.astream("hi"): + assert isinstance(token, AIMessageChunk) + + async for token in model.astream("hi"): + assert isinstance(token, AIMessageChunk) + + async def test_abatch() -> None: """Test streaming tokens from ChatAnthropicMessages.""" llm = ChatAnthropicMessages(model_name=MODEL_NAME) # type: ignore[call-arg, call-arg]