Haystack/backend/agents/event_location_agent.go

package agents

import (
	"context"
	"encoding/json"
	"io"
	"log"
	"screenmark/screenmark/models"

	"github.com/google/uuid"
)

const eventLocationPrompt = `
You are an agent that extracts events and locations from an image.
Your job is to check if an image has an event or a location and use the correct tools to extract this information.

If you find an event, you should look for a location for this event on the image, it is possible an event doesn't have a location.

It is possible that there is no location or event on an image.

You should ask for a list of locations, as the user is likely to have this location saved. Reuse existing locations where possible.
`

// TODO: this should be read directly from a file on load.
const TOOLS = `
[
    {
        "type": "function",
        "function": {
            "name": "createLocation",
            "description": "Creates a location",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string"
                    },
                    "coordinates": {
                        "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"
                    },
                    "datetime": {
                        "type": "string"
                    },
                    "locationId": {
                        "type": "string",
                        "description": "The ID of the location, available by listLocations"
                    }
                },
                "required": ["name"]
            }
        }
    }
]
`

type EventLocationAgent struct {
	client AgentClient

	eventModel    models.EventModel
	locationModel models.LocationModel
}

func (agent EventLocationAgent) GetLocations(userId uuid.UUID, imageName string, imageData []byte) error {
	var tools any
	err := json.Unmarshal([]byte(TOOLS), &tools)

	auto := "auto"

	request := AgentRequestBody{
		Tools:       &tools,
		ToolChoice:  &auto,
		Model:       "pixtral-12b-2409",
		Temperature: 0.3,
		ResponseFormat: ResponseFormat{
			Type: "text",
		},
	}

	err = request.AddSystem(eventLocationPrompt)
	if err != nil {
		return err
	}

	log.Println(request)

	request.AddImage(imageName, imageData)

	jsonAiRequest, err := json.Marshal(request)
	if err != nil {
		return err
	}

	httpRequest, err := agent.client.getRequest(jsonAiRequest)
	if err != nil {
		return err
	}

	resp, err := agent.client.Do(httpRequest)
	if err != nil {
		return err
	}

	response, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	log.Println(string(response))

	agentResponse := AgentResponse{}
	err = json.Unmarshal(response, &agentResponse)

	toolCalls := agentResponse.Choices[0].Message.ToolCalls[0]

	if toolCalls.Function.Name == "listLocations" {
		locations, err := agent.locationModel.List(context.Background(), userId)
		if err != nil {
			return err
		}

		jsonLocations, err := json.Marshal(locations)
		if err != nil {
			return err
		}

		request.AddToolCall(AgentAssistantToolCall{
			Role:      "assistant",
			Content:   "",
			ToolCalls: []ToolCall{toolCalls},
		})

		request.AddText(AgentTextMessage{
			Role:       "tool",
			Name:       "listLocations",
			Content:    string(jsonLocations),
			ToolCallId: toolCalls.Id,
		})

		jsonAiRequest, err := json.Marshal(request)
		if err != nil {
			return err
		}

		httpRequest, err := agent.client.getRequest(jsonAiRequest)
		if err != nil {
			return err
		}

		resp, err := agent.client.Do(httpRequest)
		if err != nil {
			return err
		}

		response, err := io.ReadAll(resp.Body)
		if err != nil {
			return err
		}

		log.Println(string(response))
	}

	return nil
}

func NewLocationEventAgent(locationModel models.LocationModel, eventModel models.EventModel) (EventLocationAgent, error) {
	client, err := CreateAgentClient(eventLocationPrompt)
	if err != nil {
		return EventLocationAgent{}, err
	}

	return EventLocationAgent{
		client:        client,
		locationModel: locationModel,
		eventModel:    eventModel,
	}, nil
}