Extended middleware and Mistral Chat Completions support

microsoft · Aug 16, 2024 · bbfebd6 · bbfebd6
1 parent b36f741
commit bbfebd6
Show file tree

Hide file tree

Showing 16 changed files with 346 additions and 103 deletions.
diff --git a/src/AzureAIProxy/Middleware/ApiKeyAuthenticationHandler.cs b/src/AzureAIProxy/Middleware/ApiKeyAuthenticationHandler.cs
@@ -20,15 +20,13 @@ protected override async Task<AuthenticateResult> HandleAuthenticateAsync()
 
         var apiKey = apiKeyValues.ToString(); // Convert StringValues to string
         if (string.IsNullOrWhiteSpace(apiKey))
-            return AuthenticateResult.Fail("Missing API key is empty.");
+            return AuthenticateResult.Fail("API key is empty.");
 
         var requestContext = await authorizeService.IsUserAuthorizedAsync(apiKey);
         if (requestContext is null)
             return AuthenticateResult.Fail("Authentication failed.");
 
         Context.Items["RequestContext"] = requestContext;
-        Context.Items["RateLimited"] = requestContext.RateLimitExceed;
-        Context.Items["DailyRequestCap"] = requestContext.DailyRequestCap;
 
         var identity = new ClaimsIdentity(null, nameof(ApiKeyAuthenticationHandler));
         var ticket = new AuthenticationTicket(new ClaimsPrincipal(identity), Scheme.Name);

diff --git a/src/AzureAIProxy/Middleware/BearerTokenAuthenticationHandler.cs b/src/AzureAIProxy/Middleware/BearerTokenAuthenticationHandler.cs
@@ -0,0 +1,37 @@
+using System.Security.Claims;
+using System.Text.Encodings.Web;
+using Microsoft.AspNetCore.Authentication;
+using Microsoft.Extensions.Options;
+using AzureAIProxy.Services;
+
+namespace AzureAIProxy.Middleware;
+
+public class BearerTokenAuthenticationHandler(
+    IOptionsMonitor<ProxyAuthenticationOptions> options,
+    IAuthorizeService authorizeService,
+    ILoggerFactory logger,
+    UrlEncoder encoder
+) : AuthenticationHandler<ProxyAuthenticationOptions>(options, logger, encoder)
+{
+    protected override async Task<AuthenticateResult> HandleAuthenticateAsync()
+    {
+        if (!Request.Headers.TryGetValue("Authorization", out var apiKeyValues))
+            return AuthenticateResult.Fail("Missing API key is empty.");
+
+        // Extract the API key from the Authorization header
+        var apiKey = apiKeyValues.ToString().Split(" ").Last(); // Convert StringValues to string
+        if (string.IsNullOrWhiteSpace(apiKey))
+            return AuthenticateResult.Fail("API key is empty.");
+
+        var requestContext = await authorizeService.IsUserAuthorizedAsync(apiKey);
+        if (requestContext is null)
+            return AuthenticateResult.Fail("Authentication failed.");
+
+        Context.Items["RequestContext"] = requestContext;
+
+        var identity = new ClaimsIdentity(null, nameof(BearerTokenAuthenticationHandler));
+        var ticket = new AuthenticationTicket(new ClaimsPrincipal(identity), Scheme.Name);
+
+        return AuthenticateResult.Success(ticket);
+    }
+}
diff --git a/src/AzureAIProxy/Middleware/BearerTokenAuthorizeAttribute.cs b/src/AzureAIProxy/Middleware/BearerTokenAuthorizeAttribute.cs
@@ -0,0 +1,12 @@
+using Microsoft.AspNetCore.Authorization;
+using AzureAIProxy.Middleware;
+using System.Diagnostics;
+
+namespace AzureAIProxy;
+
+[AttributeUsage(AttributeTargets.Class | AttributeTargets.Method, AllowMultiple = true, Inherited = true)]
+[DebuggerDisplay("{ToString(),nq}")]
+public class BearerTokenAuthorizeAttribute : AuthorizeAttribute
+{
+    public BearerTokenAuthorizeAttribute() => AuthenticationSchemes = ProxyAuthenticationOptions.BearerTokenScheme;
+}
diff --git a/src/AzureAIProxy/Middleware/JwtAuthenticationHandler.cs b/src/AzureAIProxy/Middleware/JwtAuthenticationHandler.cs
@@ -28,7 +28,7 @@ protected override async Task<AuthenticateResult> HandleAuthenticateAsync()
 
         Context.Items["RequestContext"] = requestContext;
 
-        var identity = new ClaimsIdentity(null, nameof(ApiKeyAuthenticationHandler));
+        var identity = new ClaimsIdentity(null, nameof(JwtAuthenticationHandler));
         var ticket = new AuthenticationTicket(new ClaimsPrincipal(identity), Scheme.Name);
 
         return AuthenticateResult.Success(ticket);

diff --git a/src/AzureAIProxy/Middleware/LoadProperties.cs b/src/AzureAIProxy/Middleware/LoadProperties.cs
@@ -0,0 +1,68 @@
+using AzureAIProxy.Shared.Database;
+using System.Text.Json;
+
+namespace AzureAIProxy.Middleware;
+
+public class LoadProperties(RequestDelegate next)
+{
+    private readonly RequestDelegate _next = next;
+
+    public async Task InvokeAsync(HttpContext context)
+    {
+        JsonDocument? jsonDoc = null;
+        try
+        {
+            if (!context.Request.HasFormContentType &&
+                context.Request.ContentType != null &&
+                context.Request.ContentType.Contains("application/json", StringComparison.InvariantCultureIgnoreCase))
+            {
+                using var reader = new StreamReader(context.Request.Body);
+                string json = await reader.ReadToEndAsync();
+                if (!string.IsNullOrWhiteSpace(json))
+                {
+                    try
+                    {
+                        jsonDoc = JsonDocument.Parse(json);
+                    }
+                    catch (JsonException)
+                    {
+                        await OpenAIErrorResponse.BadRequest($"Invalid JSON in request body: {json}").WriteAsync(context);
+                        return;
+                    }
+
+                    jsonDoc = JsonDocument.Parse(json);
+
+                    context.Items["IsStreaming"] = IsStreaming(jsonDoc);
+                    context.Items["ModelName"] = GetModelName(jsonDoc);
+                }
+            }
+
+            context.Items["requestPath"]= context.Request.Path.Value!.Split("/api/v1/").Last();
+            context.Items["jsonDoc"] = jsonDoc;
+
+            await _next(context);
+        }
+        finally
+        {
+            jsonDoc?.Dispose();
+        }
+    }
+
+    private static bool IsStreaming(JsonDocument requestJsonDoc)
+    {
+        return requestJsonDoc.RootElement.ValueKind == JsonValueKind.Object
+            && requestJsonDoc.RootElement.TryGetProperty("stream", out JsonElement streamElement)
+            && (
+                streamElement.ValueKind == JsonValueKind.True
+                || streamElement.ValueKind == JsonValueKind.False
+            )
+            && streamElement.GetBoolean();
+    }
+
+    private static string? GetModelName(JsonDocument requestJsonDoc)
+    {
+        return requestJsonDoc.RootElement.TryGetProperty("model", out JsonElement modelElement) && modelElement.ValueKind == JsonValueKind.String
+            ? modelElement.GetString()
+            : null;
+    }
+}
diff --git a/src/AzureAIProxy/Middleware/MaxTokensHandler.cs b/src/AzureAIProxy/Middleware/MaxTokensHandler.cs
@@ -0,0 +1,40 @@
+using AzureAIProxy.Shared.Database;
+using System.Text.Json;
+
+namespace AzureAIProxy.Middleware;
+
+public class MaxTokensHandler(RequestDelegate next)
+{
+    private readonly RequestDelegate _next = next;
+
+    public async Task InvokeAsync(HttpContext context)
+    {
+        RequestContext? requestContext = context.Items["RequestContext"]! as RequestContext;
+        JsonDocument? jsonDoc = context.Items["jsonDoc"]! as JsonDocument;
+
+        if (requestContext is not null && jsonDoc is not null)
+        {
+            int? maxTokens = GetMaxTokens(jsonDoc);
+            if (maxTokens.HasValue && maxTokens > requestContext.MaxTokenCap && requestContext.MaxTokenCap > 0)
+            {
+                await OpenAIErrorResponse.BadRequest(
+                    $"max_tokens exceeds the event max token cap of {requestContext.MaxTokenCap}"
+                ).WriteAsync(context);
+                return;
+            }
+        }
+
+        await _next(context);
+    }
+
+    private static int? GetMaxTokens(JsonDocument requestJsonDoc)
+    {
+        return
+            requestJsonDoc.RootElement.ValueKind == JsonValueKind.Object
+            && requestJsonDoc.RootElement.TryGetProperty("max_tokens", out var maxTokensElement)
+            && maxTokensElement.ValueKind == JsonValueKind.Number
+            && maxTokensElement.TryGetInt32(out int maxTokens)
+            ? maxTokens
+            : null;
+    }
+}
diff --git a/src/AzureAIProxy/Middleware/OpenAIErrorResponse.cs b/src/AzureAIProxy/Middleware/OpenAIErrorResponse.cs
@@ -0,0 +1,32 @@
+using System.Net;
+using Microsoft.AspNetCore.Http.HttpResults;
+
+namespace AzureAIProxy.Middleware;
+
+public class OpenAIErrorResponse(string value, HttpStatusCode statusCode)
+{
+    private readonly JsonHttpResult<OpenAIErrorPayload> innerResult = TypedResults.Json(
+        new OpenAIErrorPayload((int)statusCode, value),
+        statusCode: (int)statusCode
+    );
+
+    public static OpenAIErrorResponse BadRequest(string message) =>
+        new(message, HttpStatusCode.BadRequest);
+
+    public static OpenAIErrorResponse TooManyRequests(string message) => new(message, HttpStatusCode.TooManyRequests);
+
+    public static OpenAIErrorResponse Unauthorized(string message) =>
+        new(message, HttpStatusCode.Unauthorized);
+
+    public async Task WriteAsync(HttpContext httpContext)
+    {
+        httpContext.Response.StatusCode = (int)innerResult.StatusCode!;
+        await httpContext.Response.WriteAsJsonAsync(new
+        {
+            code = innerResult.StatusCode,
+            message = innerResult.Value
+        });
+    }
+
+    record OpenAIErrorPayload(int Code, string Message);
+}
diff --git a/src/AzureAIProxy/Middleware/ProxyAuthenticationOptions.cs b/src/AzureAIProxy/Middleware/ProxyAuthenticationOptions.cs
@@ -6,4 +6,5 @@ public class ProxyAuthenticationOptions : AuthenticationSchemeOptions
 {
     public const string ApiKeyScheme = "ApiKeyScheme";
     public const string JwtScheme = "JwtScheme";
+    public const string BearerTokenScheme = "BearerTokenScheme";
 }
diff --git a/src/AzureAIProxy/Middleware/RateLimiterHandler.cs b/src/AzureAIProxy/Middleware/RateLimiterHandler.cs
@@ -1,23 +1,20 @@
+using AzureAIProxy.Shared.Database;
+
 namespace AzureAIProxy.Middleware;
 
 public class RateLimiterHandler(RequestDelegate next)
 {
-    private const int RateLimitStatusCode = 429;
     private readonly RequestDelegate _next = next;
 
     public async Task InvokeAsync(HttpContext context)
     {
-        if (context.Items.TryGetValue("RateLimited", out var rateLimit) && rateLimit is true)
+        RequestContext? requestContext = context.Items["RequestContext"] as RequestContext;
+
+        if (requestContext is not null && requestContext.RateLimitExceed)
         {
-            var dailyRequestCap = context.Items["DailyRequestCap"] ?? 0;
-            context.Response.StatusCode = RateLimitStatusCode; // Too Many Requests
-            await context.Response.WriteAsJsonAsync(
-                new
-                {
-                    code = RateLimitStatusCode,
-                    message = $"The event daily request rate of {dailyRequestCap} calls has been exceeded. Requests are disabled until UTC midnight."
-                }
-            );
+            await OpenAIErrorResponse.TooManyRequests(
+                $"The event daily request rate of {requestContext.DailyRequestCap} calls has been exceeded. Requests are disabled until UTC midnight."
+            ).WriteAsync(context);
         }
         else
         {

diff --git a/src/AzureAIProxy/Program.cs b/src/AzureAIProxy/Program.cs
@@ -18,6 +18,10 @@
     .AddScheme<ProxyAuthenticationOptions, JwtAuthenticationHandler>(
         ProxyAuthenticationOptions.JwtScheme,
         _ => { }
+    )
+    .AddScheme<ProxyAuthenticationOptions, BearerTokenAuthenticationHandler>(
+        ProxyAuthenticationOptions.BearerTokenScheme,
+        _ => { }
     );
 
 builder.Services.AddMemoryCache();
@@ -31,6 +35,8 @@
 app.UseAuthentication();
 app.UseAuthorization();
 app.UseMiddleware<RateLimiterHandler>();
+app.UseMiddleware<LoadProperties>();
+app.UseMiddleware<MaxTokensHandler>();
 app.MapProxyRoutes();
 
 app.Run();