Haystack/backend/agents/list_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 listPrompt = `
**You are an AI used to classify what list a certain image belongs in**

You will need to decide using tool calls, if you must create a new list, or use an existing one.
You must be specific enough so it is useful, but not too specific such that all images belong on seperate lists.

An example of lists are:
- Locations
- Events
- TV Shows
- Movies
- Books

You must call "listLists" to see which available lists are already available.

*Important*
You must not create lists with the names Locations, Events, Contacts or Notes. You can create lists adjacent to those, but
those lists are dealt with seperately.

**Tools:**
* think: Internal reasoning/planning step.
* listLists: Get existing lists
* createList: Creates a new list with a name and description.
* addToList: Add to an existing list.
* stopAgent: Signal task completion.
`

const listTools = `
[
	{
			"type": "function",
			"function": {
				"name": "think",
				"description": "Use this tool to think through the image, evaluating the event and whether or not it exists in the users listEvents.",
				"parameters": {
					"type": "object",
					"properties": {
						"thought": {
							"type": "string",
							"description": "A singular thought about the image."
						}
					},
					"required": ["thought"]
				}
			}
	  },
    {
        "type": "function",
        "function": {
            "name": "listLists",
            "description": "Retrieves the list of the user's existing lists.",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "createList",
            "description": "Creates a new list",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": "The name of this new list."
                    },
                    "description": {
                        "type": "string",
                        "description": "A simple description of this list."
                    }
                },
                "required": ["name", "description"]
            }
        }
    },
	{
		"type": "function",
		"function": {
			"name": "addToList",
			"description": "Adds an image to a list, this could be a new one you just created or not.",
            "parameters": {
                "type": "object",
                "properties": {
                    "listId": {
                        "type": "string",
                        "description": "The UUID of the existing list"
                    }
                },
                "required": ["listId"]
            }

		}
	},
	  {
		"type": "function",
		"function": {
		  "name": "stopAgent",
		  "description": "Use this tool to signal that the contact processing for the current image is complete. Call this *only* when: 1) No contact info was found initially, OR 2) All found contacts were confirmed to already exist after calling listContacts, OR 3) All necessary createContact calls for new individuals have been completed.",
		  "parameters": {
			"type": "object",
			"properties": {},
			"required": []
		  }
		}
	  }
]`

type listListsArguments struct{}
type createListArguments struct {
	Name       string `json:"name"`
	Desription string `json:"description"`
}
type addToListArguments struct {
	ListID string `json:"listId"`
}

func NewListAgent(log *log.Logger, listModel models.ListModel) client.AgentClient {
	agentClient := client.CreateAgentClient(client.CreateAgentClientOptions{
		SystemPrompt: listPrompt,
		JsonTools:    listTools,
		Log:          log,
		EndToolCall:  "stopAgent",
	})

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

	agentClient.ToolHandler.AddTool("listLists", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		return listModel.List(context.Background(), info.UserId)
	})

	agentClient.ToolHandler.AddTool("createList", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
		args := createListArguments{}
		err := json.Unmarshal([]byte(_args), &args)
		if err != nil {
			return model.Events{}, err
		}

		ctx := context.Background()
		savedList, err := listModel.Save(ctx, info.UserId, args.Name, args.Desription)

		if err != nil {
			return model.Lists{}, err
		}

		return savedList, nil
	})

	agentClient.ToolHandler.AddTool("addToList", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
		args := addToListArguments{}
		err := json.Unmarshal([]byte(_args), &args)
		if err != nil {
			return "", err
		}

		ctx := context.Background()

		listUuid, err := uuid.Parse(args.ListID)
		if err != nil {
			return "", err
		}

		if err := listModel.SaveInto(ctx, listUuid, info.ImageId); err != nil {
			return "", err
		}

		return "Saved", nil
	})

	return agentClient
}