From a23a07505138d1cef960bf4cd2928d6c94dccb47 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 25 Mar 2026 12:42:47 -0500 Subject: [PATCH 1/8] feat: Introduce ManagedModel and ModelRunner feat: Add ModelRunner ABC with invoke_model() and invoke_structured_model() feat: Add ManagedModel replacing Chat; expose get_model_runner() escape hatch feat!: Rename ChatResponse to ModelResponse in providers/types.py feat!: Extract OpenAIModelRunner from OpenAIRunnerFactory; factory is now model-creation-only feat!: Extract LangChainModelRunner from LangChainRunnerFactory; factory is now model-creation-only feat: Add OpenAIHelper with shared utilities for model and future agent runners feat: Add LangChainHelper with shared utilities for model and future agent runners feat!: LangChainRunnerFactory is now a no-arg factory; static helpers moved to LangChainHelper fix: LDClient.create_chat() is deprecated in favour of create_model() fix: Chat alias in ldai.chat is deprecated in favour of ManagedModel fix: Rename ai_provider param to model_runner in Judge and ManagedModel --- .../src/ldai_langchain/__init__.py | 6 + .../src/ldai_langchain/langchain_helper.py | 106 ++++++++ .../ldai_langchain/langchain_model_runner.py | 114 +++++++++ .../langchain_runner_factory.py | 239 +----------------- .../tests/test_langchain_provider.py | 73 +++--- .../src/ldai_openai/__init__.py | 6 + .../src/ldai_openai/openai_helper.py | 46 ++++ .../src/ldai_openai/openai_model_runner.py | 128 ++++++++++ .../src/ldai_openai/openai_runner_factory.py | 189 ++------------ .../tests/test_openai_provider.py | 33 ++- packages/sdk/server-ai/src/ldai/__init__.py | 7 +- .../sdk/server-ai/src/ldai/chat/__init__.py | 186 +------------- packages/sdk/server-ai/src/ldai/client.py | 49 ++-- .../sdk/server-ai/src/ldai/judge/__init__.py | 22 +- .../sdk/server-ai/src/ldai/managed_model.py | 127 ++++++++++ .../server-ai/src/ldai/providers/__init__.py | 2 + .../src/ldai/providers/ai_provider.py | 12 +- .../src/ldai/providers/model_runner.py | 40 +++ .../sdk/server-ai/src/ldai/providers/types.py | 6 +- packages/sdk/server-ai/tests/test_judge.py | 70 ++--- 20 files changed, 758 insertions(+), 703 deletions(-) create mode 100644 packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py create mode 100644 packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py create mode 100644 packages/ai-providers/server-ai-openai/src/ldai_openai/openai_helper.py create mode 100644 packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py create mode 100644 packages/sdk/server-ai/src/ldai/managed_model.py create mode 100644 packages/sdk/server-ai/src/ldai/providers/model_runner.py diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py index fab6cc59..9624606e 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,7 @@ +"""LaunchDarkly AI SDK - LangChain Connector.""" + +from ldai_langchain.langchain_helper import LangChainHelper +from ldai_langchain.langchain_model_runner import LangChainModelRunner from ldai_langchain.langchain_runner_factory import LangChainRunnerFactory __version__ = "0.1.0" @@ -5,4 +9,6 @@ __all__ = [ '__version__', 'LangChainRunnerFactory', + 'LangChainHelper', + 'LangChainModelRunner', ] 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 new file mode 100644 index 00000000..a4fb9e10 --- /dev/null +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py @@ -0,0 +1,106 @@ +"""Shared LangChain utilities for the LaunchDarkly AI SDK.""" + +from typing import Any, Dict, List, Optional, Union + +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage +from ldai import LDMessage +from ldai.models import AIConfigKind +from ldai.providers.types import LDAIMetrics +from ldai.tracker import TokenUsage + + +class LangChainHelper: + """ + Shared utilities for LangChain-based runners (model, agent, agent graph). + + All methods are static — this class is a namespace, not meant to be instantiated. + """ + + @staticmethod + def map_provider(ld_provider_name: str) -> str: + """ + Map a LaunchDarkly provider name to its LangChain equivalent. + + :param ld_provider_name: LaunchDarkly provider name + :return: LangChain-compatible provider name + """ + lowercased_name = ld_provider_name.lower() + # Bedrock is the only provider that uses "provider:model_family" (e.g. Bedrock:Anthropic). + if lowercased_name.startswith('bedrock:'): + return 'bedrock_converse' + + mapping: Dict[str, str] = { + 'gemini': 'google-genai', + 'bedrock': 'bedrock_converse', + } + return mapping.get(lowercased_name, lowercased_name) + + @staticmethod + def convert_messages( + messages: List[LDMessage], + ) -> List[Union[HumanMessage, SystemMessage, AIMessage]]: + """ + Convert LaunchDarkly messages to LangChain message objects. + + :param messages: List of LDMessage objects + :return: List of LangChain message objects + :raises ValueError: If an unsupported message role is encountered + """ + result: List[Union[HumanMessage, SystemMessage, AIMessage]] = [] + for msg in messages: + if msg.role == 'system': + result.append(SystemMessage(content=msg.content)) + elif msg.role == 'user': + result.append(HumanMessage(content=msg.content)) + elif msg.role == 'assistant': + result.append(AIMessage(content=msg.content)) + else: + raise ValueError(f'Unsupported message role: {msg.role}') + return result + + @staticmethod + def create_langchain_model(ai_config: AIConfigKind) -> BaseChatModel: + """ + Create a LangChain BaseChatModel from a LaunchDarkly AI configuration. + + :param ai_config: The LaunchDarkly AI configuration + :return: A configured LangChain BaseChatModel + """ + from langchain.chat_models import init_chat_model + + config_dict = ai_config.to_dict() + model_dict = config_dict.get('model') or {} + provider_dict = config_dict.get('provider') or {} + + model_name = model_dict.get('name', '') + provider = provider_dict.get('name', '') + parameters = model_dict.get('parameters') or {} + + return init_chat_model( + model_name, + model_provider=LangChainHelper.map_provider(provider), + **parameters, + ) + + @staticmethod + def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: + """ + Extract LaunchDarkly AI metrics from a LangChain response. + + :param response: The response from a LangChain model (BaseMessage or similar) + :return: LDAIMetrics with success status and token usage + """ + usage: Optional[TokenUsage] = None + if hasattr(response, 'response_metadata') and response.response_metadata: + token_usage = ( + response.response_metadata.get('tokenUsage') + or response.response_metadata.get('token_usage') + ) + if token_usage: + usage = TokenUsage( + total=token_usage.get('totalTokens', 0) or token_usage.get('total_tokens', 0), + input=token_usage.get('promptTokens', 0) or token_usage.get('prompt_tokens', 0), + output=token_usage.get('completionTokens', 0) or token_usage.get('completion_tokens', 0), + ) + return LDAIMetrics(success=True, usage=usage) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py new file mode 100644 index 00000000..113670ae --- /dev/null +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py @@ -0,0 +1,114 @@ +"""LangChain model runner for LaunchDarkly AI SDK.""" + +from typing import Any, Dict, List + +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import BaseMessage +from ldai import LDMessage, log +from ldai.providers.model_runner import ModelRunner +from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse +from ldai_langchain.langchain_helper import LangChainHelper + + +class LangChainModelRunner(ModelRunner): + """ + ModelRunner implementation for LangChain. + + Holds a fully-configured BaseChatModel. + Returned by LangChainConnector.create_model(config). + """ + + def __init__(self, llm: BaseChatModel): + self._llm = llm + + def get_llm(self) -> BaseChatModel: + """ + Return the underlying LangChain BaseChatModel. + + :return: The BaseChatModel instance + """ + return self._llm + + async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse: + """ + Invoke the LangChain model with an array of messages. + + :param messages: Array of LDMessage objects representing the conversation + :return: ModelResponse containing the model's response and metrics + """ + try: + langchain_messages = LangChainHelper.convert_messages(messages) + response: BaseMessage = await self._llm.ainvoke(langchain_messages) + metrics = LangChainHelper.get_ai_metrics_from_response(response) + + content: str = '' + if isinstance(response.content, str): + content = response.content + else: + log.warning( + f'Multimodal response not supported, expecting a string. ' + f'Content type: {type(response.content)}, Content: {response.content}' + ) + metrics = LDAIMetrics(success=False, usage=metrics.usage) + + return ModelResponse( + message=LDMessage(role='assistant', content=content), + metrics=metrics, + ) + except Exception as error: + log.warning(f'LangChain model invocation failed: {error}') + return ModelResponse( + message=LDMessage(role='assistant', content=''), + metrics=LDAIMetrics(success=False, usage=None), + ) + + async def invoke_structured_model( + self, + messages: List[LDMessage], + response_structure: Dict[str, Any], + ) -> StructuredResponse: + """ + Invoke the LangChain model with structured output support. + + :param messages: Array of LDMessage objects representing the conversation + :param response_structure: Dictionary defining the output structure + :return: StructuredResponse containing the structured data + """ + try: + langchain_messages = LangChainHelper.convert_messages(messages) + structured_llm = self._llm.with_structured_output(response_structure, include_raw=True) + response = await structured_llm.ainvoke(langchain_messages) + + if not isinstance(response, dict): + log.warning(f'Structured output did not return a dict. Got: {type(response)}') + return StructuredResponse( + data={}, + raw_response='', + metrics=LDAIMetrics(success=False, usage=None), + ) + + raw_response = response.get('raw') + usage = LangChainHelper.get_ai_usage_from_response(raw_response) if raw_response is not None else None + raw_content = raw_response.content if hasattr(raw_response, 'content') else '' + + if response.get('parsing_error'): + log.warning('LangChain structured model invocation had a parsing error') + return StructuredResponse( + data={}, + raw_response=raw_content, + metrics=LDAIMetrics(success=False, usage=usage), + ) + + return StructuredResponse( + data=response.get('parsed') or {}, + raw_response=raw_content, + metrics=LDAIMetrics(success=True, usage=usage), + ) + except Exception as error: + log.warning(f'LangChain structured model invocation failed: {error}') + return StructuredResponse( + data={}, + raw_response='', + metrics=LDAIMetrics(success=False, usage=None), + ) + 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 afb5c8b1..41c8a145 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,238 +1,25 @@ -from typing import Any, Dict, List, Optional, Union +"""LangChain connector for LaunchDarkly AI SDK.""" -from langchain_core.language_models.chat_models import BaseChatModel -from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage -from ldai import LDMessage, log from ldai.models import AIConfigKind from ldai.providers import AIProvider -from ldai.providers.types import ChatResponse, LDAIMetrics, StructuredResponse -from ldai.tracker import TokenUsage +from ldai_langchain.langchain_helper import LangChainHelper +from ldai_langchain.langchain_model_runner import LangChainModelRunner class LangChainRunnerFactory(AIProvider): - """LangChain provider for the LaunchDarkly AI SDK.""" + """ + LangChain connector for the LaunchDarkly AI SDK. - def __init__(self, llm: Optional[BaseChatModel] = None): - """ - :param llm: A LangChain BaseChatModel instance (optional) - """ - self._llm = llm + Acts as a per-provider factory. Instantiate with no arguments, then call + ``create_model(config)`` to obtain a configured ``LangChainModelRunner``. + """ - def create_model(self, config: AIConfigKind) -> 'LangChainRunnerFactory': + def create_model(self, config: AIConfigKind) -> LangChainModelRunner: """ - Create a configured LangChain model provider for the given AI config. + Create a configured LangChainModelRunner for the given AI config. :param config: The LaunchDarkly AI configuration - :return: Configured LangChainRunnerFactory ready to invoke the model - """ - llm = LangChainRunnerFactory.create_langchain_model(config) - return LangChainRunnerFactory(llm) - - async def invoke_model(self, messages: List[LDMessage]) -> ChatResponse: - """ - Invoke the LangChain model with an array of messages. - - :param messages: Array of LDMessage objects representing the conversation - :return: ChatResponse containing the model's response and metrics - """ - try: - assert self._llm is not None - langchain_messages = LangChainRunnerFactory.convert_messages_to_langchain(messages) - response: BaseMessage = await self._llm.ainvoke(langchain_messages) - metrics = LangChainRunnerFactory.get_ai_metrics_from_response(response) - - content: str = '' - if isinstance(response.content, str): - content = response.content - else: - log.warning( - f'Multimodal response not supported, expecting a string. ' - f'Content type: {type(response.content)}, Content: {response.content}' - ) - metrics = LDAIMetrics(success=False, usage=metrics.usage) - - return ChatResponse( - message=LDMessage(role='assistant', content=content), - metrics=metrics, - ) - except Exception as error: - log.warning(f'LangChain model invocation failed: {error}') - - return ChatResponse( - message=LDMessage(role='assistant', content=''), - metrics=LDAIMetrics(success=False, usage=None), - ) - - async def invoke_structured_model( - self, - messages: List[LDMessage], - response_structure: Dict[str, Any], - ) -> StructuredResponse: - """ - Invoke the LangChain model with structured output support. - - :param messages: Array of LDMessage objects representing the conversation - :param response_structure: Dictionary defining the output structure - :return: StructuredResponse containing the structured data + :return: LangChainModelRunner ready to invoke the model """ - structured_response = StructuredResponse( - data={}, - raw_response='', - metrics=LDAIMetrics(success=False, usage=None), - ) - try: - assert self._llm is not None - langchain_messages = LangChainRunnerFactory.convert_messages_to_langchain(messages) - structured_llm = self._llm.with_structured_output(response_structure, include_raw=True) - response = await structured_llm.ainvoke(langchain_messages) - - if not isinstance(response, dict): - log.warning( - f'Structured output did not return a dict. ' - f'Got: {type(response)}' - ) - return structured_response - - raw_response = response.get('raw') - if raw_response is not None: - if hasattr(raw_response, 'content'): - structured_response.raw_response = raw_response.content - structured_response.metrics.usage = LangChainRunnerFactory.get_ai_usage_from_response(raw_response) - - if response.get('parsing_error'): - log.warning(f'LangChain structured model invocation had a parsing error') - return structured_response - - structured_response.metrics.success = True - structured_response.data = response.get('parsed') or {} - return structured_response - except Exception as error: - log.warning(f'LangChain structured model invocation failed: {error}') - return structured_response - - def get_chat_model(self) -> Optional[BaseChatModel]: - """ - Get the underlying LangChain model instance. - - :return: The underlying BaseChatModel, or None if not yet configured - """ - return self._llm - - @staticmethod - def map_provider(ld_provider_name: str) -> str: - """ - Map LaunchDarkly provider names to LangChain provider names. - - :param ld_provider_name: LaunchDarkly provider name - :return: LangChain-compatible provider name - """ - lowercased_name = ld_provider_name.lower() - # Bedrock is the only provider that uses "provider:model_family" (e.g. Bedrock:Anthropic). - if lowercased_name.startswith('bedrock:'): - return 'bedrock_converse' - - mapping: Dict[str, str] = { - 'gemini': 'google-genai', - 'bedrock': 'bedrock_converse', - } - return mapping.get(lowercased_name, lowercased_name) - - @staticmethod - def get_ai_usage_from_response(response: BaseMessage) -> TokenUsage: - """ - Get token usage from a LangChain provider response. - - :param response: The response from the LangChain model - :return: TokenUsage with success status and token usage - """ - # Extract token usage if available - usage: Optional[TokenUsage] = None - if hasattr(response, 'usage_metadata') and response.usage_metadata: - usage = 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), - ) - if not usage and hasattr(response, 'response_metadata') and response.response_metadata: - token_usage = response.response_metadata.get('tokenUsage') or response.response_metadata.get('token_usage') - if token_usage: - usage = TokenUsage( - total=token_usage.get('totalTokens', 0) or token_usage.get('total_tokens', 0), - input=token_usage.get('promptTokens', 0) or token_usage.get('prompt_tokens', 0), - output=token_usage.get('completionTokens', 0) or token_usage.get('completion_tokens', 0), - ) - - return usage - - @staticmethod - def get_ai_metrics_from_response(response: BaseMessage) -> LDAIMetrics: - """ - Extract LaunchDarkly AI metrics from a LangChain response. - - :param response: The response from the LangChain model - :return: LDAIMetrics with success status and token usage - - Example:: - - response = await tracker.track_metrics_of( - lambda: llm.ainvoke(messages), - LangChainRunnerFactory.get_ai_metrics_from_response - ) - """ - usage = LangChainRunnerFactory.get_ai_usage_from_response(response) - - return LDAIMetrics(success=True, usage=usage) - - @staticmethod - def convert_messages_to_langchain( - messages: List[LDMessage], - ) -> List[Union[HumanMessage, SystemMessage, AIMessage]]: - """ - Convert LaunchDarkly messages to LangChain messages. - - :param messages: List of LDMessage objects - :return: List of LangChain message objects - :raises ValueError: If an unsupported message role is encountered - """ - result: List[Union[HumanMessage, SystemMessage, AIMessage]] = [] - - for msg in messages: - if msg.role == 'system': - result.append(SystemMessage(content=msg.content)) - elif msg.role == 'user': - result.append(HumanMessage(content=msg.content)) - elif msg.role == 'assistant': - result.append(AIMessage(content=msg.content)) - else: - raise ValueError(f'Unsupported message role: {msg.role}') - - return result - - @staticmethod - def create_langchain_model(ai_config: AIConfigKind) -> BaseChatModel: - """ - Create a LangChain model from a LaunchDarkly AI configuration. - - :param ai_config: The LaunchDarkly AI configuration - :return: A configured LangChain BaseChatModel - """ - from langchain.chat_models import init_chat_model - - config_dict = ai_config.to_dict() - model_dict = config_dict.get('model') or {} - provider_dict = config_dict.get('provider') or {} - - model_name = model_dict.get('name', '') - provider = provider_dict.get('name', '') - parameters = dict(model_dict.get('parameters') or {}) - mapped_provider = LangChainRunnerFactory.map_provider(provider) - - # Bedrock requires the foundation provider (e.g. Bedrock:Anthropic) passed in - # parameters separately from model_provider, which is used for LangChain routing. - if mapped_provider == 'bedrock_converse' and 'provider' not in parameters: - parameters['provider'] = provider.removeprefix('bedrock:') - return init_chat_model( - model_name, - model_provider=mapped_provider, - **parameters, - ) + llm = LangChainHelper.create_langchain_model(config) + return LangChainModelRunner(llm) 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 9abc047c..611b59fc 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 @@ -7,16 +7,16 @@ from ldai import LDMessage -from ldai_langchain import LangChainRunnerFactory +from ldai_langchain import LangChainHelper, LangChainModelRunner, LangChainRunnerFactory -class TestConvertMessagesToLangchain: - """Tests for convert_messages_to_langchain static method.""" +class TestConvertMessages: + """Tests for LangChainHelper.convert_messages.""" def test_converts_system_messages_to_system_message(self): """Should convert system messages to SystemMessage.""" messages = [LDMessage(role='system', content='You are a helpful assistant.')] - result = LangChainRunnerFactory.convert_messages_to_langchain(messages) + result = LangChainHelper.convert_messages(messages) assert len(result) == 1 assert isinstance(result[0], SystemMessage) @@ -25,7 +25,7 @@ def test_converts_system_messages_to_system_message(self): def test_converts_user_messages_to_human_message(self): """Should convert user messages to HumanMessage.""" messages = [LDMessage(role='user', content='Hello, how are you?')] - result = LangChainRunnerFactory.convert_messages_to_langchain(messages) + result = LangChainHelper.convert_messages(messages) assert len(result) == 1 assert isinstance(result[0], HumanMessage) @@ -34,7 +34,7 @@ def test_converts_user_messages_to_human_message(self): def test_converts_assistant_messages_to_ai_message(self): """Should convert assistant messages to AIMessage.""" messages = [LDMessage(role='assistant', content='I am doing well, thank you!')] - result = LangChainRunnerFactory.convert_messages_to_langchain(messages) + result = LangChainHelper.convert_messages(messages) assert len(result) == 1 assert isinstance(result[0], AIMessage) @@ -47,7 +47,7 @@ def test_converts_multiple_messages_in_order(self): LDMessage(role='user', content='What is the weather like?'), LDMessage(role='assistant', content='I cannot check the weather.'), ] - result = LangChainRunnerFactory.convert_messages_to_langchain(messages) + result = LangChainHelper.convert_messages(messages) assert len(result) == 3 assert isinstance(result[0], SystemMessage) @@ -56,22 +56,21 @@ def test_converts_multiple_messages_in_order(self): def test_throws_error_for_unsupported_message_role(self): """Should throw error for unsupported message role.""" - # Create a mock message with unsupported role class MockMessage: role = 'unknown' content = 'Test message' - + with pytest.raises(ValueError, match='Unsupported message role: unknown'): - LangChainRunnerFactory.convert_messages_to_langchain([MockMessage()]) # type: ignore + LangChainHelper.convert_messages([MockMessage()]) # type: ignore def test_handles_empty_message_array(self): """Should handle empty message array.""" - result = LangChainRunnerFactory.convert_messages_to_langchain([]) + result = LangChainHelper.convert_messages([]) assert len(result) == 0 class TestGetAIMetricsFromResponse: - """Tests for get_ai_metrics_from_response static method.""" + """Tests for LangChainHelper.get_ai_metrics_from_response.""" def test_creates_metrics_with_success_true_and_token_usage(self): """Should create metrics with success=True and token usage.""" @@ -84,7 +83,7 @@ def test_creates_metrics_with_success_true_and_token_usage(self): }, } - result = LangChainRunnerFactory.get_ai_metrics_from_response(mock_response) + result = LangChainHelper.get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is not None @@ -103,7 +102,7 @@ def test_creates_metrics_with_snake_case_token_usage(self): }, } - result = LangChainRunnerFactory.get_ai_metrics_from_response(mock_response) + result = LangChainHelper.get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is not None @@ -115,34 +114,34 @@ def test_creates_metrics_with_success_true_and_no_usage_when_metadata_missing(se """Should create metrics with success=True and no usage when metadata is missing.""" mock_response = AIMessage(content='Test response') - result = LangChainRunnerFactory.get_ai_metrics_from_response(mock_response) + result = LangChainHelper.get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is None class TestMapProvider: - """Tests for map_provider static method.""" + """Tests for LangChainHelper.map_provider.""" def test_maps_gemini_to_google_genai(self): """Should map gemini to google-genai.""" - assert LangChainRunnerFactory.map_provider('gemini') == 'google-genai' - assert LangChainRunnerFactory.map_provider('Gemini') == 'google-genai' - assert LangChainRunnerFactory.map_provider('GEMINI') == 'google-genai' + assert LangChainHelper.map_provider('gemini') == 'google-genai' + assert LangChainHelper.map_provider('Gemini') == 'google-genai' + assert LangChainHelper.map_provider('GEMINI') == 'google-genai' def test_maps_bedrock_and_model_families_to_bedrock_converse(self): """Should map bedrock and bedrock:model_family to bedrock_converse.""" - assert LangChainRunnerFactory.map_provider('bedrock') == 'bedrock_converse' - assert LangChainRunnerFactory.map_provider('Bedrock:Anthropic') == 'bedrock_converse' - assert LangChainRunnerFactory.map_provider('bedrock:anthropic') == 'bedrock_converse' - assert LangChainRunnerFactory.map_provider('bedrock:amazon') == 'bedrock_converse' - assert LangChainRunnerFactory.map_provider('bedrock:cohere') == 'bedrock_converse' + assert LangChainHelper.map_provider('bedrock') == 'bedrock_converse' + assert LangChainHelper.map_provider('Bedrock:Anthropic') == 'bedrock_converse' + assert LangChainHelper.map_provider('bedrock:anthropic') == 'bedrock_converse' + assert LangChainHelper.map_provider('bedrock:amazon') == 'bedrock_converse' + assert LangChainHelper.map_provider('bedrock:cohere') == 'bedrock_converse' def test_returns_provider_name_unchanged_for_unmapped_providers(self): """Should return provider name unchanged for unmapped providers.""" - assert LangChainRunnerFactory.map_provider('openai') == 'openai' - assert LangChainRunnerFactory.map_provider('anthropic') == 'anthropic' - assert LangChainRunnerFactory.map_provider('unknown') == 'unknown' + assert LangChainHelper.map_provider('openai') == 'openai' + assert LangChainHelper.map_provider('anthropic') == 'anthropic' + assert LangChainHelper.map_provider('unknown') == 'unknown' class TestInvokeModel: @@ -158,7 +157,7 @@ async def test_returns_success_true_for_string_content(self, mock_llm): """Should return success=True for string content.""" mock_response = AIMessage(content='Test response') mock_llm.ainvoke = AsyncMock(return_value=mock_response) - provider = LangChainRunnerFactory(mock_llm) + provider = LangChainModelRunner(mock_llm) messages = [LDMessage(role='user', content='Hello')] result = await provider.invoke_model(messages) @@ -171,7 +170,7 @@ async def test_returns_success_false_for_non_string_content_and_logs_warning(sel """Should return success=False for non-string content and log warning.""" mock_response = AIMessage(content=[{'type': 'image', 'data': 'base64data'}]) mock_llm.ainvoke = AsyncMock(return_value=mock_response) - provider = LangChainRunnerFactory(mock_llm) + provider = LangChainModelRunner(mock_llm) messages = [LDMessage(role='user', content='Hello')] result = await provider.invoke_model(messages) @@ -184,7 +183,7 @@ async def test_returns_success_false_when_model_invocation_throws_error(self, mo """Should return success=False when model invocation throws an error.""" error = Exception('Model invocation failed') mock_llm.ainvoke = AsyncMock(side_effect=error) - provider = LangChainRunnerFactory(mock_llm) + provider = LangChainModelRunner(mock_llm) messages = [LDMessage(role='user', content='Hello')] result = await provider.invoke_model(messages) @@ -210,7 +209,7 @@ async def test_returns_success_true_for_successful_invocation(self, mock_llm): mock_structured_llm = MagicMock() mock_structured_llm.ainvoke = AsyncMock(return_value=mock_response) mock_llm.with_structured_output = MagicMock(return_value=mock_structured_llm) - provider = LangChainRunnerFactory(mock_llm) + provider = LangChainModelRunner(mock_llm) messages = [LDMessage(role='user', content='Hello')] response_structure = {'type': 'object', 'properties': {}} @@ -226,7 +225,7 @@ async def test_returns_success_false_when_structured_model_invocation_throws_err mock_structured_llm = MagicMock() mock_structured_llm.ainvoke = AsyncMock(side_effect=error) mock_llm.with_structured_output = MagicMock(return_value=mock_structured_llm) - provider = LangChainRunnerFactory(mock_llm) + provider = LangChainModelRunner(mock_llm) messages = [LDMessage(role='user', content='Hello')] response_structure = {'type': 'object', 'properties': {}} @@ -238,14 +237,12 @@ async def test_returns_success_false_when_structured_model_invocation_throws_err assert result.metrics.usage is None -class TestGetChatModel: - """Tests for get_chat_model instance method.""" +class TestGetLlm: + """Tests for LangChainModelRunner.get_llm.""" def test_returns_underlying_llm(self): """Should return the underlying LLM.""" mock_llm = MagicMock() - provider = LangChainRunnerFactory(mock_llm) - - assert provider.get_chat_model() is mock_llm - + runner = LangChainModelRunner(mock_llm) + assert runner.get_llm() is mock_llm 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 a24989f0..51c1c404 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,5 +1,11 @@ +"""LaunchDarkly AI SDK OpenAI Connector.""" + +from ldai_openai.openai_helper import OpenAIHelper +from ldai_openai.openai_model_runner import OpenAIModelRunner from ldai_openai.openai_runner_factory import OpenAIRunnerFactory __all__ = [ 'OpenAIRunnerFactory', + 'OpenAIHelper', + 'OpenAIModelRunner', ] 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 new file mode 100644 index 00000000..b868a86d --- /dev/null +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_helper.py @@ -0,0 +1,46 @@ +"""Shared OpenAI utilities for the LaunchDarkly AI SDK.""" + +from typing import Any, Iterable, List, Optional, cast + +from ldai import LDMessage +from ldai.providers.types import LDAIMetrics +from ldai.tracker import TokenUsage +from openai.types.chat import ChatCompletionMessageParam + + +class OpenAIHelper: + """ + Shared utilities for OpenAI-based runners (model, agent, agent graph). + + All methods are static — this class is a namespace, not meant to be instantiated. + """ + + @staticmethod + def convert_messages(messages: List[LDMessage]) -> Iterable[ChatCompletionMessageParam]: + """ + Convert LaunchDarkly messages to OpenAI chat completion message format. + + :param messages: List of LDMessage objects + :return: Iterable of OpenAI ChatCompletionMessageParam dicts + """ + return cast( + Iterable[ChatCompletionMessageParam], + [{'role': msg.role, 'content': msg.content} for msg in messages], + ) + + @staticmethod + 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 + :return: LDAIMetrics with success status and token usage + """ + usage: Optional[TokenUsage] = None + if hasattr(response, 'usage') and response.usage: + usage = TokenUsage( + total=response.usage.total_tokens or 0, + input=response.usage.prompt_tokens or 0, + output=response.usage.completion_tokens or 0, + ) + return LDAIMetrics(success=True, usage=usage) diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py new file mode 100644 index 00000000..9ccb4d4c --- /dev/null +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py @@ -0,0 +1,128 @@ +"""OpenAI model runner for LaunchDarkly AI SDK.""" + +import json +from typing import Any, Dict, List + +from ldai import LDMessage, log +from ldai.providers.model_runner import ModelRunner +from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse +from ldai.tracker import TokenUsage +from openai import AsyncOpenAI +from ldai_openai.openai_helper import OpenAIHelper + + +class OpenAIModelRunner(ModelRunner): + """ + ModelRunner implementation for OpenAI. + + Holds a fully-configured AsyncOpenAI client, model name, and parameters. + Returned by OpenAIConnector.create_model(config). + """ + + def __init__( + self, + client: AsyncOpenAI, + model_name: str, + parameters: Dict[str, Any], + ): + self._client = client + self._model_name = model_name + self._parameters = parameters + + async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse: + """ + Invoke the OpenAI model with an array of messages. + + :param messages: Array of LDMessage objects representing the conversation + :return: ModelResponse containing the model's response and metrics + """ + try: + response = await self._client.chat.completions.create( + model=self._model_name, + messages=OpenAIHelper.convert_messages(messages), + **self._parameters, + ) + + metrics = OpenAIHelper.get_ai_metrics_from_response(response) + + content = '' + if response.choices and len(response.choices) > 0: + message = response.choices[0].message + if message and message.content: + content = message.content + + if not content: + log.warning('OpenAI response has no content available') + metrics = LDAIMetrics(success=False, usage=metrics.usage) + + return ModelResponse( + message=LDMessage(role='assistant', content=content), + metrics=metrics, + ) + except Exception as error: + log.warning(f'OpenAI model invocation failed: {error}') + return ModelResponse( + message=LDMessage(role='assistant', content=''), + metrics=LDAIMetrics(success=False, usage=None), + ) + + async def invoke_structured_model( + self, + messages: List[LDMessage], + response_structure: Dict[str, Any], + ) -> StructuredResponse: + """ + Invoke the OpenAI model with structured output support. + + :param messages: Array of LDMessage objects representing the conversation + :param response_structure: Dictionary defining the JSON schema for output structure + :return: StructuredResponse containing the structured data + """ + try: + response = await self._client.chat.completions.create( + model=self._model_name, + messages=OpenAIHelper.convert_messages(messages), + response_format={ # type: ignore[arg-type] + 'type': 'json_schema', + 'json_schema': { + 'name': 'structured_output', + 'schema': response_structure, + 'strict': True, + }, + }, + **self._parameters, + ) + + metrics = OpenAIHelper.get_ai_metrics_from_response(response) + + content = '' + if response.choices and len(response.choices) > 0: + message = response.choices[0].message + if message and message.content: + content = message.content + + if not content: + log.warning('OpenAI structured response has no content available') + return StructuredResponse( + data={}, + raw_response='', + metrics=LDAIMetrics(success=False, usage=metrics.usage), + ) + + try: + data = json.loads(content) + return StructuredResponse(data=data, raw_response=content, metrics=metrics) + except json.JSONDecodeError as parse_error: + log.warning(f'OpenAI structured response contains invalid JSON: {parse_error}') + return StructuredResponse( + data={}, + raw_response=content, + metrics=LDAIMetrics(success=False, usage=metrics.usage), + ) + except Exception as error: + log.warning(f'OpenAI structured model invocation failed: {error}') + return StructuredResponse( + data={}, + raw_response='', + metrics=LDAIMetrics(success=False, usage=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 a5ae0ff0..c3ee773a 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,197 +1,54 @@ -import json +"""OpenAI connector for LaunchDarkly AI SDK.""" + import os -from typing import Any, Dict, Iterable, List, Optional, cast +from typing import Optional -from ldai import LDMessage, log from ldai.models import AIConfigKind from ldai.providers import AIProvider -from ldai.providers.types import ChatResponse, LDAIMetrics, StructuredResponse -from ldai.tracker import TokenUsage +from ldai_openai.openai_model_runner import OpenAIModelRunner from openai import AsyncOpenAI -from openai.types.chat import ChatCompletionMessageParam class OpenAIRunnerFactory(AIProvider): - """OpenAI provider for the LaunchDarkly AI SDK.""" + """ + OpenAI connector for the LaunchDarkly AI SDK. + + Acts as a per-provider factory. Instantiate with no arguments to read + credentials from the environment (``OPENAI_API_KEY``), then call + ``create_model(config)`` to obtain a configured ``OpenAIModelRunner``. - def __init__( - self, - client: Optional[AsyncOpenAI] = None, - model_name: str = '', - parameters: Optional[Dict[str, Any]] = None, - ): + For advanced use, pass an explicit ``AsyncOpenAI`` client. + """ + + def __init__(self, client: Optional[AsyncOpenAI] = None): """ + Initialize the OpenAI connector. + :param client: An AsyncOpenAI client instance (created from env if omitted) - :param model_name: The name of the model to use - :param parameters: Additional model parameters """ self._client = client if client is not None else AsyncOpenAI( api_key=os.environ.get('OPENAI_API_KEY'), ) - self._model_name = model_name - self._parameters = parameters or {} - def create_model(self, config: AIConfigKind) -> 'OpenAIRunnerFactory': + def create_model(self, config: AIConfigKind) -> OpenAIModelRunner: """ - Create a configured OpenAI model provider for the given AI config. + Create a configured OpenAIModelRunner for the given AI config. - Reuses the underlying AsyncOpenAI client so that connection pooling is - preserved across calls. + Reuses the underlying AsyncOpenAI client so connection pooling is preserved. :param config: The LaunchDarkly AI configuration - :return: Configured OpenAIRunnerFactory ready to invoke the model + :return: OpenAIModelRunner ready to invoke the model """ config_dict = config.to_dict() model_dict = config_dict.get('model') or {} model_name = model_dict.get('name', '') parameters = model_dict.get('parameters') or {} - return OpenAIRunnerFactory(self._client, model_name, parameters) - - async def invoke_model(self, messages: List[LDMessage]) -> ChatResponse: - """ - Invoke the OpenAI model with an array of messages. - - :param messages: Array of LDMessage objects representing the conversation - :return: ChatResponse containing the model's response and metrics - """ - try: - openai_messages: Iterable[ChatCompletionMessageParam] = cast( - Iterable[ChatCompletionMessageParam], - [{'role': msg.role, 'content': msg.content} for msg in messages] - ) - - response = await self._client.chat.completions.create( - model=self._model_name, - messages=openai_messages, - **self._parameters, - ) - - metrics = OpenAIRunnerFactory.get_ai_metrics_from_response(response) - - content = '' - if response.choices and len(response.choices) > 0: - message = response.choices[0].message - if message and message.content: - content = message.content - - if not content: - log.warning('OpenAI response has no content available') - metrics = LDAIMetrics(success=False, usage=metrics.usage) - - return ChatResponse( - message=LDMessage(role='assistant', content=content), - metrics=metrics, - ) - except Exception as error: - log.warning(f'OpenAI model invocation failed: {error}') - - return ChatResponse( - message=LDMessage(role='assistant', content=''), - metrics=LDAIMetrics(success=False, usage=None), - ) - - async def invoke_structured_model( - self, - messages: List[LDMessage], - response_structure: Dict[str, Any], - ) -> StructuredResponse: - """ - Invoke the OpenAI model with structured output support. - - :param messages: Array of LDMessage objects representing the conversation - :param response_structure: Dictionary defining the JSON schema for output structure - :return: StructuredResponse containing the structured data - """ - try: - openai_messages: Iterable[ChatCompletionMessageParam] = cast( - Iterable[ChatCompletionMessageParam], - [{'role': msg.role, 'content': msg.content} for msg in messages] - ) - - response = await self._client.chat.completions.create( - model=self._model_name, - messages=openai_messages, - response_format={ # type: ignore[arg-type] - 'type': 'json_schema', - 'json_schema': { - 'name': 'structured_output', - 'schema': response_structure, - 'strict': True, - }, - }, - **self._parameters, - ) - - metrics = OpenAIRunnerFactory.get_ai_metrics_from_response(response) - - content = '' - if response.choices and len(response.choices) > 0: - message = response.choices[0].message - if message and message.content: - content = message.content - - if not content: - log.warning('OpenAI structured response has no content available') - metrics = LDAIMetrics(success=False, usage=metrics.usage) - return StructuredResponse( - data={}, - raw_response='', - metrics=metrics, - ) - - try: - data = json.loads(content) - return StructuredResponse( - data=data, - raw_response=content, - metrics=metrics, - ) - except json.JSONDecodeError as parse_error: - log.warning(f'OpenAI structured response contains invalid JSON: {parse_error}') - metrics = LDAIMetrics(success=False, usage=metrics.usage) - return StructuredResponse( - data={}, - raw_response=content, - metrics=metrics, - ) - except Exception as error: - log.warning(f'OpenAI structured model invocation failed: {error}') - - return StructuredResponse( - data={}, - raw_response='', - metrics=LDAIMetrics(success=False, usage=None), - ) + return OpenAIModelRunner(self._client, model_name, parameters) def get_client(self) -> AsyncOpenAI: """ - Get the underlying OpenAI client instance. + Return the underlying AsyncOpenAI client. - :return: The underlying AsyncOpenAI client + :return: The AsyncOpenAI client instance """ return self._client - - @staticmethod - def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: - """ - Extract LaunchDarkly AI metrics from an OpenAI response. - - :param response: The response from OpenAI chat completions API - :return: LDAIMetrics with success status and token usage - - Example:: - - response = await tracker.track_metrics_of( - lambda: client.chat.completions.create(config), - OpenAIRunnerFactory.get_ai_metrics_from_response - ) - """ - usage: Optional[TokenUsage] = None - if hasattr(response, 'usage') and response.usage: - usage = TokenUsage( - total=response.usage.total_tokens or 0, - input=response.usage.prompt_tokens or 0, - output=response.usage.completion_tokens or 0, - ) - - return LDAIMetrics(success=True, usage=usage) 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 385f74a1..595fe9fb 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,11 +5,11 @@ from ldai import LDMessage -from ldai_openai import OpenAIRunnerFactory +from ldai_openai import OpenAIHelper, OpenAIModelRunner, OpenAIRunnerFactory class TestGetAIMetricsFromResponse: - """Tests for get_ai_metrics_from_response static method.""" + """Tests for OpenAIHelper.get_ai_metrics_from_response.""" def test_creates_metrics_with_success_true_and_token_usage(self): """Should create metrics with success=True and token usage.""" @@ -19,7 +19,7 @@ def test_creates_metrics_with_success_true_and_token_usage(self): mock_response.usage.completion_tokens = 50 mock_response.usage.total_tokens = 100 - result = OpenAIRunnerFactory.get_ai_metrics_from_response(mock_response) + result = OpenAIHelper.get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is not None @@ -32,7 +32,7 @@ def test_creates_metrics_with_success_true_and_no_usage_when_usage_missing(self) mock_response = MagicMock() mock_response.usage = None - result = OpenAIRunnerFactory.get_ai_metrics_from_response(mock_response) + result = OpenAIHelper.get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is None @@ -45,7 +45,7 @@ def test_handles_partial_usage_data(self): mock_response.usage.completion_tokens = None mock_response.usage.total_tokens = None - result = OpenAIRunnerFactory.get_ai_metrics_from_response(mock_response) + result = OpenAIHelper.get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is not None @@ -78,7 +78,7 @@ async def test_invokes_openai_chat_completions_and_returns_response(self, mock_c mock_client.chat.completions = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIModelRunner(mock_client, 'gpt-3.5-turbo', {}) messages = [LDMessage(role='user', content='Hello!')] result = await provider.invoke_model(messages) @@ -108,7 +108,7 @@ async def test_returns_unsuccessful_response_when_no_content(self, mock_client): mock_client.chat.completions = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIModelRunner(mock_client, 'gpt-3.5-turbo', {}) messages = [LDMessage(role='user', content='Hello!')] result = await provider.invoke_model(messages) @@ -127,7 +127,7 @@ async def test_returns_unsuccessful_response_when_choices_empty(self, mock_clien mock_client.chat.completions = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIModelRunner(mock_client, 'gpt-3.5-turbo', {}) messages = [LDMessage(role='user', content='Hello!')] result = await provider.invoke_model(messages) @@ -142,7 +142,7 @@ async def test_returns_unsuccessful_response_when_exception_thrown(self, mock_cl mock_client.chat.completions = MagicMock() mock_client.chat.completions.create = AsyncMock(side_effect=Exception('API Error')) - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIModelRunner(mock_client, 'gpt-3.5-turbo', {}) messages = [LDMessage(role='user', content='Hello!')] result = await provider.invoke_model(messages) @@ -175,7 +175,7 @@ async def test_invokes_openai_with_structured_output(self, mock_client): mock_client.chat.completions = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIModelRunner(mock_client, 'gpt-3.5-turbo', {}) messages = [LDMessage(role='user', content='Tell me about a person')] response_structure = { 'type': 'object', @@ -210,7 +210,7 @@ async def test_returns_unsuccessful_when_no_content_in_structured_response(self, mock_client.chat.completions = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIModelRunner(mock_client, 'gpt-3.5-turbo', {}) messages = [LDMessage(role='user', content='Tell me about a person')] response_structure = {'type': 'object'} @@ -236,7 +236,7 @@ async def test_handles_json_parsing_errors(self, mock_client): mock_client.chat.completions = MagicMock() mock_client.chat.completions.create = AsyncMock(return_value=mock_response) - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIModelRunner(mock_client, 'gpt-3.5-turbo', {}) messages = [LDMessage(role='user', content='Tell me about a person')] response_structure = {'type': 'object'} @@ -255,7 +255,7 @@ async def test_returns_unsuccessful_response_when_exception_thrown(self, mock_cl mock_client.chat.completions = MagicMock() mock_client.chat.completions.create = AsyncMock(side_effect=Exception('API Error')) - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIModelRunner(mock_client, 'gpt-3.5-turbo', {}) messages = [LDMessage(role='user', content='Tell me about a person')] response_structure = {'type': 'object'} @@ -272,7 +272,7 @@ class TestGetClient: def test_returns_underlying_client(self): """Should return the underlying OpenAI client.""" mock_client = MagicMock() - provider = OpenAIRunnerFactory(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIRunnerFactory(mock_client) assert provider.get_client() is mock_client @@ -300,7 +300,7 @@ def test_creates_connector_with_correct_model_and_parameters(self): result = OpenAIRunnerFactory().create_model(mock_ai_config) - assert isinstance(result, OpenAIRunnerFactory) + assert isinstance(result, OpenAIModelRunner) assert result._model_name == 'gpt-4' assert result._parameters == {'temperature': 0.7, 'max_tokens': 1000} @@ -315,7 +315,6 @@ def test_handles_missing_model_config(self): result = OpenAIRunnerFactory().create_model(mock_ai_config) - assert isinstance(result, OpenAIRunnerFactory) + assert isinstance(result, OpenAIModelRunner) assert result._model_name == '' assert result._parameters == {} - diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index d7863f60..718d6921 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -3,7 +3,8 @@ from ldclient import log from ldai.agent_graph import AgentGraphDefinition -from ldai.chat import Chat +from ldai.managed_model import ManagedModel +from ldai.chat import Chat # Deprecated — use ManagedModel from ldai.client import LDAIClient from ldai.judge import Judge from ldai.models import ( # Deprecated aliases for backward compatibility @@ -42,8 +43,10 @@ 'AICompletionConfigDefault', 'AIJudgeConfig', 'AIJudgeConfigDefault', - 'Chat', + 'ManagedModel', 'EvalScore', + # Deprecated — use ManagedModel + 'Chat', 'AgentGraphDefinition', 'Judge', 'JudgeConfiguration', diff --git a/packages/sdk/server-ai/src/ldai/chat/__init__.py b/packages/sdk/server-ai/src/ldai/chat/__init__.py index c826fed6..cc3dbff9 100644 --- a/packages/sdk/server-ai/src/ldai/chat/__init__.py +++ b/packages/sdk/server-ai/src/ldai/chat/__init__.py @@ -1,184 +1,8 @@ -"""Chat implementation for managing AI chat conversations.""" +"""Backward-compatibility shim — use ldai.managed_model.ManagedModel instead.""" -import asyncio -from typing import Any, Dict, List, Optional +from ldai.managed_model import ManagedModel -from ldai import log -from ldai.judge import Judge -from ldai.models import AICompletionConfig, LDMessage -from ldai.providers.ai_provider import AIProvider -from ldai.providers.types import ChatResponse, JudgeResponse -from ldai.tracker import LDAIConfigTracker +# Deprecated alias +Chat = ManagedModel - -class Chat: - """ - Concrete implementation of Chat that provides chat functionality - by delegating to an AIProvider implementation. - - This class handles conversation management and tracking, while delegating - the actual model invocation to the provider. - """ - - def __init__( - self, - ai_config: AICompletionConfig, - tracker: LDAIConfigTracker, - provider: AIProvider, - judges: Optional[Dict[str, Judge]] = None, - ): - """ - Initialize the Chat. - - :param ai_config: The completion AI configuration - :param tracker: The tracker for the completion configuration - :param provider: The AI provider to use for chat - :param judges: Optional dictionary of judge instances keyed by their configuration keys - """ - self._ai_config = ai_config - self._tracker = tracker - self._provider = provider - self._judges = judges or {} - self._messages: List[LDMessage] = [] - - async def invoke(self, prompt: str) -> ChatResponse: - """ - Invoke the chat model with a prompt string. - - This method handles conversation management and tracking, delegating to the provider's invoke_model method. - - :param prompt: The user prompt to send to the chat model - :return: ChatResponse containing the model's response and metrics - """ - # Convert prompt string to LDMessage with role 'user' and add to conversation history - user_message: LDMessage = LDMessage(role='user', content=prompt) - self._messages.append(user_message) - - # Prepend config messages to conversation history for model invocation - config_messages = self._ai_config.messages or [] - all_messages = config_messages + self._messages - - # Delegate to provider-specific implementation with tracking - response = await self._tracker.track_metrics_of( - lambda: self._provider.invoke_model(all_messages), - lambda result: result.metrics, - ) - - # Start judge evaluations as async tasks (don't await them) - if ( - self._ai_config.judge_configuration - and self._ai_config.judge_configuration.judges - and len(self._ai_config.judge_configuration.judges) > 0 - ): - response.evaluations = self._start_judge_evaluations(self._messages, response) - - # Add the response message to conversation history - self._messages.append(response.message) - return response - - def _start_judge_evaluations( - self, - messages: List[LDMessage], - response: ChatResponse, - ) -> List[asyncio.Task[Optional[JudgeResponse]]]: - """ - Start judge evaluations as async tasks without awaiting them. - - Returns a list of async tasks that can be awaited later. - - :param messages: Array of messages representing the conversation history - :param response: The AI response to be evaluated - :return: List of async tasks that will return judge evaluation results - """ - if not self._ai_config.judge_configuration or not self._ai_config.judge_configuration.judges: - return [] - - judge_configs = self._ai_config.judge_configuration.judges - - # Start all judge evaluations as tasks - async def evaluate_judge(judge_config): - judge = self._judges.get(judge_config.key) - if not judge: - log.warn( - f"Judge configuration is not enabled: {judge_config.key}", - ) - return None - - eval_result = await judge.evaluate_messages( - messages, response, judge_config.sampling_rate - ) - - if eval_result and eval_result.success: - self._tracker.track_judge_response(eval_result) - - return eval_result - - # Create tasks for each judge evaluation - tasks = [ - asyncio.create_task(evaluate_judge(judge_config)) - for judge_config in judge_configs - ] - - return tasks - - def get_config(self) -> AICompletionConfig: - """ - Get the underlying AI configuration used to initialize this Chat. - - :return: The AI completion configuration - """ - return self._ai_config - - def get_tracker(self) -> LDAIConfigTracker: - """ - Get the underlying AI configuration tracker used to initialize this Chat. - - :return: The tracker instance - """ - return self._tracker - - def get_provider(self) -> AIProvider: - """ - Get the underlying AI provider instance. - - This provides direct access to the provider for advanced use cases. - - :return: The AI provider instance - """ - return self._provider - - def get_judges(self) -> Dict[str, Judge]: - """ - Get the judges associated with this Chat. - - Returns a dictionary of judge instances keyed by their configuration keys. - - :return: Dictionary of judge instances - """ - return self._judges - - def append_messages(self, messages: List[LDMessage]) -> None: - """ - Append messages to the conversation history. - - Adds messages to the conversation history without invoking the model, - which is useful for managing multi-turn conversations or injecting context. - - :param messages: Array of messages to append to the conversation history - """ - self._messages.extend(messages) - - def get_messages(self, include_config_messages: bool = False) -> List[LDMessage]: - """ - Get all messages in the conversation history. - - :param include_config_messages: Whether to include the config messages from the AIConfig. - Defaults to False. - :return: Array of messages. When include_config_messages is True, returns both config - messages and conversation history with config messages prepended. When False, - returns only the conversation history messages. - """ - if include_config_messages: - config_messages = self._ai_config.messages or [] - return config_messages + self._messages - return list(self._messages) +__all__ = ['ManagedModel', 'Chat'] diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index d6801993..b6229268 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -6,7 +6,7 @@ from ldai import log from ldai.agent_graph import AgentGraphDefinition -from ldai.chat import Chat +from ldai.managed_model import ManagedModel from ldai.judge import Judge from ldai.models import ( AIAgentConfig, @@ -30,7 +30,7 @@ _TRACK_SDK_INFO = '$ld:ai:sdk:info' _TRACK_USAGE_COMPLETION_CONFIG = '$ld:ai:usage:completion-config' -_TRACK_USAGE_CREATE_CHAT = '$ld:ai:usage:create-chat' +_TRACK_USAGE_CREATE_MODEL = '$ld:ai:usage:create-model' _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' @@ -298,16 +298,16 @@ async def create_judge_for_config(judge_key: str): return judges - async def create_chat( + async def create_model( self, key: str, context: Context, default: Optional[AICompletionConfigDefault] = None, variables: Optional[Dict[str, Any]] = None, default_ai_provider: Optional[str] = None, - ) -> Optional[Chat]: + ) -> Optional[ManagedModel]: """ - Creates and returns a new Chat instance for AI conversations. + Creates and returns a new ManagedModel for AI conversations. :param key: The key identifying the AI completion configuration to use :param context: Standard Context used when evaluating flags @@ -315,11 +315,11 @@ async def create_chat( 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: Chat instance or None if disabled/unsupported + :return: ManagedModel instance or None if disabled/unsupported Example:: - chat = await client.create_chat( + model = await client.create_model( "customer-support-chat", context, AICompletionConfigDefault( @@ -331,23 +331,19 @@ async def create_chat( variables={'customerName': 'John'} ) - if chat: - response = await chat.invoke("I need help with my order") + if model: + response = await model.invoke("I need help with my order") print(response.message.content) - - # Access conversation history - messages = chat.get_messages() - print(f"Conversation has {len(messages)} messages") """ - self._client.track(_TRACK_USAGE_CREATE_CHAT, context, key, 1) - log.debug(f"Creating chat for key: {key}") + self._client.track(_TRACK_USAGE_CREATE_MODEL, context, key, 1) + log.debug(f"Creating managed model for key: {key}") config = self._completion_config(key, context, default or AICompletionConfigDefault.disabled(), variables) if not config.enabled or not config.tracker: return None - provider = RunnerFactory.create_model(config, default_ai_provider) - if not provider: + runner = RunnerFactory.create_model(config, default_ai_provider) + if not runner: return None judges = {} @@ -359,7 +355,24 @@ async def create_chat( default_ai_provider, ) - return Chat(config, config.tracker, provider, judges) + return ManagedModel(config, config.tracker, runner, judges) + + async def create_chat( + self, + key: str, + context: Context, + default: Optional[AICompletionConfigDefault] = None, + variables: Optional[Dict[str, Any]] = None, + default_ai_provider: Optional[str] = None, + ) -> Optional[ManagedModel]: + """ + .. deprecated:: Use :meth:`create_model` instead. + + Creates and returns a ManagedModel for AI conversations. + This method is a deprecated alias for :meth:`create_model`. + """ + log.warn('create_chat() is deprecated, use create_model() instead') + return await self.create_model(key, context, default, variables, default_ai_provider) def agent_config( self, diff --git a/packages/sdk/server-ai/src/ldai/judge/__init__.py b/packages/sdk/server-ai/src/ldai/judge/__init__.py index 280dc791..3a247310 100644 --- a/packages/sdk/server-ai/src/ldai/judge/__init__.py +++ b/packages/sdk/server-ai/src/ldai/judge/__init__.py @@ -8,8 +8,8 @@ from ldai import log from ldai.judge.evaluation_schema_builder import EvaluationSchemaBuilder from ldai.models import AIJudgeConfig, LDMessage -from ldai.providers.ai_provider import AIProvider -from ldai.providers.types import ChatResponse, EvalScore, JudgeResponse +from ldai.providers.model_runner import ModelRunner +from ldai.providers.types import EvalScore, JudgeResponse, ModelResponse from ldai.tracker import LDAIConfigTracker @@ -25,18 +25,18 @@ def __init__( self, ai_config: AIJudgeConfig, ai_config_tracker: LDAIConfigTracker, - ai_provider: AIProvider, + model_runner: ModelRunner, ): """ Initialize the Judge. :param ai_config: The judge AI configuration :param ai_config_tracker: The tracker for the judge configuration - :param ai_provider: The AI provider to use for evaluation + :param model_runner: The model runner to use for evaluation """ self._ai_config = ai_config self._ai_config_tracker = ai_config_tracker - self._ai_provider = ai_provider + self._model_runner = model_runner self._evaluation_response_structure = EvaluationSchemaBuilder.build() async def evaluate( @@ -72,7 +72,7 @@ async def evaluate( assert self._evaluation_response_structure is not None response = await self._ai_config_tracker.track_metrics_of( - lambda: self._ai_provider.invoke_structured_model(messages, self._evaluation_response_structure), + lambda: self._model_runner.invoke_structured_model(messages, self._evaluation_response_structure), lambda result: result.metrics, ) @@ -99,7 +99,7 @@ async def evaluate( async def evaluate_messages( self, messages: list[LDMessage], - response: ChatResponse, + response: ModelResponse, sampling_ratio: float = 1.0, ) -> Optional[JudgeResponse]: """ @@ -131,13 +131,13 @@ def get_tracker(self) -> LDAIConfigTracker: """ return self._ai_config_tracker - def get_provider(self) -> AIProvider: + def get_model_runner(self) -> ModelRunner: """ - Returns the AI provider used by this judge. + Returns the model runner used by this judge. - :return: The AI provider + :return: The model runner """ - return self._ai_provider + return self._model_runner def _construct_evaluation_messages(self, input_text: str, output_text: str) -> list[LDMessage]: """ diff --git a/packages/sdk/server-ai/src/ldai/managed_model.py b/packages/sdk/server-ai/src/ldai/managed_model.py new file mode 100644 index 00000000..93c3c675 --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/managed_model.py @@ -0,0 +1,127 @@ +"""ManagedModel — LaunchDarkly managed wrapper for model invocations.""" + +import asyncio +from typing import Any, Dict, List, Optional + +from ldai import log +from ldai.judge import Judge +from ldai.models import AICompletionConfig, LDMessage +from ldai.providers.model_runner import ModelRunner +from ldai.providers.types import JudgeResponse, ModelResponse +from ldai.tracker import LDAIConfigTracker + + +class ManagedModel: + """ + LaunchDarkly managed wrapper for AI model invocations. + + Holds a ModelRunner and an LDAIConfigTracker. Handles conversation + management, judge evaluation dispatch, and tracking automatically. + Obtain an instance via ``LDAIClient.create_model()``. + """ + + def __init__( + self, + ai_config: AICompletionConfig, + tracker: LDAIConfigTracker, + model_runner: ModelRunner, + judges: Optional[Dict[str, Judge]] = None, + ): + self._ai_config = ai_config + self._tracker = tracker + self._model_runner = model_runner + self._judges = judges or {} + self._messages: List[LDMessage] = [] + + async def invoke(self, prompt: str) -> ModelResponse: + """ + Invoke the model with a prompt string. + + Appends the prompt to the conversation history, prepends any + system messages from the config, delegates to the runner, and + appends the response to the history. + + :param prompt: The user prompt to send to the model + :return: ModelResponse containing the model's response and metrics + """ + user_message = LDMessage(role='user', content=prompt) + self._messages.append(user_message) + + config_messages = self._ai_config.messages or [] + all_messages = config_messages + self._messages + + response = await self._tracker.track_metrics_of( + lambda: self._model_runner.invoke_model(all_messages), + lambda result: result.metrics, + ) + + if ( + self._ai_config.judge_configuration + and self._ai_config.judge_configuration.judges + ): + response.evaluations = self._start_judge_evaluations(self._messages, response) + + self._messages.append(response.message) + return response + + def _start_judge_evaluations( + self, + messages: List[LDMessage], + response: ModelResponse, + ) -> List[asyncio.Task[Optional[JudgeResponse]]]: + if not self._ai_config.judge_configuration or not self._ai_config.judge_configuration.judges: + return [] + + async def evaluate_judge(judge_config: Any) -> Optional[JudgeResponse]: + judge = self._judges.get(judge_config.key) + if not judge: + log.warn(f'Judge configuration is not enabled: {judge_config.key}') + return None + eval_result = await judge.evaluate_messages(messages, response, judge_config.sampling_rate) + if eval_result and eval_result.success: + self._tracker.track_judge_response(eval_result) + return eval_result + + return [ + asyncio.create_task(evaluate_judge(jc)) + for jc in self._ai_config.judge_configuration.judges + ] + + def get_messages(self, include_config_messages: bool = False) -> List[LDMessage]: + """ + Get all messages in the conversation history. + + :param include_config_messages: When True, prepends config messages. + :return: List of conversation messages. + """ + if include_config_messages: + return (self._ai_config.messages or []) + self._messages + return list(self._messages) + + def append_messages(self, messages: List[LDMessage]) -> None: + """ + Append messages to the conversation history without invoking the model. + + :param messages: Messages to append. + """ + self._messages.extend(messages) + + def get_model_runner(self) -> ModelRunner: + """ + Return the underlying ModelRunner for advanced use. + + :return: The ModelRunner instance. + """ + return self._model_runner + + def get_config(self) -> AICompletionConfig: + """Return the AI completion config.""" + return self._ai_config + + def get_tracker(self) -> LDAIConfigTracker: + """Return the config tracker.""" + return self._tracker + + def get_judges(self) -> Dict[str, Judge]: + """Return the judges associated with this model.""" + return self._judges diff --git a/packages/sdk/server-ai/src/ldai/providers/__init__.py b/packages/sdk/server-ai/src/ldai/providers/__init__.py index cbf2b5f5..db1c8836 100644 --- a/packages/sdk/server-ai/src/ldai/providers/__init__.py +++ b/packages/sdk/server-ai/src/ldai/providers/__init__.py @@ -1,7 +1,9 @@ from ldai.providers.ai_provider import AIProvider +from ldai.providers.model_runner import ModelRunner from ldai.providers.runner_factory import RunnerFactory __all__ = [ 'AIProvider', + 'ModelRunner', 'RunnerFactory', ] 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 fe2d0fb5..576f1c19 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 ChatResponse, StructuredResponse +from ldai.providers.types import ModelResponse, StructuredResponse class AIProvider(ABC): @@ -16,7 +16,7 @@ class AIProvider(ABC): create_model(), create_agent(), and create_agent_graph(). """ - async def invoke_model(self, messages: List[LDMessage]) -> ChatResponse: + async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse: """ Invoke the chat model with an array of messages. @@ -24,14 +24,14 @@ async def invoke_model(self, messages: List[LDMessage]) -> ChatResponse: Provider implementations should override this method. :param messages: Array of LDMessage objects representing the conversation - :return: ChatResponse containing the model's response + :return: ModelResponse containing the model's response """ log.warning('invoke_model not implemented by this provider') from ldai.models import LDMessage from ldai.providers.types import LDAIMetrics - return ChatResponse( + return ModelResponse( message=LDMessage(role='assistant', content=''), metrics=LDAIMetrics(success=False, usage=None), ) @@ -61,14 +61,14 @@ async def invoke_structured_model( metrics=LDAIMetrics(success=False, usage=None), ) - def create_model(self, config: Any) -> Optional['AIProvider']: + def create_model(self, config: Any) -> Optional[Any]: """ Create a configured model executor for the given AI config. Default implementation warns. Provider implementations should override this method. :param config: The LaunchDarkly AI configuration - :return: Configured AIProvider instance, or None if unsupported + :return: Configured model runner instance, or None if unsupported """ log.warning('create_model not implemented by this provider') return None diff --git a/packages/sdk/server-ai/src/ldai/providers/model_runner.py b/packages/sdk/server-ai/src/ldai/providers/model_runner.py new file mode 100644 index 00000000..6365203e --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/providers/model_runner.py @@ -0,0 +1,40 @@ +"""Abstract base class for model runners.""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, List + +from ldai.models import LDMessage +from ldai.providers.types import ModelResponse, StructuredResponse + + +class ModelRunner(ABC): + """ + Runtime capability interface for model invocation. + + A ModelRunner is a focused, configured object returned by + AIConnector.create_model(). It knows exactly which model to call + and with what parameters — the caller just passes messages. + """ + + @abstractmethod + async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse: + """ + Invoke the model with an array of messages. + + :param messages: Array of LDMessage objects representing the conversation + :return: ModelResponse containing the model's response and metrics + """ + + @abstractmethod + async def invoke_structured_model( + self, + messages: List[LDMessage], + response_structure: Dict[str, Any], + ) -> StructuredResponse: + """ + Invoke the model with structured output support. + + :param messages: Array of LDMessage objects representing the conversation + :param response_structure: Dictionary defining the JSON schema for output structure + :return: StructuredResponse containing the structured data + """ diff --git a/packages/sdk/server-ai/src/ldai/providers/types.py b/packages/sdk/server-ai/src/ldai/providers/types.py index e9160cc8..0a07151f 100644 --- a/packages/sdk/server-ai/src/ldai/providers/types.py +++ b/packages/sdk/server-ai/src/ldai/providers/types.py @@ -34,13 +34,13 @@ def to_dict(self) -> Dict[str, Any]: @dataclass -class ChatResponse: +class ModelResponse: """ - Chat response structure. + Response from a model invocation. """ message: LDMessage metrics: LDAIMetrics - evaluations: Optional[List[JudgeResponse]] = None # List of JudgeResponse, will be populated later + evaluations: Optional[List[JudgeResponse]] = None @dataclass diff --git a/packages/sdk/server-ai/tests/test_judge.py b/packages/sdk/server-ai/tests/test_judge.py index 7c0d3785..8326b721 100644 --- a/packages/sdk/server-ai/tests/test_judge.py +++ b/packages/sdk/server-ai/tests/test_judge.py @@ -39,7 +39,7 @@ def client(td: TestData) -> LDClient: @pytest.fixture -def mock_ai_provider(): +def mock_runner(): """Create a mock AI provider.""" provider = MagicMock() provider.invoke_structured_model = AsyncMock() @@ -101,10 +101,10 @@ class TestJudgeInitialization: """Tests for Judge initialization.""" def test_judge_initializes_with_evaluation_metric_key( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Judge should initialize successfully with evaluation_metric_key.""" - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) assert judge._ai_config == judge_config_with_key assert judge._evaluation_response_structure is not None @@ -119,31 +119,31 @@ class TestJudgeEvaluate: @pytest.mark.asyncio async def test_evaluate_returns_none_when_evaluation_metric_key_missing( - self, judge_config_without_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_without_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should return None when evaluation_metric_key is missing.""" - judge = Judge(judge_config_without_key, tracker, mock_ai_provider) + judge = Judge(judge_config_without_key, tracker, mock_runner) result = await judge.evaluate("input text", "output text") assert result is None - mock_ai_provider.invoke_structured_model.assert_not_called() + mock_runner.invoke_structured_model.assert_not_called() @pytest.mark.asyncio async def test_evaluate_returns_none_when_messages_missing( - self, judge_config_without_messages: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_without_messages: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should return None when messages are missing.""" - judge = Judge(judge_config_without_messages, tracker, mock_ai_provider) + judge = Judge(judge_config_without_messages, tracker, mock_runner) result = await judge.evaluate("input text", "output text") assert result is None - mock_ai_provider.invoke_structured_model.assert_not_called() + mock_runner.invoke_structured_model.assert_not_called() @pytest.mark.asyncio async def test_evaluate_success_with_valid_response( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should return JudgeResponse with valid evaluation.""" mock_response = StructuredResponse( @@ -155,10 +155,10 @@ async def test_evaluate_success_with_valid_response( metrics=LDAIMetrics(success=True) ) - mock_ai_provider.invoke_structured_model.return_value = mock_response + mock_runner.invoke_structured_model.return_value = mock_response tracker.track_metrics_of = AsyncMock(return_value=mock_response) - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) result = await judge.evaluate("What is AI?", "AI is artificial intelligence.") @@ -171,7 +171,7 @@ async def test_evaluate_success_with_valid_response( @pytest.mark.asyncio async def test_evaluate_success_with_evaluation_response_shape( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should accept shape { score, reasoning } and key by metric.""" mock_response = StructuredResponse( @@ -182,10 +182,10 @@ async def test_evaluate_success_with_evaluation_response_shape( raw_response='{"score": 0.9, "reasoning": "..."}', metrics=LDAIMetrics(success=True), ) - mock_ai_provider.invoke_structured_model.return_value = mock_response + mock_runner.invoke_structured_model.return_value = mock_response tracker.track_metrics_of = AsyncMock(return_value=mock_response) - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) result = await judge.evaluate("What is feature flagging?", "Feature flagging is...") assert result is not None @@ -196,7 +196,7 @@ async def test_evaluate_success_with_evaluation_response_shape( @pytest.mark.asyncio async def test_evaluate_handles_missing_evaluation_in_response( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should handle missing score/reasoning in response.""" mock_response = StructuredResponse( @@ -205,10 +205,10 @@ async def test_evaluate_handles_missing_evaluation_in_response( metrics=LDAIMetrics(success=True) ) - mock_ai_provider.invoke_structured_model.return_value = mock_response + mock_runner.invoke_structured_model.return_value = mock_response tracker.track_metrics_of = AsyncMock(return_value=mock_response) - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) result = await judge.evaluate("input", "output") @@ -218,7 +218,7 @@ async def test_evaluate_handles_missing_evaluation_in_response( @pytest.mark.asyncio async def test_evaluate_handles_invalid_score( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should handle invalid score values.""" mock_response = StructuredResponse( @@ -230,10 +230,10 @@ async def test_evaluate_handles_invalid_score( metrics=LDAIMetrics(success=True) ) - mock_ai_provider.invoke_structured_model.return_value = mock_response + mock_runner.invoke_structured_model.return_value = mock_response tracker.track_metrics_of = AsyncMock(return_value=mock_response) - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) result = await judge.evaluate("input", "output") @@ -243,7 +243,7 @@ async def test_evaluate_handles_invalid_score( @pytest.mark.asyncio async def test_evaluate_handles_missing_reasoning( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should handle missing reasoning.""" mock_response = StructuredResponse( @@ -252,10 +252,10 @@ async def test_evaluate_handles_missing_reasoning( metrics=LDAIMetrics(success=True) ) - mock_ai_provider.invoke_structured_model.return_value = mock_response + mock_runner.invoke_structured_model.return_value = mock_response tracker.track_metrics_of = AsyncMock(return_value=mock_response) - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) result = await judge.evaluate("input", "output") @@ -265,13 +265,13 @@ async def test_evaluate_handles_missing_reasoning( @pytest.mark.asyncio async def test_evaluate_handles_exception( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should handle exceptions gracefully.""" - mock_ai_provider.invoke_structured_model.side_effect = Exception("Provider error") + mock_runner.invoke_structured_model.side_effect = Exception("Provider error") tracker.track_metrics_of = AsyncMock(side_effect=Exception("Provider error")) - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) result = await judge.evaluate("input", "output") @@ -283,15 +283,15 @@ async def test_evaluate_handles_exception( @pytest.mark.asyncio async def test_evaluate_respects_sampling_rate( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Evaluate should respect sampling rate.""" - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) result = await judge.evaluate("input", "output", sampling_rate=0.0) assert result is None - mock_ai_provider.invoke_structured_model.assert_not_called() + mock_runner.invoke_structured_model.assert_not_called() class TestJudgeEvaluateMessages: @@ -299,10 +299,10 @@ class TestJudgeEvaluateMessages: @pytest.mark.asyncio async def test_evaluate_messages_calls_evaluate( - self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """evaluate_messages should call evaluate with constructed input/output.""" - from ldai.providers.types import ChatResponse + from ldai.providers.types import ModelResponse mock_response = StructuredResponse( data={'score': 0.9, 'reasoning': 'Very relevant'}, @@ -310,16 +310,16 @@ async def test_evaluate_messages_calls_evaluate( metrics=LDAIMetrics(success=True) ) - mock_ai_provider.invoke_structured_model.return_value = mock_response + mock_runner.invoke_structured_model.return_value = mock_response tracker.track_metrics_of = AsyncMock(return_value=mock_response) - judge = Judge(judge_config_with_key, tracker, mock_ai_provider) + judge = Judge(judge_config_with_key, tracker, mock_runner) messages = [ LDMessage(role='user', content='Question 1'), LDMessage(role='assistant', content='Answer 1'), ] - chat_response = ChatResponse( + chat_response = ModelResponse( message=LDMessage(role='assistant', content='Answer 2'), metrics=LDAIMetrics(success=True) ) From fc3880c8c836b12e5e88053d7825d48f193a2e9c Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 25 Mar 2026 12:47:23 -0500 Subject: [PATCH 2/8] Remove module docstrings and use log.warning instead of log.warn --- .../src/ldai_langchain/__init__.py | 2 -- .../src/ldai_langchain/langchain_helper.py | 2 -- .../src/ldai_langchain/langchain_model_runner.py | 2 -- .../src/ldai_langchain/langchain_runner_factory.py | 2 -- .../server-ai-openai/src/ldai_openai/__init__.py | 2 -- .../src/ldai_openai/openai_helper.py | 2 -- .../src/ldai_openai/openai_model_runner.py | 2 -- .../src/ldai_openai/openai_runner_factory.py | 2 -- packages/sdk/server-ai/src/ldai/client.py | 2 +- packages/sdk/server-ai/src/ldai/judge/__init__.py | 14 +++++++------- packages/sdk/server-ai/src/ldai/managed_model.py | 4 +--- .../server-ai/src/ldai/providers/model_runner.py | 2 -- 12 files changed, 9 insertions(+), 29 deletions(-) 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 9624606e..40e1d967 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,5 +1,3 @@ -"""LaunchDarkly AI SDK - LangChain Connector.""" - from ldai_langchain.langchain_helper import LangChainHelper from ldai_langchain.langchain_model_runner import LangChainModelRunner from ldai_langchain.langchain_runner_factory import LangChainRunnerFactory 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 a4fb9e10..8c0c5907 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 @@ -1,5 +1,3 @@ -"""Shared LangChain utilities for the LaunchDarkly AI SDK.""" - from typing import Any, Dict, List, Optional, Union from langchain_core.language_models.chat_models import BaseChatModel diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py index 113670ae..5986b05c 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py @@ -1,5 +1,3 @@ -"""LangChain model runner for LaunchDarkly AI SDK.""" - from typing import Any, Dict, List from langchain_core.language_models.chat_models import BaseChatModel 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 41c8a145..aa69d737 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,5 +1,3 @@ -"""LangChain connector for LaunchDarkly AI SDK.""" - from ldai.models import AIConfigKind from ldai.providers import AIProvider from ldai_langchain.langchain_helper import LangChainHelper 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 51c1c404..59e0b313 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,5 +1,3 @@ -"""LaunchDarkly AI SDK OpenAI Connector.""" - from ldai_openai.openai_helper import OpenAIHelper from ldai_openai.openai_model_runner import OpenAIModelRunner from ldai_openai.openai_runner_factory import OpenAIRunnerFactory 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 b868a86d..4e61de8f 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,5 +1,3 @@ -"""Shared OpenAI utilities for the LaunchDarkly AI SDK.""" - from typing import Any, Iterable, List, Optional, cast from ldai import LDMessage diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py index 9ccb4d4c..871436cf 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py @@ -1,5 +1,3 @@ -"""OpenAI model runner for LaunchDarkly AI SDK.""" - import json from typing import Any, Dict, List 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 c3ee773a..95aec0a8 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,3 @@ -"""OpenAI connector for LaunchDarkly AI SDK.""" - import os from typing import Optional diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index b6229268..9f4e86c4 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -371,7 +371,7 @@ async def create_chat( Creates and returns a ManagedModel for AI conversations. This method is a deprecated alias for :meth:`create_model`. """ - log.warn('create_chat() is deprecated, use create_model() instead') + log.warning('create_chat() is deprecated, use create_model() instead') return await self.create_model(key, context, default, variables, default_ai_provider) def agent_config( diff --git a/packages/sdk/server-ai/src/ldai/judge/__init__.py b/packages/sdk/server-ai/src/ldai/judge/__init__.py index 3a247310..b364fb0d 100644 --- a/packages/sdk/server-ai/src/ldai/judge/__init__.py +++ b/packages/sdk/server-ai/src/ldai/judge/__init__.py @@ -55,13 +55,13 @@ async def evaluate( """ try: if not self._ai_config.evaluation_metric_key: - log.warn( + log.warning( 'Judge configuration is missing required evaluationMetricKey' ) return None if not self._ai_config.messages: - log.warn('Judge configuration must include messages') + log.warning('Judge configuration must include messages') return None if random.random() > sampling_rate: @@ -80,7 +80,7 @@ async def evaluate( evals = self._parse_evaluation_response(response.data) if not evals: - log.warn('Judge evaluation did not return the expected evaluation') + log.warning('Judge evaluation did not return the expected evaluation') success = False return JudgeResponse( @@ -179,20 +179,20 @@ def _parse_evaluation_response(self, data: Dict[str, Any]) -> Dict[str, EvalScor results: Dict[str, EvalScore] = {} metric_key = self._ai_config.evaluation_metric_key if not metric_key: - log.warn('Evaluation metric key is missing') + log.warning('Evaluation metric key is missing') return results if not isinstance(data, dict): - log.warn('Invalid response: missing or invalid evaluation') + log.warning('Invalid response: missing or invalid evaluation') return results score = data.get('score') reasoning = data.get('reasoning') if not isinstance(score, (int, float)) or score < 0 or score > 1: - log.warn(f'Invalid score: {score}. Score must be a number between 0 and 1 inclusive') + log.warning(f'Invalid score: {score}. Score must be a number between 0 and 1 inclusive') return results if not isinstance(reasoning, str): - log.warn('Invalid reasoning: must be a string') + log.warning('Invalid reasoning: must be a string') return results results[metric_key] = EvalScore(score=float(score), reasoning=reasoning) diff --git a/packages/sdk/server-ai/src/ldai/managed_model.py b/packages/sdk/server-ai/src/ldai/managed_model.py index 93c3c675..28bab2f0 100644 --- a/packages/sdk/server-ai/src/ldai/managed_model.py +++ b/packages/sdk/server-ai/src/ldai/managed_model.py @@ -1,5 +1,3 @@ -"""ManagedModel — LaunchDarkly managed wrapper for model invocations.""" - import asyncio from typing import Any, Dict, List, Optional @@ -75,7 +73,7 @@ def _start_judge_evaluations( async def evaluate_judge(judge_config: Any) -> Optional[JudgeResponse]: judge = self._judges.get(judge_config.key) if not judge: - log.warn(f'Judge configuration is not enabled: {judge_config.key}') + log.warning(f'Judge configuration is not enabled: {judge_config.key}') return None eval_result = await judge.evaluate_messages(messages, response, judge_config.sampling_rate) if eval_result and eval_result.success: diff --git a/packages/sdk/server-ai/src/ldai/providers/model_runner.py b/packages/sdk/server-ai/src/ldai/providers/model_runner.py index 6365203e..79309f33 100644 --- a/packages/sdk/server-ai/src/ldai/providers/model_runner.py +++ b/packages/sdk/server-ai/src/ldai/providers/model_runner.py @@ -1,5 +1,3 @@ -"""Abstract base class for model runners.""" - from abc import ABC, abstractmethod from typing import Any, Dict, List From 61f825b25e3804999f7779ed81b50ce5b7b99d6f Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 25 Mar 2026 12:54:35 -0500 Subject: [PATCH 3/8] Fix runner factory docstrings --- .../src/ldai_langchain/langchain_runner_factory.py | 7 +------ .../src/ldai_openai/openai_runner_factory.py | 10 +--------- 2 files changed, 2 insertions(+), 15 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py index aa69d737..700d1887 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 @@ -5,12 +5,7 @@ class LangChainRunnerFactory(AIProvider): - """ - LangChain connector for the LaunchDarkly AI SDK. - - Acts as a per-provider factory. Instantiate with no arguments, then call - ``create_model(config)`` to obtain a configured ``LangChainModelRunner``. - """ + """LangChain ``AIProvider`` implementation for the LaunchDarkly AI SDK.""" def create_model(self, config: AIConfigKind) -> LangChainModelRunner: """ 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 95aec0a8..0e542087 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 @@ -8,15 +8,7 @@ class OpenAIRunnerFactory(AIProvider): - """ - OpenAI connector for the LaunchDarkly AI SDK. - - Acts as a per-provider factory. Instantiate with no arguments to read - credentials from the environment (``OPENAI_API_KEY``), then call - ``create_model(config)`` to obtain a configured ``OpenAIModelRunner``. - - For advanced use, pass an explicit ``AsyncOpenAI`` client. - """ + """OpenAI ``AIProvider`` implementation for the LaunchDarkly AI SDK.""" def __init__(self, client: Optional[AsyncOpenAI] = None): """ From 9b72435e1016c9719bb96bc46f8fcbb7e863eeb9 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 25 Mar 2026 13:06:47 -0500 Subject: [PATCH 4/8] Fix LangChainHelper: add get_ai_usage_from_response, usage_metadata check, and Bedrock provider param --- .../src/ldai_langchain/langchain_helper.py | 37 +++++++++++++++---- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py index 8c0c5907..600dbe41 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 @@ -73,32 +73,53 @@ def create_langchain_model(ai_config: AIConfigKind) -> BaseChatModel: model_name = model_dict.get('name', '') provider = provider_dict.get('name', '') - parameters = model_dict.get('parameters') or {} + parameters = dict(model_dict.get('parameters') or {}) + mapped_provider = LangChainHelper.map_provider(provider) + + # Bedrock requires the foundation provider (e.g. Bedrock:Anthropic) passed in + # parameters separately from model_provider, which is used for LangChain routing. + if mapped_provider == 'bedrock_converse' and 'provider' not in parameters: + parameters['provider'] = provider.removeprefix('bedrock:') return init_chat_model( model_name, - model_provider=LangChainHelper.map_provider(provider), + model_provider=mapped_provider, **parameters, ) @staticmethod - def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: + def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: """ - Extract LaunchDarkly AI metrics from a LangChain response. + Extract token usage from a LangChain response. :param response: The response from a LangChain model (BaseMessage or similar) - :return: LDAIMetrics with success status and token usage + :return: TokenUsage or None if unavailable """ - usage: Optional[TokenUsage] = None + 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), + ) if hasattr(response, 'response_metadata') and response.response_metadata: token_usage = ( response.response_metadata.get('tokenUsage') or response.response_metadata.get('token_usage') ) if token_usage: - usage = TokenUsage( + return TokenUsage( total=token_usage.get('totalTokens', 0) or token_usage.get('total_tokens', 0), input=token_usage.get('promptTokens', 0) or token_usage.get('prompt_tokens', 0), output=token_usage.get('completionTokens', 0) or token_usage.get('completion_tokens', 0), ) - return LDAIMetrics(success=True, usage=usage) + return None + + @staticmethod + def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: + """ + Extract LaunchDarkly AI metrics from a LangChain response. + + :param response: The response from a LangChain model (BaseMessage or similar) + :return: LDAIMetrics with success status and token usage + """ + return LDAIMetrics(success=True, usage=LangChainHelper.get_ai_usage_from_response(response)) From bdc2aada6cfa21b59b37faa1e5a4d8c40a6211a7 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 25 Mar 2026 13:35:58 -0500 Subject: [PATCH 5/8] fix lint errors --- .../ldai_langchain/langchain_model_runner.py | 39 ++++++++----------- .../langchain_runner_factory.py | 1 + .../src/ldai_openai/openai_model_runner.py | 1 + .../src/ldai_openai/openai_runner_factory.py | 3 +- packages/sdk/server-ai/src/ldai/__init__.py | 5 +-- packages/sdk/server-ai/src/ldai/client.py | 2 +- .../src/ldai/providers/runner_factory.py | 5 ++- 7 files changed, 26 insertions(+), 30 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py index 5986b05c..e1208ad9 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py @@ -5,6 +5,7 @@ from ldai import LDMessage, log from ldai.providers.model_runner import ModelRunner from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse + from ldai_langchain.langchain_helper import LangChainHelper @@ -72,6 +73,11 @@ async def invoke_structured_model( :param response_structure: Dictionary defining the output structure :return: StructuredResponse containing the structured data """ + structured_response = StructuredResponse( + data={}, + raw_response='', + metrics=LDAIMetrics(success=False, usage=None), + ) try: langchain_messages = LangChainHelper.convert_messages(messages) structured_llm = self._llm.with_structured_output(response_structure, include_raw=True) @@ -79,34 +85,21 @@ async def invoke_structured_model( if not isinstance(response, dict): log.warning(f'Structured output did not return a dict. Got: {type(response)}') - return StructuredResponse( - data={}, - raw_response='', - metrics=LDAIMetrics(success=False, usage=None), - ) + return structured_response raw_response = response.get('raw') - usage = LangChainHelper.get_ai_usage_from_response(raw_response) if raw_response is not None else None - raw_content = raw_response.content if hasattr(raw_response, 'content') else '' + if raw_response is not None: + if hasattr(raw_response, 'content'): + structured_response.raw_response = raw_response.content + structured_response.metrics.usage = LangChainHelper.get_ai_usage_from_response(raw_response) if response.get('parsing_error'): log.warning('LangChain structured model invocation had a parsing error') - return StructuredResponse( - data={}, - raw_response=raw_content, - metrics=LDAIMetrics(success=False, usage=usage), - ) + return structured_response - return StructuredResponse( - data=response.get('parsed') or {}, - raw_response=raw_content, - metrics=LDAIMetrics(success=True, usage=usage), - ) + structured_response.metrics.success = True + structured_response.data = response.get('parsed') or {} + return structured_response except Exception as error: log.warning(f'LangChain structured model invocation failed: {error}') - return StructuredResponse( - data={}, - raw_response='', - metrics=LDAIMetrics(success=False, usage=None), - ) - + return structured_response diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_runner_factory.py index 700d1887..bf5383d5 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,5 +1,6 @@ from ldai.models import AIConfigKind from ldai.providers import AIProvider + from ldai_langchain.langchain_helper import LangChainHelper from ldai_langchain.langchain_model_runner import LangChainModelRunner diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py index 871436cf..9f6d5dfe 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py @@ -6,6 +6,7 @@ from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse from ldai.tracker import TokenUsage from openai import AsyncOpenAI + from ldai_openai.openai_helper import OpenAIHelper 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 0e542087..d80fc01d 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 @@ -3,9 +3,10 @@ from ldai.models import AIConfigKind from ldai.providers import AIProvider -from ldai_openai.openai_model_runner import OpenAIModelRunner from openai import AsyncOpenAI +from ldai_openai.openai_model_runner import OpenAIModelRunner + class OpenAIRunnerFactory(AIProvider): """OpenAI ``AIProvider`` implementation for the LaunchDarkly AI SDK.""" diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 718d6921..e8ef83b4 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -3,10 +3,10 @@ from ldclient import log from ldai.agent_graph import AgentGraphDefinition -from ldai.managed_model import ManagedModel from ldai.chat import Chat # Deprecated — use ManagedModel from ldai.client import LDAIClient from ldai.judge import Judge +from ldai.managed_model import ManagedModel from ldai.models import ( # Deprecated aliases for backward compatibility AIAgentConfig, AIAgentConfigDefault, @@ -45,8 +45,6 @@ 'AIJudgeConfigDefault', 'ManagedModel', 'EvalScore', - # Deprecated — use ManagedModel - 'Chat', 'AgentGraphDefinition', 'Judge', 'JudgeConfiguration', @@ -57,6 +55,7 @@ 'log', # Deprecated exports 'AIConfig', + 'Chat', 'LDAIAgent', 'LDAIAgentConfig', 'LDAIAgentDefaults', diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 9f4e86c4..358f9eb5 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -6,8 +6,8 @@ from ldai import log from ldai.agent_graph import AgentGraphDefinition -from ldai.managed_model import ManagedModel from ldai.judge import Judge +from ldai.managed_model import ManagedModel from ldai.models import ( AIAgentConfig, AIAgentConfigDefault, diff --git a/packages/sdk/server-ai/src/ldai/providers/runner_factory.py b/packages/sdk/server-ai/src/ldai/providers/runner_factory.py index adbcb256..3612baca 100644 --- a/packages/sdk/server-ai/src/ldai/providers/runner_factory.py +++ b/packages/sdk/server-ai/src/ldai/providers/runner_factory.py @@ -4,6 +4,7 @@ from ldai import log from ldai.models import AIConfigKind from ldai.providers.ai_provider import AIProvider +from ldai.providers.model_runner import ModelRunner T = TypeVar('T') @@ -115,13 +116,13 @@ def _get_providers_to_try( def create_model( config: AIConfigKind, default_ai_provider: Optional[str] = None, - ) -> Optional[AIProvider]: + ) -> Optional[ModelRunner]: """ Create a model executor for the given AI completion config. :param config: LaunchDarkly AI config (completion or judge) :param default_ai_provider: Optional provider override ('openai', 'langchain', …) - :return: Configured AIProvider that can invoke_model(), or None + :return: Configured ModelRunner ready to invoke the model, or None """ provider_name = config.provider.name.lower() if config.provider else None providers = RunnerFactory._get_providers_to_try(default_ai_provider, provider_name) From c2f7475cee8485b0bd524aeb6a9ef31485a781b1 Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Wed, 25 Mar 2026 13:47:00 -0500 Subject: [PATCH 6/8] fix: restore structured output pattern, rename convert_messages, remove unused import - Restore mutable structured_response pattern from f951dac so raw_response and usage are preserved even when an exception occurs after partial extraction - Rename LangChainHelper.convert_messages -> convert_messages_to_langchain - Rename OpenAIHelper.convert_messages -> convert_messages_to_openai - Remove unused TokenUsage import from openai_model_runner.py - Fix runner_factory.create_model return type to Optional[ModelRunner] - Fix isort in langchain/openai runner and helper files - Fix isort in ldai/__init__.py and client.py Co-Authored-By: Claude Sonnet 4.6 --- .../src/ldai_langchain/langchain_helper.py | 2 +- .../src/ldai_langchain/langchain_model_runner.py | 4 ++-- .../tests/test_langchain_provider.py | 14 +++++++------- .../src/ldai_openai/openai_helper.py | 2 +- .../src/ldai_openai/openai_model_runner.py | 5 ++--- 5 files changed, 13 insertions(+), 14 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py index 600dbe41..143a0980 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 @@ -35,7 +35,7 @@ def map_provider(ld_provider_name: str) -> str: return mapping.get(lowercased_name, lowercased_name) @staticmethod - def convert_messages( + def convert_messages_to_langchain( messages: List[LDMessage], ) -> List[Union[HumanMessage, SystemMessage, AIMessage]]: """ diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py index e1208ad9..370c72e2 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py @@ -36,7 +36,7 @@ async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse: :return: ModelResponse containing the model's response and metrics """ try: - langchain_messages = LangChainHelper.convert_messages(messages) + langchain_messages = LangChainHelper.convert_messages_to_langchain(messages) response: BaseMessage = await self._llm.ainvoke(langchain_messages) metrics = LangChainHelper.get_ai_metrics_from_response(response) @@ -79,7 +79,7 @@ async def invoke_structured_model( metrics=LDAIMetrics(success=False, usage=None), ) try: - langchain_messages = LangChainHelper.convert_messages(messages) + langchain_messages = LangChainHelper.convert_messages_to_langchain(messages) structured_llm = self._llm.with_structured_output(response_structure, include_raw=True) response = await structured_llm.ainvoke(langchain_messages) 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 611b59fc..19682e41 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 @@ -11,12 +11,12 @@ class TestConvertMessages: - """Tests for LangChainHelper.convert_messages.""" + """Tests for LangChainHelper.convert_messages_to_langchain.""" def test_converts_system_messages_to_system_message(self): """Should convert system messages to SystemMessage.""" messages = [LDMessage(role='system', content='You are a helpful assistant.')] - result = LangChainHelper.convert_messages(messages) + result = LangChainHelper.convert_messages_to_langchain(messages) assert len(result) == 1 assert isinstance(result[0], SystemMessage) @@ -25,7 +25,7 @@ def test_converts_system_messages_to_system_message(self): def test_converts_user_messages_to_human_message(self): """Should convert user messages to HumanMessage.""" messages = [LDMessage(role='user', content='Hello, how are you?')] - result = LangChainHelper.convert_messages(messages) + result = LangChainHelper.convert_messages_to_langchain(messages) assert len(result) == 1 assert isinstance(result[0], HumanMessage) @@ -34,7 +34,7 @@ def test_converts_user_messages_to_human_message(self): def test_converts_assistant_messages_to_ai_message(self): """Should convert assistant messages to AIMessage.""" messages = [LDMessage(role='assistant', content='I am doing well, thank you!')] - result = LangChainHelper.convert_messages(messages) + result = LangChainHelper.convert_messages_to_langchain(messages) assert len(result) == 1 assert isinstance(result[0], AIMessage) @@ -47,7 +47,7 @@ def test_converts_multiple_messages_in_order(self): LDMessage(role='user', content='What is the weather like?'), LDMessage(role='assistant', content='I cannot check the weather.'), ] - result = LangChainHelper.convert_messages(messages) + result = LangChainHelper.convert_messages_to_langchain(messages) assert len(result) == 3 assert isinstance(result[0], SystemMessage) @@ -61,11 +61,11 @@ class MockMessage: content = 'Test message' with pytest.raises(ValueError, match='Unsupported message role: unknown'): - LangChainHelper.convert_messages([MockMessage()]) # type: ignore + LangChainHelper.convert_messages_to_langchain([MockMessage()]) # type: ignore def test_handles_empty_message_array(self): """Should handle empty message array.""" - result = LangChainHelper.convert_messages([]) + result = LangChainHelper.convert_messages_to_langchain([]) assert len(result) == 0 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 4e61de8f..96334bb5 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 @@ -14,7 +14,7 @@ class OpenAIHelper: """ @staticmethod - def convert_messages(messages: List[LDMessage]) -> Iterable[ChatCompletionMessageParam]: + def convert_messages_to_openai(messages: List[LDMessage]) -> Iterable[ChatCompletionMessageParam]: """ Convert LaunchDarkly messages to OpenAI chat completion message format. diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py index 9f6d5dfe..47cc5c15 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py @@ -4,7 +4,6 @@ from ldai import LDMessage, log from ldai.providers.model_runner import ModelRunner from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse -from ldai.tracker import TokenUsage from openai import AsyncOpenAI from ldai_openai.openai_helper import OpenAIHelper @@ -38,7 +37,7 @@ async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse: try: response = await self._client.chat.completions.create( model=self._model_name, - messages=OpenAIHelper.convert_messages(messages), + messages=OpenAIHelper.convert_messages_to_openai(messages), **self._parameters, ) @@ -80,7 +79,7 @@ async def invoke_structured_model( try: response = await self._client.chat.completions.create( model=self._model_name, - messages=OpenAIHelper.convert_messages(messages), + messages=OpenAIHelper.convert_messages_to_openai(messages), response_format={ # type: ignore[arg-type] 'type': 'json_schema', 'json_schema': { From 4cd5800058324cfcb19d8651feae4f587f2d2a4c Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Thu, 26 Mar 2026 08:50:47 -0500 Subject: [PATCH 7/8] refactor!: convert LangChainHelper and OpenAIHelper to pure module-level functions Co-Authored-By: Claude Sonnet 4.6 --- .../src/ldai_langchain/__init__.py | 14 +- .../src/ldai_langchain/langchain_helper.py | 206 +++++++++--------- .../ldai_langchain/langchain_model_runner.py | 10 +- .../langchain_runner_factory.py | 4 +- .../tests/test_langchain_provider.py | 48 ++-- .../src/ldai_openai/__init__.py | 6 +- .../src/ldai_openai/openai_helper.py | 52 ++--- .../src/ldai_openai/openai_model_runner.py | 10 +- .../tests/test_openai_provider.py | 10 +- 9 files changed, 178 insertions(+), 182 deletions(-) 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 40e1d967..2b88026c 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,4 +1,10 @@ -from ldai_langchain.langchain_helper import LangChainHelper +from ldai_langchain.langchain_helper import ( + convert_messages_to_langchain, + create_langchain_model, + get_ai_metrics_from_response, + get_ai_usage_from_response, + map_provider, +) from ldai_langchain.langchain_model_runner import LangChainModelRunner from ldai_langchain.langchain_runner_factory import LangChainRunnerFactory @@ -7,6 +13,10 @@ __all__ = [ '__version__', 'LangChainRunnerFactory', - 'LangChainHelper', 'LangChainModelRunner', + 'convert_messages_to_langchain', + 'create_langchain_model', + 'get_ai_metrics_from_response', + 'get_ai_usage_from_response', + 'map_provider', ] 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 143a0980..5061a1b2 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 @@ -8,118 +8,110 @@ from ldai.tracker import TokenUsage -class LangChainHelper: +def map_provider(ld_provider_name: str) -> str: """ - Shared utilities for LangChain-based runners (model, agent, agent graph). + Map a LaunchDarkly provider name to its LangChain equivalent. - All methods are static — this class is a namespace, not meant to be instantiated. + :param ld_provider_name: LaunchDarkly provider name + :return: LangChain-compatible provider name """ + lowercased_name = ld_provider_name.lower() + # Bedrock is the only provider that uses "provider:model_family" (e.g. Bedrock:Anthropic). + if lowercased_name.startswith('bedrock:'): + return 'bedrock_converse' - @staticmethod - def map_provider(ld_provider_name: str) -> str: - """ - Map a LaunchDarkly provider name to its LangChain equivalent. - - :param ld_provider_name: LaunchDarkly provider name - :return: LangChain-compatible provider name - """ - lowercased_name = ld_provider_name.lower() - # Bedrock is the only provider that uses "provider:model_family" (e.g. Bedrock:Anthropic). - if lowercased_name.startswith('bedrock:'): - return 'bedrock_converse' - - mapping: Dict[str, str] = { - 'gemini': 'google-genai', - 'bedrock': 'bedrock_converse', - } - return mapping.get(lowercased_name, lowercased_name) - - @staticmethod - def convert_messages_to_langchain( - messages: List[LDMessage], - ) -> List[Union[HumanMessage, SystemMessage, AIMessage]]: - """ - Convert LaunchDarkly messages to LangChain message objects. - - :param messages: List of LDMessage objects - :return: List of LangChain message objects - :raises ValueError: If an unsupported message role is encountered - """ - result: List[Union[HumanMessage, SystemMessage, AIMessage]] = [] - for msg in messages: - if msg.role == 'system': - result.append(SystemMessage(content=msg.content)) - elif msg.role == 'user': - result.append(HumanMessage(content=msg.content)) - elif msg.role == 'assistant': - result.append(AIMessage(content=msg.content)) - else: - raise ValueError(f'Unsupported message role: {msg.role}') - return result - - @staticmethod - def create_langchain_model(ai_config: AIConfigKind) -> BaseChatModel: - """ - Create a LangChain BaseChatModel from a LaunchDarkly AI configuration. - - :param ai_config: The LaunchDarkly AI configuration - :return: A configured LangChain BaseChatModel - """ - from langchain.chat_models import init_chat_model - - config_dict = ai_config.to_dict() - model_dict = config_dict.get('model') or {} - provider_dict = config_dict.get('provider') or {} - - model_name = model_dict.get('name', '') - provider = provider_dict.get('name', '') - parameters = dict(model_dict.get('parameters') or {}) - mapped_provider = LangChainHelper.map_provider(provider) - - # Bedrock requires the foundation provider (e.g. Bedrock:Anthropic) passed in - # parameters separately from model_provider, which is used for LangChain routing. - if mapped_provider == 'bedrock_converse' and 'provider' not in parameters: - parameters['provider'] = provider.removeprefix('bedrock:') - - return init_chat_model( - model_name, - model_provider=mapped_provider, - **parameters, - ) + mapping: Dict[str, str] = { + 'gemini': 'google-genai', + 'bedrock': 'bedrock_converse', + } + return mapping.get(lowercased_name, lowercased_name) + + +def convert_messages_to_langchain( + messages: List[LDMessage], +) -> List[Union[HumanMessage, SystemMessage, AIMessage]]: + """ + Convert LaunchDarkly messages to LangChain message objects. + + :param messages: List of LDMessage objects + :return: List of LangChain message objects + :raises ValueError: If an unsupported message role is encountered + """ + result: List[Union[HumanMessage, SystemMessage, AIMessage]] = [] + for msg in messages: + if msg.role == 'system': + result.append(SystemMessage(content=msg.content)) + elif msg.role == 'user': + result.append(HumanMessage(content=msg.content)) + elif msg.role == 'assistant': + result.append(AIMessage(content=msg.content)) + else: + raise ValueError(f'Unsupported message role: {msg.role}') + return result + + +def create_langchain_model(ai_config: AIConfigKind) -> BaseChatModel: + """ + Create a LangChain BaseChatModel from a LaunchDarkly AI configuration. + + :param ai_config: The LaunchDarkly AI configuration + :return: A configured LangChain BaseChatModel + """ + from langchain.chat_models import init_chat_model + + config_dict = ai_config.to_dict() + model_dict = config_dict.get('model') or {} + provider_dict = config_dict.get('provider') or {} + + model_name = model_dict.get('name', '') + provider = provider_dict.get('name', '') + parameters = dict(model_dict.get('parameters') or {}) + mapped_provider = map_provider(provider) + + # Bedrock requires the foundation provider (e.g. Bedrock:Anthropic) passed in + # parameters separately from model_provider, which is used for LangChain routing. + if mapped_provider == 'bedrock_converse' and 'provider' not in parameters: + parameters['provider'] = provider.removeprefix('bedrock:') - @staticmethod - def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: - """ - Extract token usage from a LangChain response. + return init_chat_model( + model_name, + model_provider=mapped_provider, + **parameters, + ) - :param response: The response from a LangChain model (BaseMessage or similar) - :return: TokenUsage or None if unavailable - """ - if hasattr(response, 'usage_metadata') and response.usage_metadata: + +def get_ai_usage_from_response(response: Any) -> Optional[TokenUsage]: + """ + Extract token usage from a LangChain response. + + :param response: The response from a LangChain model (BaseMessage or similar) + :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), + ) + if hasattr(response, 'response_metadata') and response.response_metadata: + token_usage = ( + response.response_metadata.get('tokenUsage') + or response.response_metadata.get('token_usage') + ) + if token_usage: 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=token_usage.get('totalTokens', 0) or token_usage.get('total_tokens', 0), + input=token_usage.get('promptTokens', 0) or token_usage.get('prompt_tokens', 0), + output=token_usage.get('completionTokens', 0) or token_usage.get('completion_tokens', 0), ) - if hasattr(response, 'response_metadata') and response.response_metadata: - token_usage = ( - response.response_metadata.get('tokenUsage') - or response.response_metadata.get('token_usage') - ) - if token_usage: - return TokenUsage( - total=token_usage.get('totalTokens', 0) or token_usage.get('total_tokens', 0), - input=token_usage.get('promptTokens', 0) or token_usage.get('prompt_tokens', 0), - output=token_usage.get('completionTokens', 0) or token_usage.get('completion_tokens', 0), - ) - return None - - @staticmethod - def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: - """ - Extract LaunchDarkly AI metrics from a LangChain response. - - :param response: The response from a LangChain model (BaseMessage or similar) - :return: LDAIMetrics with success status and token usage - """ - return LDAIMetrics(success=True, usage=LangChainHelper.get_ai_usage_from_response(response)) + return None + + +def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: + """ + Extract LaunchDarkly AI metrics from a LangChain response. + + :param response: The response from a LangChain model (BaseMessage or similar) + :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-langchain/src/ldai_langchain/langchain_model_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py index 370c72e2..a933fc4a 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py @@ -6,7 +6,7 @@ from ldai.providers.model_runner import ModelRunner from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse -from ldai_langchain.langchain_helper import LangChainHelper +from ldai_langchain.langchain_helper import convert_messages_to_langchain, get_ai_metrics_from_response, get_ai_usage_from_response class LangChainModelRunner(ModelRunner): @@ -36,9 +36,9 @@ async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse: :return: ModelResponse containing the model's response and metrics """ try: - langchain_messages = LangChainHelper.convert_messages_to_langchain(messages) + langchain_messages = convert_messages_to_langchain(messages) response: BaseMessage = await self._llm.ainvoke(langchain_messages) - metrics = LangChainHelper.get_ai_metrics_from_response(response) + metrics = get_ai_metrics_from_response(response) content: str = '' if isinstance(response.content, str): @@ -79,7 +79,7 @@ async def invoke_structured_model( metrics=LDAIMetrics(success=False, usage=None), ) try: - langchain_messages = LangChainHelper.convert_messages_to_langchain(messages) + langchain_messages = convert_messages_to_langchain(messages) structured_llm = self._llm.with_structured_output(response_structure, include_raw=True) response = await structured_llm.ainvoke(langchain_messages) @@ -91,7 +91,7 @@ async def invoke_structured_model( if raw_response is not None: if hasattr(raw_response, 'content'): structured_response.raw_response = raw_response.content - structured_response.metrics.usage = LangChainHelper.get_ai_usage_from_response(raw_response) + structured_response.metrics.usage = get_ai_usage_from_response(raw_response) if response.get('parsing_error'): log.warning('LangChain structured model invocation had a parsing error') 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 bf5383d5..402e2951 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,7 @@ from ldai.models import AIConfigKind from ldai.providers import AIProvider -from ldai_langchain.langchain_helper import LangChainHelper +from ldai_langchain.langchain_helper import create_langchain_model from ldai_langchain.langchain_model_runner import LangChainModelRunner @@ -15,5 +15,5 @@ def create_model(self, config: AIConfigKind) -> LangChainModelRunner: :param config: The LaunchDarkly AI configuration :return: LangChainModelRunner ready to invoke the model """ - llm = LangChainHelper.create_langchain_model(config) + llm = create_langchain_model(config) return LangChainModelRunner(llm) 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 19682e41..b78fde84 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 @@ -7,16 +7,16 @@ from ldai import LDMessage -from ldai_langchain import LangChainHelper, LangChainModelRunner, LangChainRunnerFactory +from ldai_langchain import LangChainModelRunner, LangChainRunnerFactory, convert_messages_to_langchain, get_ai_metrics_from_response, map_provider class TestConvertMessages: - """Tests for LangChainHelper.convert_messages_to_langchain.""" + """Tests for convert_messages_to_langchain.""" def test_converts_system_messages_to_system_message(self): """Should convert system messages to SystemMessage.""" messages = [LDMessage(role='system', content='You are a helpful assistant.')] - result = LangChainHelper.convert_messages_to_langchain(messages) + result = convert_messages_to_langchain(messages) assert len(result) == 1 assert isinstance(result[0], SystemMessage) @@ -25,7 +25,7 @@ def test_converts_system_messages_to_system_message(self): def test_converts_user_messages_to_human_message(self): """Should convert user messages to HumanMessage.""" messages = [LDMessage(role='user', content='Hello, how are you?')] - result = LangChainHelper.convert_messages_to_langchain(messages) + result = convert_messages_to_langchain(messages) assert len(result) == 1 assert isinstance(result[0], HumanMessage) @@ -34,7 +34,7 @@ def test_converts_user_messages_to_human_message(self): def test_converts_assistant_messages_to_ai_message(self): """Should convert assistant messages to AIMessage.""" messages = [LDMessage(role='assistant', content='I am doing well, thank you!')] - result = LangChainHelper.convert_messages_to_langchain(messages) + result = convert_messages_to_langchain(messages) assert len(result) == 1 assert isinstance(result[0], AIMessage) @@ -47,7 +47,7 @@ def test_converts_multiple_messages_in_order(self): LDMessage(role='user', content='What is the weather like?'), LDMessage(role='assistant', content='I cannot check the weather.'), ] - result = LangChainHelper.convert_messages_to_langchain(messages) + result = convert_messages_to_langchain(messages) assert len(result) == 3 assert isinstance(result[0], SystemMessage) @@ -61,16 +61,16 @@ class MockMessage: content = 'Test message' with pytest.raises(ValueError, match='Unsupported message role: unknown'): - LangChainHelper.convert_messages_to_langchain([MockMessage()]) # type: ignore + convert_messages_to_langchain([MockMessage()]) # type: ignore def test_handles_empty_message_array(self): """Should handle empty message array.""" - result = LangChainHelper.convert_messages_to_langchain([]) + result = convert_messages_to_langchain([]) assert len(result) == 0 class TestGetAIMetricsFromResponse: - """Tests for LangChainHelper.get_ai_metrics_from_response.""" + """Tests for get_ai_metrics_from_response.""" def test_creates_metrics_with_success_true_and_token_usage(self): """Should create metrics with success=True and token usage.""" @@ -83,7 +83,7 @@ def test_creates_metrics_with_success_true_and_token_usage(self): }, } - result = LangChainHelper.get_ai_metrics_from_response(mock_response) + result = get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is not None @@ -102,7 +102,7 @@ def test_creates_metrics_with_snake_case_token_usage(self): }, } - result = LangChainHelper.get_ai_metrics_from_response(mock_response) + result = get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is not None @@ -114,34 +114,34 @@ def test_creates_metrics_with_success_true_and_no_usage_when_metadata_missing(se """Should create metrics with success=True and no usage when metadata is missing.""" mock_response = AIMessage(content='Test response') - result = LangChainHelper.get_ai_metrics_from_response(mock_response) + result = get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is None class TestMapProvider: - """Tests for LangChainHelper.map_provider.""" + """Tests for map_provider.""" def test_maps_gemini_to_google_genai(self): """Should map gemini to google-genai.""" - assert LangChainHelper.map_provider('gemini') == 'google-genai' - assert LangChainHelper.map_provider('Gemini') == 'google-genai' - assert LangChainHelper.map_provider('GEMINI') == 'google-genai' + assert map_provider('gemini') == 'google-genai' + assert map_provider('Gemini') == 'google-genai' + assert map_provider('GEMINI') == 'google-genai' def test_maps_bedrock_and_model_families_to_bedrock_converse(self): """Should map bedrock and bedrock:model_family to bedrock_converse.""" - assert LangChainHelper.map_provider('bedrock') == 'bedrock_converse' - assert LangChainHelper.map_provider('Bedrock:Anthropic') == 'bedrock_converse' - assert LangChainHelper.map_provider('bedrock:anthropic') == 'bedrock_converse' - assert LangChainHelper.map_provider('bedrock:amazon') == 'bedrock_converse' - assert LangChainHelper.map_provider('bedrock:cohere') == 'bedrock_converse' + assert map_provider('bedrock') == 'bedrock_converse' + assert map_provider('Bedrock:Anthropic') == 'bedrock_converse' + assert map_provider('bedrock:anthropic') == 'bedrock_converse' + assert map_provider('bedrock:amazon') == 'bedrock_converse' + assert map_provider('bedrock:cohere') == 'bedrock_converse' def test_returns_provider_name_unchanged_for_unmapped_providers(self): """Should return provider name unchanged for unmapped providers.""" - assert LangChainHelper.map_provider('openai') == 'openai' - assert LangChainHelper.map_provider('anthropic') == 'anthropic' - assert LangChainHelper.map_provider('unknown') == 'unknown' + assert map_provider('openai') == 'openai' + assert map_provider('anthropic') == 'anthropic' + assert map_provider('unknown') == 'unknown' class TestInvokeModel: 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 59e0b313..ee982af7 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,9 +1,11 @@ -from ldai_openai.openai_helper import OpenAIHelper +from ldai_openai.openai_helper import convert_messages_to_openai, get_ai_metrics_from_response, get_ai_usage_from_response from ldai_openai.openai_model_runner import OpenAIModelRunner from ldai_openai.openai_runner_factory import OpenAIRunnerFactory __all__ = [ 'OpenAIRunnerFactory', - 'OpenAIHelper', 'OpenAIModelRunner', + '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_helper.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_helper.py index 96334bb5..0e5c0848 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 @@ -6,39 +6,31 @@ from openai.types.chat import ChatCompletionMessageParam -class OpenAIHelper: +def convert_messages_to_openai(messages: List[LDMessage]) -> Iterable[ChatCompletionMessageParam]: """ - Shared utilities for OpenAI-based runners (model, agent, agent graph). + Convert LaunchDarkly messages to OpenAI chat completion message format. - All methods are static — this class is a namespace, not meant to be instantiated. + :param messages: List of LDMessage objects + :return: Iterable of OpenAI ChatCompletionMessageParam dicts """ + return cast( + Iterable[ChatCompletionMessageParam], + [{'role': msg.role, 'content': msg.content} for msg in messages], + ) - @staticmethod - def convert_messages_to_openai(messages: List[LDMessage]) -> Iterable[ChatCompletionMessageParam]: - """ - Convert LaunchDarkly messages to OpenAI chat completion message format. - :param messages: List of LDMessage objects - :return: Iterable of OpenAI ChatCompletionMessageParam dicts - """ - return cast( - Iterable[ChatCompletionMessageParam], - [{'role': msg.role, 'content': msg.content} for msg in messages], - ) - - @staticmethod - def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: - """ - Extract LaunchDarkly AI metrics from an OpenAI response. +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 - :return: LDAIMetrics with success status and token usage - """ - usage: Optional[TokenUsage] = None - if hasattr(response, 'usage') and response.usage: - usage = TokenUsage( - total=response.usage.total_tokens or 0, - input=response.usage.prompt_tokens or 0, - output=response.usage.completion_tokens or 0, - ) - return LDAIMetrics(success=True, usage=usage) + :param response: The response from the OpenAI chat completions API + :return: LDAIMetrics with success status and token usage + """ + usage: Optional[TokenUsage] = None + if hasattr(response, 'usage') and response.usage: + usage = TokenUsage( + total=response.usage.total_tokens or 0, + input=response.usage.prompt_tokens or 0, + output=response.usage.completion_tokens or 0, + ) + return LDAIMetrics(success=True, usage=usage) diff --git a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py index 47cc5c15..30e0a5c7 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py @@ -6,7 +6,7 @@ from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse from openai import AsyncOpenAI -from ldai_openai.openai_helper import OpenAIHelper +from ldai_openai.openai_helper import convert_messages_to_openai, get_ai_metrics_from_response class OpenAIModelRunner(ModelRunner): @@ -37,11 +37,11 @@ async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse: try: response = await self._client.chat.completions.create( model=self._model_name, - messages=OpenAIHelper.convert_messages_to_openai(messages), + messages=convert_messages_to_openai(messages), **self._parameters, ) - metrics = OpenAIHelper.get_ai_metrics_from_response(response) + metrics = get_ai_metrics_from_response(response) content = '' if response.choices and len(response.choices) > 0: @@ -79,7 +79,7 @@ async def invoke_structured_model( try: response = await self._client.chat.completions.create( model=self._model_name, - messages=OpenAIHelper.convert_messages_to_openai(messages), + messages=convert_messages_to_openai(messages), response_format={ # type: ignore[arg-type] 'type': 'json_schema', 'json_schema': { @@ -91,7 +91,7 @@ async def invoke_structured_model( **self._parameters, ) - metrics = OpenAIHelper.get_ai_metrics_from_response(response) + metrics = get_ai_metrics_from_response(response) content = '' if response.choices and len(response.choices) > 0: 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 595fe9fb..3a9de4dc 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,11 +5,11 @@ from ldai import LDMessage -from ldai_openai import OpenAIHelper, OpenAIModelRunner, OpenAIRunnerFactory +from ldai_openai import OpenAIModelRunner, OpenAIRunnerFactory, get_ai_metrics_from_response class TestGetAIMetricsFromResponse: - """Tests for OpenAIHelper.get_ai_metrics_from_response.""" + """Tests for get_ai_metrics_from_response.""" def test_creates_metrics_with_success_true_and_token_usage(self): """Should create metrics with success=True and token usage.""" @@ -19,7 +19,7 @@ def test_creates_metrics_with_success_true_and_token_usage(self): mock_response.usage.completion_tokens = 50 mock_response.usage.total_tokens = 100 - result = OpenAIHelper.get_ai_metrics_from_response(mock_response) + result = get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is not None @@ -32,7 +32,7 @@ def test_creates_metrics_with_success_true_and_no_usage_when_usage_missing(self) mock_response = MagicMock() mock_response.usage = None - result = OpenAIHelper.get_ai_metrics_from_response(mock_response) + result = get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is None @@ -45,7 +45,7 @@ def test_handles_partial_usage_data(self): mock_response.usage.completion_tokens = None mock_response.usage.total_tokens = None - result = OpenAIHelper.get_ai_metrics_from_response(mock_response) + result = get_ai_metrics_from_response(mock_response) assert result.success is True assert result.usage is not None From 9285ffd1159da43f2a6f46b7c3b70fcf9ffa84fc Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Thu, 26 Mar 2026 08:58:47 -0500 Subject: [PATCH 8/8] fix: add get_ai_usage_from_response to openai_helper and fix import sort Co-Authored-By: Claude Sonnet 4.6 --- .../ldai_langchain/langchain_model_runner.py | 6 ++++- .../src/ldai_openai/__init__.py | 6 ++++- .../src/ldai_openai/openai_helper.py | 26 +++++++++++++------ .../src/ldai_openai/openai_model_runner.py | 5 +++- 4 files changed, 32 insertions(+), 11 deletions(-) diff --git a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py index a933fc4a..d504030b 100644 --- a/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py @@ -6,7 +6,11 @@ from ldai.providers.model_runner import ModelRunner from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse -from ldai_langchain.langchain_helper import convert_messages_to_langchain, get_ai_metrics_from_response, get_ai_usage_from_response +from ldai_langchain.langchain_helper import ( + convert_messages_to_langchain, + get_ai_metrics_from_response, + get_ai_usage_from_response, +) class LangChainModelRunner(ModelRunner): 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 ee982af7..8a8199b3 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,8 @@ -from ldai_openai.openai_helper import convert_messages_to_openai, get_ai_metrics_from_response, get_ai_usage_from_response +from ldai_openai.openai_helper import ( + convert_messages_to_openai, + get_ai_metrics_from_response, + get_ai_usage_from_response, +) from ldai_openai.openai_model_runner import OpenAIModelRunner from ldai_openai.openai_runner_factory import OpenAIRunnerFactory 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 0e5c0848..3cc41e48 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 @@ -19,6 +19,23 @@ def convert_messages_to_openai(messages: List[LDMessage]) -> Iterable[ChatComple ) +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 + :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, + ) + return None + + def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: """ Extract LaunchDarkly AI metrics from an OpenAI response. @@ -26,11 +43,4 @@ def get_ai_metrics_from_response(response: Any) -> LDAIMetrics: :param response: The response from the OpenAI chat completions API :return: LDAIMetrics with success status and token usage """ - usage: Optional[TokenUsage] = None - if hasattr(response, 'usage') and response.usage: - usage = TokenUsage( - total=response.usage.total_tokens or 0, - input=response.usage.prompt_tokens or 0, - output=response.usage.completion_tokens or 0, - ) - return LDAIMetrics(success=True, usage=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_model_runner.py b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py index 30e0a5c7..9c4a34d8 100644 --- a/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py +++ b/packages/ai-providers/server-ai-openai/src/ldai_openai/openai_model_runner.py @@ -6,7 +6,10 @@ from ldai.providers.types import LDAIMetrics, ModelResponse, StructuredResponse from openai import AsyncOpenAI -from ldai_openai.openai_helper import convert_messages_to_openai, get_ai_metrics_from_response +from ldai_openai.openai_helper import ( + convert_messages_to_openai, + get_ai_metrics_from_response, +) class OpenAIModelRunner(ModelRunner):