JohnCosta27/Haystack
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

WIP: image processing is back and working

Browse Source
This commit is contained in:
John Costa
2025-09-21 22:07:56 +01:00
parent 013447fa90
commit 1fb9616aa7
11 changed files with 54 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
backend/.gen/haystack/haystack/model/image.go
View File

@ -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

2
backend/agents/create_list_agent.go
View File

@ -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)
}

9
backend/agents/description_agent.go
View File

@ -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
}

2
backend/agents/list_agent.go
View File

@ -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

2
backend/images/handler.go
View File

@ -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
}

20
backend/models/image.go
View File

@ -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).

12
backend/models/lists.go
View File

@ -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}
}

24
backend/processor/image.go
View File

@ -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)

4
backend/router.go
View File

@ -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)

2
backend/schema.sql
View File

@ -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,

6
backend/stacks/handler.go
View File

@ -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")