kiln-ai 0.11.1__py3-none-any.whl → 0.12.0__py3-none-any.whl

kiln_ai/adapters/model_adapters/test_structured_output.py

@@ -2,8 +2,6 @@ import json
 from pathlib import Path
 from typing import Dict
 
-import jsonschema
-import jsonschema.exceptions
 import pytest
 
 import kiln_ai.datamodel as datamodel
@@ -12,16 +10,13 @@ from kiln_ai.adapters.ml_model_list import (
     built_in_models,
 )
 from kiln_ai.adapters.model_adapters.base_adapter import (
-    AdapterInfo,
     BaseAdapter,
     RunOutput,
 )
 from kiln_ai.adapters.ollama_tools import ollama_online
-from kiln_ai.adapters.prompt_builders import (
-    BasePromptBuilder,
-    SimpleChainOfThoughtPromptBuilder,
-)
 from kiln_ai.adapters.test_prompt_adaptors import get_all_models_and_providers
+from kiln_ai.datamodel import PromptId
+from kiln_ai.datamodel.task import RunConfig
 from kiln_ai.datamodel.test_json_schema import json_joke_schema, json_triangle_schema
 
 
@@ -39,9 +34,9 @@ async def test_structured_output_gpt_4o_mini(tmp_path):
     await run_structured_output_test(tmp_path, "gpt_4o_mini", "openai")
 
 
-@pytest.mark.parametrize("model_name", ["llama_3_1_8b"])
+@pytest.mark.parametrize("model_name", ["llama_3_1_8b", "gemma_2_2b"])
 @pytest.mark.ollama
-async def test_structured_output_ollama_llama(tmp_path, model_name):
+async def test_structured_output_ollama(tmp_path, model_name):
     if not await ollama_online():
         pytest.skip("Ollama API not running. Expect it running on localhost:11434")
     await run_structured_output_test(tmp_path, model_name, "ollama")
@@ -49,19 +44,21 @@ async def test_structured_output_ollama_llama(tmp_path, model_name):
 
 class MockAdapter(BaseAdapter):
     def __init__(self, kiln_task: datamodel.Task, response: Dict | str | None):
-        super().__init__(kiln_task, model_name="phi_3_5", model_provider_name="ollama")
+        super().__init__(
+            run_config=RunConfig(
+                task=kiln_task,
+                model_name="phi_3_5",
+                model_provider_name="ollama",
+                prompt_id="simple_chain_of_thought_prompt_builder",
+            ),
+        )
         self.response = response
 
     async def _run(self, input: str) -> RunOutput:
         return RunOutput(output=self.response, intermediate_outputs=None)
 
-    def adapter_info(self) -> AdapterInfo:
-        return AdapterInfo(
-            adapter_name="mock_adapter",
-            model_name="mock_model",
-            model_provider="mock_provider",
-            prompt_builder_name="mock_prompt_builder",
-        )
+    def adapter_name(self) -> str:
+        return "mock_adapter"
 
 
 async def test_mock_unstructred_response(tmp_path):
@@ -204,15 +201,21 @@ async def run_structured_input_task(
     task: datamodel.Task,
     model_name: str,
     provider: str,
-    pb: BasePromptBuilder | None = None,
+    prompt_id: PromptId | None = None,
 ):
     a = adapter_for_task(
-        task, model_name=model_name, provider=provider, prompt_builder=pb
+        task,
+        model_name=model_name,
+        provider=provider,
+        prompt_id=prompt_id,
     )
     with pytest.raises(ValueError):
         # not structured input in dictionary
         await a.invoke("a=1, b=2, c=3")
-    with pytest.raises(jsonschema.exceptions.ValidationError):
+    with pytest.raises(
+        ValueError,
+        match="This task requires a specific output schema. While the model produced JSON, that JSON didn't meet the schema.",
+    ):
         # invalid structured input
         await a.invoke({"a": 1, "b": 2, "d": 3})
 
@@ -229,13 +232,14 @@ async def run_structured_input_task(
         assert "[[equilateral]]" in response
     else:
         assert response["is_equilateral"] is True
-    adapter_info = a.adapter_info()
+
     expected_pb_name = "simple_prompt_builder"
-    if pb is not None:
-        expected_pb_name = pb.__class__.prompt_builder_name()
-    assert adapter_info.prompt_builder_name == expected_pb_name
-    assert adapter_info.model_name == model_name
-    assert adapter_info.model_provider == provider
+    if prompt_id is not None:
+        expected_pb_name = prompt_id
+    assert a.run_config.prompt_id == expected_pb_name
+
+    assert a.run_config.model_name == model_name
+    assert a.run_config.model_provider_name == provider
 
 
 @pytest.mark.paid
@@ -257,8 +261,9 @@ async def test_all_built_in_models_structured_input(
 @pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
 async def test_structured_input_cot_prompt_builder(tmp_path, model_name, provider_name):
     task = build_structured_input_test_task(tmp_path)
-    pb = SimpleChainOfThoughtPromptBuilder(task)
-    await run_structured_input_task(task, model_name, provider_name, pb)
+    await run_structured_input_task(
+        task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
+    )
 
 
 @pytest.mark.paid
@@ -302,5 +307,6 @@ When asked for a final result, this is the format (for an equilateral example):
 """
     task.output_json_schema = json.dumps(triangle_schema)
     task.save_to_file()
-    pb = SimpleChainOfThoughtPromptBuilder(task)
-    await run_structured_input_task(task, model_name, provider_name, pb)
+    await run_structured_input_task(
+        task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
+    )

kiln_ai/adapters/ollama_tools.py

@@ -1,4 +1,3 @@
-import os
 from typing import Any, List
 
 import httpx

kiln_ai/adapters/prompt_builders.py

@@ -2,8 +2,8 @@ import json
 from abc import ABCMeta, abstractmethod
 from typing import Dict
 
-from kiln_ai.datamodel import Task, TaskRun
-from kiln_ai.utils.formatting import snake_case
+from kiln_ai.datamodel import PromptGenerators, PromptId, Task, TaskRun
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
 class BasePromptBuilder(metaclass=ABCMeta):
@@ -53,17 +53,6 @@ class BasePromptBuilder(metaclass=ABCMeta):
         """
         pass
 
-    @classmethod
-    def prompt_builder_name(cls) -> str:
-        """Returns the name of the prompt builder, to be used for persisting into the datastore.
-
-        Default implementation gets the name of the prompt builder in snake case. If you change the class name, you should override this so prior saved data is compatible.
-
-        Returns:
-            str: The prompt builder name in snake_case format.
-        """
-        return snake_case(cls.__name__)
-
     def build_user_message(self, input: Dict | str) -> str:
         """Build a user message from the input.
 
@@ -300,6 +289,57 @@ class SavedPromptBuilder(BasePromptBuilder):
         return self.prompt_model.chain_of_thought_instructions
 
 
+class TaskRunConfigPromptBuilder(BasePromptBuilder):
+    """A prompt builder that looks up a static prompt in a task run config."""
+
+    def __init__(self, task: Task, run_config_prompt_id: str):
+        parts = run_config_prompt_id.split("::")
+        if len(parts) != 4:
+            raise ValueError(
+                f"Invalid task run config prompt ID: {run_config_prompt_id}. Expected format: 'task_run_config::[project_id]::[task_id]::[run_config_id]'."
+            )
+
+        task_id = parts[2]
+        if task_id != task.id:
+            raise ValueError(
+                f"Task run config prompt ID: {run_config_prompt_id}. Task ID mismatch. Expected: {task.id}, got: {task_id}."
+            )
+
+        run_config_id = parts[3]
+        run_config = next(
+            (
+                run_config
+                for run_config in task.run_configs(readonly=True)
+                if run_config.id == run_config_id
+            ),
+            None,
+        )
+        if not run_config:
+            raise ValueError(
+                f"Task run config ID not found: {run_config_id} for prompt id {run_config_prompt_id}"
+            )
+        if run_config.prompt is None:
+            raise ValueError(
+                f"Task run config ID {run_config_id} does not have a stored prompt. Used as prompt id {run_config_prompt_id}"
+            )
+
+        # Load the prompt from the model
+        self.prompt = run_config.prompt.prompt
+        self.cot_prompt = run_config.prompt.chain_of_thought_instructions
+        self.id = run_config_prompt_id
+
+        super().__init__(task)
+
+    def prompt_id(self) -> str | None:
+        return self.id
+
+    def build_base_prompt(self) -> str:
+        return self.prompt
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return self.cot_prompt
+
+
 class FineTunePromptBuilder(BasePromptBuilder):
     """A prompt builder that looks up a fine-tune prompt."""
 
@@ -337,25 +377,12 @@ class FineTunePromptBuilder(BasePromptBuilder):
         return self.fine_tune_model.thinking_instructions
 
 
-# TODO P2: we end up with 2 IDs for these: the keys here (ui_name) and the prompt_builder_name from the class
-# We end up maintaining this in _prompt_generators as well.
-prompt_builder_registry = {
-    "simple_prompt_builder": SimplePromptBuilder,
-    "multi_shot_prompt_builder": MultiShotPromptBuilder,
-    "few_shot_prompt_builder": FewShotPromptBuilder,
-    "repairs_prompt_builder": RepairsPromptBuilder,
-    "simple_chain_of_thought_prompt_builder": SimpleChainOfThoughtPromptBuilder,
-    "few_shot_chain_of_thought_prompt_builder": FewShotChainOfThoughtPromptBuilder,
-    "multi_shot_chain_of_thought_prompt_builder": MultiShotChainOfThoughtPromptBuilder,
-}
-
-
 # Our UI has some names that are not the same as the class names, which also hint parameters.
-def prompt_builder_from_ui_name(ui_name: str, task: Task) -> BasePromptBuilder:
+def prompt_builder_from_id(prompt_id: PromptId, task: Task) -> BasePromptBuilder:
     """Convert a name used in the UI to the corresponding prompt builder class.
 
     Args:
-        ui_name (str): The UI name for the prompt builder type.
+        prompt_id (PromptId): The prompt ID.
 
     Returns:
         type[BasePromptBuilder]: The corresponding prompt builder class.
@@ -365,29 +392,40 @@ def prompt_builder_from_ui_name(ui_name: str, task: Task) -> BasePromptBuilder:
     """
 
     # Saved prompts are prefixed with "id::"
-    if ui_name.startswith("id::"):
-        prompt_id = ui_name[4:]
+    if prompt_id.startswith("id::"):
+        prompt_id = prompt_id[4:]
         return SavedPromptBuilder(task, prompt_id)
 
+    # Task run config prompts are prefixed with "task_run_config::"
+    # task_run_config::[project_id]::[task_id]::[run_config_id]
+    if prompt_id.startswith("task_run_config::"):
+        return TaskRunConfigPromptBuilder(task, prompt_id)
+
     # Fine-tune prompts are prefixed with "fine_tune_prompt::"
-    if ui_name.startswith("fine_tune_prompt::"):
-        fine_tune_id = ui_name[18:]
-        return FineTunePromptBuilder(task, fine_tune_id)
+    if prompt_id.startswith("fine_tune_prompt::"):
+        prompt_id = prompt_id[18:]
+        return FineTunePromptBuilder(task, prompt_id)
+
+    # Check if the prompt_id matches any enum value
+    if prompt_id not in [member.value for member in PromptGenerators]:
+        raise ValueError(f"Unknown prompt generator: {prompt_id}")
+    typed_prompt_generator = PromptGenerators(prompt_id)
 
-    match ui_name:
-        case "basic":
+    match typed_prompt_generator:
+        case PromptGenerators.SIMPLE:
             return SimplePromptBuilder(task)
-        case "few_shot":
+        case PromptGenerators.FEW_SHOT:
             return FewShotPromptBuilder(task)
-        case "many_shot":
+        case PromptGenerators.MULTI_SHOT:
             return MultiShotPromptBuilder(task)
-        case "repairs":
+        case PromptGenerators.REPAIRS:
             return RepairsPromptBuilder(task)
-        case "simple_chain_of_thought":
+        case PromptGenerators.SIMPLE_CHAIN_OF_THOUGHT:
             return SimpleChainOfThoughtPromptBuilder(task)
-        case "few_shot_chain_of_thought":
+        case PromptGenerators.FEW_SHOT_CHAIN_OF_THOUGHT:
             return FewShotChainOfThoughtPromptBuilder(task)
-        case "multi_shot_chain_of_thought":
+        case PromptGenerators.MULTI_SHOT_CHAIN_OF_THOUGHT:
             return MultiShotChainOfThoughtPromptBuilder(task)
         case _:
-            raise ValueError(f"Unknown prompt builder: {ui_name}")
+            # Type checking will find missing cases
+            raise_exhaustive_enum_error(typed_prompt_generator)

kiln_ai/adapters/repair/repair_task.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel, Field
 from kiln_ai.adapters.prompt_builders import (
     BasePromptBuilder,
     SavedPromptBuilder,
-    prompt_builder_registry,
+    prompt_builder_from_id,
 )
 from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
 
@@ -49,28 +49,16 @@ feedback describing what should be improved. Your job is to understand the evalu
         if run.output.source is None or run.output.source.properties is None:
             raise ValueError("No source properties found")
 
-        # Try ID first, then builder name
-        prompt_id = run.output.source.properties.get("prompt_id", None)
+        # Get the prompt builder id. Need the second check because we used to store this in a prompt_builder_name field, so loading legacy runs will need this.
+        prompt_id = run.output.source.properties.get(
+            "prompt_id"
+        ) or run.output.source.properties.get("prompt_builder_name", None)
         if prompt_id is not None and isinstance(prompt_id, str):
-            static_prompt_builder = SavedPromptBuilder(task, prompt_id)
-            return static_prompt_builder.build_prompt(include_json_instructions=False)
+            prompt_builder = prompt_builder_from_id(prompt_id, task)
+            if isinstance(prompt_builder, BasePromptBuilder):
+                return prompt_builder.build_prompt(include_json_instructions=False)
 
-        prompt_builder_class: Type[BasePromptBuilder] | None = None
-        prompt_builder_name = run.output.source.properties.get(
-            "prompt_builder_name", None
-        )
-        if prompt_builder_name is not None and isinstance(prompt_builder_name, str):
-            prompt_builder_class = prompt_builder_registry.get(
-                prompt_builder_name, None
-            )
-        if prompt_builder_class is None:
-            raise ValueError(f"No prompt builder found for name: {prompt_builder_name}")
-        prompt_builder = prompt_builder_class(task=task)
-        if not isinstance(prompt_builder, BasePromptBuilder):
-            raise ValueError(
-                f"Prompt builder {prompt_builder_name} is not a valid prompt builder"
-            )
-        return prompt_builder.build_prompt(include_json_instructions=False)
+        raise ValueError(f"Prompt builder '{prompt_id}' is not a valid prompt builder")
 
     @classmethod
     def build_repair_task_input(

kiln_ai/adapters/repair/test_repair_task.py

@@ -95,7 +95,7 @@ def sample_task_run(sample_task):
                     "model_name": "gpt_4o",
                     "model_provider": "openai",
                     "adapter_name": "langchain_adapter",
-                    "prompt_builder_name": "simple_prompt_builder",
+                    "prompt_id": "simple_prompt_builder",
                 },
             ),
         ),
@@ -201,7 +201,7 @@ async def test_live_run(sample_task, sample_task_run, sample_repair_data):
         "adapter_name": "kiln_langchain_adapter",
         "model_name": "llama_3_1_8b",
         "model_provider": "groq",
-        "prompt_builder_name": "simple_prompt_builder",
+        "prompt_id": "simple_prompt_builder",
     }
 
 
@@ -238,7 +238,7 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
         "adapter_name": "kiln_langchain_adapter",
         "model_name": "llama_3_1_8b",
         "model_provider": "ollama",
-        "prompt_builder_name": "simple_prompt_builder",
+        "prompt_id": "simple_prompt_builder",
     }
     assert run.input_source.type == DataSourceType.human
     assert "created_by" in run.input_source.properties

kiln_ai/adapters/run_output.py

@@ -1,8 +1,11 @@
 from dataclasses import dataclass
 from typing import Dict
 
+from openai.types.chat.chat_completion import ChoiceLogprobs
+
 
 @dataclass
 class RunOutput:
     output: Dict | str
     intermediate_outputs: Dict[str, str] | None
+    output_logprobs: ChoiceLogprobs | None = None

kiln_ai/adapters/test_adapter_registry.py

@@ -5,6 +5,7 @@ import pytest
 from kiln_ai import datamodel
 from kiln_ai.adapters.adapter_registry import adapter_for_task
 from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
 from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
 from kiln_ai.adapters.model_adapters.openai_model_adapter import OpenAICompatibleAdapter
 from kiln_ai.adapters.prompt_builders import BasePromptBuilder
@@ -84,24 +85,19 @@ def test_langchain_adapter_creation(mock_config, basic_task, provider):
     )
 
     assert isinstance(adapter, LangchainAdapter)
-    assert adapter.model_name == "test-model"
+    assert adapter.run_config.model_name == "test-model"
 
 
 # TODO should run for all cases
 def test_custom_prompt_builder(mock_config, basic_task):
-    class TestPromptBuilder(BasePromptBuilder):
-        def build_base_prompt(self, kiln_task) -> str:
-            return "test-prompt"
-
-    prompt_builder = TestPromptBuilder(basic_task)
     adapter = adapter_for_task(
         kiln_task=basic_task,
         model_name="gpt-4",
         provider=ModelProviderName.openai,
-        prompt_builder=prompt_builder,
+        prompt_id="simple_chain_of_thought_prompt_builder",
     )
 
-    assert adapter.prompt_builder == prompt_builder
+    assert adapter.run_config.prompt_id == "simple_chain_of_thought_prompt_builder"
 
 
 # TODO should run for all cases
@@ -111,10 +107,12 @@ def test_tags_passed_through(mock_config, basic_task):
         kiln_task=basic_task,
         model_name="gpt-4",
         provider=ModelProviderName.openai,
-        tags=tags,
+        base_adapter_config=AdapterConfig(
+            default_tags=tags,
+        ),
     )
 
-    assert adapter.default_tags == tags
+    assert adapter.base_adapter_config.default_tags == tags
 
 
 def test_invalid_provider(mock_config, basic_task):
@@ -129,6 +127,7 @@ def test_openai_compatible_adapter(mock_compatible_config, mock_config, basic_ta
     mock_compatible_config.return_value.model_name = "test-model"
     mock_compatible_config.return_value.api_key = "test-key"
     mock_compatible_config.return_value.base_url = "https://test.com/v1"
+    mock_compatible_config.return_value.provider_name = "CustomProvider99"
 
     adapter = adapter_for_task(
         kiln_task=basic_task,
@@ -141,6 +140,7 @@ def test_openai_compatible_adapter(mock_compatible_config, mock_config, basic_ta
     assert adapter.config.model_name == "test-model"
     assert adapter.config.api_key == "test-key"
     assert adapter.config.base_url == "https://test.com/v1"
+    assert adapter.config.provider_name == "CustomProvider99"
 
 
 def test_custom_openai_compatible_provider(mock_config, basic_task):

kiln_ai/adapters/test_generate_docs.py

@@ -1,13 +1,13 @@
+import logging
 from typing import List
 
 import pytest
 
-from libs.core.kiln_ai.adapters.ml_model_list import (
-    KilnModelProvider,
-    built_in_models,
-)
+from libs.core.kiln_ai.adapters.ml_model_list import KilnModelProvider, built_in_models
 from libs.core.kiln_ai.adapters.provider_tools import provider_name_from_id
 
+logger = logging.getLogger(__name__)
+
 
 def _all_providers_support(providers: List[KilnModelProvider], attribute: str) -> bool:
     """Check if all providers support a given feature"""
@@ -58,8 +58,8 @@ def test_generate_model_table():
         table.append(row)
 
     # Print the table (useful for documentation)
-    print("\nModel Capability Matrix:\n")
-    print("\n".join(table))
+    logger.info("\nModel Capability Matrix:\n")
+    logger.info("\n".join(table))
 
     # Basic assertions to ensure the table is well-formed
     assert len(table) > 2, "Table should have header and at least one row"

kiln_ai/adapters/test_ollama_tools.py

@@ -10,7 +10,6 @@ from kiln_ai.adapters.ollama_tools import (
 def test_parse_ollama_tags_no_models():
     json_response = '{"models":[{"name":"scosman_net","model":"scosman_net:latest"},{"name":"phi3.5:latest","model":"phi3.5:latest","modified_at":"2024-10-02T12:04:35.191519822-04:00","size":2176178843,"digest":"61819fb370a3c1a9be6694869331e5f85f867a079e9271d66cb223acb81d04ba","details":{"parent_model":"","format":"gguf","family":"phi3","families":["phi3"],"parameter_size":"3.8B","quantization_level":"Q4_0"}},{"name":"gemma2:2b","model":"gemma2:2b","modified_at":"2024-09-09T16:46:38.64348929-04:00","size":1629518495,"digest":"8ccf136fdd5298f3ffe2d69862750ea7fb56555fa4d5b18c04e3fa4d82ee09d7","details":{"parent_model":"","format":"gguf","family":"gemma2","families":["gemma2"],"parameter_size":"2.6B","quantization_level":"Q4_0"}},{"name":"llama3.1:latest","model":"llama3.1:latest","modified_at":"2024-09-01T17:19:43.481523695-04:00","size":4661230720,"digest":"f66fc8dc39ea206e03ff6764fcc696b1b4dfb693f0b6ef751731dd4e6269046e","details":{"parent_model":"","format":"gguf","family":"llama","families":["llama"],"parameter_size":"8.0B","quantization_level":"Q4_0"}}]}'
     tags = json.loads(json_response)
-    print(json.dumps(tags, indent=2))
     conn = parse_ollama_tags(tags)
     assert "phi3.5:latest" in conn.supported_models
     assert "gemma2:2b" in conn.supported_models

kiln_ai/adapters/test_prompt_adaptors.py

@@ -13,6 +13,7 @@ from kiln_ai.adapters.prompt_builders import (
     BasePromptBuilder,
     SimpleChainOfThoughtPromptBuilder,
 )
+from kiln_ai.datamodel import PromptId
 
 
 def get_all_models_and_providers():
@@ -132,7 +133,7 @@ async def test_mock_returning_run(tmp_path):
         "adapter_name": "kiln_langchain_adapter",
         "model_name": "custom.langchain:unknown_model",
         "model_provider": "ollama",
-        "prompt_builder_name": "simple_prompt_builder",
+        "prompt_id": "simple_prompt_builder",
     }
 
 
@@ -149,8 +150,9 @@ async def test_all_models_providers_plaintext(tmp_path, model_name, provider_nam
 @pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
 async def test_cot_prompt_builder(tmp_path, model_name, provider_name):
     task = build_test_task(tmp_path)
-    pb = SimpleChainOfThoughtPromptBuilder(task)
-    await run_simple_task(task, model_name, provider_name, pb)
+    await run_simple_task(
+        task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
+    )
 
 
 def build_test_task(tmp_path: Path):
@@ -186,20 +188,20 @@ async def run_simple_test(
     tmp_path: Path,
     model_name: str,
     provider: str | None = None,
-    prompt_builder: BasePromptBuilder | None = None,
+    prompt_id: PromptId | None = None,
 ):
     task = build_test_task(tmp_path)
-    return await run_simple_task(task, model_name, provider, prompt_builder)
+    return await run_simple_task(task, model_name, provider, prompt_id)
 
 
 async def run_simple_task(
     task: datamodel.Task,
     model_name: str,
     provider: str,
-    prompt_builder: BasePromptBuilder | None = None,
+    prompt_id: PromptId | None = None,
 ) -> datamodel.TaskRun:
     adapter = adapter_for_task(
-        task, model_name=model_name, provider=provider, prompt_builder=prompt_builder
+        task, model_name=model_name, provider=provider, prompt_id=prompt_id
     )
 
     run = await adapter.invoke(
@@ -212,13 +214,14 @@ async def run_simple_task(
     )
     assert "64" in run.output.output
     source_props = run.output.source.properties
-    assert source_props["adapter_name"] == "kiln_langchain_adapter"
+    assert source_props["adapter_name"] in [
+        "kiln_langchain_adapter",
+        "kiln_openai_compatible_adapter",
+    ]
     assert source_props["model_name"] == model_name
     assert source_props["model_provider"] == provider
-    expected_prompt_builder_name = (
-        prompt_builder.__class__.prompt_builder_name()
-        if prompt_builder
-        else "simple_prompt_builder"
-    )
-    assert source_props["prompt_builder_name"] == expected_prompt_builder_name
+    if prompt_id is None:
+        assert source_props["prompt_id"] == "simple_prompt_builder"
+    else:
+        assert source_props["prompt_id"] == prompt_id
     return run