From 0e0c30f8b2a8e34e5aa4cc9231cbd7a9fef0f8b8 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 11 Mar 2026 14:49:02 -0500 Subject: [PATCH 1/5] feat: Introduce ManagedAgent and AgentRunner implementations feat: Add OpenAIAgentRunner with agentic tool-calling loop feat: Add LangChainAgentRunner with agentic tool-calling loop feat: Add OpenAIRunnerFactory.create_agent(config, tools) -> OpenAIAgentRunner feat: Add LangChainRunnerFactory.create_agent(config, tools) -> LangChainAgentRunner feat: Add ManagedAgent wrapper holding AgentRunner and LDAIConfigTracker feat: Add LDAIClient.create_agent() returning ManagedAgent --- .../src/ldai_langchain/__init__.py | 2 + .../ldai_langchain/langchain_agent_runner.py | 121 ++++++++++++ .../langchain_runner_factory.py | 21 ++ .../tests/test_langchain_provider.py | 122 ++++++++++++ .../src/ldai_openai/__init__.py | 2 + .../src/ldai_openai/openai_agent_runner.py | 182 ++++++++++++++++++ .../src/ldai_openai/openai_runner_factory.py | 28 ++- .../tests/test_openai_provider.py | 137 +++++++++++++ packages/sdk/server-ai/src/ldai/__init__.py | 2 + packages/sdk/server-ai/src/ldai/client.py | 55 ++++++ .../sdk/server-ai/src/ldai/managed_agent.py | 53 +++++ .../sdk/server-ai/tests/test_managed_agent.py | 147 ++++++++++++++ 12 files changed, 871 insertions(+), 1 deletion(-) create mode 100644 packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py create mode 100644 packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_runner.py create mode 100644 packages/sdk/server-ai/src/ldai/managed_agent.py create mode 100644 packages/sdk/server-ai/tests/test_managed_agent.py diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py index cb455e5..b2f6233 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py @@ -1,3 +1,4 @@ +from ldai_langchain.langchain_agent_runner import LangChainAgentRunner from ldai_langchain.langchain_helper import ( convert_messages_to_langchain, create_langchain_model, @@ -16,6 +17,7 @@ '__version__', 'LangChainRunnerFactory', 'LangChainModelRunner', + 'LangChainAgentRunner', 'convert_messages_to_langchain', 'create_langchain_model', 'get_ai_metrics_from_response', diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py new file mode 100644 index 0000000..18700e4 --- /dev/null +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py @@ -0,0 +1,121 @@ +"""LangChain agent runner for LaunchDarkly AI SDK.""" + +from typing import Any, Dict, List + +from ldai import log +from ldai.providers import AgentResult, AgentRunner, ToolRegistry +from ldai.providers.types import LDAIMetrics +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage + +from ldai_langchain.langchain_helper import get_ai_metrics_from_response + + +class LangChainAgentRunner(AgentRunner): + """ + AgentRunner implementation for LangChain. + + Executes a single-agent loop using a LangChain BaseChatModel with tool calling. + Returned by LangChainRunnerFactory.create_agent(config, tools). + """ + + def __init__( + self, + llm: Any, + instructions: str, + tool_definitions: List[Dict[str, Any]], + tools: ToolRegistry, + ): + self._llm = llm + self._instructions = instructions + self._tool_definitions = tool_definitions + self._tools = tools + + 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. + + :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))) + + openai_tools = self._build_openai_tools() + model = self._llm.bind_tools(openai_tools) if openai_tools else self._llm + + raw_response = None + + try: + while True: + response: AIMessage = await model.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)) + + except Exception as error: + log.warning(f"LangChain agent run failed: {error}") + return AgentResult( + output="", + raw=raw_response, + metrics=LDAIMetrics(success=False, usage=None), + ) + + def _build_openai_tools(self) -> List[Dict[str, Any]]: + """Convert LD tool definitions to OpenAI function-calling format for bind_tools.""" + tools = [] + for td in self._tool_definitions: + if not isinstance(td, dict): + continue + if "type" in td: + tools.append(td) + elif "name" in td: + tools.append({ + "type": "function", + "function": { + "name": td["name"], + "description": td.get("description", ""), + "parameters": td.get("parameters", {"type": "object", "properties": {}}), + }, + }) + return tools + + def get_llm(self) -> Any: + """Return the underlying LangChain LLM.""" + return self._llm diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py index 402e295..e027845 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py @@ -1,3 +1,5 @@ +from typing import Any + from ldai.models import AIConfigKind from ldai.providers import AIProvider @@ -17,3 +19,22 @@ def create_model(self, config: AIConfigKind) -> LangChainModelRunner: """ llm = create_langchain_model(config) return LangChainModelRunner(llm) + + def create_agent(self, config: Any, tools: Any) -> '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 + + config_dict = config.to_dict() + model_dict = config_dict.get('model') or {} + parameters = dict(model_dict.get('parameters') or {}) + tool_definitions = parameters.pop('tools', []) or [] + instructions = config.instructions or '' if hasattr(config, 'instructions') else '' + + llm = LangChainHelper.create_langchain_model(config) + return LangChainAgentRunner(llm, instructions, tool_definitions, tools or {}) diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py index 9ce4e88..9bf2faf 100644 --- a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py +++ b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py @@ -330,3 +330,125 @@ def test_returns_underlying_llm(self): runner = LangChainModelRunner(mock_llm) assert runner.get_llm() is mock_llm + + +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.""" + from unittest.mock import patch + from ldai_langchain import LangChainAgentRunner + + mock_ai_config = MagicMock() + mock_ai_config.instructions = "You are a helpful assistant." + mock_ai_config.to_dict.return_value = { + 'model': { + 'name': 'gpt-4', + 'parameters': { + 'tools': [ + {'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}, + ], + }, + }, + 'provider': {'name': 'openai'}, + } + + with patch.object(LangChainHelper, 'create_langchain_model') as mock_create: + mock_llm = MagicMock() + mock_create.return_value = mock_llm + + 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 len(result._tool_definitions) == 1 + + def test_creates_agent_runner_with_no_tools(self): + """Should create LangChainAgentRunner with no tool definitions.""" + from unittest.mock import patch + from ldai_langchain import LangChainAgentRunner + + mock_ai_config = MagicMock() + mock_ai_config.instructions = "You are a helpful assistant." + mock_ai_config.to_dict.return_value = { + 'model': {'name': 'gpt-4', 'parameters': {}}, + 'provider': {'name': 'openai'}, + } + + with patch.object(LangChainHelper, 'create_langchain_model') as mock_create: + mock_create.return_value = MagicMock() + + factory = LangChainRunnerFactory() + result = factory.create_agent(mock_ai_config, {}) + + assert isinstance(result, LangChainAgentRunner) + assert result._tool_definitions == [] + + +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.""" + from ldai_langchain import LangChainAgentRunner + from langchain_core.messages import AIMessage + + mock_llm = MagicMock() + mock_response = AIMessage(content="The answer is 42.") + mock_llm.bind_tools = MagicMock(return_value=mock_llm) + mock_llm.ainvoke = AsyncMock(return_value=mock_response) + + runner = LangChainAgentRunner(mock_llm, "You are helpful.", [], {}) + result = await runner.run("What is the answer?") + + assert result.output == "The answer is 42." + assert result.metrics.success is True + + @pytest.mark.asyncio + async def test_executes_tool_calls_and_returns_final_response(self): + """Should execute tool calls and continue loop until final response.""" + 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.") + + mock_llm = MagicMock() + mock_llm.bind_tools = MagicMock(return_value=mock_llm) + mock_llm.ainvoke = AsyncMock(side_effect=[first_response, second_response]) + + weather_fn = MagicMock(return_value="Sunny, 25°C") + runner = LangChainAgentRunner( + mock_llm, "You are helpful.", + [{'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}], + {'get-weather': weather_fn}, + ) + result = await runner.run("What is the weather in Paris?") + + assert result.output == "It is sunny in Paris." + assert result.metrics.success is True + weather_fn.assert_called_once_with(location="Paris") + + @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.bind_tools = MagicMock(return_value=mock_llm) + mock_llm.ainvoke = AsyncMock(side_effect=Exception("LLM Error")) + + runner = LangChainAgentRunner(mock_llm, "", [], {}) + result = await runner.run("Hello") + + assert result.output == "" + assert result.metrics.success is False diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py index 8a8199b..a64f0b3 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py @@ -1,3 +1,4 @@ +from ldai_openai.openai_agent_runner import OpenAIAgentRunner from ldai_openai.openai_helper import ( convert_messages_to_openai, get_ai_metrics_from_response, @@ -9,6 +10,7 @@ __all__ = [ 'OpenAIRunnerFactory', 'OpenAIModelRunner', + 'OpenAIAgentRunner', 'convert_messages_to_openai', 'get_ai_metrics_from_response', 'get_ai_usage_from_response', diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_runner.py new file mode 100644 index 0000000..1a588d8 --- /dev/null +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_runner.py @@ -0,0 +1,182 @@ +"""OpenAI agent runner for LaunchDarkly AI SDK.""" + +import json +from typing import Any, Dict, List + +from ldai import log +from ldai.providers import AgentResult, AgentRunner, ToolRegistry +from ldai.providers.types import LDAIMetrics +from ldai.tracker import TokenUsage +from openai import AsyncOpenAI + +from ldai_openai.openai_helper import get_ai_metrics_from_response + + +class OpenAIAgentRunner(AgentRunner): + """ + AgentRunner implementation for OpenAI. + + Executes a single-agent loop using OpenAI Chat Completions with tool calling. + Returned by OpenAIRunnerFactory.create_agent(config, tools). + """ + + def __init__( + self, + client: AsyncOpenAI, + model_name: str, + parameters: Dict[str, Any], + instructions: str, + tool_definitions: List[Dict[str, Any]], + tools: ToolRegistry, + ): + self._client = client + self._model_name = model_name + self._parameters = parameters + self._instructions = instructions + self._tool_definitions = tool_definitions + self._tools = tools + + 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. + + :param input: The user prompt or input to the agent + :return: AgentResult with output, raw response, and aggregated metrics + """ + messages: List[Dict[str, Any]] = [] + if self._instructions: + messages.append({"role": "system", "content": self._instructions}) + messages.append({"role": "user", "content": str(input)}) + + total_input = 0 + total_output = 0 + raw_response = None + + try: + while True: + create_kwargs: Dict[str, Any] = { + "model": self._model_name, + "messages": messages, + **self._parameters, + } + openai_tools = self._build_openai_tools() + if openai_tools: + create_kwargs["tools"] = openai_tools + create_kwargs["tool_choice"] = "auto" + + response = await self._client.chat.completions.create(**create_kwargs) # type: ignore[arg-type] + raw_response = response + metrics = get_ai_metrics_from_response(response) + + if metrics.usage: + total_input += metrics.usage.input + total_output += metrics.usage.output + + if not response.choices: + break + + message = response.choices[0].message + + # Add assistant message to history + assistant_msg: Dict[str, Any] = { + "role": "assistant", + "content": message.content, + } + if message.tool_calls: + assistant_msg["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in message.tool_calls + ] + messages.append(assistant_msg) + + if not message.tool_calls: + total_tokens = total_input + total_output + return AgentResult( + output=message.content or "", + raw=raw_response, + metrics=LDAIMetrics( + success=True, + usage=TokenUsage( + total=total_tokens, + input=total_input, + output=total_output, + ) if total_tokens > 0 else None, + ), + ) + + # Execute tool calls and append results + for tool_call in message.tool_calls: + result = await self._call_tool( + tool_call.function.name, + tool_call.function.arguments, + ) + messages.append({ + "role": "tool", + "tool_call_id": tool_call.id, + "content": result, + }) + + except Exception as error: + log.warning(f"OpenAI agent run failed: {error}") + return AgentResult( + output="", + raw=raw_response, + metrics=LDAIMetrics(success=False, usage=None), + ) + + return AgentResult( + output="", + raw=raw_response, + metrics=LDAIMetrics(success=False, usage=None), + ) + + async def _call_tool(self, name: str, arguments_json: str) -> str: + """Execute a tool by name, returning the result as a string.""" + tool_fn = self._tools.get(name) + if not tool_fn: + log.warning(f"Tool '{name}' not found in registry") + return f"Tool '{name}' not found" + try: + args = json.loads(arguments_json) if arguments_json else {} + result = tool_fn(**args) + if hasattr(result, "__await__"): + result = await result + return str(result) + except Exception as error: + log.warning(f"Tool '{name}' execution failed: {error}") + return f"Tool execution failed: {error}" + + def _build_openai_tools(self) -> List[Dict[str, Any]]: + """Convert LD tool definitions to OpenAI function-calling format.""" + tools = [] + for td in self._tool_definitions: + if not isinstance(td, dict): + continue + if "type" in td: + # Already in OpenAI format + tools.append(td) + elif "name" in td: + # LD simplified format: {name, description, parameters} + tools.append({ + "type": "function", + "function": { + "name": td["name"], + "description": td.get("description", ""), + "parameters": td.get("parameters", {"type": "object", "properties": {}}), + }, + }) + return tools + + def get_client(self) -> AsyncOpenAI: + """Return the underlying AsyncOpenAI client.""" + return self._client diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py index d80fc01..0de2d03 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py @@ -1,5 +1,5 @@ import os -from typing import Optional +from typing import Any, Optional from ldai.models import AIConfigKind from ldai.providers import AIProvider @@ -36,6 +36,32 @@ def create_model(self, config: AIConfigKind) -> OpenAIModelRunner: parameters = model_dict.get('parameters') or {} return OpenAIModelRunner(self._client, model_name, parameters) + def create_agent(self, config: Any, tools: Any) -> 'OpenAIAgentRunner': + """ + Create a configured OpenAIAgentRunner for the given AI agent config. + + :param config: The LaunchDarkly AI agent configuration + :param tools: ToolRegistry mapping tool names to callables + :return: OpenAIAgentRunner ready to run the agent + """ + from ldai_openai.openai_agent_runner import OpenAIAgentRunner + + config_dict = config.to_dict() + model_dict = config_dict.get('model') or {} + model_name = model_dict.get('name', '') + parameters = dict(model_dict.get('parameters') or {}) + tool_definitions = parameters.pop('tools', []) or [] + instructions = config.instructions or '' if hasattr(config, 'instructions') else '' + + return OpenAIAgentRunner( + self._client, + model_name, + parameters, + instructions, + tool_definitions, + tools or {}, + ) + def get_client(self) -> AsyncOpenAI: """ Return the underlying AsyncOpenAI client. diff --git a/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py b/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py index 3a9de4d..5a4b850 100644 --- a/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py +++ b/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py @@ -318,3 +318,140 @@ def test_handles_missing_model_config(self): assert isinstance(result, OpenAIModelRunner) assert result._model_name == '' assert result._parameters == {} + + +class TestCreateAgent: + """Tests for OpenAIRunnerFactory.create_agent.""" + + def test_creates_agent_runner_with_instructions_and_tool_definitions(self): + """Should create OpenAIAgentRunner with instructions and tool definitions.""" + mock_ai_config = MagicMock() + mock_ai_config.instructions = "You are a helpful assistant." + mock_ai_config.to_dict.return_value = { + 'model': { + 'name': 'gpt-4', + 'parameters': { + 'temperature': 0.7, + 'tools': [ + {'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}, + ], + }, + }, + } + + mock_client = MagicMock() + factory = OpenAIRunnerFactory(mock_client) + result = factory.create_agent(mock_ai_config, {'get-weather': lambda loc: 'sunny'}) + + from ldai_openai import OpenAIAgentRunner + assert isinstance(result, OpenAIAgentRunner) + assert result._model_name == 'gpt-4' + assert result._instructions == "You are a helpful assistant." + assert result._parameters == {'temperature': 0.7} + assert len(result._tool_definitions) == 1 + assert result._tool_definitions[0]['name'] == 'get-weather' + + def test_creates_agent_runner_with_no_tools(self): + """Should create OpenAIAgentRunner with no tool definitions.""" + mock_ai_config = MagicMock() + mock_ai_config.instructions = "You are a helpful assistant." + mock_ai_config.to_dict.return_value = { + 'model': {'name': 'gpt-4', 'parameters': {}}, + } + + mock_client = MagicMock() + factory = OpenAIRunnerFactory(mock_client) + result = factory.create_agent(mock_ai_config, {}) + + from ldai_openai import OpenAIAgentRunner + assert isinstance(result, OpenAIAgentRunner) + assert result._tool_definitions == [] + + +class TestOpenAIAgentRunner: + """Tests for OpenAIAgentRunner.run.""" + + @pytest.fixture + def mock_client(self): + return MagicMock() + + @pytest.mark.asyncio + async def test_runs_agent_and_returns_result_with_no_tool_calls(self, mock_client): + """Should return AgentResult when model responds with no tool calls.""" + from ldai_openai import OpenAIAgentRunner + + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.content = "The answer is 42." + mock_response.choices[0].message.tool_calls = None + mock_response.usage = MagicMock() + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 5 + mock_response.usage.total_tokens = 15 + mock_client.chat.completions.create = AsyncMock(return_value=mock_response) + + runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, 'You are helpful.', [], {}) + result = await runner.run("What is the answer?") + + assert result.output == "The answer is 42." + assert result.metrics.success is True + + @pytest.mark.asyncio + async def test_executes_tool_calls_and_returns_final_response(self, mock_client): + """Should execute tool calls and continue loop until final response.""" + from ldai_openai import OpenAIAgentRunner + + # First response: has a tool call + tool_call = MagicMock() + tool_call.id = "call_123" + tool_call.function.name = "get-weather" + tool_call.function.arguments = '{"location": "Paris"}' + + first_response = MagicMock() + first_response.choices = [MagicMock()] + first_response.choices[0].message.content = None + first_response.choices[0].message.tool_calls = [tool_call] + first_response.usage = MagicMock() + first_response.usage.prompt_tokens = 10 + first_response.usage.completion_tokens = 5 + first_response.usage.total_tokens = 15 + + # Second response: final answer + second_response = MagicMock() + second_response.choices = [MagicMock()] + second_response.choices[0].message.content = "It is sunny in Paris." + second_response.choices[0].message.tool_calls = None + second_response.usage = MagicMock() + second_response.usage.prompt_tokens = 20 + second_response.usage.completion_tokens = 8 + second_response.usage.total_tokens = 28 + + mock_client.chat.completions.create = AsyncMock( + side_effect=[first_response, second_response] + ) + + weather_fn = MagicMock(return_value="Sunny, 25°C") + runner = OpenAIAgentRunner( + mock_client, 'gpt-4', {}, 'You are helpful.', + [{'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}], + {'get-weather': weather_fn}, + ) + result = await runner.run("What is the weather in Paris?") + + assert result.output == "It is sunny in Paris." + assert result.metrics.success is True + weather_fn.assert_called_once_with(location="Paris") + assert mock_client.chat.completions.create.call_count == 2 + + @pytest.mark.asyncio + async def test_returns_failure_when_exception_thrown(self, mock_client): + """Should return unsuccessful AgentResult when exception is thrown.""" + from ldai_openai import OpenAIAgentRunner + + mock_client.chat.completions.create = AsyncMock(side_effect=Exception("API Error")) + + runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, '', [], {}) + result = await runner.run("Hello") + + assert result.output == "" + assert result.metrics.success is False diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 944a0cb..1131eb4 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -6,6 +6,7 @@ from ldai.chat import Chat # Deprecated — use ManagedModel from ldai.client import LDAIClient from ldai.judge import Judge +from ldai.managed_agent import ManagedAgent from ldai.managed_model import ManagedModel from ldai.models import ( # Deprecated aliases for backward compatibility AIAgentConfig, @@ -55,6 +56,7 @@ 'AICompletionConfigDefault', 'AIJudgeConfig', 'AIJudgeConfigDefault', + 'ManagedAgent', 'ManagedModel', 'EvalScore', 'AgentGraphDefinition', diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 358f9eb..606cbaa 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -7,6 +7,7 @@ from ldai import log from ldai.agent_graph import AgentGraphDefinition from ldai.judge import Judge +from ldai.managed_agent import ManagedAgent from ldai.managed_model import ManagedModel from ldai.models import ( AIAgentConfig, @@ -25,12 +26,14 @@ ProviderConfig, ) from ldai.providers.runner_factory import RunnerFactory +from ldai.runners.types import ToolRegistry from ldai.sdk_info import AI_SDK_LANGUAGE, AI_SDK_NAME, AI_SDK_VERSION from ldai.tracker import AIGraphTracker, LDAIConfigTracker _TRACK_SDK_INFO = '$ld:ai:sdk:info' _TRACK_USAGE_COMPLETION_CONFIG = '$ld:ai:usage:completion-config' _TRACK_USAGE_CREATE_MODEL = '$ld:ai:usage:create-model' +_TRACK_USAGE_CREATE_AGENT = '$ld:ai:usage:create-agent' _TRACK_USAGE_JUDGE_CONFIG = '$ld:ai:usage:judge-config' _TRACK_USAGE_CREATE_JUDGE = '$ld:ai:usage:create-judge' _TRACK_USAGE_AGENT_CONFIG = '$ld:ai:usage:agent-config' @@ -374,6 +377,58 @@ async def create_chat( log.warning('create_chat() is deprecated, use create_model() instead') return await self.create_model(key, context, default, variables, default_ai_provider) + async def create_agent( + self, + key: str, + context: Context, + tools: Optional[ToolRegistry] = None, + default: Optional[AIAgentConfigDefault] = None, + variables: Optional[Dict[str, Any]] = None, + default_ai_provider: Optional[str] = None, + ) -> Optional[ManagedAgent]: + """ + Creates and returns a new ManagedAgent for AI agent invocations. + + :param key: The key identifying the AI agent configuration to use + :param context: Standard Context used when evaluating flags + :param tools: ToolRegistry mapping tool names to callable implementations + :param default: A default value representing a standard AI agent config result. + When not provided, a disabled config is used as the fallback. + :param variables: Dictionary of values for instruction interpolation + :param default_ai_provider: Optional default AI provider to use + :return: ManagedAgent instance or None if disabled/unsupported + + Example:: + + agent = await client.create_agent( + "customer-support-agent", + context, + tools={"get-order": fetch_order_fn}, + default=AIAgentConfigDefault( + enabled=True, + model=ModelConfig("gpt-4"), + provider=ProviderConfig("openai"), + instructions="You are a helpful customer support agent." + ), + ) + + if agent: + result = await agent.run("Where is my order?") + print(result.output) + """ + self._client.track(_TRACK_USAGE_CREATE_AGENT, context, key, 1) + log.debug(f"Creating managed agent for key: {key}") + config = self.__evaluate_agent(key, context, default or AIAgentConfigDefault.disabled(), variables) + + if not config.enabled or not config.tracker: + return None + + runner = await RunnerFactory.create_agent(config, tools or {}, default_ai_provider) + if not runner: + return None + + return ManagedAgent(config, config.tracker, runner) + def agent_config( self, key: str, diff --git a/packages/sdk/server-ai/src/ldai/managed_agent.py b/packages/sdk/server-ai/src/ldai/managed_agent.py new file mode 100644 index 0000000..6c2d17e --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/managed_agent.py @@ -0,0 +1,53 @@ +"""ManagedAgent — LaunchDarkly managed wrapper for agent invocations.""" + +from ldai.models import AIAgentConfig +from ldai.runners.agent_runner import AgentRunner +from ldai.runners.types import AgentResult +from ldai.tracker import LDAIConfigTracker + + +class ManagedAgent: + """ + LaunchDarkly managed wrapper for AI agent invocations. + + Holds an AgentRunner and an LDAIConfigTracker. Handles tracking automatically. + Obtain an instance via ``LDAIClient.create_agent()``. + """ + + def __init__( + self, + ai_config: AIAgentConfig, + tracker: LDAIConfigTracker, + agent_runner: AgentRunner, + ): + self._ai_config = ai_config + self._tracker = tracker + self._agent_runner = agent_runner + + async def run(self, input: str) -> AgentResult: + """ + Run the agent with the given input string. + + :param input: The user prompt or input to the agent + :return: AgentResult containing the agent's output and metrics + """ + return await self._tracker.track_metrics_of( + lambda: self._agent_runner.run(input), + lambda result: result.metrics, + ) + + def get_agent_runner(self) -> AgentRunner: + """ + Return the underlying AgentRunner for advanced use. + + :return: The AgentRunner instance. + """ + return self._agent_runner + + def get_config(self) -> AIAgentConfig: + """Return the AI agent config.""" + return self._ai_config + + def get_tracker(self) -> LDAIConfigTracker: + """Return the config tracker.""" + return self._tracker diff --git a/packages/sdk/server-ai/tests/test_managed_agent.py b/packages/sdk/server-ai/tests/test_managed_agent.py new file mode 100644 index 0000000..699cbef --- /dev/null +++ b/packages/sdk/server-ai/tests/test_managed_agent.py @@ -0,0 +1,147 @@ +"""Tests for ManagedAgent.""" + +import pytest +from unittest.mock import AsyncMock, MagicMock + +from ldai import LDAIClient, ManagedAgent +from ldai.managed_agent import ManagedAgent +from ldai.models import AIAgentConfig, AIAgentConfigDefault, ModelConfig, ProviderConfig +from ldai.runners.types import AgentResult +from ldai.providers.types import LDAIMetrics + +from ldclient import Config, Context, LDClient +from ldclient.integrations.test_data import TestData + + +@pytest.fixture +def td() -> TestData: + td = TestData.data_source() + td.update( + td.flag('customer-support-agent') + .variations({ + 'model': {'name': 'gpt-4', 'parameters': {'temperature': 0.3}}, + 'provider': {'name': 'openai'}, + 'instructions': 'You are a helpful customer support agent.', + '_ldMeta': {'enabled': True, 'variationKey': 'agent-v1', 'version': 1}, + }) + .variation_for_all(0) + ) + td.update( + td.flag('disabled-agent') + .variations({ + 'model': {'name': 'gpt-4'}, + '_ldMeta': {'enabled': False, 'variationKey': 'disabled-v1', 'version': 1}, + }) + .variation_for_all(0) + ) + return td + + +@pytest.fixture +def client(td: TestData) -> LDClient: + config = Config('sdk-key', update_processor_class=td, send_events=False) + return LDClient(config=config) + + +@pytest.fixture +def ldai_client(client: LDClient) -> LDAIClient: + return LDAIClient(client) + + +class TestManagedAgentRun: + """Tests for ManagedAgent.run.""" + + @pytest.mark.asyncio + async def test_run_delegates_to_agent_runner(self): + """Should delegate run() to the underlying AgentRunner.""" + mock_config = MagicMock(spec=AIAgentConfig) + mock_tracker = MagicMock() + mock_tracker.track_metrics_of = AsyncMock( + return_value=AgentResult( + output="Test response", + raw=None, + metrics=LDAIMetrics(success=True, usage=None), + ) + ) + mock_runner = MagicMock() + mock_runner.run = AsyncMock( + return_value=AgentResult( + output="Test response", + raw=None, + metrics=LDAIMetrics(success=True, usage=None), + ) + ) + + agent = ManagedAgent(mock_config, mock_tracker, mock_runner) + result = await agent.run("Hello") + + assert result.output == "Test response" + assert result.metrics.success is True + mock_tracker.track_metrics_of.assert_called_once() + + def test_get_agent_runner_returns_runner(self): + """Should return the underlying AgentRunner.""" + mock_runner = MagicMock() + agent = ManagedAgent(MagicMock(), MagicMock(), mock_runner) + + assert agent.get_agent_runner() is mock_runner + + def test_get_config_returns_config(self): + """Should return the AI agent config.""" + mock_config = MagicMock() + agent = ManagedAgent(mock_config, MagicMock(), MagicMock()) + + assert agent.get_config() is mock_config + + def test_get_tracker_returns_tracker(self): + """Should return the tracker.""" + mock_tracker = MagicMock() + agent = ManagedAgent(MagicMock(), mock_tracker, MagicMock()) + + assert agent.get_tracker() is mock_tracker + + +class TestLDAIClientCreateAgent: + """Tests for LDAIClient.create_agent.""" + + @pytest.mark.asyncio + async def test_returns_none_when_agent_is_disabled(self, ldai_client: LDAIClient): + """Should return None when agent config is disabled.""" + context = Context.create('user-key') + result = await ldai_client.create_agent('disabled-agent', context) + + assert result is None + + @pytest.mark.asyncio + async def test_returns_none_when_provider_unavailable(self, ldai_client: LDAIClient): + """Should return None when no AI provider is available.""" + import ldai.providers.runner_factory as rf + context = Context.create('user-key') + + original = rf.RunnerFactory.create_agent + rf.RunnerFactory.create_agent = AsyncMock(return_value=None) + try: + result = await ldai_client.create_agent('customer-support-agent', context) + assert result is None + finally: + rf.RunnerFactory.create_agent = original + + @pytest.mark.asyncio + async def test_returns_managed_agent_when_runner_available(self, ldai_client: LDAIClient): + """Should return ManagedAgent when runner is successfully created.""" + import ldai.providers.runner_factory as rf + context = Context.create('user-key') + + mock_runner = MagicMock() + mock_runner.run = AsyncMock( + return_value=AgentResult(output="Hello!", raw=None, metrics=LDAIMetrics(success=True, usage=None)) + ) + + original = rf.RunnerFactory.create_agent + rf.RunnerFactory.create_agent = AsyncMock(return_value=mock_runner) + try: + result = await ldai_client.create_agent('customer-support-agent', context) + assert isinstance(result, ManagedAgent) + assert result.get_agent_runner() is mock_runner + finally: + rf.RunnerFactory.create_agent = original From c1b87a6ba34d6da763c1f7fca4e9338d00792334 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Thu, 26 Mar 2026 08:27:42 -0500 Subject: [PATCH 2/5] feat: update managed-agent to use track_metrics_of_async and add provider helper tests feat: add TestGetAIUsageFromResponse and TestGetToolCallsFromResponse test coverage for LangChainHelper feat: add TestGetAIUsageFromResponse test coverage for OpenAIHelper fix: update ManagedAgent.invoke to use track_metrics_of_async --- .../tests/test_langchain_provider.py | 65 +++++++++++++++++++ .../tests/test_openai_provider.py | 29 +++++++++ .../sdk/server-ai/src/ldai/managed_agent.py | 2 +- 3 files changed, 95 insertions(+), 1 deletion(-) diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py index 9bf2faf..484e952 100644 --- a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py +++ b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py @@ -127,6 +127,71 @@ def test_creates_metrics_with_success_true_and_no_usage_when_metadata_missing(se assert result.success is True assert result.usage is None + def test_usage_metadata_preferred_over_response_metadata(self): + """usage_metadata should be used when it has non-zero counts.""" + mock_response = AIMessage(content='Test') + mock_response.usage_metadata = { + 'total_tokens': 10, + 'input_tokens': 4, + 'output_tokens': 6, + } + mock_response.response_metadata = { + 'tokenUsage': { + 'totalTokens': 999, + 'promptTokens': 500, + 'completionTokens': 499, + }, + } + usage = LangChainHelper.get_ai_usage_from_response(mock_response) + assert usage is not None + assert usage.total == 10 + assert usage.input == 4 + assert usage.output == 6 + + +class TestGetAIUsageFromResponse: + """Tests for LangChainHelper.get_ai_usage_from_response.""" + + def test_returns_none_when_no_usage(self): + msg = AIMessage(content='hi') + assert LangChainHelper.get_ai_usage_from_response(msg) is None + + def test_returns_none_when_all_zeros_in_metadata(self): + msg = AIMessage(content='hi') + msg.usage_metadata = {'total_tokens': 0, 'input_tokens': 0, 'output_tokens': 0} + assert LangChainHelper.get_ai_usage_from_response(msg) is None + + +class TestGetToolCallsFromResponse: + """Tests for LangChainHelper.get_tool_calls_from_response.""" + + def test_returns_empty_when_no_tool_calls(self): + msg = AIMessage(content='hi') + assert LangChainHelper.get_tool_calls_from_response(msg) == [] + + def test_returns_empty_when_tool_calls_not_a_sequence(self): + msg = AIMessage(content='hi') + msg.tool_calls = None # type: ignore + assert LangChainHelper.get_tool_calls_from_response(msg) == [] + + def test_extracts_names_from_dict_tool_calls(self): + msg = AIMessage(content='') + msg.tool_calls = [ # type: ignore + {'name': 'search', 'args': {}, 'id': '1'}, + {'name': 'calc', 'args': {}, 'id': '2'}, + ] + assert LangChainHelper.get_tool_calls_from_response(msg) == ['search', 'calc'] + + def test_returns_empty_when_tool_calls_is_not_a_list(self): + msg = AIMessage(content='hi') + msg.tool_calls = () # type: ignore + assert LangChainHelper.get_tool_calls_from_response(msg) == [] + + def test_skips_entries_without_name(self): + msg = AIMessage(content='') + msg.tool_calls = [{'name': 'a', 'id': '1'}, {}, {'name': 'b', 'id': '2'}] # type: ignore + assert LangChainHelper.get_tool_calls_from_response(msg) == ['a', 'b'] + class TestMapProvider: """Tests for map_provider.""" diff --git a/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py b/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py index 5a4b850..927cf73 100644 --- a/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py +++ b/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py @@ -8,6 +8,35 @@ from ldai_openai import OpenAIModelRunner, OpenAIRunnerFactory, get_ai_metrics_from_response +class TestGetAIUsageFromResponse: + """Tests for OpenAIHelper.get_ai_usage_from_response.""" + + def test_returns_usage_when_present(self): + mock_response = MagicMock() + mock_response.usage = MagicMock() + mock_response.usage.prompt_tokens = 50 + mock_response.usage.completion_tokens = 50 + mock_response.usage.total_tokens = 100 + u = OpenAIHelper.get_ai_usage_from_response(mock_response) + assert u is not None + assert u.total == 100 + assert u.input == 50 + assert u.output == 50 + + def test_returns_none_when_usage_missing(self): + mock_response = MagicMock() + mock_response.usage = None + assert OpenAIHelper.get_ai_usage_from_response(mock_response) is None + + def test_returns_none_when_all_counts_zero(self): + mock_response = MagicMock() + mock_response.usage = MagicMock() + mock_response.usage.total_tokens = 0 + mock_response.usage.prompt_tokens = 0 + mock_response.usage.completion_tokens = 0 + assert OpenAIHelper.get_ai_usage_from_response(mock_response) is None + + class TestGetAIMetricsFromResponse: """Tests for get_ai_metrics_from_response.""" diff --git a/packages/sdk/server-ai/src/ldai/managed_agent.py b/packages/sdk/server-ai/src/ldai/managed_agent.py index 6c2d17e..8d0a72e 100644 --- a/packages/sdk/server-ai/src/ldai/managed_agent.py +++ b/packages/sdk/server-ai/src/ldai/managed_agent.py @@ -31,7 +31,7 @@ async def run(self, input: str) -> AgentResult: :param input: The user prompt or input to the agent :return: AgentResult containing the agent's output and metrics """ - return await self._tracker.track_metrics_of( + return await self._tracker.track_metrics_of_async( lambda: self._agent_runner.run(input), lambda result: result.metrics, ) From 90b548f2f9e80f064ad4596bd3103dddbd56bea5 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Thu, 26 Mar 2026 14:47:36 -0500 Subject: [PATCH 3/5] fix: resolve lint errors from rebase onto main Co-Authored-By: Claude Sonnet 4.6 --- .../ldai_langchain/langchain_agent_runner.py | 8 ++++++- .../src/ldai_langchain/langchain_helper.py | 10 ++++----- .../langchain_runner_factory.py | 7 +++++-- .../tests/test_langchain_provider.py | 21 ++++++++++--------- .../src/ldai_openai/openai_helper.py | 10 ++++----- .../src/ldai_openai/openai_runner_factory.py | 5 ++++- .../tests/test_openai_provider.py | 8 +++---- packages/sdk/server-ai/src/ldai/client.py | 4 ++-- .../sdk/server-ai/src/ldai/managed_agent.py | 3 +-- .../sdk/server-ai/tests/test_managed_agent.py | 10 ++++----- 10 files changed, 49 insertions(+), 37 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py index 18700e4..27327ea 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py @@ -2,10 +2,16 @@ from typing import Any, Dict, List +from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) from ldai import log from ldai.providers import AgentResult, AgentRunner, ToolRegistry from ldai.providers.types import LDAIMetrics -from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage from ldai_langchain.langchain_helper import get_ai_metrics_from_response diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py index e160061..29c2f14 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py @@ -88,11 +88,11 @@ def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: :return: TokenUsage or None if unavailable """ if hasattr(response, 'usage_metadata') and response.usage_metadata: - return TokenUsage( - total=response.usage_metadata.get('total_tokens', 0), - input=response.usage_metadata.get('input_tokens', 0), - output=response.usage_metadata.get('output_tokens', 0), - ) + total = response.usage_metadata.get('total_tokens', 0) + inp = response.usage_metadata.get('input_tokens', 0) + out = response.usage_metadata.get('output_tokens', 0) + if total or inp or out: + return TokenUsage(total=total, input=inp, output=out) if hasattr(response, 'response_metadata') and response.response_metadata: token_usage = ( response.response_metadata.get('tokenUsage') diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py index e027845..3a2c10b 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py @@ -1,6 +1,9 @@ -from typing import Any +from typing import TYPE_CHECKING, Any from ldai.models import AIConfigKind + +if TYPE_CHECKING: + from ldai_langchain.langchain_agent_runner import LangChainAgentRunner from ldai.providers import AIProvider from ldai_langchain.langchain_helper import create_langchain_model @@ -36,5 +39,5 @@ def create_agent(self, config: Any, tools: Any) -> 'LangChainAgentRunner': tool_definitions = parameters.pop('tools', []) or [] instructions = config.instructions or '' if hasattr(config, 'instructions') else '' - llm = LangChainHelper.create_langchain_model(config) + llm = create_langchain_model(config) return LangChainAgentRunner(llm, instructions, tool_definitions, tools or {}) diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py index 484e952..38ad72f 100644 --- a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py +++ b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py @@ -12,6 +12,7 @@ LangChainRunnerFactory, convert_messages_to_langchain, get_ai_metrics_from_response, + get_ai_usage_from_response, get_tool_calls_from_response, map_provider, sum_token_usage_from_messages, @@ -142,7 +143,7 @@ def test_usage_metadata_preferred_over_response_metadata(self): 'completionTokens': 499, }, } - usage = LangChainHelper.get_ai_usage_from_response(mock_response) + usage = get_ai_usage_from_response(mock_response) assert usage is not None assert usage.total == 10 assert usage.input == 4 @@ -154,12 +155,12 @@ class TestGetAIUsageFromResponse: def test_returns_none_when_no_usage(self): msg = AIMessage(content='hi') - assert LangChainHelper.get_ai_usage_from_response(msg) is None + assert get_ai_usage_from_response(msg) is None def test_returns_none_when_all_zeros_in_metadata(self): msg = AIMessage(content='hi') msg.usage_metadata = {'total_tokens': 0, 'input_tokens': 0, 'output_tokens': 0} - assert LangChainHelper.get_ai_usage_from_response(msg) is None + assert get_ai_usage_from_response(msg) is None class TestGetToolCallsFromResponse: @@ -167,12 +168,12 @@ class TestGetToolCallsFromResponse: def test_returns_empty_when_no_tool_calls(self): msg = AIMessage(content='hi') - assert LangChainHelper.get_tool_calls_from_response(msg) == [] + assert get_tool_calls_from_response(msg) == [] def test_returns_empty_when_tool_calls_not_a_sequence(self): msg = AIMessage(content='hi') msg.tool_calls = None # type: ignore - assert LangChainHelper.get_tool_calls_from_response(msg) == [] + assert get_tool_calls_from_response(msg) == [] def test_extracts_names_from_dict_tool_calls(self): msg = AIMessage(content='') @@ -180,17 +181,17 @@ def test_extracts_names_from_dict_tool_calls(self): {'name': 'search', 'args': {}, 'id': '1'}, {'name': 'calc', 'args': {}, 'id': '2'}, ] - assert LangChainHelper.get_tool_calls_from_response(msg) == ['search', 'calc'] + assert get_tool_calls_from_response(msg) == ['search', 'calc'] def test_returns_empty_when_tool_calls_is_not_a_list(self): msg = AIMessage(content='hi') msg.tool_calls = () # type: ignore - assert LangChainHelper.get_tool_calls_from_response(msg) == [] + assert get_tool_calls_from_response(msg) == [] def test_skips_entries_without_name(self): msg = AIMessage(content='') msg.tool_calls = [{'name': 'a', 'id': '1'}, {}, {'name': 'b', 'id': '2'}] # type: ignore - assert LangChainHelper.get_tool_calls_from_response(msg) == ['a', 'b'] + assert get_tool_calls_from_response(msg) == ['a', 'b'] class TestMapProvider: @@ -419,7 +420,7 @@ def test_creates_agent_runner_with_instructions_and_tool_definitions(self): 'provider': {'name': 'openai'}, } - with patch.object(LangChainHelper, 'create_langchain_model') as mock_create: + with patch('ldai_langchain.langchain_runner_factory.create_langchain_model') as mock_create: mock_llm = MagicMock() mock_create.return_value = mock_llm @@ -442,7 +443,7 @@ def test_creates_agent_runner_with_no_tools(self): 'provider': {'name': 'openai'}, } - with patch.object(LangChainHelper, 'create_langchain_model') as mock_create: + with patch('ldai_langchain.langchain_runner_factory.create_langchain_model') as mock_create: mock_create.return_value = MagicMock() factory = LangChainRunnerFactory() diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_helper.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_helper.py index 3cc41e4..dc40ec6 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_helper.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_helper.py @@ -28,11 +28,11 @@ def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: """ if hasattr(response, 'usage') and response.usage: u = response.usage - return TokenUsage( - total=getattr(u, 'total_tokens', None) or 0, - input=getattr(u, 'prompt_tokens', None) or 0, - output=getattr(u, 'completion_tokens', None) or 0, - ) + total = getattr(u, 'total_tokens', None) or 0 + inp = getattr(u, 'prompt_tokens', None) or 0 + out = getattr(u, 'completion_tokens', None) or 0 + if total or inp or out: + return TokenUsage(total=total, input=inp, output=out) return None diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py index 0de2d03..f7781ad 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py @@ -1,5 +1,5 @@ import os -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from ldai.models import AIConfigKind from ldai.providers import AIProvider @@ -7,6 +7,9 @@ from ldai_openai.openai_model_runner import OpenAIModelRunner +if TYPE_CHECKING: + from ldai_openai.openai_agent_runner import OpenAIAgentRunner + class OpenAIRunnerFactory(AIProvider): """OpenAI ``AIProvider`` implementation for the LaunchDarkly AI SDK.""" diff --git a/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py b/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py index 927cf73..afdb524 100644 --- a/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py +++ b/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py @@ -5,7 +5,7 @@ from ldai import LDMessage -from ldai_openai import OpenAIModelRunner, OpenAIRunnerFactory, get_ai_metrics_from_response +from ldai_openai import OpenAIModelRunner, OpenAIRunnerFactory, get_ai_metrics_from_response, get_ai_usage_from_response class TestGetAIUsageFromResponse: @@ -17,7 +17,7 @@ def test_returns_usage_when_present(self): mock_response.usage.prompt_tokens = 50 mock_response.usage.completion_tokens = 50 mock_response.usage.total_tokens = 100 - u = OpenAIHelper.get_ai_usage_from_response(mock_response) + u = get_ai_usage_from_response(mock_response) assert u is not None assert u.total == 100 assert u.input == 50 @@ -26,7 +26,7 @@ def test_returns_usage_when_present(self): def test_returns_none_when_usage_missing(self): mock_response = MagicMock() mock_response.usage = None - assert OpenAIHelper.get_ai_usage_from_response(mock_response) is None + assert get_ai_usage_from_response(mock_response) is None def test_returns_none_when_all_counts_zero(self): mock_response = MagicMock() @@ -34,7 +34,7 @@ def test_returns_none_when_all_counts_zero(self): mock_response.usage.total_tokens = 0 mock_response.usage.prompt_tokens = 0 mock_response.usage.completion_tokens = 0 - assert OpenAIHelper.get_ai_usage_from_response(mock_response) is None + assert get_ai_usage_from_response(mock_response) is None class TestGetAIMetricsFromResponse: diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 606cbaa..21bcf54 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -25,8 +25,8 @@ ModelConfig, ProviderConfig, ) +from ldai.providers import ToolRegistry from ldai.providers.runner_factory import RunnerFactory -from ldai.runners.types import ToolRegistry from ldai.sdk_info import AI_SDK_LANGUAGE, AI_SDK_NAME, AI_SDK_VERSION from ldai.tracker import AIGraphTracker, LDAIConfigTracker @@ -423,7 +423,7 @@ async def create_agent( if not config.enabled or not config.tracker: return None - runner = await RunnerFactory.create_agent(config, tools or {}, default_ai_provider) + runner = RunnerFactory.create_agent(config, tools or {}, default_ai_provider) if not runner: return None diff --git a/packages/sdk/server-ai/src/ldai/managed_agent.py b/packages/sdk/server-ai/src/ldai/managed_agent.py index 8d0a72e..12c4d9b 100644 --- a/packages/sdk/server-ai/src/ldai/managed_agent.py +++ b/packages/sdk/server-ai/src/ldai/managed_agent.py @@ -1,8 +1,7 @@ """ManagedAgent — LaunchDarkly managed wrapper for agent invocations.""" from ldai.models import AIAgentConfig -from ldai.runners.agent_runner import AgentRunner -from ldai.runners.types import AgentResult +from ldai.providers import AgentResult, AgentRunner from ldai.tracker import LDAIConfigTracker diff --git a/packages/sdk/server-ai/tests/test_managed_agent.py b/packages/sdk/server-ai/tests/test_managed_agent.py index 699cbef..60cf7db 100644 --- a/packages/sdk/server-ai/tests/test_managed_agent.py +++ b/packages/sdk/server-ai/tests/test_managed_agent.py @@ -6,7 +6,7 @@ from ldai import LDAIClient, ManagedAgent from ldai.managed_agent import ManagedAgent from ldai.models import AIAgentConfig, AIAgentConfigDefault, ModelConfig, ProviderConfig -from ldai.runners.types import AgentResult +from ldai.providers import AgentResult from ldai.providers.types import LDAIMetrics from ldclient import Config, Context, LDClient @@ -56,7 +56,7 @@ async def test_run_delegates_to_agent_runner(self): """Should delegate run() to the underlying AgentRunner.""" mock_config = MagicMock(spec=AIAgentConfig) mock_tracker = MagicMock() - mock_tracker.track_metrics_of = AsyncMock( + mock_tracker.track_metrics_of_async = AsyncMock( return_value=AgentResult( output="Test response", raw=None, @@ -77,7 +77,7 @@ async def test_run_delegates_to_agent_runner(self): assert result.output == "Test response" assert result.metrics.success is True - mock_tracker.track_metrics_of.assert_called_once() + mock_tracker.track_metrics_of_async.assert_called_once() def test_get_agent_runner_returns_runner(self): """Should return the underlying AgentRunner.""" @@ -119,7 +119,7 @@ async def test_returns_none_when_provider_unavailable(self, ldai_client: LDAICli context = Context.create('user-key') original = rf.RunnerFactory.create_agent - rf.RunnerFactory.create_agent = AsyncMock(return_value=None) + rf.RunnerFactory.create_agent = MagicMock(return_value=None) try: result = await ldai_client.create_agent('customer-support-agent', context) assert result is None @@ -138,7 +138,7 @@ async def test_returns_managed_agent_when_runner_available(self, ldai_client: LD ) original = rf.RunnerFactory.create_agent - rf.RunnerFactory.create_agent = AsyncMock(return_value=mock_runner) + rf.RunnerFactory.create_agent = MagicMock(return_value=mock_runner) try: result = await ldai_client.create_agent('customer-support-agent', context) assert isinstance(result, ManagedAgent) From ab9c4bf029d32a6fd12de6ad7a146917d039ff95 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 11:20:10 -0500 Subject: [PATCH 4/5] simplifying tool configuration --- .../ldai_langchain/langchain_agent_runner.py | 29 +------- .../src/ldai_langchain/langchain_helper.py | 66 +++++++++++++++++-- .../langchain_runner_factory.py | 17 ++--- .../tests/test_langchain_provider.py | 11 +--- .../src/ldai_openai/openai_runner_factory.py | 6 +- .../src/ldai/providers/ai_provider.py | 4 +- 6 files changed, 79 insertions(+), 54 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py index 27327ea..369b291 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py @@ -1,6 +1,6 @@ """LangChain agent runner for LaunchDarkly AI SDK.""" -from typing import Any, Dict, List +from typing import Any, List from langchain_core.messages import ( AIMessage, @@ -21,6 +21,7 @@ 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. Returned by LangChainRunnerFactory.create_agent(config, tools). """ @@ -28,12 +29,10 @@ def __init__( self, llm: Any, instructions: str, - tool_definitions: List[Dict[str, Any]], tools: ToolRegistry, ): self._llm = llm self._instructions = instructions - self._tool_definitions = tool_definitions self._tools = tools async def run(self, input: Any) -> AgentResult: @@ -51,14 +50,11 @@ async def run(self, input: Any) -> AgentResult: messages.append(SystemMessage(content=self._instructions)) messages.append(HumanMessage(content=str(input))) - openai_tools = self._build_openai_tools() - model = self._llm.bind_tools(openai_tools) if openai_tools else self._llm - raw_response = None try: while True: - response: AIMessage = await model.ainvoke(messages) + response: AIMessage = await self._llm.ainvoke(messages) raw_response = response messages.append(response) @@ -103,25 +99,6 @@ async def run(self, input: Any) -> AgentResult: metrics=LDAIMetrics(success=False, usage=None), ) - def _build_openai_tools(self) -> List[Dict[str, Any]]: - """Convert LD tool definitions to OpenAI function-calling format for bind_tools.""" - tools = [] - for td in self._tool_definitions: - if not isinstance(td, dict): - continue - if "type" in td: - tools.append(td) - elif "name" in td: - tools.append({ - "type": "function", - "function": { - "name": td["name"], - "description": td.get("description", ""), - "parameters": td.get("parameters", {"type": "object", "properties": {}}), - }, - }) - return tools - def get_llm(self) -> Any: """Return the underlying LangChain LLM.""" return self._llm diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py index 29c2f14..5bcc6d8 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py @@ -2,8 +2,9 @@ from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage -from ldai import LDMessage +from ldai import LDMessage, log from ldai.models import AIConfigKind +from ldai.providers import ToolRegistry from ldai.providers.types import LDAIMetrics from ldai.tracker import TokenUsage @@ -50,12 +51,18 @@ def convert_messages_to_langchain( return result -def create_langchain_model(ai_config: AIConfigKind) -> BaseChatModel: +def create_langchain_model(ai_config: AIConfigKind, tool_registry: Optional[ToolRegistry] = None) -> BaseChatModel: """ Create a LangChain BaseChatModel from a LaunchDarkly AI configuration. + If the config includes tool definitions and a tool_registry is provided, tools found + in the registry are bound to the model. Tools not found in the registry are skipped + with a warning. Built-in provider tools (e.g. code_interpreter) are not supported + via LangChain's bind_tools abstraction and are skipped with a warning. + :param ai_config: The LaunchDarkly AI configuration - :return: A configured LangChain BaseChatModel + :param tool_registry: Optional registry mapping tool names to callable implementations + :return: A configured LangChain BaseChatModel, with tools bound if applicable """ from langchain.chat_models import init_chat_model @@ -66,6 +73,7 @@ def create_langchain_model(ai_config: AIConfigKind) -> BaseChatModel: model_name = model_dict.get('name', '') provider = provider_dict.get('name', '') parameters = dict(model_dict.get('parameters') or {}) + tool_definitions = parameters.pop('tools', []) or [] mapped_provider = map_provider(provider) # Bedrock requires the foundation provider (e.g. Bedrock:Anthropic) passed in @@ -73,12 +81,62 @@ def create_langchain_model(ai_config: AIConfigKind) -> BaseChatModel: if mapped_provider == 'bedrock_converse' and 'provider' not in parameters: parameters['provider'] = provider.removeprefix('bedrock:') - return init_chat_model( + model = init_chat_model( model_name, model_provider=mapped_provider, **parameters, ) + if tool_definitions: + bindable = _resolve_tools_for_langchain(tool_definitions, tool_registry or {}) + if bindable: + model = model.bind_tools(bindable) + + return model + + +def _resolve_tools_for_langchain( + tool_definitions: List[Dict[str, Any]], + tool_registry: ToolRegistry, +) -> List[Dict[str, Any]]: + """ + Match LD tool definitions against a registry, returning function-calling tool dicts + for tools that have a callable implementation. Built-in provider tools and tools + missing from the registry are skipped with a warning. + """ + bindable = [] + 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's " + "bind_tools abstraction 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 + + bindable.append({ + 'type': 'function', + 'function': { + 'name': name, + 'description': td.get('description', ''), + 'parameters': td.get('parameters', {'type': 'object', 'properties': {}}), + }, + }) + + return bindable + def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: """ diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py index 3a2c10b..a858635 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py @@ -1,10 +1,10 @@ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Optional from ldai.models import AIConfigKind if TYPE_CHECKING: from ldai_langchain.langchain_agent_runner import LangChainAgentRunner -from ldai.providers import AIProvider +from ldai.providers import AIProvider, ToolRegistry from ldai_langchain.langchain_helper import create_langchain_model from ldai_langchain.langchain_model_runner import LangChainModelRunner @@ -23,7 +23,7 @@ def create_model(self, config: AIConfigKind) -> LangChainModelRunner: llm = create_langchain_model(config) return LangChainModelRunner(llm) - def create_agent(self, config: Any, tools: Any) -> 'LangChainAgentRunner': + def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> 'LangChainAgentRunner': """ Create a configured LangChainAgentRunner for the given AI agent config. @@ -33,11 +33,6 @@ def create_agent(self, config: Any, tools: Any) -> 'LangChainAgentRunner': """ from ldai_langchain.langchain_agent_runner import LangChainAgentRunner - config_dict = config.to_dict() - model_dict = config_dict.get('model') or {} - parameters = dict(model_dict.get('parameters') or {}) - tool_definitions = parameters.pop('tools', []) or [] - instructions = config.instructions or '' if hasattr(config, 'instructions') else '' - - llm = create_langchain_model(config) - return LangChainAgentRunner(llm, instructions, tool_definitions, tools or {}) + instructions = (config.instructions or '') if hasattr(config, 'instructions') else '' + llm = create_langchain_model(config, tool_registry=tools or {}) + return LangChainAgentRunner(llm, instructions, tools or {}) diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py index 38ad72f..dd88b3e 100644 --- a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py +++ b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py @@ -429,7 +429,6 @@ def test_creates_agent_runner_with_instructions_and_tool_definitions(self): assert isinstance(result, LangChainAgentRunner) assert result._instructions == "You are a helpful assistant." - assert len(result._tool_definitions) == 1 def test_creates_agent_runner_with_no_tools(self): """Should create LangChainAgentRunner with no tool definitions.""" @@ -450,7 +449,7 @@ def test_creates_agent_runner_with_no_tools(self): result = factory.create_agent(mock_ai_config, {}) assert isinstance(result, LangChainAgentRunner) - assert result._tool_definitions == [] + assert result._tools == {} class TestLangChainAgentRunner: @@ -464,10 +463,9 @@ async def test_runs_agent_and_returns_result_with_no_tool_calls(self): mock_llm = MagicMock() mock_response = AIMessage(content="The answer is 42.") - mock_llm.bind_tools = MagicMock(return_value=mock_llm) mock_llm.ainvoke = AsyncMock(return_value=mock_response) - runner = LangChainAgentRunner(mock_llm, "You are helpful.", [], {}) + runner = LangChainAgentRunner(mock_llm, "You are helpful.", {}) result = await runner.run("What is the answer?") assert result.output == "The answer is 42." @@ -489,13 +487,11 @@ async def test_executes_tool_calls_and_returns_final_response(self): second_response = AIMessage(content="It is sunny in Paris.") mock_llm = MagicMock() - mock_llm.bind_tools = MagicMock(return_value=mock_llm) mock_llm.ainvoke = AsyncMock(side_effect=[first_response, second_response]) weather_fn = MagicMock(return_value="Sunny, 25°C") runner = LangChainAgentRunner( mock_llm, "You are helpful.", - [{'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}], {'get-weather': weather_fn}, ) result = await runner.run("What is the weather in Paris?") @@ -510,10 +506,9 @@ async def test_returns_failure_when_exception_thrown(self): from ldai_langchain import LangChainAgentRunner mock_llm = MagicMock() - mock_llm.bind_tools = MagicMock(return_value=mock_llm) mock_llm.ainvoke = AsyncMock(side_effect=Exception("LLM Error")) - runner = LangChainAgentRunner(mock_llm, "", [], {}) + runner = LangChainAgentRunner(mock_llm, "", {}) result = await runner.run("Hello") assert result.output == "" diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py index f7781ad..f69e1fb 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Any, Optional from ldai.models import AIConfigKind -from ldai.providers import AIProvider +from ldai.providers import AIProvider, ToolRegistry from openai import AsyncOpenAI from ldai_openai.openai_model_runner import OpenAIModelRunner @@ -39,7 +39,7 @@ def create_model(self, config: AIConfigKind) -> OpenAIModelRunner: parameters = model_dict.get('parameters') or {} return OpenAIModelRunner(self._client, model_name, parameters) - def create_agent(self, config: Any, tools: Any) -> 'OpenAIAgentRunner': + def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> 'OpenAIAgentRunner': """ Create a configured OpenAIAgentRunner for the given AI agent config. @@ -54,7 +54,7 @@ def create_agent(self, config: Any, tools: Any) -> 'OpenAIAgentRunner': model_name = model_dict.get('name', '') parameters = dict(model_dict.get('parameters') or {}) tool_definitions = parameters.pop('tools', []) or [] - instructions = config.instructions or '' if hasattr(config, 'instructions') else '' + instructions = (config.instructions or '') if hasattr(config, 'instructions') else '' return OpenAIAgentRunner( self._client, diff --git a/packages/sdk/server-ai/src/ldai/providers/ai_provider.py b/packages/sdk/server-ai/src/ldai/providers/ai_provider.py index 576f1c1..171c50c 100644 --- a/packages/sdk/server-ai/src/ldai/providers/ai_provider.py +++ b/packages/sdk/server-ai/src/ldai/providers/ai_provider.py @@ -3,7 +3,7 @@ from ldai import log from ldai.models import LDMessage -from ldai.providers.types import ModelResponse, StructuredResponse +from ldai.providers.types import ModelResponse, StructuredResponse, ToolRegistry class AIProvider(ABC): @@ -73,7 +73,7 @@ def create_model(self, config: Any) -> Optional[Any]: log.warning('create_model not implemented by this provider') return None - def create_agent(self, config: Any, tools: Any) -> Optional[Any]: + def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> Optional[Any]: """ Create a configured agent executor for the given AI config and tool registry. From e4b3830341913489fded1df808d91408cea6b4a4 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 14:55:43 -0500 Subject: [PATCH 5/5] simplify agent loop to use built-ins --- .../ldai_langchain/langchain_agent_runner.py | 103 ++++++------------ .../src/ldai_langchain/langchain_helper.py | 53 ++++++++- .../langchain_runner_factory.py | 27 +++-- .../tests/test_langchain_provider.py | 78 ++++++------- 4 files changed, 139 insertions(+), 122 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py index 369b291..b0a1c85 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py +++ b/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 diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py index 5bcc6d8..4357cad 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py +++ b/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. diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py index a858635..ba5087f 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py +++ b/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 - -if TYPE_CHECKING: - from ldai_langchain.langchain_agent_runner import LangChainAgentRunner from ldai.providers import AIProvider, ToolRegistry -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 @@ -23,7 +25,7 @@ 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. @@ -31,8 +33,13 @@ def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> 'La :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) diff --git a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py b/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py index dd88b3e..3aa04a9 100644 --- a/packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py +++ b/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 == ""