fix(langchain): keep trace attributes on non-chain llm runs

[hassiebp]

Important

Fixes in CallbackHandler.py ensure trace attributes are preserved for isolated non-chain LLM runs in Langchain integration.

• Behavior:
  • Modify _register_langfuse_prompt in CallbackHandler.py to return early if run_id is not provided.
  • Update __on_llm_action to keep trace attributes if LLM is run isolated and outside a chain.
• Functions:
  • Change _deregister_langfuse_prompt to check for run_id being None before accessing prompt_to_parent_run_map.
  • Update __join_tags_and_metadata to accept keep_langfuse_trace_attributes and pass it to _strip_langfuse_keys_from_dict.
  • Modify _strip_langfuse_keys_from_dict to conditionally remove trace attributes based on keep_langfuse_trace_attributes.

^{This description was created by for 8b5d63a. 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 Langfuse Langchain callback handler where trace attributes (session_id, user_id, tags) were being incorrectly stripped from LLM runs executed outside of chains. The issue occurred because the existing logic was designed to remove all Langfuse-specific metadata keys from chain runs but was also removing valuable trace context from standalone LLM calls.

The changes introduce conditional logic to distinguish between two scenarios:

1. Chain-based LLM runs: Continue stripping all Langfuse keys as before to prevent metadata pollution
2. Isolated LLM runs: Preserve trace attributes while still removing internal metadata keys

Key technical changes include:

• Adding null safety checks in prompt registration methods (_register_langfuse_prompt and _deregister_langfuse_prompt)
• Modifying __on_llm_action to detect root-level LLM runs (where parent_run_id is None) and preserve trace attributes
• Refactoring _strip_langfuse_keys_from_dict to separate internal metadata keys from user-facing trace attribute keys
• Adding a keep_langfuse_trace_attributes parameter to __join_tags_and_metadata for conditional behavior

This change integrates well with the existing Langfuse callback architecture by extending rather than replacing the current metadata handling logic. The conditional approach ensures backward compatibility for chain-based workflows while fixing the data loss issue for standalone LLM invocations.

Confidence score: 4/5

• This PR addresses a clear bug with a focused solution that maintains backward compatibility
• Score reflects solid logic and defensive programming practices, though the complexity of conditional metadata handling introduces some risk
• Pay close attention to the metadata stripping logic in _strip_langfuse_keys_from_dict to ensure the key categorization is correct

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Bot Settings | Greptile}