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

kiln_ai/adapters/ollama_tools.py

@@ -0,0 +1,114 @@
+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
+
+
+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
+    """
+    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() -> 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,295 @@
+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.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
+
+    # 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 = next((t for t in project.tasks() if t.id == task_id), None)
+    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)
+    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 _:
+                # 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_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/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"
         )
 

kiln_ai/adapters/test_langchain_adapter.py

@@ -3,16 +3,17 @@ from unittest.mock import AsyncMock, MagicMock, patch
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_groq import ChatGroq
 
-from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
+from kiln_ai.adapters.langchain_adapters import (
+    LangchainAdapter,
+    get_structured_output_options,
+)
 from kiln_ai.adapters.prompt_builders import SimpleChainOfThoughtPromptBuilder
 from kiln_ai.adapters.test_prompt_adaptors import build_test_task
 
 
 def test_langchain_adapter_munge_response(tmp_path):
     task = build_test_task(tmp_path)
-    lca = LangChainPromptAdapter(
-        kiln_task=task, model_name="llama_3_1_8b", provider="ollama"
-    )
+    lca = LangchainAdapter(kiln_task=task, model_name="llama_3_1_8b", provider="ollama")
     # Mistral Large tool calling format is a bit different
     response = {
         "name": "task_response",
@@ -35,7 +36,7 @@ def test_langchain_adapter_infer_model_name(tmp_path):
     task = build_test_task(tmp_path)
     custom = ChatGroq(model="llama-3.1-8b-instant", groq_api_key="test")
 
-    lca = LangChainPromptAdapter(kiln_task=task, custom_model=custom)
+    lca = LangchainAdapter(kiln_task=task, custom_model=custom)
 
     model_info = lca.adapter_info()
     assert model_info.model_name == "custom.langchain:llama-3.1-8b-instant"
@@ -45,9 +46,7 @@ def test_langchain_adapter_infer_model_name(tmp_path):
 def test_langchain_adapter_info(tmp_path):
     task = build_test_task(tmp_path)
 
-    lca = LangChainPromptAdapter(
-        kiln_task=task, model_name="llama_3_1_8b", provider="ollama"
-    )
+    lca = LangchainAdapter(kiln_task=task, model_name="llama_3_1_8b", provider="ollama")
 
     model_info = lca.adapter_info()
     assert model_info.adapter_name == "kiln_langchain_adapter"
@@ -60,7 +59,7 @@ async def test_langchain_adapter_with_cot(tmp_path):
     task.output_json_schema = (
         '{"type": "object", "properties": {"count": {"type": "integer"}}}'
     )
-    lca = LangChainPromptAdapter(
+    lca = LangchainAdapter(
         kiln_task=task,
         model_name="llama_3_1_8b",
         provider="ollama",
@@ -69,13 +68,13 @@ async def test_langchain_adapter_with_cot(tmp_path):
 
     # Mock the base model and its invoke method
     mock_base_model = MagicMock()
-    mock_base_model.invoke.return_value = AIMessage(
-        content="Chain of thought reasoning..."
+    mock_base_model.ainvoke = AsyncMock(
+        return_value=AIMessage(content="Chain of thought reasoning...")
     )
 
     # Create a separate mock for self.model()
     mock_model_instance = MagicMock()
-    mock_model_instance.invoke.return_value = {"parsed": {"count": 1}}
+    mock_model_instance.ainvoke = AsyncMock(return_value={"parsed": {"count": 1}})
 
     # Mock the langchain_model_from function to return the base model
     mock_model_from = AsyncMock(return_value=mock_base_model)
@@ -85,14 +84,14 @@ async def test_langchain_adapter_with_cot(tmp_path):
         patch(
             "kiln_ai.adapters.langchain_adapters.langchain_model_from", mock_model_from
         ),
-        patch.object(LangChainPromptAdapter, "model", return_value=mock_model_instance),
+        patch.object(LangchainAdapter, "model", return_value=mock_model_instance),
     ):
         response = await lca._run("test input")
 
     # First 3 messages are the same for both calls
     for invoke_args in [
-        mock_base_model.invoke.call_args[0][0],
-        mock_model_instance.invoke.call_args[0][0],
+        mock_base_model.ainvoke.call_args[0][0],
+        mock_model_instance.ainvoke.call_args[0][0],
     ]:
         assert isinstance(
             invoke_args[0], SystemMessage
@@ -107,11 +106,11 @@ async def test_langchain_adapter_with_cot(tmp_path):
         assert "step by step" in invoke_args[2].content
 
     # the COT should only have 3 messages
-    assert len(mock_base_model.invoke.call_args[0][0]) == 3
-    assert len(mock_model_instance.invoke.call_args[0][0]) == 5
+    assert len(mock_base_model.ainvoke.call_args[0][0]) == 3
+    assert len(mock_model_instance.ainvoke.call_args[0][0]) == 5
 
     # the final response should have the COT content and the final instructions
-    invoke_args = mock_model_instance.invoke.call_args[0][0]
+    invoke_args = mock_model_instance.ainvoke.call_args[0][0]
     assert isinstance(invoke_args[3], AIMessage)
     assert "Chain of thought reasoning..." in invoke_args[3].content
     assert isinstance(invoke_args[4], SystemMessage)
@@ -122,3 +121,32 @@ async def test_langchain_adapter_with_cot(tmp_path):
         == "Chain of thought reasoning..."
     )
     assert response.output == {"count": 1}
+
+
+async def test_get_structured_output_options():
+    # Mock the provider response
+    mock_provider = MagicMock()
+    mock_provider.adapter_options = {
+        "langchain": {
+            "with_structured_output_options": {
+                "force_json_response": True,
+                "max_retries": 3,
+            }
+        }
+    }
+
+    # Test with provider that has options
+    with patch(
+        "kiln_ai.adapters.langchain_adapters.kiln_model_provider_from",
+        AsyncMock(return_value=mock_provider),
+    ):
+        options = await get_structured_output_options("model_name", "provider")
+        assert options == {"force_json_response": True, "max_retries": 3}
+
+    # Test with provider that has no options
+    with patch(
+        "kiln_ai.adapters.langchain_adapters.kiln_model_provider_from",
+        AsyncMock(return_value=None),
+    ):
+        options = await get_structured_output_options("model_name", "provider")
+        assert options == {}

kiln_ai/adapters/test_ollama_tools.py

@@ -0,0 +1,42 @@
+import json
+
+from kiln_ai.adapters.ollama_tools import (
+    OllamaConnection,
+    ollama_model_installed,
+    parse_ollama_tags,
+)
+
+
+def test_parse_ollama_tags_no_models():
+    json_response = '{"models":[{"name":"scosman_net","model":"scosman_net:latest"},{"name":"phi3.5:latest","model":"phi3.5:latest","modified_at":"2024-10-02T12:04:35.191519822-04:00","size":2176178843,"digest":"61819fb370a3c1a9be6694869331e5f85f867a079e9271d66cb223acb81d04ba","details":{"parent_model":"","format":"gguf","family":"phi3","families":["phi3"],"parameter_size":"3.8B","quantization_level":"Q4_0"}},{"name":"gemma2:2b","model":"gemma2:2b","modified_at":"2024-09-09T16:46:38.64348929-04:00","size":1629518495,"digest":"8ccf136fdd5298f3ffe2d69862750ea7fb56555fa4d5b18c04e3fa4d82ee09d7","details":{"parent_model":"","format":"gguf","family":"gemma2","families":["gemma2"],"parameter_size":"2.6B","quantization_level":"Q4_0"}},{"name":"llama3.1:latest","model":"llama3.1:latest","modified_at":"2024-09-01T17:19:43.481523695-04:00","size":4661230720,"digest":"f66fc8dc39ea206e03ff6764fcc696b1b4dfb693f0b6ef751731dd4e6269046e","details":{"parent_model":"","format":"gguf","family":"llama","families":["llama"],"parameter_size":"8.0B","quantization_level":"Q4_0"}}]}'
+    tags = json.loads(json_response)
+    print(json.dumps(tags, indent=2))
+    conn = parse_ollama_tags(tags)
+    assert "phi3.5:latest" in conn.supported_models
+    assert "gemma2:2b" in conn.supported_models
+    assert "llama3.1:latest" in conn.supported_models
+    assert "scosman_net:latest" in conn.untested_models
+
+
+def test_parse_ollama_tags_only_untested_models():
+    json_response = '{"models":[{"name":"scosman_net","model":"scosman_net:latest"}]}'
+    tags = json.loads(json_response)
+    conn = parse_ollama_tags(tags)
+    assert conn.supported_models == []
+    assert conn.untested_models == ["scosman_net:latest"]
+
+
+def test_ollama_model_installed():
+    conn = OllamaConnection(
+        supported_models=["phi3.5:latest", "gemma2:2b", "llama3.1:latest"],
+        message="Connected",
+        untested_models=["scosman_net:latest"],
+    )
+    assert ollama_model_installed(conn, "phi3.5:latest")
+    assert ollama_model_installed(conn, "phi3.5")
+    assert ollama_model_installed(conn, "gemma2:2b")
+    assert ollama_model_installed(conn, "llama3.1:latest")
+    assert ollama_model_installed(conn, "llama3.1")
+    assert ollama_model_installed(conn, "scosman_net:latest")
+    assert ollama_model_installed(conn, "scosman_net")
+    assert not ollama_model_installed(conn, "unknown_model")

kiln_ai/adapters/test_prompt_adaptors.py

@@ -5,8 +5,10 @@ import pytest
 from langchain_core.language_models.fake_chat_models import FakeListChatModel
 
 import kiln_ai.datamodel as datamodel
-from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
-from kiln_ai.adapters.ml_model_list import built_in_models, ollama_online
+from kiln_ai.adapters.adapter_registry import adapter_for_task
+from kiln_ai.adapters.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.ml_model_list import built_in_models
+from kiln_ai.adapters.ollama_tools import ollama_online
 from kiln_ai.adapters.prompt_builders import (
     BasePromptBuilder,
     SimpleChainOfThoughtPromptBuilder,
@@ -106,7 +108,7 @@ async def test_amazon_bedrock(tmp_path):
 async def test_mock(tmp_path):
     task = build_test_task(tmp_path)
     mockChatModel = FakeListChatModel(responses=["mock response"])
-    adapter = LangChainPromptAdapter(task, custom_model=mockChatModel)
+    adapter = LangchainAdapter(task, custom_model=mockChatModel)
     run = await adapter.invoke("You are a mock, send me the response!")
     assert "mock response" in run.output.output
 
@@ -114,7 +116,7 @@ async def test_mock(tmp_path):
 async def test_mock_returning_run(tmp_path):
     task = build_test_task(tmp_path)
     mockChatModel = FakeListChatModel(responses=["mock response"])
-    adapter = LangChainPromptAdapter(task, custom_model=mockChatModel)
+    adapter = LangchainAdapter(task, custom_model=mockChatModel)
     run = await adapter.invoke("You are a mock, send me the response!")
     assert run.output.output == "mock response"
     assert run is not None
@@ -192,7 +194,7 @@ async def run_simple_task(
     provider: str,
     prompt_builder: BasePromptBuilder | None = None,
 ) -> datamodel.TaskRun:
-    adapter = LangChainPromptAdapter(
+    adapter = adapter_for_task(
         task, model_name=model_name, provider=provider, prompt_builder=prompt_builder
     )