feat(callbacks): add MermaidTrace handler for sequence diagram visualization#519
feat(callbacks): add MermaidTrace handler for sequence diagram visualization#519xt765 wants to merge 7 commits intolangchain-ai:mainfrom
Conversation
- Implement MermaidTraceCallbackHandler in callbacks/mermaid_trace.py - Register handler in callbacks/__init__.py - Add unit tests in tests/unit_tests/callbacks/test_mermaid_trace.py - Update extended_testing_deps.txt with mermaid-trace dependency - Ensure compatibility with langchain-core BaseCallbackHandler signatures - Fix linting and type checking issues
…or MermaidTraceCallbackHandler
…TraceCallbackHandler - Implement thread-safe participant stack using threading.local() - Cache mermaid-trace module to avoid repetitive imports - Refactor methods to use centralized trace ID and source retrieval - Optimize performance by reducing redundant function calls
![graphite-app[bot]](https://avatars.githubusercontent.com/in/158384?s=60&v=4){kind=small}
| chain.invoke({"input": "hello"}, config={"callbacks": [handler]}) | ||
|
|
||
| # Verify stack is empty after error | ||
| assert len(handler._participant_stack) == 0 |
There was a problem hiding this comment.
Look for Suspicious or Risky Code: The code references handler._participant_stack which does not exist in the MermaidTraceCallbackHandler class. The actual implementation uses _root_to_stack dictionary to track participant stacks per root run. This will cause an AttributeError at runtime. The test should either use the public API or reference the correct internal attribute like handler._root_to_stack.
| assert len(handler._participant_stack) == 0 | |
| assert all(len(stack) == 0 for stack in handler._root_to_stack.values()) |
Spotted by Graphite Agent (based on custom rule: Code quality)
Is this helpful? React 👍 or 👎 to let us know.
| target = self._end_run(run_id) | ||
| if not target: | ||
| return | ||
|
|
||
| source = self._get_current_source(run_id) |
There was a problem hiding this comment.
Critical Bug: Use-after-delete causes incorrect source participant
The _end_run() method deletes _run_to_root[run_id] on line 112, but then _get_current_source(run_id) is called which needs this mapping to find the participant stack. This causes _get_participant_stack() to return an empty list, making _get_current_source() fall back to the LogContext default instead of using the correct parent from the stack.
Impact: Return/error messages in the sequence diagram will have incorrect source participants, breaking the visualization.
Fix: Call _get_current_source(run_id) BEFORE _end_run(run_id):
def on_chain_end(self, outputs: Dict[str, Any], *, run_id: uuid.UUID, ...) -> None:
source = self._get_current_source(run_id) # Move before _end_run
target = self._end_run(run_id)
if not target:
return
# rest of method...This affects all on_*_end and on_*_error methods (lines 194-260, 346-412, 458-524, 568-634).
| target = self._end_run(run_id) | |
| if not target: | |
| return | |
| source = self._get_current_source(run_id) | |
| source = self._get_current_source(run_id) | |
| target = self._end_run(run_id) | |
| if not target: | |
| return | |
Spotted by Graphite Agent
Is this helpful? React 👍 or 👎 to let us know.
… name as suggested by Graphite
| def _get_participant_stack(self, run_id: uuid.UUID) -> List[str]: | ||
| """Get the participant stack for the given run_id. | ||
|
|
||
| Args: | ||
| run_id: The run ID. | ||
|
|
||
| Returns: | ||
| The participant stack. | ||
| """ | ||
| with self._lock: | ||
| root_id = self._run_to_root.get(run_id) | ||
| if root_id is None: | ||
| return [] | ||
| return self._root_to_stack.get(root_id, []) |
There was a problem hiding this comment.
Critical Race Condition: Returns mutable list reference without lock protection
The method returns a direct reference to the internal stack list, which can be modified by other threads after the lock is released. This causes race conditions when callers use the returned stack.
Impact: In _get_current_source (line 144-146), after getting the stack reference and releasing the lock, another thread could call _end_run and pop from the stack, leading to:
IndexErrorif stack becomes empty between theif stack:check andstack[-1]access- Wrong participant name if a different element is now at the top
Fix:
def _get_participant_stack(self, run_id: uuid.UUID) -> List[str]:
with self._lock:
root_id = self._run_to_root.get(run_id)
if root_id is None:
return []
return list(self._root_to_stack.get(root_id, [])) # Return a copy| def _get_participant_stack(self, run_id: uuid.UUID) -> List[str]: | |
| """Get the participant stack for the given run_id. | |
| Args: | |
| run_id: The run ID. | |
| Returns: | |
| The participant stack. | |
| """ | |
| with self._lock: | |
| root_id = self._run_to_root.get(run_id) | |
| if root_id is None: | |
| return [] | |
| return self._root_to_stack.get(root_id, []) | |
| def _get_participant_stack(self, run_id: uuid.UUID) -> List[str]: | |
| """Get the participant stack for the given run_id. | |
| Args: | |
| run_id: The run ID. | |
| Returns: | |
| The participant stack. | |
| """ | |
| with self._lock: | |
| root_id = self._run_to_root.get(run_id) | |
| if root_id is None: | |
| return [] | |
| return list(self._root_to_stack.get(root_id, [])) # Return a copy |
Spotted by Graphite Agent
Is this helpful? React 👍 or 👎 to let us know.
…fix source participant identification
|
Updated the PR with fixes for all critical issues identified by Graphite Agent. Local tests and linting are all passing. |
1 participant
Description
Disclaimer: This contribution was developed with the assistance of an AI programming agent (Trae/Gemini).
This PR introduces the
MermaidTraceCallbackHandlerinlangchain-community, allowing users to visualize LangChain execution flows as Mermaid sequence diagrams.Why this solution is right:
_run_to_rootmapping andthreading.Lockto correctly isolate concurrent requests.guard_importfor themermaid-tracedependency, ensuring it doesn't affect users who don't use this feature.Key Fixes & Improvements:
_end_runto safely extract parent metadata before cleaning up internal mappings, preventing potential crashes or data corruption in end/error events.Areas for careful review:
root_run_idinmermaid_trace.py.Issue
N/A
Dependencies
This feature requires the optional
mermaid-tracepackage for diagram generation. It is handled via lazy import.Testing
All tests passed locally using
pytest,ruff, andmypy.libs/community/tests/unit_tests/callbacks/test_mermaid_trace.py.libs/community/tests/unit_tests/callbacks/test_mermaid_trace_concurrency.py.libs/community/tests/integration_tests/callbacks/test_mermaid_trace.py.ruff check,ruff format, andmypy --disallow-untyped-defs.