diff options
Diffstat (limited to 'models')
| -rw-r--r-- | models/models.go | 134 | ||||
| -rw-r--r-- | models/models_test.go | 161 |
2 files changed, 16 insertions, 279 deletions
diff --git a/models/models.go b/models/models.go index 7ac9ec9..200d898 100644 --- a/models/models.go +++ b/models/models.go @@ -5,22 +5,9 @@ import ( "encoding/json" "fmt" "os" - "path/filepath" "strings" ) -var ( - // imageBaseDir is the base directory for displaying image paths. - // If set, image paths will be shown relative to this directory. - imageBaseDir = "" -) - -// SetImageBaseDir sets the base directory for displaying image paths. -// If dir is empty, full paths will be shown. -func SetImageBaseDir(dir string) { - imageBaseDir = dir -} - type FuncCall struct { ID string `json:"id,omitempty"` Name string `json:"name"` @@ -119,14 +106,14 @@ type RoleMsg struct { IsShellCommand bool `json:"is_shell_command,omitempty"` // True for shell command outputs (always shown) KnownTo []string `json:"known_to,omitempty"` Stats *ResponseStats `json:"stats"` - hasContentParts bool // Flag to indicate which content type to marshal + HasContentParts bool // Flag to indicate which content type to marshal } // MarshalJSON implements custom JSON marshaling for RoleMsg // //nolint:gocritic func (m RoleMsg) MarshalJSON() ([]byte, error) { - if m.hasContentParts { + if m.HasContentParts { // Use structured content format aux := struct { Role string `json:"role"` @@ -189,7 +176,7 @@ func (m *RoleMsg) UnmarshalJSON(data []byte) error { m.IsShellCommand = structured.IsShellCommand m.KnownTo = structured.KnownTo m.Stats = structured.Stats - m.hasContentParts = true + m.HasContentParts = true return nil } @@ -213,77 +200,13 @@ func (m *RoleMsg) UnmarshalJSON(data []byte) error { m.IsShellCommand = simple.IsShellCommand m.KnownTo = simple.KnownTo m.Stats = simple.Stats - m.hasContentParts = false + m.HasContentParts = false return nil } -func (m *RoleMsg) ToText(i int) string { - var contentStr string - var imageIndicators []string - if !m.hasContentParts { - contentStr = m.Content - } else { - var textParts []string - for _, part := range m.ContentParts { - switch p := part.(type) { - case TextContentPart: - if p.Type == "text" { - textParts = append(textParts, p.Text) - } - case ImageContentPart: - displayPath := p.Path - if displayPath == "" { - displayPath = "image" - } else { - displayPath = extractDisplayPath(displayPath) - } - imageIndicators = append(imageIndicators, fmt.Sprintf("[orange::i][image: %s][-:-:-]", displayPath)) - case map[string]any: - if partType, exists := p["type"]; exists { - switch partType { - case "text": - if textVal, textExists := p["text"]; textExists { - if textStr, isStr := textVal.(string); isStr { - textParts = append(textParts, textStr) - } - } - case "image_url": - var displayPath string - if pathVal, pathExists := p["path"]; pathExists { - if pathStr, isStr := pathVal.(string); isStr && pathStr != "" { - displayPath = extractDisplayPath(pathStr) - } - } - if displayPath == "" { - displayPath = "image" - } - imageIndicators = append(imageIndicators, fmt.Sprintf("[orange::i][image: %s][-:-:-]", displayPath)) - } - } - } - } - contentStr = strings.Join(textParts, " ") + " " - } - contentStr, _ = strings.CutPrefix(contentStr, m.Role+":") - icon := fmt.Sprintf("(%d) <%s>: ", i, m.Role) - var finalContent strings.Builder - if len(imageIndicators) > 0 { - for _, indicator := range imageIndicators { - finalContent.WriteString(indicator) - finalContent.WriteString("\n") - } - } - finalContent.WriteString(contentStr) - if m.Stats != nil { - fmt.Fprintf(&finalContent, "\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") -} - func (m *RoleMsg) ToPrompt() string { var contentStr string - if !m.hasContentParts { + if !m.HasContentParts { contentStr = m.Content } else { // For structured content, just take the text parts @@ -316,7 +239,7 @@ func NewRoleMsg(role, content string) RoleMsg { return RoleMsg{ Role: role, Content: content, - hasContentParts: false, + HasContentParts: false, } } @@ -325,7 +248,7 @@ func NewMultimodalMsg(role string, contentParts []any) RoleMsg { return RoleMsg{ Role: role, ContentParts: contentParts, - hasContentParts: true, + HasContentParts: true, } } @@ -334,7 +257,7 @@ func (m *RoleMsg) HasContent() bool { if m.Content != "" { return true } - if m.hasContentParts && len(m.ContentParts) > 0 { + if m.HasContentParts && len(m.ContentParts) > 0 { return true } return false @@ -342,7 +265,7 @@ func (m *RoleMsg) HasContent() bool { // IsContentParts returns true if the message uses structured content parts func (m *RoleMsg) IsContentParts() bool { - return m.hasContentParts + return m.HasContentParts } // GetContentParts returns the content parts of the message @@ -359,14 +282,14 @@ func (m *RoleMsg) Copy() RoleMsg { ToolCallID: m.ToolCallID, KnownTo: m.KnownTo, Stats: m.Stats, - hasContentParts: m.hasContentParts, + HasContentParts: m.HasContentParts, } } // GetText returns the text content of the message, handling both // simple Content and multimodal ContentParts formats. func (m *RoleMsg) GetText() string { - if !m.hasContentParts { + if !m.HasContentParts { return m.Content } var textParts []string @@ -395,7 +318,7 @@ func (m *RoleMsg) GetText() string { // ContentParts (multimodal), it updates the text parts while preserving // images. If not, it sets the simple Content field. func (m *RoleMsg) SetText(text string) { - if !m.hasContentParts { + if !m.HasContentParts { m.Content = text return } @@ -425,14 +348,14 @@ func (m *RoleMsg) SetText(text string) { // AddTextPart adds a text content part to the message func (m *RoleMsg) AddTextPart(text string) { - if !m.hasContentParts { + if !m.HasContentParts { // Convert to content parts format if m.Content != "" { m.ContentParts = []any{TextContentPart{Type: "text", Text: m.Content}} } else { m.ContentParts = []any{} } - m.hasContentParts = true + m.HasContentParts = true } textPart := TextContentPart{Type: "text", Text: text} m.ContentParts = append(m.ContentParts, textPart) @@ -440,14 +363,14 @@ func (m *RoleMsg) AddTextPart(text string) { // AddImagePart adds an image content part to the message func (m *RoleMsg) AddImagePart(imageURL, imagePath string) { - if !m.hasContentParts { + if !m.HasContentParts { // Convert to content parts format if m.Content != "" { m.ContentParts = []any{TextContentPart{Type: "text", Text: m.Content}} } else { m.ContentParts = []any{} } - m.hasContentParts = true + m.HasContentParts = true } imagePart := ImageContentPart{ Type: "image_url", @@ -491,31 +414,6 @@ func CreateImageURLFromPath(imagePath string) (string, error) { return fmt.Sprintf("data:%s;base64,%s", mimeType, encoded), nil } -// extractDisplayPath returns a path suitable for display, potentially relative to imageBaseDir -func extractDisplayPath(p string) string { - if p == "" { - return "" - } - - // If base directory is set, try to make path relative to it - if imageBaseDir != "" { - if rel, err := filepath.Rel(imageBaseDir, p); err == nil { - // Check if relative path doesn't start with ".." (meaning it's within base dir) - // If it starts with "..", we might still want to show it as relative - // but for now we show full path if it goes outside base dir - if !strings.HasPrefix(rel, "..") { - p = rel - } - } - } - - // Truncate long paths to last 60 characters if needed - if len(p) > 60 { - return "..." + p[len(p)-60:] - } - return p -} - type ChatBody struct { Model string `json:"model"` Stream bool `json:"stream"` diff --git a/models/models_test.go b/models/models_test.go deleted file mode 100644 index 3b6476a..0000000 --- a/models/models_test.go +++ /dev/null @@ -1,161 +0,0 @@ -package models -import ( - "strings" - "testing" -) -func TestRoleMsgToTextWithImages(t *testing.T) { - tests := []struct { - name string - msg RoleMsg - index int - expected string // substring to check - }{ - { - name: "text and image", - index: 0, - msg: func() RoleMsg { - msg := NewMultimodalMsg("user", []interface{}{}) - msg.AddTextPart("Look at this picture") - msg.AddImagePart("data:image/jpeg;base64,abc123", "/home/user/Pictures/cat.jpg") - return msg - }(), - expected: "[orange::i][image: /home/user/Pictures/cat.jpg][-:-:-]", - }, - { - name: "image only", - index: 1, - msg: func() RoleMsg { - msg := NewMultimodalMsg("user", []interface{}{}) - msg.AddImagePart("data:image/png;base64,xyz789", "/tmp/screenshot_20250217_123456.png") - return msg - }(), - expected: "[orange::i][image: /tmp/screenshot_20250217_123456.png][-:-:-]", - }, - { - name: "long filename truncated", - index: 2, - msg: func() RoleMsg { - msg := NewMultimodalMsg("user", []interface{}{}) - msg.AddTextPart("Check this") - msg.AddImagePart("data:image/jpeg;base64,foo", "/very/long/path/to/a/really_long_filename_that_exceeds_forty_characters.jpg") - return msg - }(), - expected: "[orange::i][image: .../to/a/really_long_filename_that_exceeds_forty_characters.jpg][-:-:-]", - }, - { - name: "multiple images", - index: 3, - msg: func() RoleMsg { - msg := NewMultimodalMsg("user", []interface{}{}) - msg.AddTextPart("Multiple images") - msg.AddImagePart("data:image/jpeg;base64,a", "/path/img1.jpg") - msg.AddImagePart("data:image/png;base64,b", "/path/img2.png") - return msg - }(), - expected: "[orange::i][image: /path/img1.jpg][-:-:-]\n[orange::i][image: /path/img2.png][-:-:-]", - }, - { - name: "old format without path", - index: 4, - msg: RoleMsg{ - Role: "user", - hasContentParts: true, - ContentParts: []interface{}{ - map[string]interface{}{ - "type": "image_url", - "image_url": map[string]interface{}{ - "url": "data:image/jpeg;base64,old", - }, - }, - }, - }, - expected: "[orange::i][image: image][-:-:-]", - }, - { - name: "old format with path", - index: 5, - msg: RoleMsg{ - Role: "user", - hasContentParts: true, - ContentParts: []interface{}{ - map[string]interface{}{ - "type": "image_url", - "path": "/old/path/photo.jpg", - "image_url": map[string]interface{}{ - "url": "data:image/jpeg;base64,old", - }, - }, - }, - }, - expected: "[orange::i][image: /old/path/photo.jpg][-:-:-]", - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := tt.msg.ToText(tt.index) - if !strings.Contains(result, tt.expected) { - t.Errorf("ToText() result does not contain expected indicator\ngot: %s\nwant substring: %s", result, tt.expected) - } - // Ensure the indicator appears before text content - if strings.Contains(tt.expected, "cat.jpg") && strings.Contains(result, "Look at this picture") { - indicatorPos := strings.Index(result, "[orange::i][image: /home/user/Pictures/cat.jpg][-:-:-]") - textPos := strings.Index(result, "Look at this picture") - if indicatorPos == -1 || textPos == -1 || indicatorPos >= textPos { - t.Errorf("image indicator should appear before text") - } - } - }) - } -} -func TestExtractDisplayPath(t *testing.T) { - // Save original base dir - originalBaseDir := imageBaseDir - defer func() { imageBaseDir = originalBaseDir }() - tests := []struct { - name string - baseDir string - path string - expected string - }{ - { - name: "no base dir shows full path", - baseDir: "", - path: "/home/user/images/cat.jpg", - expected: "/home/user/images/cat.jpg", - }, - { - name: "relative path within base dir", - baseDir: "/home/user", - path: "/home/user/images/cat.jpg", - expected: "images/cat.jpg", - }, - { - name: "path outside base dir shows full path", - baseDir: "/home/user", - path: "/tmp/test.jpg", - expected: "/tmp/test.jpg", - }, - { - name: "same directory", - baseDir: "/home/user/images", - path: "/home/user/images/cat.jpg", - expected: "cat.jpg", - }, - { - name: "long path truncated", - baseDir: "", - path: "/very/long/path/to/a/really_long_filename_that_exceeds_sixty_characters_limit_yes_it_is_very_long.jpg", - expected: "..._that_exceeds_sixty_characters_limit_yes_it_is_very_long.jpg", - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - imageBaseDir = tt.baseDir - result := extractDisplayPath(tt.path) - if result != tt.expected { - t.Errorf("extractDisplayPath(%q) with baseDir=%q = %q, want %q", - tt.path, tt.baseDir, result, tt.expected) - } - }) - } -} |