119 lines
2.6 KiB
Go
119 lines
2.6 KiB
Go
|
package api
|
||
|
|
|
||
|
|
import (
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
)
|
||
|
|
|
||
|
|
type MessageRole string
|
||
|
|
|
||
|
|
const (
|
||
|
|
MessageRoleSystem MessageRole = "system"
|
||
|
|
MessageRoleUser MessageRole = "user"
|
||
|
|
MessageRoleAssistant MessageRole = "assistant"
|
||
|
|
MessageRoleToolCall MessageRole = "tool_call"
|
||
|
|
MessageRoleToolResult MessageRole = "tool_result"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Message struct {
|
||
|
|
Content string // TODO: support multi-part messages
|
||
|
|
Role MessageRole
|
||
|
|
ToolCalls []ToolCall
|
||
|
|
ToolResults []ToolResult
|
||
|
|
}
|
||
|
|
|
||
|
|
type ToolSpec struct {
|
||
|
|
Name string
|
||
|
|
Description string
|
||
|
|
Parameters []ToolParameter
|
||
|
|
Impl func(*ToolSpec, map[string]interface{}) (string, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
type ToolParameter struct {
|
||
|
|
Name string `json:"name"`
|
||
|
|
Type string `json:"type"` // "string", "integer", "boolean"
|
||
|
|
Required bool `json:"required"`
|
||
|
|
Description string `json:"description"`
|
||
|
|
Enum []string `json:"enum,omitempty"`
|
||
|
|
}
|
||
|
|
|
||
|
|
type ToolCall struct {
|
||
|
|
ID string `json:"id" yaml:"-"`
|
||
|
|
Name string `json:"name" yaml:"tool"`
|
||
|
|
Parameters map[string]interface{} `json:"parameters" yaml:"parameters"`
|
||
|
|
}
|
||
|
|
|
||
|
|
type ToolResult struct {
|
||
|
|
ToolCallID string `json:"toolCallID" yaml:"-"`
|
||
|
|
ToolName string `json:"toolName,omitempty" yaml:"tool"`
|
||
|
|
Result string `json:"result,omitempty" yaml:"result"`
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewMessageWithAssistant(content string) *Message {
|
||
|
|
return &Message{
|
||
|
|
Role: MessageRoleAssistant,
|
||
|
|
Content: content,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewMessageWithToolCalls(content string, toolCalls []ToolCall) *Message {
|
||
|
|
return &Message{
|
||
|
|
Role: MessageRoleToolCall,
|
||
|
|
Content: content,
|
||
|
|
ToolCalls: toolCalls,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m MessageRole) IsAssistant() bool {
|
||
|
|
switch m {
|
||
|
|
case MessageRoleAssistant, MessageRoleToolCall:
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m MessageRole) IsSystem() bool {
|
||
|
|
switch m {
|
||
|
|
case MessageRoleSystem:
|
||
|
|
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)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TODO: remove this
|
||
|
|
type CallResult struct {
|
||
|
|
Message string `json:"message"`
|
||
|
|
Result any `json:"result,omitempty"`
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r CallResult) ToJson() (string, error) {
|
||
|
|
if r.Message == "" {
|
||
|
|
// When message not supplied, assume success
|
||
|
|
r.Message = "success"
|
||
|
|
}
|
||
|
|
|
||
|
|
jsonBytes, err := json.Marshal(r)
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("Could not marshal CallResult to JSON: %v\n", err)
|
||
|
|
}
|
||
|
|
return string(jsonBytes), nil
|
||
|
|
}
|