diff --git a/core/http/app.go b/core/http/app.go index 88cd3ffa8a18..81031c493b63 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -121,6 +121,9 @@ func API(application *application.Application) (*echo.Echo, error) { } }) + //Job tracking middleware (Fix for #7906) + e.Use(httpMiddleware.JobTracker()) + // Recover middleware if !application.ApplicationConfig().Debug { e.Use(middleware.Recover()) @@ -222,7 +225,7 @@ func API(application *application.Application) (*echo.Echo, error) { routes.RegisterUIRoutes(e, application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig(), application.GalleryService()) } routes.RegisterJINARoutes(e, requestExtractor, application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()) - + routes.RegisterJobRoutes(e) // Note: 404 handling is done via HTTPErrorHandler above, no need for catch-all route // Log startup message diff --git a/core/http/jobs/store.go b/core/http/jobs/store.go new file mode 100644 index 000000000000..c9310e2bc46c --- /dev/null +++ b/core/http/jobs/store.go @@ -0,0 +1,85 @@ +package jobs + +import ( + "sync" + "time" + + "github.com/google/uuid" +) + +//A task running in LocalAI +type Job struct{ + ID string `json:"job_id"` + Type string `json:"type"` + Model string `json:"model"` + StartTime time.Time `json:"start_time"` + Status string `json:"status"` + ClientIP string `json:"client_ip"` +} + +//All jobs +type JobStore struct{ + jobs map[string]*Job + mu sync.RWMutex +} + +var currentStore *JobStore +var once sync.Once + +//Return singleton instance of the JobStore +func GetStore() *JobStore{ + once.Do(func(){ + currentStore=&JobStore{ + jobs: make(map[string]*Job), + } + }) + return currentStore +} + +//Add new job to the store +func (s *JobStore) AddJob(j *Job){ + s.mu.Lock() + defer s.mu.Unlock() + s.jobs[j.ID]=j +} + +//Return specific job +func (s *JobStore) GetJob(id string) *Job{ + s.mu.RLock() + defer s.mu.RUnlock() + return s.jobs[id] +} + +//Return a list of all jobs +func (s *JobStore) GetAllJobs() []*Job{ + s.mu.RLock() + defer s.mu.RUnlock() + + var list []*Job + for _,job:=range s.jobs{ + list=append(list,job) + } + return list +} + +//Update status of a job +func (s *JobStore) UpdateStatus(id string,status string){ + s.mu.Lock() + defer s.mu.Unlock() + + if job,exists:=s.jobs[id]; exists{ + job.Status=status + } +} + +//Helper function to generate a new job +func CreateJob(jobType,model,clientIP string) *Job{ + return &Job{ + ID: uuid.New().String(), + Type: jobType, + Model: model, + StartTime: time.Now(), + Status: "executing", + ClientIP: clientIP, + } +} \ No newline at end of file diff --git a/core/http/middleware/job_middleware.go b/core/http/middleware/job_middleware.go new file mode 100644 index 000000000000..a7002b0e40b7 --- /dev/null +++ b/core/http/middleware/job_middleware.go @@ -0,0 +1,58 @@ +package middleware + +import ( + "bytes" + "encoding/json" + "io" + "strings" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/http/jobs" +) + +//Find model name from JSON body +type simpleModelRequest struct{ + Model string `json:"model"` +} + +//Track jobs by intercepting requests +func JobTracker() echo.MiddlewareFunc{ + return func(next echo.HandlerFunc) echo.HandlerFunc{ + return func(c echo.Context) error{ + req:=c.Request() + urlPath:=req.URL.Path + + //Filter + if !strings.HasPrefix(urlPath,"/v1/") && !strings.HasPrefix(urlPath,"/api/"){ + return next(c) + } + if strings.Contains(urlPath,"/jobs"){ + return next(c) + } + + modelName:="unknown" + + if req.Method=="POST"{ + bodyBytes,_:=io.ReadAll(req.Body) + req.Body=io.NopCloser(bytes.NewBuffer(bodyBytes)) + + var tmp simpleModelRequest + if err:=json.Unmarshal(bodyBytes,&tmp);err==nil && tmp.Model!=""{ + modelName=tmp.Model + } + } + job:=jobs.CreateJob(urlPath,modelName,c.RealIP()) + store:=jobs.GetStore() + store.AddJob(job) + + err:=next(c) + + if err!=nil{ + store.UpdateStatus(job.ID,"error") + }else{ + store.UpdateStatus(job.ID,"finished") + } + return err + } + } +} \ No newline at end of file diff --git a/core/http/routes/jobs.go b/core/http/routes/jobs.go new file mode 100644 index 000000000000..a3a2621e577a --- /dev/null +++ b/core/http/routes/jobs.go @@ -0,0 +1,18 @@ +package routes + +import ( + "net/http" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/http/jobs" +) + +func RegisterJobRoutes(e *echo.Echo){ + e.GET("/backends/jobs",listJobs) +} + +func listJobs(c echo.Context) error{ + store:=jobs.GetStore() + runningJobs:=store.GetAllJobs() + return c.JSON(http.StatusOK,runningJobs) +} \ No newline at end of file