176 lines
4.3 KiB
Go
176 lines
4.3 KiB
Go
package agents
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"screenmark/screenmark/.gen/haystack/haystack/model"
|
|
"screenmark/screenmark/agents/client"
|
|
"screenmark/screenmark/models"
|
|
|
|
"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 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"`
|
|
}
|
|
|
|
// Yeah this is just a copy of the other one.
|
|
func (agent ContactAgent) GetContacts(userId uuid.UUID, imageId uuid.UUID, imageName string, imageData []byte) error {
|
|
var tools any
|
|
err := json.Unmarshal([]byte(contactTools), &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 NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {
|
|
agentClient, err := client.CreateAgentClient()
|
|
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
|
|
})
|
|
|
|
return agent, nil
|
|
}
|