Haystack/backend/agents/client/chat.go

package client

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"path/filepath"
)

type Chat struct {
	Messages []ChatMessage `json:"messages"`
}

type ChatMessage interface {
	IsResponse() bool
}

// TODO: the role could be inferred from the type.
// This would solve some bugs.

/*

Is there a world where this actually becomes the product?
Where we build such a resilient system of AI calls that we
can build some app builder, or even just an API system,
with a fancy UI?

Manage all the complexity for the user?

*/

// =============================================
// Messages from us to the AI.
// =============================================

type UserRole = string

const (
	User   UserRole = "user"
	System UserRole = "system"
)

type ToolRole = string

const (
	Tool ToolRole = "tool"
)

type ChatUserMessage struct {
	Role UserRole `json:"role"`

	MessageContent `json:"MessageContent"`
}

func (m ChatUserMessage) MarshalJSON() ([]byte, error) {
	switch t := m.MessageContent.(type) {
	case SingleMessage:
		return json.Marshal(&struct {
			Role    UserRole `json:"role"`
			Content string   `json:"content"`
		}{
			Role:    User,
			Content: t.Content,
		})
	case ArrayMessage:
		return json.Marshal(&struct {
			Role    UserRole                `json:"role"`
			Content []MessageContentMessage `json:"content"`
		}{
			Role:    User,
			Content: t.Content,
		})
	}

	return []byte{}, errors.New("Unreachable")
}

func (r ChatUserMessage) IsResponse() bool {
	return false
}

type ChatUserToolResponse struct {
	Role ToolRole `json:"role"`

	// The name of the function we are responding to.
	Name       string `json:"name"`
	Content    string `json:"content"`
	ToolCallId string `json:"tool_call_id"`
}

func (r ChatUserToolResponse) IsResponse() bool {
	return false
}

type ChatAiMessage struct {
	Role      string      `json:"role"`
	ToolCalls *[]ToolCall `json:"tool_calls,omitempty"`

	Content string `json:"content"`
}

func (m ChatAiMessage) IsResponse() bool {
	return true
}

// =============================================
// Unique interface for message content.
// =============================================

type MessageContent interface {
	IsSingleMessage() bool
}

type SingleMessage struct {
	Content string `json:"content"`
}

func (m SingleMessage) IsSingleMessage() bool {
	return true
}

type ArrayMessage struct {
	Content []MessageContentMessage `json:"content"`
}

func (m ArrayMessage) IsSingleMessage() bool {
	return false
}

type MessageContentMessage interface {
	IsImageMessage() bool
}

type TextMessageContent struct {
	TextType string `json:"type"`
	Text     string `json:"text"`
}

func (m TextMessageContent) IsImageMessage() bool {
	return false
}

type ImageMessageContent struct {
	ImageType string `json:"type"`
	ImageUrl  string `json:"image_url"`
}

func (m ImageMessageContent) IsImageMessage() bool {
	return true
}

type ImageContentUrl struct {
	Url string `json:"url"`
}

// =============================================
// Adjacent interfaces.
// =============================================

type ToolCall struct {
	Index    int          `json:"index"`
	Id       string       `json:"id"`
	Function FunctionCall `json:"function"`
}

type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

// =============================================
// Chat methods
// =============================================

func (chat *Chat) AddSystem(prompt string) {
	chat.Messages = append(chat.Messages, ChatUserMessage{
		Role: System,
		MessageContent: SingleMessage{
			Content: prompt,
		},
	})
}

func (chat *Chat) AddImage(imageName string, image []byte, query *string) error {
	extension := filepath.Ext(imageName)
	if len(extension) == 0 {
		// TODO: could also validate for image types we support.
		return errors.New("Image does not have extension")
	}

	extension = extension[1:]
	encodedString := base64.StdEncoding.EncodeToString(image)

	contentLength := 1
	if query != nil {
		contentLength += 1
	}

	messageContent := ArrayMessage{
		Content: make([]MessageContentMessage, contentLength),
	}

	index := 0

	if query != nil {
		messageContent.Content[index] = TextMessageContent{
			TextType: "text",
			Text:     *query,
		}
		index += 1
	}

	messageContent.Content[index] = ImageMessageContent{
		ImageType: "image_url",
		ImageUrl:  fmt.Sprintf("data:image/%s;base64,%s", extension, encodedString),
	}

	arrayMessage := ChatUserMessage{Role: User, MessageContent: messageContent}
	chat.Messages = append(chat.Messages, arrayMessage)

	return nil
}

func (chat *Chat) AddAiResponse(res ChatAiMessage) {
	chat.Messages = append(chat.Messages, res)
}

func (chat *Chat) AddToolResponse(res ChatUserToolResponse) {
	chat.Messages = append(chat.Messages, res)
}

func (chat Chat) GetLatest() (ChatMessage, error) {
	if len(chat.Messages) == 0 {
		return nil, errors.New("Not enough messages")
	}

	return chat.Messages[len(chat.Messages)-1], nil
}