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

Another one of your tasks is to create a schema for this list. This should contain information that this, and following
pictures contain. Be specific but also generic. You should use the parameters in "createList" to create this schema.

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

**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."
                    },
					"schema": {
						"type": "array",
						"items": {
							"type": "object",
							"properties": {
								"key": {
									"type": "string",
									"description": "The name of the key for this specific field. Similar to a column in a database"
								},
								"value": {
									"type": "string",
									"enum": ["string", "number", "boolean"]
								}
							}
						}
					}
                },
                "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"`
	Schema     []struct {
		Key   string `json:"key"`
		Value string `json:"value"`
	}
}
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
		}

		log.Debug("Create list", "schema", args.Schema)

		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
}