data-designer-engine 0.4.0__tar.gz → 0.4.0rc2__tar.gz

{data_designer_engine-0.4.0 → data_designer_engine-0.4.0rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer-engine
-Version: 0.4.0
+Version: 0.4.0rc2
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta

{data_designer_engine-0.4.0 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.0'
-__version_tuple__ = version_tuple = (0, 4, 0)
+__version__ = version = '0.4.0rc2'
+__version_tuple__ = version_tuple = (0, 4, 0, 'rc2')
 
 __commit_id__ = commit_id = None

{data_designer_engine-0.4.0 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/column_generators/generators/llm_completion.py

@@ -12,7 +12,7 @@ from data_designer.config.column_configs import (
     LLMStructuredColumnConfig,
     LLMTextColumnConfig,
 )
-from data_designer.config.utils.constants import TRACE_COLUMN_POSTFIX
+from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
 from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModel, GenerationStrategy
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
@@ -66,7 +66,7 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
             for context in self.config.multi_modal_context:
                 multi_modal_context.extend(context.get_contexts(deserialized_record))
 
-        response, trace = self.model.generate(
+        response, reasoning_trace = self.model.generate(
             prompt=self.prompt_renderer.render(
                 record=deserialized_record,
                 prompt_template=self.config.prompt,
@@ -87,11 +87,8 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
         serialized_output = self.response_recipe.serialize_output(response)
         data[self.config.name] = self._process_serialized_output(serialized_output)
 
-        should_save_trace = (
-            self.config.with_trace or self.resource_provider.run_config.debug_override_save_all_column_traces
-        )
-        if should_save_trace:
-            data[self.config.name + TRACE_COLUMN_POSTFIX] = [message.to_dict() for message in trace]
+        if reasoning_trace:
+            data[self.config.name + REASONING_TRACE_COLUMN_POSTFIX] = reasoning_trace
 
         return data
 

{data_designer_engine-0.4.0 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/dataset_builders/column_wise_builder.py

@@ -34,7 +34,6 @@ from data_designer.engine.dataset_builders.multi_column_configs import MultiColu
 from data_designer.engine.dataset_builders.utils.concurrency import ConcurrentThreadExecutor
 from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
 from data_designer.engine.dataset_builders.utils.dataset_batch_manager import DatasetBatchManager
-from data_designer.engine.dataset_builders.utils.progress_tracker import ProgressTracker
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
@@ -222,18 +221,16 @@ class ColumnWiseDatasetBuilder:
                 "generator so concurrency through threads is not supported."
             )
 
-        progress_tracker = ProgressTracker(
-            total_records=self.batch_manager.num_records_batch,
-            label=f"{generator.config.column_type} column '{generator.config.name}'",
+        logger.info(
+            f"🐙 Processing {generator.config.column_type} column '{generator.config.name}' "
+            f"with {max_workers} concurrent workers"
         )
-        progress_tracker.log_start(max_workers)
-
         settings = self._resource_provider.run_config
         with ConcurrentThreadExecutor(
             max_workers=max_workers,
             column_name=generator.config.name,
-            result_callback=self._make_result_callback(progress_tracker),
-            error_callback=self._make_error_callback(progress_tracker),
+            result_callback=self._worker_result_callback,
+            error_callback=self._worker_error_callback,
             shutdown_error_rate=settings.shutdown_error_rate,
             shutdown_error_window=settings.shutdown_error_window,
             disable_early_shutdown=settings.disable_early_shutdown,
@@ -241,26 +238,10 @@ class ColumnWiseDatasetBuilder:
             for i, record in self.batch_manager.iter_current_batch():
                 executor.submit(lambda record: generator.generate(record), record, context={"index": i})
 
-        progress_tracker.log_final()
-
         if len(self._records_to_drop) > 0:
             self.batch_manager.drop_records(self._records_to_drop)
             self._records_to_drop.clear()
 
-    def _make_result_callback(self, progress_tracker: ProgressTracker) -> Callable[[dict], None]:
-        def callback(result: dict, *, context: dict | None = None) -> None:
-            self._worker_result_callback(result, context=context)
-            progress_tracker.record_success()
-
-        return callback
-
-    def _make_error_callback(self, progress_tracker: ProgressTracker) -> Callable[[Exception], None]:
-        def callback(exc: Exception, *, context: dict | None = None) -> None:
-            self._worker_error_callback(exc, context=context)
-            progress_tracker.record_failure()
-
-        return callback
-
     def _write_processed_batch(self, dataframe: pd.DataFrame) -> None:
         self.batch_manager.update_records(dataframe.to_dict(orient="records"))
         self.batch_manager.write()

{data_designer_engine-0.4.0 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/models/facade.py

@@ -18,7 +18,7 @@ from data_designer.engine.models.errors import (
 from data_designer.engine.models.litellm_overrides import CustomRouter, LiteLLMRouterDefaultKwargs
 from data_designer.engine.models.parsers.errors import ParserException
 from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
-from data_designer.engine.models.utils import ChatMessage, prompt_to_messages
+from data_designer.engine.models.utils import prompt_to_messages, str_to_message
 from data_designer.engine.secret_resolver import SecretResolver
 from data_designer.lazy_heavy_imports import litellm
 
@@ -67,17 +67,16 @@ class ModelFacade:
         return self._usage_stats
 
     def completion(
-        self, messages: list[ChatMessage], skip_usage_tracking: bool = False, **kwargs
+        self, messages: list[dict[str, str]], skip_usage_tracking: bool = False, **kwargs
     ) -> litellm.ModelResponse:
-        message_payloads = [message.to_dict() for message in messages]
         logger.debug(
             f"Prompting model {self.model_name!r}...",
-            extra={"model": self.model_name, "messages": message_payloads},
+            extra={"model": self.model_name, "messages": messages},
         )
         response = None
         kwargs = self.consolidate_kwargs(**kwargs)
         try:
-            response = self._router.completion(model=self.model_name, messages=message_payloads, **kwargs)
+            response = self._router.completion(model=self.model_name, messages=messages, **kwargs)
             logger.debug(
                 f"Received completion from model {self.model_name!r}",
                 extra={
@@ -150,7 +149,7 @@ class ModelFacade:
         skip_usage_tracking: bool = False,
         purpose: str | None = None,
         **kwargs,
-    ) -> tuple[Any, list[ChatMessage]]:
+    ) -> tuple[Any, str | None]:
         """Generate a parsed output with correction steps.
 
         This generation call will attempt to generate an output which is
@@ -183,12 +182,6 @@ class ModelFacade:
                 It is expected to be used by the @catch_llm_exceptions decorator.
             **kwargs: Additional arguments to pass to the model.
 
-        Returns:
-            A tuple containing:
-                - The parsed output object from the parser.
-                - The full trace of ChatMessage entries in the conversation, including any
-                  corrections and reasoning traces. Callers can decide whether to store this.
-
         Raises:
             GenerationValidationFailureError: If the maximum number of retries or
                 correction steps are met and the last response failures on
@@ -197,17 +190,29 @@ class ModelFacade:
         output_obj = None
         curr_num_correction_steps = 0
         curr_num_restarts = 0
+        curr_generation_attempt = 0
+        max_generation_attempts = (max_correction_steps + 1) * (max_conversation_restarts + 1)
 
         starting_messages = prompt_to_messages(
             user_prompt=prompt, system_prompt=system_prompt, multi_modal_context=multi_modal_context
         )
-        messages: list[ChatMessage] = deepcopy(starting_messages)
+        messages = deepcopy(starting_messages)
 
         while True:
+            curr_generation_attempt += 1
+            logger.debug(
+                f"Starting generation attempt {curr_generation_attempt} of {max_generation_attempts} attempts."
+            )
+
             completion_response = self.completion(messages, skip_usage_tracking=skip_usage_tracking, **kwargs)
             response = completion_response.choices[0].message.content or ""
             reasoning_trace = getattr(completion_response.choices[0].message, "reasoning_content", None)
-            messages.append(ChatMessage.as_assistant(content=response, reasoning_content=reasoning_trace or None))
+
+            if reasoning_trace:
+                ## There are generally some extra newlines with how these get parsed.
+                response = response.strip()
+                reasoning_trace = reasoning_trace.strip()
+
             curr_num_correction_steps += 1
 
             try:
@@ -218,23 +223,21 @@ class ModelFacade:
                     raise GenerationValidationFailureError(
                         "Unsuccessful generation attempt. No retries were attempted."
                     ) from exc
-
                 if curr_num_correction_steps <= max_correction_steps:
-                    # Add user message with error for correction
-                    messages.append(ChatMessage.as_user(content=str(get_exception_primary_cause(exc))))
-
+                    ## Add turns to loop-back errors for correction
+                    messages += [
+                        str_to_message(content=response, role="assistant"),
+                        str_to_message(content=str(get_exception_primary_cause(exc)), role="user"),
+                    ]
                 elif curr_num_restarts < max_conversation_restarts:
                     curr_num_correction_steps = 0
                     curr_num_restarts += 1
                     messages = deepcopy(starting_messages)
-
                 else:
                     raise GenerationValidationFailureError(
-                        f"Unsuccessful generation despite {max_correction_steps} correction steps "
-                        f"and {max_conversation_restarts} conversation restarts."
+                        f"Unsuccessful generation attempt despite {max_generation_attempts} attempts."
                     ) from exc
-
-        return output_obj, messages
+        return output_obj, reasoning_trace
 
     def _get_litellm_deployment(self, model_config: ModelConfig) -> litellm.DeploymentTypedDict:
         provider = self._model_provider_registry.get_provider(model_config.provider)

data_designer_engine-0.4.0rc2/src/data_designer/engine/models/utils.py

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def prompt_to_messages(
+    *,
+    user_prompt: str,
+    system_prompt: str | None = None,
+    multi_modal_context: list[dict[str, Any]] | None = None,
+) -> list[dict[str, str | list[dict]]]:
+    """Convert a user and system prompt into Messages format.
+
+    Args:
+        user_prompt (str): A user prompt.
+        system_prompt (str, optional): An optional system prompt.
+    """
+    user_content = user_prompt
+    if multi_modal_context and len(multi_modal_context) > 0:
+        user_content = []
+        for context in multi_modal_context:
+            user_content.append(context)
+        user_content.append({"type": "text", "text": user_prompt})
+    return (
+        [
+            str_to_message(content=system_prompt, role="system"),
+            str_to_message(content=user_content, role="user"),
+        ]
+        if system_prompt
+        else [str_to_message(content=user_content, role="user")]
+    )
+
+
+def str_to_message(content: str | list[dict], role: str = "user") -> dict[str, str | list[dict]]:
+    return {"content": content, "role": role}

{data_designer_engine-0.4.0 → data_designer_engine-0.4.0rc2}/tests/engine/column_generators/generators/test_llm_completion_generators.py

@@ -12,15 +12,14 @@ from data_designer.config.column_configs import (
     LLMTextColumnConfig,
 )
 from data_designer.config.run_config import RunConfig
-from data_designer.config.utils.constants import TRACE_COLUMN_POSTFIX
 from data_designer.engine.column_generators.generators.base import GenerationStrategy
 from data_designer.engine.column_generators.generators.llm_completion import (
+    REASONING_TRACE_COLUMN_POSTFIX,
     LLMCodeCellGenerator,
     LLMJudgeCellGenerator,
     LLMStructuredCellGenerator,
     LLMTextCellGenerator,
 )
-from data_designer.engine.models.utils import ChatMessage
 
 
 def _create_generator_with_mocks(config_class=LLMTextColumnConfig, **config_kwargs):
@@ -68,14 +67,14 @@ def _create_generator_with_mocks(config_class=LLMTextColumnConfig, **config_kwar
     )
 
 
-def _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, output="test_output"):
+def _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, output="test_output", reasoning=None):
     """Helper function to setup common generate method mocks."""
     mock_prompt_renderer.render.side_effect = ["rendered_user_prompt", "rendered_system_prompt"]
     mock_response_recipe.serialize_output.return_value = {"result": output}
-    mock_model.generate.return_value = ({"result": output}, [])
+    mock_model.generate.return_value = ({"result": output}, reasoning)
 
 
-def test_generate_method() -> None:
+def test_generate_method():
     generator, _, mock_model, _, _, mock_prompt_renderer, mock_response_recipe = _create_generator_with_mocks()
 
     # Test basic generation
@@ -88,19 +87,16 @@ def test_generate_method() -> None:
     assert mock_model.generate.call_args[1]["max_correction_steps"] == 2
     assert mock_model.generate.call_args[1]["max_conversation_restarts"] == 7
     assert result["test_column"] == {"result": "test_output"}
-    assert "test_column" + TRACE_COLUMN_POSTFIX not in result
+    assert "test_column" + REASONING_TRACE_COLUMN_POSTFIX not in result
 
-    # Test with full trace enabled
+    # Test with reasoning trace
     mock_model.reset_mock()
     mock_prompt_renderer.reset_mock()
-    generator.resource_provider.run_config.debug_override_save_all_column_traces = True
-    mock_prompt_renderer.render.side_effect = ["rendered_user_prompt", "rendered_system_prompt"]
-    mock_response_recipe.serialize_output.return_value = {"result": "test_output"}
-    mock_model.generate.return_value = ({"result": "test_output"}, [ChatMessage.as_user("x")])
+    _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, reasoning="reasoning_trace")
     result = generator.generate(data)
 
     assert result["test_column"] == {"result": "test_output"}
-    assert result["test_column" + TRACE_COLUMN_POSTFIX] == [{"role": "user", "content": "x"}]
+    assert result["test_column" + REASONING_TRACE_COLUMN_POSTFIX] == "reasoning_trace"
 
     # Test multi-modal context is None
     call_args = mock_model.generate.call_args
@@ -239,7 +235,7 @@ def test_generate_with_errors(error_type, error_message):
 
     if error_type == "serialization":
         mock_response_recipe.serialize_output.side_effect = Exception(error_message)
-        mock_model.generate.return_value = ({"result": "test_output"}, [])
+        mock_model.generate.return_value = ({"result": "test_output"}, None)
     elif error_type == "model":
         mock_model.generate.side_effect = Exception(error_message)
     elif error_type == "prompt_render":
@@ -253,12 +249,13 @@ def test_generate_with_errors(error_type, error_message):
 
 def test_generate_with_complex_data():
     generator, _, mock_model, _, _, mock_prompt_renderer, mock_response_recipe = _create_generator_with_mocks()
-    _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, "complex_output")
+    _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, "complex_output", "complex_reasoning")
 
     data = {"input": "test_input", "nested": {"key": "value"}, "list": [1, 2, 3], "json_string": '{"key": "value"}'}
     result = generator.generate(data)
 
     assert result["test_column"] == {"result": "complex_output"}
+    assert result["test_column" + REASONING_TRACE_COLUMN_POSTFIX] == "complex_reasoning"
     assert result["input"] == "test_input"
     assert result["nested"] == {"key": "value"}
     assert result["list"] == [1, 2, 3]
@@ -344,7 +341,7 @@ def test_generator_output_type_handling(
     mock_response_recipe.serialize_output.return_value = serialized_output
     stub_resource_provider.model_registry.get_model.return_value.generate.return_value = (
         {"result": "raw_output"},
-        [],
+        None,
     )
 
     data = {"input": "test_input"}

{data_designer_engine-0.4.0 → data_designer_engine-0.4.0rc2}/tests/engine/dataset_builders/test_column_wise_builder.py

@@ -378,7 +378,6 @@ def test_fan_out_with_threads_uses_early_shutdown_settings_from_resource_provide
     mock_generator.config.column_type = "llm_text"
 
     builder.batch_manager = Mock()
-    builder.batch_manager.num_records_batch = 10
     builder.batch_manager.iter_current_batch.return_value = []
 
     builder._fan_out_with_threads(mock_generator, max_workers=4)

{data_designer_engine-0.4.0 → data_designer_engine-0.4.0rc2}/tests/engine/models/test_facade.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any
+from collections import namedtuple
 from unittest.mock import patch
 
 import pytest
@@ -10,33 +10,14 @@ from litellm.types.utils import Choices, EmbeddingResponse, Message, ModelRespon
 from data_designer.engine.models.errors import ModelGenerationValidationFailureError
 from data_designer.engine.models.facade import ModelFacade
 from data_designer.engine.models.parsers.errors import ParserException
-from data_designer.engine.models.utils import ChatMessage
 
+MockMessage = namedtuple("MockMessage", ["content"])
+MockChoice = namedtuple("MockChoice", ["message"])
+MockCompletion = namedtuple("MockCompletion", ["choices"])
 
-class FakeMessage:
-    """Unified fake message class for mocking LLM completion responses."""
 
-    def __init__(
-        self,
-        content: str | None,
-        reasoning_content: str | None = None,
-    ) -> None:
-        self.content = content
-        self.reasoning_content = reasoning_content
-
-
-class FakeChoice:
-    def __init__(self, message: FakeMessage) -> None:
-        self.message = message
-
-
-class FakeResponse:
-    def __init__(self, message: FakeMessage) -> None:
-        self.choices = [FakeChoice(message)]
-
-
-def mock_oai_response_object(response_text: str) -> FakeResponse:
-    return FakeResponse(FakeMessage(content=response_text))
+def mock_oai_response_object(response_text: str) -> MockCompletion:
+    return MockCompletion(choices=[MockChoice(message=MockMessage(content=response_text))])
 
 
 @pytest.fixture
@@ -49,8 +30,8 @@ def stub_model_facade(stub_model_configs, stub_secrets_resolver, stub_model_prov
 
 
 @pytest.fixture
-def stub_completion_messages() -> list[ChatMessage]:
-    return [ChatMessage.as_user("test")]
+def stub_completion_messages():
+    return [{"role": "user", "content": "test"}]
 
 
 @pytest.fixture
@@ -112,29 +93,17 @@ def test_generate(
 @pytest.mark.parametrize(
     "system_prompt,expected_messages",
     [
-        ("", [ChatMessage.as_user("does not matter")]),
-        ("hello!", [ChatMessage.as_system("hello!"), ChatMessage.as_user("does not matter")]),
+        ("", [{"role": "user", "content": "does not matter"}]),
+        ("hello!", [{"content": "hello!", "role": "system"}, {"role": "user", "content": "does not matter"}]),
     ],
 )
 @patch("data_designer.engine.models.facade.ModelFacade.completion", autospec=True)
-def test_generate_with_system_prompt(
-    mock_completion: Any,
-    stub_model_facade: ModelFacade,
-    system_prompt: str,
-    expected_messages: list[ChatMessage],
-) -> None:
-    # Capture messages at call time since they get mutated after the call
-    captured_messages = []
-
-    def capture_and_return(*args: Any, **kwargs: Any) -> ModelResponse:
-        captured_messages.append(list(args[1]))  # Copy the messages list
-        return ModelResponse(choices=Choices(message=Message(content="Hello!")))
-
-    mock_completion.side_effect = capture_and_return
+def test_generate_with_system_prompt(mock_completion, stub_model_facade, system_prompt, expected_messages):
+    mock_completion.return_value = ModelResponse(choices=Choices(message=Message(content="Hello!")))
 
     stub_model_facade.generate(prompt="does not matter", system_prompt=system_prompt, parser=lambda x: x)
     assert mock_completion.call_count == 1
-    assert captured_messages[0] == expected_messages
+    assert mock_completion.call_args[0][1] == expected_messages
 
 
 def test_model_alias_property(stub_model_facade, stub_model_configs):
@@ -182,31 +151,26 @@ def test_consolidate_kwargs(stub_model_configs, stub_model_facade):
 )
 @patch("data_designer.engine.models.facade.CustomRouter.completion", autospec=True)
 def test_completion_success(
-    mock_router_completion: Any,
-    stub_completion_messages: list[ChatMessage],
-    stub_model_configs: Any,
-    stub_model_facade: ModelFacade,
-    stub_expected_completion_response: ModelResponse,
-    skip_usage_tracking: bool,
-) -> None:
+    mock_router_completion,
+    stub_completion_messages,
+    stub_model_configs,
+    stub_model_facade,
+    stub_expected_completion_response,
+    skip_usage_tracking,
+):
     mock_router_completion.side_effect = lambda self, model, messages, **kwargs: stub_expected_completion_response
     result = stub_model_facade.completion(stub_completion_messages, skip_usage_tracking=skip_usage_tracking)
-    expected_messages = [message.to_dict() for message in stub_completion_messages]
     assert result == stub_expected_completion_response
     assert mock_router_completion.call_count == 1
     assert mock_router_completion.call_args[1] == {
         "model": "stub-model-text",
-        "messages": expected_messages,
+        "messages": stub_completion_messages,
         **stub_model_configs[0].inference_parameters.generate_kwargs,
     }
 
 
 @patch("data_designer.engine.models.facade.CustomRouter.completion", autospec=True)
-def test_completion_with_exception(
-    mock_router_completion: Any,
-    stub_completion_messages: list[ChatMessage],
-    stub_model_facade: ModelFacade,
-) -> None:
+def test_completion_with_exception(mock_router_completion, stub_completion_messages, stub_model_facade):
     mock_router_completion.side_effect = Exception("Router error")
 
     with pytest.raises(Exception, match="Router error"):
@@ -215,15 +179,15 @@ def test_completion_with_exception(
 
 @patch("data_designer.engine.models.facade.CustomRouter.completion", autospec=True)
 def test_completion_with_kwargs(
-    mock_router_completion: Any,
-    stub_completion_messages: list[ChatMessage],
-    stub_model_configs: Any,
-    stub_model_facade: ModelFacade,
-    stub_expected_completion_response: ModelResponse,
-) -> None:
+    mock_router_completion,
+    stub_completion_messages,
+    stub_model_configs,
+    stub_model_facade,
+    stub_expected_completion_response,
+):
     captured_kwargs = {}
 
-    def mock_completion(self: Any, model: str, messages: list[dict[str, Any]], **kwargs: Any) -> ModelResponse:
+    def mock_completion(self, model, messages, **kwargs):
         captured_kwargs.update(kwargs)
         return stub_expected_completion_response
 

data_designer_engine-0.4.0rc2/tests/engine/models/test_model_utils.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.engine.models.utils import prompt_to_messages, str_to_message
+
+
+def test_str_to_message():
+    assert str_to_message("hello") == {"content": "hello", "role": "user"}
+    assert str_to_message("hello", role="system") == {"content": "hello", "role": "system"}
+    assert str_to_message([{"type": "text", "text": "hello"}]) == {
+        "content": [{"type": "text", "text": "hello"}],
+        "role": "user",
+    }
+    assert str_to_message([{"type": "text", "text": "hello"}], role="system") == {
+        "content": [{"type": "text", "text": "hello"}],
+        "role": "system",
+    }
+
+
+def test_prompt_to_messages():
+    stub_system_prompt = "some system prompt"
+    mult_modal_context = {"type": "image_url", "image_url": {"url": "http://example.com/image.png"}}
+    assert prompt_to_messages(user_prompt="hello") == [{"content": "hello", "role": "user"}]
+    assert prompt_to_messages(user_prompt="hello", system_prompt=stub_system_prompt) == [
+        {"content": stub_system_prompt, "role": "system"},
+        {"content": "hello", "role": "user"},
+    ]
+    assert prompt_to_messages(user_prompt="hello", multi_modal_context=[mult_modal_context]) == [
+        {"content": [mult_modal_context, {"type": "text", "text": "hello"}], "role": "user"}
+    ]
+    assert prompt_to_messages(
+        user_prompt="hello", system_prompt=stub_system_prompt, multi_modal_context=[mult_modal_context]
+    ) == [
+        {"content": stub_system_prompt, "role": "system"},
+        {"content": [mult_modal_context, {"type": "text", "text": "hello"}], "role": "user"},
+    ]