crewAI/lib/crewai/tests/llms/snowflake/test_snowflake.py

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import Mock, patch

import pytest

from crewai.llm import LLM
from crewai.llms.providers.snowflake.completion import (
    SNOWFLAKE_CORTEX_PATH,
    SnowflakeCompletion,
    _normalize_snowflake_base_url,
)


def _snowflake_env(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("SNOWFLAKE_PAT", "test-pat")
    monkeypatch.setenv("SNOWFLAKE_ACCOUNT_URL", "https://org-account.snowflakecomputing.com")
    monkeypatch.delenv("SNOWFLAKE_TOKEN", raising=False)
    monkeypatch.delenv("SNOWFLAKE_JWT", raising=False)
    monkeypatch.delenv("SNOWFLAKE_ACCOUNT", raising=False)
    monkeypatch.delenv("SNOWFLAKE_ACCOUNT_ID", raising=False)
    monkeypatch.delenv("SNOWFLAKE_ACCOUNT_IDENTIFIER", raising=False)


class TestSnowflakeConfig:
    def test_normalizes_account_url_to_cortex_base_url(self):
        assert (
            _normalize_snowflake_base_url("https://org-account.snowflakecomputing.com")
            == f"https://org-account.snowflakecomputing.com{SNOWFLAKE_CORTEX_PATH}"
        )

    def test_preserves_existing_cortex_base_url(self):
        base_url = f"https://org-account.snowflakecomputing.com{SNOWFLAKE_CORTEX_PATH}"
        assert _normalize_snowflake_base_url(base_url) == base_url

    def test_rejects_endpoint_path_in_base_url(self):
        with pytest.raises(ValueError, match="do not include endpoint paths"):
            _normalize_snowflake_base_url(
                "https://org-account.snowflakecomputing.com"
                f"{SNOWFLAKE_CORTEX_PATH}/chat/completions"
            )

    def test_empty_api_key_falls_back_to_env_token(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)

        llm = SnowflakeCompletion(model="openai-gpt-4.1", api_key="")

        assert llm.api_key == "test-pat"

    def test_uses_env_token_and_account_url(self, monkeypatch: pytest.MonkeyPatch):
        _snowflake_env(monkeypatch)

        llm = SnowflakeCompletion(model="openai-gpt-4.1")

        assert llm.api_key == "test-pat"
        assert llm.base_url == (
            f"https://org-account.snowflakecomputing.com{SNOWFLAKE_CORTEX_PATH}"
        )
        assert llm.account_url == llm.base_url

    def test_strips_litellm_pat_prefix_for_compatibility(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        monkeypatch.setenv("SNOWFLAKE_PAT", "pat/test-pat")
        monkeypatch.setenv("SNOWFLAKE_ACCOUNT", "org-account")

        llm = SnowflakeCompletion(model="openai-gpt-4.1")

        assert llm.api_key == "test-pat"

    def test_missing_token_raises_clear_error(self, monkeypatch: pytest.MonkeyPatch):
        monkeypatch.delenv("SNOWFLAKE_PAT", raising=False)
        monkeypatch.delenv("SNOWFLAKE_TOKEN", raising=False)
        monkeypatch.delenv("SNOWFLAKE_JWT", raising=False)
        monkeypatch.setenv("SNOWFLAKE_ACCOUNT_URL", "https://org-account.snowflakecomputing.com")

        with pytest.raises(ValueError, match="Snowflake token is required"):
            SnowflakeCompletion(model="openai-gpt-4.1")

    def test_missing_account_raises_clear_error(self, monkeypatch: pytest.MonkeyPatch):
        monkeypatch.setenv("SNOWFLAKE_PAT", "test-pat")
        monkeypatch.delenv("SNOWFLAKE_ACCOUNT_URL", raising=False)
        monkeypatch.delenv("SNOWFLAKE_ACCOUNT", raising=False)
        monkeypatch.delenv("SNOWFLAKE_ACCOUNT_ID", raising=False)
        monkeypatch.delenv("SNOWFLAKE_ACCOUNT_IDENTIFIER", raising=False)

        with pytest.raises(ValueError, match="Snowflake account URL is required"):
            SnowflakeCompletion(model="openai-gpt-4.1")

    def test_responses_api_is_rejected(self, monkeypatch: pytest.MonkeyPatch):
        _snowflake_env(monkeypatch)

        with pytest.raises(ValueError, match="supports only the Chat Completions API"):
            SnowflakeCompletion(model="openai-gpt-4.1", api="responses")


class TestSnowflakeFactory:
    def test_llm_creates_native_snowflake_provider(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)

        llm = LLM(model="snowflake/openai-gpt-4.1")

        assert isinstance(llm, SnowflakeCompletion)
        assert llm.provider == "snowflake"
        assert llm.model == "openai-gpt-4.1"
        assert llm.is_litellm is False

    def test_explicit_provider_creates_native_snowflake_provider(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)

        llm = LLM(model="claude-sonnet-4-5", provider="snowflake")

        assert isinstance(llm, SnowflakeCompletion)
        assert llm.model == "claude-sonnet-4-5"


class TestSnowflakeRequests:
    def test_prepare_completion_params_uses_snowflake_model_name(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(
            model="openai-gpt-4.1",
            temperature=0.2,
            max_completion_tokens=128,
        )

        params = llm._prepare_completion_params(
            [{"role": "user", "content": "Hello"}]
        )

        assert params["model"] == "openai-gpt-4.1"
        assert params["temperature"] == 0.2
        assert params["max_completion_tokens"] == 128
        assert params["messages"] == [{"role": "user", "content": "Hello"}]

    def test_claude_model_removes_trailing_assistant_prefill(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="claude-sonnet-4-5")

        messages = llm._format_messages(
            [
                {"role": "user", "content": "Write a summary."},
                {"role": "assistant", "content": "Here is"},
            ]
        )

        assert messages == [{"role": "user", "content": "Write a summary."}]

    def test_claude_model_normalizes_stringified_tool_calls_with_results(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="claude-sonnet-4-5")

        messages = llm._format_messages(
            [
                {"role": "user", "content": "Use the tools."},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        "{'id': 'toolu_1', 'type': 'function', 'function': {'name': \"'search_the_internet_with_serper'\", 'arguments': '\\\'{\"search_query\":\"CrewAI tools\"}\\\''}}",
                        "{'id': 'toolu_2', 'type': 'function', 'function': {'name': \"'search_the_internet_with_serper'\", 'arguments': '\\\'{\"search_query\":\"CrewAI demos\"}\\\''}}",
                    ],
                },
                {
                    "role": "tool",
                    "tool_call_id": "toolu_1",
                    "name": "search_the_internet_with_serper",
                    "content": "result 1",
                },
                {
                    "role": "tool",
                    "tool_call_id": "toolu_2",
                    "name": "search_the_internet_with_serper",
                    "content": "result 2",
                },
            ]
        )

        assert messages[-2] == {"role": "user", "content": "Use the tools."}
        assert messages[-1]["role"] == "user"
        assert "result 1" in messages[-1]["content"]
        assert "result 2" in messages[-1]["content"]
        assert all("tool_calls" not in message for message in messages)

    def test_claude_model_removes_dangling_tool_call_without_result(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="claude-sonnet-4-5")

        messages = llm._format_messages(
            [
                {"role": "user", "content": "Use the tool."},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "call_1",
                            "type": "function",
                            "function": {"name": "lookup", "arguments": "{}"},
                        }
                    ],
                },
            ]
        )

        assert messages == [{"role": "user", "content": "Use the tool."}]

    def test_claude_model_preserves_complete_tool_call_result_pair(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="claude-sonnet-4-5")

        messages = llm._format_messages(
            [
                {"role": "user", "content": "Use the tool."},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "call_1",
                            "type": "function",
                            "function": {"name": "lookup", "arguments": "{}"},
                        }
                    ],
                },
                {
                    "role": "tool",
                    "tool_call_id": "call_1",
                    "content": "result",
                },
            ]
        )

        assert messages[-2] == {"role": "user", "content": "Use the tool."}
        assert messages[-1]["role"] == "user"
        assert "result" in messages[-1]["content"]
        assert all("tool_calls" not in message for message in messages)

    def test_claude_model_drops_unrelated_tool_results_from_preserved_pair(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="claude-sonnet-4-5")

        messages = llm._format_messages(
            [
                {"role": "user", "content": "Use the tool."},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "call_1",
                            "type": "function",
                            "function": {"name": "lookup", "arguments": "{}"},
                        }
                    ],
                },
                {
                    "role": "tool",
                    "tool_call_id": "call_1",
                    "content": "valid result",
                },
                {
                    "role": "tool",
                    "tool_call_id": "unrelated_call",
                    "content": "unrelated result",
                },
            ]
        )

        assert messages[-2] == {"role": "user", "content": "Use the tool."}
        assert messages[-1]["role"] == "user"
        assert "valid result" in messages[-1]["content"]
        assert "unrelated result" not in messages[-1]["content"]
        assert all("tool_call_id" not in message for message in messages)

    def test_claude_model_removes_dangling_tool_use_content_block(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="claude-sonnet-4-5")

        messages = llm._format_messages(
            [
                {"role": "user", "content": "Use the tool."},
                {
                    "role": "assistant",
                    "content": [
                        {
                            "toolUse": {
                                "toolUseId": "tooluse_1",
                                "name": "lookup",
                                "input": {},
                            }
                        }
                    ],
                },
                {"role": "user", "content": "Continue."},
            ]
        )

        assert messages == [
            {"role": "user", "content": "Use the tool."},
            {"role": "user", "content": "Continue."},
        ]

    def test_claude_model_preserves_complete_tool_use_content_block_pair(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="claude-sonnet-4-5")

        messages = llm._format_messages(
            [
                {"role": "user", "content": "Use the tool."},
                {
                    "role": "assistant",
                    "content": [
                        {
                            "toolUse": {
                                "toolUseId": "tooluse_1",
                                "name": "lookup",
                                "input": {},
                            }
                        }
                    ],
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "toolResult": {
                                "toolUseId": "tooluse_1",
                                "content": [{"text": "result"}],
                            }
                        }
                    ],
                },
            ]
        )

        assert messages[-2] == {"role": "user", "content": "Use the tool."}
        assert messages[-1]["role"] == "user"
        assert "result" in messages[-1]["content"]
        assert "toolResult" not in messages[-1]["content"]
        assert all(
            not (
                message.get("role") == "assistant"
                and isinstance(message.get("content"), list)
            )
            for message in messages
        )

    def test_claude_model_maps_max_tokens_to_max_completion_tokens(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="claude-sonnet-4-5", max_tokens=256)

        params = llm._prepare_completion_params(
            [{"role": "user", "content": "Hello"}]
        )

        assert "max_tokens" not in params
        assert params["max_completion_tokens"] == 256

    def test_streaming_params_include_usage(self, monkeypatch: pytest.MonkeyPatch):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="openai-gpt-4.1", stream=True)

        params = llm._prepare_completion_params(
            [{"role": "user", "content": "Hello"}]
        )

        assert params["stream"] is True
        assert params["stream_options"] == {"include_usage": True}

    def test_non_streaming_call_uses_native_openai_client(
        self, monkeypatch: pytest.MonkeyPatch
    ):
        _snowflake_env(monkeypatch)
        llm = SnowflakeCompletion(model="openai-gpt-4.1")
        fake_response = SimpleNamespace(
            usage=SimpleNamespace(
                prompt_tokens=3,
                completion_tokens=2,
                total_tokens=5,
                prompt_tokens_details=None,
                completion_tokens_details=None,
            ),
            choices=[
                SimpleNamespace(
                    message=SimpleNamespace(content="Snowflake response", tool_calls=None)
                )
            ],
        )
        create = Mock(return_value=fake_response)
        fake_client = SimpleNamespace(
            chat=SimpleNamespace(completions=SimpleNamespace(create=create))
        )

        with patch.object(llm, "_get_sync_client", return_value=fake_client):
            response = llm.call([{"role": "user", "content": "Hello"}])

        assert response == "Snowflake response"
        create.assert_called_once()
        assert create.call_args.kwargs["model"] == "openai-gpt-4.1"
        assert create.call_args.kwargs["messages"] == [
            {"role": "user", "content": "Hello"}
        ]