kiln-ai 0.12.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

kiln_ai/adapters/model_adapters/openai_model_adapter.py

@@ -1,289 +0,0 @@
-from typing import Any, Dict
-
-from openai import AsyncOpenAI
-from openai.types.chat import (
-    ChatCompletion,
-    ChatCompletionAssistantMessageParam,
-    ChatCompletionSystemMessageParam,
-    ChatCompletionUserMessageParam,
-)
-
-import kiln_ai.datamodel as datamodel
-from kiln_ai.adapters.ml_model_list import (
-    KilnModelProvider,
-    ModelProviderName,
-    StructuredOutputMode,
-)
-from kiln_ai.adapters.model_adapters.base_adapter import (
-    COT_FINAL_ANSWER_PROMPT,
-    AdapterConfig,
-    BaseAdapter,
-    RunOutput,
-)
-from kiln_ai.adapters.model_adapters.openai_compatible_config import (
-    OpenAICompatibleConfig,
-)
-from kiln_ai.adapters.parsers.json_parser import parse_json_string
-from kiln_ai.datamodel import PromptGenerators, PromptId
-from kiln_ai.datamodel.task import RunConfig
-from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
-
-
-class OpenAICompatibleAdapter(BaseAdapter):
-    def __init__(
-        self,
-        config: OpenAICompatibleConfig,
-        kiln_task: datamodel.Task,
-        prompt_id: PromptId | None = None,
-        base_adapter_config: AdapterConfig | None = None,
-    ):
-        self.config = config
-        self.client = AsyncOpenAI(
-            api_key=config.api_key,
-            base_url=config.base_url,
-            default_headers=config.default_headers,
-        )
-
-        run_config = RunConfig(
-            task=kiln_task,
-            model_name=config.model_name,
-            model_provider_name=config.provider_name,
-            prompt_id=prompt_id or PromptGenerators.SIMPLE,
-        )
-
-        super().__init__(
-            run_config=run_config,
-            config=base_adapter_config,
-        )
-
-    async def _run(self, input: Dict | str) -> RunOutput:
-        provider = self.model_provider()
-        intermediate_outputs: dict[str, str] = {}
-        prompt = self.build_prompt()
-        user_msg = self.prompt_builder.build_user_message(input)
-        messages = [
-            ChatCompletionSystemMessageParam(role="system", content=prompt),
-            ChatCompletionUserMessageParam(role="user", content=user_msg),
-        ]
-
-        run_strategy, cot_prompt = self.run_strategy()
-
-        if run_strategy == "cot_as_message":
-            if not cot_prompt:
-                raise ValueError("cot_prompt is required for cot_as_message strategy")
-            messages.append(
-                ChatCompletionSystemMessageParam(role="system", content=cot_prompt)
-            )
-        elif run_strategy == "cot_two_call":
-            if not cot_prompt:
-                raise ValueError("cot_prompt is required for cot_two_call strategy")
-            messages.append(
-                ChatCompletionSystemMessageParam(role="system", content=cot_prompt)
-            )
-
-            # First call for chain of thought
-            cot_response = await self.client.chat.completions.create(
-                model=provider.provider_options["model"],
-                messages=messages,
-            )
-            cot_content = cot_response.choices[0].message.content
-            if cot_content is not None:
-                intermediate_outputs["chain_of_thought"] = cot_content
-
-            messages.extend(
-                [
-                    ChatCompletionAssistantMessageParam(
-                        role="assistant", content=cot_content
-                    ),
-                    ChatCompletionUserMessageParam(
-                        role="user",
-                        content=COT_FINAL_ANSWER_PROMPT,
-                    ),
-                ]
-            )
-
-        # Build custom request params based on model provider
-        extra_body = self.build_extra_body(provider)
-
-        # Main completion call
-        response_format_options = await self.response_format_options()
-        response = await self.client.chat.completions.create(
-            model=provider.provider_options["model"],
-            messages=messages,
-            extra_body=extra_body,
-            logprobs=self.base_adapter_config.top_logprobs is not None,
-            top_logprobs=self.base_adapter_config.top_logprobs,
-            **response_format_options,
-        )
-
-        if not isinstance(response, ChatCompletion):
-            raise RuntimeError(
-                f"Expected ChatCompletion response, got {type(response)}."
-            )
-
-        if hasattr(response, "error") and response.error:  # pyright: ignore
-            raise RuntimeError(
-                f"OpenAI compatible API returned status code {response.error.get('code')}: {response.error.get('message') or 'Unknown error'}.\nError: {response.error}"  # pyright: ignore
-            )
-        if not response.choices or len(response.choices) == 0:
-            raise RuntimeError(
-                "No message content returned in the response from OpenAI compatible API"
-            )
-
-        message = response.choices[0].message
-        logprobs = response.choices[0].logprobs
-
-        # Check logprobs worked, if requested
-        if self.base_adapter_config.top_logprobs is not None and logprobs is None:
-            raise RuntimeError("Logprobs were required, but no logprobs were returned.")
-
-        # Save reasoning if it exists (OpenRouter specific api response field)
-        if provider.require_openrouter_reasoning:
-            if (
-                hasattr(message, "reasoning") and message.reasoning  # pyright: ignore
-            ):
-                intermediate_outputs["reasoning"] = message.reasoning  # pyright: ignore
-            else:
-                raise RuntimeError(
-                    "Reasoning is required for this model, but no reasoning was returned from OpenRouter."
-                )
-
-        # the string content of the response
-        response_content = message.content
-
-        # Fallback: Use args of first tool call to task_response if it exists
-        if not response_content and message.tool_calls:
-            tool_call = next(
-                (
-                    tool_call
-                    for tool_call in message.tool_calls
-                    if tool_call.function.name == "task_response"
-                ),
-                None,
-            )
-            if tool_call:
-                response_content = tool_call.function.arguments
-
-        if not isinstance(response_content, str):
-            raise RuntimeError(f"response is not a string: {response_content}")
-
-        # Parse to dict if we have structured output
-        output: Dict | str = response_content
-        if self.has_structured_output():
-            output = parse_json_string(response_content)
-
-        return RunOutput(
-            output=output,
-            intermediate_outputs=intermediate_outputs,
-            output_logprobs=logprobs,
-        )
-
-    def adapter_name(self) -> str:
-        return "kiln_openai_compatible_adapter"
-
-    async def response_format_options(self) -> dict[str, Any]:
-        # Unstructured if task isn't structured
-        if not self.has_structured_output():
-            return {}
-
-        provider = self.model_provider()
-        match provider.structured_output_mode:
-            case StructuredOutputMode.json_mode:
-                return {"response_format": {"type": "json_object"}}
-            case StructuredOutputMode.json_schema:
-                output_schema = self.task().output_schema()
-                return {
-                    "response_format": {
-                        "type": "json_schema",
-                        "json_schema": {
-                            "name": "task_response",
-                            "schema": output_schema,
-                        },
-                    }
-                }
-            case StructuredOutputMode.function_calling_weak:
-                return self.tool_call_params(strict=False)
-            case StructuredOutputMode.function_calling:
-                return self.tool_call_params(strict=True)
-            case StructuredOutputMode.json_instructions:
-                # JSON done via instructions in prompt, not the API response format. Do not ask for json_object (see option below).
-                return {}
-            case StructuredOutputMode.json_instruction_and_object:
-                # We set response_format to json_object and also set json instructions in the prompt
-                return {"response_format": {"type": "json_object"}}
-            case StructuredOutputMode.default:
-                # Default to function calling -- it's older than the other modes. Higher compatibility.
-                return self.tool_call_params(strict=True)
-            case _:
-                raise_exhaustive_enum_error(provider.structured_output_mode)
-
-    def tool_call_params(self, strict: bool) -> dict[str, Any]:
-        # Add additional_properties: false to the schema (OpenAI requires this for some models)
-        output_schema = self.task().output_schema()
-        if not isinstance(output_schema, dict):
-            raise ValueError(
-                "Invalid output schema for this task. Can not use tool calls."
-            )
-        output_schema["additionalProperties"] = False
-
-        function_params = {
-            "name": "task_response",
-            "parameters": output_schema,
-        }
-        # This should be on, but we allow setting function_calling_weak for APIs that don't support it.
-        if strict:
-            function_params["strict"] = True
-
-        return {
-            "tools": [
-                {
-                    "type": "function",
-                    "function": function_params,
-                }
-            ],
-            "tool_choice": {
-                "type": "function",
-                "function": {"name": "task_response"},
-            },
-        }
-
-    def build_extra_body(self, provider: KilnModelProvider) -> dict[str, Any]:
-        # TODO P1: Don't love having this logic here. But it's a usability improvement
-        # so better to keep it than exclude it. Should figure out how I want to isolate
-        # this sort of logic so it's config driven and can be overridden
-
-        extra_body = {}
-        provider_options = {}
-
-        if provider.require_openrouter_reasoning:
-            # https://openrouter.ai/docs/use-cases/reasoning-tokens
-            extra_body["reasoning"] = {
-                "exclude": False,
-            }
-
-        if provider.r1_openrouter_options:
-            # Require providers that support the reasoning parameter
-            provider_options["require_parameters"] = True
-            # Prefer R1 providers with reasonable perf/quants
-            provider_options["order"] = ["Fireworks", "Together"]
-            # R1 providers with unreasonable quants
-            provider_options["ignore"] = ["DeepInfra"]
-
-        # Only set of this request is to get logprobs.
-        if (
-            provider.logprobs_openrouter_options
-            and self.base_adapter_config.top_logprobs is not None
-        ):
-            # Don't let OpenRouter choose a provider that doesn't support logprobs.
-            provider_options["require_parameters"] = True
-            # DeepInfra silently fails to return logprobs consistently.
-            provider_options["ignore"] = ["DeepInfra"]
-
-        if provider.openrouter_skip_required_parameters:
-            # Oddball case, R1 14/8/1.5B fail with this param, even though they support thinking params.
-            provider_options["require_parameters"] = False
-
-        if len(provider_options) > 0:
-            extra_body["provider"] = provider_options
-
-        return extra_body

kiln_ai/adapters/model_adapters/test_langchain_adapter.py

@@ -1,343 +0,0 @@
-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 kiln_ai.adapters.ml_model_list import (
-    KilnModelProvider,
-    ModelProviderName,
-    StructuredOutputMode,
-)
-from kiln_ai.adapters.model_adapters.base_adapter import COT_FINAL_ANSWER_PROMPT
-from kiln_ai.adapters.model_adapters.langchain_adapters import (
-    LangchainAdapter,
-    langchain_model_from_provider,
-)
-from kiln_ai.adapters.test_prompt_adaptors import build_test_task
-from kiln_ai.datamodel.task import RunConfig
-
-
-@pytest.fixture
-def mock_adapter(tmp_path):
-    return LangchainAdapter(
-        kiln_task=build_test_task(tmp_path),
-        model_name="llama_3_1_8b",
-        provider="ollama",
-    )
-
-
-def test_langchain_adapter_munge_response(mock_adapter):
-    # Mistral Large tool calling format is a bit different
-    response = {
-        "name": "task_response",
-        "arguments": {
-            "setup": "Why did the cow join a band?",
-            "punchline": "Because she wanted to be a moo-sician!",
-        },
-    }
-    munged = mock_adapter._munge_response(response)
-    assert munged["setup"] == "Why did the cow join a band?"
-    assert munged["punchline"] == "Because she wanted to be a moo-sician!"
-
-    # non mistral format should continue to work
-    munged = mock_adapter._munge_response(response["arguments"])
-    assert munged["setup"] == "Why did the cow join a band?"
-    assert munged["punchline"] == "Because she wanted to be a moo-sician!"
-
-
-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 = LangchainAdapter(kiln_task=task, custom_model=custom)
-
-    assert lca.run_config.model_name == "custom.langchain:llama-3.1-8b-instant"
-    assert lca.run_config.model_provider_name == "custom.langchain:ChatGroq"
-
-
-def test_langchain_adapter_info(tmp_path):
-    task = build_test_task(tmp_path)
-
-    lca = LangchainAdapter(kiln_task=task, model_name="llama_3_1_8b", provider="ollama")
-
-    assert lca.adapter_name() == "kiln_langchain_adapter"
-    assert lca.run_config.model_name == "llama_3_1_8b"
-    assert lca.run_config.model_provider_name == "ollama"
-
-
-async def test_langchain_adapter_with_cot(tmp_path):
-    task = build_test_task(tmp_path)
-    task.output_json_schema = (
-        '{"type": "object", "properties": {"count": {"type": "integer"}}}'
-    )
-    lca = LangchainAdapter(
-        kiln_task=task,
-        model_name="llama_3_1_8b",
-        provider="ollama",
-        prompt_id="simple_chain_of_thought_prompt_builder",
-    )
-
-    # Mock the base model and its invoke method
-    mock_base_model = MagicMock()
-    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.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)
-
-    # Patch both the langchain_model_from function and self.model()
-    with (
-        patch.object(LangchainAdapter, "langchain_model_from", mock_model_from),
-        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.ainvoke.call_args[0][0],
-        mock_model_instance.ainvoke.call_args[0][0],
-    ]:
-        assert isinstance(
-            invoke_args[0], SystemMessage
-        )  # First message should be system prompt
-        assert (
-            "You are an assistant which performs math tasks provided in plain text."
-            in invoke_args[0].content
-        )
-        assert isinstance(invoke_args[1], HumanMessage)
-        assert "test input" in invoke_args[1].content
-        assert isinstance(invoke_args[2], SystemMessage)
-        assert "step by step" in invoke_args[2].content
-
-    # the COT should only have 3 messages
-    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.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], HumanMessage)
-    assert COT_FINAL_ANSWER_PROMPT in invoke_args[4].content
-
-    assert (
-        response.intermediate_outputs["chain_of_thought"]
-        == "Chain of thought reasoning..."
-    )
-    assert response.output == {"count": 1}
-
-
-@pytest.mark.parametrize(
-    "structured_output_mode,expected_method",
-    [
-        (StructuredOutputMode.function_calling, "function_calling"),
-        (StructuredOutputMode.json_mode, "json_mode"),
-        (StructuredOutputMode.json_schema, "json_schema"),
-        (StructuredOutputMode.json_instruction_and_object, "json_mode"),
-        (StructuredOutputMode.default, None),
-    ],
-)
-async def test_get_structured_output_options(
-    mock_adapter, structured_output_mode, expected_method
-):
-    # Mock the provider response
-    mock_provider = MagicMock()
-    mock_provider.structured_output_mode = structured_output_mode
-
-    # Mock adapter.model_provider()
-    mock_adapter.model_provider = MagicMock(return_value=mock_provider)
-
-    options = mock_adapter.get_structured_output_options("model_name", "provider")
-    assert options.get("method") == expected_method
-
-
-@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.model_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.model_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.model_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.model_adapters.langchain_adapters.get_ollama_connection",
-            return_value=AsyncMock(return_value=mock_connection),
-        ),
-        patch(
-            "kiln_ai.adapters.model_adapters.langchain_adapters.ollama_model_installed",
-            return_value=True,
-        ),
-        patch(
-            "kiln_ai.adapters.model_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="ollama"
-    )
-    adapter.get_structured_output_options = MagicMock(
-        return_value={"option1": "value1"}
-    )
-    adapter.langchain_model_from = AsyncMock(return_value=mock_model)
-
-    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="ollama"
-    )
-    adapter.langchain_model_from = AsyncMock(return_value=mock_model)
-
-    with pytest.raises(ValueError, match="does not support structured output"):
-        await adapter.model()
-
-
-import pytest
-
-
-@pytest.mark.parametrize(
-    "provider_name",
-    [
-        (ModelProviderName.openai),
-        (ModelProviderName.openai_compatible),
-        (ModelProviderName.openrouter),
-    ],
-)
-@pytest.mark.asyncio
-async def test_langchain_model_from_provider_unsupported_providers(provider_name):
-    # Arrange
-    provider = KilnModelProvider(
-        name=provider_name, provider_options={}, structured_output_mode="default"
-    )
-
-    # Assert unsupported providers raise an error
-    with pytest.raises(ValueError):
-        await langchain_model_from_provider(provider, "test-model")