Haystack/backend/agents/client/client.go

package client

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"os"

	"github.com/charmbracelet/log"
	"github.com/google/uuid"
)

type ResponseFormat struct {
	Type       string `json:"type"`
	JsonSchema any    `json:"json_schema"`
}

type AgentRequestBody struct {
	Model          string         `json:"model"`
	Temperature    float64        `json:"temperature"`
	ResponseFormat ResponseFormat `json:"response_format"`

	Tools      *any    `json:"tools,omitempty"`
	ToolChoice *string `json:"tool_choice,omitempty"`

	EndToolCall string `json:"-"`

	Chat *Chat `json:"messages"`
}

func (req AgentRequestBody) MarshalJSON() ([]byte, error) {
	return json.Marshal(&struct {
		Model          string         `json:"model"`
		Temperature    float64        `json:"temperature"`
		ResponseFormat ResponseFormat `json:"response_format"`

		Tools      *any          `json:"tools,omitempty"`
		ToolChoice *string       `json:"tool_choice,omitempty"`
		Messages   []ChatMessage `json:"messages"`
	}{
		Model:          req.Model,
		Temperature:    req.Temperature,
		ResponseFormat: req.ResponseFormat,

		Tools:      req.Tools,
		ToolChoice: req.ToolChoice,

		Messages: req.Chat.Messages,
	})
}

type ResponseChoice struct {
	Index        int           `json:"index"`
	Message      ChatAiMessage `json:"message"`
	FinishReason string        `json:"finish_reason"`
}

type AgentResponse struct {
	Id      string           `json:"id"`
	Object  string           `json:"object"`
	Choices []ResponseChoice `json:"choices"`
	Created int              `json:"created"`
}

type AgentClient struct {
	url            string
	apiKey         string
	responseFormat string

	ToolHandler ToolsHandlers

	Log *log.Logger

	Reply string

	Do func(req *http.Request) (*http.Response, error)
}

const OPENAI_API_KEY = "OPENAI_API_KEY"

func CreateAgentClient(log *log.Logger) (AgentClient, error) {
	apiKey := os.Getenv(OPENAI_API_KEY)

	if len(apiKey) == 0 {
		return AgentClient{}, errors.New(OPENAI_API_KEY + " was not found.")
	}

	return AgentClient{
		apiKey: apiKey,
		url:    "https://api.mistral.ai/v1/chat/completions",
		Do: func(req *http.Request) (*http.Response, error) {
			client := &http.Client{}
			return client.Do(req)
		},

		Log: log,

		ToolHandler: ToolsHandlers{
			handlers: map[string]ToolHandler{},
		},
	}, nil
}

func (client AgentClient) getRequest(body []byte) (*http.Request, error) {
	req, err := http.NewRequest("POST", client.url, bytes.NewBuffer(body))
	if err != nil {
		return req, err
	}

	req.Header.Add("Authorization", "Bearer "+client.apiKey)
	req.Header.Add("Content-Type", "application/json")

	return req, nil
}

func (client AgentClient) Request(req *AgentRequestBody) (AgentResponse, error) {
	jsonAiRequest, err := json.Marshal(req)
	if err != nil {
		return AgentResponse{}, err
	}

	httpRequest, err := client.getRequest(jsonAiRequest)
	if err != nil {
		return AgentResponse{}, err
	}

	resp, err := client.Do(httpRequest)
	if err != nil {
		return AgentResponse{}, err
	}

	response, err := io.ReadAll(resp.Body)
	if err != nil {
		return AgentResponse{}, err
	}

	agentResponse := AgentResponse{}
	err = json.Unmarshal(response, &agentResponse)

	if err != nil {
		return AgentResponse{}, err
	}

	if len(agentResponse.Choices) != 1 {
		client.Log.Errorf("Received more than 1 choice from AI \n %s\n", string(response))
		return AgentResponse{}, errors.New("Unsupported. We currently only accept 1 choice from AI.")
	}

	client.Log.SetLevel(log.DebugLevel)

	msg := agentResponse.Choices[0].Message

	if len(msg.Content) > 0 {
		client.Log.Debugf("Content: %s", msg.Content)
	}

	if msg.ToolCalls != nil && len(*msg.ToolCalls) > 0 {
		client.Log.Debugf("Tool Call: %s", (*msg.ToolCalls)[0].Function.Name)

		prettyJson, err := json.MarshalIndent((*msg.ToolCalls)[0].Function.Arguments, "", "	")
		if err != nil {
			return AgentResponse{}, err
		}

		client.Log.Debugf("Arguments: %s", string(prettyJson))
	}

	req.Chat.AddAiResponse(agentResponse.Choices[0].Message)

	return agentResponse, nil
}

func (client AgentClient) ToolLoop(info ToolHandlerInfo, req *AgentRequestBody) error {
	for {
		err := client.Process(info, req)
		if err != nil {
			return err
		}

		_, err = client.Request(req)
		if err != nil {
			return err
		}
	}
}

var FinishedCall = errors.New("Last tool tool was called")

func (client *AgentClient) Process(info ToolHandlerInfo, req *AgentRequestBody) error {
	var err error

	message, err := req.Chat.GetLatest()
	if err != nil {
		return err
	}

	aiMessage, ok := message.(ChatAiMessage)
	if !ok {
		return errors.New("Latest message isnt an AI message")
	}

	if aiMessage.ToolCalls == nil {
		// Not an error, we just dont have any tool calls to process.
		return nil
	}

	for _, toolCall := range *aiMessage.ToolCalls {
		if toolCall.Function.Name == req.EndToolCall {
			return FinishedCall
		}

		toolResponse := client.ToolHandler.Handle(info, toolCall)

		if toolCall.Function.Name == "reply" {
			client.Reply = toolCall.Function.Arguments
		}

		client.Log.SetLevel(log.DebugLevel)
		client.Log.Debugf("Response: %s", toolResponse.Content)

		req.Chat.AddToolResponse(toolResponse)
	}

	return err
}

func (client AgentClient) RunAgent(systemPrompt string, jsonTools string, endToolCall string, query *string, userId uuid.UUID, imageId uuid.UUID, imageName string, imageData []byte) error {
	var tools any
	err := json.Unmarshal([]byte(jsonTools), &tools)

	toolChoice := "any"

	request := AgentRequestBody{
		Tools:       &tools,
		ToolChoice:  &toolChoice,
		Model:       "pixtral-12b-2409",
		Temperature: 0.3,
		EndToolCall: endToolCall,
		ResponseFormat: ResponseFormat{
			Type: "text",
		},
		Chat: &Chat{
			Messages: make([]ChatMessage, 0),
		},
	}

	request.Chat.AddSystem(systemPrompt)
	request.Chat.AddImage(imageName, imageData, query)

	_, err = client.Request(&request)
	if err != nil {
		return err
	}

	toolHandlerInfo := ToolHandlerInfo{
		ImageId:   imageId,
		ImageName: imageName,
		UserId:    userId,
		Image:     &imageData,
	}

	return client.ToolLoop(toolHandlerInfo, &request)
}