From e192fccb524cfffa440ec875dc11722a51572b58 Mon Sep 17 00:00:00 2001
From: Divya Chitimalla <divya@getgarak.com>
Date: Fri, 12 Sep 2025 13:27:00 -0700
Subject: [PATCH] Fix Langfuse callback passthrough in RunnableRails

Fixes #472: The langchain runnable was not called with the runnable
config that defines callbacks, breaking Langfuse tracing integration.

Changes:
- Store RunnableConfig and kwargs from invoke() for use in passthrough function
- Pass stored config to underlying runnable to preserve callbacks
- Add test to verify callback passthrough functionality

This ensures that callbacks (like Langfuse tracing) are properly
propagated to the underlying LLM runnable wrapped in RunnableRails.
---
 .../integrations/langchain/runnable_rails.py  | 11 ++++++-
 tests/test_runnable_rails.py                  | 29 +++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/nemoguardrails/integrations/langchain/runnable_rails.py b/nemoguardrails/integrations/langchain/runnable_rails.py
index 1eb282848..6d9916386 100644
--- a/nemoguardrails/integrations/langchain/runnable_rails.py
+++ b/nemoguardrails/integrations/langchain/runnable_rails.py
@@ -49,6 +49,8 @@ def __init__(
         self.passthrough_bot_output_key = output_key
         self.verbose = verbose
         self.config: Optional[RunnableConfig] = None
+        self._current_config: Optional[RunnableConfig] = None
+        self._current_kwargs: dict = {}
 
         # We override the config passthrough.
         config.passthrough = passthrough
@@ -74,7 +76,10 @@ async def passthrough_fn(context: dict, events: List[dict]):
             # First, we fetch the input from the context
             _input = context.get("passthrough_input")
             async_wrapped_invoke = async_wrap(self.passthrough_runnable.invoke)
-            _output = await async_wrapped_invoke(_input, self.config, **self.kwargs)
+            
+            # Pass the config and kwargs that were captured in the invoke method
+            # This ensures that callbacks (like Langfuse tracing) are properly propagated
+            _output = await async_wrapped_invoke(_input, self._current_config, **self._current_kwargs)
 
             # If the output is a string, we consider it to be the output text
             if isinstance(_output, str):
@@ -188,8 +193,12 @@ def invoke(
     ) -> Output:
         """Invoke this runnable synchronously."""
         input_messages = self._transform_input_to_rails_format(input)
+        # Store config and kwargs for use in passthrough function
+        # This ensures callbacks are properly passed to the underlying runnable
         self.config = config
         self.kwargs = kwargs
+        self._current_config = config
+        self._current_kwargs = kwargs
         res = self.rails.generate(
             messages=input_messages, options=GenerationOptions(output_vars=True)
         )
diff --git a/tests/test_runnable_rails.py b/tests/test_runnable_rails.py
index 10b33c056..38239f7ca 100644
--- a/tests/test_runnable_rails.py
+++ b/tests/test_runnable_rails.py
@@ -658,3 +658,32 @@ def log(x):
     print(result)
     assert "LOL" not in result["output"]
     assert "can't respond" in result["output"]
+
+
+def test_runnable_config_callback_passthrough():
+    """Test that RunnableConfig with callbacks is properly passed to passthrough runnable."""
+    config_received = []
+    
+    class CallbackTestRunnable(Runnable):
+        def invoke(self, input: Input, config: Optional[RunnableConfig] = None) -> Output:
+            # Capture the config to verify callbacks were passed
+            config_received.append(config)
+            return {"output": "Test response"}
+    
+    # Create a mock callback for testing
+    mock_callbacks = ["mock_callback"]
+    test_config = RunnableConfig(callbacks=mock_callbacks)
+    
+    rails_config = RailsConfig.from_content(config={"models": []})
+    runnable_with_rails = RunnableRails(
+        rails_config, passthrough=True, runnable=CallbackTestRunnable()
+    )
+    
+    # Invoke with the config containing callbacks
+    result = runnable_with_rails.invoke("test input", config=test_config)
+    
+    # Verify that the config with callbacks was passed through
+    assert len(config_received) == 1
+    assert config_received[0] is not None
+    assert config_received[0].get("callbacks") == mock_callbacks
+    assert result == {"output": "Test response"}