From 1fb9616aa792e37ab9c91743f4fc81de4ebe9325 Mon Sep 17 00:00:00 2001 From: John Costa Date: Sun, 21 Sep 2025 22:07:56 +0100 Subject: [PATCH] WIP: image processing is back and working --- backend/.gen/haystack/haystack/model/image.go | 2 +- backend/agents/create_list_agent.go | 2 +- backend/agents/description_agent.go | 9 ++----- backend/agents/list_agent.go | 2 +- backend/images/handler.go | 2 +- backend/models/image.go | 20 ++++++++++++++++ backend/models/lists.go | 12 +++++----- backend/processor/image.go | 24 +++++++++++-------- backend/router.go | 4 +++- backend/schema.sql | 2 +- backend/stacks/handler.go | 6 +++-- 11 files changed, 54 insertions(+), 31 deletions(-) diff --git a/backend/.gen/haystack/haystack/model/image.go b/backend/.gen/haystack/haystack/model/image.go index b0df734..2493ab5 100644 --- a/backend/.gen/haystack/haystack/model/image.go +++ b/backend/.gen/haystack/haystack/model/image.go @@ -14,7 +14,7 @@ import ( type Image struct { ID uuid.UUID `sql:"primary_key"` - UserID *uuid.UUID + UserID uuid.UUID ImageName string Description string Status Progress diff --git a/backend/agents/create_list_agent.go b/backend/agents/create_list_agent.go index 78011e1..2d43548 100644 --- a/backend/agents/create_list_agent.go +++ b/backend/agents/create_list_agent.go @@ -120,7 +120,7 @@ func (agent *CreateListAgent) CreateList(log *log.Logger, userID uuid.UUID, user }) } - _, err = agent.listModel.Save(ctx, userID, createListArgs.Title, createListArgs.Description) + _, err = agent.listModel.Save(ctx, userID, createListArgs.Title, createListArgs.Description, model.Progress_Complete) if err != nil { return fmt.Errorf("creating list agent, saving list: %w", err) } diff --git a/backend/agents/description_agent.go b/backend/agents/description_agent.go index 3587483..446a17b 100644 --- a/backend/agents/description_agent.go +++ b/backend/agents/description_agent.go @@ -3,7 +3,6 @@ package agents import ( "context" "fmt" - "screenmark/screenmark/.gen/haystack/haystack/model" "screenmark/screenmark/agents/client" "screenmark/screenmark/models" @@ -50,13 +49,9 @@ func (agent DescriptionAgent) Describe(log *log.Logger, imageID uuid.UUID, image ctx := context.Background() - markdown := resp.Choices[0].Message.Content - - _, err = agent.imageModel.Update(ctx, model.Image{ - ID: imageID, - Description: markdown, - }) + description := resp.Choices[0].Message.Content + err = agent.imageModel.UpdateDescription(ctx, imageID, description) if err != nil { return err } diff --git a/backend/agents/list_agent.go b/backend/agents/list_agent.go index 215ffc5..a31ce30 100644 --- a/backend/agents/list_agent.go +++ b/backend/agents/list_agent.go @@ -206,7 +206,7 @@ func NewListAgent(log *log.Logger, listModel models.StackModel, limitsMethods li } ctx := context.Background() - savedList, err := listModel.Save(ctx, info.UserId, args.Name, args.Desription) + savedList, err := listModel.Save(ctx, info.UserId, args.Name, args.Desription, model.Progress_Complete) if err != nil { log.Error("saving list", "err", err) return "", err diff --git a/backend/images/handler.go b/backend/images/handler.go index 7c6c866..db1a2fe 100644 --- a/backend/images/handler.go +++ b/backend/images/handler.go @@ -60,7 +60,7 @@ func (h *ImageHandler) serveImage(w http.ResponseWriter, r *http.Request) { } // Do not leak that this ID exists. - if !exists || *image.UserID != userID { + if !exists || image.UserID != userID { w.WriteHeader(http.StatusNotFound) return } diff --git a/backend/models/image.go b/backend/models/image.go index 24116b8..302d9fe 100644 --- a/backend/models/image.go +++ b/backend/models/image.go @@ -37,6 +37,26 @@ func (m ImageModel) Get(ctx context.Context, imageID uuid.UUID) (model.Image, bo return image, err != qrm.ErrNoRows, err } +func (m ImageModel) UpdateDescription(ctx context.Context, imageID uuid.UUID, description string) error { + updateImageDescriptionStmt := Image.UPDATE(Image.Description). + SET(Image.Description.SET(String(description))). + WHERE(Image.ID.EQ(UUID(imageID))) + + _, err := updateImageDescriptionStmt.ExecContext(ctx, m.dbPool) + + return err +} + +func (m ImageModel) UpdateProcess(ctx context.Context, imageID uuid.UUID, process model.Progress) error { + updateImageDescriptionStmt := Image.UPDATE(Image.Status). + SET(process). + WHERE(Image.ID.EQ(UUID(imageID))) + + _, err := updateImageDescriptionStmt.ExecContext(ctx, m.dbPool) + + return err +} + func (m ImageModel) Update(ctx context.Context, image model.Image) (model.Image, error) { updateImageStmt := Image.UPDATE(Image.MutableColumns.Except(Image.Image)). MODEL(image). diff --git a/backend/models/lists.go b/backend/models/lists.go index 78dd01c..d8931a5 100644 --- a/backend/models/lists.go +++ b/backend/models/lists.go @@ -83,11 +83,11 @@ func (m StackModel) Get(ctx context.Context, listID uuid.UUID) (model.Stacks, er // INSERT methods // ======================================== -func (m StackModel) Save(ctx context.Context, userID uuid.UUID, name string, description string) (model.Stacks, error) { +func (m StackModel) Save(ctx context.Context, userID uuid.UUID, name string, description string, status model.Progress) (model.Stacks, error) { saveListStmt := Stacks. - INSERT(Stacks.UserID, Stacks.Name, Stacks.Description). - VALUES(userID, name, description). - RETURNING(Stacks.ID, Stacks.UserID, Stacks.Name, Stacks.Description, Stacks.CreatedAt) + INSERT(Stacks.UserID, Stacks.Name, Stacks.Description, Stacks.Status). + VALUES(userID, name, description, status). + RETURNING(Stacks.ID, Stacks.UserID, Stacks.Name, Stacks.Description, Stacks.Status, Stacks.CreatedAt) list := model.Stacks{} err := saveListStmt.QueryContext(ctx, m.dbPool, &list) @@ -96,7 +96,7 @@ func (m StackModel) Save(ctx context.Context, userID uuid.UUID, name string, des } func (m StackModel) SaveItems(ctx context.Context, items []model.SchemaItems) error { - saveItemsStmt := SchemaItems.INSERT(SchemaItems.AllColumns).MODELS(items) + saveItemsStmt := SchemaItems.INSERT(SchemaItems.MutableColumns).MODELS(items) _, err := saveItemsStmt.ExecContext(ctx, m.dbPool) @@ -137,6 +137,6 @@ func (m StackModel) Delete(ctx context.Context, listID uuid.UUID, userID uuid.UU return err } -func NewListModel(db *sql.DB) StackModel { +func NewStackModel(db *sql.DB) StackModel { return StackModel{dbPool: db} } diff --git a/backend/processor/image.go b/backend/processor/image.go index ee41b52..ab5a302 100644 --- a/backend/processor/image.go +++ b/backend/processor/image.go @@ -5,6 +5,7 @@ import ( "screenmark/screenmark/.gen/haystack/haystack/model" "screenmark/screenmark/agents" "screenmark/screenmark/agents/client" + "screenmark/screenmark/limits" "screenmark/screenmark/models" "sync" @@ -18,7 +19,7 @@ type ImageProcessor struct { logger *log.Logger descriptionAgent agents.DescriptionAgent - listAgent client.AgentClient + stackAgent client.AgentClient // TODO: add the notifier here @@ -26,12 +27,7 @@ type ImageProcessor struct { } func (p *ImageProcessor) setImageToProcess(ctx context.Context, image model.Image) { - updatedImage := model.Image{ - ID: image.ID, - Status: model.Progress_InProgress, - } - - _, err := p.imageModel.Update(ctx, updatedImage) + err := p.imageModel.UpdateProcess(ctx, image.ID, model.Progress_InProgress) if err != nil { // TODO: what can we actually do here for the errors? // We can't stop the work for the others @@ -56,7 +52,7 @@ func (p *ImageProcessor) describe(ctx context.Context, image model.Image) { } func (p *ImageProcessor) extractInfo(ctx context.Context, image model.Image) { - err := p.listAgent.RunAgent(*image.UserID, image.ID, image.ImageName, image.Image) + err := p.stackAgent.RunAgent(image.UserID, image.ID, image.ImageName, image.Image) if err != nil { // Again, wtf do we do? // Although i think the agent actually returns an error when it's finished @@ -88,8 +84,16 @@ func (p *ImageProcessor) processImage(image model.Image) { wg.Wait() } -func NewImageProcessor(logger *log.Logger, imageModel models.ImageModel) ImageProcessor { - imageProcessor := ImageProcessor{imageModel: imageModel, logger: logger} +func NewImageProcessor(logger *log.Logger, imageModel models.ImageModel, listModel models.StackModel, limitsManager limits.LimitsManagerMethods) ImageProcessor { + descriptionAgent := agents.NewDescriptionAgent(logger, imageModel) + stackAgent := agents.NewListAgent(logger, listModel, limitsManager) + + imageProcessor := ImageProcessor{ + imageModel: imageModel, + logger: logger, + descriptionAgent: descriptionAgent, + stackAgent: stackAgent, + } imageProcessor.Processor = NewProcessor(int(IMAGE_PROCESS_AT_A_TIME), imageProcessor.processImage) diff --git a/backend/router.go b/backend/router.go index 66752d5..1661b93 100644 --- a/backend/router.go +++ b/backend/router.go @@ -29,8 +29,10 @@ func setupRouter(db *sql.DB, jwtManager *ourmiddleware.JwtManager) chi.Router { limitsManager := limits.CreateLimitsManager(db) imageModel := models.NewImageModel(db) + stackModel := models.NewStackModel(db) + imageProcessorLogger := createLogger("Image Processor", os.Stdout) - imageProcessor := processor.NewImageProcessor(imageProcessorLogger, imageModel) + imageProcessor := processor.NewImageProcessor(imageProcessorLogger, imageModel, stackModel, limitsManager) go imageProcessor.Processor.Work() stackHandler := stacks.CreateStackHandler(db, limitsManager, jwtManager) diff --git a/backend/schema.sql b/backend/schema.sql index 5c972e9..587f2ba 100644 --- a/backend/schema.sql +++ b/backend/schema.sql @@ -17,7 +17,7 @@ CREATE TABLE haystack.users ( CREATE TABLE haystack.image ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), - user_id UUID REFERENCES haystack.users (id), + user_id UUID NOT NULL REFERENCES haystack.users (id), image_name TEXT NOT NULL, description TEXT NOT NULL, diff --git a/backend/stacks/handler.go b/backend/stacks/handler.go index 0f18dbf..9adde25 100644 --- a/backend/stacks/handler.go +++ b/backend/stacks/handler.go @@ -4,6 +4,7 @@ import ( "database/sql" "net/http" "os" + "screenmark/screenmark/.gen/haystack/haystack/model" "screenmark/screenmark/limits" "screenmark/screenmark/middleware" "screenmark/screenmark/models" @@ -155,7 +156,8 @@ func (h *StackHandler) createStack(body CreateStackBody, w http.ResponseWriter, return } - _, err = h.stackModel.Save(ctx, userID, body.Title, body.Description) + // TODO: Add the stack processor here + _, err = h.stackModel.Save(ctx, userID, body.Title, body.Description, model.Progress_NotStarted) if err != nil { h.logger.Warn("could not save stack", "err", err) w.WriteHeader(http.StatusInternalServerError) @@ -183,7 +185,7 @@ func (h *StackHandler) CreateRoutes(r chi.Router) { } func CreateStackHandler(db *sql.DB, limitsManager limits.LimitsManagerMethods, jwtManager *middleware.JwtManager) StackHandler { - stackModel := models.NewListModel(db) + stackModel := models.NewStackModel(db) imageModel := models.NewImageModel(db) logger := log.New(os.Stdout).WithPrefix("Stacks")