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..b0a1c85 --- /dev/null +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py @@ -0,0 +1,63 @@ +"""LangChain agent runner for LaunchDarkly AI SDK.""" + +from typing import Any + +from ldai import log +from ldai.providers import AgentResult, AgentRunner +from ldai.providers.types import LDAIMetrics + +from ldai_langchain.langchain_helper import sum_token_usage_from_messages + + +class LangChainAgentRunner(AgentRunner): + """ + AgentRunner implementation for LangChain. + + 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, agent: Any): + self._agent = agent + + async def run(self, input: Any) -> AgentResult: + """ + Run the agent with the given input string. + + 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 + """ + try: + 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=None, + metrics=LDAIMetrics(success=False, usage=None), + ) + + 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 e160061..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 @@ -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,105 @@ 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 and tool_registry is not None: + bindable = _resolve_tools_for_langchain(tool_definitions, tool_registry) + if bindable: + model = model.bind_tools(bindable) + + return model + + +def _iter_valid_tools( + tool_definitions: List[Dict[str, Any]], + tool_registry: ToolRegistry, +) -> List[tuple]: + """ + 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. + """ + valid = [] + 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 + + 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': {}}), + }, + } + for name, td in _iter_valid_tools(tool_definitions, tool_registry) + ] + + +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 [] + + return [ + StructuredTool.from_function( + func=tool_registry[name], + name=name, + description=td.get('description', ''), + ) + for name, td in _iter_valid_tools(tool_definitions, tool_registry) + ] + def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: """ @@ -88,11 +189,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 43febb5..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,9 +1,13 @@ -from typing import Any +from typing import Any, Optional from ldai.models import AIConfigKind from ldai.providers import AIProvider, ToolRegistry -from ldai_langchain.langchain_helper import create_langchain_model +from ldai_langchain.langchain_agent_runner import LangChainAgentRunner +from ldai_langchain.langchain_helper import ( + build_structured_tools, + create_langchain_model, +) from ldai_langchain.langchain_model_runner import LangChainModelRunner @@ -32,3 +36,23 @@ 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: + """ + 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 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 {}) + + 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 9ce4e88..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 @@ -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, @@ -127,6 +128,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 = 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 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 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 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 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 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 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 get_tool_calls_from_response(msg) == ['a', 'b'] + class TestMapProvider: """Tests for map_provider.""" @@ -330,3 +396,122 @@ 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 wrapping a compiled graph.""" + 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'}, + } + + 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('langchain.agents.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._agent is mock_agent + + 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'}, + } + + 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('langchain.agents.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._agent is mock_agent + + +class TestLangChainAgentRunner: + """Tests for LangChainAgentRunner.run.""" + + @pytest.mark.asyncio + 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 + + final_msg = AIMessage(content="The answer is 42.") + mock_agent = MagicMock() + mock_agent.ainvoke = AsyncMock(return_value={"messages": [final_msg]}) + + 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_aggregates_token_usage_across_messages(self): + """Should sum token usage from all messages in the graph result.""" + from ldai_langchain import LangChainAgentRunner + + 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_agent = MagicMock() + mock_agent.ainvoke = AsyncMock(return_value={"messages": [msg1, msg2]}) + + runner = LangChainAgentRunner(mock_agent) + result = await runner.run("Hello") + + assert result.output == "final answer" + assert result.metrics.success is True + 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_agent = MagicMock() + mock_agent.ainvoke = AsyncMock(side_effect=Exception("Graph Error")) + + runner = LangChainAgentRunner(mock_agent) + 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_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 new file mode 100644 index 0000000..199bef4 --- /dev/null +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_runner.py @@ -0,0 +1,154 @@ +"""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_openai.openai_helper import ( + NATIVE_OPENAI_TOOLS, + get_ai_usage_from_response, +) + + +class OpenAIAgentRunner(AgentRunner): + """ + AgentRunner implementation for OpenAI. + + 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__( + self, + model_name: str, + parameters: Dict[str, Any], + instructions: str, + tool_definitions: List[Dict[str, Any]], + tools: ToolRegistry, + ): + 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. + + 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 + """ + 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: + 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) + + return AgentResult( + output=str(result.final_output), + raw=result, + metrics=LDAIMetrics( + success=True, + usage=get_ai_usage_from_response(result), + ), + ) + except Exception as 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 + + tools = [] + for td in self._tool_definitions: + if not isinstance(td, dict): + continue + 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, 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 {} + 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 '{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, name), + )) + 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 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..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 @@ -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,31 @@ 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``). + + :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 - 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, - ) + 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 + + 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) + return None @@ -40,7 +76,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 ae4094c..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 @@ -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.""" @@ -47,6 +50,31 @@ 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( + 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..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 @@ -1,11 +1,85 @@ """Tests for OpenAI Provider.""" import pytest +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch 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 + + +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_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): + 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: @@ -13,14 +87,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 @@ -31,6 +98,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) @@ -40,10 +108,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) @@ -66,13 +132,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() @@ -163,13 +227,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() @@ -224,13 +286,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() @@ -318,3 +378,159 @@ 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 == [] + + +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_tool_context_module = MagicMock() + mock_tool_context_module.ToolContext = MagicMock() + + agents_mock = MagicMock() + agents_mock.Agent = MagicMock() + agents_mock.Runner = mock_runner_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 + + +class TestOpenAIAgentRunner: + """Tests for OpenAIAgentRunner.run.""" + + 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): + """Should return AgentResult when Runner.run returns a final output.""" + import sys + + from ldai_openai import OpenAIAgentRunner + + 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('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?") + + 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): + """Should delegate tool-calling loop to Runner.run and return final output.""" + import sys + + from ldai_openai import OpenAIAgentRunner + + 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( + 'gpt-4', {}, 'You are helpful.', + [{'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}], + {'get-weather': weather_fn}, + ) + 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 + assert result.metrics.usage.total == 43 + + @pytest.mark.asyncio + async def test_returns_failure_when_exception_thrown(self): + """Should return unsuccessful AgentResult when Runner.run raises.""" + import sys + + from ldai_openai import OpenAIAgentRunner + + agents_mock, tc_mock = _make_agents_mock(AsyncMock(side_effect=Exception("API Error"))) + + runner = OpenAIAgentRunner('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): + """Should return unsuccessful AgentResult when openai-agents is not installed.""" + import sys + + from ldai_openai import OpenAIAgentRunner + + runner = OpenAIAgentRunner('gpt-4', {}, '', [], {}) + with patch.dict(sys.modules, {'agents': None}): + 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..d77974e 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 = 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..12c4d9b --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/managed_agent.py @@ -0,0 +1,52 @@ +"""ManagedAgent — LaunchDarkly managed wrapper for agent invocations.""" + +from ldai.models import AIAgentConfig +from ldai.providers import AgentResult, AgentRunner +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_async( + 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/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. 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..60cf7db --- /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.providers 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_async = 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_async.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 = MagicMock(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 = MagicMock(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