Spaces:
Running
Running
| from __future__ import annotations | |
| from langchain_core.messages import AIMessageChunk | |
| from livekit.plugins.langchain import langgraph as lk_langgraph | |
| from src.agent._langchain_usage_patch import ( | |
| _PATCH_FLAG, | |
| _ORIGINAL_TO_CHAT_CHUNK, | |
| _completion_usage_from_message_chunk, | |
| apply_langchain_usage_patch, | |
| ) | |
| def test_completion_usage_is_extracted_from_usage_metadata() -> None: | |
| chunk = AIMessageChunk( | |
| content="hello", | |
| usage_metadata={ | |
| "input_tokens": 12, | |
| "output_tokens": 24, | |
| "total_tokens": 36, | |
| "input_token_details": {"cache_read": 7}, | |
| }, | |
| ) | |
| usage = _completion_usage_from_message_chunk(chunk) | |
| assert usage is not None | |
| assert usage.prompt_tokens == 12 | |
| assert usage.completion_tokens == 24 | |
| assert usage.total_tokens == 36 | |
| assert usage.prompt_cached_tokens == 7 | |
| def test_completion_usage_falls_back_to_response_metadata_token_usage() -> None: | |
| chunk = AIMessageChunk( | |
| content="hello", | |
| response_metadata={ | |
| "token_usage": { | |
| "prompt_tokens": 10, | |
| "completion_tokens": 5, | |
| "total_tokens": 15, | |
| "prompt_tokens_details": {"cached_tokens": 2}, | |
| } | |
| }, | |
| ) | |
| usage = _completion_usage_from_message_chunk(chunk) | |
| assert usage is not None | |
| assert usage.prompt_tokens == 10 | |
| assert usage.completion_tokens == 5 | |
| assert usage.total_tokens == 15 | |
| assert usage.prompt_cached_tokens == 2 | |
| def test_patch_preserves_usage_only_chunks(monkeypatch) -> None: # type: ignore[no-untyped-def] | |
| original = lk_langgraph._to_chat_chunk | |
| monkeypatch.setattr(lk_langgraph, _PATCH_FLAG, False, raising=False) | |
| monkeypatch.delattr(lk_langgraph, _ORIGINAL_TO_CHAT_CHUNK, raising=False) | |
| monkeypatch.setattr(lk_langgraph, "_to_chat_chunk", original) | |
| assert apply_langchain_usage_patch() is True | |
| assert apply_langchain_usage_patch() is False | |
| usage_only_chunk = AIMessageChunk( | |
| content="", | |
| usage_metadata={ | |
| "input_tokens": 8, | |
| "output_tokens": 3, | |
| "total_tokens": 11, | |
| }, | |
| ) | |
| chat_chunk = lk_langgraph._to_chat_chunk(usage_only_chunk) | |
| assert chat_chunk is not None | |
| assert chat_chunk.delta is None | |
| assert chat_chunk.usage is not None | |
| assert chat_chunk.usage.prompt_tokens == 8 | |
| assert chat_chunk.usage.completion_tokens == 3 | |
| # Restore to avoid bleeding monkeypatches between tests. | |
| monkeypatch.setattr(lk_langgraph, "_to_chat_chunk", original) | |
| monkeypatch.setattr(lk_langgraph, _PATCH_FLAG, False, raising=False) | |
| monkeypatch.delattr(lk_langgraph, _ORIGINAL_TO_CHAT_CHUNK, raising=False) | |