probo/pkg/llm/message.go

// Copyright (c) 2026 Probo Inc <hello@probo.com>.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package llm

import (
	"encoding/json"
	"fmt"
	"strings"
)

type (
	Message struct {
		Role       Role
		Parts      []Part
		ToolCalls  []ToolCall
		ToolCallID string // set when Role is RoleTool
	}

	ToolCall struct {
		ID       string
		Function FunctionCall
	}

	FunctionCall struct {
		Name      string
		Arguments string
	}

	Tool struct {
		Name        string
		Description string
		Parameters  json.RawMessage // JSON Schema
	}
)

type partEnvelope struct {
	Type string `json:"type"`
	// TextPart and ThinkingPart share the Text field.
	Text string `json:"text,omitempty"`
	// ImagePart fields
	URL string `json:"url,omitempty"`
	// FilePart fields
	Data     string `json:"data,omitempty"`
	MimeType string `json:"mime_type,omitempty"`
	Filename string `json:"filename,omitempty"`
	// ThinkingPart fields
	Signature string `json:"signature,omitempty"`
}

type messageJSON struct {
	Role       Role           `json:"role"`
	Parts      []partEnvelope `json:"parts,omitempty"`
	ToolCalls  []toolCallJSON `json:"tool_calls,omitempty"`
	ToolCallID string         `json:"tool_call_id,omitempty"`
}

type toolCallJSON struct {
	ID       string           `json:"id"`
	Function functionCallJSON `json:"function"`
}

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

func (m Message) MarshalJSON() ([]byte, error) {
	mj := messageJSON{
		Role:       m.Role,
		ToolCallID: m.ToolCallID,
	}

	if len(m.ToolCalls) > 0 {
		mj.ToolCalls = make([]toolCallJSON, len(m.ToolCalls))
		for i, tc := range m.ToolCalls {
			mj.ToolCalls[i] = toolCallJSON{
				ID: tc.ID,
				Function: functionCallJSON{
					Name:      tc.Function.Name,
					Arguments: tc.Function.Arguments,
				},
			}
		}
	}

	for _, p := range m.Parts {
		switch v := p.(type) {
		case TextPart:
			mj.Parts = append(mj.Parts, partEnvelope{Type: "text", Text: v.Text})
		case ImagePart:
			mj.Parts = append(mj.Parts, partEnvelope{Type: "image", URL: v.URL})
		case FilePart:
			mj.Parts = append(mj.Parts, partEnvelope{
				Type: "file", Data: v.Data, MimeType: v.MimeType, Filename: v.Filename,
			})
		case ThinkingPart:
			mj.Parts = append(mj.Parts, partEnvelope{Type: "thinking", Text: v.Text, Signature: v.Signature})
		default:
			return nil, fmt.Errorf("cannot marshal unknown Part type %T", p)
		}
	}

	return json.Marshal(mj)
}

func (m *Message) UnmarshalJSON(data []byte) error {
	var mj messageJSON
	if err := json.Unmarshal(data, &mj); err != nil {
		return err
	}

	m.Role = mj.Role
	m.ToolCallID = mj.ToolCallID

	if len(mj.ToolCalls) > 0 {
		m.ToolCalls = make([]ToolCall, len(mj.ToolCalls))
		for i, tc := range mj.ToolCalls {
			m.ToolCalls[i] = ToolCall{
				ID: tc.ID,
				Function: FunctionCall{
					Name:      tc.Function.Name,
					Arguments: tc.Function.Arguments,
				},
			}
		}
	} else {
		m.ToolCalls = nil
	}

	if len(mj.Parts) == 0 {
		m.Parts = nil
		return nil
	}

	m.Parts = make([]Part, len(mj.Parts))
	for i, env := range mj.Parts {
		switch env.Type {
		case "text":
			m.Parts[i] = TextPart{Text: env.Text}
		case "image":
			m.Parts[i] = ImagePart{URL: env.URL}
		case "file":
			m.Parts[i] = FilePart{Data: env.Data, MimeType: env.MimeType, Filename: env.Filename}
		case "thinking":
			m.Parts[i] = ThinkingPart{Text: env.Text, Signature: env.Signature}
		default:
			return fmt.Errorf("cannot unmarshal unknown Part type %q", env.Type)
		}
	}

	return nil
}

func (m Message) Text() string {
	var s strings.Builder

	for _, p := range m.Parts {
		if tp, ok := p.(TextPart); ok {
			s.WriteString(tp.Text)
		}
	}

	return s.String()
}

func (m Message) Thinking() string {
	var s strings.Builder

	for _, p := range m.Parts {
		if tp, ok := p.(ThinkingPart); ok {
			s.WriteString(tp.Text)
		}
	}

	return s.String()
}