feat(location): agent to create locations
This commit is contained in:
179
backend/agents/location_agent.go
Normal file
179
backend/agents/location_agent.go
Normal file
@@ -0,0 +1,179 @@
|
||||
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, you should do this before using createLocation to avoid creating duplicates.
|
||||
|
||||
createLocation
|
||||
Use this to create a new location. Avoid making duplicates and only create a new location if listLocations doesnt contain the location on the image.
|
||||
|
||||
linkLocation
|
||||
Links an image to 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": "linkLocation",
|
||||
"description": "Use to link an already existing location to the image you were sent",
|
||||
"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
|
||||
})
|
||||
|
||||
return agent, nil
|
||||
}
|
||||
@@ -20,17 +20,15 @@ The agents are available as tool calls.
|
||||
|
||||
Agents available:
|
||||
|
||||
eventLocationAgent
|
||||
Use it when you think the image contains an event or a location of any sort. This can be an event page, a map, an address or a date.
|
||||
This could also be a conversation describing an event.
|
||||
|
||||
noteAgent
|
||||
Use when there is ANY text on the image.
|
||||
|
||||
contactAgent
|
||||
|
||||
Use it when the image contains information relating a person.
|
||||
|
||||
locationAgent
|
||||
Use it when the image contains some address or a place.
|
||||
|
||||
noAction
|
||||
When you think there is no more information to extract from the image.
|
||||
|
||||
@@ -41,18 +39,6 @@ Do not call the agent if you do not think it is relevant for the image.
|
||||
|
||||
const OrchestratorTools = `
|
||||
[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "eventLocationAgent",
|
||||
"description": "Use when there is an event or location on the image. This could be in writing form",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
@@ -76,6 +62,18 @@ const OrchestratorTools = `
|
||||
"required": []
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "locationAgent",
|
||||
"description": "Use when then image contains some place, location or address",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
@@ -101,7 +99,7 @@ type Status struct {
|
||||
Ok bool `json:"ok"`
|
||||
}
|
||||
|
||||
func NewOrchestratorAgent(eventLocationAgent EventLocationAgent, noteAgent NoteAgent, contactAgent ContactAgent, imageName string, imageData []byte) (OrchestratorAgent, error) {
|
||||
func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locationAgent LocationAgent, imageName string, imageData []byte) (OrchestratorAgent, error) {
|
||||
agent, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
|
||||
ReportTimestamp: true,
|
||||
TimeFormat: time.Kitchen,
|
||||
@@ -112,17 +110,6 @@ func NewOrchestratorAgent(eventLocationAgent EventLocationAgent, noteAgent NoteA
|
||||
return OrchestratorAgent{}, err
|
||||
}
|
||||
|
||||
agent.ToolHandler.AddTool("eventLocationAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
|
||||
// We need a way to keep track of this async?
|
||||
// Probably just a DB, because we don't want to wait. The orchistrator shouldnt wait for this stuff to finish.
|
||||
|
||||
go eventLocationAgent.client.RunAgent(eventLocationPrompt, eventLocationTools, "finish", info.UserId, info.ImageId, imageName, imageData)
|
||||
|
||||
return Status{
|
||||
Ok: true,
|
||||
}, nil
|
||||
})
|
||||
|
||||
agent.ToolHandler.AddTool("noteAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
|
||||
go noteAgent.GetNotes(info.UserId, info.ImageId, imageName, imageData)
|
||||
|
||||
@@ -139,6 +126,14 @@ func NewOrchestratorAgent(eventLocationAgent EventLocationAgent, noteAgent NoteA
|
||||
}, nil
|
||||
})
|
||||
|
||||
agent.ToolHandler.AddTool("locationAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
|
||||
go locationAgent.client.RunAgent(locationPrompt, locationTools, "finish", info.UserId, info.ImageId, imageName, imageData)
|
||||
|
||||
return Status{
|
||||
Ok: true,
|
||||
}, nil
|
||||
})
|
||||
|
||||
agent.ToolHandler.AddTool("noAction", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
|
||||
// To nothing
|
||||
|
||||
|
||||
Reference in New Issue
Block a user