Skip to content

Commit 55cf945

Browse files
committed
Updated
1 parent 989fedb commit 55cf945

File tree

2 files changed

+36
-9
lines changed

2 files changed

+36
-9
lines changed

attachment.go

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package llm
22

33
import (
4+
"crypto/md5"
45
"encoding/base64"
56
"encoding/json"
7+
"fmt"
68
"io"
79
"mime"
810
"net/http"
@@ -88,14 +90,16 @@ func (a *Attachment) MarshalJSON() ([]byte, error) {
8890
Filename string `json:"filename,omitempty"`
8991
Type string `json:"type"`
9092
Bytes uint64 `json:"bytes"`
93+
Hash string `json:"hash,omitempty"`
9194
Caption string `json:"caption,omitempty"`
9295
}
9396

9497
j.Type = a.Type()
9598
j.Caption = a.Caption()
99+
j.Hash = a.Hash()
100+
j.Filename = a.Filename()
96101
if a.meta != nil {
97102
j.Id = a.meta.Id
98-
j.Filename = a.meta.Filename
99103
j.Bytes = uint64(len(a.meta.Data))
100104
} else if a.image != nil {
101105
j.Bytes = uint64(len(a.image.Data))
@@ -116,6 +120,13 @@ func (a *Attachment) String() string {
116120
////////////////////////////////////////////////////////////////////////////////
117121
// PUBLIC METHODS
118122

123+
// Compute and print the MD5 hash
124+
func (a *Attachment) Hash() string {
125+
hash := md5.New()
126+
hash.Write(a.Data())
127+
return fmt.Sprintf("%x", hash.Sum(nil))
128+
}
129+
119130
// Write out attachment
120131
func (a *Attachment) Write(w io.Writer) (int, error) {
121132
if a.meta != nil {
@@ -129,11 +140,14 @@ func (a *Attachment) Write(w io.Writer) (int, error) {
129140

130141
// Return the filename of an attachment
131142
func (a *Attachment) Filename() string {
132-
if a.meta != nil {
143+
if a.meta != nil && a.meta.Filename != "" {
133144
return a.meta.Filename
134-
} else {
135-
return ""
136145
}
146+
// Obtain filename from MD5
147+
if ext, err := mime.ExtensionsByType(a.Type()); err == nil && len(ext) > 0 {
148+
return a.Hash() + ext[0]
149+
}
150+
return ""
137151
}
138152

139153
// Return the raw attachment data

cmd/llm/complete.go

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -94,14 +94,27 @@ func (cmd *CompleteCmd) Run(globals *Globals) error {
9494
// Print the completion - text
9595
if cmd.NoStream {
9696
fmt.Println(completion.Text(0))
97-
} else {
98-
fmt.Println("")
9997
}
10098

101-
// Print the completion - attachments
99+
// Output completion attachments
102100
for i := 0; i < completion.Num(); i++ {
103-
if attachment := completion.Attachment(i); attachment != nil {
104-
fmt.Println(attachment)
101+
attachment := completion.Attachment(i)
102+
if attachment == nil {
103+
continue
104+
}
105+
if attachment.Filename() == "" {
106+
continue
107+
}
108+
f, err := os.Create(attachment.Filename())
109+
if err != nil {
110+
return err
111+
}
112+
defer f.Close()
113+
114+
if _, err := f.Write(attachment.Data()); err != nil {
115+
return err
116+
} else {
117+
fmt.Printf("%q written to %s\n", attachment.Caption(), attachment.Filename())
105118
}
106119
}
107120

0 commit comments

Comments
 (0)