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

kiln_ai/adapters/adapter_registry.py

@@ -9,6 +9,7 @@ def adapter_for_task(
     model_name: str | None = None,
     provider: str | None = None,
     prompt_builder: BasePromptBuilder | None = None,
+    tags: list[str] | None = None,
 ) -> BaseAdapter:
     # We use langchain for everything right now, but can add any others here
     return LangchainAdapter(
@@ -16,4 +17,5 @@ def adapter_for_task(
         model_name=model_name,
         provider=provider,
         prompt_builder=prompt_builder,
+        tags=tags,
     )

kiln_ai/adapters/base_adapter.py

@@ -45,12 +45,16 @@ class BaseAdapter(metaclass=ABCMeta):
     """
 
     def __init__(
-        self, kiln_task: Task, prompt_builder: BasePromptBuilder | None = None
+        self,
+        kiln_task: Task,
+        prompt_builder: BasePromptBuilder | None = None,
+        tags: list[str] | None = None,
     ):
         self.prompt_builder = prompt_builder or SimplePromptBuilder(kiln_task)
         self.kiln_task = kiln_task
         self.output_schema = self.kiln_task.output_json_schema
         self.input_schema = self.kiln_task.input_json_schema
+        self.default_tags = tags
 
     async def invoke_returning_raw(
         self,
@@ -148,6 +152,7 @@ class BaseAdapter(metaclass=ABCMeta):
                 ),
             ),
             intermediate_outputs=run_output.intermediate_outputs,
+            tags=self.default_tags or [],
         )
 
         exclude_fields = {

kiln_ai/adapters/langchain_adapters.py

@@ -39,8 +39,9 @@ class LangchainAdapter(BaseAdapter):
         model_name: str | None = None,
         provider: str | None = None,
         prompt_builder: BasePromptBuilder | None = None,
+        tags: list[str] | None = None,
     ):
-        super().__init__(kiln_task, prompt_builder=prompt_builder)
+        super().__init__(kiln_task, prompt_builder=prompt_builder, tags=tags)
         if custom_model is not None:
             self._model = custom_model
 
@@ -198,6 +199,9 @@ async def langchain_model_from_provider(
     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.openai_compatible:
+        # See provider_tools.py for how base_url, key and other parameters are set
+        return ChatOpenAI(**provider.provider_options)  # type: ignore[arg-type]
     elif provider.name == ModelProviderName.groq:
         api_key = Config.shared().groq_api_key
         if api_key is None:

kiln_ai/adapters/ml_model_list.py

@@ -22,6 +22,8 @@ class ModelProviderName(str, Enum):
     openrouter = "openrouter"
     fireworks_ai = "fireworks_ai"
     kiln_fine_tune = "kiln_fine_tune"
+    kiln_custom_registry = "kiln_custom_registry"
+    openai_compatible = "openai_compatible"
 
 
 class ModelFamily(str, Enum):
@@ -54,6 +56,7 @@ class ModelName(str, Enum):
     llama_3_2_3b = "llama_3_2_3b"
     llama_3_2_11b = "llama_3_2_11b"
     llama_3_2_90b = "llama_3_2_90b"
+    llama_3_3_70b = "llama_3_3_70b"
     gpt_4o_mini = "gpt_4o_mini"
     gpt_4o = "gpt_4o"
     phi_3_5 = "phi_3_5"
@@ -502,6 +505,46 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
+    # Llama 3.3 70B
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.llama_3_3_70b,
+        friendly_name="Llama 3.3 70B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "meta-llama/llama-3.3-70b-instruct"},
+                # Openrouter not supporing tools yet. Once they do probably can remove. JSON mode sometimes works, but not consistently.
+                supports_structured_output=False,
+                supports_data_gen=False,
+                adapter_options={
+                    "langchain": {
+                        "with_structured_output_options": {"method": "json_mode"}
+                    }
+                },
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.groq,
+                supports_structured_output=True,
+                supports_data_gen=True,
+                provider_options={"model": "llama-3.3-70b-versatile"},
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={"model": "llama3.3"},
+            ),
+            KilnModelProvider(
+                name=ModelProviderName.fireworks_ai,
+                # Finetuning not live yet
+                # provider_finetune_id="accounts/fireworks/models/llama-v3p3-70b-instruct",
+                supports_structured_output=True,
+                supports_data_gen=True,
+                provider_options={
+                    "model": "accounts/fireworks/models/llama-v3p3-70b-instruct"
+                },
+            ),
+        ],
+    ),
     # Phi 3.5
     KilnModel(
         family=ModelFamily.phi,
@@ -598,18 +641,6 @@ built_in_models: List[KilnModel] = [
         name=ModelName.mixtral_8x7b,
         friendly_name="Mixtral 8x7B",
         providers=[
-            KilnModelProvider(
-                name=ModelProviderName.fireworks_ai,
-                provider_options={
-                    "model": "accounts/fireworks/models/mixtral-8x7b-instruct-hf",
-                },
-                provider_finetune_id="accounts/fireworks/models/mixtral-8x7b-instruct-hf",
-                adapter_options={
-                    "langchain": {
-                        "with_structured_output_options": {"method": "json_mode"}
-                    }
-                },
-            ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 provider_options={"model": "mistralai/mixtral-8x7b-instruct"},

kiln_ai/adapters/ollama_tools.py

@@ -6,6 +6,7 @@ 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:
@@ -16,9 +17,9 @@ def ollama_base_url() -> str:
         The base URL to use for Ollama API calls, using environment variable if set
         or falling back to localhost default
     """
-    env_base_url = os.getenv("OLLAMA_BASE_URL")
-    if env_base_url is not None:
-        return env_base_url
+    config_base_url = Config.shared().ollama_base_url
+    if config_base_url:
+        return config_base_url
     return "http://localhost:11434"
 
 

kiln_ai/adapters/provider_tools.py

@@ -11,6 +11,7 @@ from kiln_ai.adapters.ml_model_list import (
 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
@@ -107,10 +108,18 @@ async def kiln_model_provider_from(
     if provider_name == ModelProviderName.kiln_fine_tune:
         return finetune_provider_model(name)
 
+    if provider_name == ModelProviderName.openai_compatible:
+        return openai_compatible_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")
@@ -130,6 +139,45 @@ async def kiln_model_provider_from(
 finetune_cache: dict[str, KilnModelProvider] = {}
 
 
+def openai_compatible_provider_model(
+    model_id: str,
+) -> KilnModelProvider:
+    try:
+        openai_provider_name, model_id = model_id.split("::")
+    except Exception:
+        raise ValueError(f"Invalid openai compatible model ID: {model_id}")
+
+    openai_compatible_providers = Config.shared().openai_compatible_providers or []
+    provider = next(
+        filter(
+            lambda p: p.get("name") == openai_provider_name, openai_compatible_providers
+        ),
+        None,
+    )
+    if provider is None:
+        raise ValueError(f"OpenAI compatible provider {openai_provider_name} not found")
+
+    # API key optional some providers don't use it
+    api_key = provider.get("api_key")
+    base_url = provider.get("base_url")
+    if base_url is None:
+        raise ValueError(
+            f"OpenAI compatible provider {openai_provider_name} has no base URL"
+        )
+
+    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,
+        untested_model=True,
+    )
+
+
 def finetune_provider_model(
     model_id: str,
 ) -> KilnModelProvider:
@@ -143,10 +191,10 @@ def finetune_provider_model(
     project = project_from_id(project_id)
     if project is None:
         raise ValueError(f"Project {project_id} not found")
-    task = next((t for t in project.tasks() if t.id == task_id), None)
+    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 = next((f for f in task.finetunes() if f.id == fine_tune_id), None)
+    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:
@@ -220,6 +268,10 @@ def provider_name_from_id(id: str) -> str:
                 return "Fine Tuned Models"
             case ModelProviderName.fireworks_ai:
                 return "Fireworks AI"
+            case ModelProviderName.kiln_custom_registry:
+                return "Custom Models"
+            case ModelProviderName.openai_compatible:
+                return "OpenAI Compatible"
             case _:
                 # triggers pyright warning if I miss a case
                 raise_exhaustive_error(enum_id)
@@ -233,6 +285,7 @@ def provider_options_for_custom_model(
     """
     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}")
 
@@ -249,10 +302,18 @@ def provider_options_for_custom_model(
             | 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_error(enum_id)

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

@@ -1,12 +1,20 @@
+import os
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import pytest
+from langchain_aws import ChatBedrockConverse
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_fireworks import ChatFireworks
 from langchain_groq import ChatGroq
+from langchain_ollama import ChatOllama
+from langchain_openai import ChatOpenAI
 
 from kiln_ai.adapters.langchain_adapters import (
     LangchainAdapter,
     get_structured_output_options,
+    langchain_model_from_provider,
 )
+from kiln_ai.adapters.ml_model_list import KilnModelProvider, ModelProviderName
 from kiln_ai.adapters.prompt_builders import SimpleChainOfThoughtPromptBuilder
 from kiln_ai.adapters.test_prompt_adaptors import build_test_task
 
@@ -150,3 +158,178 @@ async def test_get_structured_output_options():
     ):
         options = await get_structured_output_options("model_name", "provider")
         assert options == {}
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_openai():
+    provider = KilnModelProvider(
+        name=ModelProviderName.openai, provider_options={"model": "gpt-4"}
+    )
+
+    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+        mock_config.return_value.open_ai_api_key = "test_key"
+        model = await langchain_model_from_provider(provider, "gpt-4")
+        assert isinstance(model, ChatOpenAI)
+        assert model.model_name == "gpt-4"
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_groq():
+    provider = KilnModelProvider(
+        name=ModelProviderName.groq, provider_options={"model": "mixtral-8x7b"}
+    )
+
+    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+        mock_config.return_value.groq_api_key = "test_key"
+        model = await langchain_model_from_provider(provider, "mixtral-8x7b")
+        assert isinstance(model, ChatGroq)
+        assert model.model_name == "mixtral-8x7b"
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_bedrock():
+    provider = KilnModelProvider(
+        name=ModelProviderName.amazon_bedrock,
+        provider_options={"model": "anthropic.claude-v2", "region_name": "us-east-1"},
+    )
+
+    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+        mock_config.return_value.bedrock_access_key = "test_access"
+        mock_config.return_value.bedrock_secret_key = "test_secret"
+        model = await langchain_model_from_provider(provider, "anthropic.claude-v2")
+        assert isinstance(model, ChatBedrockConverse)
+        assert os.environ.get("AWS_ACCESS_KEY_ID") == "test_access"
+        assert os.environ.get("AWS_SECRET_ACCESS_KEY") == "test_secret"
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_fireworks():
+    provider = KilnModelProvider(
+        name=ModelProviderName.fireworks_ai, provider_options={"model": "mixtral-8x7b"}
+    )
+
+    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+        mock_config.return_value.fireworks_api_key = "test_key"
+        model = await langchain_model_from_provider(provider, "mixtral-8x7b")
+        assert isinstance(model, ChatFireworks)
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_ollama():
+    provider = KilnModelProvider(
+        name=ModelProviderName.ollama,
+        provider_options={"model": "llama2", "model_aliases": ["llama2-uncensored"]},
+    )
+
+    mock_connection = MagicMock()
+    with (
+        patch(
+            "kiln_ai.adapters.langchain_adapters.get_ollama_connection",
+            return_value=AsyncMock(return_value=mock_connection),
+        ),
+        patch(
+            "kiln_ai.adapters.langchain_adapters.ollama_model_installed",
+            return_value=True,
+        ),
+        patch(
+            "kiln_ai.adapters.langchain_adapters.ollama_base_url",
+            return_value="http://localhost:11434",
+        ),
+    ):
+        model = await langchain_model_from_provider(provider, "llama2")
+        assert isinstance(model, ChatOllama)
+        assert model.model == "llama2"
+
+
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_invalid():
+    provider = KilnModelProvider.model_construct(
+        name="invalid_provider", provider_options={}
+    )
+
+    with pytest.raises(ValueError, match="Invalid model or provider"):
+        await langchain_model_from_provider(provider, "test_model")
+
+
+@pytest.mark.asyncio
+async def test_langchain_adapter_model_caching(tmp_path):
+    task = build_test_task(tmp_path)
+    custom_model = ChatGroq(model="mixtral-8x7b", groq_api_key="test")
+
+    adapter = LangchainAdapter(kiln_task=task, custom_model=custom_model)
+
+    # First call should return the cached model
+    model1 = await adapter.model()
+    assert model1 is custom_model
+
+    # Second call should return the same cached instance
+    model2 = await adapter.model()
+    assert model2 is model1
+
+
+@pytest.mark.asyncio
+async def test_langchain_adapter_model_structured_output(tmp_path):
+    task = build_test_task(tmp_path)
+    task.output_json_schema = """
+    {
+        "type": "object",
+        "properties": {
+            "count": {"type": "integer"}
+        }
+    }
+    """
+
+    mock_model = MagicMock()
+    mock_model.with_structured_output = MagicMock(return_value="structured_model")
+
+    adapter = LangchainAdapter(
+        kiln_task=task, model_name="test_model", provider="test_provider"
+    )
+
+    with (
+        patch(
+            "kiln_ai.adapters.langchain_adapters.langchain_model_from",
+            AsyncMock(return_value=mock_model),
+        ),
+        patch(
+            "kiln_ai.adapters.langchain_adapters.get_structured_output_options",
+            AsyncMock(return_value={"option1": "value1"}),
+        ),
+    ):
+        model = await adapter.model()
+
+        # Verify the model was configured with structured output
+        mock_model.with_structured_output.assert_called_once_with(
+            {
+                "type": "object",
+                "properties": {"count": {"type": "integer"}},
+                "title": "task_response",
+                "description": "A response from the task",
+            },
+            include_raw=True,
+            option1="value1",
+        )
+        assert model == "structured_model"
+
+
+@pytest.mark.asyncio
+async def test_langchain_adapter_model_no_structured_output_support(tmp_path):
+    task = build_test_task(tmp_path)
+    task.output_json_schema = (
+        '{"type": "object", "properties": {"count": {"type": "integer"}}}'
+    )
+
+    mock_model = MagicMock()
+    # Remove with_structured_output method
+    del mock_model.with_structured_output
+
+    adapter = LangchainAdapter(
+        kiln_task=task, model_name="test_model", provider="test_provider"
+    )
+
+    with patch(
+        "kiln_ai.adapters.langchain_adapters.langchain_model_from",
+        AsyncMock(return_value=mock_model),
+    ):
+        with pytest.raises(ValueError, match="does not support structured output"):
+            await adapter.model()