feat: Mistral: support auth via managed identity (#265) (#266)

Author: adubovik

aidial_adapter_openai/endpoints/chat_completion.py

@@ -116,7 +116,7 @@ async def call_chat_completion(
             )
 
         case ChatCompletionDeploymentType.MISTRAL:
-            return await mistral_chat_completion(data, upstream_endpoint, creds)
+            return await mistral_chat_completion(data, endpoint, creds)
         case ChatCompletionDeploymentType.DATABRICKS:
             return await databricks_chat_completion(data, endpoint, creds)
 

aidial_adapter_openai/mistral.py

@@ -1,23 +1,25 @@
-from typing import Any
+from typing import Any, cast
 
-from openai import AsyncOpenAI, AsyncStream
+from openai import AsyncStream
 from openai.types.chat.chat_completion import ChatCompletion
 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 
 from aidial_adapter_openai.utils.auth import OpenAICreds
-from aidial_adapter_openai.utils.http_client import get_http_client
+from aidial_adapter_openai.utils.parsers import (
+    AzureOpenAIEndpoint,
+    OpenAIEndpoint,
+    OpenAIParams,
+)
 from aidial_adapter_openai.utils.reflection import call_with_extra_body
 from aidial_adapter_openai.utils.streaming import chunk_to_dict, map_stream
 
 
 async def chat_completion(
-    data: Any, upstream_endpoint: str, creds: OpenAICreds
+    data: Any,
+    endpoint: AzureOpenAIEndpoint | OpenAIEndpoint,
+    creds: OpenAICreds,
 ):
-    client = AsyncOpenAI(
-        base_url=upstream_endpoint,
-        api_key=creds.get("api_key"),
-        http_client=get_http_client(),
-    )
+    client = endpoint.get_client(cast(OpenAIParams, creds))
 
     response: AsyncStream[ChatCompletionChunk] | ChatCompletion = (
         await call_with_extra_body(client.chat.completions.create, data)

aidial_adapter_openai/utils/parsers.py

@@ -6,7 +6,6 @@
 from fastapi import Request
 from openai import AsyncAzureOpenAI, AsyncOpenAI, Timeout
 
-from aidial_adapter_openai.utils.auth import OpenAICreds
 from aidial_adapter_openai.utils.http_client import get_http_client
 from aidial_adapter_openai.utils.pydantic import ExtraForbidModel
 
@@ -44,33 +43,21 @@ def get_client(self, params: OpenAIParams) -> AsyncAzureOpenAI:
             http_client=get_http_client(),
         )
 
-    def get_auth_headers(self, creds: OpenAICreds) -> dict[str, str]:
-        if key := creds.get("api_key"):
-            return {"api-key": key}
-
-        if token := creds.get("azure_ad_token"):
-            return {"Authorization": f"Bearer {token}"}
-
-        raise ValueError("Invalid credentials")
-
 
 class OpenAIEndpoint(ExtraForbidModel):
     base_url: str
 
     def get_client(self, params: OpenAIParams) -> AsyncOpenAI:
+        api_key = params.get("api_key") or params.get("azure_ad_token")
+
         return AsyncOpenAI(
             base_url=self.base_url,
-            api_key=params.get("api_key"),
+            api_key=api_key,
             timeout=params.get("timeout"),
             max_retries=_MAX_RETRIES,
             http_client=get_http_client(),
         )
 
-    def get_auth_headers(self, creds: OpenAICreds) -> dict[str, str]:
-        if key := (creds.get("api_key") or creds.get("azure_ad_token")):
-            return {"Authorization": f"Bearer {key}"}
-        raise ValueError("Invalid credentials")
-
 
 def _parse_endpoint(
     name: str | None, endpoint: str