opencode/internal/message/content.go
2025-05-13 11:15:14 -05:00

325 lines
6.7 KiB
Go

package message
import (
"encoding/base64"
"slices"
"time"
"github.com/sst/opencode/internal/llm/models"
)
type MessageRole string
const (
Assistant MessageRole = "assistant"
User MessageRole = "user"
System MessageRole = "system"
Tool MessageRole = "tool"
)
type FinishReason string
const (
FinishReasonEndTurn FinishReason = "end_turn"
FinishReasonMaxTokens FinishReason = "max_tokens"
FinishReasonToolUse FinishReason = "tool_use"
FinishReasonCanceled FinishReason = "canceled"
FinishReasonError FinishReason = "error"
FinishReasonPermissionDenied FinishReason = "permission_denied"
// Should never happen
FinishReasonUnknown FinishReason = "unknown"
)
type ContentPart interface {
isPart()
}
type ReasoningContent struct {
Thinking string `json:"thinking"`
}
func (tc ReasoningContent) String() string {
return tc.Thinking
}
func (ReasoningContent) isPart() {}
type TextContent struct {
Text string `json:"text"`
}
func (tc *TextContent) String() string {
if tc == nil {
return ""
}
return tc.Text
}
func (TextContent) isPart() {}
type ImageURLContent struct {
URL string `json:"url"`
Detail string `json:"detail,omitempty"`
}
func (iuc ImageURLContent) String() string {
return iuc.URL
}
func (ImageURLContent) isPart() {}
type BinaryContent struct {
Path string
MIMEType string
Data []byte
}
func (bc BinaryContent) String(provider models.ModelProvider) string {
base64Encoded := base64.StdEncoding.EncodeToString(bc.Data)
if provider == models.ProviderOpenAI {
return "data:" + bc.MIMEType + ";base64," + base64Encoded
}
return base64Encoded
}
func (BinaryContent) isPart() {}
type ToolCall struct {
ID string `json:"id"`
Name string `json:"name"`
Input string `json:"input"`
Type string `json:"type"`
Finished bool `json:"finished"`
}
func (ToolCall) isPart() {}
type ToolResult struct {
ToolCallID string `json:"tool_call_id"`
Name string `json:"name"`
Content string `json:"content"`
Metadata string `json:"metadata"`
IsError bool `json:"is_error"`
}
func (ToolResult) isPart() {}
type Finish struct {
Reason FinishReason `json:"reason"`
Time time.Time `json:"time"`
}
type DBFinish struct {
Reason FinishReason `json:"reason"`
Time int64 `json:"time"`
}
func (Finish) isPart() {}
func (m *Message) Content() *TextContent {
for _, part := range m.Parts {
if c, ok := part.(TextContent); ok {
return &c
}
}
return nil
}
func (m *Message) ReasoningContent() ReasoningContent {
for _, part := range m.Parts {
if c, ok := part.(ReasoningContent); ok {
return c
}
}
return ReasoningContent{}
}
func (m *Message) ImageURLContent() []ImageURLContent {
imageURLContents := make([]ImageURLContent, 0)
for _, part := range m.Parts {
if c, ok := part.(ImageURLContent); ok {
imageURLContents = append(imageURLContents, c)
}
}
return imageURLContents
}
func (m *Message) BinaryContent() []BinaryContent {
binaryContents := make([]BinaryContent, 0)
for _, part := range m.Parts {
if c, ok := part.(BinaryContent); ok {
binaryContents = append(binaryContents, c)
}
}
return binaryContents
}
func (m *Message) ToolCalls() []ToolCall {
toolCalls := make([]ToolCall, 0)
for _, part := range m.Parts {
if c, ok := part.(ToolCall); ok {
toolCalls = append(toolCalls, c)
}
}
return toolCalls
}
func (m *Message) ToolResults() []ToolResult {
toolResults := make([]ToolResult, 0)
for _, part := range m.Parts {
if c, ok := part.(ToolResult); ok {
toolResults = append(toolResults, c)
}
}
return toolResults
}
func (m *Message) IsFinished() bool {
for _, part := range m.Parts {
if _, ok := part.(Finish); ok {
return true
}
}
return false
}
func (m *Message) FinishPart() *Finish {
for _, part := range m.Parts {
if c, ok := part.(Finish); ok {
return &c
}
}
return nil
}
func (m *Message) FinishReason() FinishReason {
for _, part := range m.Parts {
if c, ok := part.(Finish); ok {
return c.Reason
}
}
return ""
}
func (m *Message) IsThinking() bool {
if m.ReasoningContent().Thinking != "" && m.Content().Text == "" && !m.IsFinished() {
return true
}
return false
}
func (m *Message) AppendContent(delta string) {
found := false
for i, part := range m.Parts {
if c, ok := part.(TextContent); ok {
m.Parts[i] = TextContent{Text: c.Text + delta}
found = true
}
}
if !found {
m.Parts = append(m.Parts, TextContent{Text: delta})
}
}
func (m *Message) AppendReasoningContent(delta string) {
found := false
for i, part := range m.Parts {
if c, ok := part.(ReasoningContent); ok {
m.Parts[i] = ReasoningContent{Thinking: c.Thinking + delta}
found = true
}
}
if !found {
m.Parts = append(m.Parts, ReasoningContent{Thinking: delta})
}
}
func (m *Message) FinishToolCall(toolCallID string) {
for i, part := range m.Parts {
if c, ok := part.(ToolCall); ok {
if c.ID == toolCallID {
m.Parts[i] = ToolCall{
ID: c.ID,
Name: c.Name,
Input: c.Input,
Type: c.Type,
Finished: true,
}
return
}
}
}
}
func (m *Message) AppendToolCallInput(toolCallID string, inputDelta string) {
for i, part := range m.Parts {
if c, ok := part.(ToolCall); ok {
if c.ID == toolCallID {
m.Parts[i] = ToolCall{
ID: c.ID,
Name: c.Name,
Input: c.Input + inputDelta,
Type: c.Type,
Finished: c.Finished,
}
return
}
}
}
}
func (m *Message) AddToolCall(tc ToolCall) {
for i, part := range m.Parts {
if c, ok := part.(ToolCall); ok {
if c.ID == tc.ID {
m.Parts[i] = tc
return
}
}
}
m.Parts = append(m.Parts, tc)
}
func (m *Message) SetToolCalls(tc []ToolCall) {
// remove any existing tool call part it could have multiple
parts := make([]ContentPart, 0)
for _, part := range m.Parts {
if _, ok := part.(ToolCall); ok {
continue
}
parts = append(parts, part)
}
m.Parts = parts
for _, toolCall := range tc {
m.Parts = append(m.Parts, toolCall)
}
}
func (m *Message) AddToolResult(tr ToolResult) {
m.Parts = append(m.Parts, tr)
}
func (m *Message) SetToolResults(tr []ToolResult) {
for _, toolResult := range tr {
m.Parts = append(m.Parts, toolResult)
}
}
func (m *Message) AddFinish(reason FinishReason) {
// remove any existing finish part
for i, part := range m.Parts {
if _, ok := part.(Finish); ok {
m.Parts = slices.Delete(m.Parts, i, i+1)
break
}
}
m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now()})
}
func (m *Message) AddImageURL(url, detail string) {
m.Parts = append(m.Parts, ImageURLContent{URL: url, Detail: detail})
}
func (m *Message) AddBinary(mimeType string, data []byte) {
m.Parts = append(m.Parts, BinaryContent{MIMEType: mimeType, Data: data})
}