Haystack/backend/agents/event_location_agent.go

140 lines
3.2 KiB
Go

package agents
import (
"encoding/json"
"io"
"log"
)
const eventLocationPrompt = `
You are an agent that extracts events and locations from an image.
Your job is to check if an image has an event or a location and use the correct tools to extract this information.
If you find an event, you should look for a location for this event on the image, it is possible an event doesn't have a location.
It is possible that there is no location or event on an image.
`
// TODO: this should be read directly from a file on load.
const TOOLS = `
[
{
"type": "function",
"function": {
"name": "createLocation",
"description": "Creates a location",
"parameters": {
"type": "object",
"properties": {
"name": {
"type": "string"
},
"coordinates": {
"type": "string"
},
"address": {
"type": "string"
}
},
"required": ["name"]
}
}
},
{
"type": "function",
"function": {
"name": "listLocations",
"description": "Lists the locations available",
"parameters": {
"type": "object",
"properties": {}
}
}
},
{
"type": "function",
"function": {
"name": "createEvent",
"description": "Creates a new event",
"parameters": {
"type": "object",
"properties": {
"name": {
"type": "string"
},
"datetime": {
"type": "string"
},
"locationId": {
"type": "string",
"description": "The ID of the location, available by listLocations"
}
},
"required": ["name"]
}
}
}
]
`
type EventLocationAgent = AgentClient
func (agent EventLocationAgent) GetLocations(imageName string, imageData []byte) error {
var tools any
err := json.Unmarshal([]byte(TOOLS), &tools)
auto := "auto"
request := AgentRequestBody{
Tools: &tools,
ToolChoice: &auto,
Model: "pixtral-12b-2409",
Temperature: 0.3,
ResponseFormat: ResponseFormat{
Type: "text",
},
}
err = request.AddSystem(eventLocationPrompt)
if err != nil {
return err
}
log.Println(request)
request.AddImage(imageName, imageData)
jsonAiRequest, err := json.Marshal(request)
if err != nil {
return err
}
httpRequest, err := agent.getRequest(jsonAiRequest)
if err != nil {
return err
}
resp, err := agent.Do(httpRequest)
if err != nil {
return err
}
response, err := io.ReadAll(resp.Body)
if err != nil {
return err
}
log.Println(string(response))
return nil
}
func NewLocationEventAgent() (EventLocationAgent, error) {
agent, err := CreateAgentClient(eventLocationPrompt)
if err != nil {
return EventLocationAgent{}, err
}
return agent, nil
}