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 ee58f4e..9624606 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,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" @@ -7,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 0000000..afdf007 --- /dev/null +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py @@ -0,0 +1,98 @@ +"""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 + """ + mapping: Dict[str, str] = {'gemini': 'google-genai'} + return mapping.get(ld_provider_name.lower(), ld_provider_name.lower()) + + @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 0000000..683be75 --- /dev/null +++ b/packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_model_runner.py @@ -0,0 +1,103 @@ +"""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.tracker import TokenUsage +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) + 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=TokenUsage(total=0, input=0, output=0)), + ) + + return StructuredResponse( + data=response, + raw_response=str(response), + metrics=LDAIMetrics(success=True, usage=TokenUsage(total=0, input=0, output=0)), + ) + except Exception as error: + log.warning(f'LangChain structured model invocation failed: {error}') + return StructuredResponse( + data={}, + raw_response='', + metrics=LDAIMetrics(success=False, usage=TokenUsage(total=0, input=0, output=0)), + ) + 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 4bfea17..41c8a14 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,244 +1,25 @@ """LangChain connector for 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, 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 connector for the LaunchDarkly AI SDK. - Can be used in two ways: - - Transparently via ExecutorFactory (pass ``default_ai_provider='langchain'`` to - ``create_model()`` / ``create_chat()``). - - Directly for full control: instantiate with a ``BaseChatModel``, then call - ``invoke_model()`` yourself and use the static convenience methods - (``get_ai_metrics_from_response``, ``convert_messages_to_langchain``, - ``map_provider``, ``create_langchain_model``). + Acts as a per-provider factory. Instantiate with no arguments, then call + ``create_model(config)`` to obtain a configured ``LangChainModelRunner``. """ - def __init__(self, llm: Optional[BaseChatModel] = None): - """ - Initialize the LangChain connector. - - When called with no arguments the connector acts as a per-provider factory - — call ``create_model(config)`` to obtain a configured instance. - - When called with an explicit ``llm`` the connector is ready to invoke - the model immediately. - - :param llm: A LangChain BaseChatModel instance (optional) - """ - self._llm = llm - - # --- AIProvider factory methods --- - - def create_model(self, config: AIConfigKind) -> 'LangChainRunnerFactory': + def create_model(self, config: AIConfigKind) -> LangChainModelRunner: """ - Create a configured LangChain model connector 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) - - # --- Model invocation --- - - 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: - 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 - """ - try: - langchain_messages = LangChainRunnerFactory.convert_messages_to_langchain(messages) - structured_llm = self._llm.with_structured_output(response_structure) - 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 StructuredResponse( - data={}, - raw_response='', - metrics=LDAIMetrics( - success=False, - usage=TokenUsage(total=0, input=0, output=0), - ), - ) - - return StructuredResponse( - data=response, - raw_response=str(response), - metrics=LDAIMetrics( - success=True, - usage=TokenUsage(total=0, input=0, output=0), - ), - ) - except Exception as error: - log.warning(f'LangChain structured model invocation failed: {error}') - - return StructuredResponse( - data={}, - raw_response='', - metrics=LDAIMetrics( - success=False, - usage=TokenUsage(total=0, input=0, output=0), - ), - ) - - # --- Convenience accessors --- - - def get_chat_model(self) -> Optional[BaseChatModel]: - """ - Get the underlying LangChain model instance. - - :return: The underlying BaseChatModel, or None if not yet configured + :return: LangChainModelRunner ready to invoke the model """ - 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() - - mapping: Dict[str, str] = { - 'gemini': 'google-genai', - } - - return mapping.get(lowercased_name, lowercased_name) - - @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: 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) - - @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 = model_dict.get('parameters') or {} - - return init_chat_model( - model_name, - model_provider=LangChainRunnerFactory.map_provider(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 d74f8c3..851187e 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,26 +114,26 @@ 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_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: @@ -150,7 +149,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) @@ -163,7 +162,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) @@ -176,7 +175,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) @@ -201,7 +200,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': {}} @@ -217,7 +216,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': {}} @@ -230,14 +229,12 @@ async def test_returns_success_false_when_structured_model_invocation_throws_err assert result.metrics.usage.total == 0 -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 1284f48..51c1c40 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,7 +1,11 @@ """LaunchDarkly AI SDK OpenAI Connector.""" -from ldai_openai.openai_runner_factory import OpenAIProvider +from ldai_openai.openai_helper import OpenAIHelper +from ldai_openai.openai_model_runner import OpenAIModelRunner +from ldai_openai.openai_runner_factory import OpenAIRunnerFactory __all__ = [ - 'OpenAIProvider', + '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 0000000..b868a86 --- /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 0000000..9ccb4d4 --- /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 1313404..086b9f9 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,223 +1,55 @@ """OpenAI connector for LaunchDarkly AI SDK.""" -import json import os -from typing import Any, Dict, Iterable, List, Optional, cast +from typing import Optional -from ldai import LDMessage, log +from ldai import 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 OpenAIProvider(AIProvider): +class OpenAIRunnerFactory(AIProvider): """ OpenAI connector for the LaunchDarkly AI SDK. - Can be used in two ways: - - Transparently via ExecutorFactory (pass ``default_ai_provider='openai'`` to - ``create_model()`` / ``create_chat()``). - - Directly for full control: instantiate with an ``AsyncOpenAI`` client, - model name, and parameters, then call ``invoke_model()`` yourself. + 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. """ - def __init__( - self, - client: Optional[AsyncOpenAI] = None, - model_name: str = '', - parameters: Optional[Dict[str, Any]] = None, - ): + def __init__(self, client: Optional[AsyncOpenAI] = None): """ Initialize the OpenAI connector. - When called with no arguments the connector reads credentials from the - environment (``OPENAI_API_KEY``) and acts as a per-provider factory — - call ``create_model(config)`` to obtain a configured instance. - - When called with explicit arguments the connector is ready to invoke - the model immediately. - :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 {} - # --- AIProvider factory methods --- - - def create_model(self, config: AIConfigKind) -> 'OpenAIProvider': + def create_model(self, config: AIConfigKind) -> OpenAIModelRunner: """ - Create a configured OpenAI model connector 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 OpenAIProvider 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 OpenAIProvider(self._client, model_name, parameters) - - # --- Model invocation --- - - 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 = OpenAIProvider.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 = OpenAIProvider.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), - ) - - # --- Convenience accessors --- + 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), - OpenAIProvider.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 d684df0..595fe9f 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 OpenAIProvider +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 = OpenAIProvider.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 = OpenAIProvider.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 = OpenAIProvider.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 = OpenAIProvider(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 = OpenAIProvider(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 = OpenAIProvider(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 = OpenAIProvider(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 = OpenAIProvider(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 = OpenAIProvider(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 = OpenAIProvider(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 = OpenAIProvider(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 = OpenAIProvider(mock_client, 'gpt-3.5-turbo', {}) + provider = OpenAIRunnerFactory(mock_client) assert provider.get_client() is mock_client @@ -281,7 +281,7 @@ class TestCreateModel: """Tests for create_model instance method.""" def test_creates_connector_with_correct_model_and_parameters(self): - """Should create OpenAIProvider with correct model and parameters.""" + """Should create OpenAIRunnerFactory with correct model and parameters.""" mock_ai_config = MagicMock() mock_ai_config.to_dict.return_value = { 'model': { @@ -298,9 +298,9 @@ def test_creates_connector_with_correct_model_and_parameters(self): mock_client = MagicMock() mock_openai_class.return_value = mock_client - result = OpenAIProvider().create_model(mock_ai_config) + result = OpenAIRunnerFactory().create_model(mock_ai_config) - assert isinstance(result, OpenAIProvider) + assert isinstance(result, OpenAIModelRunner) assert result._model_name == 'gpt-4' assert result._parameters == {'temperature': 0.7, 'max_tokens': 1000} @@ -313,9 +313,8 @@ def test_handles_missing_model_config(self): mock_client = MagicMock() mock_openai_class.return_value = mock_client - result = OpenAIProvider().create_model(mock_ai_config) + result = OpenAIRunnerFactory().create_model(mock_ai_config) - assert isinstance(result, OpenAIProvider) + 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 da2340c..8077b74 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 @@ -27,10 +28,17 @@ ProviderConfig, ) from ldai.providers.types import EvalScore, JudgeResponse +from ldai.runners import AgentGraphRunner, AgentRunner +from ldai.runners.types import AgentGraphResult, AgentResult, ToolRegistry from ldai.tracker import AIGraphTracker __all__ = [ 'LDAIClient', + 'AgentRunner', + 'AgentGraphRunner', + 'AgentResult', + 'AgentGraphResult', + 'ToolRegistry', 'AIAgentConfig', 'AIAgentConfigDefault', 'AIAgentConfigRequest', @@ -42,8 +50,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 c826fed..cc3dbff 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 a87bef0..ba345ed 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 = await RunnerFactory.create_model(config, default_ai_provider) - if not provider: + runner = await 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 0ca402a..7e4c610 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(ai_config.evaluation_metric_key) 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, ) @@ -100,7 +100,7 @@ async def evaluate( async def evaluate_messages( self, messages: list[LDMessage], - response: ChatResponse, + response: ModelResponse, sampling_ratio: float = 1.0, ) -> Optional[JudgeResponse]: """ @@ -132,13 +132,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 0000000..93c3c67 --- /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 4cebeea..14896e1 100644 --- a/packages/sdk/server-ai/src/ldai/providers/__init__.py +++ b/packages/sdk/server-ai/src/ldai/providers/__init__.py @@ -1,9 +1,11 @@ """AI Connector interfaces and factory for LaunchDarkly AI SDK.""" 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 a675eda..fa65e5b 100644 --- a/packages/sdk/server-ai/src/ldai/providers/ai_provider.py +++ b/packages/sdk/server-ai/src/ldai/providers/ai_provider.py @@ -5,7 +5,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,12 +16,9 @@ class AIProvider(ABC): (with no arguments — credentials are read from environment variables) and is responsible for constructing focused runtime capability objects via create_model(), create_agent(), and create_agent_graph(). - - The invoke_model() / invoke_structured_model() methods remain on this base - class for compatibility and will migrate to ModelExecutor in PR 2. """ - 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. @@ -29,14 +26,14 @@ async def invoke_model(self, messages: List[LDMessage]) -> ChatResponse: Connector 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.warn('invoke_model not implemented by this connector') 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), ) @@ -66,14 +63,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 connectors 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.warn('create_model not implemented by this connector') return None @@ -103,4 +100,3 @@ def create_agent_graph(self, graph_def: Any, tools: Any) -> Optional[Any]: """ log.warn('create_agent_graph not implemented by this connector') 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 0000000..6365203 --- /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/runner_factory.py b/packages/sdk/server-ai/src/ldai/providers/runner_factory.py index f2e5fca..9496b0b 100644 --- a/packages/sdk/server-ai/src/ldai/providers/runner_factory.py +++ b/packages/sdk/server-ai/src/ldai/providers/runner_factory.py @@ -42,8 +42,8 @@ def _get_ai_adapter(provider_type: str) -> Optional[AIProvider]: if provider_type == 'openai': RunnerFactory._pkg_exists('ldai_openai') - from ldai_openai import OpenAIProvider - return OpenAIProvider() + from ldai_openai import OpenAIRunnerFactory + return OpenAIRunnerFactory() log.warn( f"Provider '{provider_type}' is not supported. " @@ -178,4 +178,3 @@ def _pkg_exists(package_name: str) -> None: """ if util.find_spec(package_name) is None: raise ImportError(f"Package '{package_name}' not found") - diff --git a/packages/sdk/server-ai/src/ldai/providers/types.py b/packages/sdk/server-ai/src/ldai/providers/types.py index e9160cc..0a07151 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/src/ldai/runners/__init__.py b/packages/sdk/server-ai/src/ldai/runners/__init__.py new file mode 100644 index 0000000..a9a834c --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/runners/__init__.py @@ -0,0 +1,13 @@ +"""Runner ABCs and result types for LaunchDarkly AI SDK.""" + +from ldai.runners.agent_graph_runner import AgentGraphRunner +from ldai.runners.agent_runner import AgentRunner +from ldai.runners.types import AgentGraphResult, AgentResult, ToolRegistry + +__all__ = [ + 'AgentRunner', + 'AgentGraphRunner', + 'AgentResult', + 'AgentGraphResult', + 'ToolRegistry', +] diff --git a/packages/sdk/server-ai/src/ldai/runners/agent_graph_runner.py b/packages/sdk/server-ai/src/ldai/runners/agent_graph_runner.py new file mode 100644 index 0000000..3cfdd80 --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/runners/agent_graph_runner.py @@ -0,0 +1,26 @@ +"""Abstract base class for agent graph runners.""" + +from abc import ABC, abstractmethod +from typing import Any + +from ldai.runners.types import AgentGraphResult + + +class AgentGraphRunner(ABC): + """ + Abstract base class for agent graph runners. + + An AgentGraphRunner encapsulates multi-agent graph execution. + Provider-specific implementations (e.g. OpenAIAgentGraphRunner) are + returned by RunnerFactory.create_agent_graph() and hold all provider + wiring internally. + """ + + @abstractmethod + async def run(self, input: Any) -> AgentGraphResult: + """ + Run the agent graph with the given input. + + :param input: The input to the agent graph (string prompt or structured input) + :return: AgentGraphResult containing the output, raw response, and metrics + """ diff --git a/packages/sdk/server-ai/src/ldai/runners/agent_runner.py b/packages/sdk/server-ai/src/ldai/runners/agent_runner.py new file mode 100644 index 0000000..063198d --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/runners/agent_runner.py @@ -0,0 +1,25 @@ +"""Abstract base class for agent runners.""" + +from abc import ABC, abstractmethod +from typing import Any + +from ldai.runners.types import AgentResult + + +class AgentRunner(ABC): + """ + Abstract base class for single-agent runners. + + An AgentRunner encapsulates the execution of a single AI agent. + Provider-specific implementations (e.g. OpenAIAgentRunner) are returned + by RunnerFactory.create_agent() and hold all provider wiring internally. + """ + + @abstractmethod + async def run(self, input: Any) -> AgentResult: + """ + Run the agent with the given input. + + :param input: The input to the agent (string prompt or structured input) + :return: AgentResult containing the output, raw response, and metrics + """ diff --git a/packages/sdk/server-ai/src/ldai/runners/types.py b/packages/sdk/server-ai/src/ldai/runners/types.py new file mode 100644 index 0000000..f4db9a6 --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/runners/types.py @@ -0,0 +1,32 @@ +"""Result types and type aliases for agent and agent graph runners.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Dict + +from ldai.providers.types import LDAIMetrics + +# Type alias for a registry of tools available to an agent. +# Keys are tool names; values are the callable implementations. +ToolRegistry = Dict[str, Callable] + + +@dataclass +class AgentResult: + """ + Result from a single-agent run. + """ + output: str + raw: Any + metrics: LDAIMetrics + + +@dataclass +class AgentGraphResult: + """ + Result from an agent graph run. + """ + output: str + raw: Any + metrics: LDAIMetrics diff --git a/packages/sdk/server-ai/tests/test_judge.py b/packages/sdk/server-ai/tests/test_judge.py index d386b92..9d9f2d2 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 @@ -112,10 +112,10 @@ def test_judge_initializes_with_evaluation_metric_key( assert '$ld:ai:judge:relevance' in judge._evaluation_response_structure['properties']['evaluations']['required'] def test_judge_initializes_without_evaluation_metric_key( - self, judge_config_without_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_ai_provider + self, judge_config_without_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner ): """Judge should initialize but have None for evaluation_response_structure.""" - judge = Judge(judge_config_without_key, tracker, mock_ai_provider) + judge = Judge(judge_config_without_key, tracker, mock_runner) assert judge._ai_config == judge_config_without_key assert judge._evaluation_response_structure is None @@ -126,31 +126,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( @@ -166,10 +166,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.") @@ -182,7 +182,7 @@ async def test_evaluate_success_with_valid_response( @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 evaluation in response.""" mock_response = StructuredResponse( @@ -198,10 +198,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") @@ -211,7 +211,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( @@ -227,10 +227,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") @@ -240,7 +240,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( @@ -255,10 +255,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") @@ -268,13 +268,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") @@ -286,15 +286,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: @@ -302,10 +302,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={ @@ -320,16 +320,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) ) diff --git a/packages/sdk/server-ai/tests/test_runner_abcs.py b/packages/sdk/server-ai/tests/test_runner_abcs.py new file mode 100644 index 0000000..1fc03ae --- /dev/null +++ b/packages/sdk/server-ai/tests/test_runner_abcs.py @@ -0,0 +1,100 @@ +import pytest + +from ldai.providers.types import LDAIMetrics +from ldai.runners.agent_graph_runner import AgentGraphRunner +from ldai.runners.agent_runner import AgentRunner +from ldai.runners.types import AgentGraphResult, AgentResult, ToolRegistry + + +# --- Concrete test doubles --- + +class ConcreteAgentRunner(AgentRunner): + async def run(self, input): + return AgentResult( + output=f"agent response to: {input}", + raw={"raw": input}, + metrics=LDAIMetrics(success=True), + ) + + +class ConcreteAgentGraphRunner(AgentGraphRunner): + async def run(self, input): + return AgentGraphResult( + output=f"graph response to: {input}", + raw={"raw": input}, + metrics=LDAIMetrics(success=True), + ) + + +# --- AgentRunner --- + +def test_agent_runner_is_abstract(): + with pytest.raises(TypeError): + AgentRunner() # type: ignore[abstract] + + +@pytest.mark.asyncio +async def test_agent_runner_run_returns_agent_result(): + runner = ConcreteAgentRunner() + result = await runner.run("hello") + assert isinstance(result, AgentResult) + assert result.output == "agent response to: hello" + assert result.raw == {"raw": "hello"} + assert result.metrics.success is True + + +@pytest.mark.asyncio +async def test_agent_result_fields(): + metrics = LDAIMetrics(success=True) + result = AgentResult(output="done", raw={"key": "val"}, metrics=metrics) + assert result.output == "done" + assert result.raw == {"key": "val"} + assert result.metrics is metrics + + +# --- AgentGraphRunner --- + +def test_agent_graph_runner_is_abstract(): + with pytest.raises(TypeError): + AgentGraphRunner() # type: ignore[abstract] + + +@pytest.mark.asyncio +async def test_agent_graph_runner_run_returns_agent_graph_result(): + runner = ConcreteAgentGraphRunner() + result = await runner.run("hello graph") + assert isinstance(result, AgentGraphResult) + assert result.output == "graph response to: hello graph" + assert result.raw == {"raw": "hello graph"} + assert result.metrics.success is True + + +@pytest.mark.asyncio +async def test_agent_graph_result_fields(): + metrics = LDAIMetrics(success=False) + result = AgentGraphResult(output="", raw=None, metrics=metrics) + assert result.output == "" + assert result.raw is None + assert result.metrics.success is False + + +# --- ToolRegistry --- + +def test_tool_registry_is_dict_of_callables(): + tools: ToolRegistry = { + "search": lambda q: f"results for {q}", + "calculator": lambda x: x * 2, + } + assert tools["search"]("python") == "results for python" + assert tools["calculator"](21) == 42 + + +# --- Top-level exports --- + +def test_top_level_exports(): + import ldai + assert hasattr(ldai, 'AgentRunner') + assert hasattr(ldai, 'AgentGraphRunner') + assert hasattr(ldai, 'AgentResult') + assert hasattr(ldai, 'AgentGraphResult') + assert hasattr(ldai, 'ToolRegistry')