diff --git a/src/openlayer/lib/integrations/langchain_callback.py b/src/openlayer/lib/integrations/langchain_callback.py index 3d66e742..f907beb1 100644 --- a/src/openlayer/lib/integrations/langchain_callback.py +++ b/src/openlayer/lib/integrations/langchain_callback.py @@ -562,6 +562,11 @@ def _handle_chain_start( if tags and "langsmith:hidden" in tags: return + if chain_name == "LangGraph" and inputs.get("messages"): + inputs = { + "prompt": inputs.get("messages"), + } + self._start_step( run_id=run_id, parent_run_id=parent_run_id, @@ -571,7 +576,6 @@ def _handle_chain_start( metadata={ "tags": tags, "serialized": serialized, - "is_chain": True, **(metadata or {}), **kwargs, }, @@ -607,6 +611,15 @@ def _handle_chain_end( if current_trace: current_trace.update_metadata(context=context_list) + # Parse output for LangGraph + step = self.steps[run_id] + if step.name == "LangGraph" and outputs.get("messages"): + if isinstance(outputs.get("messages"), list): + if isinstance( + outputs.get("messages")[-1], langchain_schema.BaseMessage + ): + outputs = outputs.get("messages")[-1].content + self._end_step( run_id=run_id, parent_run_id=parent_run_id,