simplify agent loop to use built-ins

Author: jsonbailey

packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py

@@ -1,104 +1,63 @@
 """LangChain agent runner for LaunchDarkly AI SDK."""
 
-from typing import Any, List
+from typing import Any
 
-from langchain_core.messages import (
-    AIMessage,
-    BaseMessage,
-    HumanMessage,
-    SystemMessage,
-    ToolMessage,
-)
 from ldai import log
-from ldai.providers import AgentResult, AgentRunner, ToolRegistry
+from ldai.providers import AgentResult, AgentRunner
 from ldai.providers.types import LDAIMetrics
 
-from ldai_langchain.langchain_helper import get_ai_metrics_from_response
+from ldai_langchain.langchain_helper import sum_token_usage_from_messages
 
 
 class LangChainAgentRunner(AgentRunner):
     """
     AgentRunner implementation for LangChain.
 
-    Executes a single-agent loop using a LangChain BaseChatModel with tool calling.
-    The model is expected to have tools already bound to it.
+    Wraps a compiled LangChain agent graph (from ``langchain.agents.create_agent``)
+    and delegates execution to it. Tool calling and loop management are handled
+    internally by the graph.
     Returned by LangChainRunnerFactory.create_agent(config, tools).
     """
 
-    def __init__(
-        self,
-        llm: Any,
-        instructions: str,
-        tools: ToolRegistry,
-    ):
-        self._llm = llm
-        self._instructions = instructions
-        self._tools = tools
+    def __init__(self, agent: Any):
+        self._agent = agent
 
     async def run(self, input: Any) -> AgentResult:
         """
         Run the agent with the given input string.
 
-        Executes an agentic loop: calls the model, handles tool calls,
-        and continues until the model produces a final response.
+        Delegates to the compiled LangChain agent, which handles
+        the tool-calling loop internally.
 
         :param input: The user prompt or input to the agent
         :return: AgentResult with output, raw response, and aggregated metrics
         """
-        messages: List[BaseMessage] = []
-        if self._instructions:
-            messages.append(SystemMessage(content=self._instructions))
-        messages.append(HumanMessage(content=str(input)))
-
-        raw_response = None
-
         try:
-            while True:
-                response: AIMessage = await self._llm.ainvoke(messages)
-                raw_response = response
-                messages.append(response)
-
-                tool_calls = getattr(response, 'tool_calls', None)
-
-                if not tool_calls:
-                    metrics = get_ai_metrics_from_response(response)
-                    content = response.content if isinstance(response.content, str) else ""
-                    return AgentResult(
-                        output=content,
-                        raw=raw_response,
-                        metrics=metrics,
-                    )
-
-                # Execute tool calls and append results
-                for tool_call in tool_calls:
-                    tool_name = tool_call["name"]
-                    tool_args = tool_call.get("args", {})
-                    tool_id = tool_call.get("id", "")
-
-                    tool_fn = self._tools.get(tool_name)
-                    if tool_fn:
-                        try:
-                            result = tool_fn(**tool_args)
-                            if hasattr(result, "__await__"):
-                                result = await result
-                            result_str = str(result)
-                        except Exception as error:
-                            log.warning(f"Tool '{tool_name}' execution failed: {error}")
-                            result_str = f"Tool execution failed: {error}"
-                    else:
-                        log.warning(f"Tool '{tool_name}' not found in registry")
-                        result_str = f"Tool '{tool_name}' not found"
-
-                    messages.append(ToolMessage(content=result_str, tool_call_id=tool_id))
-
+            result = await self._agent.ainvoke({
+                "messages": [{"role": "user", "content": str(input)}]
+            })
+            messages = result.get("messages", [])
+            output = ""
+            if messages:
+                last = messages[-1]
+                if hasattr(last, 'content') and isinstance(last.content, str):
+                    output = last.content
+            return AgentResult(
+                output=output,
+                raw=result,
+                metrics=LDAIMetrics(
+                    success=True,
+                    usage=sum_token_usage_from_messages(messages),
+                ),
+            )
         except Exception as error:
             log.warning(f"LangChain agent run failed: {error}")
             return AgentResult(
                 output="",
-                raw=raw_response,
+                raw=None,
                 metrics=LDAIMetrics(success=False, usage=None),
             )
 
-    def get_llm(self) -> Any:
-        """Return the underlying LangChain LLM."""
-        return self._llm
+    def get_agent(self) -> Any:
+        """Return the underlying compiled LangChain agent."""
+        return self._agent

packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py

@@ -87,8 +87,8 @@ def create_langchain_model(ai_config: AIConfigKind, tool_registry: Optional[Tool
         **parameters,
     )
 
-    if tool_definitions:
-        bindable = _resolve_tools_for_langchain(tool_definitions, tool_registry or {})
+    if tool_definitions and tool_registry is not None:
+        bindable = _resolve_tools_for_langchain(tool_definitions, tool_registry)
         if bindable:
             model = model.bind_tools(bindable)
 
@@ -138,6 +138,55 @@ def _resolve_tools_for_langchain(
     return bindable
 
 
+def build_structured_tools(ai_config: AIConfigKind, tool_registry: ToolRegistry) -> List[Any]:
+    """
+    Build a list of LangChain StructuredTool instances from LD tool definitions and a registry.
+
+    Tools found in the registry are wrapped as StructuredTool with the name and description
+    from the LD config. Built-in provider tools and tools missing from the registry are
+    skipped with a warning.
+
+    :param ai_config: The LaunchDarkly AI configuration
+    :param tool_registry: Registry mapping tool names to callable implementations
+    :return: List of StructuredTool instances ready to pass to langchain.agents.create_agent
+    """
+    from langchain_core.tools import StructuredTool
+
+    config_dict = ai_config.to_dict()
+    model_dict = config_dict.get('model') or {}
+    parameters = dict(model_dict.get('parameters') or {})
+    tool_definitions = parameters.pop('tools', []) or []
+
+    structured = []
+    for td in tool_definitions:
+        if not isinstance(td, dict):
+            continue
+
+        tool_type = td.get('type')
+        if tool_type and tool_type != 'function':
+            log.warning(
+                f"Built-in tool '{tool_type}' is not reliably supported via LangChain and will be skipped. "
+                "Use a provider-specific runner to use built-in provider tools."
+            )
+            continue
+
+        name = td.get('name')
+        if not name:
+            continue
+
+        if name not in tool_registry:
+            log.warning(f"Tool '{name}' is defined in the AI config but was not found in the tool registry; skipping.")
+            continue
+
+        structured.append(StructuredTool.from_function(
+            func=tool_registry[name],
+            name=name,
+            description=td.get('description', ''),
+        ))
+
+    return structured
+
+
 def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]:
     """
     Extract token usage from a LangChain response.

packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py

@@ -1,12 +1,14 @@
-from typing import TYPE_CHECKING, Any, Optional
+from typing import Any, Optional
 
+from langchain.agents import create_agent as lc_create_agent
 from ldai.models import AIConfigKind
 from ldai.providers import AIProvider, ToolRegistry
 
-if TYPE_CHECKING:
-    from ldai_langchain.langchain_agent_runner import LangChainAgentRunner
-
-from ldai_langchain.langchain_helper import create_langchain_model
+from ldai_langchain.langchain_agent_runner import LangChainAgentRunner
+from ldai_langchain.langchain_helper import (
+    build_structured_tools,
+    create_langchain_model,
+)
 from ldai_langchain.langchain_model_runner import LangChainModelRunner
 
 
@@ -36,16 +38,21 @@ def create_model(self, config: AIConfigKind) -> LangChainModelRunner:
         llm = create_langchain_model(config)
         return LangChainModelRunner(llm)
 
-    def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> 'LangChainAgentRunner':
+    def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> LangChainAgentRunner:
         """
         Create a configured LangChainAgentRunner for the given AI agent config.
 
         :param config: The LaunchDarkly AI agent configuration
         :param tools: ToolRegistry mapping tool names to callables
         :return: LangChainAgentRunner ready to run the agent
         """
-        from ldai_langchain.langchain_agent_runner import LangChainAgentRunner
-
         instructions = (config.instructions or '') if hasattr(config, 'instructions') else ''
-        llm = create_langchain_model(config, tool_registry=tools or {})
-        return LangChainAgentRunner(llm, instructions, tools or {})
+        llm = create_langchain_model(config)
+        lc_tools = build_structured_tools(config, tools or {})
+
+        agent = lc_create_agent(
+            llm,
+            tools=lc_tools or None,
+            system_prompt=instructions or None,
+        )
+        return LangChainAgentRunner(agent)

packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py

@@ -402,7 +402,7 @@ class TestCreateAgent:
     """Tests for LangChainRunnerFactory.create_agent."""
 
     def test_creates_agent_runner_with_instructions_and_tool_definitions(self):
-        """Should create LangChainAgentRunner with instructions and tool definitions."""
+        """Should create LangChainAgentRunner wrapping a compiled graph."""
         from unittest.mock import patch
         from ldai_langchain import LangChainAgentRunner
 
@@ -420,15 +420,18 @@ def test_creates_agent_runner_with_instructions_and_tool_definitions(self):
             'provider': {'name': 'openai'},
         }
 
-        with patch('ldai_langchain.langchain_runner_factory.create_langchain_model') as mock_create:
-            mock_llm = MagicMock()
-            mock_create.return_value = mock_llm
+        mock_agent = MagicMock()
+        with patch('ldai_langchain.langchain_runner_factory.create_langchain_model') as mock_create, \
+             patch('ldai_langchain.langchain_runner_factory.build_structured_tools') as mock_tools, \
+             patch('ldai_langchain.langchain_runner_factory.lc_create_agent', return_value=mock_agent):
+            mock_create.return_value = MagicMock()
+            mock_tools.return_value = [MagicMock()]
 
             factory = LangChainRunnerFactory()
             result = factory.create_agent(mock_ai_config, {'get-weather': lambda loc: 'sunny'})
 
             assert isinstance(result, LangChainAgentRunner)
-            assert result._instructions == "You are a helpful assistant."
+            assert result._agent is mock_agent
 
     def test_creates_agent_runner_with_no_tools(self):
         """Should create LangChainAgentRunner with no tool definitions."""
@@ -442,73 +445,72 @@ def test_creates_agent_runner_with_no_tools(self):
             'provider': {'name': 'openai'},
         }
 
-        with patch('ldai_langchain.langchain_runner_factory.create_langchain_model') as mock_create:
+        mock_agent = MagicMock()
+        with patch('ldai_langchain.langchain_runner_factory.create_langchain_model') as mock_create, \
+             patch('ldai_langchain.langchain_runner_factory.build_structured_tools', return_value=[]), \
+             patch('ldai_langchain.langchain_runner_factory.lc_create_agent', return_value=mock_agent):
             mock_create.return_value = MagicMock()
 
             factory = LangChainRunnerFactory()
             result = factory.create_agent(mock_ai_config, {})
 
             assert isinstance(result, LangChainAgentRunner)
-            assert result._tools == {}
+            assert result._agent is mock_agent
 
 
 class TestLangChainAgentRunner:
     """Tests for LangChainAgentRunner.run."""
 
     @pytest.mark.asyncio
-    async def test_runs_agent_and_returns_result_with_no_tool_calls(self):
-        """Should return AgentResult when model responds with no tool calls."""
+    async def test_runs_agent_and_returns_result(self):
+        """Should return AgentResult with the last message content from the graph."""
         from ldai_langchain import LangChainAgentRunner
-        from langchain_core.messages import AIMessage
 
-        mock_llm = MagicMock()
-        mock_response = AIMessage(content="The answer is 42.")
-        mock_llm.ainvoke = AsyncMock(return_value=mock_response)
+        final_msg = AIMessage(content="The answer is 42.")
+        mock_agent = MagicMock()
+        mock_agent.ainvoke = AsyncMock(return_value={"messages": [final_msg]})
 
-        runner = LangChainAgentRunner(mock_llm, "You are helpful.", {})
+        runner = LangChainAgentRunner(mock_agent)
         result = await runner.run("What is the answer?")
 
         assert result.output == "The answer is 42."
         assert result.metrics.success is True
+        mock_agent.ainvoke.assert_called_once_with(
+            {"messages": [{"role": "user", "content": "What is the answer?"}]}
+        )
 
     @pytest.mark.asyncio
-    async def test_executes_tool_calls_and_returns_final_response(self):
-        """Should execute tool calls and continue loop until final response."""
+    async def test_aggregates_token_usage_across_messages(self):
+        """Should sum token usage from all messages in the graph result."""
         from ldai_langchain import LangChainAgentRunner
-        from langchain_core.messages import AIMessage
-
-        # First response: has a tool call
-        first_response = AIMessage(content="")
-        first_response.tool_calls = [
-            {"name": "get-weather", "args": {"location": "Paris"}, "id": "call_123"}
-        ]
 
-        # Second response: final answer
-        second_response = AIMessage(content="It is sunny in Paris.")
+        msg1 = AIMessage(content="intermediate")
+        msg1.usage_metadata = {'total_tokens': 10, 'input_tokens': 6, 'output_tokens': 4}
+        msg2 = AIMessage(content="final answer")
+        msg2.usage_metadata = {'total_tokens': 20, 'input_tokens': 12, 'output_tokens': 8}
 
-        mock_llm = MagicMock()
-        mock_llm.ainvoke = AsyncMock(side_effect=[first_response, second_response])
+        mock_agent = MagicMock()
+        mock_agent.ainvoke = AsyncMock(return_value={"messages": [msg1, msg2]})
 
-        weather_fn = MagicMock(return_value="Sunny, 25°C")
-        runner = LangChainAgentRunner(
-            mock_llm, "You are helpful.",
-            {'get-weather': weather_fn},
-        )
-        result = await runner.run("What is the weather in Paris?")
+        runner = LangChainAgentRunner(mock_agent)
+        result = await runner.run("Hello")
 
-        assert result.output == "It is sunny in Paris."
+        assert result.output == "final answer"
         assert result.metrics.success is True
-        weather_fn.assert_called_once_with(location="Paris")
+        assert result.metrics.usage is not None
+        assert result.metrics.usage.total == 30
+        assert result.metrics.usage.input == 18
+        assert result.metrics.usage.output == 12
 
     @pytest.mark.asyncio
     async def test_returns_failure_when_exception_thrown(self):
         """Should return unsuccessful AgentResult when exception is thrown."""
         from ldai_langchain import LangChainAgentRunner
 
-        mock_llm = MagicMock()
-        mock_llm.ainvoke = AsyncMock(side_effect=Exception("LLM Error"))
+        mock_agent = MagicMock()
+        mock_agent.ainvoke = AsyncMock(side_effect=Exception("Graph Error"))
 
-        runner = LangChainAgentRunner(mock_llm, "", {})
+        runner = LangChainAgentRunner(mock_agent)
         result = await runner.run("Hello")
 
         assert result.output == ""