143 lines
3.8 KiB
Go
143 lines
3.8 KiB
Go
package processor
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"screenmark/screenmark/.gen/haystack/haystack/model"
|
|
"screenmark/screenmark/agents"
|
|
"screenmark/screenmark/agents/client"
|
|
"screenmark/screenmark/limits"
|
|
"screenmark/screenmark/models"
|
|
"screenmark/screenmark/notifications"
|
|
"sync"
|
|
|
|
"github.com/charmbracelet/log"
|
|
)
|
|
|
|
const IMAGE_PROCESS_AT_A_TIME = 10
|
|
|
|
type ImageProcessor struct {
|
|
imageModel models.ImageModel
|
|
logger *log.Logger
|
|
|
|
descriptionAgent agents.DescriptionAgent
|
|
stackAgent client.AgentClient
|
|
|
|
Processor *Processor[model.Image]
|
|
notifier *notifications.Notifier[notifications.Notification]
|
|
}
|
|
|
|
func (p *ImageProcessor) setImageToProcess(ctx context.Context, image model.Image) {
|
|
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
|
|
|
|
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.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
|
|
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)
|
|
|
|
imageNotification := notifications.GetImageNotification(notifications.ImageNotification{
|
|
Type: notifications.IMAGE_TYPE,
|
|
ImageID: image.ID,
|
|
ImageName: image.ImageName,
|
|
Status: string(model.Progress_InProgress),
|
|
})
|
|
|
|
err := p.notifier.SendAndCreate(image.UserID.String(), imageNotification)
|
|
if err != nil {
|
|
p.logger.Error("sending in progress notification", "err", err)
|
|
return
|
|
}
|
|
|
|
go func() {
|
|
p.describe(ctx, image)
|
|
wg.Done()
|
|
}()
|
|
|
|
go func() {
|
|
p.extractInfo(ctx, image)
|
|
wg.Done()
|
|
}()
|
|
|
|
wg.Wait()
|
|
|
|
// TODO: there is some repeated code here. The ergonomicts of the notifications,
|
|
// isn't the best.
|
|
imageNotification = notifications.GetImageNotification(notifications.ImageNotification{
|
|
Type: notifications.IMAGE_TYPE,
|
|
ImageID: image.ID,
|
|
ImageName: image.ImageName,
|
|
Status: string(model.Progress_Complete),
|
|
})
|
|
|
|
err = p.notifier.SendAndCreate(image.UserID.String(), imageNotification)
|
|
if err != nil {
|
|
p.logger.Error("sending done notification", "err", err)
|
|
return
|
|
}
|
|
}
|
|
|
|
func NewImageProcessor(
|
|
logger *log.Logger,
|
|
imageModel models.ImageModel,
|
|
listModel models.StackModel,
|
|
limitsManager limits.LimitsManagerMethods,
|
|
notifier *notifications.Notifier[notifications.Notification],
|
|
) (ImageProcessor, error) {
|
|
if notifier == nil {
|
|
return ImageProcessor{}, fmt.Errorf("notifier is nil")
|
|
}
|
|
|
|
descriptionAgent := agents.NewDescriptionAgent(logger, imageModel)
|
|
stackAgent := agents.NewStackAgent(logger, listModel, limitsManager)
|
|
|
|
imageProcessor := ImageProcessor{
|
|
imageModel: imageModel,
|
|
logger: logger,
|
|
descriptionAgent: descriptionAgent,
|
|
stackAgent: stackAgent,
|
|
|
|
notifier: notifier,
|
|
}
|
|
|
|
imageProcessor.Processor = NewProcessor(int(IMAGE_PROCESS_AT_A_TIME), imageProcessor.processImage)
|
|
|
|
return imageProcessor, nil
|
|
}
|