224 lines
8.4 KiB
Go
224 lines
8.4 KiB
Go
package agents
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"screenmark/screenmark/.gen/haystack/haystack/model"
|
|
"screenmark/screenmark/agents/client"
|
|
"screenmark/screenmark/models"
|
|
|
|
"github.com/charmbracelet/log"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
const locationPrompt = `
|
|
Role: Location AI Assistant
|
|
|
|
Objective: Identify locations from images/text, manage a saved list (create, update), and answer user queries about saved locations using the provided tools.
|
|
|
|
Core Logic:
|
|
|
|
**Handle Image/Text Location (if no query was handled in Step 2):**
|
|
* If location details (InputName, InputAddress, etc.) were successfully extracted from the input in Step 1:
|
|
* Use listLocations to check if a location matching InputName or InputAddress already exists in the saved list.
|
|
* **If *no match*** is found:
|
|
* Use createLocation, providing the extracted InputName (required) and any other details like InputAddress.
|
|
* *(Proceed to step 5)*
|
|
* **If a *match*** is found (meaning it's a potential duplicate/update candidate - let the matching saved location have ExistingLocationId):
|
|
* Use updateLocation. Provide the locationId = ExistingLocationId.
|
|
* Also provide any *new or potentially refined* details extracted from the current input (e.g., name = InputName, address = InputAddress). The updateLocation tool should handle updating the record with these details and/or linking the new input context (like the image) to this existing location.
|
|
* *(Proceed to step 5)*
|
|
* If no location details could be extracted from the input in Step 1.
|
|
* *(Proceed to step 5)*
|
|
|
|
**Summarize & Stop:** Always finish by writing a concise message explaining the outcome of the turn. Examples:
|
|
* "Okay, I've answered your query about [Location Name]." (after calling reply)
|
|
* "I couldn't find [Queried Location Name] in my saved list." (after failing to find a match for a query)
|
|
* "I've saved '[InputName]' as a new location." (after calling createLocation)
|
|
* "I found that '[InputName]' was already saved, so I've updated its information/context based on your latest input." (after calling updateLocation)
|
|
* "I couldn't identify a specific location from your input." (if Step 1 failed or no action was taken)
|
|
* After providing the summary message, call stopAgent to signal the end of processing for this turn.
|
|
|
|
Tool Usage:
|
|
|
|
* listLocations: Check saved locations. Used to find matches for user queries or to detect existing entries before creating/updating. Returns matching location(s) including their locationId.
|
|
* createLocation: Save a *new* location. Requires name, can include address, etc.
|
|
* updateLocation: Update an *existing* location. Requires locationId. Can include name, address, etc., to update specific fields or simply to associate the new input context with the existing location.
|
|
* stopAgent: Signals the end of the agent's processing for the current turn. Call this *after* providing the summary message.
|
|
`
|
|
|
|
const replyTool = `
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "reply",
|
|
"description": "Signals intent to provide information about a specific known location in response to a user's query. Use only if the user asked a question and the location's ID was found via listLocations.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"locationId": {
|
|
"type": "string",
|
|
"description": "The unique identifier of the saved location that the user is asking about."
|
|
}
|
|
},
|
|
"required": ["locationId"]
|
|
}
|
|
}
|
|
},`
|
|
|
|
const locationTools = `
|
|
[
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "listLocations",
|
|
"description": "Retrieves the list of the user's currently saved locations (names, addresses, IDs). Use this first to check if a location from an image already exists, or to find the ID of a location the user is asking about.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {},
|
|
"required": []
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "createLocation",
|
|
"description": "Creates a new location entry in the user's saved list. Use only after listLocations confirms the location does not already exist.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {
|
|
"type": "string",
|
|
"description": "The primary name of the location (e.g., 'Eiffel Tower', 'Mom's House', 'Acme Corp HQ'). This field is mandatory."
|
|
},
|
|
"address": {
|
|
"type": "string",
|
|
"description": "The full street address of the location, if available (e.g., 'Champ de Mars, 5 Av. Anatole France, 75007 Paris, France'). Include if extracted."
|
|
}
|
|
},
|
|
"required": ["name"]
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "updateLocation",
|
|
"description": "Updates an existing saved location identified by its locationId. Use when input matches a pre-existing location. Pass locationId and any new details (name, address) to update.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"locationId": {
|
|
"type": "string",
|
|
"description": "The UUID of the location you are trying to update"
|
|
}
|
|
},
|
|
"required": ["locationId"]
|
|
}
|
|
}
|
|
},
|
|
%s
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "stopAgent",
|
|
"description": "Use this tool to signal that the contact processing for the current image is complete.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {},
|
|
"required": []
|
|
}
|
|
}
|
|
}
|
|
]`
|
|
|
|
func getLocationAgentTools(allowReply bool) string {
|
|
if allowReply {
|
|
return fmt.Sprintf(locationTools, replyTool)
|
|
} else {
|
|
return fmt.Sprintf(locationTools, "")
|
|
}
|
|
}
|
|
|
|
type listLocationArguments struct{}
|
|
type createLocationArguments struct {
|
|
Name string `json:"name"`
|
|
Address *string `json:"address"`
|
|
}
|
|
type updateLocationArguments struct {
|
|
LocationID string `json:"locationId"`
|
|
}
|
|
|
|
func NewLocationAgentWithComm(log *log.Logger, locationModel models.LocationModel) client.AgentClient {
|
|
client := NewLocationAgent(log, locationModel)
|
|
|
|
client.Options.JsonTools = getLocationAgentTools(true)
|
|
|
|
return client
|
|
}
|
|
|
|
func NewLocationAgent(log *log.Logger, locationModel models.LocationModel) client.AgentClient {
|
|
agentClient := client.CreateAgentClient(client.CreateAgentClientOptions{
|
|
SystemPrompt: locationPrompt,
|
|
JsonTools: getLocationAgentTools(false),
|
|
Log: log,
|
|
EndToolCall: "stopAgent",
|
|
})
|
|
|
|
agentClient.ToolHandler.AddTool("listLocations", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
|
|
return 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 := locationModel.Save(ctx, info.UserId, model.Locations{
|
|
Name: args.Name,
|
|
Address: args.Address,
|
|
})
|
|
|
|
if err != nil {
|
|
return model.Locations{}, err
|
|
}
|
|
|
|
_, err = locationModel.SaveToImage(ctx, info.ImageId, location.ID)
|
|
if err != nil {
|
|
return model.Locations{}, err
|
|
}
|
|
|
|
return location, nil
|
|
})
|
|
|
|
agentClient.ToolHandler.AddTool("updateLocation", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
|
|
args := updateLocationArguments{}
|
|
err := json.Unmarshal([]byte(_args), &args)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
ctx := context.Background()
|
|
|
|
locationId, err := uuid.Parse(args.LocationID)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
locationModel.SaveToImage(ctx, info.ImageId, locationId)
|
|
return "Saved", nil
|
|
})
|
|
|
|
agentClient.ToolHandler.AddTool("reply", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
|
|
return "ok", nil
|
|
})
|
|
|
|
return agentClient
|
|
}
|