fix(langchain): do not end span on agent callbacks

[hassiebp]

Important

Improved handling of agent callbacks in CallbackHandler.py by removing exceptions for missing run IDs and preventing premature span endings.

• Behavior:
  • Removed exceptions for missing run_id in on_retriever_error, on_agent_action, on_agent_finish, on_chain_end, and on_llm_end in CallbackHandler.py.
  • Modified on_agent_action and on_agent_finish to update spans without ending them.
  • Updated on_llm_end to handle response extraction and span updates without checking for run_id existence.
• Misc:
  • Added comments in on_agent_finish to clarify handling of run ID overlap between agent finish and chain end.

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

Disclaimer: Experimental PR review

Greptile Summary

This PR fixes a critical issue in the LangChain integration's callback handler related to agent execution flow. The changes address premature span ending and improve robustness when handling LangChain's agent callbacks.

Key Changes Made:

1. Removed premature span ending: The .end() calls have been removed from both on_agent_action and on_agent_finish methods. This prevents spans from being terminated too early in the agent execution lifecycle.

2. Improved error handling: Replaced strict dictionary access (self.runs[run_id]) with graceful lookups (self.runs.get(run_id, None)) throughout the callback handler to handle cases where run IDs might not be present.

3. Eliminated strict validation: Removed several exception-raising checks when run IDs are not found, making the handler more resilient to unexpected callback ordering or missing runs.

4. Clarified execution flow: Added documentation explaining that LangChain sends the same run ID for both agent finish and chain end events, so cleanup is properly deferred to the chain end callback.

The root issue stems from LangChain's agent implementation sending duplicate run IDs across different callback events, which was causing double cleanup attempts and premature span termination. By deferring cleanup to the appropriate chain end callback and using defensive programming practices, this fix ensures proper observation lifecycle management within the Langfuse tracing system.

This change integrates with the existing callback system by maintaining the same interface while making the internal handling more robust against LangChain's complex execution patterns.

Confidence score: 4/5

• This PR addresses a specific, well-understood issue with defensive programming changes that reduce crash potential
• Score reflects solid understanding of the LangChain callback flow and appropriate fixes for the identified timing issues
• Pay close attention to CallbackHandler.py to ensure agent callback flows work correctly in complex LangChain scenarios

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Bot Settings | Greptile}

[qnnn]

Hi @hassiebp , I have a question. Should we update the input and output during the agent action or finish process? It’s an intermediate step and may not accurately represent the true input and output of the span.

[hassiebp]

Could you please try with the latest SDK version and report whether you see unexpected results?

[qnnn]

Could you please try with the latest SDK version and report whether you see unexpected results?

This appears to be a different issue that I didn't explain clearly before (or it might not be an issue at all).

• In the agent action callback when setting output, if the tool execution fails, the current output is the action's output, but I expect the output to be empty in such cases.