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. createEvent Use this to create a new events. getEventLocationId Use this if the image contains a location or place. This tool will return the locationId. 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": "getEventLocationId", "description": "Get the ID of the location on the image, only use if the event contains a location or place.", "parameters": { "type": "object", "properties": {}, "required": [] } } }, { "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 locationAgent LocationAgent } 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, locationAgent LocationAgent) (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, locationAgent: locationAgent, } 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 }) 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) return locationAgent.client.Reply, nil }) return agent, nil }