""" Snowflake Cortex REST API — Chat Transformation Routes to native Cortex REST API endpoints based on model: - Claude models → POST /api/v2/cortex/v1/messages (Anthropic format) - All other models → POST /api/v2/cortex/v1/chat/completions (OpenAI format) Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api """ import json from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict import httpx from typing_extensions import ReadOnly from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.utils import ( ChatCompletionMessageToolCall, ChatCompletionUsageBlock, Choices, Function, GenericStreamingChunk, Message, ModelResponse, Usage, ) from ...base_llm.base_model_iterator import BaseModelResponseIterator from ...openai_like.chat.transformation import OpenAIGPTConfig from ..utils import SnowflakeBaseConfig if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj else: LiteLLMLoggingObj = Any ANTHROPIC_VERSION: Final = "2023-06-01" _CLAUDE_MODEL_PREFIXES: Final = ( "claude-", "claude_", ) class _AnthropicContentBlock(TypedDict, total=False): type: ReadOnly[str] text: ReadOnly[str] id: ReadOnly[str] name: ReadOnly[str] input: ReadOnly[Mapping[str, object]] class _AnthropicUsageBlock(TypedDict, total=False): input_tokens: ReadOnly[int] output_tokens: ReadOnly[int] class _AnthropicMessagesResponse(TypedDict, total=False): id: ReadOnly[str] model: ReadOnly[str] stop_reason: ReadOnly[str] content: ReadOnly[Sequence[_AnthropicContentBlock]] usage: ReadOnly[_AnthropicUsageBlock] class _ChatCompletionsResponse(Protocol): """Response view that decodes the Cortex chat-completions body as a field mapping.""" def json(self) -> Mapping[str, object]: ... class _MessagesResponse(Protocol): """Response view that decodes the Cortex messages body in Anthropic shape.""" def json(self) -> _AnthropicMessagesResponse: ... def _decoded_chat_completions(response: _ChatCompletionsResponse) -> Mapping[str, object]: return response.json() def _decoded_messages(response: _MessagesResponse) -> _AnthropicMessagesResponse: return response.json() def _is_claude_model(model: str) -> bool: """Return True if model name (after stripping snowflake/ prefix) is a Claude model.""" name: Final = model.lower().removeprefix("snowflake/") return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES) class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): """ Snowflake Cortex REST API — unified provider. Auto-routes based on model name: - Claude models → /api/v2/cortex/v1/messages (Anthropic Messages format) - All others → /api/v2/cortex/v1/chat/completions (OpenAI format) Auth: PAT: api_key="pat/" → X-Snowflake-Authorization-Token-Type: PROGRAMMATIC_ACCESS_TOKEN JWT: api_key="" → X-Snowflake-Authorization-Token-Type: KEYPAIR_JWT """ @classmethod def get_config(cls): return super().get_config() def get_supported_openai_params(self, model: str) -> list[str]: params: Final = [ "temperature", "max_tokens", "max_completion_tokens", "top_p", "stream", "tools", "tool_choice", ] if _is_claude_model(model): params.append("thinking") return params def get_complete_url( self, api_base: str | None, api_key: str | None, model: str, optional_params: dict, litellm_params: dict, stream: bool | None = None, ) -> str: api_base = self._get_api_base(api_base, optional_params) if _is_claude_model(model): return f"{api_base}/cortex/v1/messages" return f"{api_base}/cortex/v1/chat/completions" def validate_environment( self, headers: dict, model: str, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, api_base: str | None = None, ) -> dict: headers = super().validate_environment( headers=headers, model=model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, api_key=api_key, api_base=api_base, ) if _is_claude_model(model): headers["anthropic-version"] = ANTHROPIC_VERSION return headers def _transform_tools_to_anthropic(self, tools: list[dict]) -> list[dict]: """ Convert tools from OpenAI format to Anthropic format. OpenAI: {"type": "function", "function": {"name": ..., "parameters": {...}}} Anthropic: {"name": ..., "description": ..., "input_schema": {...}} """ anthropic_tools: Final = [] for tool in tools: if tool.get("type") == "function" and "function" in tool: func = tool["function"] anthropic_tool: dict[str, object] = { "name": func.get("name", ""), } if "description" in func: anthropic_tool["description"] = func["description"] if "parameters" in func: anthropic_tool["input_schema"] = func["parameters"] else: anthropic_tool["input_schema"] = { "type": "object", "properties": {}, } anthropic_tools.append(anthropic_tool) else: anthropic_tools.append(tool) return anthropic_tools def _extract_system_and_messages(self, messages: list[AllMessageValues]) -> tuple[str | None, list[dict]]: """ Split messages into system prompt and conversation turns for Anthropic format. - system messages → collected and joined (preserves guardrail prompts) - assistant messages with tool_calls → tool_use content blocks - tool role messages → user role with tool_result content blocks """ system_parts: Final[list[str]] = [] conversation: Final[list[dict]] = [] for msg in messages: if isinstance(msg, dict): role = msg.get("role", "") content: Any = msg.get("content", "") else: role = getattr(msg, "role", "") content = getattr(msg, "content", "") if role == "system": if isinstance(content, str) and content: system_parts.append(content) elif isinstance(content, list): system_parts.append("\n".join(b.get("text", "") for b in content if b.get("type") == "text")) elif role == "assistant": tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None) if tool_calls: content_blocks: list[dict[str, object]] = [] if content: content_blocks.append({"type": "text", "text": content}) for tc in tool_calls: func = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", {}) tc_id = tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", "") func_name = func.get("name", "") if isinstance(func, dict) else getattr(func, "name", "") func_args = ( func.get("arguments", "{}") if isinstance(func, dict) else getattr(func, "arguments", "{}") ) try: input_data = json.loads(func_args) if isinstance(func_args, str) else func_args except (json.JSONDecodeError, TypeError): input_data = {} content_blocks.append( { "type": "tool_use", "id": tc_id, "name": func_name, "input": input_data, } ) conversation.append({"role": "assistant", "content": content_blocks}) else: conversation.append({"role": "assistant", "content": content}) elif role == "tool": tool_call_id = ( msg.get("tool_call_id", "") if isinstance(msg, dict) else getattr(msg, "tool_call_id", "") ) tool_content = content if isinstance(content, str) else json.dumps(content) tool_result_block = { "type": "tool_result", "tool_use_id": tool_call_id, "content": tool_content, } if ( conversation and conversation[-1]["role"] == "user" and isinstance(conversation[-1]["content"], list) and conversation[-1]["content"] and conversation[-1]["content"][0].get("type") == "tool_result" ): conversation[-1]["content"].append(tool_result_block) else: conversation.append({"role": "user", "content": [tool_result_block]}) else: conversation.append({"role": role, "content": content}) system: Final[str | None] = "\n\n".join(system_parts) if system_parts else None return system, conversation def transform_request( self, model: str, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, headers: dict, ) -> dict: stream: Final[bool] = optional_params.pop("stream", False) or False extra_body: Final = optional_params.pop("extra_body", {}) if _is_claude_model(model): return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body) return self._transform_request_openai(model, messages, optional_params, stream, extra_body) def _transform_request_openai( self, model: str, messages: list[AllMessageValues], optional_params: dict, stream: bool, extra_body: dict, ) -> dict: """OpenAI format for /chat/completions endpoint.""" max_tokens: Final = optional_params.pop("max_tokens", None) max_completion_tokens: Final = optional_params.pop("max_completion_tokens", None) resolved_max: Final = max_completion_tokens or max_tokens body: Final[dict] = { "model": model.removeprefix("snowflake/"), "messages": messages, "stream": stream, **optional_params, **extra_body, } if resolved_max is not None: body["max_completion_tokens"] = resolved_max return body def _transform_tool_choice_to_anthropic(self, tool_choice: Any) -> dict[str, Any]: """ Convert tool_choice from OpenAI format to Anthropic format. OpenAI string values: "auto", "required", "none" OpenAI dict: {"type": "function", "function": {"name": "..."}} Anthropic: {"type": "auto"}, {"type": "any"}, {"type": "tool", "name": "..."} """ if isinstance(tool_choice, str): mapping: Final = { "auto": {"type": "auto"}, "required": {"type": "any"}, "none": {"type": "none"}, } return mapping.get(tool_choice, {"type": "auto"}) elif isinstance(tool_choice, dict): if tool_choice.get("type") == "function": func: Final = tool_choice.get("function", {}) return {"type": "tool", "name": func.get("name", "")} return tool_choice return {"type": "auto"} def _transform_request_anthropic( self, model: str, messages: list[AllMessageValues], optional_params: dict, stream: bool, extra_body: dict, ) -> dict: """Anthropic Messages format for /messages endpoint.""" system, conversation = self._extract_system_and_messages(messages) if "tools" in optional_params: optional_params["tools"] = self._transform_tools_to_anthropic(optional_params["tools"]) if "tool_choice" in optional_params: optional_params["tool_choice"] = self._transform_tool_choice_to_anthropic(optional_params["tool_choice"]) max_completion_tokens: Final = optional_params.pop("max_completion_tokens", None) if max_completion_tokens and "max_tokens" not in optional_params: optional_params["max_tokens"] = max_completion_tokens model_name: Final = model.removeprefix("snowflake/") body: Final[dict[str, object]] = { "model": model_name, "messages": conversation, "stream": stream, **optional_params, **extra_body, } if system is not None: body["system"] = system if "max_tokens" not in body: body["max_tokens"] = 4096 # reasonable default; Anthropic API max varies by model return body def transform_response( self, model: str, raw_response: httpx.Response, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, request_data: dict, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: if _is_claude_model(model): return self._transform_response_anthropic( model, raw_response, model_response, logging_obj, request_data, messages ) return self._transform_response_openai(model, raw_response, model_response, logging_obj, request_data, messages) def _transform_response_openai( self, model: str, raw_response: httpx.Response, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, request_data: dict, messages: list[AllMessageValues], ) -> ModelResponse: """Parse standard OpenAI chat completions response.""" response_json: Final = _decoded_chat_completions(raw_response) logging_obj.post_call( input=messages, api_key="", original_response=response_json, additional_args={"complete_input_dict": request_data}, ) returned_response: Final = ModelResponse(**response_json) returned_response.model = "snowflake/" + (returned_response.model or "") if model is not None: returned_response._hidden_params["model"] = model return returned_response def _transform_response_anthropic( self, model: str, raw_response: httpx.Response, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, request_data: dict, messages: list[AllMessageValues], ) -> ModelResponse: """Parse Anthropic Messages response into OpenAI format.""" response_json: Final = _decoded_messages(raw_response) logging_obj.post_call( input=messages, api_key="", original_response=response_json, additional_args={"complete_input_dict": request_data}, ) text_content = "" tool_calls: Final = [] for block in response_json.get("content", []): if block.get("type") == "text": text_content += block.get("text", "") elif block.get("type") == "tool_use": tool_calls.append( ChatCompletionMessageToolCall( id=block.get("id", ""), type="function", function=Function( name=block.get("name", ""), arguments=json.dumps(block.get("input", {})), ), ) ) _stop_reason_map: Final = { "end_turn": "stop", "max_tokens": "length", "tool_use": "tool_calls", "stop_sequence": "stop", } finish_reason: Final = _stop_reason_map.get(response_json.get("stop_reason", "end_turn"), "stop") message: Final = Message(content=text_content or None, role="assistant") if tool_calls: message.tool_calls = tool_calls choice: Final = Choices( finish_reason=finish_reason, index=0, message=message, ) usage_data: Final = response_json.get("usage", {}) usage: Final = Usage( prompt_tokens=usage_data.get("input_tokens", 0), completion_tokens=usage_data.get("output_tokens", 0), total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0), ) model_response.choices = [choice] model_response.usage = usage model_response.model = "snowflake/" + response_json.get("model", model) model_response.id = response_json.get("id", "") if model is not None: model_response._hidden_params["model"] = model return model_response def get_model_response_iterator( self, streaming_response: object, sync_stream: bool, json_mode: bool | None = False, ) -> "SnowflakeStreamingHandler": return SnowflakeStreamingHandler( streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, ) class SnowflakeStreamingHandler(BaseModelResponseIterator): """ Parse streaming events from both Snowflake endpoints. - /chat/completions: OpenAI SSE format (has "choices" key) - /messages: Anthropic SSE format (has "type" key like content_block_delta) """ def __init__( self, streaming_response: object, sync_stream: bool, json_mode: bool | None = False, ): super().__init__(streaming_response=streaming_response, sync_stream=sync_stream) self._tool_index = 0 self._tool_id = "" self._tool_name = "" self._input_tokens = 0 def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: if "choices" in chunk: return self._parse_openai_chunk(chunk) return self._parse_anthropic_chunk(chunk) def _parse_openai_chunk(self, chunk: dict) -> GenericStreamingChunk: choices: Final = chunk.get("choices", []) if not choices: return GenericStreamingChunk( text="", is_finished=False, finish_reason="", usage=None, index=0, tool_use=None, ) choice: Final = choices[0] delta: Final = choice.get("delta", {}) finish_reason: Final = choice.get("finish_reason") or "" text: Final = delta.get("content") or "" tool_use = None tool_calls: Final = delta.get("tool_calls") if tool_calls: tc: Final = tool_calls[0] func: Final = tc.get("function", {}) tool_use = ChatCompletionToolCallChunk( id=tc.get("id", ""), type="function", function={ "name": func.get("name", ""), "arguments": func.get("arguments", ""), }, index=tc.get("index", 0), ) return GenericStreamingChunk( text=text, is_finished=finish_reason != "", finish_reason=finish_reason, usage=None, index=choice.get("index", 0), tool_use=tool_use, ) def _parse_anthropic_chunk(self, chunk: dict) -> GenericStreamingChunk: event_type: Final = chunk.get("type", "") if event_type == "message_start": message: Final = chunk.get("message", {}) usage_data = message.get("usage", {}) self._input_tokens = usage_data.get("input_tokens", 0) return GenericStreamingChunk( text="", is_finished=False, finish_reason="", usage=None, index=0, tool_use=None, ) elif event_type == "content_block_delta": delta = chunk.get("delta", {}) delta_type: Final = delta.get("type", "") if delta_type == "text_delta": return GenericStreamingChunk( text=delta.get("text", ""), is_finished=False, finish_reason="", usage=None, index=chunk.get("index", 0), tool_use=None, ) elif delta_type == "input_json_delta": return GenericStreamingChunk( text="", is_finished=False, finish_reason="", usage=None, index=chunk.get("index", 0), tool_use=ChatCompletionToolCallChunk( id=self._tool_id, type="function", function={ "name": self._tool_name, "arguments": delta.get("partial_json", ""), }, index=self._tool_index, ), ) elif event_type == "content_block_start": content_block: Final = chunk.get("content_block", {}) if content_block.get("type") == "tool_use": self._tool_id = content_block.get("id", "") self._tool_name = content_block.get("name", "") self._tool_index = chunk.get("index", 0) return GenericStreamingChunk( text="", is_finished=False, finish_reason="", usage=None, index=chunk.get("index", 0), tool_use=ChatCompletionToolCallChunk( id=self._tool_id, type="function", function={"name": self._tool_name, "arguments": ""}, index=self._tool_index, ), ) elif event_type == "message_delta": delta = chunk.get("delta", {}) stop_reason: Final = delta.get("stop_reason", "") usage_data = chunk.get("usage", {}) _stop_map: Final = { "end_turn": "stop", "max_tokens": "length", "tool_use": "tool_calls", "stop_sequence": "stop", } usage = None if usage_data or self._input_tokens: output_t: Final = usage_data.get("output_tokens", 0) input_t: Final = self._input_tokens or usage_data.get("input_tokens", 0) usage = ChatCompletionUsageBlock( prompt_tokens=input_t, completion_tokens=output_t, total_tokens=input_t + output_t, ) return GenericStreamingChunk( text="", is_finished=True, finish_reason=_stop_map.get(stop_reason, "stop"), usage=usage, index=0, tool_use=None, ) elif event_type == "message_stop": return GenericStreamingChunk( text="", is_finished=True, finish_reason="stop", usage=None, index=0, tool_use=None, ) return GenericStreamingChunk( text="", is_finished=False, finish_reason="", usage=None, index=0, tool_use=None, )