From 494780ee8f657c4426e0290f3441380ee446527e Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 11 Mar 2026 14:49:02 -0500 Subject: [PATCH 1/9] 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 | 19 ++ .../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 | 26 +++ .../tests/test_openai_provider.py | 137 +++++++++++++ packages/sdk/server-ai/src/ldai/__init__.py | 2 + packages/sdk/server-ai/src/ldai/client.py | 54 ++++++ .../sdk/server-ai/src/ldai/managed_agent.py | 53 +++++ .../sdk/server-ai/tests/test_managed_agent.py | 147 ++++++++++++++ 12 files changed, 867 insertions(+) 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 04e299d..b7a9d62 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, @@ -18,6 +19,7 @@ 'LangChainRunnerFactory', 'LangGraphAgentGraphRunner', '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 43febb5..19022df 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 @@ -32,3 +32,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 422c059..e1856e2 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,4 +1,5 @@ from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner +from ldai_openai.openai_agent_runner import OpenAIAgentRunner from ldai_openai.openai_helper import ( convert_messages_to_openai, get_ai_metrics_from_response, @@ -11,6 +12,7 @@ 'OpenAIRunnerFactory', 'OpenAIAgentGraphRunner', '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 ae4094c..65ccf55 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 @@ -47,6 +47,32 @@ def create_agent_graph(self, graph_def: Any, tools: ToolRegistry) -> Any: from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner return OpenAIAgentGraphRunner(graph_def, tools) + def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> '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 25295a7..5d5ee77 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_agent_graph import ManagedAgentGraph from ldai.managed_model import ManagedModel from ldai.models import ( # Deprecated aliases for backward compatibility @@ -56,6 +57,7 @@ 'AICompletionConfigDefault', 'AIJudgeConfig', 'AIJudgeConfigDefault', + 'ManagedAgent', 'ManagedModel', 'ManagedAgentGraph', 'EvalScore', diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 6aef7c6..aab8ccb 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_agent_graph import ManagedAgentGraph from ldai.managed_model import ManagedModel from ldai.models import ( @@ -33,6 +34,7 @@ _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' @@ -377,6 +379,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 0f7752d7303b3cc7ec3b65e698e29a8e3242ea8d Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Thu, 26 Mar 2026 08:27:42 -0500 Subject: [PATCH 2/9] 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 9523fa78599145281a1f328da75f06d0ffef6167 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Thu, 26 Mar 2026 14:47:36 -0500 Subject: [PATCH 3/9] 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 | 2 +- .../sdk/server-ai/src/ldai/managed_agent.py | 3 +-- .../sdk/server-ai/tests/test_managed_agent.py | 10 ++++----- 10 files changed, 48 insertions(+), 36 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 19022df..1c03199 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,8 +1,11 @@ -from typing import Any +from typing import TYPE_CHECKING, Any 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_model_runner import LangChainModelRunner @@ -49,5 +52,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 65ccf55..85cf523 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, ToolRegistry @@ -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 aab8ccb..d77974e 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -425,7 +425,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 b1712dc8f60e2e19a8950a28c77c81f2761fe0bf Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 11:20:10 -0500 Subject: [PATCH 4/9] simplifying tool configuration --- .../ldai_langchain/langchain_agent_runner.py | 29 +------- .../src/ldai_langchain/langchain_helper.py | 66 +++++++++++++++++-- .../langchain_runner_factory.py | 15 ++--- .../tests/test_langchain_provider.py | 11 +--- .../src/ldai_openai/openai_runner_factory.py | 1 + .../src/ldai/providers/ai_provider.py | 4 +- 6 files changed, 76 insertions(+), 50 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 1c03199..e119ac0 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,4 +1,4 @@ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Optional from ldai.models import AIConfigKind from ldai.providers import AIProvider, ToolRegistry @@ -36,7 +36,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. @@ -46,11 +46,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 85cf523..6e0ac12 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 @@ -50,6 +50,7 @@ def create_agent_graph(self, graph_def: Any, tools: ToolRegistry) -> Any: from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner return OpenAIAgentGraphRunner(graph_def, tools) + def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> 'OpenAIAgentRunner': """ Create a configured OpenAIAgentRunner for the given AI agent config. 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 afe344eef7d1cbc730d135d74abca2255748738c Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Fri, 27 Mar 2026 14:55:43 -0500 Subject: [PATCH 5/9] 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 e119ac0..a8445c0 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 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,7 +38,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. @@ -44,8 +46,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 == "" From d9ab2a8d7721b680609817c3325c7b6497a304f1 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Mon, 30 Mar 2026 13:42:03 -0500 Subject: [PATCH 6/9] address feedback --- .../src/ldai_langchain/langchain_helper.py | 70 +++--- .../ldai_openai/openai_agent_graph_runner.py | 58 +---- .../src/ldai_openai/openai_agent_runner.py | 225 ++++++++---------- .../src/ldai_openai/openai_helper.py | 57 ++++- .../src/ldai_openai/openai_runner_factory.py | 1 - .../tests/test_openai_provider.py | 217 ++++++++++------- 6 files changed, 328 insertions(+), 300 deletions(-) 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 4357cad..677d243 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 @@ -95,16 +95,16 @@ def create_langchain_model(ai_config: AIConfigKind, tool_registry: Optional[Tool return model -def _resolve_tools_for_langchain( +def _iter_valid_tools( tool_definitions: List[Dict[str, Any]], tool_registry: ToolRegistry, -) -> List[Dict[str, Any]]: +) -> List[tuple]: """ - 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. + Filter LD tool definitions against a registry, returning (name, td) pairs for each + valid function tool that has a callable implementation. Built-in provider tools and + tools missing from the registry are skipped with a warning. """ - bindable = [] + valid = [] for td in tool_definitions: if not isinstance(td, dict): continue @@ -112,9 +112,8 @@ def _resolve_tools_for_langchain( 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." + 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 @@ -126,16 +125,31 @@ def _resolve_tools_for_langchain( log.warning(f"Tool '{name}' is defined in the AI config but was not found in the tool registry; skipping.") continue - bindable.append({ + valid.append((name, td)) + + return valid + + +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. + """ + return [ + { 'type': 'function', 'function': { 'name': name, 'description': td.get('description', ''), 'parameters': td.get('parameters', {'type': 'object', 'properties': {}}), }, - }) - - return bindable + } + for name, td in _iter_valid_tools(tool_definitions, tool_registry) + ] def build_structured_tools(ai_config: AIConfigKind, tool_registry: ToolRegistry) -> List[Any]: @@ -157,34 +171,14 @@ def build_structured_tools(ai_config: AIConfigKind, tool_registry: ToolRegistry) 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( + return [ + StructuredTool.from_function( func=tool_registry[name], name=name, description=td.get('description', ''), - )) - - return structured + ) + for name, td in _iter_valid_tools(tool_definitions, tool_registry) + ] def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py index df2acf6..a0840c2 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_graph_runner.py @@ -9,31 +9,10 @@ from ldai.providers.types import LDAIMetrics from ldai.tracker import TokenUsage - -def _to_openai_name(name: str) -> str: - """Convert a hyphenated tool/node name to an underscore-separated OpenAI function name.""" - return name.replace('-', '_') - - -def _build_native_tool_map() -> dict: - try: - from agents import ( - CodeInterpreterTool, - FileSearchTool, - ImageGenerationTool, - WebSearchTool, - ) - return { - 'web_search_tool': lambda _: WebSearchTool(), - 'file_search_tool': lambda _: FileSearchTool(), - 'code_interpreter': lambda _: CodeInterpreterTool(), - 'image_generation': lambda _: ImageGenerationTool(), - } - except ImportError: - return {} - - -_NATIVE_OPENAI_TOOLS = _build_native_tool_map() +from ldai_openai.openai_helper import ( + NATIVE_OPENAI_TOOLS, + get_ai_usage_from_response, +) class _RunState: @@ -97,17 +76,9 @@ async def run(self, input: Any) -> AgentGraphResult: tracker.track_path(path) tracker.track_latency(duration) tracker.track_invocation_success() - try: - usage = result.context_wrapper.usage - tracker.track_total_tokens( - TokenUsage( - total=usage.total_tokens, - input=usage.input_tokens, - output=usage.output_tokens, - ) - ) - except Exception: - pass + token_usage = get_ai_usage_from_response(result) + if token_usage is not None: + tracker.track_total_tokens(token_usage) return AgentGraphResult( output=str(result.final_output), @@ -281,21 +252,19 @@ def build_node(node: AgentGraphNode, ctx: dict) -> Any: # --- tools --- agent_tools: List[Tool] = [] for tool_def in tool_defs: - tool_name_raw = tool_def.get('name', '') - tool_name = _to_openai_name(tool_name_raw) + tool_name = tool_def.get('name', '') # Check native OpenAI tools first, then fall back to ToolRegistry - if tool_name in _NATIVE_OPENAI_TOOLS: - agent_tools.append(_NATIVE_OPENAI_TOOLS[tool_name](tool_def)) + if tool_name in NATIVE_OPENAI_TOOLS: + agent_tools.append(NATIVE_OPENAI_TOOLS[tool_name](tool_def)) continue - tool_fn = self._tools.get(tool_name) or self._tools.get(tool_name_raw) + tool_fn = self._tools.get(tool_name) if not tool_fn: continue def _make_tool( name: str, - raw_name: str, fn: Any, description: str, params_schema: dict, @@ -306,7 +275,7 @@ def wrapped(tool_ctx: ToolContext, tool_args: str) -> Any: args = json.loads(tool_args) except Exception: args = {} - path.append(raw_name) + path.append(name) if config_tracker is not None: config_tracker.track_tool_call( name, @@ -324,7 +293,6 @@ def wrapped(tool_ctx: ToolContext, tool_args: str) -> Any: agent_tools.append( _make_tool( tool_name, - tool_name_raw, tool_fn, tool_def.get('description', ''), tool_def.get('parameters', {}), @@ -332,7 +300,7 @@ def wrapped(tool_ctx: ToolContext, tool_args: str) -> Any: ) return Agent( - name=_to_openai_name(node_config.key), + name=node_config.key, instructions=f'{RECOMMENDED_PROMPT_PREFIX} {node_config.instructions or ""}', handoffs=list(agent_handoffs), tools=list(agent_tools), 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 index 1a588d8..ac15f88 100644 --- 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 @@ -6,18 +6,23 @@ 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 +from ldai_openai.openai_helper import ( + NATIVE_OPENAI_TOOLS, + get_ai_usage_from_response, +) class OpenAIAgentRunner(AgentRunner): """ AgentRunner implementation for OpenAI. - Executes a single-agent loop using OpenAI Chat Completions with tool calling. + Executes a single agent using the OpenAI Agents SDK (``openai-agents``). + Tool calling and the agentic loop are handled internally by ``Runner.run``. Returned by OpenAIRunnerFactory.create_agent(config, tools). + + Requires ``openai-agents`` to be installed. """ def __init__( @@ -40,143 +45,117 @@ 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 OpenAI Agents SDK ``Runner.run``, 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[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: + from agents import Agent, Runner + except ImportError: + log.warning( + "openai-agents is required for OpenAIAgentRunner. " + "Install it with: pip install openai-agents" + ) + return AgentResult(output="", raw=None, metrics=LDAIMetrics(success=False, usage=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, - }) + agent_tools = self._build_agent_tools() + model_settings = self._build_model_settings() + + agent = Agent( + name="ldai-agent", + instructions=self._instructions or None, + model=self._model_name, + tools=agent_tools, + model_settings=model_settings, + ) + + result = await Runner.run(agent, str(input), max_turns=25) - except Exception as error: - log.warning(f"OpenAI agent run failed: {error}") return AgentResult( - output="", - raw=raw_response, - metrics=LDAIMetrics(success=False, usage=None), + output=str(result.final_output), + raw=result, + metrics=LDAIMetrics( + success=True, + usage=get_ai_usage_from_response(result), + ), ) - - 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}" + log.warning(f"OpenAI agent run failed: {error}") + return AgentResult(output="", raw=None, metrics=LDAIMetrics(success=False, usage=None)) + + def _build_agent_tools(self) -> List[Any]: + """Build tool instances from LD tool definitions and registry.""" + from agents import FunctionTool + from agents.tool_context import ToolContext - 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": {}}), - }, - }) + name = td.get("name", "") + + # Native OpenAI tools run on OpenAI's infrastructure — no local fn required. + if name and name in NATIVE_OPENAI_TOOLS: + tools.append(NATIVE_OPENAI_TOOLS[name](td)) + continue + + tool_type = td.get("type") + if tool_type and tool_type != "function": + log.warning( + f"Built-in tool '{tool_type}' is not supported and will be skipped. " + "Use the OpenAIAgentGraphRunner for built-in provider tools." + ) + continue + + if not name: + continue + + tool_fn = self._tools.get(name) + if not tool_fn: + log.warning( + f"Tool '{name}' is defined in the AI config but was not found in " + "the tool registry; skipping." + ) + continue + + def _make_invoker(fn: Any) -> Any: + async def on_invoke_tool(tool_ctx: ToolContext, args_json: str) -> str: + try: + args = json.loads(args_json) if args_json else {} + except Exception: + args = {} + try: + res = fn(**args) + if hasattr(res, "__await__"): + res = await res + return str(res) + except Exception as e: + log.warning(f"Tool '{name}' execution failed: {e}") + return f"Tool execution failed: {e}" + return on_invoke_tool + + tools.append(FunctionTool( + name=name, + description=td.get("description", ""), + params_json_schema=td.get("parameters", {}), + on_invoke_tool=_make_invoker(tool_fn), + )) return tools + def _build_model_settings(self) -> Any: + """Map LD model parameters to an openai-agents ModelSettings instance.""" + from agents import ModelSettings + + known = { + "temperature", "top_p", "max_tokens", + "frequency_penalty", "presence_penalty", + } + kwargs = {k: v for k, v in self._parameters.items() if k in known} + return ModelSettings(**kwargs) if kwargs else None + 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_helper.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_helper.py index dc40ec6..2fc5412 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 @@ -1,4 +1,4 @@ -from typing import Any, Iterable, List, Optional, cast +from typing import Any, Dict, Iterable, List, Optional, cast from ldai import LDMessage from ldai.providers.types import LDAIMetrics @@ -6,6 +6,27 @@ from openai.types.chat import ChatCompletionMessageParam +def _build_native_tool_map() -> Dict[str, Any]: + try: + from agents import ( + CodeInterpreterTool, + FileSearchTool, + ImageGenerationTool, + WebSearchTool, + ) + return { + 'web_search_tool': lambda _: WebSearchTool(), + 'file_search_tool': lambda _: FileSearchTool(), + 'code_interpreter': lambda _: CodeInterpreterTool(), + 'image_generation': lambda _: ImageGenerationTool(), + } + except ImportError: + return {} + + +NATIVE_OPENAI_TOOLS: Dict[str, Any] = _build_native_tool_map() + + def convert_messages_to_openai(messages: List[LDMessage]) -> Iterable[ChatCompletionMessageParam]: """ Convert LaunchDarkly messages to OpenAI chat completion message format. @@ -23,16 +44,32 @@ def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: """ Extract token usage from an OpenAI response. - :param response: The response from the OpenAI chat completions API + Handles both chat completions responses (``response.usage``) and + openai-agents ``RunResult`` objects (``response.context_wrapper.usage``). + Field names are normalised: agents SDK uses ``input_tokens``/``output_tokens`` + while the chat completions API uses ``prompt_tokens``/``completion_tokens``. + + :param response: An OpenAI chat completions response or openai-agents RunResult :return: TokenUsage or None if unavailable """ - if hasattr(response, 'usage') and response.usage: - u = response.usage - 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) + usage = None + try: + usage = response.context_wrapper.usage + except Exception: + pass + + if usage is None and hasattr(response, 'usage') and response.usage: + usage = response.usage + + if usage is None: + return None + + total = getattr(usage, 'total_tokens', None) or 0 + inp = getattr(usage, 'input_tokens', None) or getattr(usage, 'prompt_tokens', None) or 0 + out = getattr(usage, 'output_tokens', None) or getattr(usage, 'completion_tokens', None) or 0 + + if total or inp or out: + return TokenUsage(total=total, input=inp, output=out) return None @@ -40,7 +77,7 @@ def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: """ Extract LaunchDarkly AI metrics from an OpenAI response. - :param response: The response from the OpenAI chat completions API + :param response: An OpenAI chat completions response or openai-agents RunResult :return: LDAIMetrics with success status and token usage """ return LDAIMetrics(success=True, usage=get_ai_usage_from_response(response)) 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 6e0ac12..85cf523 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 @@ -50,7 +50,6 @@ def create_agent_graph(self, graph_def: Any, tools: ToolRegistry) -> Any: from ldai_openai.openai_agent_graph_runner import OpenAIAgentGraphRunner return OpenAIAgentGraphRunner(graph_def, tools) - def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> 'OpenAIAgentRunner': """ Create a configured OpenAIAgentRunner for the given AI agent config. 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 afdb524..3f3a7e3 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 @@ -1,6 +1,7 @@ """Tests for OpenAI Provider.""" import pytest +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch from ldai import LDMessage @@ -8,33 +9,68 @@ from ldai_openai import OpenAIModelRunner, OpenAIRunnerFactory, get_ai_metrics_from_response, get_ai_usage_from_response +def _make_usage(attrs: dict): + """Build a simple namespace with only the given attributes (no MagicMock auto-attrs).""" + class _Usage: + pass + u = _Usage() + for k, v in attrs.items(): + setattr(u, k, v) + return u + + +def _make_completions_response(total=100, prompt=50, completion=50): + """Build a mock chat completions response with .usage (chat completions field names).""" + mock = MagicMock() + mock.context_wrapper = None + mock.usage = _make_usage({ + 'total_tokens': total, + 'prompt_tokens': prompt, + 'completion_tokens': completion, + }) + return mock + + +def _make_runner_result(total=100, input_tokens=50, output_tokens=50): + """Build a mock openai-agents RunResult with .context_wrapper.usage (agents field names).""" + mock_ctx = MagicMock() + mock_ctx.usage = _make_usage({ + 'total_tokens': total, + 'input_tokens': input_tokens, + 'output_tokens': output_tokens, + }) + mock = MagicMock() + mock.context_wrapper = mock_ctx + mock.usage = None + return mock + + 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 = get_ai_usage_from_response(mock_response) + def test_returns_usage_from_chat_completions_response(self): + u = get_ai_usage_from_response(_make_completions_response(total=100, prompt=50, completion=50)) assert u is not None assert u.total == 100 assert u.input == 50 assert u.output == 50 + def test_returns_usage_from_runner_result(self): + u = get_ai_usage_from_response(_make_runner_result(total=43, input_tokens=30, output_tokens=13)) + assert u is not None + assert u.total == 43 + assert u.input == 30 + assert u.output == 13 + def test_returns_none_when_usage_missing(self): mock_response = MagicMock() mock_response.usage = None + mock_response.context_wrapper = None assert 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 get_ai_usage_from_response(mock_response) is None + u = get_ai_usage_from_response(_make_completions_response(total=0, prompt=0, completion=0)) + assert u is None class TestGetAIMetricsFromResponse: @@ -42,14 +78,7 @@ class TestGetAIMetricsFromResponse: def test_creates_metrics_with_success_true_and_token_usage(self): """Should create metrics with success=True and token usage.""" - 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 - - result = get_ai_metrics_from_response(mock_response) - + result = get_ai_metrics_from_response(_make_completions_response(total=100, prompt=50, completion=50)) assert result.success is True assert result.usage is not None assert result.usage.total == 100 @@ -60,6 +89,7 @@ def test_creates_metrics_with_success_true_and_no_usage_when_usage_missing(self) """Should create metrics with success=True and no usage when usage is missing.""" mock_response = MagicMock() mock_response.usage = None + mock_response.context_wrapper = None result = get_ai_metrics_from_response(mock_response) @@ -69,10 +99,8 @@ def test_creates_metrics_with_success_true_and_no_usage_when_usage_missing(self) def test_handles_partial_usage_data(self): """Should handle partial usage data.""" mock_response = MagicMock() - mock_response.usage = MagicMock() - mock_response.usage.prompt_tokens = 30 - mock_response.usage.completion_tokens = None - mock_response.usage.total_tokens = None + mock_response.context_wrapper = None + mock_response.usage = _make_usage({'prompt_tokens': 30, 'completion_tokens': None, 'total_tokens': None}) result = get_ai_metrics_from_response(mock_response) @@ -95,13 +123,11 @@ def mock_client(self): async def test_invokes_openai_chat_completions_and_returns_response(self, mock_client): """Should invoke OpenAI chat completions and return response.""" mock_response = MagicMock() + mock_response.context_wrapper = None mock_response.choices = [MagicMock()] mock_response.choices[0].message = MagicMock() mock_response.choices[0].message.content = 'Hello! How can I help you today?' - mock_response.usage = MagicMock() - mock_response.usage.prompt_tokens = 10 - mock_response.usage.completion_tokens = 15 - mock_response.usage.total_tokens = 25 + mock_response.usage = _make_usage({'total_tokens': 25, 'prompt_tokens': 10, 'completion_tokens': 15}) mock_client.chat = MagicMock() mock_client.chat.completions = MagicMock() @@ -192,13 +218,11 @@ def mock_client(self): async def test_invokes_openai_with_structured_output(self, mock_client): """Should invoke OpenAI with structured output and return parsed response.""" mock_response = MagicMock() + mock_response.context_wrapper = None mock_response.choices = [MagicMock()] mock_response.choices[0].message = MagicMock() mock_response.choices[0].message.content = '{"name": "John", "age": 30, "city": "New York"}' - mock_response.usage = MagicMock() - mock_response.usage.prompt_tokens = 20 - mock_response.usage.completion_tokens = 10 - mock_response.usage.total_tokens = 30 + mock_response.usage = _make_usage({'total_tokens': 30, 'prompt_tokens': 20, 'completion_tokens': 10}) mock_client.chat = MagicMock() mock_client.chat.completions = MagicMock() @@ -253,13 +277,11 @@ async def test_returns_unsuccessful_when_no_content_in_structured_response(self, async def test_handles_json_parsing_errors(self, mock_client): """Should handle JSON parsing errors gracefully.""" mock_response = MagicMock() + mock_response.context_wrapper = None mock_response.choices = [MagicMock()] mock_response.choices[0].message = MagicMock() mock_response.choices[0].message.content = 'invalid json content' - mock_response.usage = MagicMock() - mock_response.usage.prompt_tokens = 10 - mock_response.usage.completion_tokens = 5 - mock_response.usage.total_tokens = 15 + mock_response.usage = _make_usage({'total_tokens': 15, 'prompt_tokens': 10, 'completion_tokens': 5}) mock_client.chat = MagicMock() mock_client.chat.completions = MagicMock() @@ -397,6 +419,29 @@ def test_creates_agent_runner_with_no_tools(self): assert result._tool_definitions == [] +def _make_agents_mock(runner_run_mock: Any) -> MagicMock: + """Build a mock ``agents`` module with Agent, Runner, FunctionTool, ModelSettings.""" + mock_runner_cls = MagicMock() + mock_runner_cls.run = runner_run_mock + + mock_agent_cls = MagicMock() + + mock_function_tool_cls = MagicMock(side_effect=lambda **kw: MagicMock(**kw)) + + mock_model_settings_cls = MagicMock(side_effect=lambda **kw: MagicMock(**kw)) + + mock_tool_context_module = MagicMock() + mock_tool_context_module.ToolContext = MagicMock() + + agents_mock = MagicMock() + agents_mock.Agent = mock_agent_cls + agents_mock.Runner = mock_runner_cls + agents_mock.FunctionTool = mock_function_tool_cls + agents_mock.ModelSettings = mock_model_settings_cls + + return agents_mock, mock_tool_context_module + + class TestOpenAIAgentRunner: """Tests for OpenAIAgentRunner.run.""" @@ -404,60 +449,49 @@ class TestOpenAIAgentRunner: def mock_client(self): return MagicMock() + def _make_run_result(self, output: str, total: int = 15, input_tokens: int = 10, output_tokens: int = 5): + """Build a mock RunResult with final_output and context_wrapper.usage.""" + mock_usage = MagicMock() + mock_usage.total_tokens = total + mock_usage.input_tokens = input_tokens + mock_usage.output_tokens = output_tokens + + mock_ctx = MagicMock() + mock_ctx.usage = mock_usage + + mock_result = MagicMock() + mock_result.final_output = output + mock_result.context_wrapper = mock_ctx + return mock_result + @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.""" + """Should return AgentResult when Runner.run returns a final output.""" + import sys + 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) + mock_run_result = self._make_run_result("The answer is 42.", total=15, input_tokens=10, output_tokens=5) + agents_mock, tc_mock = _make_agents_mock(AsyncMock(return_value=mock_run_result)) runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, 'You are helpful.', [], {}) - result = await runner.run("What is the answer?") + with patch.dict(sys.modules, {'agents': agents_mock, 'agents.tool_context': tc_mock}): + result = await runner.run("What is the answer?") assert result.output == "The answer is 42." assert result.metrics.success is True + assert result.metrics.usage is not None + assert result.metrics.usage.total == 15 @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.""" + """Should delegate tool-calling loop to Runner.run and return final output.""" + import sys + 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] - ) + mock_run_result = self._make_run_result("It is sunny in Paris.", total=43, input_tokens=30, output_tokens=13) + agents_mock, tc_mock = _make_agents_mock(AsyncMock(return_value=mock_run_result)) weather_fn = MagicMock(return_value="Sunny, 25°C") runner = OpenAIAgentRunner( @@ -465,22 +499,39 @@ async def test_executes_tool_calls_and_returns_final_response(self, mock_client) [{'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}], {'get-weather': weather_fn}, ) - result = await runner.run("What is the weather in Paris?") + with patch.dict(sys.modules, {'agents': agents_mock, 'agents.tool_context': tc_mock}): + 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 + assert result.metrics.usage.total == 43 @pytest.mark.asyncio async def test_returns_failure_when_exception_thrown(self, mock_client): - """Should return unsuccessful AgentResult when exception is thrown.""" + """Should return unsuccessful AgentResult when Runner.run raises.""" + import sys + from ldai_openai import OpenAIAgentRunner - mock_client.chat.completions.create = AsyncMock(side_effect=Exception("API Error")) + agents_mock, tc_mock = _make_agents_mock(AsyncMock(side_effect=Exception("API Error"))) + + runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, '', [], {}) + with patch.dict(sys.modules, {'agents': agents_mock, 'agents.tool_context': tc_mock}): + result = await runner.run("Hello") + + assert result.output == "" + assert result.metrics.success is False + + @pytest.mark.asyncio + async def test_returns_failure_when_openai_agents_not_installed(self, mock_client): + """Should return unsuccessful AgentResult when openai-agents is not installed.""" + import sys + + from ldai_openai import OpenAIAgentRunner runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, '', [], {}) - result = await runner.run("Hello") + with patch.dict(sys.modules, {'agents': None}): + result = await runner.run("Hello") assert result.output == "" assert result.metrics.success is False From aca83ea7df655ee91d2339c00305ef8466219e45 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Mon, 30 Mar 2026 14:10:20 -0500 Subject: [PATCH 7/9] address PR feedback --- .../src/ldai_langchain/langchain_runner_factory.py | 2 +- .../server-ai-langchain/tests/test_langchain_provider.py | 4 ++-- .../server-ai-openai/src/ldai_openai/openai_agent_runner.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) 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 a8445c0..29ab68b 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,5 @@ 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 @@ -46,6 +45,7 @@ def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> Lan :param tools: ToolRegistry mapping tool names to callables :return: LangChainAgentRunner ready to run the agent """ + from langchain.agents import create_agent as lc_create_agent instructions = (config.instructions or '') if hasattr(config, 'instructions') else '' llm = create_langchain_model(config) lc_tools = build_structured_tools(config, 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 3aa04a9..52cca65 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 @@ -423,7 +423,7 @@ def test_creates_agent_runner_with_instructions_and_tool_definitions(self): 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): + patch('langchain.agents.create_agent', return_value=mock_agent): mock_create.return_value = MagicMock() mock_tools.return_value = [MagicMock()] @@ -448,7 +448,7 @@ def test_creates_agent_runner_with_no_tools(self): 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): + patch('langchain.agents.create_agent', return_value=mock_agent): mock_create.return_value = MagicMock() factory = LangChainRunnerFactory() 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 index ac15f88..962221b 100644 --- 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 @@ -121,7 +121,7 @@ def _build_agent_tools(self) -> List[Any]: ) continue - def _make_invoker(fn: Any) -> Any: + def _make_invoker(fn: Any, tool_name: str) -> Any: async def on_invoke_tool(tool_ctx: ToolContext, args_json: str) -> str: try: args = json.loads(args_json) if args_json else {} @@ -133,7 +133,7 @@ async def on_invoke_tool(tool_ctx: ToolContext, args_json: str) -> str: res = await res return str(res) except Exception as e: - log.warning(f"Tool '{name}' execution failed: {e}") + log.warning(f"Tool '{tool_name}' execution failed: {e}") return f"Tool execution failed: {e}" return on_invoke_tool @@ -141,7 +141,7 @@ async def on_invoke_tool(tool_ctx: ToolContext, args_json: str) -> str: name=name, description=td.get("description", ""), params_json_schema=td.get("parameters", {}), - on_invoke_tool=_make_invoker(tool_fn), + on_invoke_tool=_make_invoker(tool_fn, name), )) return tools From 268b6696ff11558b7a207011cbdb3f7648629cd3 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Mon, 30 Mar 2026 17:11:38 -0500 Subject: [PATCH 8/9] remove unused client --- .../src/ldai_openai/openai_agent_runner.py | 7 ---- .../src/ldai_openai/openai_runner_factory.py | 1 - .../tests/test_openai_provider.py | 32 +++++++------------ 3 files changed, 11 insertions(+), 29 deletions(-) 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 index 962221b..199bef4 100644 --- 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 @@ -6,7 +6,6 @@ from ldai import log from ldai.providers import AgentResult, AgentRunner, ToolRegistry from ldai.providers.types import LDAIMetrics -from openai import AsyncOpenAI from ldai_openai.openai_helper import ( NATIVE_OPENAI_TOOLS, @@ -27,14 +26,12 @@ class OpenAIAgentRunner(AgentRunner): 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 @@ -155,7 +152,3 @@ def _build_model_settings(self) -> Any: } kwargs = {k: v for k, v in self._parameters.items() if k in known} return ModelSettings(**kwargs) if kwargs else None - - 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 85cf523..644d12f 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 @@ -68,7 +68,6 @@ def create_agent(self, config: Any, tools: Optional[ToolRegistry] = None) -> 'Op instructions = (config.instructions or '') if hasattr(config, 'instructions') else '' return OpenAIAgentRunner( - self._client, model_name, parameters, instructions, 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 3f3a7e3..0d47e28 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 @@ -424,20 +424,14 @@ def _make_agents_mock(runner_run_mock: Any) -> MagicMock: mock_runner_cls = MagicMock() mock_runner_cls.run = runner_run_mock - mock_agent_cls = MagicMock() - - mock_function_tool_cls = MagicMock(side_effect=lambda **kw: MagicMock(**kw)) - - mock_model_settings_cls = MagicMock(side_effect=lambda **kw: MagicMock(**kw)) - mock_tool_context_module = MagicMock() mock_tool_context_module.ToolContext = MagicMock() agents_mock = MagicMock() - agents_mock.Agent = mock_agent_cls + agents_mock.Agent = MagicMock() agents_mock.Runner = mock_runner_cls - agents_mock.FunctionTool = mock_function_tool_cls - agents_mock.ModelSettings = mock_model_settings_cls + agents_mock.FunctionTool = MagicMock(side_effect=lambda **kw: MagicMock(**kw)) + agents_mock.ModelSettings = MagicMock(side_effect=lambda **kw: MagicMock(**kw)) return agents_mock, mock_tool_context_module @@ -445,10 +439,6 @@ def _make_agents_mock(runner_run_mock: Any) -> MagicMock: class TestOpenAIAgentRunner: """Tests for OpenAIAgentRunner.run.""" - @pytest.fixture - def mock_client(self): - return MagicMock() - def _make_run_result(self, output: str, total: int = 15, input_tokens: int = 10, output_tokens: int = 5): """Build a mock RunResult with final_output and context_wrapper.usage.""" mock_usage = MagicMock() @@ -465,7 +455,7 @@ def _make_run_result(self, output: str, total: int = 15, input_tokens: int = 10, return mock_result @pytest.mark.asyncio - async def test_runs_agent_and_returns_result_with_no_tool_calls(self, mock_client): + async def test_runs_agent_and_returns_result_with_no_tool_calls(self): """Should return AgentResult when Runner.run returns a final output.""" import sys @@ -474,7 +464,7 @@ async def test_runs_agent_and_returns_result_with_no_tool_calls(self, mock_clien mock_run_result = self._make_run_result("The answer is 42.", total=15, input_tokens=10, output_tokens=5) agents_mock, tc_mock = _make_agents_mock(AsyncMock(return_value=mock_run_result)) - runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, 'You are helpful.', [], {}) + runner = OpenAIAgentRunner('gpt-4', {}, 'You are helpful.', [], {}) with patch.dict(sys.modules, {'agents': agents_mock, 'agents.tool_context': tc_mock}): result = await runner.run("What is the answer?") @@ -484,7 +474,7 @@ async def test_runs_agent_and_returns_result_with_no_tool_calls(self, mock_clien assert result.metrics.usage.total == 15 @pytest.mark.asyncio - async def test_executes_tool_calls_and_returns_final_response(self, mock_client): + async def test_executes_tool_calls_and_returns_final_response(self): """Should delegate tool-calling loop to Runner.run and return final output.""" import sys @@ -495,7 +485,7 @@ async def test_executes_tool_calls_and_returns_final_response(self, mock_client) weather_fn = MagicMock(return_value="Sunny, 25°C") runner = OpenAIAgentRunner( - mock_client, 'gpt-4', {}, 'You are helpful.', + 'gpt-4', {}, 'You are helpful.', [{'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}], {'get-weather': weather_fn}, ) @@ -507,7 +497,7 @@ async def test_executes_tool_calls_and_returns_final_response(self, mock_client) assert result.metrics.usage.total == 43 @pytest.mark.asyncio - async def test_returns_failure_when_exception_thrown(self, mock_client): + async def test_returns_failure_when_exception_thrown(self): """Should return unsuccessful AgentResult when Runner.run raises.""" import sys @@ -515,7 +505,7 @@ async def test_returns_failure_when_exception_thrown(self, mock_client): agents_mock, tc_mock = _make_agents_mock(AsyncMock(side_effect=Exception("API Error"))) - runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, '', [], {}) + runner = OpenAIAgentRunner('gpt-4', {}, '', [], {}) with patch.dict(sys.modules, {'agents': agents_mock, 'agents.tool_context': tc_mock}): result = await runner.run("Hello") @@ -523,13 +513,13 @@ async def test_returns_failure_when_exception_thrown(self, mock_client): assert result.metrics.success is False @pytest.mark.asyncio - async def test_returns_failure_when_openai_agents_not_installed(self, mock_client): + async def test_returns_failure_when_openai_agents_not_installed(self): """Should return unsuccessful AgentResult when openai-agents is not installed.""" import sys from ldai_openai import OpenAIAgentRunner - runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, '', [], {}) + runner = OpenAIAgentRunner('gpt-4', {}, '', [], {}) with patch.dict(sys.modules, {'agents': None}): result = await runner.run("Hello") From 9295cc3a01da62959d1686a5f5f44eedb2627108 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Tue, 31 Mar 2026 09:20:50 -0500 Subject: [PATCH 9/9] fix token tracking --- .../src/ldai_openai/openai_helper.py | 27 +++++++++---------- .../tests/test_openai_provider.py | 9 +++++++ 2 files changed, 22 insertions(+), 14 deletions(-) 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 2fc5412..6088e89 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 @@ -46,30 +46,29 @@ def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: Handles both chat completions responses (``response.usage``) and openai-agents ``RunResult`` objects (``response.context_wrapper.usage``). - Field names are normalised: agents SDK uses ``input_tokens``/``output_tokens`` - while the chat completions API uses ``prompt_tokens``/``completion_tokens``. :param response: An OpenAI chat completions response or openai-agents RunResult :return: TokenUsage or None if unavailable """ - usage = None try: usage = response.context_wrapper.usage + if usage is not None: + total = getattr(usage, 'total_tokens', None) or 0 + inp = getattr(usage, 'input_tokens', None) or 0 + out = getattr(usage, 'output_tokens', None) or 0 + if total or inp or out: + return TokenUsage(total=total, input=inp, output=out) except Exception: pass - if usage is None and hasattr(response, 'usage') and response.usage: - usage = response.usage + usage = getattr(response, 'usage', None) + if usage is not None: + total = getattr(usage, 'total_tokens', None) or 0 + inp = getattr(usage, 'prompt_tokens', None) or 0 + out = getattr(usage, 'completion_tokens', None) or 0 + if total or inp or out: + return TokenUsage(total=total, input=inp, output=out) - if usage is None: - return None - - total = getattr(usage, 'total_tokens', None) or 0 - inp = getattr(usage, 'input_tokens', None) or getattr(usage, 'prompt_tokens', None) or 0 - out = getattr(usage, 'output_tokens', None) or getattr(usage, '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/tests/test_openai_provider.py b/packages/ai-providers/server-ai-openai/tests/test_openai_provider.py index 0d47e28..19d2cff 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 @@ -72,6 +72,15 @@ def test_returns_none_when_all_counts_zero(self): u = get_ai_usage_from_response(_make_completions_response(total=0, prompt=0, completion=0)) assert u is None + def test_zero_input_tokens_not_conflated_with_missing(self): + """input_tokens=0 should be used as-is, not fall through to prompt_tokens.""" + u = get_ai_usage_from_response( + _make_runner_result(total=10, input_tokens=0, output_tokens=10) + ) + assert u is not None + assert u.input == 0 + assert u.output == 10 + class TestGetAIMetricsFromResponse: """Tests for get_ai_metrics_from_response."""