fix(langchain): capture retriever steps in sync handler and tool-call outputs

- Wire on_retriever_start/end/error into the sync OpenlayerHandler (previously
  only AsyncOpenlayerHandler exposed them), so synchronous RAG pipelines produce
  a RETRIEVER step and populate the context column.
- Fall back to serializing AIMessage.tool_calls in _extract_output when the
  generation text is empty, so tool-only agent turns no longer record empty
  output; preserve tool_calls in _message_to_dict for assistant messages.
- Fix context metadata being dropped when a sync root-level retriever has no
  external trace context (external trace still takes priority).
- Add offline unit tests for the sync retriever path and tool-call capture.

Implements the high-priority items of OPEN-11315; medium/low items tracked
separately on the issue.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: viniciusdsmello

src/openlayer/lib/integrations/langchain_callback.py

@@ -364,7 +364,7 @@ def _convert_langchain_objects(self, obj: Any) -> Any:
 
     def _message_to_dict(
         self, message: "langchain_schema.BaseMessage"
-    ) -> Dict[str, str]:
+    ) -> Dict[str, Any]:
         """Convert a LangChain message to a JSON-serializable dictionary."""
         message_type = getattr(message, "type", "user")
 
@@ -374,7 +374,17 @@ def _message_to_dict(
         elif message_type == "system":
             role = "system"
 
-        return {"role": role, "content": str(message.content)}
+        result: Dict[str, Any] = {"role": role, "content": str(message.content)}
+
+        # Preserve tool calls on assistant messages (e.g. AIMessage.tool_calls).
+        # Without this, agent turns that only emit tool calls would drop them
+        # from the recorded inputs. Done defensively so it never raises on
+        # messages that lack the attribute.
+        tool_calls = getattr(message, "tool_calls", None)
+        if tool_calls:
+            result["tool_calls"] = tool_calls
+
+        return result
 
     def _messages_to_prompt_format(
         self, messages: List[List["langchain_schema.BaseMessage"]]
@@ -480,13 +490,45 @@ def _extract_token_info(
         }
 
     def _extract_output(self, response: "langchain_schema.LLMResult") -> str:
-        """Extract output text from LLM response."""
+        """Extract output text from LLM response.
+
+        When a generation carries no text (e.g. an agent turn whose model
+        response is *only* tool calls, as in LangGraph / ``create_agent``),
+        fall back to serializing the message's tool calls so the step output
+        is not empty.
+        """
         output = ""
         for generations in response.generations:
             for generation in generations:
-                output += generation.text.replace("\n", " ")
+                text = generation.text or ""
+                if text:
+                    output += text.replace("\n", " ")
+                else:
+                    output += self._tool_calls_to_str(generation)
         return output
 
+    def _tool_calls_to_str(self, generation: Any) -> str:
+        """Serialize tool calls from a generation's message, if any.
+
+        Done defensively (getattr) so it never raises on generations or
+        messages that lack a ``message`` / ``tool_calls`` attribute, and works
+        whether or not langchain uses the v1 message shape.
+        """
+        message = getattr(generation, "message", None)
+        if message is None:
+            return ""
+
+        tool_calls = getattr(message, "tool_calls", None)
+        if not tool_calls:
+            return ""
+
+        import json
+
+        try:
+            return json.dumps(tool_calls, default=str)
+        except (TypeError, ValueError):
+            return str(tool_calls)
+
     def _safe_parse_json(self, input_str: str) -> Any:
         """Safely parse JSON string, returning the string if parsing fails."""
         try:
@@ -822,6 +864,33 @@ def _handle_agent_finish(
             output=finish.return_values,
         )
 
+    def _find_owning_trace(self, run_id: UUID) -> Optional["traces.Trace"]:
+        """Find the standalone trace that contains the step for ``run_id``.
+
+        Used to attach trace-level data (e.g. retriever ``context``) when no
+        external trace context is active. Returns ``None`` if the run's step is
+        not part of any standalone trace tracked by this handler.
+        """
+        step = self.steps.get(run_id)
+        if step is None:
+            return None
+
+        for trace in self._traces_by_root.values():
+            if self._trace_contains_step(trace, step):
+                return trace
+        return None
+
+    @staticmethod
+    def _trace_contains_step(trace: "traces.Trace", target: steps.Step) -> bool:
+        """Return True if ``target`` is anywhere in ``trace`` (including nested)."""
+        stack = list(trace.steps)
+        while stack:
+            step = stack.pop()
+            if step is target:
+                return True
+            stack.extend(step.steps)
+        return False
+
     def _handle_retriever_start(
         self,
         serialized: Dict[str, Any],
@@ -875,7 +944,12 @@ def _handle_retriever_end(
             else:
                 doc_contents.append(str(doc))
 
-        current_trace = tracer.get_current_trace()
+        # Populate the trace-level `context` so the reserved column is set.
+        # Prefer an active (external) trace context; otherwise fall back to the
+        # standalone trace this handler owns for the run. Without the fallback,
+        # synchronous RAG pipelines (where the retriever is the root step and no
+        # external @trace context exists) would silently drop the context.
+        current_trace = tracer.get_current_trace() or self._find_owning_trace(run_id)
         if current_trace:
             current_trace.update_metadata(context=doc_contents)
 
@@ -1062,6 +1136,26 @@ def on_tool_error(
             return
         return self._handle_tool_error(error, **kwargs)
 
+    def on_retriever_start(
+        self, serialized: Dict[str, Any], query: str, **kwargs: Any
+    ) -> Any:
+        """Run when retriever starts running."""
+        if self.ignore_retriever:
+            return
+        return self._handle_retriever_start(serialized, query, **kwargs)
+
+    def on_retriever_end(self, documents: List[Any], **kwargs: Any) -> Any:
+        """Run when retriever ends running."""
+        if self.ignore_retriever:
+            return
+        return self._handle_retriever_end(documents, **kwargs)
+
+    def on_retriever_error(self, error: Exception, **kwargs: Any) -> Any:
+        """Run when retriever errors."""
+        if self.ignore_retriever:
+            return
+        return self._handle_retriever_error(error, **kwargs)
+
     def on_text(self, text: str, **kwargs: Any) -> Any:
         """Run on arbitrary text."""
         pass

tests/lib/__init__.py

@@ -0,0 +1,212 @@
+"""Tests for the Openlayer LangChain callback handler.
+
+Regression coverage for OPEN-11315:
+
+1. The synchronous ``OpenlayerHandler`` must wire up the retriever callbacks
+   (``on_retriever_start`` / ``on_retriever_end`` / ``on_retriever_error``) so
+   synchronous RAG pipelines produce a RETRIEVER step and the auto-populated
+   ``context`` column on the trace.
+
+2. A chat-completion turn whose response contains ONLY tool calls (every agent
+   iteration in LangGraph / ``create_agent``) must still produce a non-empty
+   step output: ``_extract_output`` falls back to serializing the message's
+   tool calls, and ``_message_to_dict`` preserves ``tool_calls`` on the
+   assistant message.
+"""
+
+import uuid
+
+import pytest
+
+pytest.importorskip("langchain_core")
+
+from langchain_core.documents import Document
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langchain_core.outputs import ChatGeneration, LLMResult
+
+from openlayer.lib.integrations.langchain_callback import OpenlayerHandler
+from openlayer.lib.tracing import enums, steps
+
+
+@pytest.fixture(autouse=True)
+def _disable_publish(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Keep all tracer publish paths off during the test (fully offline)."""
+    monkeypatch.setenv("OPENLAYER_DISABLE_PUBLISH", "true")
+    monkeypatch.setenv("OPENLAYER_API_KEY", "fake")
+
+    from openlayer.lib.tracing import tracer as _tracer
+
+    monkeypatch.setattr(_tracer, "_publish", False, raising=False)
+
+
+def _tool_call(name: str, args: dict, call_id: str) -> dict:
+    """Build a langchain_core tool-call dict."""
+    return {"name": name, "args": args, "id": call_id, "type": "tool_call"}
+
+
+# --------------------------------------------------------------------------- #
+# Fix 2 (b): _message_to_dict preserves tool_calls
+# --------------------------------------------------------------------------- #
+class TestMessageToDict:
+    def test_preserves_tool_calls_on_ai_message(self) -> None:
+        handler = OpenlayerHandler()
+        tool_calls = [_tool_call("search", {"query": "openlayer"}, "call_1")]
+        message = AIMessage(content="", tool_calls=tool_calls)
+
+        result = handler._message_to_dict(message)
+
+        assert result["role"] == "assistant"
+        assert "tool_calls" in result, "tool_calls should be preserved"
+        assert result["tool_calls"] == tool_calls
+
+    def test_human_message_without_tool_calls_is_backwards_compatible(self) -> None:
+        handler = OpenlayerHandler()
+        message = HumanMessage(content="hello there")
+
+        result = handler._message_to_dict(message)
+
+        assert result == {"role": "user", "content": "hello there"}
+        assert "tool_calls" not in result
+
+    def test_ai_message_without_tool_calls_omits_key(self) -> None:
+        handler = OpenlayerHandler()
+        message = AIMessage(content="plain answer")
+
+        result = handler._message_to_dict(message)
+
+        assert result == {"role": "assistant", "content": "plain answer"}
+        assert "tool_calls" not in result
+
+    def test_tool_message_does_not_raise(self) -> None:
+        handler = OpenlayerHandler()
+        message = ToolMessage(content="42", tool_call_id="call_1")
+
+        result = handler._message_to_dict(message)
+
+        assert result["content"] == "42"
+        assert "tool_calls" not in result
+
+
+# --------------------------------------------------------------------------- #
+# Fix 2 (a): _extract_output falls back to tool calls
+# --------------------------------------------------------------------------- #
+class TestExtractOutput:
+    def test_text_generation_returns_text(self) -> None:
+        handler = OpenlayerHandler()
+        gen = ChatGeneration(message=AIMessage(content="Hello world"))
+        response = LLMResult(generations=[[gen]])
+
+        assert handler._extract_output(response) == "Hello world"
+
+    def test_tool_only_generation_returns_tool_calls(self) -> None:
+        handler = OpenlayerHandler()
+        tool_calls = [_tool_call("get_weather", {"city": "Rio"}, "call_42")]
+        message = AIMessage(content="", tool_calls=tool_calls)
+        gen = ChatGeneration(message=message)
+        response = LLMResult(generations=[[gen]])
+
+        output = handler._extract_output(response)
+
+        assert output, "tool-only output must not be empty"
+        assert "get_weather" in output
+
+    def test_tool_only_generation_via_callbacks(self) -> None:
+        """Drive on_chat_model_start / on_llm_end and assert step output is set."""
+        handler = OpenlayerHandler()
+        run_id = uuid.uuid4()
+
+        handler.on_chat_model_start(
+            serialized={"name": "gpt-4o"},
+            messages=[[HumanMessage(content="What is the weather in Rio?")]],
+            run_id=run_id,
+            invocation_params={"model_name": "gpt-4o", "_type": "openai-chat"},
+        )
+
+        # Capture the standalone trace before the root step is ended/popped.
+        trace = handler._traces_by_root[run_id]
+        step = handler.steps[run_id]
+        assert isinstance(step, steps.ChatCompletionStep)
+
+        tool_calls = [_tool_call("get_weather", {"city": "Rio"}, "call_42")]
+        message = AIMessage(content="", tool_calls=tool_calls)
+        response = LLMResult(generations=[[ChatGeneration(message=message)]])
+
+        handler.on_llm_end(response, run_id=run_id)
+
+        assert step.step_type == enums.StepType.CHAT_COMPLETION
+        assert step.output, "tool-only chat completion output must not be empty"
+        assert "get_weather" in step.output
+        assert trace.steps[0] is step
+
+
+# --------------------------------------------------------------------------- #
+# Fix 1: sync handler wires retriever callbacks
+# --------------------------------------------------------------------------- #
+class TestSyncRetrieverCallbacks:
+    def test_handler_exposes_retriever_callbacks(self) -> None:
+        handler = OpenlayerHandler()
+        assert hasattr(handler, "on_retriever_start")
+        assert hasattr(handler, "on_retriever_end")
+        assert hasattr(handler, "on_retriever_error")
+
+    def test_sync_retriever_run_produces_step_and_context(self) -> None:
+        handler = OpenlayerHandler()
+        run_id = uuid.uuid4()
+
+        handler.on_retriever_start(
+            serialized={"id": ["langchain", "retrievers", "VectorStoreRetriever"]},
+            query="what is openlayer?",
+            run_id=run_id,
+        )
+
+        trace = handler._traces_by_root[run_id]
+        step = handler.steps[run_id]
+        assert isinstance(step, steps.RetrieverStep)
+        assert step.step_type == enums.StepType.RETRIEVER
+        assert step.inputs == {"query": "what is openlayer?"}
+
+        documents = [
+            Document(page_content="Openlayer is an evaluation platform."),
+            Document(page_content="It supports LangChain tracing."),
+        ]
+        handler.on_retriever_end(documents, run_id=run_id)
+
+        # The retriever step captured the documents...
+        assert step.documents == [
+            "Openlayer is an evaluation platform.",
+            "It supports LangChain tracing.",
+        ]
+        # ...and the trace gained the auto-populated `context` metadata.
+        assert trace.metadata is not None
+        assert trace.metadata.get("context") == [
+            "Openlayer is an evaluation platform.",
+            "It supports LangChain tracing.",
+        ]
+
+    def test_sync_retriever_respects_ignore_flag(self) -> None:
+        handler = OpenlayerHandler(ignore_retriever=True)
+        run_id = uuid.uuid4()
+
+        handler.on_retriever_start(
+            serialized={"id": ["VectorStoreRetriever"]},
+            query="ignored",
+            run_id=run_id,
+        )
+
+        assert run_id not in handler.steps
+
+    def test_sync_retriever_error(self) -> None:
+        handler = OpenlayerHandler()
+        run_id = uuid.uuid4()
+
+        handler.on_retriever_start(
+            serialized={"id": ["VectorStoreRetriever"]},
+            query="boom",
+            run_id=run_id,
+        )
+        step = handler.steps[run_id]
+
+        handler.on_retriever_error(ValueError("retrieval failed"), run_id=run_id)
+
+        assert step.metadata.get("error") == "retrieval failed"
+        assert run_id not in handler.steps