diff options
Diffstat (limited to 'models/models.go')
| -rw-r--r-- | models/models.go | 35 |
1 files changed, 18 insertions, 17 deletions
diff --git a/models/models.go b/models/models.go index c5e9435..1e65b25 100644 --- a/models/models.go +++ b/models/models.go @@ -105,12 +105,13 @@ type ImageContentPart struct { // RoleMsg represents a message with content that can be either a simple string or structured content parts type RoleMsg struct { - Role string `json:"role"` - Content string `json:"-"` - ContentParts []any `json:"-"` - ToolCallID string `json:"tool_call_id,omitempty"` // For tool response messages - KnownTo []string `json:"known_to,omitempty"` - hasContentParts bool // Flag to indicate which content type to marshal + Role string `json:"role"` + Content string `json:"-"` + ContentParts []any `json:"-"` + ToolCallID string `json:"tool_call_id,omitempty"` // For tool response messages + KnownTo []string `json:"known_to,omitempty"` + Stats *ResponseStats `json:"-"` // Display-only, not persisted + hasContentParts bool // Flag to indicate which content type to marshal } // MarshalJSON implements custom JSON marshaling for RoleMsg @@ -183,13 +184,11 @@ func (m *RoleMsg) UnmarshalJSON(data []byte) error { } func (m *RoleMsg) ToText(i int) string { - // Convert content to string representation var contentStr string var imageIndicators []string if !m.hasContentParts { contentStr = m.Content } else { - // For structured content, collect text parts and image indicators var textParts []string for _, part := range m.ContentParts { switch p := part.(type) { @@ -198,7 +197,6 @@ func (m *RoleMsg) ToText(i int) string { textParts = append(textParts, p.Text) } case ImageContentPart: - // Collect image indicator displayPath := p.Path if displayPath == "" { displayPath = "image" @@ -216,7 +214,6 @@ func (m *RoleMsg) ToText(i int) string { } } case "image_url": - // Handle unmarshaled image content var displayPath string if pathVal, pathExists := p["path"]; pathExists { if pathStr, isStr := pathVal.(string); isStr && pathStr != "" { @@ -233,23 +230,20 @@ func (m *RoleMsg) ToText(i int) string { } contentStr = strings.Join(textParts, " ") + " " } - // check if already has role annotation (/completion makes them) - // in that case remove it, and then add to icon - // since icon and content are separated by \n contentStr, _ = strings.CutPrefix(contentStr, m.Role+":") - // if !strings.HasPrefix(contentStr, m.Role+":") { icon := fmt.Sprintf("(%d) <%s>: ", i, m.Role) - // } - // Build final message with image indicators before text var finalContent strings.Builder if len(imageIndicators) > 0 { - // Add each image indicator on its own line for _, indicator := range imageIndicators { finalContent.WriteString(indicator) finalContent.WriteString("\n") } } finalContent.WriteString(contentStr) + if m.Stats != nil { + finalContent.WriteString(fmt.Sprintf("\n[gray::i][%d tok, %.1fs, %.1f t/s][-:-:-]", + m.Stats.Tokens, m.Stats.Duration, m.Stats.TokensPerSec)) + } textMsg := fmt.Sprintf("[-:-:b]%s[-:-:-]\n%s\n", icon, finalContent.String()) return strings.ReplaceAll(textMsg, "\n\n", "\n") } @@ -331,6 +325,7 @@ func (m *RoleMsg) Copy() RoleMsg { ContentParts: m.ContentParts, ToolCallID: m.ToolCallID, KnownTo: m.KnownTo, + Stats: m.Stats, hasContentParts: m.hasContentParts, } } @@ -643,6 +638,12 @@ func (lcp *LCPModels) ListModels() []string { return resp } +type ResponseStats struct { + Tokens int + Duration float64 + TokensPerSec float64 +} + type ChatRoundReq struct { UserMsg string Role string |