195 lines
7.8 KiB
Go
195 lines
7.8 KiB
Go
package agent
|
||
|
||
import (
|
||
"archive/zip"
|
||
"bytes"
|
||
"context"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"image"
|
||
"image/png"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"strings"
|
||
"testing"
|
||
)
|
||
|
||
func uploadFixture(t *testing.T, s *Service, id, name string, data []byte) Attachment {
|
||
t.Helper()
|
||
a, err := s.UploadAttachment(context.Background(), id, AttachmentInput{Name: name, Data: base64.StdEncoding.EncodeToString(data)})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return a
|
||
}
|
||
|
||
func officeFixture(t *testing.T, files map[string]string) []byte {
|
||
t.Helper()
|
||
var buf bytes.Buffer
|
||
z := zip.NewWriter(&buf)
|
||
for name, text := range files {
|
||
w, err := z.Create(name)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, err = w.Write([]byte(text)); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
}
|
||
if err := z.Close(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return buf.Bytes()
|
||
}
|
||
|
||
func pdfFixture() []byte {
|
||
stream := "BT /F1 12 Tf 40 70 Td (Budget 3500 quantity 500) Tj ET"
|
||
objects := []string{"<< /Type /Catalog /Pages 2 0 R >>", "<< /Type /Pages /Kids [3 0 R] /Count 1 >>", "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 300 144] /Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>", "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", fmt.Sprintf("<< /Length %d >>\nstream\n%s\nendstream", len(stream), stream)}
|
||
var buf bytes.Buffer
|
||
buf.WriteString("%PDF-1.4\n")
|
||
var offsets []int
|
||
for i, obj := range objects {
|
||
offsets = append(offsets, buf.Len())
|
||
fmt.Fprintf(&buf, "%d 0 obj\n%s\nendobj\n", i+1, obj)
|
||
}
|
||
xref := buf.Len()
|
||
fmt.Fprintf(&buf, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1)
|
||
for _, offset := range offsets {
|
||
fmt.Fprintf(&buf, "%010d 00000 n \n", offset)
|
||
}
|
||
fmt.Fprintf(&buf, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref)
|
||
return buf.Bytes()
|
||
}
|
||
|
||
func TestAttachmentFormatsAndEncoding(t *testing.T) {
|
||
docx := officeFixture(t, map[string]string{"word/document.xml": `<w:document xmlns:w="urn:test"><w:body><w:p><w:r><w:t>预算3500元</w:t></w:r></w:p><w:p><w:r><w:t>数量500个</w:t></w:r></w:p></w:body></w:document>`})
|
||
xlsx := officeFixture(t, map[string]string{"xl/sharedStrings.xml": `<sst><si><t>数量</t></si></sst>`, "xl/worksheets/sheet1.xml": `<worksheet><sheetData><row><c r="A1" t="s"><v>0</v></c><c r="B1"><v>500</v></c><c r="C1" t="inlineStr"><is><t>预算3500</t></is></c></row></sheetData></worksheet>`})
|
||
for _, tc := range []struct {
|
||
name string
|
||
data []byte
|
||
want string
|
||
}{{"资料.txt", []byte("预算3500元"), "预算3500"}, {"utf16.txt", []byte{0xff, 0xfe, '5', 0, '0', 0, '0', 0}, "500"}, {"报价.docx", docx, "数量500"}, {"客户.xlsx", xlsx, "B1=500"}, {"报价.pdf", pdfFixture(), "3500"}} {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
result, err := extractAttachment(context.Background(), tc.name, tc.data)
|
||
if err != nil || !strings.Contains(result.Text, tc.want) {
|
||
t.Fatalf("extract %s: %q %v", tc.name, result.Text, err)
|
||
}
|
||
})
|
||
}
|
||
var picture bytes.Buffer
|
||
png.Encode(&picture, image.NewRGBA(image.Rect(0, 0, 2, 2)))
|
||
result, err := extractAttachment(context.Background(), "图.png", picture.Bytes())
|
||
if err != nil || result.Kind != "image" || result.MIMEType != "image/png" {
|
||
t.Fatalf("image: %+v %v", result, err)
|
||
}
|
||
for _, tc := range []struct {
|
||
name string
|
||
data []byte
|
||
}{{"a.exe", []byte("hello")}, {"a.txt", []byte{0, 1, 2}}, {"a.png", []byte("not image")}, {"a.pdf", []byte("invalid PDF")}, {"a.docx", []byte("not zip")}} {
|
||
if _, err := extractAttachment(context.Background(), tc.name, tc.data); err == nil {
|
||
t.Fatalf("invalid file accepted: %s", tc.name)
|
||
}
|
||
}
|
||
result, err = extractAttachment(context.Background(), "long.txt", []byte(strings.Repeat("内容", 40000)))
|
||
if err != nil || !result.Truncated || len(result.Text) > maxAttachmentText {
|
||
t.Fatal("long text limit", err)
|
||
}
|
||
}
|
||
|
||
func TestAttachmentsPersistAndStayWithinConversation(t *testing.T) {
|
||
s, secrets := testService(t)
|
||
c, _ := s.CreateConversation()
|
||
other, _ := s.CreateConversation()
|
||
a := uploadFixture(t, s, c.ID, "需求.md", []byte("预算3500，数量500，只起草不发送"))
|
||
c.DraftAttachments = []Attachment{a}
|
||
c.Draft = "分析这份资料"
|
||
if _, err := s.SaveConversation(c); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
restarted := newService(s.configPath, secrets)
|
||
loaded, err := restarted.GetConversation(c.ID)
|
||
if err != nil || len(loaded.DraftAttachments) != 1 || loaded.DraftAttachments[0].Name != a.Name {
|
||
t.Fatalf("restore: %+v %v", loaded, err)
|
||
}
|
||
other.DraftAttachments = []Attachment{a}
|
||
if _, err := s.SaveConversation(other); err == nil {
|
||
t.Fatal("cross-conversation file reference accepted")
|
||
}
|
||
forged := a
|
||
forged.Name = "forged"
|
||
forged.Size = -1
|
||
c.DraftAttachments = []Attachment{forged}
|
||
s.SaveConversation(c)
|
||
loaded, _ = s.GetConversation(c.ID)
|
||
if loaded.DraftAttachments[0].Name != a.Name || loaded.DraftAttachments[0].Size <= 0 {
|
||
t.Fatal("client metadata was trusted")
|
||
}
|
||
if err := s.RemoveAttachment(c.ID, a.ID); err == nil {
|
||
t.Fatal("removed referenced draft")
|
||
}
|
||
c.DraftAttachments = nil
|
||
c.Messages = []ConversationMessage{{ID: "u", Role: "user", Content: "分析附件", Attachments: []Attachment{a}}, {ID: "a", Role: "assistant", Content: "已读取", Status: "complete"}}
|
||
s.SaveConversation(c)
|
||
if err := s.RemoveAttachment(c.ID, a.ID); err == nil {
|
||
t.Fatal("removed sent file")
|
||
}
|
||
message, err := s.messageForModel(c.ID, c.Messages[0])
|
||
if err != nil || !strings.Contains(fmt.Sprint(message.Content), "预算3500") {
|
||
t.Fatalf("missing file context: %+v %v", message, err)
|
||
}
|
||
source, err := s.memorySource(context.Background(), c.ID, c.Messages)
|
||
if err != nil || !strings.Contains(source, "预算3500") {
|
||
t.Fatal("file missing in memory source", err)
|
||
}
|
||
dir, _ := s.attachmentDir(c.ID, "")
|
||
if err := s.DeleteConversation(c.ID); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, err := os.Stat(dir); !os.IsNotExist(err) {
|
||
t.Fatal("conversation deletion left attachments")
|
||
}
|
||
if _, err := s.UploadAttachment(context.Background(), "../escape", AttachmentInput{Name: "safe.txt", Data: base64.StdEncoding.EncodeToString([]byte("text"))}); err == nil {
|
||
t.Fatal("invalid conversation accepted")
|
||
}
|
||
}
|
||
|
||
func TestChatStreamReceivesAttachmentContentAndImageParts(t *testing.T) {
|
||
s, _ := testService(t)
|
||
c, _ := s.CreateConversation()
|
||
a := uploadFixture(t, s, c.ID, "budget.csv", []byte("budget,quantity\n3500,500"))
|
||
var received []modelMessage
|
||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
var body struct {
|
||
Messages []modelMessage `json:"messages"`
|
||
}
|
||
json.NewDecoder(r.Body).Decode(&body)
|
||
received = body.Messages
|
||
w.Header().Set("Content-Type", "text/event-stream")
|
||
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"预算3500，数量500\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")
|
||
}))
|
||
defer server.Close()
|
||
s.SaveModelConfig(ModelConfigInput{Config: testConfig(server.URL), APIKey: "key"})
|
||
result, err := s.ChatStreamWithFiles(context.Background(), "files-test", c.ID, "分析预算", []string{a.ID})
|
||
if err != nil || !strings.Contains(result["content"], "3500") || !strings.Contains(fmt.Sprint(received[len(received)-1].Content), "3500,500") {
|
||
t.Fatalf("model file context: %v %v", received, err)
|
||
}
|
||
var picture bytes.Buffer
|
||
png.Encode(&picture, image.NewRGBA(image.Rect(0, 0, 2, 2)))
|
||
img := uploadFixture(t, s, c.ID, "图.png", picture.Bytes())
|
||
message, err := s.messageForModel(c.ID, ConversationMessage{Role: "user", Content: "看图", Attachments: []Attachment{img}})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
raw, _ := json.Marshal(message)
|
||
if !strings.Contains(string(raw), "data:image/png;base64,") {
|
||
t.Fatal("image not adapted for model")
|
||
}
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
cancel()
|
||
if _, err := s.UploadAttachment(ctx, c.ID, AttachmentInput{Name: "cancel.txt", Data: base64.StdEncoding.EncodeToString([]byte("cancel"))}); err == nil {
|
||
t.Fatal("cancelled upload continued")
|
||
}
|
||
}
|