diff --git a/libs/partners/anthropic/Makefile b/libs/partners/anthropic/Makefile index 43404f5c284..136c4debc85 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 tests/integration_tests/test_chat_models.py::test_async_stream_twice + uv run --group test --group test_integration pytest -vvv --timeout 30 $(TEST_FILE) test_watch: uv run --group test ptw --snapshot-update --now . -- -vv $(TEST_FILE) diff --git a/libs/partners/anthropic/profile.yappi b/libs/partners/anthropic/profile.yappi new file mode 100644 index 00000000000..21443b3120f Binary files /dev/null and b/libs/partners/anthropic/profile.yappi differ diff --git a/libs/partners/anthropic/profile_master.yappi b/libs/partners/anthropic/profile_master.yappi new file mode 100644 index 00000000000..ceed3938204 Binary files /dev/null and b/libs/partners/anthropic/profile_master.yappi differ diff --git a/libs/partners/anthropic/test.py b/libs/partners/anthropic/test.py new file mode 100644 index 00000000000..8e155cfd661 --- /dev/null +++ b/libs/partners/anthropic/test.py @@ -0,0 +1,46 @@ +import cProfile +import asyncio + +from langchain_anthropic import ChatAnthropic +from langchain_core.messages import AIMessageChunk + + +MODEL_NAME = "claude-3-5-haiku-latest" + + +async def test_stream_usage() -> None: + """Test usage metadata can be excluded.""" + model = ChatAnthropic(model_name=MODEL_NAME, stream_usage=False) + async for token in model.astream("hi"): + assert isinstance(token, AIMessageChunk) + assert token.usage_metadata is None + # check we override with kwarg + model = ChatAnthropic(model_name=MODEL_NAME) + assert model.stream_usage + async for token in model.astream("hi", stream_usage=False): + assert isinstance(token, AIMessageChunk) + assert token.usage_metadata is None + + +# def run_async(func, *args, **kwargs): +# return asyncio.get_event_loop().run_until_complete(func(*args, **kwargs)) + +# if __name__ == "__main__": +# profiler = cProfile.Profile() +# profiler.enable() +# run_async(test_stream_usage) # your async function here +# profiler.disable() +# profiler.dump_stats('test_stream_usage_master.prof') # or 'test_stream_usage.out' + + + +import yappi + + +def run_async(func, *args, **kwargs): + return asyncio.get_event_loop().run_until_complete(func(*args, **kwargs)) + +yappi.start() +run_async(test_stream_usage) +yappi.stop() +yappi.get_func_stats().save('profile.yappi', type='pstat') # pstat format diff --git a/libs/partners/anthropic/test_stream_usage.prof b/libs/partners/anthropic/test_stream_usage.prof new file mode 100644 index 00000000000..7db1f50c92d Binary files /dev/null and b/libs/partners/anthropic/test_stream_usage.prof differ diff --git a/libs/partners/anthropic/test_stream_usage_master.prof b/libs/partners/anthropic/test_stream_usage_master.prof new file mode 100644 index 00000000000..2d065aa738a Binary files /dev/null and b/libs/partners/anthropic/test_stream_usage_master.prof differ