Add anthropic_version header for AnthropicProvider (mlflow#11056)

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>
StateFarmIns · Feb 20, 2024 · 0b9aa67 · 0b9aa67
1 parent e2a3e59
commit 0b9aa67
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/mlflow/gateway/config.py b/mlflow/gateway/config.py
@@ -167,6 +167,7 @@ def validate_field_compatibility(cls, config: Dict[str, Any]):
 
 class AnthropicConfig(ConfigModel):
     anthropic_api_key: str
+    anthropic_version: str = "2023-06-01"
 
     # pylint: disable=no-self-argument
     @validator("anthropic_api_key", pre=True)

diff --git a/mlflow/gateway/providers/anthropic.py b/mlflow/gateway/providers/anthropic.py
@@ -104,7 +104,10 @@ def __init__(self, config: RouteConfig) -> None:
         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
-        self.headers = {"x-api-key": self.anthropic_config.anthropic_api_key}
+        self.headers = {
+            "x-api-key": self.anthropic_config.anthropic_api_key,
+            "anthropic-version": self.anthropic_config.anthropic_version,
+        }
         self.base_url = "https://api.anthropic.com/v1/"
 
     async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload: