package middleware import ( "context" "encoding/json" "fmt" "strconv" "strings" "github.com/google/uuid" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services" "github.com/mudler/LocalAI/pkg/functions" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/templates" "github.com/mudler/LocalAI/pkg/utils" "github.com/gofiber/fiber/v2" "github.com/rs/zerolog/log" ) type correlationIDKeyType string // CorrelationIDKey to track request across process boundary const CorrelationIDKey correlationIDKeyType = "correlationID" type RequestExtractor struct { backendConfigLoader *config.BackendConfigLoader modelLoader *model.ModelLoader applicationConfig *config.ApplicationConfig } func NewRequestExtractor(backendConfigLoader *config.BackendConfigLoader, modelLoader *model.ModelLoader, applicationConfig *config.ApplicationConfig) *RequestExtractor { return &RequestExtractor{ backendConfigLoader: backendConfigLoader, modelLoader: modelLoader, applicationConfig: applicationConfig, } } const CONTEXT_LOCALS_KEY_MODEL_NAME = "MODEL_NAME" const CONTEXT_LOCALS_KEY_LOCALAI_REQUEST = "LOCALAI_REQUEST" const CONTEXT_LOCALS_KEY_MODEL_CONFIG = "MODEL_CONFIG" // TODO: Refactor to not return error if unchanged func (re *RequestExtractor) setModelNameFromRequest(ctx *fiber.Ctx) { model, ok := ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_NAME).(string) if ok && model != "" { return } model = ctx.Params("model") if (model == "") && ctx.Query("model") != "" { model = ctx.Query("model") } if model == "" { // Set model from bearer token, if available bearer := strings.TrimLeft(ctx.Get("authorization"), "Bear ") // "Bearer " => "Bear" to please go-staticcheck. It looks dumb but we might as well take free performance on something called for nearly every request. if bearer != "" { exists, err := services.CheckIfModelExists(re.backendConfigLoader, re.modelLoader, bearer, services.ALWAYS_INCLUDE) if err == nil && exists { model = bearer } } } ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_NAME, model) } func (re *RequestExtractor) BuildConstantDefaultModelNameMiddleware(defaultModelName string) fiber.Handler { return func(ctx *fiber.Ctx) error { re.setModelNameFromRequest(ctx) localModelName, ok := ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_NAME).(string) if !ok || localModelName == "" { ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_NAME, defaultModelName) log.Debug().Str("defaultModelName", defaultModelName).Msg("context local model name not found, setting to default") } return ctx.Next() } } func (re *RequestExtractor) BuildFilteredFirstAvailableDefaultModel(filterFn config.BackendConfigFilterFn) fiber.Handler { return func(ctx *fiber.Ctx) error { re.setModelNameFromRequest(ctx) localModelName := ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_NAME).(string) if localModelName != "" { // Don't overwrite existing values return ctx.Next() } modelNames, err := services.ListModels(re.backendConfigLoader, re.modelLoader, filterFn, services.SKIP_IF_CONFIGURED) if err != nil { log.Error().Err(err).Msg("non-fatal error calling ListModels during SetDefaultModelNameToFirstAvailable()") return ctx.Next() } if len(modelNames) == 0 { log.Warn().Msg("SetDefaultModelNameToFirstAvailable used with no matching models installed") // This is non-fatal - making it so was breaking the case of direct installation of raw models // return errors.New("this endpoint requires at least one model to be installed") return ctx.Next() } ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_NAME, modelNames[0]) log.Debug().Str("first model name", modelNames[0]).Msg("context local model name not found, setting to the first model") return ctx.Next() } } // TODO: If context and cancel above belong on all methods, move that part of above into here! // Otherwise, it's in its own method below for now func (re *RequestExtractor) SetModelAndConfig(initializer func() schema.LocalAIRequest) fiber.Handler { return func(ctx *fiber.Ctx) error { input := initializer() if input == nil { return fmt.Errorf("unable to initialize body") } if err := ctx.BodyParser(input); err != nil { return fmt.Errorf("failed parsing request body: %w", err) } // If this request doesn't have an associated model name, fetch it from earlier in the middleware chain if input.ModelName(nil) == "" { localModelName, ok := ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_NAME).(string) if ok && localModelName != "" { log.Debug().Str("context localModelName", localModelName).Msg("overriding empty model name in request body with value found earlier in middleware chain") input.ModelName(&localModelName) } } cfg, err := re.backendConfigLoader.LoadBackendConfigFileByNameDefaultOptions(input.ModelName(nil), re.applicationConfig) if err != nil { log.Err(err) log.Warn().Msgf("Model Configuration File not found for %q", input.ModelName(nil)) } else if cfg.Model == "" && input.ModelName(nil) != "" { log.Debug().Str("input.ModelName", input.ModelName(nil)).Msg("config does not include model, using input") cfg.Model = input.ModelName(nil) } ctx.Locals(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input) ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_CONFIG, cfg) return ctx.Next() } } func (re *RequestExtractor) SetOpenAIRequest(ctx *fiber.Ctx) error { input, ok := ctx.Locals(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest) if !ok || input.Model == "" { return fiber.ErrBadRequest } cfg, ok := ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.BackendConfig) if !ok || cfg == nil { return fiber.ErrBadRequest } // Extract or generate the correlation ID correlationID := ctx.Get("X-Correlation-ID", uuid.New().String()) ctx.Set("X-Correlation-ID", correlationID) c1, cancel := context.WithCancel(re.applicationConfig.Context) // Add the correlation ID to the new context ctxWithCorrelationID := context.WithValue(c1, CorrelationIDKey, correlationID) input.Context = ctxWithCorrelationID input.Cancel = cancel err := mergeOpenAIRequestAndBackendConfig(cfg, input) if err != nil { return err } if cfg.Model == "" { log.Debug().Str("input.Model", input.Model).Msg("replacing empty cfg.Model with input value") cfg.Model = input.Model } ctx.Locals(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input) ctx.Locals(CONTEXT_LOCALS_KEY_MODEL_CONFIG, cfg) return ctx.Next() } func mergeOpenAIRequestAndBackendConfig(config *config.BackendConfig, input *schema.OpenAIRequest) error { if input.Echo { config.Echo = input.Echo } if input.TopK != nil { config.TopK = input.TopK } if input.TopP != nil { config.TopP = input.TopP } if input.Backend != "" { config.Backend = input.Backend } if input.ClipSkip != 0 { config.Diffusers.ClipSkip = input.ClipSkip } if input.ModelBaseName != "" { config.AutoGPTQ.ModelBaseName = input.ModelBaseName } if input.NegativePromptScale != 0 { config.NegativePromptScale = input.NegativePromptScale } if input.UseFastTokenizer { config.UseFastTokenizer = input.UseFastTokenizer } if input.NegativePrompt != "" { config.NegativePrompt = input.NegativePrompt } if input.RopeFreqBase != 0 { config.RopeFreqBase = input.RopeFreqBase } if input.RopeFreqScale != 0 { config.RopeFreqScale = input.RopeFreqScale } if input.Grammar != "" { config.Grammar = input.Grammar } if input.Temperature != nil { config.Temperature = input.Temperature } if input.Maxtokens != nil { config.Maxtokens = input.Maxtokens } if input.ResponseFormat != nil { switch responseFormat := input.ResponseFormat.(type) { case string: config.ResponseFormat = responseFormat case map[string]interface{}: config.ResponseFormatMap = responseFormat } } switch stop := input.Stop.(type) { case string: if stop != "" { config.StopWords = append(config.StopWords, stop) } case []interface{}: for _, pp := range stop { if s, ok := pp.(string); ok { config.StopWords = append(config.StopWords, s) } } } if len(input.Tools) > 0 { for _, tool := range input.Tools { input.Functions = append(input.Functions, tool.Function) } } if input.ToolsChoice != nil { var toolChoice functions.Tool switch content := input.ToolsChoice.(type) { case string: _ = json.Unmarshal([]byte(content), &toolChoice) case map[string]interface{}: dat, _ := json.Marshal(content) _ = json.Unmarshal(dat, &toolChoice) } input.FunctionCall = map[string]interface{}{ "name": toolChoice.Function.Name, } } // Decode each request's message content imgIndex, vidIndex, audioIndex := 0, 0, 0 for i, m := range input.Messages { nrOfImgsInMessage := 0 nrOfVideosInMessage := 0 nrOfAudiosInMessage := 0 switch content := m.Content.(type) { case string: input.Messages[i].StringContent = content case []interface{}: dat, _ := json.Marshal(content) c := []schema.Content{} json.Unmarshal(dat, &c) textContent := "" // we will template this at the end CONTENT: for _, pp := range c { switch pp.Type { case "text": textContent += pp.Text //input.Messages[i].StringContent = pp.Text case "video", "video_url": // Decode content as base64 either if it's an URL or base64 text base64, err := utils.GetContentURIAsBase64(pp.VideoURL.URL) if err != nil { log.Error().Msgf("Failed encoding video: %s", err) continue CONTENT } input.Messages[i].StringVideos = append(input.Messages[i].StringVideos, base64) // TODO: make sure that we only return base64 stuff vidIndex++ nrOfVideosInMessage++ case "audio_url", "audio": // Decode content as base64 either if it's an URL or base64 text base64, err := utils.GetContentURIAsBase64(pp.AudioURL.URL) if err != nil { log.Error().Msgf("Failed encoding image: %s", err) continue CONTENT } input.Messages[i].StringAudios = append(input.Messages[i].StringAudios, base64) // TODO: make sure that we only return base64 stuff audioIndex++ nrOfAudiosInMessage++ case "image_url", "image": // Decode content as base64 either if it's an URL or base64 text base64, err := utils.GetContentURIAsBase64(pp.ImageURL.URL) if err != nil { log.Error().Msgf("Failed encoding image: %s", err) continue CONTENT } input.Messages[i].StringImages = append(input.Messages[i].StringImages, base64) // TODO: make sure that we only return base64 stuff imgIndex++ nrOfImgsInMessage++ } } input.Messages[i].StringContent, _ = templates.TemplateMultiModal(config.TemplateConfig.Multimodal, templates.MultiModalOptions{ TotalImages: imgIndex, TotalVideos: vidIndex, TotalAudios: audioIndex, ImagesInMessage: nrOfImgsInMessage, VideosInMessage: nrOfVideosInMessage, AudiosInMessage: nrOfAudiosInMessage, }, textContent) } } if input.RepeatPenalty != 0 { config.RepeatPenalty = input.RepeatPenalty } if input.FrequencyPenalty != 0 { config.FrequencyPenalty = input.FrequencyPenalty } if input.PresencePenalty != 0 { config.PresencePenalty = input.PresencePenalty } if input.Keep != 0 { config.Keep = input.Keep } if input.Batch != 0 { config.Batch = input.Batch } if input.IgnoreEOS { config.IgnoreEOS = input.IgnoreEOS } if input.Seed != nil { config.Seed = input.Seed } if input.TypicalP != nil { config.TypicalP = input.TypicalP } log.Debug().Str("input.Input", fmt.Sprintf("%+v", input.Input)) switch inputs := input.Input.(type) { case string: if inputs != "" { config.InputStrings = append(config.InputStrings, inputs) } case []interface{}: for _, pp := range inputs { switch i := pp.(type) { case string: config.InputStrings = append(config.InputStrings, i) case []interface{}: tokens := []int{} for _, ii := range i { tokens = append(tokens, int(ii.(float64))) } config.InputToken = append(config.InputToken, tokens) } } } // Can be either a string or an object switch fnc := input.FunctionCall.(type) { case string: if fnc != "" { config.SetFunctionCallString(fnc) } case map[string]interface{}: var name string n, exists := fnc["name"] if exists { nn, e := n.(string) if e { name = nn } } config.SetFunctionCallNameString(name) } switch p := input.Prompt.(type) { case string: config.PromptStrings = append(config.PromptStrings, p) case []interface{}: for _, pp := range p { if s, ok := pp.(string); ok { config.PromptStrings = append(config.PromptStrings, s) } } } // If a quality was defined as number, convert it to step if input.Quality != "" { q, err := strconv.Atoi(input.Quality) if err == nil { config.Step = q } } if config.Validate() { return nil } return fmt.Errorf("unable to validate configuration after merging") }