kiln-ai 0.11.1__py3-none-any.whl → 0.13.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):
@@ -69,7 +66,8 @@ async def test_mock_unstructred_response(tmp_path):
 
     # don't error on valid response
     adapter = MockAdapter(task, response={"setup": "asdf", "punchline": "asdf"})
-    answer = await adapter.invoke_returning_raw("You are a mock, send me the response!")
+    run = await adapter.invoke("You are a mock, send me the response!")
+    answer = json.loads(run.output.output)
     assert answer["setup"] == "asdf"
     assert answer["punchline"] == "asdf"
 
@@ -79,9 +77,12 @@ async def test_mock_unstructred_response(tmp_path):
         answer = await adapter.invoke("You are a mock, send me the response!")
 
     adapter = MockAdapter(task, response="string instead of dict")
-    with pytest.raises(RuntimeError):
+    with pytest.raises(
+        ValueError,
+        match="This task requires JSON output but the model didn't return valid JSON",
+    ):
         # Not a structed response so should error
-        answer = await adapter.invoke("You are a mock, send me the response!")
+        run = await adapter.invoke("You are a mock, send me the response!")
 
     # Should error, expecting a string, not a dict
     project = datamodel.Project(name="test", path=tmp_path / "test.kiln")
@@ -146,7 +147,8 @@ async def run_structured_output_test(tmp_path: Path, model_name: str, provider:
     task = build_structured_output_test_task(tmp_path)
     a = adapter_for_task(task, model_name=model_name, provider=provider)
     try:
-        parsed = await a.invoke_returning_raw("Cows")  # a joke about cows
+        run = await a.invoke("Cows")  # a joke about cows
+        parsed = json.loads(run.output.output)
     except ValueError as e:
         if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
             pytest.skip(
@@ -165,6 +167,12 @@ async def run_structured_output_test(tmp_path: Path, model_name: str, provider:
         assert rating >= 0
         assert rating <= 10
 
+    # Check reasoning models
+    assert a._model_provider is not None
+    if a._model_provider.reasoning_capable:
+        assert "reasoning" in run.intermediate_outputs
+        assert isinstance(run.intermediate_outputs["reasoning"], str)
+
 
 def build_structured_input_test_task(tmp_path: Path):
     project = datamodel.Project(name="test", path=tmp_path / "test.kiln")
@@ -204,20 +212,27 @@ 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})
 
     try:
-        response = await a.invoke_returning_raw({"a": 2, "b": 2, "c": 2})
+        run = await a.invoke({"a": 2, "b": 2, "c": 2})
+        response = run.output.output
     except ValueError as e:
         if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
             pytest.skip(
@@ -229,13 +244,20 @@ 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
+
+    # Check reasoning models
+    assert a._model_provider is not None
+    if a._model_provider.reasoning_capable:
+        assert "reasoning" in run.intermediate_outputs
+        assert isinstance(run.intermediate_outputs["reasoning"], str)
 
 
 @pytest.mark.paid
@@ -257,8 +279,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 +325,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
@@ -39,6 +38,7 @@ async def ollama_online() -> bool:
 
 class OllamaConnection(BaseModel):
     message: str
+    version: str | None = None
     supported_models: List[str]
     untested_models: List[str] = Field(default_factory=list)
 
@@ -50,7 +50,7 @@ class OllamaConnection(BaseModel):
 def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
     # Build a list of models we support for Ollama from the built-in model list
     supported_ollama_models = [
-        provider.provider_options["model"]
+        provider.model_id
         for model in built_in_models
         for provider in model.providers
         if provider.name == ModelProviderName.ollama
@@ -61,7 +61,7 @@ def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
             alias
             for model in built_in_models
             for provider in model.providers
-            for alias in provider.provider_options.get("model_aliases", [])
+            for alias in provider.ollama_model_aliases or []
         ]
     )
 

kiln_ai/adapters/parsers/r1_parser.py

@@ -20,21 +20,33 @@ class R1ThinkingParser(BaseParser):
         Raises:
             ValueError: If response format is invalid (missing tags, multiple tags, or no content after closing tag)
         """
+
+        # The upstream providers (litellm, openrouter, fireworks) all keep changing their response formats, sometimes adding reasoning parsing where it didn't previously exist.
+        # If they do it already, great just return. If not we parse it ourselves. Not ideal, but better than upstream changes breaking the app.
+        if (
+            original_output.intermediate_outputs is not None
+            and "reasoning" in original_output.intermediate_outputs
+        ):
+            return original_output
+
         # This parser only works for strings
         if not isinstance(original_output.output, str):
             raise ValueError("Response must be a string for R1 parser")
 
         # Strip whitespace and validate basic structure
         cleaned_response = original_output.output.strip()
-        if not cleaned_response.startswith(self.START_TAG):
-            raise ValueError("Response must start with <think> tag")
 
         # Find the thinking tags
-        think_start = cleaned_response.find(self.START_TAG)
         think_end = cleaned_response.find(self.END_TAG)
+        if think_end == -1:
+            raise ValueError("Missing </think> tag")
 
-        if think_start == -1 or think_end == -1:
-            raise ValueError("Missing thinking tags")
+        think_tag_start = cleaned_response.find(self.START_TAG)
+        if think_tag_start == -1:
+            # We allow no start <think>, thinking starts on first char. QwQ does this.
+            think_start = 0
+        else:
+            think_start = think_tag_start + len(self.START_TAG)
 
         # Check for multiple tags
         if (
@@ -44,9 +56,7 @@ class R1ThinkingParser(BaseParser):
             raise ValueError("Multiple thinking tags found")
 
         # Extract thinking content
-        thinking_content = cleaned_response[
-            think_start + len(self.START_TAG) : think_end
-        ].strip()
+        thinking_content = cleaned_response[think_start:think_end].strip()
 
         # Extract result (everything after </think>)
         result = cleaned_response[think_end + len(self.END_TAG) :].strip()
@@ -54,16 +64,11 @@ class R1ThinkingParser(BaseParser):
         if not result or len(result) == 0:
             raise ValueError("No content found after </think> tag")
 
-        # Parse JSON if needed
-        output = result
-        if self.structured_output:
-            output = parse_json_string(result)
-
         # Add thinking content to intermediate outputs if it exists
         intermediate_outputs = original_output.intermediate_outputs or {}
         intermediate_outputs["reasoning"] = thinking_content
 
         return RunOutput(
-            output=output,
+            output=result,
             intermediate_outputs=intermediate_outputs,
         )

kiln_ai/adapters/parsers/test_r1_parser.py

@@ -19,6 +19,16 @@ def test_valid_response(parser):
     assert parsed.output == "This is the result"
 
 
+def test_already_parsed_response(parser):
+    response = RunOutput(
+        output="This is the result",
+        intermediate_outputs={"reasoning": "This is thinking content"},
+    )
+    parsed = parser.parse_output(response)
+    assert parsed.intermediate_outputs["reasoning"] == "This is thinking content"
+    assert parsed.output == "This is the result"
+
+
 def test_response_with_whitespace(parser):
     response = RunOutput(
         output="""
@@ -37,14 +47,16 @@ def test_response_with_whitespace(parser):
 
 
 def test_missing_start_tag(parser):
-    with pytest.raises(ValueError, match="Response must start with <think> tag"):
-        parser.parse_output(
-            RunOutput(output="Some content</think>result", intermediate_outputs=None)
-        )
+    parsed = parser.parse_output(
+        RunOutput(output="Some content</think>result", intermediate_outputs=None)
+    )
+
+    assert parsed.intermediate_outputs["reasoning"] == "Some content"
+    assert parsed.output == "result"
 
 
 def test_missing_end_tag(parser):
-    with pytest.raises(ValueError, match="Missing thinking tags"):
+    with pytest.raises(ValueError, match="Missing </think> tag"):
         parser.parse_output(
             RunOutput(output="<think>Some content", intermediate_outputs=None)
         )

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/provider_tools.py

@@ -9,8 +9,8 @@ from kiln_ai.adapters.ml_model_list import (
     StructuredOutputMode,
     built_in_models,
 )
-from kiln_ai.adapters.model_adapters.openai_compatible_config import (
-    OpenAICompatibleConfig,
+from kiln_ai.adapters.model_adapters.litellm_config import (
+    LiteLlmConfig,
 )
 from kiln_ai.adapters.ollama_tools import (
     get_ollama_connection,
@@ -153,7 +153,7 @@ def kiln_model_provider_from(
         return finetune_provider_model(name)
 
     if provider_name == ModelProviderName.openai_compatible:
-        return openai_compatible_provider_model(name)
+        return lite_llm_provider_model(name)
 
     built_in_model = builtin_model_from(name, provider_name)
     if built_in_model:
@@ -175,13 +175,13 @@ def kiln_model_provider_from(
         supports_structured_output=False,
         supports_data_gen=False,
         untested_model=True,
-        provider_options=provider_options_for_custom_model(name, provider_name),
+        model_id=name,
     )
 
 
-def openai_compatible_config(
+def lite_llm_config(
     model_id: str,
-) -> OpenAICompatibleConfig:
+) -> LiteLlmConfig:
     try:
         openai_provider_name, model_id = model_id.split("::")
     except Exception:
@@ -205,22 +205,23 @@ def openai_compatible_config(
             f"OpenAI compatible provider {openai_provider_name} has no base URL"
         )
 
-    return OpenAICompatibleConfig(
-        api_key=api_key,
+    return LiteLlmConfig(
+        # OpenAI compatible, with a custom base URL
         model_name=model_id,
         provider_name=ModelProviderName.openai_compatible,
         base_url=base_url,
+        additional_body_options={
+            "api_key": api_key,
+        },
     )
 
 
-def openai_compatible_provider_model(
+def lite_llm_provider_model(
     model_id: str,
 ) -> KilnModelProvider:
     return KilnModelProvider(
         name=ModelProviderName.openai_compatible,
-        provider_options={
-            "model": model_id,
-        },
+        model_id=model_id,
         supports_structured_output=False,
         supports_data_gen=False,
         untested_model=True,
@@ -264,9 +265,7 @@ def finetune_provider_model(
     provider = ModelProviderName[fine_tune.provider]
     model_provider = KilnModelProvider(
         name=provider,
-        provider_options={
-            "model": fine_tune.fine_tune_model_id,
-        },
+        model_id=fine_tune.fine_tune_model_id,
     )
 
     if fine_tune.structured_output_mode is not None:
@@ -331,6 +330,18 @@ def provider_name_from_id(id: str) -> str:
                 return "Custom Models"
             case ModelProviderName.openai_compatible:
                 return "OpenAI Compatible"
+            case ModelProviderName.azure_openai:
+                return "Azure OpenAI"
+            case ModelProviderName.gemini_api:
+                return "Gemini API"
+            case ModelProviderName.anthropic:
+                return "Anthropic"
+            case ModelProviderName.huggingface:
+                return "Hugging Face"
+            case ModelProviderName.vertex:
+                return "Google Vertex AI"
+            case ModelProviderName.together_ai:
+                return "Together AI"
             case _:
                 # triggers pyright warning if I miss a case
                 raise_exhaustive_enum_error(enum_id)
@@ -338,49 +349,6 @@ def provider_name_from_id(id: str) -> str:
     return "Unknown provider: " + id
 
 
-def provider_options_for_custom_model(
-    model_name: str, provider_name: str
-) -> Dict[str, str]:
-    """
-    Generated model provider options for a custom model. Each has their own format/options.
-    """
-
-    if provider_name not in ModelProviderName.__members__:
-        raise ValueError(f"Invalid provider name: {provider_name}")
-
-    enum_id = ModelProviderName(provider_name)
-    match enum_id:
-        case ModelProviderName.amazon_bedrock:
-            # us-west-2 is the only region consistently supported by Bedrock
-            return {"model": model_name, "region_name": "us-west-2"}
-        case (
-            ModelProviderName.openai
-            | ModelProviderName.ollama
-            | ModelProviderName.fireworks_ai
-            | ModelProviderName.openrouter
-            | ModelProviderName.groq
-        ):
-            return {"model": model_name}
-        case ModelProviderName.kiln_custom_registry:
-            raise ValueError(
-                "Custom models from registry should be parsed into provider/model before calling this."
-            )
-        case ModelProviderName.kiln_fine_tune:
-            raise ValueError(
-                "Fine tuned models should populate provider options via another path"
-            )
-        case ModelProviderName.openai_compatible:
-            raise ValueError(
-                "OpenAI compatible models should populate provider options via another path"
-            )
-        case _:
-            # triggers pyright warning if I miss a case
-            raise_exhaustive_enum_error(enum_id)
-
-    # Won't reach this, type checking will catch missed values
-    return {"model": model_name}
-
-
 @dataclass
 class ModelProviderWarning:
     required_config_keys: List[str]
@@ -408,4 +376,28 @@ provider_warnings: Dict[ModelProviderName, ModelProviderWarning] = {
         required_config_keys=["fireworks_api_key", "fireworks_account_id"],
         message="Attempted to use Fireworks without an API key and account ID set. \nGet your API key from https://fireworks.ai/account/api-keys and your account ID from https://fireworks.ai/account/profile",
     ),
+    ModelProviderName.anthropic: ModelProviderWarning(
+        required_config_keys=["anthropic_api_key"],
+        message="Attempted to use Anthropic without an API key set. \nGet your API key from https://console.anthropic.com/settings/keys",
+    ),
+    ModelProviderName.gemini_api: ModelProviderWarning(
+        required_config_keys=["gemini_api_key"],
+        message="Attempted to use Gemini without an API key set. \nGet your API key from https://aistudio.google.com/app/apikey",
+    ),
+    ModelProviderName.azure_openai: ModelProviderWarning(
+        required_config_keys=["azure_openai_api_key", "azure_openai_endpoint"],
+        message="Attempted to use Azure OpenAI without an API key and endpoint set. Configure these in settings.",
+    ),
+    ModelProviderName.huggingface: ModelProviderWarning(
+        required_config_keys=["huggingface_api_key"],
+        message="Attempted to use Hugging Face without an API key set. \nGet your API key from https://huggingface.co/settings/tokens",
+    ),
+    ModelProviderName.vertex: ModelProviderWarning(
+        required_config_keys=["vertex_project_id"],
+        message="Attempted to use Vertex without a project ID set. \nGet your project ID from the Vertex AI console.",
+    ),
+    ModelProviderName.together_ai: ModelProviderWarning(
+        required_config_keys=["together_api_key"],
+        message="Attempted to use Together without an API key set. \nGet your API key from https://together.ai/settings/keys",
+    ),
 }