revert: "feat: Add tracing to LLM invoke"

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
 

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)
 
 

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)