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, 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
}