import json import time from typing import AsyncIterable from mlflow.gateway.config import AnthropicConfig, RouteConfig from mlflow.gateway.constants import ( MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS, MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS, ) from mlflow.gateway.exceptions import AIGatewayException from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter from mlflow.gateway.providers.utils import rename_payload_keys, send_request, send_stream_request from mlflow.gateway.schemas import chat, completions class AnthropicAdapter(ProviderAdapter): @classmethod def chat_to_model(cls, payload, config): key_mapping = {"stop": "stop_sequences"} payload["model"] = config.model.name payload = rename_payload_keys(payload, key_mapping) if "top_p" in payload and "temperature" in payload: raise AIGatewayException( status_code=422, detail="Cannot set both 'temperature' and 'top_p' parameters." ) max_tokens = payload.get("max_tokens", MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS) if max_tokens > MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS: raise AIGatewayException( status_code=422, detail="Invalid value for max_tokens: cannot exceed " f"{MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS}.", ) payload["max_tokens"] = max_tokens if payload.pop("n", 1) != 1: raise AIGatewayException( status_code=422, detail="'n' must be '1' for the Anthropic provider. Received value: '{n}'.", ) # Cohere uses `system` to set the system message # we concatenate all system messages from the user with a newline system_messages = [m for m in payload["messages"] if m["role"] == "system"] if system_messages: payload["system"] = "\n".join(m["content"] for m in system_messages) # remaining messages are chat history # we want to include only user and assistant messages payload["messages"] = [m for m in payload["messages"] if m["role"] in ("user", "assistant")] # The range of Anthropic's temperature is 0-1, but ours is 0-2, so we halve it if "temperature" in payload: payload["temperature"] = 0.5 * payload["temperature"] return payload @classmethod def model_to_chat(cls, resp, config): # API reference: https://docs.anthropic.com/en/api/messages#body-messages # # Example response: # ``` # { # "content": [ # { # "text": "Blue is often seen as a calming and soothing color.", # "type": "text" # }, # { # "source": { # "type": "base64", # "media_type": "image/jpeg", # "data": "/9j/4AAQSkZJRg...", # "type": "image", # } # } # ], # "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF", # "model": "claude-2.1", # "role": "assistant", # "stop_reason": "end_turn", # "stop_sequence": null, # "type": "message", # "usage": { # "input_tokens": 10, # "output_tokens": 25 # } # } # ``` from mlflow.anthropic.chat import convert_message_to_mlflow_chat stop_reason = "length" if resp["stop_reason"] == "max_tokens" else "stop" return chat.ResponsePayload( id=resp["id"], created=int(time.time()), object="chat.completion", model=resp["model"], choices=[ chat.Choice( index=0, # TODO: Remove this casting once # https://github.com/mlflow/mlflow/pull/14160 is merged message=chat.ResponseMessage( **convert_message_to_mlflow_chat(resp).model_dump_compat() ), finish_reason=stop_reason, ) ], usage=chat.ChatUsage( prompt_tokens=resp["usage"]["input_tokens"], completion_tokens=resp["usage"]["output_tokens"], total_tokens=resp["usage"]["input_tokens"] + resp["usage"]["output_tokens"], ), ) @classmethod def chat_streaming_to_model(cls, payload, config): return cls.chat_to_model(payload, config) @classmethod def model_to_chat_streaming(cls, resp, config): content = resp.get("delta") or resp.get("content_block") or {} if (stop_reason := content.get("stop_reason")) is not None: stop_reason = "length" if stop_reason == "max_tokens" else "stop" return chat.StreamResponsePayload( id=resp["id"], created=int(time.time()), model=resp["model"], choices=[ chat.StreamChoice( index=resp["index"], finish_reason=stop_reason, delta=chat.StreamDelta( role=None, content=content.get("text"), ), ) ], ) @classmethod def model_to_completions(cls, resp, config): stop_reason = "stop" if resp["stop_reason"] == "stop_sequence" else "length" return completions.ResponsePayload( created=int(time.time()), object="text_completion", model=resp["model"], choices=[ completions.Choice( index=0, text=resp["completion"], finish_reason=stop_reason, ) ], usage=completions.CompletionsUsage( prompt_tokens=None, completion_tokens=None, total_tokens=None, ), ) @classmethod def completions_to_model(cls, payload, config): key_mapping = {"max_tokens": "max_tokens_to_sample", "stop": "stop_sequences"} payload["model"] = config.model.name if "top_p" in payload: raise AIGatewayException( status_code=422, detail="Cannot set both 'temperature' and 'top_p' parameters. " "Please use only the temperature parameter for your query.", ) max_tokens = payload.get("max_tokens", MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS) if max_tokens > MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS: raise AIGatewayException( status_code=422, detail="Invalid value for max_tokens: cannot exceed " f"{MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS}.", ) payload["max_tokens"] = max_tokens if payload.get("stream", False): raise AIGatewayException( status_code=422, detail="Setting the 'stream' parameter to 'true' is not supported with the MLflow " "Gateway.", ) n = payload.pop("n", 1) if n != 1: raise AIGatewayException( status_code=422, detail=f"'n' must be '1' for the Anthropic provider. Received value: '{n}'.", ) payload = rename_payload_keys(payload, key_mapping) if payload["prompt"].startswith("Human: "): payload["prompt"] = "\n\n" + payload["prompt"] if not payload["prompt"].startswith("\n\nHuman: "): payload["prompt"] = "\n\nHuman: " + payload["prompt"] if not payload["prompt"].endswith("\n\nAssistant:"): payload["prompt"] = payload["prompt"] + "\n\nAssistant:" # The range of Anthropic's temperature is 0-1, but ours is 0-2, so we halve it if "temperature" in payload: payload["temperature"] = 0.5 * payload["temperature"] return payload @classmethod def embeddings_to_model(cls, payload, config): raise NotImplementedError @classmethod def model_to_embeddings(cls, resp, config): raise NotImplementedError class AnthropicProvider(BaseProvider, AnthropicAdapter): NAME = "Anthropic" CONFIG_TYPE = AnthropicConfig def __init__(self, config: RouteConfig) -> None: super().__init__(config) if config.model.config is None or not isinstance(config.model.config, AnthropicConfig): raise TypeError(f"Invalid config type {config.model.config}") self.anthropic_config: AnthropicConfig = config.model.config @property def headers(self) -> dict[str, str]: return { "x-api-key": self.anthropic_config.anthropic_api_key, "anthropic-version": self.anthropic_config.anthropic_version, } @property def base_url(self) -> str: return "https://api.anthropic.com/v1" @property def adapter_class(self) -> type[ProviderAdapter]: return AnthropicAdapter def get_endpoint_url(self, route_type: str) -> str: if route_type == "llm/v1/chat": return f"{self.base_url}/messages" elif route_type == "llm/v1/completions": return f"{self.base_url}/complete" else: raise ValueError(f"Invalid route type {route_type}") async def chat_stream( self, payload: chat.RequestPayload ) -> AsyncIterable[chat.StreamResponsePayload]: from fastapi.encoders import jsonable_encoder payload = jsonable_encoder(payload, exclude_none=True) self.check_for_model_field(payload) stream = send_stream_request( headers=self.headers, base_url=self.base_url, path="messages", payload=AnthropicAdapter.chat_streaming_to_model(payload, self.config), ) indices = [] metadata = {} async for chunk in stream: chunk = chunk.strip() if not chunk: continue # No handling on "event" lines prefix, content = chunk.split(b":", 1) if prefix != b"data": continue # See https://docs.anthropic.com/claude/reference/messages-streaming resp = json.loads(content.decode("utf-8")) # response id and model are only present in `message_start` if resp["type"] == "message_start": metadata["id"] = resp["message"]["id"] metadata["model"] = resp["message"]["model"] continue if resp["type"] not in ( "message_delta", "content_block_start", "content_block_delta", ): continue index = resp.get("index") if index is not None and index not in indices: indices.append(index) resp.update(metadata) if resp["type"] == "message_delta": for index in indices: yield AnthropicAdapter.model_to_chat_streaming( {**resp, "index": index}, self.config, ) else: yield AnthropicAdapter.model_to_chat_streaming(resp, self.config) async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload: from fastapi.encoders import jsonable_encoder payload = jsonable_encoder(payload, exclude_none=True) self.check_for_model_field(payload) resp = await send_request( headers=self.headers, base_url=self.base_url, path="messages", payload=AnthropicAdapter.chat_to_model(payload, self.config), ) return AnthropicAdapter.model_to_chat(resp, self.config) async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload: from fastapi.encoders import jsonable_encoder payload = jsonable_encoder(payload, exclude_none=True) self.check_for_model_field(payload) resp = await send_request( headers=self.headers, base_url=self.base_url, path="complete", payload=AnthropicAdapter.completions_to_model(payload, self.config), ) # Example response: # Documentation: https://docs.anthropic.com/claude/reference/complete_post # ``` # { # "completion": " Hello! My name is Claude." # "stop_reason": "stop_sequence", # "model": "claude-instant-1.1", # "truncated": False, # "stop": None, # "log_id": "dee173f87ddf1357da639dee3c38d833", # "exception": None, # } # ``` return AnthropicAdapter.model_to_completions(resp, self.config)