149 lines
3.8 KiB
Go
149 lines
3.8 KiB
Go
package agents
|
|
|
|
import (
|
|
"errors"
|
|
"os"
|
|
"screenmark/screenmark/agents/client"
|
|
"time"
|
|
|
|
"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.
|
|
|
|
contactAgent
|
|
Use it when the image contains information relating a person.
|
|
|
|
locationAgent
|
|
Use it when the image contains some address or a place.
|
|
|
|
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": "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(noteAgent NoteAgent, contactAgent ContactAgent, locationAgent LocationAgent, imageName string, imageData []byte) (OrchestratorAgent, error) {
|
|
agent, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
|
|
ReportTimestamp: true,
|
|
TimeFormat: time.Kitchen,
|
|
Prefix: "Orchestrator 🎼",
|
|
}))
|
|
|
|
if err != nil {
|
|
return OrchestratorAgent{}, err
|
|
}
|
|
|
|
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.client.RunAgent(contactPrompt, contactTools, "finish", 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.client.RunAgent(locationPrompt, locationTools, "finish", 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 OrchestratorAgent{
|
|
Client: agent,
|
|
}, nil
|
|
}
|