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.

This schema should not be super specific. You must be able to understand the image, and if the content of the image doesnt seem relevant, try
and extract some meaning about what the image is.

You must call "listLists" to see which available lists are already available.
Use "createList" only once, don't create multiple lists for one image.

You can add an image to multiple lists, this is also true if you already created a list. But only add to a list if it makes sense to do so.

**Tools:**
* think: Internal reasoning/planning step.
* listLists: Get existing lists
* createList: Creates a new list with a name and description. Only use this once.
* addToList: Add to an existing list. This will also mean extracting information from this image, and inserting it, fitting the schema.
* 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": {
								"Item": {
									"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"]
								},
								"Description": {
									"type": "string",
									"description": "The description for this item"
								}
							}
						}
					}
                },
                "required": ["name", "description", "schema"]
            }
        }
    },
	{
		"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"
                    },
					"schema": {
						"type": "array",
						"items": {
							"type": "object",
							"description": "A key-value of ID - value from this image to fit the schema. any of the values can be null",
							"properties": {
								"id": {
									"type": "string",
									"description": "The UUID of the schema item."
								},
								"value": {
									"type": "string",
									"description": "the concrete value for this field"
								}
							}
						}
					}
                },
                "required": ["listId", "schema"]
            }

		}
	},
	  {
		"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 createListArguments struct {
	Name       string `json:"name"`
	Desription string `json:"description"`
	Schema     []model.SchemaItems
}
type addToListArguments struct {
	ListID string `json:"listId"`
	Schema []models.IDValue
}

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 "", err
		}

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

		if err != nil {
			log.Error(err)
			return "", err
		}

		log.Debug(savedList)

		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, args.Schema); err != nil {
			return "", err
		}

		return "Saved", nil
	})

	return agentClient
}