Haystack/backend/agents/contact_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 contactPrompt = `
You are an agent that performs actions on contacts and people you find on an image.

You can use tools to achieve your task.

You should use listContacts to make sure that you don't create duplicate contacts.

Call createContact when you see there is a new contact on this image. Do not create duplicate contacts.
Or call linkContact when you think this image contains an existing contact.

Call finish if you dont think theres anything else to do.
`

const contactTools = `
[
    {
        "type": "function",
        "function": {
            "name": "listContacts",
            "description": "List the users existing contacts",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "createContact",
            "description": "Creates a new contact",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
						"description": "the name of the person"
                    },
                    "phoneNumber": {
                        "type": "string"
                    },
                    "address": {
                        "type": "string",
						"description": "their physical address"
                    },
                    "email": {
                        "type": "string"
                    }
                },
                "required": ["name"]
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "linkContact",
            "description": "Links an existing contact with this image",
            "parameters": {
                "type": "object",
                "properties": {
                    "contactId": {
                        "type": "string",
						"description": "The UUID of the existing contact"
                    }
                },
                "required": ["contactId"]
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "finish",
            "description": "Call when you dont think theres anything to do",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    }
]
`

type ContactAgent struct {
	client client.AgentClient

	contactModel models.ContactModel
}

type listContactsArguments struct{}
type createContactsArguments struct {
	Name        string  `json:"name"`
	PhoneNumber *string `json:"phoneNumber"`
	Address     *string `json:"address"`
	Email       *string `json:"email"`
}
type linkContactArguments struct {
	ContactID string `json:"contactId"`
}

func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {
	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
		ReportTimestamp: true,
		TimeFormat:      time.Kitchen,
		Prefix:          "Contacts 👥",
	}))
	if err != nil {
		return ContactAgent{}, err
	}

	agent := ContactAgent{
		client:       agentClient,
		contactModel: contactModel,
	}

	agentClient.ToolHandler.AddTool("listContacts", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		return agent.contactModel.List(context.Background(), info.UserId)
	})

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

		ctx := context.Background()

		contact, err := agent.contactModel.Save(ctx, info.UserId, model.Contacts{
			Name:        args.Name,
			PhoneNumber: args.PhoneNumber,
			Email:       args.Email,
		})

		if err != nil {
			return model.Contacts{}, err
		}

		_, err = agent.contactModel.SaveToImage(ctx, info.ImageId, contact.ID)
		if err != nil {
			return model.Contacts{}, err
		}

		return contact, nil
	})

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

		ctx := context.Background()

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

		_, err = agent.contactModel.SaveToImage(ctx, info.ImageId, contactUuid)
		if err != nil {
			return "", err
		}

		return "Saved", nil
	})

	return agent, nil
}