154 lines
3.9 KiB
Go
154 lines
3.9 KiB
Go
package agents
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"screenmark/screenmark/.gen/haystack/haystack/model"
|
|
"screenmark/screenmark/agents/client"
|
|
"screenmark/screenmark/models"
|
|
|
|
"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.
|
|
|
|
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": "finish",
|
|
"description": "Call when you dont think theres anything to do",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {},
|
|
"required": []
|
|
}
|
|
}
|
|
}
|
|
]
|
|
`
|
|
|
|
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(log *log.Logger, contactModel models.ContactModel) (client.AgentClient, error) {
|
|
agentClient, err := client.CreateAgentClient(log)
|
|
if err != nil {
|
|
return client.AgentClient{}, err
|
|
}
|
|
|
|
agentClient.ToolHandler.AddTool("listContacts", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
|
|
return 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 := contactModel.Save(ctx, info.UserId, model.Contacts{
|
|
Name: args.Name,
|
|
PhoneNumber: args.PhoneNumber,
|
|
Email: args.Email,
|
|
})
|
|
|
|
if err != nil {
|
|
return model.Contacts{}, err
|
|
}
|
|
|
|
_, err = 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 = contactModel.SaveToImage(ctx, info.ImageId, contactUuid)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
return "Saved", nil
|
|
})
|
|
|
|
return agentClient, nil
|
|
}
|