Haystack/backend/agents/orchestrator.go

package agents

import (
	"errors"
	"screenmark/screenmark/agents/client"

	"github.com/charmbracelet/log"
)

const orchestratorPrompt = `
You are an Orchestrator for various AI agents.

The user will send you images and you have to determine which agents you have to call, in order to best help the user.

You might decide no agent needs to be called.

The agents are available as tool calls.

Agents available:

noteAgent
Use when there is ANY text on the image. You should call this almost all the time.

contactAgent
Use it when the image contains information relating a person.

locationAgent
Use it when the image contains some address or a place.

eventAgent
Use it when the image contains an event, this can be a date, a message suggesting an event.

noAction
When you think there is no more information to extract from the image.

Always call agents in parallel if you need to call more than 1.

Do not call the agent if you do not think it is relevant for the image.
`

const orchestratorTools = `
[
	{
        "type": "function",
        "function": {
            "name": "noteAgent",
            "description": "Use when there is any text on the image, this can be code/text/formulas any writing",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "contactAgent",
            "description": "Use when then image contains some person or contact",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "locationAgent",
            "description": "Use when then image contains some place, location or address",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "eventAgent",
            "description": "Use when then image contains some event",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "noAction",
            "description": "Use when you are sure nothing can be done about this image anymore",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    }
]`

type OrchestratorAgent struct {
	Client client.AgentClient

	log log.Logger
}

type Status struct {
	Ok bool `json:"ok"`
}

func NewOrchestratorAgent(log *log.Logger, noteAgent NoteAgent, contactAgent client.AgentClient, locationAgent client.AgentClient, eventAgent client.AgentClient, imageName string, imageData []byte) client.AgentClient {
	agent := client.CreateAgentClient(client.CreateAgentClientOptions{
		SystemPrompt: orchestratorPrompt,
		JsonTools:    orchestratorTools,
		Log:          log,
		EndToolCall:  "noAction",
	})

	agent.ToolHandler.AddTool("noteAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		go noteAgent.GetNotes(info.UserId, info.ImageId, imageName, imageData)

		return Status{
			Ok: true,
		}, nil
	})

	agent.ToolHandler.AddTool("contactAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		go contactAgent.RunAgent(info.UserId, info.ImageId, imageName, imageData)

		return Status{
			Ok: true,
		}, nil
	})

	agent.ToolHandler.AddTool("locationAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		go locationAgent.RunAgent(info.UserId, info.ImageId, imageName, imageData)

		return Status{
			Ok: true,
		}, nil
	})

	agent.ToolHandler.AddTool("eventAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		go eventAgent.RunAgent(info.UserId, info.ImageId, imageName, imageData)

		return Status{
			Ok: true,
		}, nil
	})

	agent.ToolHandler.AddTool("noAction", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		// To nothing

		return Status{
			Ok: true,
		}, errors.New("Finished! Kinda bad return type but...")
	})

	return agent
}