use middleware for Machine-Tag only if tag is specified

Signed-off-by: mintyleaf <mintyleafdev@gmail.com>
This commit is contained in:
mintyleaf 2025-01-17 11:30:06 +04:00
parent 625029cc96
commit 4f9ed3af26
15 changed files with 9 additions and 37 deletions

View file

@ -70,8 +70,8 @@ type RunCMD struct {
WatchdogBusyTimeout string `env:"LOCALAI_WATCHDOG_BUSY_TIMEOUT,WATCHDOG_BUSY_TIMEOUT" default:"5m" help:"Threshold beyond which a busy backend should be stopped" group:"backends"`
Federated bool `env:"LOCALAI_FEDERATED,FEDERATED" help:"Enable federated instance" group:"federated"`
DisableGalleryEndpoint bool `env:"LOCALAI_DISABLE_GALLERY_ENDPOINT,DISABLE_GALLERY_ENDPOINT" help:"Disable the gallery endpoints" group:"api"`
MachineTag string `env:"LOCALAI_MACHINE_TAG" help:"Add Machine-Tag header to each response which is useful to track the machine in the P2P network" group:"api"`
LoadToMemory []string `env:"LOCALAI_LOAD_TO_MEMORY,LOAD_TO_MEMORY" help:"A list of models to load into memory at startup" group:"models"`
MachineTag string `env:"LOCALAI_MACHINE_TAG" help:"TODO: write a help string"`
}
func (r *RunCMD) Run(ctx *cliContext.Context) error {

View file

@ -4,7 +4,6 @@ import (
"context"
"embed"
"encoding/json"
"os"
"regexp"
"time"
@ -99,10 +98,6 @@ func WithModelPath(path string) AppOption {
func WithMachineTag(tag string) AppOption {
return func(o *ApplicationConfig) {
if tag == "" {
hostname, _ := os.Hostname()
tag = hostname
}
o.MachineTag = tag
}
}

View file

@ -89,6 +89,14 @@ func API(application *application.Application) (*fiber.App, error) {
router.Use(middleware.StripPathPrefix())
if application.ApplicationConfig().MachineTag != "" {
router.Use(func(c *fiber.Ctx) error {
c.Response().Header.Set("Machine-Tag", application.ApplicationConfig().MachineTag)
return c.Next()
})
}
router.Hooks().OnListen(func(listenData fiber.ListenData) error {
scheme := "http"
if listenData.TLS {

View file

@ -19,8 +19,6 @@ import (
// @Router /v1/rerank [post]
func JINARerankEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
req := new(schema.JINARerankRequest)
if err := c.BodyParser(req); err != nil {
return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{

View file

@ -24,8 +24,6 @@ import (
// @Router /tts [post]
func TTSEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
input := new(schema.TTSRequest)
// Get input data from the request body

View file

@ -19,8 +19,6 @@ import (
// @Router /vad [post]
func VADEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
input := new(schema.VADRequest)
// Get input data from the request body

View file

@ -76,7 +76,6 @@ type AssistantRequest struct {
// @Router /v1/assistants [post]
func CreateAssistantEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
request := new(AssistantRequest)
if err := c.BodyParser(request); err != nil {
log.Warn().AnErr("Unable to parse AssistantRequest", err)
@ -138,7 +137,6 @@ func generateRandomID() int64 {
// @Router /v1/assistants [get]
func ListAssistantsEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
// Because we're altering the existing assistants list we should just duplicate it for now.
returnAssistants := Assistants
// Parse query parameters
@ -248,7 +246,6 @@ func modelExists(cl *config.BackendConfigLoader, ml *model.ModelLoader, modelNam
// @Router /v1/assistants/{assistant_id} [delete]
func DeleteAssistantEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
assistantID := c.Params("assistant_id")
if assistantID == "" {
return c.Status(fiber.StatusBadRequest).SendString("parameter assistant_id is required")
@ -281,7 +278,6 @@ func DeleteAssistantEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoad
// @Router /v1/assistants/{assistant_id} [get]
func GetAssistantEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
assistantID := c.Params("assistant_id")
if assistantID == "" {
return c.Status(fiber.StatusBadRequest).SendString("parameter assistant_id is required")
@ -311,7 +307,6 @@ var (
func CreateAssistantFileEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
request := new(schema.AssistantFileRequest)
if err := c.BodyParser(request); err != nil {
return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{"error": "Cannot parse JSON"})
@ -358,7 +353,6 @@ func ListAssistantFilesEndpoint(cl *config.BackendConfigLoader, ml *model.ModelL
}
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
assistantID := c.Params("assistant_id")
if assistantID == "" {
return c.Status(fiber.StatusBadRequest).SendString("parameter assistant_id is required")
@ -416,7 +410,6 @@ func ListAssistantFilesEndpoint(cl *config.BackendConfigLoader, ml *model.ModelL
func ModifyAssistantEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
request := new(AssistantRequest)
if err := c.BodyParser(request); err != nil {
log.Warn().AnErr("Unable to parse AssistantRequest", err)
@ -456,7 +449,6 @@ func ModifyAssistantEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoad
func DeleteAssistantFileEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
assistantID := c.Params("assistant_id")
fileId := c.Params("file_id")
if assistantID == "" {
@ -511,7 +503,6 @@ func DeleteAssistantFileEndpoint(cl *config.BackendConfigLoader, ml *model.Model
func GetAssistantFileEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
assistantID := c.Params("assistant_id")
fileId := c.Params("file_id")
if assistantID == "" {

View file

@ -171,8 +171,6 @@ func ChatEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, evaluat
}
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", startupOptions.MachineTag)
textContentToReturn = ""
id = uuid.New().String()
created = int(time.Now().Unix())
@ -327,7 +325,6 @@ func ChatEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, evaluat
c.Context().SetContentType("text/event-stream")
//c.Response().Header.SetContentType(fiber.MIMETextHTMLCharsetUTF8)
// c.Set("Content-Type", "text/event-stream")
c.Set("LocalAI-Machine-Tag", startupOptions.MachineTag)
c.Set("Cache-Control", "no-cache")
c.Set("Connection", "keep-alive")
c.Set("Transfer-Encoding", "chunked")

View file

@ -63,7 +63,6 @@ func CompletionEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, e
}
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
// Add Correlation
c.Set("X-Correlation-ID", id)

View file

@ -25,8 +25,6 @@ import (
func EditEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
// Opt-in extra usage flag
extraUsage := c.Get("LocalAI-Extra-Usage", "") != ""

View file

@ -23,7 +23,6 @@ import (
// @Router /v1/embeddings [post]
func EmbeddingsEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
model, input, err := readRequest(c, cl, ml, appConfig, true)
if err != nil {
return fmt.Errorf("failed reading parameters from request:%w", err)

View file

@ -23,7 +23,6 @@ const UploadedFilesFile = "uploadedFiles.json"
// UploadFilesEndpoint https://platform.openai.com/docs/api-reference/files/create
func UploadFilesEndpoint(cm *config.BackendConfigLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
file, err := c.FormFile("file")
if err != nil {
return err
@ -83,7 +82,6 @@ func getNextFileId() int64 {
func ListFilesEndpoint(cm *config.BackendConfigLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
var listFiles schema.ListFiles
purpose := c.Query("purpose")
@ -122,7 +120,6 @@ func getFileFromRequest(c *fiber.Ctx) (*schema.File, error) {
// @Router /v1/files/{file_id} [get]
func GetFilesEndpoint(cm *config.BackendConfigLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
file, err := getFileFromRequest(c)
if err != nil {
return c.Status(fiber.StatusInternalServerError).SendString(bluemonday.StrictPolicy().Sanitize(err.Error()))
@ -145,7 +142,6 @@ type DeleteStatus struct {
func DeleteFilesEndpoint(cm *config.BackendConfigLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
file, err := getFileFromRequest(c)
if err != nil {
return c.Status(fiber.StatusInternalServerError).SendString(bluemonday.StrictPolicy().Sanitize(err.Error()))
@ -183,7 +179,6 @@ func DeleteFilesEndpoint(cm *config.BackendConfigLoader, appConfig *config.Appli
// GetFilesContentsEndpoint
func GetFilesContentsEndpoint(cm *config.BackendConfigLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
file, err := getFileFromRequest(c)
if err != nil {
return c.Status(fiber.StatusInternalServerError).SendString(bluemonday.StrictPolicy().Sanitize(err.Error()))

View file

@ -66,7 +66,6 @@ func downloadFile(url string) (string, error) {
// @Router /v1/images/generations [post]
func ImageEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
m, input, err := readRequest(c, cl, ml, appConfig, false)
if err != nil {
return fmt.Errorf("failed reading parameters from request:%w", err)

View file

@ -14,8 +14,6 @@ import (
// @Router /v1/models [get]
func ListModelsEndpoint(bcl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(ctx *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
// If blank, no filter is applied.
filter := c.Query("filter")

View file

@ -25,7 +25,6 @@ import (
// @Router /v1/audio/transcriptions [post]
func TranscriptEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
return func(c *fiber.Ctx) error {
c.Set("LocalAI-Machine-Tag", appConfig.MachineTag)
m, input, err := readRequest(c, cl, ml, appConfig, false)
if err != nil {
return fmt.Errorf("failed reading parameters from request:%w", err)