Haystack/backend/agents/location_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 locationPrompt = `
You are an agent.

The user will send you images and you have to identify if they have any location or a place. 
This could a picture of a real place, an address, or it's name.

There are various tools you can use to perform this task.

listLocations
Lists the users already existing locations.

createLocation
Use this to create a new location, when you don't see a matching one from listLocations call.

reply
Use this only if the user has asked a question about a location.

finish
Call when there is nothing else to do.
`

const locationTools = `
[
    {
        "type": "function",
        "function": {
            "name": "listLocations",
            "description": "List the locations the user already has.",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "createLocation",
            "description": "Use to create a new location",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string"
                    },
                    "address": {
                        "type": "string"
                    }
				},
                "required": ["name"]
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "reply",
            "description": "Reply to a user query, only if the user has asked something",
            "parameters": {
                "type": "object",
                "properties": {
					"locationId": {
						"type": "string"
					}
				},
                "required": ["locationId"]
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "finish",
            "description": "Call this when there is nothing left to do.",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    }
]`

type LocationAgent struct {
	client client.AgentClient

	locationModel models.LocationModel
}

type listLocationArguments struct{}
type createLocationArguments struct {
	Name    string  `json:"name"`
	Address *string `json:"address"`
}
type linkLocationArguments struct {
	LocationID string `json:"locationId"`
}

func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error) {
	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
		ReportTimestamp: true,
		TimeFormat:      time.Kitchen,
		Prefix:          "Locations 📍",
	}))

	if err != nil {
		return LocationAgent{}, err
	}

	agent := LocationAgent{
		client:        agentClient,
		locationModel: locationModel,
	}

	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 model.Locations{}, err
		}

		_, err = agent.locationModel.SaveToImage(ctx, info.ImageId, location.ID)
		if err != nil {
			return model.Locations{}, err
		}

		return location, nil
	})

	agentClient.ToolHandler.AddTool("linkLocation", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
		args := linkLocationArguments{}
		err := json.Unmarshal([]byte(_args), &args)
		if err != nil {
			return "", err
		}

		ctx := context.Background()

		contactUuid, err := uuid.Parse(args.LocationID)
		if err != nil {
			return "", err
		}

		agent.locationModel.SaveToImage(ctx, info.ImageId, contactUuid)
		return "Saved", nil
	})

	agentClient.ToolHandler.AddTool("reply", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		agent.client.Log.Debug(args)
		return "ok", nil
	})

	return agent, nil
}