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

kiln_ai/adapters/test_langchain_adapter.py

@@ -1,18 +1,27 @@
+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 LangChainPromptAdapter
+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
 
 
 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 +44,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 +54,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 +67,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 +76,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 +92,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 +114,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 +129,207 @@ 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 == {}
+
+
+@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()

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
     )