Haystack/backend/agents/location_agent.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 locationPrompt = `
Role: Location AI Assistant

Objective: Identify locations from images/text, manage a saved list, and answer user queries about saved locations using the provided tools.

Core Logic:

1.  **Analyze Input:** Look for location details (Name, Address) in the image and check for any user query about a location.

2.  **Handle User Query First:**
    * If the user asks about a *specific* location:
        * Use listLocations to find its locationId.
        * If found, use reply with the locationId.
        * If not found, prepare summary (no matching location).
    * *(Proceed to step 4)*

3.  **Handle Image Location (if no query was handled):**
    * If location details were found in the image:
        * Use listLocations to check if it's already saved.
        * If *new*, use createLocation (Name is required).
        * If *duplicate*, prepare summary (location already exists).
    * *(Proceed to step 4)*

4.  **Summarize & Stop:** Always finish by writing a message explaining what you did (e.g., called reply, called createLocation, found a duplicate, couldn't find a match) or if no location information was found. After providing the summary message, call stopAgent to signal the end of processing for this turn.

Tool Usage:

* listLocations: Check saved locations (for queries or before saving).
* createLocation: Save a *new* location (requires Name).
* reply: Answer a query about a *known*, *saved* location using its locationId.
* stopAgent: Signals the end of the agent's processing for the current turn. Call this *after* providing the summary message.
* **Constraint:** Typically, only one main action tool (listLocations, createLocation, or reply) will be called per turn before summarizing and stopping. listLocations might precede createLocation or reply within the logic.
`

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": "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"]
            }
        }
    },
	{
		"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": []
		  }
		}
	  }
]`

type listLocationArguments struct{}
type createLocationArguments struct {
	Name    string  `json:"name"`
	Address *string `json:"address"`
}
type linkLocationArguments struct {
	LocationID string `json:"locationId"`
}

func NewLocationAgent(log *log.Logger, locationModel models.LocationModel) client.AgentClient {
	agentClient := client.CreateAgentClient(client.CreateAgentClientOptions{
		SystemPrompt: locationPrompt,
		JsonTools:    locationTools,
		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("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
		}

		locationModel.SaveToImage(ctx, info.ImageId, contactUuid)
		return "Saved", nil
	})

	agentClient.ToolHandler.AddTool("reply", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		return "ok", nil
	})

	return agentClient
}