Haystack/backend/agents/event_agent.go

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
}