diff --git a/backend/agents/contact_agent.go b/backend/agents/contact_agent.go index 7dd45ce..eacbb28 100644 --- a/backend/agents/contact_agent.go +++ b/backend/agents/contact_agent.go @@ -3,11 +3,9 @@ package agents import ( "context" "encoding/json" - "os" "screenmark/screenmark/.gen/haystack/haystack/model" "screenmark/screenmark/agents/client" "screenmark/screenmark/models" - "time" "github.com/charmbracelet/log" "github.com/google/uuid" @@ -81,12 +79,6 @@ const contactTools = ` ] ` -type ContactAgent struct { - client client.AgentClient - - contactModel models.ContactModel -} - type listContactsArguments struct{} type createContactsArguments struct { Name string `json:"name"` @@ -98,23 +90,14 @@ type linkContactArguments struct { ContactID string `json:"contactId"` } -func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) { - agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{ - ReportTimestamp: true, - TimeFormat: time.Kitchen, - Prefix: "Contacts 👥", - })) +func NewContactAgent(log *log.Logger, contactModel models.ContactModel) (client.AgentClient, error) { + agentClient, err := client.CreateAgentClient(log) if err != nil { - return ContactAgent{}, err - } - - agent := ContactAgent{ - client: agentClient, - contactModel: contactModel, + return client.AgentClient{}, err } agentClient.ToolHandler.AddTool("listContacts", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { - return agent.contactModel.List(context.Background(), info.UserId) + return contactModel.List(context.Background(), info.UserId) }) agentClient.ToolHandler.AddTool("createContact", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) { @@ -126,7 +109,7 @@ func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) { ctx := context.Background() - contact, err := agent.contactModel.Save(ctx, info.UserId, model.Contacts{ + contact, err := contactModel.Save(ctx, info.UserId, model.Contacts{ Name: args.Name, PhoneNumber: args.PhoneNumber, Email: args.Email, @@ -136,7 +119,7 @@ func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) { return model.Contacts{}, err } - _, err = agent.contactModel.SaveToImage(ctx, info.ImageId, contact.ID) + _, err = contactModel.SaveToImage(ctx, info.ImageId, contact.ID) if err != nil { return model.Contacts{}, err } @@ -158,7 +141,7 @@ func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) { return "", err } - _, err = agent.contactModel.SaveToImage(ctx, info.ImageId, contactUuid) + _, err = contactModel.SaveToImage(ctx, info.ImageId, contactUuid) if err != nil { return "", err } @@ -166,5 +149,5 @@ func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) { return "Saved", nil }) - return agent, nil + return agentClient, nil } diff --git a/backend/agents/event_agent.go b/backend/agents/event_agent.go index 0e1c6f9..609f1f0 100644 --- a/backend/agents/event_agent.go +++ b/backend/agents/event_agent.go @@ -3,7 +3,6 @@ package agents import ( "context" "encoding/json" - "os" "screenmark/screenmark/.gen/haystack/haystack/model" "screenmark/screenmark/agents/client" "screenmark/screenmark/models" @@ -98,14 +97,6 @@ const eventTools = ` } ]` -type EventAgent struct { - client client.AgentClient - - eventsModel models.EventModel - - locationAgent LocationAgent -} - type listEventArguments struct{} type createEventArguments struct { Name string `json:"name"` @@ -117,25 +108,15 @@ type linkEventArguments struct { EventID string `json:"eventId"` } -func NewEventAgent(eventsModel models.EventModel, locationAgent LocationAgent) (EventAgent, error) { - agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{ - ReportTimestamp: true, - TimeFormat: time.Kitchen, - Prefix: "Events 📍", - })) +func NewEventAgent(log *log.Logger, eventsModel models.EventModel, locationAgent client.AgentClient) (client.AgentClient, error) { + agentClient, err := client.CreateAgentClient(log) if err != nil { - return EventAgent{}, err - } - - agent := EventAgent{ - client: agentClient, - eventsModel: eventsModel, - locationAgent: locationAgent, + return client.AgentClient{}, err } agentClient.ToolHandler.AddTool("listEvents", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { - return agent.eventsModel.List(context.Background(), info.UserId) + return eventsModel.List(context.Background(), info.UserId) }) agentClient.ToolHandler.AddTool("createEvent", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) { @@ -159,7 +140,7 @@ func NewEventAgent(eventsModel models.EventModel, locationAgent LocationAgent) ( return model.Events{}, err } - events, err := agent.eventsModel.Save(ctx, info.UserId, model.Events{ + events, err := eventsModel.Save(ctx, info.UserId, model.Events{ Name: args.Name, StartDateTime: &startTime, EndDateTime: &endTime, @@ -169,7 +150,7 @@ func NewEventAgent(eventsModel models.EventModel, locationAgent LocationAgent) ( return model.Events{}, err } - _, err = agent.eventsModel.SaveToImage(ctx, info.ImageId, events.ID) + _, err = eventsModel.SaveToImage(ctx, info.ImageId, events.ID) if err != nil { return model.Events{}, err } @@ -191,16 +172,16 @@ func NewEventAgent(eventsModel models.EventModel, locationAgent LocationAgent) ( return "", err } - agent.eventsModel.SaveToImage(ctx, info.ImageId, contactUuid) + eventsModel.SaveToImage(ctx, info.ImageId, contactUuid) return "Saved", nil }) agentClient.ToolHandler.AddTool("getEventLocationId", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { query := "Can you get me the ID of the location present in this image?" - locationAgent.client.RunAgent(locationPrompt, locationTools, "finish", &query, info.UserId, info.ImageId, info.ImageName, *info.Image) + locationAgent.RunAgent(locationPrompt, locationTools, "finish", &query, info.UserId, info.ImageId, info.ImageName, *info.Image) - return locationAgent.client.Reply, nil + return locationAgent.Reply, nil }) - return agent, nil + return agentClient, nil } diff --git a/backend/agents/location_agent.go b/backend/agents/location_agent.go index 24f8d41..0c1b4be 100644 --- a/backend/agents/location_agent.go +++ b/backend/agents/location_agent.go @@ -3,11 +3,9 @@ package agents import ( "context" "encoding/json" - "os" "screenmark/screenmark/.gen/haystack/haystack/model" "screenmark/screenmark/agents/client" "screenmark/screenmark/models" - "time" "github.com/charmbracelet/log" "github.com/google/uuid" @@ -97,12 +95,6 @@ const locationTools = ` } ]` -type LocationAgent struct { - client client.AgentClient - - locationModel models.LocationModel -} - type listLocationArguments struct{} type createLocationArguments struct { Name string `json:"name"` @@ -112,24 +104,15 @@ type linkLocationArguments struct { LocationID string `json:"locationId"` } -func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error) { - agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{ - ReportTimestamp: true, - TimeFormat: time.Kitchen, - Prefix: "Locations 📍", - })) +func NewLocationAgent(log *log.Logger, locationModel models.LocationModel) (client.AgentClient, error) { + agentClient, err := client.CreateAgentClient(log) if err != nil { - return LocationAgent{}, err - } - - agent := LocationAgent{ - client: agentClient, - locationModel: locationModel, + return client.AgentClient{}, err } agentClient.ToolHandler.AddTool("listLocations", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { - return agent.locationModel.List(context.Background(), info.UserId) + return locationModel.List(context.Background(), info.UserId) }) agentClient.ToolHandler.AddTool("createLocation", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) { @@ -141,7 +124,7 @@ func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error) ctx := context.Background() - location, err := agent.locationModel.Save(ctx, info.UserId, model.Locations{ + location, err := locationModel.Save(ctx, info.UserId, model.Locations{ Name: args.Name, Address: args.Address, }) @@ -150,7 +133,7 @@ func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error) return model.Locations{}, err } - _, err = agent.locationModel.SaveToImage(ctx, info.ImageId, location.ID) + _, err = locationModel.SaveToImage(ctx, info.ImageId, location.ID) if err != nil { return model.Locations{}, err } @@ -172,14 +155,14 @@ func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error) return "", err } - agent.locationModel.SaveToImage(ctx, info.ImageId, contactUuid) + locationModel.SaveToImage(ctx, info.ImageId, contactUuid) return "Saved", nil }) agentClient.ToolHandler.AddTool("reply", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { - agent.client.Log.Debug(args) + agentClient.Log.Debug(args) return "ok", nil }) - return agent, nil + return agentClient, nil } diff --git a/backend/agents/note_agent.go b/backend/agents/note_agent.go index 9a69f28..81a2d57 100644 --- a/backend/agents/note_agent.go +++ b/backend/agents/note_agent.go @@ -2,11 +2,9 @@ package agents import ( "context" - "os" "screenmark/screenmark/.gen/haystack/haystack/model" "screenmark/screenmark/agents/client" "screenmark/screenmark/models" - "time" "github.com/charmbracelet/log" "github.com/google/uuid" @@ -70,12 +68,8 @@ func (agent NoteAgent) GetNotes(userId uuid.UUID, imageId uuid.UUID, imageName s return nil } -func NewNoteAgent(noteModel models.NoteModel) (NoteAgent, error) { - client, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{ - ReportTimestamp: true, - TimeFormat: time.Kitchen, - Prefix: "Notes 📝", - })) +func NewNoteAgent(log *log.Logger, noteModel models.NoteModel) (NoteAgent, error) { + client, err := client.CreateAgentClient(log) if err != nil { return NoteAgent{}, err } diff --git a/backend/agents/orchestrator.go b/backend/agents/orchestrator.go index 3778411..3b090bb 100644 --- a/backend/agents/orchestrator.go +++ b/backend/agents/orchestrator.go @@ -2,9 +2,7 @@ package agents import ( "errors" - "os" "screenmark/screenmark/agents/client" - "time" "github.com/charmbracelet/log" ) @@ -114,15 +112,11 @@ type Status struct { Ok bool `json:"ok"` } -func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locationAgent LocationAgent, eventAgent EventAgent, imageName string, imageData []byte) (OrchestratorAgent, error) { - agent, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{ - ReportTimestamp: true, - TimeFormat: time.Kitchen, - Prefix: "Orchestrator 🎼", - })) +func NewOrchestratorAgent(log *log.Logger, noteAgent NoteAgent, contactAgent client.AgentClient, locationAgent client.AgentClient, eventAgent client.AgentClient, imageName string, imageData []byte) (client.AgentClient, error) { + agent, err := client.CreateAgentClient(log) if err != nil { - return OrchestratorAgent{}, err + return client.AgentClient{}, err } agent.ToolHandler.AddTool("noteAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { @@ -134,7 +128,7 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati }) agent.ToolHandler.AddTool("contactAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { - go contactAgent.client.RunAgent(contactPrompt, contactTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData) + go contactAgent.RunAgent(contactPrompt, contactTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData) return Status{ Ok: true, @@ -142,7 +136,7 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati }) agent.ToolHandler.AddTool("locationAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { - go locationAgent.client.RunAgent(locationPrompt, locationTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData) + go locationAgent.RunAgent(locationPrompt, locationTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData) return Status{ Ok: true, @@ -150,7 +144,7 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati }) agent.ToolHandler.AddTool("eventAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { - go eventAgent.client.RunAgent(eventPrompt, eventTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData) + go agent.RunAgent(eventPrompt, eventTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData) return Status{ Ok: true, @@ -165,7 +159,5 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati }, errors.New("Finished! Kinda bad return type but...") }) - return OrchestratorAgent{ - Client: agent, - }, nil + return agent, nil } diff --git a/backend/events.go b/backend/events.go index 2c4dba5..b06174d 100644 --- a/backend/events.go +++ b/backend/events.go @@ -4,16 +4,24 @@ import ( "context" "database/sql" "fmt" - "log" "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 { @@ -28,6 +36,8 @@ func ListenNewImageEvents(db *sql.DB, eventManager *EventManager) { imageModel := models.NewImageModel(db) contactModel := models.NewContactModel(db) + databaseEventLog := createLogger("Database Events 🤖") + err := listener.Listen("new_image") if err != nil { panic(err) @@ -39,55 +49,60 @@ func ListenNewImageEvents(db *sql.DB, eventManager *EventManager) { 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(noteModel) + noteAgent, err := agents.NewNoteAgent(createLogger("Notes 📝"), noteModel) if err != nil { panic(err) } - contactAgent, err := agents.NewContactAgent(contactModel) + contactAgent, err := agents.NewContactAgent(createLogger("Contacts 👥"), contactModel) if err != nil { panic(err) } - locationAgent, err := agents.NewLocationAgent(locationModel) + locationAgent, err := agents.NewLocationAgent(createLogger("Locations 📍"), locationModel) if err != nil { panic(err) } - eventAgent, err := agents.NewEventAgent(eventModel, locationAgent) + eventAgent, err := agents.NewEventAgent(createLogger("Events 📅"), eventModel, locationAgent) if err != nil { panic(err) } image, err := imageModel.GetToProcessWithData(ctx, imageId) if err != nil { - log.Println("Failed to GetToProcessWithData") - log.Println(err) + log.Error("Failed to GetToProcessWithData", "error", err) return } if err := imageModel.StartProcessing(ctx, image.ID); err != nil { - log.Println("Failed to FinishProcessing") - log.Println(err) + log.Error("Failed to FinishProcessing", "error", err) return } - orchestrator, err := agents.NewOrchestratorAgent(noteAgent, contactAgent, locationAgent, eventAgent, image.Image.ImageName, image.Image.Image) + orchestrator, err := agents.NewOrchestratorAgent(createLogger("Orchestrator 🎼"), noteAgent, contactAgent, locationAgent, eventAgent, image.Image.ImageName, image.Image.Image) if err != nil { panic(err) } - // Still need to find some way to hide this complexity away. - // I don't think wrapping agents in structs actually works too well. - err = orchestrator.Client.RunAgent(agents.OrchestratorPrompt, agents.OrchestratorTools, "noAction", nil, image.UserID, image.ImageID, image.Image.ImageName, image.Image.Image) + err = orchestrator.RunAgent(agents.OrchestratorPrompt, agents.OrchestratorTools, "noAction", nil, image.UserID, image.ImageID, image.Image.ImageName, image.Image.Image) if err != nil { - log.Println(err) + log.Error("Orchestrator failed", "error", "err") + return } - imageModel.FinishProcessing(ctx, image.ID) + _, 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) }() } }