LLM Provider & Message Abstractions (#4826)

Summary:
Pull Request resolved: #4826

This diff implements a simple `LLMProvider` protocol (to be followed with a `LiteLLMProvider` implementation for easy access to many providers) and an `LLMMessage` dataclass. These classes support easy typing for any LLM usage within Ax, and enable storage of the inputs and conversations.

Reviewed By: lena-kashtelyan

Differential Revision: D90788634

fbshipit-source-id: 750ccbe916d57224f07a04f88e4d6648cb5fa69b

Author: saitcakmak

ax/core/__init__.py

@@ -12,6 +12,7 @@
 from ax.core.data import Data
 from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun
+from ax.core.llm_provider import LLMMessage, LLMProvider
 from ax.core.metric import Metric
 from ax.core.objective import MultiObjective, Objective
 from ax.core.observation import ObservationFeatures
@@ -46,6 +47,8 @@
     "Experiment",
     "FixedParameter",
     "GeneratorRun",
+    "LLMMessage",
+    "LLMProvider",
     "Metric",
     "MultiObjective",
     "MultiObjectiveOptimizationConfig",

ax/core/llm_provider.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+"""
+LLM Provider protocol and message types for Ax.
+
+This module defines the core abstractions for LLM integration in Ax:
+- LLMMessage: Unified message type for conversations and responses
+- LLMProvider: Protocol defining the interface for LLM providers
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Literal, Protocol, runtime_checkable
+
+
+@dataclass
+class LLMMessage:
+    """Represents a single message in a conversation.
+
+    This unified class handles both input messages and LLM responses.
+    For assistant responses (role="assistant"), the metadata field
+    captures information about the generation.
+
+    Attributes:
+        role: Message role - "system", "user", or "assistant"
+        content: Message content/text
+        metadata: Additional metadata.
+            For assistant responses, this may include:
+            - "usage": Token usage statistics
+            - "finish_reason": Reason for generation completion (e.g., "stop")
+    """
+
+    role: Literal["system", "user", "assistant"]
+    content: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@runtime_checkable
+class LLMProvider(Protocol):
+    """Protocol defining the interface for LLM providers.
+
+    Any class implementing this protocol can be used as an LLM provider in Ax.
+    This enables easy integration of custom LLM backends without requiring
+    inheritance from a base class.
+
+    Implementations must provide:
+        - generate(): method to generate responses from messages
+
+    Example:
+        >>> class MyCustomProvider:
+        ...     def generate(
+        ...         self,
+        ...         messages: list[LLMMessage],
+        ...         **kwargs: Any,
+        ...     ) -> LLMMessage:
+        ...         # Custom implementation
+        ...         return LLMMessage(role="assistant", content="response")
+        ...
+        >>> # Type checker will accept this as LLMProvider
+        >>> provider: LLMProvider = MyCustomProvider()
+    """
+
+    def generate(
+        self,
+        messages: list[LLMMessage],
+        **kwargs: Any,
+    ) -> LLMMessage:
+        """Generate a response from a sequence of messages.
+
+        Args:
+            messages: List of conversation messages with roles and content
+            **kwargs: Provider-specific parameters (e.g., temperature, max_tokens)
+
+        Returns:
+            LLMMessage with role="assistant" containing the generated response
+        """
+        ...

ax/core/tests/test_llm_provider.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from typing import Any, Literal
+
+from ax.core.llm_provider import LLMMessage, LLMProvider
+from ax.utils.common.testutils import TestCase
+
+
+class LLMMessageTest(TestCase):
+    def test_llm_message(self) -> None:
+        """Test LLMMessage creation and validation."""
+        test_cases: list[tuple[Literal["user", "system", "assistant"], str]] = [
+            ("user", "Hello"),
+            ("system", "You are helpful"),
+            ("assistant", "Hi there"),
+        ]
+        for role, content in test_cases:
+            with self.subTest(role=role):
+                msg = LLMMessage(role=role, content=content)
+                self.assertEqual(msg.role, role)
+                self.assertEqual(msg.content, content)
+                self.assertEqual(msg.metadata, {})
+
+    def test_llm_message_with_metadata(self) -> None:
+        """Test LLMMessage with assistant metadata (response case)."""
+        msg = LLMMessage(
+            role="assistant",
+            content="Hello world",
+            metadata={
+                "finish_reason": "stop",
+                "usage": {
+                    "prompt_tokens": 10,
+                    "completion_tokens": 20,
+                    "total_tokens": 30,
+                },
+            },
+        )
+        self.assertEqual(msg.role, "assistant")
+        self.assertEqual(msg.content, "Hello world")
+        self.assertEqual(msg.metadata["finish_reason"], "stop")
+        self.assertEqual(msg.metadata["usage"]["total_tokens"], 30)
+
+        # With minimal fields
+        msg_minimal = LLMMessage(role="assistant", content="Hello")
+        self.assertEqual(msg_minimal.metadata, {})
+
+
+class LLMProviderProtocolTest(TestCase):
+    def test_protocol_compliance(self) -> None:
+        """Test that custom classes can implement the LLMProvider protocol."""
+
+        class MockProvider:
+            """A mock provider that implements the LLMProvider protocol."""
+
+            def generate(
+                self,
+                messages: list[LLMMessage],
+                **kwargs: Any,
+            ) -> LLMMessage:
+                return LLMMessage(
+                    role="assistant",
+                    content=f"Mock response to: {messages[-1].content}",
+                )
+
+        provider = MockProvider()
+
+        # Test that it's recognized as implementing the protocol
+        self.assertIsInstance(provider, LLMProvider)
+
+        # Test that it works
+        response = provider.generate(
+            messages=[LLMMessage(role="user", content="Hello")]
+        )
+        self.assertEqual(response.role, "assistant")
+        self.assertIn("Hello", response.content)

ax/storage/json_store/registry.py

@@ -30,6 +30,7 @@
 from ax.core.data import Data
 from ax.core.evaluations_to_data import DataType
 from ax.core.generator_run import GeneratorRun
+from ax.core.llm_provider import LLMMessage
 from ax.core.map_metric import MapMetric
 from ax.core.metric import Metric
 from ax.core.multi_type_experiment import MultiTypeExperiment
@@ -336,6 +337,7 @@
     # name linked to the new corresponding class
     "ListSurrogate": Surrogate,
     "L2NormMetric": L2NormMetric,
+    "LLMMessage": LLMMessage,
     "LogNormalPrior": LogNormalPrior,
     "MapData": Data,
     "MapMetric": MapMetric,

ax/storage/json_store/tests/test_json_store.py

@@ -34,6 +34,7 @@
 from ax.core.auxiliary import AuxiliaryExperimentPurpose
 from ax.core.data import Data
 from ax.core.generator_run import GeneratorRun
+from ax.core.llm_provider import LLMMessage
 from ax.core.metric import Metric
 from ax.core.objective import Objective
 from ax.core.observation import ObservationFeatures
@@ -348,6 +349,14 @@
     ("HierarchicalSearchSpace", get_hierarchical_search_space),
     ("ImprovementGlobalStoppingStrategy", get_improvement_global_stopping_strategy),
     ("Interval", get_interval),
+    (
+        "LLMMessage",
+        lambda: LLMMessage(
+            role="assistant",
+            content="Hello!",
+            metadata={"finish_reason": "stop", "usage": {"total_tokens": 10}},
+        ),
+    ),
     ("MapData", get_map_data),
     ("MapMetric", partial(get_map_metric, name="test")),
     ("Metric", get_metric),