feat: Add tracing to LLM invoke (#650)

Author: Juliaj

src/rai_core/rai/agents/langchain/core/conversational_agent.py

@@ -22,12 +22,14 @@
     BaseMessage,
     SystemMessage,
 )
+from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph import START, StateGraph
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.prebuilt.tool_node import tools_condition
 
 from rai.agents.langchain.core.tool_runner import ToolRunner
+from rai.agents.langchain.invocation_helpers import invoke_llm_with_tracing
 
 
 class State(TypedDict):
@@ -39,6 +41,7 @@ def agent(
     logger: logging.Logger,
     system_prompt: str | SystemMessage,
     state: State,
+    config: RunnableConfig,
 ):
     logger.info("Running thinker")
 
@@ -54,7 +57,9 @@ def agent(
             else system_prompt
         )
         state["messages"].insert(0, system_msg)
-    ai_msg = llm.invoke(state["messages"])
+
+    # Invoke LLM with tracing if it is configured and available
+    ai_msg = invoke_llm_with_tracing(llm, state["messages"], config)
     state["messages"].append(ai_msg)
     return state
 

src/rai_core/rai/agents/langchain/core/react_agent.py

@@ -22,12 +22,13 @@
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import BaseMessage, SystemMessage
-from langchain_core.runnables import Runnable
+from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt.tool_node import tools_condition
 
 from rai.agents.langchain.core.tool_runner import ToolRunner
+from rai.agents.langchain.invocation_helpers import invoke_llm_with_tracing
 from rai.initialization import get_llm_model
 from rai.messages import SystemMultimodalMessage
 
@@ -48,6 +49,7 @@ def llm_node(
     llm: BaseChatModel,
     system_prompt: Optional[str | SystemMultimodalMessage],
     state: ReActAgentState,
+    config: RunnableConfig,
 ):
     """Process messages using the LLM.
 
@@ -57,6 +59,8 @@ def llm_node(
         The language model to use for processing
     state : ReActAgentState
         Current state containing messages
+    config : RunnableConfig
+        Configuration including callbacks for tracing
 
     Returns
     -------
@@ -75,7 +79,9 @@ def llm_node(
         # at this point, state['messages'] length should at least be 1
         if not isinstance(state["messages"][0], SystemMessage):
             state["messages"].insert(0, SystemMessage(content=system_prompt))
-    ai_msg = llm.invoke(state["messages"])
+
+    # Invoke LLM with tracing if it is configured and available
+    ai_msg = invoke_llm_with_tracing(llm, state["messages"], config)
     state["messages"].append(ai_msg)
 
 

src/rai_core/rai/agents/langchain/core/state_based_agent.py

@@ -26,12 +26,13 @@
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
-from langchain_core.runnables import Runnable
+from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt.tool_node import tools_condition
 
 from rai.agents.langchain.core.tool_runner import ToolRunner
+from rai.agents.langchain.invocation_helpers import invoke_llm_with_tracing
 from rai.initialization import get_llm_model
 from rai.messages import HumanMultimodalMessage, SystemMultimodalMessage
 
@@ -52,6 +53,7 @@ def llm_node(
     llm: BaseChatModel,
     system_prompt: Optional[str | SystemMultimodalMessage],
     state: ReActAgentState,
+    config: RunnableConfig,
 ):
     """Process messages using the LLM.
 
@@ -61,6 +63,8 @@ def llm_node(
         The language model to use for processing
     state : ReActAgentState
         Current state containing messages
+    config : RunnableConfig
+        Configuration including callbacks for tracing
 
     Returns
     -------
@@ -79,7 +83,9 @@ def llm_node(
         # at this point, state['messages'] length should at least be 1
         if not isinstance(state["messages"][0], SystemMessage):
             state["messages"].insert(0, SystemMessage(content=system_prompt))
-    ai_msg = llm.invoke(state["messages"])
+
+    # Invoke LLM with tracing if it is configured and available
+    ai_msg = invoke_llm_with_tracing(llm, state["messages"], config)
     state["messages"].append(ai_msg)
 
 

src/rai_core/rai/agents/langchain/invocation_helpers.py

@@ -0,0 +1,74 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import Any, List, Optional
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import BaseMessage
+from langchain_core.runnables import RunnableConfig
+
+from rai.initialization import get_tracing_callbacks
+
+logger = logging.getLogger(__name__)
+
+
+def invoke_llm_with_tracing(
+    llm: BaseChatModel,
+    messages: List[BaseMessage],
+    config: Optional[RunnableConfig] = None,
+) -> Any:
+    """
+    Invoke an LLM with enhanced tracing callbacks.
+
+    This function automatically adds tracing callbacks (like Langfuse) to LLM calls
+    within LangGraph nodes, solving the callback propagation issue.
+
+    Parameters
+    ----------
+    llm : BaseChatModel
+        The language model to invoke
+    messages : List[BaseMessage]
+        Messages to send to the LLM
+    config : Optional[RunnableConfig]
+        Existing configuration (may contain some callbacks)
+
+    Returns
+    -------
+    Any
+        The LLM response
+    """
+    tracing_callbacks = get_tracing_callbacks()
+
+    if len(tracing_callbacks) == 0:
+        # No tracing callbacks available, use config as-is
+        return llm.invoke(messages, config=config)
+
+    # Create enhanced config with tracing callbacks
+    enhanced_config = config.copy() if config else {}
+
+    # Add tracing callbacks to existing callbacks
+    existing_callbacks = config.get("callbacks", []) if config else []
+
+    if hasattr(existing_callbacks, "handlers"):
+        # Merge with existing CallbackManager
+        all_callbacks = existing_callbacks.handlers + tracing_callbacks
+    elif isinstance(existing_callbacks, list):
+        all_callbacks = existing_callbacks + tracing_callbacks
+    else:
+        all_callbacks = tracing_callbacks
+
+    enhanced_config["callbacks"] = all_callbacks
+
+    return llm.invoke(messages, config=enhanced_config)