diff --git a/src/rai_core/rai/agents/langchain/core/conversational_agent.py b/src/rai_core/rai/agents/langchain/core/conversational_agent.py index e008fdb7d..17fcb7271 100644 --- a/src/rai_core/rai/agents/langchain/core/conversational_agent.py +++ b/src/rai_core/rai/agents/langchain/core/conversational_agent.py @@ -23,14 +23,12 @@ 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): @@ -42,7 +40,6 @@ def agent( logger: logging.Logger, system_prompt: str | SystemMessage, state: State, - config: RunnableConfig, ): logger.info("Running thinker") @@ -58,9 +55,7 @@ def agent( else system_prompt ) state["messages"].insert(0, system_msg) - - # Invoke LLM with tracing if it is configured and available - ai_msg = invoke_llm_with_tracing(llm, state["messages"], config) + ai_msg = llm.invoke(state["messages"]) state["messages"].append(ai_msg) return state diff --git a/src/rai_core/rai/agents/langchain/core/react_agent.py b/src/rai_core/rai/agents/langchain/core/react_agent.py index f872add3f..d49fd4617 100644 --- a/src/rai_core/rai/agents/langchain/core/react_agent.py +++ b/src/rai_core/rai/agents/langchain/core/react_agent.py @@ -22,13 +22,12 @@ from langchain_core.language_models import BaseChatModel from langchain_core.messages import BaseMessage, SystemMessage -from langchain_core.runnables import Runnable, RunnableConfig +from langchain_core.runnables import Runnable 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 @@ -49,7 +48,6 @@ def llm_node( llm: BaseChatModel, system_prompt: Optional[str | SystemMultimodalMessage], state: ReActAgentState, - config: RunnableConfig, ): """Process messages using the LLM. @@ -59,8 +57,6 @@ 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,9 +75,7 @@ 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)) - - # Invoke LLM with tracing if it is configured and available - ai_msg = invoke_llm_with_tracing(llm, state["messages"], config) + ai_msg = llm.invoke(state["messages"]) state["messages"].append(ai_msg) diff --git a/src/rai_core/rai/agents/langchain/core/state_based_agent.py b/src/rai_core/rai/agents/langchain/core/state_based_agent.py index ea651acd3..0ffcff6d5 100644 --- a/src/rai_core/rai/agents/langchain/core/state_based_agent.py +++ b/src/rai_core/rai/agents/langchain/core/state_based_agent.py @@ -26,13 +26,12 @@ from langchain_core.language_models import BaseChatModel from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage -from langchain_core.runnables import Runnable, RunnableConfig +from langchain_core.runnables import Runnable 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 @@ -53,7 +52,6 @@ def llm_node( llm: BaseChatModel, system_prompt: Optional[str | SystemMultimodalMessage], state: ReActAgentState, - config: RunnableConfig, ): """Process messages using the LLM. @@ -63,8 +61,6 @@ def llm_node( The language model to use for processing state : ReActAgentState Current state containing messages - config : RunnableConfig - Configuration including callbacks for tracing Returns ------- @@ -83,9 +79,7 @@ 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)) - - # Invoke LLM with tracing if it is configured and available - ai_msg = invoke_llm_with_tracing(llm, state["messages"], config) + ai_msg = llm.invoke(state["messages"]) state["messages"].append(ai_msg) diff --git a/src/rai_core/rai/agents/langchain/invocation_helpers.py b/src/rai_core/rai/agents/langchain/invocation_helpers.py deleted file mode 100644 index c2a667f29..000000000 --- a/src/rai_core/rai/agents/langchain/invocation_helpers.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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)