98 lines
2.4 KiB
Go
98 lines
2.4 KiB
Go
package processor
|
|
|
|
import (
|
|
"context"
|
|
"screenmark/screenmark/.gen/haystack/haystack/model"
|
|
"screenmark/screenmark/agents"
|
|
"screenmark/screenmark/agents/client"
|
|
"screenmark/screenmark/models"
|
|
"sync"
|
|
|
|
"github.com/charmbracelet/log"
|
|
)
|
|
|
|
const IMAGE_PROCESS_AT_A_TIME = 10
|
|
|
|
type ImageProcessor struct {
|
|
imageModel models.ImageModel
|
|
logger *log.Logger
|
|
|
|
descriptionAgent agents.DescriptionAgent
|
|
listAgent client.AgentClient
|
|
|
|
// TODO: add the notifier here
|
|
|
|
Processor *Processor[model.Image]
|
|
}
|
|
|
|
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)
|
|
if err != nil {
|
|
// TODO: what can we actually do here for the errors?
|
|
// We can't stop the work for the others
|
|
|
|
p.logger.Error("failed to update image", "err", err)
|
|
|
|
// TODO: we can use context here to actually pass some information through
|
|
return
|
|
}
|
|
}
|
|
|
|
func (p *ImageProcessor) describe(ctx context.Context, image model.Image) {
|
|
descriptionSubLogger := p.logger.With("describe image", image.ID)
|
|
|
|
err := p.descriptionAgent.Describe(descriptionSubLogger, 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
|
|
p.logger.Error("failed to describe image", "err", err)
|
|
return
|
|
}
|
|
}
|
|
|
|
func (p *ImageProcessor) extractInfo(ctx context.Context, image model.Image) {
|
|
err := p.listAgent.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
|
|
p.logger.Error("failed to process image", "err", err)
|
|
return
|
|
}
|
|
}
|
|
|
|
func (p *ImageProcessor) processImage(image model.Image) {
|
|
p.logger.Info("Processing image", "ID", image.ID)
|
|
|
|
ctx := context.Background()
|
|
|
|
p.setImageToProcess(ctx, image)
|
|
|
|
var wg sync.WaitGroup
|
|
wg.Add(2)
|
|
|
|
go func() {
|
|
p.describe(ctx, image)
|
|
wg.Done()
|
|
}()
|
|
|
|
go func() {
|
|
p.extractInfo(ctx, image)
|
|
wg.Done()
|
|
}()
|
|
|
|
wg.Wait()
|
|
}
|
|
|
|
func NewImageProcessor(logger *log.Logger, imageModel models.ImageModel) ImageProcessor {
|
|
imageProcessor := ImageProcessor{imageModel: imageModel, logger: logger}
|
|
|
|
imageProcessor.Processor = NewProcessor(int(IMAGE_PROCESS_AT_A_TIME), imageProcessor.processImage)
|
|
|
|
return imageProcessor
|
|
}
|