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

refactor: moving image listener to own function

Browse Source
This commit is contained in:
John Costa
2025-04-09 17:20:27 +01:00
parent 11c5c8921b
commit a22c56fd2c
3 changed files with 94 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
backend/events.go Normal file
View File

@ -0,0 +1,79 @@
package main
import (
"context"
"database/sql"
"log"
"os"
"screenmark/screenmark/agents"
"screenmark/screenmark/models"
"time"
"github.com/google/uuid"
"github.com/lib/pq"
)
func ListenNewImageEvents(db *sql.DB) {
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)
err := listener.Listen("new_image")
if err != nil {
panic(err)
}
for {
select {
case parameters := <-listener.Notify:
imageId := uuid.MustParse(parameters.Extra)
ctx := context.Background()
go func() {
locationAgent, err := agents.NewLocationEventAgent(locationModel, eventModel, contactModel)
if err != nil {
panic(err)
}
noteAgent, err := agents.NewNoteAgent(noteModel)
if err != nil {
panic(err)
}
image, err := imageModel.GetToProcessWithData(ctx, imageId)
if err != nil {
log.Println("Failed to GetToProcessWithData")
log.Println(err)
return
}
_, err = imageModel.FinishProcessing(ctx, image.ID)
if err != nil {
log.Println("Failed to FinishProcessing")
log.Println(err)
return
}
orchestrator, err := agents.NewOrchestratorAgent(locationAgent, noteAgent, image.Image.ImageName, image.Image.Image)
if err != nil {
panic(err)
}
err = orchestrator.Orchestrate(image.UserID, image.ImageID, image.Image.ImageName, image.Image.Image)
if err != nil {
log.Println(err)
}
}()
}
}
}

87
backend/main.go
View File

@ -2,26 +2,21 @@ package main
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"os"
"path/filepath"
"screenmark/screenmark/.gen/haystack/haystack/model"
"screenmark/screenmark/agents"
"screenmark/screenmark/agents/client"
"screenmark/screenmark/models"
"time"
"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
"github.com/google/uuid"
"github.com/joho/godotenv"
"github.com/lib/pq"
)
type TestAiClient struct {
@ -38,84 +33,20 @@ func main() {
panic(err)
}
mode := os.Getenv("MODE")
log.Printf("Mode: %s\n", mode)
db, err := models.InitDatabase()
if err != nil {
panic(err)
}
imageModel := models.NewImageModel(db)
locationModel := models.NewLocationModel(db)
eventModel := models.NewEventModel(db)
userModel := models.NewUserModel(db)
contactModel := models.NewContactModel(db)
noteModel := models.NewNoteModel(db)
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()
go func() {
err := listener.Listen("new_image")
if err != nil {
panic(err)
}
for {
select {
case parameters := <-listener.Notify:
imageId := uuid.MustParse(parameters.Extra)
ctx := context.Background()
go func() {
locationAgent, err := agents.NewLocationEventAgent(locationModel, eventModel, contactModel)
if err != nil {
panic(err)
}
noteAgent, err := agents.NewNoteAgent(noteModel)
if err != nil {
panic(err)
}
image, err := imageModel.GetToProcessWithData(ctx, imageId)
if err != nil {
log.Println("Failed to GetToProcessWithData")
log.Println(err)
return
}
_, err = imageModel.FinishProcessing(ctx, image.ID)
if err != nil {
log.Println("Failed to FinishProcessing")
log.Println(err)
return
}
orchestrator, err := agents.NewOrchestratorAgent(locationAgent, noteAgent, image.Image.ImageName, image.Image.Image)
if err != nil {
panic(err)
}
err = orchestrator.Orchestrate(image.UserID, image.ImageID, image.Image.ImageName, image.Image.Image)
if err != nil {
fmt.Println(err)
}
}()
}
}
}()
go ListenNewImageEvents(db)
r := chi.NewRouter()
r.Use(middleware.Logger)
r.Use(CorsMiddleware)
r.Use(func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Add("Content-Type", "application/json")
@ -124,19 +55,9 @@ func main() {
})
})
r.Options("/*", func(w http.ResponseWriter, r *http.Request) {
w.Header().Add("Access-Control-Allow-Origin", "*")
w.Header().Add("Access-Control-Allow-Credentials", "*")
w.Header().Add("Access-Control-Allow-Headers", "*")
})
r.Get("/image", func(w http.ResponseWriter, r *http.Request) {
userId := r.Header.Get("userId")
w.Header().Add("Access-Control-Allow-Origin", "*")
w.Header().Add("Access-Control-Allow-Credentials", "*")
w.Header().Add("Access-Control-Allow-Headers", "*")
images, err := userModel.ListWithProperties(r.Context(), uuid.MustParse(userId))
if err != nil {
log.Println(err)
@ -190,10 +111,6 @@ func main() {
r.Get("/image/{id}", func(w http.ResponseWriter, r *http.Request) {
imageId := r.PathValue("id")
w.Header().Add("Access-Control-Allow-Origin", "*")
w.Header().Add("Access-Control-Allow-Credentials", "*")
w.Header().Add("Access-Control-Allow-Headers", "*")
// TODO: really need authorization here!
image, err := imageModel.Get(r.Context(), uuid.MustParse(imageId))
if err != nil {

13
backend/middleware.go Normal file
View File

@ -0,0 +1,13 @@
package main
import "net/http"
func CorsMiddleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Add("Access-Control-Allow-Origin", "*")
w.Header().Add("Access-Control-Allow-Credentials", "*")
w.Header().Add("Access-Control-Allow-Headers", "*")
next.ServeHTTP(w, r)
})
}