kiln-ai 0.20.1__py3-none-any.whl → 0.21.0__py3-none-any.whl

kiln_ai/adapters/fine_tune/test_dataset_formatter.py

@@ -206,7 +206,7 @@ def test_generate_chat_message_toolcall(mock_training_chat_two_step_json):
 
 def test_generate_chat_message_toolcall_invalid_json(mock_training_chat_two_step_json):
     mock_training_chat_two_step_json[-1].content = "invalid json"
-    with pytest.raises(ValueError, match="^Last message is not JSON"):
+    with pytest.raises(ValueError, match=r"^Last message is not JSON"):
         generate_chat_message_toolcall(mock_training_chat_two_step_json)
 
 
@@ -536,7 +536,7 @@ def test_generate_huggingface_chat_template_toolcall_invalid_json(
 ):
     mock_training_chat_two_step_json[-1].content = "invalid json"
 
-    with pytest.raises(ValueError, match="^Last message is not JSON"):
+    with pytest.raises(ValueError, match=r"^Last message is not JSON"):
         generate_huggingface_chat_template_toolcall(mock_training_chat_two_step_json)
 
 

kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py

@@ -14,11 +14,7 @@ from kiln_ai.adapters.fine_tune.fireworks_finetune import (
     DeployStatus,
     FireworksFinetune,
 )
-from kiln_ai.datamodel import (
-    DatasetSplit,
-    StructuredOutputMode,
-    Task,
-)
+from kiln_ai.datamodel import DatasetSplit, StructuredOutputMode, Task
 from kiln_ai.datamodel import Finetune as FinetuneModel
 from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
@@ -1053,7 +1049,7 @@ async def test_fetch_all_deployments_invalid_json(fireworks_finetune, mock_api_k
 
         with pytest.raises(
             ValueError,
-            match="Invalid response from Fireworks. Expected list of deployments in 'deployments' key",
+            match=r"Invalid response from Fireworks. Expected list of deployments in 'deployments' key",
         ):
             await fireworks_finetune._fetch_all_deployments()
 

kiln_ai/adapters/fine_tune/test_together_finetune.py

@@ -17,11 +17,7 @@ from kiln_ai.adapters.fine_tune.together_finetune import (
     _pending_statuses,
     _running_statuses,
 )
-from kiln_ai.datamodel import (
-    DatasetSplit,
-    StructuredOutputMode,
-    Task,
-)
+from kiln_ai.datamodel import DatasetSplit, StructuredOutputMode, Task
 from kiln_ai.datamodel import Finetune as FinetuneModel
 from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
 from kiln_ai.utils.config import Config
@@ -105,7 +101,7 @@ def mock_api_key():
 def test_init_missing_api_key(finetune):
     with patch.object(Config, "shared") as mock_config:
         mock_config.return_value.together_api_key = None
-        with pytest.raises(ValueError, match="Together.ai API key not set"):
+        with pytest.raises(ValueError, match=r"Together.ai API key not set"):
             TogetherFinetune(datamodel=finetune)
 
 

kiln_ai/adapters/ml_embedding_model_list.py

@@ -0,0 +1,192 @@
+from enum import Enum
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+
+
+class KilnEmbeddingModelFamily(str, Enum):
+    """
+    Enumeration of supported embedding model families.
+    """
+
+    # for bespoke proprietary models, the family tends to be the same
+    # as provider name, but it does not have to be
+    openai = "openai"
+    gemini = "gemini"
+    gemma = "gemma"
+    nomic = "nomic"
+
+
+class EmbeddingModelName(str, Enum):
+    """
+    Enumeration of specific model versions supported by the system.
+    """
+
+    # Embedding model names are often generic (e.g., "text-embedding"),
+    # so we prefix them with the provider name (e.g., "openai_") to ensure
+    # uniqueness across providers now and in the future
+    openai_text_embedding_3_small = "openai_text_embedding_3_small"
+    openai_text_embedding_3_large = "openai_text_embedding_3_large"
+    gemini_text_embedding_004 = "gemini_text_embedding_004"
+    gemini_embedding_001 = "gemini_embedding_001"
+    embedding_gemma_300m = "embedding_gemma_300m"
+    nomic_text_embedding_v1_5 = "nomic_text_embedding_v1_5"
+
+
+class KilnEmbeddingModelProvider(BaseModel):
+    name: ModelProviderName
+
+    model_id: str = Field(
+        description="The model ID for the embedding model. This is the ID used to identify the model in the provider's API.",
+    )
+
+    max_input_tokens: int | None = Field(
+        default=None,
+        description="The maximum number of tokens that can be input to the model.",
+    )
+
+    n_dimensions: int = Field(
+        description="The number of dimensions in the output embedding.",
+    )
+
+    supports_custom_dimensions: bool = Field(
+        default=False,
+        description="Whether the model supports setting a custom output dimension. If true, the user can set the output dimension in the UI.",
+    )
+
+    suggested_for_chunk_embedding: bool = Field(
+        default=False,
+        description="Whether the model is particularly good for chunk embedding.",
+    )
+
+    ollama_model_aliases: List[str] | None = None
+
+
+class KilnEmbeddingModel(BaseModel):
+    """
+    Configuration for a specific embedding model.
+    """
+
+    family: str
+    name: str
+    friendly_name: str
+    providers: List[KilnEmbeddingModelProvider]
+
+
+built_in_embedding_models: List[KilnEmbeddingModel] = [
+    # openai
+    KilnEmbeddingModel(
+        family=KilnEmbeddingModelFamily.openai,
+        name=EmbeddingModelName.openai_text_embedding_3_small,
+        friendly_name="Text Embedding 3 Small",
+        providers=[
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.openai,
+                model_id="text-embedding-3-small",
+                n_dimensions=1536,
+                max_input_tokens=8192,
+                supports_custom_dimensions=True,
+            ),
+        ],
+    ),
+    KilnEmbeddingModel(
+        family=KilnEmbeddingModelFamily.openai,
+        name=EmbeddingModelName.openai_text_embedding_3_large,
+        friendly_name="Text Embedding 3 Large",
+        providers=[
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.openai,
+                model_id="text-embedding-3-large",
+                n_dimensions=3072,
+                max_input_tokens=8192,
+                supports_custom_dimensions=True,
+                suggested_for_chunk_embedding=True,
+            ),
+        ],
+    ),
+    # gemini
+    KilnEmbeddingModel(
+        family=KilnEmbeddingModelFamily.gemini,
+        name=EmbeddingModelName.gemini_text_embedding_004,
+        friendly_name="Text Embedding 004",
+        providers=[
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.gemini_api,
+                model_id="text-embedding-004",
+                n_dimensions=768,
+                max_input_tokens=2048,
+            ),
+        ],
+    ),
+    KilnEmbeddingModel(
+        family=KilnEmbeddingModelFamily.gemini,
+        name=EmbeddingModelName.gemini_embedding_001,
+        friendly_name="Gemini Embedding 001",
+        providers=[
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.gemini_api,
+                model_id="gemini-embedding-001",
+                n_dimensions=3072,
+                max_input_tokens=2048,
+                supports_custom_dimensions=True,
+                suggested_for_chunk_embedding=True,
+            ),
+        ],
+    ),
+    # gemma
+    KilnEmbeddingModel(
+        family=KilnEmbeddingModelFamily.gemma,
+        name=EmbeddingModelName.embedding_gemma_300m,
+        friendly_name="Embedding Gemma 300m",
+        providers=[
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.ollama,
+                model_id="embeddinggemma:300m",
+                n_dimensions=768,
+                max_input_tokens=2048,
+                # the model itself does support custom dimensions, but
+                # not sure if ollama supports it
+                supports_custom_dimensions=False,
+                ollama_model_aliases=["embeddinggemma"],
+            ),
+        ],
+    ),
+    # nomic
+    KilnEmbeddingModel(
+        family=KilnEmbeddingModelFamily.nomic,
+        name=EmbeddingModelName.nomic_text_embedding_v1_5,
+        friendly_name="Nomic Embed Text v1.5",
+        providers=[
+            KilnEmbeddingModelProvider(
+                name=ModelProviderName.ollama,
+                model_id="nomic-embed-text:v1.5",
+                n_dimensions=768,
+                max_input_tokens=2048,
+                # the model itself does support custom dimensions, but
+                # not sure if ollama supports it
+                supports_custom_dimensions=False,
+                ollama_model_aliases=["nomic-embed-text"],
+            ),
+        ],
+    ),
+]
+
+
+def get_model_by_name(name: EmbeddingModelName) -> KilnEmbeddingModel:
+    for model in built_in_embedding_models:
+        if model.name == name:
+            return model
+    raise ValueError(f"Embedding model {name} not found in the list of built-in models")
+
+
+def built_in_embedding_models_from_provider(
+    provider_name: ModelProviderName, model_name: str
+) -> KilnEmbeddingModelProvider | None:
+    for model in built_in_embedding_models:
+        if model.name == model_name:
+            for p in model.providers:
+                if p.name == provider_name:
+                    return p
+    return None