1
1
package llm
2
2
3
3
import (
4
+ "crypto/md5"
4
5
"encoding/base64"
5
6
"encoding/json"
7
+ "fmt"
6
8
"io"
7
9
"mime"
8
10
"net/http"
9
11
"os"
10
12
"path/filepath"
13
+ "strings"
11
14
)
12
15
13
16
///////////////////////////////////////////////////////////////////////////////
14
17
// TYPES
15
18
19
+ // General attachment metadata
16
20
type AttachmentMeta struct {
17
21
Id string `json:"id,omitempty"`
18
22
Filename string `json:"filename,omitempty"`
19
23
ExpiresAt uint64 `json:"expires_at,omitempty"`
20
24
Caption string `json:"transcript,omitempty"`
21
25
Data []byte `json:"data"`
26
+ Type string `json:"type"`
27
+ }
28
+
29
+ // OpenAI image metadata
30
+ type ImageMeta struct {
31
+ Url string `json:"url,omitempty"`
32
+ Data []byte `json:"b64_json,omitempty"`
33
+ Prompt string `json:"revised_prompt,omitempty"`
22
34
}
23
35
24
36
// Attachment for messages
25
37
type Attachment struct {
26
- meta AttachmentMeta
38
+ meta * AttachmentMeta
39
+ image * ImageMeta
27
40
}
28
41
29
42
const (
@@ -38,21 +51,30 @@ func NewAttachment() *Attachment {
38
51
return new (Attachment )
39
52
}
40
53
54
+ // NewAttachment with OpenAI image
55
+ func NewAttachmentWithImage (image * ImageMeta ) * Attachment {
56
+ return & Attachment {image : image }
57
+ }
58
+
41
59
// ReadAttachment returns an attachment from a reader object.
42
60
// It is the responsibility of the caller to close the reader.
43
- func ReadAttachment (r io.Reader ) (* Attachment , error ) {
44
- var filename string
61
+ func ReadAttachment (r io.Reader , mimetype ... string ) (* Attachment , error ) {
62
+ var filename , typ string
45
63
data , err := io .ReadAll (r )
46
64
if err != nil {
47
65
return nil , err
48
66
}
49
67
if f , ok := r .(* os.File ); ok {
50
68
filename = f .Name ()
51
69
}
70
+ if len (mimetype ) > 0 {
71
+ typ = mimetype [0 ]
72
+ }
52
73
return & Attachment {
53
- meta : AttachmentMeta {
74
+ meta : & AttachmentMeta {
54
75
Filename : filename ,
55
76
Data : data ,
77
+ Type : typ ,
56
78
},
57
79
}, nil
58
80
}
@@ -73,19 +95,27 @@ func (a *Attachment) MarshalJSON() ([]byte, error) {
73
95
Filename string `json:"filename,omitempty"`
74
96
Type string `json:"type"`
75
97
Bytes uint64 `json:"bytes"`
76
- Caption string `json:"transcript,omitempty"`
98
+ Hash string `json:"hash,omitempty"`
99
+ Caption string `json:"caption,omitempty"`
77
100
}
78
- j .Id = a .meta .Id
79
- j .Filename = a .meta .Filename
101
+
80
102
j .Type = a .Type ()
81
- j .Bytes = uint64 (len (a .meta .Data ))
82
- j .Caption = a .meta .Caption
103
+ j .Caption = a .Caption ()
104
+ j .Hash = a .Hash ()
105
+ j .Filename = a .Filename ()
106
+ if a .meta != nil {
107
+ j .Id = a .meta .Id
108
+ j .Bytes = uint64 (len (a .meta .Data ))
109
+ } else if a .image != nil {
110
+ j .Bytes = uint64 (len (a .image .Data ))
111
+ }
112
+
83
113
return json .Marshal (j )
84
114
}
85
115
86
116
// Stringify an attachment
87
117
func (a * Attachment ) String () string {
88
- data , err := json .MarshalIndent (a . meta , "" , " " )
118
+ data , err := json .MarshalIndent (a , "" , " " )
89
119
if err != nil {
90
120
return err .Error ()
91
121
}
@@ -95,41 +125,83 @@ func (a *Attachment) String() string {
95
125
////////////////////////////////////////////////////////////////////////////////
96
126
// PUBLIC METHODS
97
127
128
+ // Compute and print the MD5 hash
129
+ func (a * Attachment ) Hash () string {
130
+ hash := md5 .New ()
131
+ hash .Write (a .Data ())
132
+ return fmt .Sprintf ("%x" , hash .Sum (nil ))
133
+ }
134
+
135
+ // Write out attachment
136
+ func (a * Attachment ) Write (w io.Writer ) (int , error ) {
137
+ if a .meta != nil {
138
+ return w .Write (a .meta .Data )
139
+ }
140
+ if a .image != nil {
141
+ return w .Write (a .image .Data )
142
+ }
143
+ return 0 , io .EOF
144
+ }
145
+
98
146
// Return the filename of an attachment
99
147
func (a * Attachment ) Filename () string {
100
- return a .meta .Filename
148
+ if a .meta != nil && a .meta .Filename != "" {
149
+ return a .meta .Filename
150
+ }
151
+ // Obtain filename from MD5
152
+ if ext , err := mime .ExtensionsByType (a .Type ()); err == nil && len (ext ) > 0 {
153
+ return a .Hash () + ext [0 ]
154
+ }
155
+ return ""
101
156
}
102
157
103
158
// Return the raw attachment data
104
159
func (a * Attachment ) Data () []byte {
105
- return a .meta .Data
160
+ if a .meta != nil {
161
+ return a .meta .Data
162
+ }
163
+ if a .image != nil {
164
+ return a .image .Data
165
+ }
166
+ return nil
106
167
}
107
168
108
169
// Return the caption for the attachment
109
170
func (a * Attachment ) Caption () string {
110
- return a .meta .Caption
171
+ if a .meta != nil {
172
+ return strings .TrimSpace (a .meta .Caption )
173
+ }
174
+ if a .image != nil {
175
+ return strings .TrimSpace (a .image .Prompt )
176
+ }
177
+ return ""
111
178
}
112
179
113
180
// Return the mime media type for the attachment, based
114
181
// on the data and/or filename extension. Returns an empty string if
115
182
// there is no data or filename
116
183
func (a * Attachment ) Type () string {
184
+ // If there's a mimetype set, use this
185
+ if a .meta != nil && a .meta .Type != "" {
186
+ return a .meta .Type
187
+ }
188
+
117
189
// If there's no data or filename, return empty
118
- if len (a .meta . Data ) == 0 && a .meta . Filename == "" {
190
+ if len (a .Data ()) == 0 && a .Filename () == "" {
119
191
return ""
120
192
}
121
193
122
194
// Mimetype based on content
123
195
mimetype := defaultMimetype
124
- if len (a .meta . Data ) > 0 {
125
- mimetype = http .DetectContentType (a .meta . Data )
196
+ if len (a .Data () ) > 0 {
197
+ mimetype = http .DetectContentType (a .Data () )
126
198
if mimetype != defaultMimetype {
127
199
return mimetype
128
200
}
129
201
}
130
202
131
203
// Mimetype based on filename
132
- if a .meta .Filename != "" {
204
+ if a .meta != nil && a . meta .Filename != "" {
133
205
// Detect mimetype from extension
134
206
mimetype = mime .TypeByExtension (filepath .Ext (a .meta .Filename ))
135
207
}
@@ -139,24 +211,27 @@ func (a *Attachment) Type() string {
139
211
}
140
212
141
213
func (a * Attachment ) Url () string {
142
- return "data:" + a .Type () + ";base64," + base64 .StdEncoding .EncodeToString (a .meta . Data )
214
+ return "data:" + a .Type () + ";base64," + base64 .StdEncoding .EncodeToString (a .Data () )
143
215
}
144
216
145
217
// Streaming includes the ability to append data
146
218
func (a * Attachment ) Append (other * Attachment ) {
147
- if other .meta .Id != "" {
148
- a .meta .Id = other .meta .Id
149
- }
150
- if other .meta .Filename != "" {
151
- a .meta .Filename = other .meta .Filename
152
- }
153
- if other .meta .ExpiresAt != 0 {
154
- a .meta .ExpiresAt = other .meta .ExpiresAt
155
- }
156
- if other .meta .Caption != "" {
157
- a .meta .Caption += other .meta .Caption
158
- }
159
- if len (other .meta .Data ) > 0 {
160
- a .meta .Data = append (a .meta .Data , other .meta .Data ... )
219
+ if a .meta != nil {
220
+ if other .meta .Id != "" {
221
+ a .meta .Id = other .meta .Id
222
+ }
223
+ if other .meta .Filename != "" {
224
+ a .meta .Filename = other .meta .Filename
225
+ }
226
+ if other .meta .ExpiresAt != 0 {
227
+ a .meta .ExpiresAt = other .meta .ExpiresAt
228
+ }
229
+ if other .meta .Caption != "" {
230
+ a .meta .Caption += other .meta .Caption
231
+ }
232
+ if len (other .meta .Data ) > 0 {
233
+ a .meta .Data = append (a .meta .Data , other .meta .Data ... )
234
+ }
161
235
}
236
+ // TODO: Append for image
162
237
}
0 commit comments