WIP

Signed-off-by: Ettore Di Giacinto <[email protected]>

Author: mudler

core/http/app.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 
 	"github.com/dave-gray101/v2keyauth"
+	"github.com/gofiber/websocket/v2"
 	"github.com/mudler/LocalAI/pkg/utils"
 
 	"github.com/mudler/LocalAI/core/http/endpoints/localai"
@@ -121,7 +122,7 @@ func App(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *confi
 		})
 	}
 
- // Health Checks should always be exempt from auth, so register these first
+	// Health Checks should always be exempt from auth, so register these first
 	routes.HealthRoutes(app)
 
 	kaConfig, err := middleware.GetKeyAuthConfig(appConfig)
@@ -178,6 +179,16 @@ func App(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *confi
 		Browse:     true,
 	}))
 
+	app.Use("/ws", func(c *fiber.Ctx) error {
+		// IsWebSocketUpgrade returns true if the client
+		// requested upgrade to the WebSocket protocol.
+		if websocket.IsWebSocketUpgrade(c) {
+			c.Locals("allowed", true)
+			return c.Next()
+		}
+		return fiber.ErrUpgradeRequired
+	})
+
 	// Define a custom 404 handler
 	// Note: keep this at the bottom!
 	app.Use(notFoundHandler)

core/http/ctx/fiber.go

@@ -19,9 +19,11 @@ func ModelFromContext(ctx *fiber.Ctx, cl *config.BackendConfigLoader, loader *mo
 	if ctx.Params("model") != "" {
 		modelInput = ctx.Params("model")
 	}
+
 	if ctx.Query("model") != "" {
 		modelInput = ctx.Query("model")
 	}
+
 	// Set model from bearer token, if available
 	bearer := strings.TrimLeft(ctx.Get("authorization"), "Bear ") // Reduced duplicate characters of Bearer
 	bearerExists := bearer != "" && loader.ExistsInModelPath(bearer)