diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py index cb455e5..b2f6233 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py @@ -1,3 +1,4 @@ +from ldai_langchain.langchain_agent_runner import LangChainAgentRunner from ldai_langchain.langchain_helper import ( convert_messages_to_langchain, create_langchain_model, @@ -16,6 +17,7 @@ '__version__', 'LangChainRunnerFactory', 'LangChainModelRunner', + 'LangChainAgentRunner', 'convert_messages_to_langchain', 'create_langchain_model', 'get_ai_metrics_from_response', diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_agent_runner.py new file mode 100644 index 0000000..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..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 @@ -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,111 @@ 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 _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 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]: """ @@ -88,11 +195,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 402e295..ba5087f 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py @@ -1,7 +1,14 @@ +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 +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 @@ -17,3 +24,22 @@ 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 + """ + 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..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 @@ -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('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._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('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._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 8a8199b..a64f0b3 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/__init__.py @@ -1,3 +1,4 @@ +from ldai_openai.openai_agent_runner import OpenAIAgentRunner from ldai_openai.openai_helper import ( convert_messages_to_openai, get_ai_metrics_from_response, @@ -9,6 +10,7 @@ __all__ = [ 'OpenAIRunnerFactory', 'OpenAIModelRunner', + 'OpenAIAgentRunner', 'convert_messages_to_openai', 'get_ai_metrics_from_response', 'get_ai_usage_from_response', diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_runner.py new file mode 100644 index 0000000..1a588d8 --- /dev/null +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_agent_runner.py @@ -0,0 +1,182 @@ +"""OpenAI agent runner for LaunchDarkly AI SDK.""" + +import json +from typing import Any, Dict, List + +from ldai import log +from ldai.providers import AgentResult, AgentRunner, ToolRegistry +from ldai.providers.types import LDAIMetrics +from ldai.tracker import TokenUsage +from openai import AsyncOpenAI + +from ldai_openai.openai_helper import get_ai_metrics_from_response + + +class OpenAIAgentRunner(AgentRunner): + """ + AgentRunner implementation for OpenAI. + + Executes a single-agent loop using OpenAI Chat Completions with tool calling. + Returned by OpenAIRunnerFactory.create_agent(config, tools). + """ + + def __init__( + self, + client: AsyncOpenAI, + model_name: str, + parameters: Dict[str, Any], + instructions: str, + tool_definitions: List[Dict[str, Any]], + tools: ToolRegistry, + ): + self._client = client + self._model_name = model_name + self._parameters = parameters + self._instructions = instructions + self._tool_definitions = tool_definitions + self._tools = tools + + async def run(self, input: Any) -> AgentResult: + """ + Run the agent with the given input string. + + Executes an agentic loop: calls the model, handles tool calls, + and continues until the model produces a final response. + + :param input: The user prompt or input to the agent + :return: AgentResult with output, raw response, and aggregated metrics + """ + messages: List[Dict[str, Any]] = [] + if self._instructions: + messages.append({"role": "system", "content": self._instructions}) + messages.append({"role": "user", "content": str(input)}) + + total_input = 0 + total_output = 0 + raw_response = None + + try: + while True: + create_kwargs: Dict[str, Any] = { + "model": self._model_name, + "messages": messages, + **self._parameters, + } + openai_tools = self._build_openai_tools() + if openai_tools: + create_kwargs["tools"] = openai_tools + create_kwargs["tool_choice"] = "auto" + + response = await self._client.chat.completions.create(**create_kwargs) # type: ignore[arg-type] + raw_response = response + metrics = get_ai_metrics_from_response(response) + + if metrics.usage: + total_input += metrics.usage.input + total_output += metrics.usage.output + + if not response.choices: + break + + message = response.choices[0].message + + # Add assistant message to history + assistant_msg: Dict[str, Any] = { + "role": "assistant", + "content": message.content, + } + if message.tool_calls: + assistant_msg["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in message.tool_calls + ] + messages.append(assistant_msg) + + if not message.tool_calls: + total_tokens = total_input + total_output + return AgentResult( + output=message.content or "", + raw=raw_response, + metrics=LDAIMetrics( + success=True, + usage=TokenUsage( + total=total_tokens, + input=total_input, + output=total_output, + ) if total_tokens > 0 else None, + ), + ) + + # Execute tool calls and append results + for tool_call in message.tool_calls: + result = await self._call_tool( + tool_call.function.name, + tool_call.function.arguments, + ) + messages.append({ + "role": "tool", + "tool_call_id": tool_call.id, + "content": result, + }) + + except Exception as error: + log.warning(f"OpenAI agent run failed: {error}") + return AgentResult( + output="", + raw=raw_response, + metrics=LDAIMetrics(success=False, usage=None), + ) + + return AgentResult( + output="", + raw=raw_response, + metrics=LDAIMetrics(success=False, usage=None), + ) + + async def _call_tool(self, name: str, arguments_json: str) -> str: + """Execute a tool by name, returning the result as a string.""" + tool_fn = self._tools.get(name) + if not tool_fn: + log.warning(f"Tool '{name}' not found in registry") + return f"Tool '{name}' not found" + try: + args = json.loads(arguments_json) if arguments_json else {} + result = tool_fn(**args) + if hasattr(result, "__await__"): + result = await result + return str(result) + except Exception as error: + log.warning(f"Tool '{name}' execution failed: {error}") + return f"Tool execution failed: {error}" + + def _build_openai_tools(self) -> List[Dict[str, Any]]: + """Convert LD tool definitions to OpenAI function-calling format.""" + tools = [] + for td in self._tool_definitions: + if not isinstance(td, dict): + continue + if "type" in td: + # Already in OpenAI format + tools.append(td) + elif "name" in td: + # LD simplified format: {name, description, parameters} + tools.append({ + "type": "function", + "function": { + "name": td["name"], + "description": td.get("description", ""), + "parameters": td.get("parameters", {"type": "object", "properties": {}}), + }, + }) + return tools + + def get_client(self) -> AsyncOpenAI: + """Return the underlying AsyncOpenAI client.""" + return self._client diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_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 d80fc01..f69e1fb 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_runner_factory.py @@ -1,12 +1,15 @@ import os -from typing import Optional +from typing import TYPE_CHECKING, Any, Optional from ldai.models import AIConfigKind -from ldai.providers import AIProvider +from ldai.providers import AIProvider, ToolRegistry from openai import AsyncOpenAI from ldai_openai.openai_model_runner import OpenAIModelRunner +if TYPE_CHECKING: + from ldai_openai.openai_agent_runner import OpenAIAgentRunner + class OpenAIRunnerFactory(AIProvider): """OpenAI ``AIProvider`` implementation for the LaunchDarkly AI SDK.""" @@ -36,6 +39,32 @@ def create_model(self, config: AIConfigKind) -> OpenAIModelRunner: parameters = model_dict.get('parameters') or {} return OpenAIModelRunner(self._client, model_name, parameters) + def create_agent(self, config: Any, tools: 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..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,36 @@ 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: + """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) + 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 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 class TestGetAIMetricsFromResponse: @@ -318,3 +347,140 @@ def test_handles_missing_model_config(self): assert isinstance(result, OpenAIModelRunner) assert result._model_name == '' assert result._parameters == {} + + +class TestCreateAgent: + """Tests for OpenAIRunnerFactory.create_agent.""" + + def test_creates_agent_runner_with_instructions_and_tool_definitions(self): + """Should create OpenAIAgentRunner with instructions and tool definitions.""" + mock_ai_config = MagicMock() + mock_ai_config.instructions = "You are a helpful assistant." + mock_ai_config.to_dict.return_value = { + 'model': { + 'name': 'gpt-4', + 'parameters': { + 'temperature': 0.7, + 'tools': [ + {'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}, + ], + }, + }, + } + + mock_client = MagicMock() + factory = OpenAIRunnerFactory(mock_client) + result = factory.create_agent(mock_ai_config, {'get-weather': lambda loc: 'sunny'}) + + from ldai_openai import OpenAIAgentRunner + assert isinstance(result, OpenAIAgentRunner) + assert result._model_name == 'gpt-4' + assert result._instructions == "You are a helpful assistant." + assert result._parameters == {'temperature': 0.7} + assert len(result._tool_definitions) == 1 + assert result._tool_definitions[0]['name'] == 'get-weather' + + def test_creates_agent_runner_with_no_tools(self): + """Should create OpenAIAgentRunner with no tool definitions.""" + mock_ai_config = MagicMock() + mock_ai_config.instructions = "You are a helpful assistant." + mock_ai_config.to_dict.return_value = { + 'model': {'name': 'gpt-4', 'parameters': {}}, + } + + mock_client = MagicMock() + factory = OpenAIRunnerFactory(mock_client) + result = factory.create_agent(mock_ai_config, {}) + + from ldai_openai import OpenAIAgentRunner + assert isinstance(result, OpenAIAgentRunner) + assert result._tool_definitions == [] + + +class TestOpenAIAgentRunner: + """Tests for OpenAIAgentRunner.run.""" + + @pytest.fixture + def mock_client(self): + return MagicMock() + + @pytest.mark.asyncio + async def test_runs_agent_and_returns_result_with_no_tool_calls(self, mock_client): + """Should return AgentResult when model responds with no tool calls.""" + from ldai_openai import OpenAIAgentRunner + + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.content = "The answer is 42." + mock_response.choices[0].message.tool_calls = None + mock_response.usage = MagicMock() + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 5 + mock_response.usage.total_tokens = 15 + mock_client.chat.completions.create = AsyncMock(return_value=mock_response) + + runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, 'You are helpful.', [], {}) + result = await runner.run("What is the answer?") + + assert result.output == "The answer is 42." + assert result.metrics.success is True + + @pytest.mark.asyncio + async def test_executes_tool_calls_and_returns_final_response(self, mock_client): + """Should execute tool calls and continue loop until final response.""" + from ldai_openai import OpenAIAgentRunner + + # First response: has a tool call + tool_call = MagicMock() + tool_call.id = "call_123" + tool_call.function.name = "get-weather" + tool_call.function.arguments = '{"location": "Paris"}' + + first_response = MagicMock() + first_response.choices = [MagicMock()] + first_response.choices[0].message.content = None + first_response.choices[0].message.tool_calls = [tool_call] + first_response.usage = MagicMock() + first_response.usage.prompt_tokens = 10 + first_response.usage.completion_tokens = 5 + first_response.usage.total_tokens = 15 + + # Second response: final answer + second_response = MagicMock() + second_response.choices = [MagicMock()] + second_response.choices[0].message.content = "It is sunny in Paris." + second_response.choices[0].message.tool_calls = None + second_response.usage = MagicMock() + second_response.usage.prompt_tokens = 20 + second_response.usage.completion_tokens = 8 + second_response.usage.total_tokens = 28 + + mock_client.chat.completions.create = AsyncMock( + side_effect=[first_response, second_response] + ) + + weather_fn = MagicMock(return_value="Sunny, 25°C") + runner = OpenAIAgentRunner( + mock_client, 'gpt-4', {}, 'You are helpful.', + [{'name': 'get-weather', 'description': 'Get weather', 'parameters': {}}], + {'get-weather': weather_fn}, + ) + result = await runner.run("What is the weather in Paris?") + + assert result.output == "It is sunny in Paris." + assert result.metrics.success is True + weather_fn.assert_called_once_with(location="Paris") + assert mock_client.chat.completions.create.call_count == 2 + + @pytest.mark.asyncio + async def test_returns_failure_when_exception_thrown(self, mock_client): + """Should return unsuccessful AgentResult when exception is thrown.""" + from ldai_openai import OpenAIAgentRunner + + mock_client.chat.completions.create = AsyncMock(side_effect=Exception("API Error")) + + runner = OpenAIAgentRunner(mock_client, 'gpt-4', {}, '', [], {}) + result = await runner.run("Hello") + + assert result.output == "" + assert result.metrics.success is False diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 944a0cb..1131eb4 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -6,6 +6,7 @@ from ldai.chat import Chat # Deprecated — use ManagedModel from ldai.client import LDAIClient from ldai.judge import Judge +from ldai.managed_agent import ManagedAgent from ldai.managed_model import ManagedModel from ldai.models import ( # Deprecated aliases for backward compatibility AIAgentConfig, @@ -55,6 +56,7 @@ 'AICompletionConfigDefault', 'AIJudgeConfig', 'AIJudgeConfigDefault', + 'ManagedAgent', 'ManagedModel', 'EvalScore', 'AgentGraphDefinition', diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 358f9eb..21bcf54 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -7,6 +7,7 @@ from ldai import log from ldai.agent_graph import AgentGraphDefinition from ldai.judge import Judge +from ldai.managed_agent import ManagedAgent from ldai.managed_model import ManagedModel from ldai.models import ( AIAgentConfig, @@ -24,6 +25,7 @@ ModelConfig, ProviderConfig, ) +from ldai.providers import ToolRegistry from ldai.providers.runner_factory import RunnerFactory from ldai.sdk_info import AI_SDK_LANGUAGE, AI_SDK_NAME, AI_SDK_VERSION from ldai.tracker import AIGraphTracker, LDAIConfigTracker @@ -31,6 +33,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' @@ -374,6 +377,58 @@ async def create_chat( log.warning('create_chat() is deprecated, use create_model() instead') return await self.create_model(key, context, default, variables, default_ai_provider) + async def create_agent( + self, + key: str, + context: Context, + tools: Optional[ToolRegistry] = None, + default: Optional[AIAgentConfigDefault] = None, + variables: Optional[Dict[str, Any]] = None, + default_ai_provider: Optional[str] = None, + ) -> Optional[ManagedAgent]: + """ + Creates and returns a new ManagedAgent for AI agent invocations. + + :param key: The key identifying the AI agent configuration to use + :param context: Standard Context used when evaluating flags + :param tools: ToolRegistry mapping tool names to callable implementations + :param default: A default value representing a standard AI agent config result. + When not provided, a disabled config is used as the fallback. + :param variables: Dictionary of values for instruction interpolation + :param default_ai_provider: Optional default AI provider to use + :return: ManagedAgent instance or None if disabled/unsupported + + Example:: + + agent = await client.create_agent( + "customer-support-agent", + context, + tools={"get-order": fetch_order_fn}, + default=AIAgentConfigDefault( + enabled=True, + model=ModelConfig("gpt-4"), + provider=ProviderConfig("openai"), + instructions="You are a helpful customer support agent." + ), + ) + + if agent: + result = await agent.run("Where is my order?") + print(result.output) + """ + self._client.track(_TRACK_USAGE_CREATE_AGENT, context, key, 1) + log.debug(f"Creating managed agent for key: {key}") + config = self.__evaluate_agent(key, context, default or AIAgentConfigDefault.disabled(), variables) + + if not config.enabled or not config.tracker: + return None + + runner = 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