kiln-ai 0.6.1__py3-none-any.whl → 0.7.1__py3-none-any.whl

kiln_ai/adapters/ollama_tools.py

@@ -0,0 +1,115 @@
+import os
+from typing import Any, List
+
+import httpx
+import requests
+from pydantic import BaseModel, Field
+
+from kiln_ai.adapters.ml_model_list import ModelProviderName, built_in_models
+from kiln_ai.utils.config import Config
+
+
+def ollama_base_url() -> str:
+    """
+    Gets the base URL for Ollama API connections.
+
+    Returns:
+        The base URL to use for Ollama API calls, using environment variable if set
+        or falling back to localhost default
+    """
+    config_base_url = Config.shared().ollama_base_url
+    if config_base_url:
+        return config_base_url
+    return "http://localhost:11434"
+
+
+async def ollama_online() -> bool:
+    """
+    Checks if the Ollama service is available and responding.
+
+    Returns:
+        True if Ollama is available and responding, False otherwise
+    """
+    try:
+        httpx.get(ollama_base_url() + "/api/tags")
+    except httpx.RequestError:
+        return False
+    return True
+
+
+class OllamaConnection(BaseModel):
+    message: str
+    supported_models: List[str]
+    untested_models: List[str] = Field(default_factory=list)
+
+    def all_models(self) -> List[str]:
+        return self.supported_models + self.untested_models
+
+
+# Parse the Ollama /api/tags response
+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"]
+        for model in built_in_models
+        for provider in model.providers
+        if provider.name == ModelProviderName.ollama
+    ]
+    # Append model_aliases to supported_ollama_models
+    supported_ollama_models.extend(
+        [
+            alias
+            for model in built_in_models
+            for provider in model.providers
+            for alias in provider.provider_options.get("model_aliases", [])
+        ]
+    )
+
+    if "models" in tags:
+        models = tags["models"]
+        if isinstance(models, list):
+            model_names = [model["model"] for model in models]
+            available_supported_models = []
+            untested_models = []
+            supported_models_latest_aliases = [
+                f"{m}:latest" for m in supported_ollama_models
+            ]
+            for model in model_names:
+                if (
+                    model in supported_ollama_models
+                    or model in supported_models_latest_aliases
+                ):
+                    available_supported_models.append(model)
+                else:
+                    untested_models.append(model)
+
+            if available_supported_models or untested_models:
+                return OllamaConnection(
+                    message="Ollama connected",
+                    supported_models=available_supported_models,
+                    untested_models=untested_models,
+                )
+
+    return OllamaConnection(
+        message="Ollama is running, but no supported models are installed. Install one or more supported model, like 'ollama pull phi3.5'.",
+        supported_models=[],
+        untested_models=[],
+    )
+
+
+async def get_ollama_connection() -> OllamaConnection | None:
+    """
+    Gets the connection status for Ollama.
+    """
+    try:
+        tags = requests.get(ollama_base_url() + "/api/tags", timeout=5).json()
+
+    except Exception:
+        return None
+
+    return parse_ollama_tags(tags)
+
+
+def ollama_model_installed(conn: OllamaConnection, model_name: str) -> bool:
+    all_models = conn.all_models()
+    return model_name in all_models or f"{model_name}:latest" in all_models

kiln_ai/adapters/provider_tools.py

@@ -0,0 +1,308 @@
+from dataclasses import dataclass
+from typing import Dict, List, NoReturn
+
+from kiln_ai.adapters.ml_model_list import (
+    KilnModel,
+    KilnModelProvider,
+    ModelName,
+    ModelProviderName,
+    built_in_models,
+)
+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
+
+
+async def provider_enabled(provider_name: ModelProviderName) -> bool:
+    if provider_name == ModelProviderName.ollama:
+        try:
+            conn = await get_ollama_connection()
+            return conn is not None and (
+                len(conn.supported_models) > 0 or len(conn.untested_models) > 0
+            )
+        except Exception:
+            return False
+
+    provider_warning = provider_warnings.get(provider_name)
+    if provider_warning is None:
+        return False
+    for required_key in provider_warning.required_config_keys:
+        if get_config_value(required_key) is None:
+            return False
+    return True
+
+
+def get_config_value(key: str):
+    try:
+        return Config.shared().__getattr__(key)
+    except AttributeError:
+        return None
+
+
+def check_provider_warnings(provider_name: ModelProviderName):
+    """
+    Validates that required configuration is present for a given provider.
+
+    Args:
+        provider_name: The provider to check
+
+    Raises:
+        ValueError: If required configuration keys are missing
+    """
+    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)
+
+
+async def builtin_model_from(
+    name: str, provider_name: str | None = None
+) -> KilnModelProvider | None:
+    """
+    Gets a model and provider from the built-in list of models.
+
+    Args:
+        name: The name of the model to get
+        provider_name: Optional specific provider to use (defaults to first available)
+
+    Returns:
+        A tuple of (provider, model)
+
+    Raises:
+        ValueError: If the model or provider is not found, or if the provider is misconfigured
+    """
+    if name not in ModelName.__members__:
+        return None
+
+    # 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:
+        provider = model.providers[0]
+    else:
+        provider = next(
+            filter(lambda p: p.name == provider_name, model.providers), None
+        )
+    if provider is None:
+        return None
+
+    check_provider_warnings(provider.name)
+    return provider
+
+
+async def kiln_model_provider_from(
+    name: str, provider_name: str | None = None
+) -> KilnModelProvider:
+    if provider_name == ModelProviderName.kiln_fine_tune:
+        return finetune_provider_model(name)
+
+    built_in_model = await 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]
+
+    # Custom/untested model. Set untested, and build a ModelProvider at runtime
+    if provider_name is None:
+        raise ValueError("Provider name is required for custom models")
+    if provider_name not in ModelProviderName.__members__:
+        raise ValueError(f"Invalid provider name: {provider_name}")
+    provider = ModelProviderName(provider_name)
+    check_provider_warnings(provider)
+    return KilnModelProvider(
+        name=provider,
+        supports_structured_output=False,
+        supports_data_gen=False,
+        untested_model=True,
+        provider_options=provider_options_for_custom_model(name, provider_name),
+    )
+
+
+finetune_cache: dict[str, KilnModelProvider] = {}
+
+
+def finetune_provider_model(
+    model_id: str,
+) -> KilnModelProvider:
+    if model_id in finetune_cache:
+        return finetune_cache[model_id]
+
+    try:
+        project_id, task_id, fine_tune_id = model_id.split("::")
+    except Exception:
+        raise ValueError(f"Invalid fine tune ID: {model_id}")
+    project = project_from_id(project_id)
+    if project is None:
+        raise ValueError(f"Project {project_id} not found")
+    task = Task.from_id_and_parent_path(task_id, project.path)
+    if task is None:
+        raise ValueError(f"Task {task_id} not found")
+    fine_tune = Finetune.from_id_and_parent_path(fine_tune_id, task.path)
+    if fine_tune is None:
+        raise ValueError(f"Fine tune {fine_tune_id} not found")
+    if fine_tune.fine_tune_model_id is None:
+        raise ValueError(
+            f"Fine tune {fine_tune_id} not completed. Refresh it's status in the fine-tune tab."
+        )
+
+    provider = ModelProviderName[fine_tune.provider]
+    model_provider = KilnModelProvider(
+        name=provider,
+        provider_options={
+            "model": fine_tune.fine_tune_model_id,
+        },
+    )
+
+    # 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
+    return model_provider
+
+
+def get_model_and_provider(
+    model_name: str, provider_name: str
+) -> tuple[KilnModel | None, KilnModelProvider | None]:
+    model = next(filter(lambda m: m.name == model_name, built_in_models), None)
+    if model is None:
+        return None, None
+    provider = next(filter(lambda p: p.name == provider_name, model.providers), None)
+    # all or nothing
+    if provider is None or model is None:
+        return None, None
+    return model, provider
+
+
+def provider_name_from_id(id: str) -> str:
+    """
+    Converts a provider ID to its human-readable name.
+
+    Args:
+        id: The provider identifier string
+
+    Returns:
+        The human-readable name of the provider
+
+    Raises:
+        ValueError: If the provider ID is invalid or unhandled
+    """
+    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 ModelProviderName.kiln_fine_tune:
+                return "Fine Tuned Models"
+            case ModelProviderName.fireworks_ai:
+                return "Fireworks AI"
+            case ModelProviderName.kiln_custom_registry:
+                return "Custom Models"
+            case _:
+                # triggers pyright warning if I miss a case
+                raise_exhaustive_error(enum_id)
+
+    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 _:
+            # triggers pyright warning if I miss a case
+            raise_exhaustive_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]
+    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",
+    ),
+    ModelProviderName.fireworks_ai: 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",
+    ),
+}

kiln_ai/adapters/repair/repair_task.py

@@ -43,8 +43,10 @@ feedback describing what should be improved. Your job is to understand the evalu
     @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
+        prompt_builder_name = (
+            run.output.source.properties.get("prompt_builder_name", None)
+            if run.output.source
+            else None
         )
         if prompt_builder_name is not None and isinstance(prompt_builder_name, str):
             prompt_builder_class = prompt_builder_registry.get(

kiln_ai/adapters/repair/test_repair_task.py

@@ -5,10 +5,9 @@ from unittest.mock import AsyncMock, patch
 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 (
-    LangChainPromptAdapter,
-)
+from kiln_ai.adapters.langchain_adapters import LangchainAdapter
 from kiln_ai.adapters.repair.repair_task import (
     RepairTaskInput,
     RepairTaskRun,
@@ -60,7 +59,7 @@ json_joke_schema = """{
 
 @pytest.fixture
 def sample_task(tmp_path):
-    task_path = tmp_path / "task.json"
+    task_path = tmp_path / "task.kiln"
     task = Task(
         name="Joke Generator",
         path=task_path,
@@ -190,9 +189,7 @@ async def test_live_run(sample_task, sample_task_run, sample_repair_data):
     repair_task_input = RepairTaskRun.build_repair_task_input(**sample_repair_data)
     assert isinstance(repair_task_input, RepairTaskInput)
 
-    adapter = LangChainPromptAdapter(
-        repair_task, model_name="llama_3_1_8b", provider="groq"
-    )
+    adapter = adapter_for_task(repair_task, model_name="llama_3_1_8b", provider="groq")
 
     run = await adapter.invoke(repair_task_input.model_dump())
     assert run is not None
@@ -220,14 +217,12 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
         "rating": 8,
     }
 
-    with patch.object(
-        LangChainPromptAdapter, "_run", new_callable=AsyncMock
-    ) as mock_run:
+    with patch.object(LangchainAdapter, "_run", new_callable=AsyncMock) as mock_run:
         mock_run.return_value = RunOutput(
             output=mocked_output, intermediate_outputs=None
         )
 
-        adapter = LangChainPromptAdapter(
+        adapter = adapter_for_task(
             repair_task, model_name="llama_3_1_8b", provider="groq"
         )