Haystack/backend/agents/location_agent.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.
    * If no location details could be extracted from the input, use stopAgent.

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. You may also use this to update the information in a location, by providing a locationId obtained from listLocations if you believe this location already exists.
* 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."
                    },
					"locationId": {
						"type": "string",
						"description": "The UUID of the location. You should only provide this IF you believe the contact already exists, from listLocation."
					},
                    "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"]
            }
        }
    },
	%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"`
	LocationID *string `json:"locationId"`
	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()

		locationId := uuid.Nil
		if args.LocationID != nil {
			locationUuid, err := uuid.Parse(*args.LocationID)
			if err != nil {
				return model.Locations{}, err
			}

			locationId = locationUuid
		}

		location, err := locationModel.Save(ctx, info.UserId, model.Locations{
			ID:      locationId,
			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("reply", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		return "ok", nil
	})

	return agentClient
}