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

kiln_ai/adapters/ollama_tools.py

@@ -4,6 +4,7 @@ import httpx
 import requests
 from pydantic import BaseModel, Field
 
+from kiln_ai.adapters.ml_embedding_model_list import built_in_embedding_models
 from kiln_ai.adapters.ml_model_list import ModelProviderName, built_in_models
 from kiln_ai.utils.config import Config
 
@@ -41,22 +42,28 @@ class OllamaConnection(BaseModel):
     version: str | None = None
     supported_models: List[str]
     untested_models: List[str] = Field(default_factory=list)
+    supported_embedding_models: List[str] = Field(default_factory=list)
 
     def all_models(self) -> List[str]:
         return self.supported_models + self.untested_models
 
+    def all_embedding_models(self) -> List[str]:
+        return self.supported_embedding_models
+
 
 # Parse the Ollama /api/tags response
-def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
+def parse_ollama_tags(tags: Any) -> OllamaConnection:
     # Build a list of models we support for Ollama from the built-in model list
-    supported_ollama_models = [
-        provider.model_id
-        for model in built_in_models
-        for provider in model.providers
-        if provider.name == ModelProviderName.ollama
-    ]
+    supported_ollama_models = set(
+        [
+            provider.model_id
+            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(
+    supported_ollama_models.update(
         [
             alias
             for model in built_in_models
@@ -65,16 +72,44 @@ def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
         ]
     )
 
+    supported_ollama_embedding_models = set(
+        [
+            provider.model_id
+            for model in built_in_embedding_models
+            for provider in model.providers
+            if provider.name == ModelProviderName.ollama
+        ]
+    )
+    supported_ollama_embedding_models.update(
+        [
+            alias
+            for model in built_in_embedding_models
+            for provider in model.providers
+            for alias in provider.ollama_model_aliases or []
+        ]
+    )
+
     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
-            ]
+            supported_models_latest_aliases = set(
+                [f"{m}:latest" for m in supported_ollama_models]
+            )
+            supported_embedding_models_latest_aliases = set(
+                [f"{m}:latest" for m in supported_ollama_embedding_models]
+            )
+
             for model in model_names:
+                # Skip embedding models - they should only appear in supported_embedding_models
+                if (
+                    model in supported_ollama_embedding_models
+                    or model in supported_embedding_models_latest_aliases
+                ):
+                    continue
+
                 if (
                     model in supported_ollama_models
                     or model in supported_models_latest_aliases
@@ -83,17 +118,31 @@ def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
                 else:
                     untested_models.append(model)
 
-            if available_supported_models or untested_models:
+            available_supported_embedding_models = []
+            for model in model_names:
+                if (
+                    model in supported_ollama_embedding_models
+                    or model in supported_embedding_models_latest_aliases
+                ):
+                    available_supported_embedding_models.append(model)
+
+            if (
+                available_supported_models
+                or untested_models
+                or available_supported_embedding_models
+            ):
                 return OllamaConnection(
                     message="Ollama connected",
                     supported_models=available_supported_models,
                     untested_models=untested_models,
+                    supported_embedding_models=available_supported_embedding_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=[],
+        supported_embedding_models=[],
     )
 
 
@@ -113,3 +162,11 @@ async def get_ollama_connection() -> OllamaConnection | None:
 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
+
+
+def ollama_embedding_model_installed(conn: OllamaConnection, model_name: str) -> bool:
+    all_embedding_models = conn.all_embedding_models()
+    return (
+        model_name in all_embedding_models
+        or f"{model_name}:latest" in all_embedding_models
+    )

kiln_ai/adapters/provider_tools.py

@@ -1,7 +1,10 @@
 import logging
+import os
 from dataclasses import dataclass
 from typing import Dict, List
 
+from pydantic import BaseModel
+
 from kiln_ai.adapters.docker_model_runner_tools import (
     get_docker_model_runner_connection,
 )
@@ -13,11 +16,9 @@ from kiln_ai.adapters.ml_model_list import (
     StructuredOutputMode,
     built_in_models,
 )
-from kiln_ai.adapters.model_adapters.litellm_config import LiteLlmConfig
 from kiln_ai.adapters.ollama_tools import get_ollama_connection
 from kiln_ai.datamodel import Finetune, Task
 from kiln_ai.datamodel.datamodel_enums import ChatStrategy
-from kiln_ai.datamodel.task import RunConfigProperties
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 from kiln_ai.utils.project_utils import project_from_id
@@ -192,50 +193,6 @@ def kiln_model_provider_from(
     )
 
 
-def lite_llm_config_for_openai_compatible(
-    run_config_properties: RunConfigProperties,
-) -> LiteLlmConfig:
-    model_id = run_config_properties.model_name
-    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 like Ollama don't use it, but LiteLLM errors without one
-    api_key = provider.get("api_key") or "NA"
-    base_url = provider.get("base_url")
-    if base_url is None:
-        raise ValueError(
-            f"OpenAI compatible provider {openai_provider_name} has no base URL"
-        )
-
-    # Update a copy of the run config properties to use the openai compatible provider
-    updated_run_config_properties = run_config_properties.model_copy(deep=True)
-    updated_run_config_properties.model_provider_name = (
-        ModelProviderName.openai_compatible
-    )
-    updated_run_config_properties.model_name = model_id
-
-    return LiteLlmConfig(
-        # OpenAI compatible, with a custom base URL
-        run_config_properties=updated_run_config_properties,
-        base_url=base_url,
-        additional_body_options={
-            "api_key": api_key,
-        },
-    )
-
-
 def lite_llm_provider_model(
     model_id: str,
 ) -> KilnModelProvider:
@@ -458,3 +415,190 @@ provider_warnings: Dict[ModelProviderName, ModelProviderWarning] = {
         message="Attempted to use Cerebras without an API key set. \nGet your API key from https://cloud.cerebras.ai/platform",
     ),
 }
+
+
+class LiteLlmCoreConfig(BaseModel):
+    base_url: str | None = None
+    default_headers: Dict[str, str] | None = None
+    additional_body_options: Dict[str, str] | None = None
+
+
+def lite_llm_core_config_for_provider(
+    provider_name: ModelProviderName,
+    openai_compatible_provider_name: str | None = None,
+) -> LiteLlmCoreConfig | None:
+    """
+    Returns a LiteLLM core config for a given provider.
+
+    Args:
+        provider_name: The provider to get the config for
+        openai_compatible_provider_name: Required for openai compatible providers, this is the name of the underlying provider
+    """
+    match provider_name:
+        case ModelProviderName.openrouter:
+            return LiteLlmCoreConfig(
+                base_url=(
+                    os.getenv("OPENROUTER_BASE_URL") or "https://openrouter.ai/api/v1"
+                ),
+                default_headers={
+                    "HTTP-Referer": "https://kiln.tech/openrouter",
+                    "X-Title": "KilnAI",
+                },
+                additional_body_options={
+                    "api_key": Config.shared().open_router_api_key,
+                },
+            )
+        case ModelProviderName.siliconflow_cn:
+            return LiteLlmCoreConfig(
+                base_url=os.getenv("SILICONFLOW_BASE_URL")
+                or "https://api.siliconflow.cn/v1",
+                default_headers={
+                    "HTTP-Referer": "https://kiln.tech/siliconflow",
+                    "X-Title": "KilnAI",
+                },
+                additional_body_options={
+                    "api_key": Config.shared().siliconflow_cn_api_key,
+                },
+            )
+        case ModelProviderName.openai:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "api_key": Config.shared().open_ai_api_key,
+                },
+            )
+        case ModelProviderName.groq:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "api_key": Config.shared().groq_api_key,
+                },
+            )
+        case ModelProviderName.amazon_bedrock:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "aws_access_key_id": Config.shared().bedrock_access_key,
+                    "aws_secret_access_key": Config.shared().bedrock_secret_key,
+                    # The only region that's widely supported for bedrock
+                    "aws_region_name": "us-west-2",
+                },
+            )
+        case ModelProviderName.ollama:
+            # Set the Ollama base URL for 2 reasons:
+            # 1. To use the correct base URL
+            # 2. We use Ollama's OpenAI compatible API (/v1), and don't just let litellm use the Ollama API. We use more advanced features like json_schema.
+            ollama_base_url = (
+                Config.shared().ollama_base_url or "http://localhost:11434"
+            )
+            return LiteLlmCoreConfig(
+                base_url=ollama_base_url + "/v1",
+                additional_body_options={
+                    # LiteLLM errors without an api_key, even though Ollama doesn't support one
+                    "api_key": "NA",
+                },
+            )
+        case ModelProviderName.docker_model_runner:
+            docker_base_url = (
+                Config.shared().docker_model_runner_base_url
+                or "http://localhost:12434/engines/llama.cpp"
+            )
+            return LiteLlmCoreConfig(
+                # Docker Model Runner uses OpenAI-compatible API at /v1 endpoint
+                base_url=docker_base_url + "/v1",
+                additional_body_options={
+                    # LiteLLM errors without an api_key, even though Docker Model Runner doesn't require one.
+                    "api_key": "DMR",
+                },
+            )
+        case ModelProviderName.fireworks_ai:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "api_key": Config.shared().fireworks_api_key,
+                },
+            )
+        case ModelProviderName.anthropic:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "api_key": Config.shared().anthropic_api_key,
+                },
+            )
+        case ModelProviderName.gemini_api:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "api_key": Config.shared().gemini_api_key,
+                },
+            )
+        case ModelProviderName.vertex:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "vertex_project": Config.shared().vertex_project_id,
+                    "vertex_location": Config.shared().vertex_location,
+                },
+            )
+        case ModelProviderName.together_ai:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "api_key": Config.shared().together_api_key,
+                },
+            )
+        case ModelProviderName.azure_openai:
+            return LiteLlmCoreConfig(
+                base_url=Config.shared().azure_openai_endpoint,
+                additional_body_options={
+                    "api_key": Config.shared().azure_openai_api_key,
+                    "api_version": "2025-02-01-preview",
+                },
+            )
+        case ModelProviderName.huggingface:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "api_key": Config.shared().huggingface_api_key,
+                },
+            )
+        case ModelProviderName.cerebras:
+            return LiteLlmCoreConfig(
+                additional_body_options={
+                    "api_key": Config.shared().cerebras_api_key,
+                },
+            )
+        case ModelProviderName.openai_compatible:
+            # openai compatible requires a model name in the format "provider::model_name"
+            if openai_compatible_provider_name is None:
+                raise ValueError("OpenAI compatible provider requires a provider name")
+
+            openai_compatible_providers = (
+                Config.shared().openai_compatible_providers or []
+            )
+
+            provider = next(
+                filter(
+                    lambda p: p.get("name") == openai_compatible_provider_name,
+                    openai_compatible_providers,
+                ),
+                None,
+            )
+
+            if provider is None:
+                raise ValueError(
+                    f"OpenAI compatible provider {openai_compatible_provider_name} not found"
+                )
+
+            # API key optional - some providers like Ollama don't use it, but LiteLLM errors without one
+            api_key = provider.get("api_key") or "NA"
+            base_url = provider.get("base_url")
+            if base_url is None:
+                raise ValueError(
+                    f"OpenAI compatible provider {openai_compatible_provider_name} has no base URL"
+                )
+
+            return LiteLlmCoreConfig(
+                base_url=base_url,
+                additional_body_options={
+                    "api_key": api_key,
+                },
+            )
+        # These are virtual providers that should have mapped to an actual provider upstream (using core_provider method)
+        case ModelProviderName.kiln_fine_tune:
+            return None
+        case ModelProviderName.kiln_custom_registry:
+            return None
+        case _:
+            raise_exhaustive_enum_error(provider_name)

kiln_ai/adapters/rag/deduplication.py

@@ -0,0 +1,49 @@
+from collections import defaultdict
+from typing import DefaultDict
+
+from kiln_ai.datamodel.chunk import ChunkedDocument
+from kiln_ai.datamodel.embedding import ChunkEmbeddings
+from kiln_ai.datamodel.extraction import Document, Extraction
+
+
+def deduplicate_extractions(items: list[Extraction]) -> list[Extraction]:
+    grouped_items: DefaultDict[str, list[Extraction]] = defaultdict(list)
+    for item in items:
+        if item.extractor_config_id is None:
+            raise ValueError("Extractor config ID is required")
+        grouped_items[item.extractor_config_id].append(item)
+    return [min(group, key=lambda x: x.created_at) for group in grouped_items.values()]
+
+
+def deduplicate_chunked_documents(
+    items: list[ChunkedDocument],
+) -> list[ChunkedDocument]:
+    grouped_items: DefaultDict[str, list[ChunkedDocument]] = defaultdict(list)
+    for item in items:
+        if item.chunker_config_id is None:
+            raise ValueError("Chunker config ID is required")
+        grouped_items[item.chunker_config_id].append(item)
+    return [min(group, key=lambda x: x.created_at) for group in grouped_items.values()]
+
+
+def deduplicate_chunk_embeddings(items: list[ChunkEmbeddings]) -> list[ChunkEmbeddings]:
+    grouped_items: DefaultDict[str, list[ChunkEmbeddings]] = defaultdict(list)
+    for item in items:
+        if item.embedding_config_id is None:
+            raise ValueError("Embedding config ID is required")
+        grouped_items[item.embedding_config_id].append(item)
+    return [min(group, key=lambda x: x.created_at) for group in grouped_items.values()]
+
+
+def filter_documents_by_tags(
+    documents: list[Document], tags: list[str] | None
+) -> list[Document]:
+    if not tags:
+        return documents
+
+    filtered_documents = []
+    for document in documents:
+        if document.tags and any(tag in document.tags for tag in tags):
+            filtered_documents.append(document)
+
+    return filtered_documents