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

kiln_ai/adapters/prompt_builders.py

@@ -20,8 +20,32 @@ class BasePromptBuilder(metaclass=ABCMeta):
         """
         self.task = task
 
+    def prompt_id(self) -> str | None:
+        """Returns the ID of the prompt, scoped to this builder.
+
+        Returns:
+            str | None: The ID of the prompt, or None if not set.
+        """
+        return None
+
+    def build_prompt(self, include_json_instructions) -> str:
+        """Build and return the complete prompt string.
+
+        Returns:
+            str: The constructed prompt.
+        """
+        prompt = self.build_base_prompt()
+
+        if include_json_instructions and self.task.output_schema():
+            prompt = (
+                prompt
+                + f"\n\n# Format Instructions\n\nReturn a JSON object conforming to the following schema:\n```\n{self.task.output_schema()}\n```"
+            )
+
+        return prompt
+
     @abstractmethod
-    def build_prompt(self) -> str:
+    def build_base_prompt(self) -> str:
         """Build and return the complete prompt string.
 
         Returns:
@@ -50,7 +74,7 @@ class BasePromptBuilder(metaclass=ABCMeta):
             str: The formatted user message.
         """
         if isinstance(input, Dict):
-            return f"The input is:\n{json.dumps(input, indent=2)}"
+            return f"The input is:\n{json.dumps(input, indent=2, ensure_ascii=False)}"
 
         return f"The input is:\n{input}"
 
@@ -70,7 +94,7 @@ class BasePromptBuilder(metaclass=ABCMeta):
         Returns:
             str: The constructed prompt string.
         """
-        base_prompt = self.build_prompt()
+        base_prompt = self.build_prompt(include_json_instructions=False)
         cot_prompt = self.chain_of_thought_prompt()
         if cot_prompt:
             base_prompt += "\n# Thinking Instructions\n\n" + cot_prompt
@@ -80,7 +104,7 @@ class BasePromptBuilder(metaclass=ABCMeta):
 class SimplePromptBuilder(BasePromptBuilder):
     """A basic prompt builder that combines task instruction with requirements."""
 
-    def build_prompt(self) -> str:
+    def build_base_prompt(self) -> str:
         """Build a simple prompt with instruction and requirements.
 
         Returns:
@@ -95,7 +119,7 @@ class SimplePromptBuilder(BasePromptBuilder):
             )
             # iterate requirements, formatting them in numbereed list like 1) task.instruction\n2)...
             for i, requirement in enumerate(self.task.requirements):
-                base_prompt += f"{i+1}) {requirement.instruction}\n"
+                base_prompt += f"{i + 1}) {requirement.instruction}\n"
 
         return base_prompt
 
@@ -112,18 +136,18 @@ class MultiShotPromptBuilder(BasePromptBuilder):
         """
         return 25
 
-    def build_prompt(self) -> str:
+    def build_base_prompt(self) -> str:
         """Build a prompt with instruction, requirements, and multiple examples.
 
         Returns:
             str: The constructed prompt string with examples.
         """
-        base_prompt = f"# Instruction\n\n{ self.task.instruction }\n\n"
+        base_prompt = f"# Instruction\n\n{self.task.instruction}\n\n"
 
         if len(self.task.requirements) > 0:
             base_prompt += "# Requirements\n\nYour response should respect the following requirements:\n"
             for i, requirement in enumerate(self.task.requirements):
-                base_prompt += f"{i+1}) {requirement.instruction}\n"
+                base_prompt += f"{i + 1}) {requirement.instruction}\n"
             base_prompt += "\n"
 
         valid_examples = self.collect_examples()
@@ -140,11 +164,11 @@ class MultiShotPromptBuilder(BasePromptBuilder):
     def prompt_section_for_example(self, index: int, example: TaskRun) -> str:
         # Prefer repaired output if it exists, otherwise use the regular output
         output = example.repaired_output or example.output
-        return f"## Example {index+1}\n\nInput: {example.input}\nOutput: {output.output}\n\n"
+        return f"## Example {index + 1}\n\nInput: {example.input}\nOutput: {output.output}\n\n"
 
     def collect_examples(self) -> list[TaskRun]:
         valid_examples: list[TaskRun] = []
-        runs = self.task.runs()
+        runs = self.task.runs(readonly=True)
 
         # first pass, we look for repaired outputs. These are the best examples.
         for run in runs:
@@ -198,7 +222,7 @@ class RepairsPromptBuilder(MultiShotPromptBuilder):
         ):
             return super().prompt_section_for_example(index, example)
 
-        prompt_section = f"## Example {index+1}\n\nInput: {example.input}\n\n"
+        prompt_section = f"## Example {index + 1}\n\nInput: {example.input}\n\n"
         prompt_section += (
             f"Initial Output Which Was Insufficient: {example.output.output}\n\n"
         )
@@ -209,7 +233,7 @@ class RepairsPromptBuilder(MultiShotPromptBuilder):
         return prompt_section
 
 
-def chain_of_thought_prompt(task: Task) -> str | None:
+def chain_of_thought_prompt(task: Task) -> str:
     """Standard implementation to build and return the chain of thought prompt string.
 
     Returns:
@@ -244,6 +268,77 @@ class MultiShotChainOfThoughtPromptBuilder(MultiShotPromptBuilder):
         return chain_of_thought_prompt(self.task)
 
 
+class SavedPromptBuilder(BasePromptBuilder):
+    """A prompt builder that looks up a static prompt."""
+
+    def __init__(self, task: Task, prompt_id: str):
+        super().__init__(task)
+        prompt_model = next(
+            (
+                prompt
+                for prompt in task.prompts(readonly=True)
+                if prompt.id == prompt_id
+            ),
+            None,
+        )
+        if not prompt_model:
+            raise ValueError(f"Prompt ID not found: {prompt_id}")
+        self.prompt_model = prompt_model
+
+    def prompt_id(self) -> str | None:
+        return self.prompt_model.id
+
+    def build_base_prompt(self) -> str:
+        """Returns a saved prompt.
+
+        Returns:
+            str: The prompt string.
+        """
+        return self.prompt_model.prompt
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return self.prompt_model.chain_of_thought_instructions
+
+
+class FineTunePromptBuilder(BasePromptBuilder):
+    """A prompt builder that looks up a fine-tune prompt."""
+
+    def __init__(self, task: Task, nested_fine_tune_id: str):
+        super().__init__(task)
+
+        # IDs are in project_id::task_id::fine_tune_id format
+        self.full_fine_tune_id = nested_fine_tune_id
+        parts = nested_fine_tune_id.split("::")
+        if len(parts) != 3:
+            raise ValueError(
+                f"Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id', got: {nested_fine_tune_id}"
+            )
+        fine_tune_id = parts[2]
+
+        fine_tune_model = next(
+            (
+                fine_tune
+                for fine_tune in task.finetunes(readonly=True)
+                if fine_tune.id == fine_tune_id
+            ),
+            None,
+        )
+        if not fine_tune_model:
+            raise ValueError(f"Fine-tune ID not found: {fine_tune_id}")
+        self.fine_tune_model = fine_tune_model
+
+    def prompt_id(self) -> str | None:
+        return self.full_fine_tune_id
+
+    def build_base_prompt(self) -> str:
+        return self.fine_tune_model.system_message
+
+    def chain_of_thought_prompt(self) -> str | None:
+        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,
@@ -256,7 +351,7 @@ prompt_builder_registry = {
 
 
 # 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) -> type[BasePromptBuilder]:
+def prompt_builder_from_ui_name(ui_name: str, task: Task) -> BasePromptBuilder:
     """Convert a name used in the UI to the corresponding prompt builder class.
 
     Args:
@@ -268,20 +363,31 @@ def prompt_builder_from_ui_name(ui_name: str) -> type[BasePromptBuilder]:
     Raises:
         ValueError: If the UI name is not recognized.
     """
+
+    # Saved prompts are prefixed with "id::"
+    if ui_name.startswith("id::"):
+        prompt_id = ui_name[4:]
+        return SavedPromptBuilder(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)
+
     match ui_name:
         case "basic":
-            return SimplePromptBuilder
+            return SimplePromptBuilder(task)
         case "few_shot":
-            return FewShotPromptBuilder
+            return FewShotPromptBuilder(task)
         case "many_shot":
-            return MultiShotPromptBuilder
+            return MultiShotPromptBuilder(task)
         case "repairs":
-            return RepairsPromptBuilder
+            return RepairsPromptBuilder(task)
         case "simple_chain_of_thought":
-            return SimpleChainOfThoughtPromptBuilder
+            return SimpleChainOfThoughtPromptBuilder(task)
         case "few_shot_chain_of_thought":
-            return FewShotChainOfThoughtPromptBuilder
+            return FewShotChainOfThoughtPromptBuilder(task)
         case "multi_shot_chain_of_thought":
-            return MultiShotChainOfThoughtPromptBuilder
+            return MultiShotChainOfThoughtPromptBuilder(task)
         case _:
             raise ValueError(f"Unknown prompt builder: {ui_name}")

kiln_ai/adapters/provider_tools.py

@@ -1,20 +1,24 @@
 from dataclasses import dataclass
-from typing import Dict, List, NoReturn
+from typing import Dict, List
 
 from kiln_ai.adapters.ml_model_list import (
     KilnModel,
     KilnModelProvider,
     ModelName,
     ModelProviderName,
+    StructuredOutputMode,
     built_in_models,
 )
+from kiln_ai.adapters.model_adapters.openai_compatible_config import (
+    OpenAICompatibleConfig,
+)
 from kiln_ai.adapters.ollama_tools import (
     get_ollama_connection,
 )
 from kiln_ai.datamodel import Finetune, Task
 from kiln_ai.datamodel.registry import project_from_id
-
-from ..utils.config import Config
+from kiln_ai.utils.config import Config
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
 async def provider_enabled(provider_name: ModelProviderName) -> bool:
@@ -61,7 +65,7 @@ def check_provider_warnings(provider_name: ModelProviderName):
             raise ValueError(warning_check.message)
 
 
-async def builtin_model_from(
+def builtin_model_from(
     name: str, provider_name: str | None = None
 ) -> KilnModelProvider | None:
     """
@@ -102,7 +106,47 @@ async def builtin_model_from(
     return provider
 
 
-async def kiln_model_provider_from(
+def core_provider(model_id: str, provider_name: ModelProviderName) -> ModelProviderName:
+    """
+    Get the provider that should be run.
+
+    Some provider IDs are wrappers (fine-tunes, custom models). This maps these to runnable providers (openai, ollama, etc)
+    """
+
+    # Custom models map to the underlying provider
+    if provider_name is ModelProviderName.kiln_custom_registry:
+        provider_name, _ = parse_custom_model_id(model_id)
+        return provider_name
+
+    # Fine-tune provider maps to an underlying provider
+    if provider_name is ModelProviderName.kiln_fine_tune:
+        finetune = finetune_from_id(model_id)
+        if finetune.provider not in ModelProviderName.__members__:
+            raise ValueError(
+                f"Finetune {model_id} has no underlying provider {finetune.provider}"
+            )
+        return ModelProviderName(finetune.provider)
+
+    return provider_name
+
+
+def parse_custom_model_id(
+    model_id: str,
+) -> tuple[ModelProviderName, str]:
+    if "::" not in model_id:
+        raise ValueError(f"Invalid custom model ID: {model_id}")
+
+    # For custom registry, get the provider name and model name from the model id
+    provider_name = model_id.split("::", 1)[0]
+    model_name = model_id.split("::", 1)[1]
+
+    if provider_name not in ModelProviderName.__members__:
+        raise ValueError(f"Invalid provider name: {provider_name}")
+
+    return ModelProviderName(provider_name), model_name
+
+
+def kiln_model_provider_from(
     name: str, provider_name: str | None = None
 ) -> KilnModelProvider:
     if provider_name == ModelProviderName.kiln_fine_tune:
@@ -111,14 +155,13 @@ async def kiln_model_provider_from(
     if provider_name == ModelProviderName.openai_compatible:
         return openai_compatible_provider_model(name)
 
-    built_in_model = await builtin_model_from(name, provider_name)
+    built_in_model = builtin_model_from(name, provider_name)
     if built_in_model:
         return built_in_model
 
     # For custom registry, get the provider name and model name from the model id
     if provider_name == ModelProviderName.kiln_custom_registry:
-        provider_name = name.split("::", 1)[0]
-        name = name.split("::", 1)[1]
+        provider_name, name = parse_custom_model_id(name)
 
     # Custom/untested model. Set untested, and build a ModelProvider at runtime
     if provider_name is None:
@@ -136,12 +179,9 @@ async def kiln_model_provider_from(
     )
 
 
-finetune_cache: dict[str, KilnModelProvider] = {}
-
-
-def openai_compatible_provider_model(
+def openai_compatible_config(
     model_id: str,
-) -> KilnModelProvider:
+) -> OpenAICompatibleConfig:
     try:
         openai_provider_name, model_id = model_id.split("::")
     except Exception:
@@ -165,12 +205,21 @@ def openai_compatible_provider_model(
             f"OpenAI compatible provider {openai_provider_name} has no base URL"
         )
 
+    return OpenAICompatibleConfig(
+        api_key=api_key,
+        model_name=model_id,
+        provider_name=ModelProviderName.openai_compatible,
+        base_url=base_url,
+    )
+
+
+def openai_compatible_provider_model(
+    model_id: str,
+) -> KilnModelProvider:
     return KilnModelProvider(
         name=ModelProviderName.openai_compatible,
         provider_options={
             "model": model_id,
-            "api_key": api_key,
-            "openai_api_base": base_url,
         },
         supports_structured_output=False,
         supports_data_gen=False,
@@ -178,9 +227,10 @@ def openai_compatible_provider_model(
     )
 
 
-def finetune_provider_model(
-    model_id: str,
-) -> KilnModelProvider:
+finetune_cache: dict[str, Finetune] = {}
+
+
+def finetune_from_id(model_id: str) -> Finetune:
     if model_id in finetune_cache:
         return finetune_cache[model_id]
 
@@ -202,6 +252,15 @@ def finetune_provider_model(
             f"Fine tune {fine_tune_id} not completed. Refresh it's status in the fine-tune tab."
         )
 
+    finetune_cache[model_id] = fine_tune
+    return fine_tune
+
+
+def finetune_provider_model(
+    model_id: str,
+) -> KilnModelProvider:
+    fine_tune = finetune_from_id(model_id)
+
     provider = ModelProviderName[fine_tune.provider]
     model_provider = KilnModelProvider(
         name=provider,
@@ -210,18 +269,18 @@ def finetune_provider_model(
         },
     )
 
-    # TODO: Don't love this abstraction/logic.
-    if fine_tune.provider == ModelProviderName.fireworks_ai:
-        # Fireworks finetunes are trained with json, not tool calling (which is LC default format)
-        model_provider.adapter_options = {
-            "langchain": {
-                "with_structured_output_options": {
-                    "method": "json_mode",
-                }
-            }
-        }
-
-    finetune_cache[model_id] = model_provider
+    if fine_tune.structured_output_mode is not None:
+        # If we know the model was trained with specific output mode, set it
+        model_provider.structured_output_mode = fine_tune.structured_output_mode
+    else:
+        # Some early adopters won't have structured_output_mode set on their fine-tunes.
+        # We know that OpenAI uses json_schema, and Fireworks (only other provider) use json_mode.
+        # This can be removed in the future
+        if provider == ModelProviderName.openai:
+            model_provider.structured_output_mode = StructuredOutputMode.json_schema
+        else:
+            model_provider.structured_output_mode = StructuredOutputMode.json_mode
+
     return model_provider
 
 
@@ -274,7 +333,7 @@ def provider_name_from_id(id: str) -> str:
                 return "OpenAI Compatible"
             case _:
                 # triggers pyright warning if I miss a case
-                raise_exhaustive_error(enum_id)
+                raise_exhaustive_enum_error(enum_id)
 
     return "Unknown provider: " + id
 
@@ -316,16 +375,12 @@ def provider_options_for_custom_model(
             )
         case _:
             # triggers pyright warning if I miss a case
-            raise_exhaustive_error(enum_id)
+            raise_exhaustive_enum_error(enum_id)
 
     # Won't reach this, type checking will catch missed values
     return {"model": model_name}
 
 
-def raise_exhaustive_error(value: NoReturn) -> NoReturn:
-    raise ValueError(f"Unhandled enum value: {value}")
-
-
 @dataclass
 class ModelProviderWarning:
     required_config_keys: List[str]

kiln_ai/adapters/repair/repair_task.py

@@ -3,7 +3,11 @@ from typing import Type
 
 from pydantic import BaseModel, Field
 
-from kiln_ai.adapters.prompt_builders import BasePromptBuilder, prompt_builder_registry
+from kiln_ai.adapters.prompt_builders import (
+    BasePromptBuilder,
+    SavedPromptBuilder,
+    prompt_builder_registry,
+)
 from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
 
 
@@ -42,11 +46,18 @@ feedback describing what should be improved. Your job is to understand the evalu
 
     @classmethod
     def _original_prompt(cls, run: TaskRun, task: Task) -> str:
+        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)
+        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_class: Type[BasePromptBuilder] | None = None
-        prompt_builder_name = (
-            run.output.source.properties.get("prompt_builder_name", None)
-            if run.output.source
-            else 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(
@@ -59,7 +70,7 @@ feedback describing what should be improved. Your job is to understand the evalu
             raise ValueError(
                 f"Prompt builder {prompt_builder_name} is not a valid prompt builder"
             )
-        return prompt_builder.build_prompt()
+        return prompt_builder.build_prompt(include_json_instructions=False)
 
     @classmethod
     def build_repair_task_input(

kiln_ai/adapters/repair/test_repair_task.py

@@ -6,8 +6,8 @@ import pytest
 from pydantic import ValidationError
 
 from kiln_ai.adapters.adapter_registry import adapter_for_task
-from kiln_ai.adapters.base_adapter import RunOutput
-from kiln_ai.adapters.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.model_adapters.base_adapter import RunOutput
+from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
 from kiln_ai.adapters.repair.repair_task import (
     RepairTaskInput,
     RepairTaskRun,
@@ -223,7 +223,7 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
         )
 
         adapter = adapter_for_task(
-            repair_task, model_name="llama_3_1_8b", provider="groq"
+            repair_task, model_name="llama_3_1_8b", provider="ollama"
         )
 
         run = await adapter.invoke(repair_task_input.model_dump())
@@ -237,7 +237,7 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
     assert run.output.source.properties == {
         "adapter_name": "kiln_langchain_adapter",
         "model_name": "llama_3_1_8b",
-        "model_provider": "groq",
+        "model_provider": "ollama",
         "prompt_builder_name": "simple_prompt_builder",
     }
     assert run.input_source.type == DataSourceType.human

kiln_ai/adapters/run_output.py

@@ -0,0 +1,8 @@
+from dataclasses import dataclass
+from typing import Dict
+
+
+@dataclass
+class RunOutput:
+    output: Dict | str
+    intermediate_outputs: Dict[str, str] | None