Haystack/backend/events.go

package main

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"screenmark/screenmark/agents"
	"screenmark/screenmark/models"
	"time"

	"github.com/charmbracelet/log"
	"github.com/google/uuid"
	"github.com/lib/pq"
)

func createLogger(prefix string) *log.Logger {
	return log.NewWithOptions(os.Stdout, log.Options{
		ReportTimestamp: true,
		TimeFormat:      time.Kitchen,
		Prefix:          prefix,
	})
}

func ListenNewImageEvents(db *sql.DB, eventManager *EventManager) {
	listener := pq.NewListener(os.Getenv("DB_CONNECTION"), time.Second, time.Second, func(event pq.ListenerEventType, err error) {
		if err != nil {
			panic(err)
		}
	})
	defer listener.Close()

	locationModel := models.NewLocationModel(db)
	eventModel := models.NewEventModel(db)
	noteModel := models.NewNoteModel(db)
	imageModel := models.NewImageModel(db)
	contactModel := models.NewContactModel(db)

	databaseEventLog := createLogger("Database Events 🤖")

	err := listener.Listen("new_image")
	if err != nil {
		panic(err)
	}

	for {
		select {
		case parameters := <-listener.Notify:
			imageId := uuid.MustParse(parameters.Extra)
			eventManager.listeners[parameters.Extra] = make(chan string)

			databaseEventLog.Debug("Starting processing image", "ImageID", imageId)

			ctx := context.Background()

			go func() {
				noteAgent, err := agents.NewNoteAgent(createLogger("Notes 📝"), noteModel)
				if err != nil {
					panic(err)
				}

				contactAgent, err := agents.NewContactAgent(createLogger("Contacts 👥"), contactModel)
				if err != nil {
					panic(err)
				}

				locationAgent, err := agents.NewLocationAgent(createLogger("Locations 📍"), locationModel)
				if err != nil {
					panic(err)
				}

				eventAgent, err := agents.NewEventAgent(createLogger("Events 📅"), eventModel, locationAgent)
				if err != nil {
					panic(err)
				}

				image, err := imageModel.GetToProcessWithData(ctx, imageId)
				if err != nil {
					log.Error("Failed to GetToProcessWithData", "error", err)
					return
				}

				if err := imageModel.StartProcessing(ctx, image.ID); err != nil {
					log.Error("Failed to FinishProcessing", "error", err)
					return
				}

				orchestrator, err := agents.NewOrchestratorAgent(createLogger("Orchestrator 🎼"), noteAgent, contactAgent, locationAgent, eventAgent, image.Image.ImageName, image.Image.Image)
				if err != nil {
					panic(err)
				}

				err = orchestrator.RunAgent(agents.OrchestratorPrompt, agents.OrchestratorTools, "noAction", nil, image.UserID, image.ImageID, image.Image.ImageName, image.Image.Image)
				if err != nil {
					log.Error("Orchestrator failed", "error", "err")
					return
				}

				_, err = imageModel.FinishProcessing(ctx, image.ID)
				if err != nil {
					log.Error("Failed to finish processing", "ImageID", imageId)
					return
				}

				databaseEventLog.Debug("Starting processing image", "ImageID", imageId)
			}()
		}
	}
}

type EventManager struct {
	// Maps processing image UUID to a channel
	listeners map[string]chan string
}

func NewEventManager() EventManager {
	return EventManager{
		listeners: make(map[string]chan string),
	}
}

func ListenProcessingImageStatus(db *sql.DB, eventManager *EventManager) {
	listener := pq.NewListener(os.Getenv("DB_CONNECTION"), time.Second, time.Second, func(event pq.ListenerEventType, err error) {
		if err != nil {
			panic(err)
		}
	})
	defer listener.Close()

	if err := listener.Listen("new_processing_image_status"); err != nil {
		panic(err)
	}

	for {
		select {
		case data := <-listener.Notify:
			stringUuid := data.Extra[0:36]
			status := data.Extra[36:]

			fmt.Printf("UUID: %s\n", stringUuid)
			fmt.Printf("Receiving :s\n", data.Extra)

			imageListener, exists := eventManager.listeners[stringUuid]
			if !exists {
				continue
			}

			imageListener <- status

			close(imageListener)
			delete(eventManager.listeners, stringUuid)
		}
	}
}