296 lines
8.0 KiB
Go
296 lines
8.0 KiB
Go
package agents
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"screenmark/screenmark/.gen/haystack/haystack/model"
|
|
"screenmark/screenmark/agents/client"
|
|
"screenmark/screenmark/models"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
// This prompt is probably shit.
|
|
const eventLocationPrompt = `
|
|
You are an agent that extracts events, locations, and organizers from an image. Your primary tasks are to identify and create locations and organizers before creating events. Follow these steps:
|
|
|
|
Identify and Create Locations:
|
|
|
|
Check if the image contains a location.
|
|
If a location is found, check if it exists in the listLocations.
|
|
If the location does not exist, create it first.
|
|
Always reuse existing locations from listLocations to avoid duplicates.
|
|
|
|
Identify and Create Events:
|
|
|
|
Check if the image contains an event. An event should have a name and a date.
|
|
If an event is found, ensure you have a location (from step 1) and an organizer (from step 2) before creating the event.
|
|
Events must have an associated location and organizer. Do not create an event without these.
|
|
If possible, return a start time and an end time as ISO datetime strings.
|
|
Handling Images Without Events or Locations:
|
|
|
|
It is possible that the image does not contain an event or a location. In such cases, do not create an event.
|
|
Always prioritize the creation of locations and organizers before events. Ensure that all events have an associated location and organizer.
|
|
`
|
|
|
|
// TODO: this should be read directly from a file on load.
|
|
const TOOLS = `
|
|
[
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "createLocation",
|
|
"description": "Creates a location. No not use if you think an existing location is suitable!",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {
|
|
"type": "string"
|
|
},
|
|
"address": {
|
|
"type": "string"
|
|
}
|
|
},
|
|
"required": ["name"]
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "listLocations",
|
|
"description": "Lists the locations available",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {}
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "createEvent",
|
|
"description": "Creates a new event",
|
|
"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"
|
|
},
|
|
"locationId": {
|
|
"type": "string",
|
|
"description": "The ID of the location, available by listLocations"
|
|
},
|
|
"organizerName": {
|
|
"type": "string",
|
|
"description": "The name of the organizer"
|
|
}
|
|
},
|
|
"required": ["name"]
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "finish",
|
|
"description": "Nothing else to do. call this function.",
|
|
"parameters": {}
|
|
}
|
|
}
|
|
]
|
|
`
|
|
|
|
type EventLocationAgent struct {
|
|
client client.AgentClient
|
|
|
|
eventModel models.EventModel
|
|
locationModel models.LocationModel
|
|
contactModel models.ContactModel
|
|
|
|
toolHandler client.ToolsHandlers
|
|
}
|
|
|
|
type ListLocationArguments struct{}
|
|
type ListOrganizerArguments struct{}
|
|
|
|
type CreateLocationArguments struct {
|
|
Name string `json:"name"`
|
|
Address *string `json:"address,omitempty"`
|
|
Coordinates *string `json:"coordinates,omitempty"`
|
|
}
|
|
|
|
type CreateOrganizerArguments struct {
|
|
Name string `json:"name"`
|
|
PhoneNumber *string `json:"phoneNumber,omitempty"`
|
|
Email *string `json:"email,omitempty"`
|
|
}
|
|
|
|
type AttachImageLocationArguments struct {
|
|
LocationId string `json:"locationId"`
|
|
}
|
|
|
|
type CreateEventArguments struct {
|
|
Name string `json:"name"`
|
|
StartDateTime string `json:"startDateTime"`
|
|
EndDateTime string `json:"endDateTime"`
|
|
LocationId string `json:"locationId"`
|
|
OrganizerName string `json:"organizerName"`
|
|
}
|
|
|
|
func (agent EventLocationAgent) GetLocations(userId uuid.UUID, imageId uuid.UUID, imageName string, imageData []byte) error {
|
|
var tools any
|
|
err := json.Unmarshal([]byte(TOOLS), &tools)
|
|
|
|
toolChoice := "any"
|
|
|
|
request := client.AgentRequestBody{
|
|
Tools: &tools,
|
|
ToolChoice: &toolChoice,
|
|
Model: "pixtral-12b-2409",
|
|
Temperature: 0.3,
|
|
EndToolCall: "finish",
|
|
ResponseFormat: client.ResponseFormat{
|
|
Type: "text",
|
|
},
|
|
Chat: &client.Chat{
|
|
Messages: make([]client.ChatMessage, 0),
|
|
},
|
|
}
|
|
|
|
request.Chat.AddSystem(eventLocationPrompt)
|
|
request.Chat.AddImage(imageName, imageData)
|
|
|
|
_, err = agent.client.Request(&request)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
toolHandlerInfo := client.ToolHandlerInfo{
|
|
ImageId: imageId,
|
|
UserId: userId,
|
|
}
|
|
|
|
return agent.client.ToolLoop(toolHandlerInfo, &request)
|
|
}
|
|
|
|
func NewLocationEventAgent(locationModel models.LocationModel, eventModel models.EventModel, contactModel models.ContactModel) (EventLocationAgent, error) {
|
|
agentClient, err := client.CreateAgentClient()
|
|
if err != nil {
|
|
return EventLocationAgent{}, err
|
|
}
|
|
|
|
agent := EventLocationAgent{
|
|
client: agentClient,
|
|
locationModel: locationModel,
|
|
eventModel: eventModel,
|
|
contactModel: contactModel,
|
|
}
|
|
|
|
agentClient.ToolHandler.AddTool("listLocations",
|
|
func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
|
|
return agent.locationModel.List(context.Background(), info.UserId)
|
|
},
|
|
)
|
|
|
|
agentClient.ToolHandler.AddTool("createLocation",
|
|
func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
|
|
args := CreateLocationArguments{}
|
|
err := json.Unmarshal([]byte(_args), &args)
|
|
if err != nil {
|
|
return model.Locations{}, err
|
|
}
|
|
|
|
ctx := context.Background()
|
|
|
|
location, err := agent.locationModel.Save(ctx, info.UserId, model.Locations{
|
|
Name: args.Name,
|
|
Address: args.Address,
|
|
})
|
|
|
|
if err != nil {
|
|
return location, err
|
|
}
|
|
|
|
_, err = agent.locationModel.SaveToImage(ctx, info.ImageId, location.ID)
|
|
|
|
return location, err
|
|
},
|
|
)
|
|
|
|
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.Locations{}, 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
|
|
}
|
|
|
|
event, err := agent.eventModel.Save(ctx, info.UserId, model.Events{
|
|
Name: args.Name,
|
|
StartDateTime: &startTime,
|
|
EndDateTime: &endTime,
|
|
})
|
|
|
|
if err != nil {
|
|
return event, err
|
|
}
|
|
|
|
organizer, err := agent.contactModel.Save(ctx, info.UserId, model.Contacts{
|
|
Name: args.Name,
|
|
})
|
|
|
|
if err != nil {
|
|
return event, err
|
|
}
|
|
|
|
_, err = agent.eventModel.SaveToImage(ctx, info.ImageId, event.ID)
|
|
if err != nil {
|
|
return event, err
|
|
}
|
|
|
|
_, err = agent.contactModel.SaveToImage(ctx, info.ImageId, organizer.ID)
|
|
if err != nil {
|
|
return event, err
|
|
}
|
|
|
|
locationId, err := uuid.Parse(args.LocationId)
|
|
if err != nil {
|
|
return event, err
|
|
}
|
|
|
|
event, err = agent.eventModel.UpdateLocation(ctx, event.ID, locationId)
|
|
if err != nil {
|
|
return event, err
|
|
}
|
|
|
|
return agent.eventModel.UpdateOrganizer(ctx, event.ID, organizer.ID)
|
|
},
|
|
)
|
|
|
|
return agent, nil
|
|
}
|