diff --git a/backend/agents/event_agent.go b/backend/agents/event_agent.go new file mode 100644 index 0000000..a375ea1 --- /dev/null +++ b/backend/agents/event_agent.go @@ -0,0 +1,200 @@ +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" +) + +const eventPrompt = ` +You are an agent. + +The user will send you images and you have to identify if they have any events or a place. +This could be a friend suggesting to meet, a conference, or anything that looks like an event. + +There are various tools you can use to perform this task. + +listEvents +Lists the users already existing events, you should do this before using createEvents to avoid creating duplicates. + +createEvent +Use this to create a new events. + +linkEvent +Links an image to a events. + +finish +Call when there is nothing else to do. +` + +const eventTools = ` +[ + { + "type": "function", + "function": { + "name": "listEvents", + "description": "List the events the user already has.", + "parameters": { + "type": "object", + "properties": {}, + "required": [] + } + } + }, + { + "type": "function", + "function": { + "name": "createEvent", + "description": "Use to create a new events", + "parameters": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "startDateTime": { + "type": "string", + "description": "The start time as an ISO string" + }, + "endDateTime": { + "type": "string", + "description": "The end time as an ISO string" + } + }, + "required": ["name"] + } + } + }, + { + "type": "function", + "function": { + "name": "linkEvent", + "description": "Use to link an already existing events to the image you were sent", + "parameters": { + "type": "object", + "properties": { + "eventId": { + "type": "string" + } + }, + "required": ["eventsId"] + } + } + }, + { + "type": "function", + "function": { + "name": "finish", + "description": "Call this when there is nothing left to do.", + "parameters": { + "type": "object", + "properties": {}, + "required": [] + } + } + } +]` + +type EventAgent struct { + client client.AgentClient + + eventsModel models.EventModel +} + +type listEventArguments struct{} +type createEventArguments struct { + Name string `json:"name"` + StartDateTime *string `json:"startDateTime"` + EndDateTime *string `json:"endDateTime"` + OrganizerName *string `json:"organizerName"` +} +type linkEventArguments struct { + EventID string `json:"eventId"` +} + +func NewEventAgent(eventsModel models.EventModel) (EventAgent, error) { + agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{ + ReportTimestamp: true, + TimeFormat: time.Kitchen, + Prefix: "Events 📍", + })) + + if err != nil { + return EventAgent{}, err + } + + agent := EventAgent{ + client: agentClient, + eventsModel: eventsModel, + } + + agentClient.ToolHandler.AddTool("listEvents", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { + return agent.eventsModel.List(context.Background(), info.UserId) + }) + + agentClient.ToolHandler.AddTool("createEvent", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) { + args := createEventArguments{} + err := json.Unmarshal([]byte(_args), &args) + if err != nil { + return model.Events{}, err + } + + ctx := context.Background() + + layout := "2006-01-02T15:04:05Z" + + startTime, err := time.Parse(layout, *args.StartDateTime) + if err != nil { + return model.Events{}, err + } + + endTime, err := time.Parse(layout, *args.EndDateTime) + if err != nil { + return model.Events{}, err + } + + events, err := agent.eventsModel.Save(ctx, info.UserId, model.Events{ + Name: args.Name, + StartDateTime: &startTime, + EndDateTime: &endTime, + }) + + if err != nil { + return model.Events{}, err + } + + _, err = agent.eventsModel.SaveToImage(ctx, info.ImageId, events.ID) + if err != nil { + return model.Events{}, err + } + + return events, nil + }) + + agentClient.ToolHandler.AddTool("linkEvent", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) { + args := linkEventArguments{} + err := json.Unmarshal([]byte(_args), &args) + if err != nil { + return "", err + } + + ctx := context.Background() + + contactUuid, err := uuid.Parse(args.EventID) + if err != nil { + return "", err + } + + agent.eventsModel.SaveToImage(ctx, info.ImageId, contactUuid) + return "Saved", nil + }) + + return agent, nil +} diff --git a/backend/agents/orchestrator.go b/backend/agents/orchestrator.go index 19a68d7..8cede86 100644 --- a/backend/agents/orchestrator.go +++ b/backend/agents/orchestrator.go @@ -2,6 +2,7 @@ package agents import ( "errors" + "fmt" "os" "screenmark/screenmark/agents/client" "time" @@ -29,6 +30,9 @@ Use it when the image contains information relating a person. locationAgent Use it when the image contains some address or a place. +eventAgent +Use it when the image contains an event, this can be a date, a message suggesting an event. + noAction When you think there is no more information to extract from the image. @@ -74,6 +78,18 @@ const OrchestratorTools = ` "required": [] } } + }, + { + "type": "function", + "function": { + "name": "eventAgent", + "description": "Use when then image contains some event", + "parameters": { + "type": "object", + "properties": {}, + "required": [] + } + } }, { "type": "function", @@ -99,7 +115,7 @@ type Status struct { Ok bool `json:"ok"` } -func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locationAgent LocationAgent, imageName string, imageData []byte) (OrchestratorAgent, error) { +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, @@ -134,6 +150,14 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati }, nil }) + agent.ToolHandler.AddTool("eventAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { + go eventAgent.client.RunAgent(eventPrompt, eventTools, "finish", info.UserId, info.ImageId, imageName, imageData) + + return Status{ + Ok: true, + }, nil + }) + agent.ToolHandler.AddTool("noAction", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) { // To nothing diff --git a/backend/events.go b/backend/events.go index 23f496a..9f7aead 100644 --- a/backend/events.go +++ b/backend/events.go @@ -23,6 +23,7 @@ func ListenNewImageEvents(db *sql.DB, eventManager *EventManager) { defer listener.Close() locationModel := models.NewLocationModel(db) + eventModel := models.NewEventModel(db) noteModel := models.NewNoteModel(db) imageModel := models.NewImageModel(db) contactModel := models.NewContactModel(db) @@ -56,6 +57,11 @@ func ListenNewImageEvents(db *sql.DB, eventManager *EventManager) { panic(err) } + eventAgent, err := agents.NewEventAgent(eventModel) + if err != nil { + panic(err) + } + image, err := imageModel.GetToProcessWithData(ctx, imageId) if err != nil { log.Println("Failed to GetToProcessWithData") @@ -69,7 +75,7 @@ func ListenNewImageEvents(db *sql.DB, eventManager *EventManager) { return } - orchestrator, err := agents.NewOrchestratorAgent(noteAgent, contactAgent, locationAgent, image.Image.ImageName, image.Image.Image) + orchestrator, err := agents.NewOrchestratorAgent(noteAgent, contactAgent, locationAgent, eventAgent, image.Image.ImageName, image.Image.Image) if err != nil { panic(err) } diff --git a/backend/models/events.go b/backend/models/events.go index c1ef38a..d473605 100644 --- a/backend/models/events.go +++ b/backend/models/events.go @@ -14,6 +14,20 @@ type EventModel struct { dbPool *sql.DB } +func (m EventModel) List(ctx context.Context, userId uuid.UUID) ([]model.Events, error) { + listEventsStmt := SELECT(Events.AllColumns). + FROM( + Events. + INNER_JOIN(UserEvents, UserEvents.EventID.EQ(Events.ID)), + ). + WHERE(UserEvents.UserID.EQ(UUID(userId))) + + events := []model.Events{} + + err := listEventsStmt.QueryContext(ctx, m.dbPool, &events) + return events, err +} + func (m EventModel) Save(ctx context.Context, userId uuid.UUID, event model.Events) (model.Events, error) { // TODO tx here insertEventStmt := Events.