fix(langchain): capture usage on streamed gemini responses

[hassiebp]

Important

Fixes usage parsing bug in CallbackHandler.py for streamed Gemini responses, ensuring valid usage data is captured.

• Behavior:
  • Fixes usage parsing bug in _parse_usage in CallbackHandler.py for streamed Gemini responses.
  • Loop now breaks only when _parse_usage_model() returns non-None, ensuring valid usage data is found.
• Tests:
  • Marks test_basic_chat_openai in test_langchain.py as flaky.

^{This description was created by for 15ff192. You can customize this summary. It will automatically update as commits are pushed.}

---

Disclaimer: Experimental PR review

Greptile Summary

This PR fixes a bug in the Langchain integration's usage parsing logic for streamed Gemini responses. The change modifies the _parse_usage function in CallbackHandler.py to improve how it handles usage metadata extraction from generation chunks.

The core issue was in the loop that iterates through generation chunks looking for usage metadata. Previously, the function would break out of the loop immediately after finding any chunk with usage_metadata, regardless of whether that metadata was valid or usable. The fix adds a conditional check that only breaks the loop when _parse_usage_model() returns a non-None value, indicating that valid usage data was actually found.

This change is specifically important for Gemini streaming responses, where early chunks may contain empty or invalid usage metadata, while later chunks contain the actual token usage information. The Langfuse Usage model (as seen in the context) tracks input/output tokens and associated costs, making accurate usage capture critical for billing and monitoring purposes. By ensuring the function continues searching until it finds valid usage data, this fix enables proper token usage tracking for streamed Gemini responses in the Langchain integration.

Confidence score: 4/5

• This PR addresses a specific, well-defined bug with a minimal and focused change
• Score reflects the targeted nature of the fix and clear understanding of the problem being solved
• Pay close attention to the modified usage parsing logic in CallbackHandler.py

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Bot Settings | Greptile}

[rvndbalaji]

Hi @hassiebp , Im eagerly waiting for this to be merged, any ETA on when this might be merged and released?

[hassiebp]

@rvndbalaji This is released in https://github.com/langfuse/langfuse-python/releases/tag/v3.3.1 👍🏾