99 lines
2.3 KiB
Go
99 lines
2.3 KiB
Go
package api
|
|
|
|
import (
|
|
"database/sql"
|
|
"database/sql/driver"
|
|
"encoding/json"
|
|
"time"
|
|
)
|
|
|
|
type Conversation struct {
|
|
ID uint `gorm:"primaryKey"`
|
|
ShortName sql.NullString
|
|
Title string
|
|
SelectedRootID *uint
|
|
SelectedRoot *Message `gorm:"foreignKey:SelectedRootID"`
|
|
}
|
|
|
|
type MessageRole string
|
|
|
|
const (
|
|
MessageRoleSystem MessageRole = "system"
|
|
MessageRoleUser MessageRole = "user"
|
|
MessageRoleAssistant MessageRole = "assistant"
|
|
MessageRoleToolCall MessageRole = "tool_call"
|
|
MessageRoleToolResult MessageRole = "tool_result"
|
|
)
|
|
|
|
type MessageMeta struct {
|
|
GenerationProvider *string `json:"generation_provider,omitempty"`
|
|
GenerationModel *string `json:"generation_model,omitempty"`
|
|
}
|
|
|
|
type Message struct {
|
|
ID uint `gorm:"primaryKey"`
|
|
CreatedAt time.Time
|
|
Metadata MessageMeta
|
|
|
|
ConversationID *uint `gorm:"index"`
|
|
Conversation *Conversation `gorm:"foreignKey:ConversationID"`
|
|
ParentID *uint
|
|
Parent *Message `gorm:"foreignKey:ParentID"`
|
|
Replies []Message `gorm:"foreignKey:ParentID"`
|
|
SelectedReplyID *uint
|
|
SelectedReply *Message `gorm:"foreignKey:SelectedReplyID"`
|
|
|
|
Role MessageRole
|
|
Content string
|
|
ToolCalls ToolCalls // a json array of tool calls (from the model)
|
|
ToolResults ToolResults // a json array of tool results
|
|
}
|
|
|
|
func (m *MessageMeta) Scan(value interface{}) error {
|
|
return json.Unmarshal(value.([]byte), m)
|
|
}
|
|
|
|
func (m MessageMeta) Value() (driver.Value, error) {
|
|
return json.Marshal(m)
|
|
}
|
|
|
|
func ApplySystemPrompt(m []Message, system string, force bool) []Message {
|
|
if len(m) > 0 && m[0].Role == MessageRoleSystem {
|
|
if force {
|
|
m[0].Content = system
|
|
}
|
|
return m
|
|
} else {
|
|
return append([]Message{{
|
|
Role: MessageRoleSystem,
|
|
Content: system,
|
|
}}, m...)
|
|
}
|
|
}
|
|
|
|
func (m *MessageRole) IsAssistant() bool {
|
|
switch *m {
|
|
case MessageRoleAssistant, MessageRoleToolCall:
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// FriendlyRole returns a human friendly signifier for the message's role.
|
|
func (m MessageRole) FriendlyRole() string {
|
|
switch m {
|
|
case MessageRoleUser:
|
|
return "You"
|
|
case MessageRoleSystem:
|
|
return "System"
|
|
case MessageRoleAssistant:
|
|
return "Assistant"
|
|
case MessageRoleToolCall:
|
|
return "Tool Call"
|
|
case MessageRoleToolResult:
|
|
return "Tool Result"
|
|
default:
|
|
return string(m)
|
|
}
|
|
}
|