kiln-ai 0.0.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

kiln_ai/adapters/ml_model_list.py

@@ -0,0 +1,436 @@
+import os
+from dataclasses import dataclass
+from enum import Enum
+from os import getenv
+from typing import Dict, List, NoReturn
+
+import httpx
+from langchain_aws import ChatBedrockConverse
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_groq import ChatGroq
+from langchain_ollama import ChatOllama
+from langchain_openai import ChatOpenAI
+from pydantic import BaseModel
+
+from ..utils.config import Config
+
+
+class ModelProviderName(str, Enum):
+    openai = "openai"
+    groq = "groq"
+    amazon_bedrock = "amazon_bedrock"
+    ollama = "ollama"
+    openrouter = "openrouter"
+
+
+class ModelFamily(str, Enum):
+    gpt = "gpt"
+    llama = "llama"
+    phi = "phi"
+    mistral = "mistral"
+    gemma = "gemma"
+
+
+# Where models have instruct and raw versions, instruct is default and raw is specified
+class ModelName(str, Enum):
+    llama_3_1_8b = "llama_3_1_8b"
+    llama_3_1_70b = "llama_3_1_70b"
+    llama_3_1_405b = "llama_3_1_405b"
+    gpt_4o_mini = "gpt_4o_mini"
+    gpt_4o = "gpt_4o"
+    phi_3_5 = "phi_3_5"
+    mistral_large = "mistral_large"
+    mistral_nemo = "mistral_nemo"
+    gemma_2_2b = "gemma_2_2b"
+    gemma_2_9b = "gemma_2_9b"
+    gemma_2_27b = "gemma_2_27b"
+
+
+class KilnModelProvider(BaseModel):
+    name: ModelProviderName
+    # Allow overriding the model level setting
+    supports_structured_output: bool = True
+    provider_options: Dict = {}
+
+
+class KilnModel(BaseModel):
+    family: str
+    name: str
+    friendly_name: str
+    providers: List[KilnModelProvider]
+    supports_structured_output: bool = True
+
+
+built_in_models: List[KilnModel] = [
+    # GPT 4o Mini
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_4o_mini,
+        friendly_name="GPT 4o Mini",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                provider_options={"model": "gpt-4o-mini"},
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "openai/gpt-4o-mini"},
+            ),
+        ],
+    ),
+    # GPT 4o
+    KilnModel(
+        family=ModelFamily.gpt,
+        name=ModelName.gpt_4o,
+        friendly_name="GPT 4o",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openai,
+                provider_options={"model": "gpt-4o"},
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "openai/gpt-4o-2024-08-06"},
+            ),
+        ],
+    ),
+    # Llama 3.1-8b
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.llama_3_1_8b,
+        friendly_name="Llama 3.1 8B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.groq,
+                provider_options={"model": "llama-3.1-8b-instant"},
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.amazon_bedrock,
+                provider_options={
+                    "model": "meta.llama3-1-8b-instruct-v1:0",
+                    "region_name": "us-west-2",  # Llama 3.1 only in west-2
+                },
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={"model": "llama3.1"},  # 8b is default
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "meta-llama/llama-3.1-8b-instruct"},
+            ),
+        ],
+    ),
+    # Llama 3.1 70b
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.llama_3_1_70b,
+        friendly_name="Llama 3.1 70B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.groq,
+                provider_options={"model": "llama-3.1-70b-versatile"},
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.amazon_bedrock,
+                # TODO: this should work but a bug in the bedrock response schema
+                supports_structured_output=False,
+                provider_options={
+                    "model": "meta.llama3-1-70b-instruct-v1:0",
+                    "region_name": "us-west-2",  # Llama 3.1 only in west-2
+                },
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "meta-llama/llama-3.1-70b-instruct"},
+            ),
+            # TODO: enable once tests update to check if model is available
+            # KilnModelProvider(
+            #     provider=ModelProviders.ollama,
+            #     provider_options={"model": "llama3.1:70b"},
+            # ),
+        ],
+    ),
+    # Llama 3.1 405b
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.llama_3_1_405b,
+        friendly_name="Llama 3.1 405B",
+        providers=[
+            # TODO: bring back when groq does: https://console.groq.com/docs/models
+            # KilnModelProvider(
+            #     name=ModelProviderName.groq,
+            #     provider_options={"model": "llama-3.1-405b-instruct-v1:0"},
+            # ),
+            KilnModelProvider(
+                name=ModelProviderName.amazon_bedrock,
+                provider_options={
+                    "model": "meta.llama3-1-405b-instruct-v1:0",
+                    "region_name": "us-west-2",  # Llama 3.1 only in west-2
+                },
+            ),
+            # TODO: enable once tests update to check if model is available
+            # KilnModelProvider(
+            #     name=ModelProviderName.ollama,
+            #     provider_options={"model": "llama3.1:405b"},
+            # ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "meta-llama/llama-3.1-405b-instruct"},
+            ),
+        ],
+    ),
+    # Mistral Nemo
+    KilnModel(
+        family=ModelFamily.mistral,
+        name=ModelName.mistral_nemo,
+        friendly_name="Mistral Nemo",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "mistralai/mistral-nemo"},
+            ),
+        ],
+    ),
+    # Mistral Large
+    KilnModel(
+        family=ModelFamily.mistral,
+        name=ModelName.mistral_large,
+        friendly_name="Mistral Large",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.amazon_bedrock,
+                provider_options={
+                    "model": "mistral.mistral-large-2407-v1:0",
+                    "region_name": "us-west-2",  # only in west-2
+                },
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "mistralai/mistral-large"},
+            ),
+            # TODO: enable once tests update to check if model is available
+            # KilnModelProvider(
+            #     provider=ModelProviders.ollama,
+            #     provider_options={"model": "mistral-large"},
+            # ),
+        ],
+    ),
+    # Phi 3.5
+    KilnModel(
+        family=ModelFamily.phi,
+        name=ModelName.phi_3_5,
+        friendly_name="Phi 3.5",
+        supports_structured_output=False,
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={"model": "phi3.5"},
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "microsoft/phi-3.5-mini-128k-instruct"},
+            ),
+        ],
+    ),
+    # Gemma 2 2.6b
+    KilnModel(
+        family=ModelFamily.gemma,
+        name=ModelName.gemma_2_2b,
+        friendly_name="Gemma 2 2B",
+        supports_structured_output=False,
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={
+                    "model": "gemma2:2b",
+                },
+            ),
+        ],
+    ),
+    # Gemma 2 9b
+    KilnModel(
+        family=ModelFamily.gemma,
+        name=ModelName.gemma_2_9b,
+        friendly_name="Gemma 2 9B",
+        supports_structured_output=False,
+        providers=[
+            # TODO: enable once tests update to check if model is available
+            # KilnModelProvider(
+            #     name=ModelProviderName.ollama,
+            #     provider_options={
+            #         "model": "gemma2:9b",
+            #     },
+            # ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "google/gemma-2-9b-it"},
+            ),
+        ],
+    ),
+    # Gemma 2 27b
+    KilnModel(
+        family=ModelFamily.gemma,
+        name=ModelName.gemma_2_27b,
+        friendly_name="Gemma 2 27B",
+        supports_structured_output=False,
+        providers=[
+            # TODO: enable once tests update to check if model is available
+            # KilnModelProvider(
+            #     name=ModelProviderName.ollama,
+            #     provider_options={
+            #         "model": "gemma2:27b",
+            #     },
+            # ),
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "google/gemma-2-27b-it"},
+            ),
+        ],
+    ),
+]
+
+
+def provider_name_from_id(id: str) -> str:
+    if id in ModelProviderName.__members__:
+        enum_id = ModelProviderName(id)
+        match enum_id:
+            case ModelProviderName.amazon_bedrock:
+                return "Amazon Bedrock"
+            case ModelProviderName.openrouter:
+                return "OpenRouter"
+            case ModelProviderName.groq:
+                return "Groq"
+            case ModelProviderName.ollama:
+                return "Ollama"
+            case ModelProviderName.openai:
+                return "OpenAI"
+            case _:
+                # triggers pyright warning if I miss a case
+                raise_exhaustive_error(enum_id)
+
+    return "Unknown provider: " + id
+
+
+def raise_exhaustive_error(value: NoReturn) -> NoReturn:
+    raise ValueError(f"Unhandled enum value: {value}")
+
+
+@dataclass
+class ModelProviderWarning:
+    required_config_keys: List[str]
+    message: str
+
+
+provider_warnings: Dict[ModelProviderName, ModelProviderWarning] = {
+    ModelProviderName.amazon_bedrock: ModelProviderWarning(
+        required_config_keys=["bedrock_access_key", "bedrock_secret_key"],
+        message="Attempted to use Amazon Bedrock without an access key and secret set. \nGet your keys from https://us-west-2.console.aws.amazon.com/bedrock/home?region=us-west-2#/overview",
+    ),
+    ModelProviderName.openrouter: ModelProviderWarning(
+        required_config_keys=["open_router_api_key"],
+        message="Attempted to use OpenRouter without an API key set. \nGet your API key from https://openrouter.ai/settings/keys",
+    ),
+    ModelProviderName.groq: ModelProviderWarning(
+        required_config_keys=["groq_api_key"],
+        message="Attempted to use Groq without an API key set. \nGet your API key from https://console.groq.com/keys",
+    ),
+    ModelProviderName.openai: ModelProviderWarning(
+        required_config_keys=["open_ai_api_key"],
+        message="Attempted to use OpenAI without an API key set. \nGet your API key from https://platform.openai.com/account/api-keys",
+    ),
+}
+
+
+def get_config_value(key: str):
+    try:
+        return Config.shared().__getattr__(key)
+    except AttributeError:
+        return None
+
+
+def check_provider_warnings(provider_name: ModelProviderName):
+    warning_check = provider_warnings.get(provider_name)
+    if warning_check is None:
+        return
+    for key in warning_check.required_config_keys:
+        if get_config_value(key) is None:
+            raise ValueError(warning_check.message)
+
+
+def langchain_model_from(name: str, provider_name: str | None = None) -> BaseChatModel:
+    if name not in ModelName.__members__:
+        raise ValueError(f"Invalid name: {name}")
+
+    # Select the model from built_in_models using the name
+    model = next(filter(lambda m: m.name == name, built_in_models))
+    if model is None:
+        raise ValueError(f"Model {name} not found")
+
+    # If a provider is provided, select the provider from the model's provider_config
+    provider: KilnModelProvider | None = None
+    if model.providers is None or len(model.providers) == 0:
+        raise ValueError(f"Model {name} has no providers")
+    elif provider_name is None:
+        # TODO: priority order
+        provider = model.providers[0]
+    else:
+        provider = next(
+            filter(lambda p: p.name == provider_name, model.providers), None
+        )
+    if provider is None:
+        raise ValueError(f"Provider {provider_name} not found for model {name}")
+
+    check_provider_warnings(provider.name)
+
+    if provider.name == ModelProviderName.openai:
+        api_key = Config.shared().open_ai_api_key
+        return ChatOpenAI(**provider.provider_options, openai_api_key=api_key)  # type: ignore[arg-type]
+    elif provider.name == ModelProviderName.groq:
+        api_key = Config.shared().groq_api_key
+        if api_key is None:
+            raise ValueError(
+                "Attempted to use Groq without an API key set. "
+                "Get your API key from https://console.groq.com/keys"
+            )
+        return ChatGroq(**provider.provider_options, groq_api_key=api_key)  # type: ignore[arg-type]
+    elif provider.name == ModelProviderName.amazon_bedrock:
+        api_key = Config.shared().bedrock_access_key
+        secret_key = Config.shared().bedrock_secret_key
+        # langchain doesn't allow passing these, so ugly hack to set env vars
+        os.environ["AWS_ACCESS_KEY_ID"] = api_key
+        os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
+        return ChatBedrockConverse(
+            **provider.provider_options,
+        )
+    elif provider.name == ModelProviderName.ollama:
+        return ChatOllama(**provider.provider_options, base_url=ollama_base_url())
+    elif provider.name == ModelProviderName.openrouter:
+        api_key = Config.shared().open_router_api_key
+        base_url = getenv("OPENROUTER_BASE_URL") or "https://openrouter.ai/api/v1"
+        return ChatOpenAI(
+            **provider.provider_options,
+            openai_api_key=api_key,  # type: ignore[arg-type]
+            openai_api_base=base_url,  # type: ignore[arg-type]
+            default_headers={
+                "HTTP-Referer": "https://kiln-ai.com/openrouter",
+                "X-Title": "KilnAI",
+            },
+        )
+    else:
+        raise ValueError(f"Invalid model or provider: {name} - {provider_name}")
+
+
+def ollama_base_url():
+    env_base_url = os.getenv("OLLAMA_BASE_URL")
+    if env_base_url is not None:
+        return env_base_url
+    return "http://localhost:11434"
+
+
+async def ollama_online():
+    try:
+        httpx.get(ollama_base_url() + "/api/tags")
+    except httpx.RequestError:
+        return False
+    return True

kiln_ai/adapters/prompt_builders.py

@@ -0,0 +1,122 @@
+import json
+from abc import ABCMeta, abstractmethod
+from typing import Dict
+
+from kiln_ai.datamodel import Task
+from kiln_ai.utils.formatting import snake_case
+
+
+class BasePromptBuilder(metaclass=ABCMeta):
+    def __init__(self, task: Task):
+        self.task = task
+
+    @abstractmethod
+    def build_prompt(self) -> str:
+        pass
+
+    # override to change the name of the prompt builder (if changing class names)
+    @classmethod
+    def prompt_builder_name(cls) -> str:
+        return snake_case(cls.__name__)
+
+    # Can be overridden to add more information to the user message
+    def build_user_message(self, input: Dict | str) -> str:
+        if isinstance(input, Dict):
+            return f"The input is:\n{json.dumps(input, indent=2)}"
+
+        return f"The input is:\n{input}"
+
+
+class SimplePromptBuilder(BasePromptBuilder):
+    def build_prompt(self) -> str:
+        base_prompt = self.task.instruction
+
+        # TODO: this is just a quick version. Formatting and best practices TBD
+        if len(self.task.requirements) > 0:
+            base_prompt += (
+                "\n\nYour response should respect the following requirements:\n"
+            )
+            # 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"
+
+        return base_prompt
+
+
+class MultiShotPromptBuilder(BasePromptBuilder):
+    @classmethod
+    def example_count(cls) -> int:
+        return 25
+
+    def build_prompt(self) -> str:
+        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 += "\n"
+
+        valid_examples: list[tuple[str, str]] = []
+        runs = self.task.runs()
+
+        # first pass, we look for repaired outputs. These are the best examples.
+        for run in runs:
+            if len(valid_examples) >= self.__class__.example_count():
+                break
+            if run.repaired_output is not None:
+                valid_examples.append((run.input, run.repaired_output.output))
+
+        # second pass, we look for high quality outputs (rating based)
+        # Minimum is "high_quality" (4 star in star rating scale), then sort by rating
+        # exclude repaired outputs as they were used above
+        runs_with_rating = [
+            run
+            for run in runs
+            if run.output.rating is not None
+            and run.output.rating.value is not None
+            and run.output.rating.is_high_quality()
+            and run.repaired_output is None
+        ]
+        runs_with_rating.sort(
+            key=lambda x: (x.output.rating and x.output.rating.value) or 0, reverse=True
+        )
+        for run in runs_with_rating:
+            if len(valid_examples) >= self.__class__.example_count():
+                break
+            valid_examples.append((run.input, run.output.output))
+
+        if len(valid_examples) > 0:
+            base_prompt += "# Example Outputs\n\n"
+            for i, example in enumerate(valid_examples):
+                base_prompt += (
+                    f"## Example {i+1}\n\nInput: {example[0]}\nOutput: {example[1]}\n\n"
+                )
+
+        return base_prompt
+
+
+class FewShotPromptBuilder(MultiShotPromptBuilder):
+    @classmethod
+    def example_count(cls) -> int:
+        return 4
+
+
+prompt_builder_registry = {
+    "simple_prompt_builder": SimplePromptBuilder,
+    "multi_shot_prompt_builder": MultiShotPromptBuilder,
+    "few_shot_prompt_builder": FewShotPromptBuilder,
+}
+
+
+# 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]:
+    match ui_name:
+        case "basic":
+            return SimplePromptBuilder
+        case "few_shot":
+            return FewShotPromptBuilder
+        case "many_shot":
+            return MultiShotPromptBuilder
+        case _:
+            raise ValueError(f"Unknown prompt builder: {ui_name}")

kiln_ai/adapters/repair/repair_task.py

@@ -0,0 +1,71 @@
+import json
+from typing import Type
+
+from kiln_ai.adapters.prompt_builders import BasePromptBuilder, prompt_builder_registry
+from kiln_ai.datamodel import Priority, Project, Task, TaskRequirement, TaskRun
+from pydantic import BaseModel, Field
+
+
+# TODO add evaluator rating
+class RepairTaskInput(BaseModel):
+    original_prompt: str
+    original_input: str
+    original_output: str
+    evaluator_feedback: str = Field(
+        min_length=1,
+        description="Feedback from an evaluator on how to repair the task run.",
+    )
+
+
+class RepairTaskRun(Task, parent_of={}):
+    def __init__(self, original_task: Task):
+        # Keep the typechecker happy
+        tmp_project = Project(name="Repair")
+        super().__init__(
+            name="Repair",
+            parent=tmp_project,
+            description="Repair a task run, given feedback from an evaluator about how the response can be improved.",
+            instruction="You are an assistant which helps improve output from another assistant (original assistant). You'll be provided a task that the original assistant executed (prompt), \
+the input it was given, and the output it generated. An evaluator has determined that the output it generated did not satisfy the task and should be improved. The evaluator will provide \
+feedback describing what should be improved. Your job is to understand the evaluator's feedback and improve the response.",
+            requirements=[
+                TaskRequirement(
+                    name="Follow Eval Feedback",
+                    instruction="The evaluator's feedback is the most important thing to consider. If it conflicts with the original task instruction or prompt, prioritize the evaluator's feedback.",
+                    priority=Priority.p0,
+                )
+            ],
+            input_json_schema=json.dumps(RepairTaskInput.model_json_schema()),
+            output_json_schema=original_task.output_json_schema,
+        )
+
+    @classmethod
+    def _original_prompt(cls, run: TaskRun, task: Task) -> str:
+        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()
+
+    @classmethod
+    def build_repair_task_input(
+        cls, original_task: Task, task_run: TaskRun, evaluator_feedback: str
+    ) -> RepairTaskInput:
+        original_prompt = cls._original_prompt(task_run, original_task)
+        return RepairTaskInput(
+            original_prompt=original_prompt,
+            original_input=task_run.input,
+            original_output=task_run.output.output,
+            evaluator_feedback=evaluator_feedback,
+        )