152 lines
3.5 KiB
Go
152 lines
3.5 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"testing"
|
|
)
|
|
|
|
func TestMessageBuilder(t *testing.T) {
|
|
content := OpenAiMessages{}
|
|
|
|
err := content.AddSystem("Some prompt")
|
|
|
|
if err != nil {
|
|
t.Log(err)
|
|
t.FailNow()
|
|
}
|
|
|
|
if len(content.Messages) != 1 {
|
|
t.Logf("Expected length 1, got %d.\n", len(content.Messages))
|
|
t.FailNow()
|
|
}
|
|
}
|
|
|
|
func TestMessageBuilderImage(t *testing.T) {
|
|
content := OpenAiMessages{}
|
|
|
|
prompt := "some prompt"
|
|
imageTitle := "image.png"
|
|
data := []byte("some data")
|
|
|
|
content.AddSystem(prompt)
|
|
content.AddImage(imageTitle, data)
|
|
|
|
if len(content.Messages) != 2 {
|
|
t.Logf("Expected length 2, got %d.\n", len(content.Messages))
|
|
t.FailNow()
|
|
}
|
|
|
|
promptMessage, ok := content.Messages[0].(OpenAiTextMessage)
|
|
if !ok {
|
|
t.Logf("Expected text content message, got %T\n", content.Messages[0])
|
|
t.FailNow()
|
|
}
|
|
|
|
if promptMessage.Role != ROLE_SYSTEM {
|
|
t.Log("Prompt message role is incorrect.")
|
|
t.FailNow()
|
|
}
|
|
|
|
if promptMessage.Content != prompt {
|
|
t.Log("Prompt message content is incorrect.")
|
|
t.FailNow()
|
|
}
|
|
|
|
arrayContentMessage, ok := content.Messages[1].(OpenAiArrayMessage)
|
|
if !ok {
|
|
t.Logf("Expected text content message, got %T\n", content.Messages[1])
|
|
t.FailNow()
|
|
}
|
|
|
|
if arrayContentMessage.Role != ROLE_USER {
|
|
t.Log("Array content message role is incorrect.")
|
|
t.FailNow()
|
|
}
|
|
|
|
if len(arrayContentMessage.Content) != 1 {
|
|
t.Logf("Expected length 1, got %d.\n", len(arrayContentMessage.Content))
|
|
t.FailNow()
|
|
}
|
|
|
|
imageContent, ok := arrayContentMessage.Content[0].(OpenAiImage)
|
|
if !ok {
|
|
t.Logf("Expected text content message, got %T\n", arrayContentMessage.Content[0])
|
|
t.FailNow()
|
|
}
|
|
|
|
base64data := base64.StdEncoding.EncodeToString(data)
|
|
url := fmt.Sprintf("data:image/%s;base64,%s", "png", base64data)
|
|
|
|
if imageContent.ImageUrl.Url != url {
|
|
t.Logf("Expected %s, but got %s.\n", url, imageContent.ImageUrl.Url)
|
|
t.FailNow()
|
|
}
|
|
}
|
|
|
|
func TestFullImageRequest(t *testing.T) {
|
|
request, err := getCompletionsForImage("model", 0.1, "You are an assistant", "image.png", []byte("some data"))
|
|
if err != nil {
|
|
t.Log(request)
|
|
t.FailNow()
|
|
}
|
|
|
|
jsonData, err := json.Marshal(request)
|
|
if err != nil {
|
|
t.Log(err)
|
|
t.FailNow()
|
|
}
|
|
|
|
expectedJson := `{"model":"model","temperature":0.1,"response_format":{"type":"","json_schema":""},"messages":[{"role":"system","content":"You are an assistant"},{"role":"user","content":[{"type":"image_url","image_url":{"url":"data:image/png;base64,c29tZSBkYXRh"}}]}]}`
|
|
|
|
if string(jsonData) != expectedJson {
|
|
t.Logf("Expected:\n%s\n Got:\n%s\n", expectedJson, string(jsonData))
|
|
t.FailNow()
|
|
}
|
|
}
|
|
|
|
func TestResponse(t *testing.T) {
|
|
testResponse := `{"tags": ["tag1", "tag2"], "text": ["text1"], "links": []}`
|
|
buffer := bytes.NewReader([]byte(testResponse))
|
|
|
|
body := io.NopCloser(buffer)
|
|
|
|
client := OpenAiClient{
|
|
url: "http://localhost:1234",
|
|
apiKey: "some-key",
|
|
Do: func(_req *http.Request) (*http.Response, error) {
|
|
return &http.Response{Body: body}, nil
|
|
},
|
|
}
|
|
|
|
info, err := client.GetImageInfo("image.png", []byte("some data"))
|
|
if err != nil {
|
|
t.Log(err)
|
|
t.FailNow()
|
|
}
|
|
|
|
if len(info.Tags) != 2 || len(info.Text) != 1 || len(info.Links) != 0 {
|
|
t.Logf("Some lengths are wrong.\nTags: %d\nText: %d\nLinks: %d\n", len(info.Tags), len(info.Text), len(info.Links))
|
|
t.FailNow()
|
|
}
|
|
|
|
if info.Tags[0] != "tag1" {
|
|
t.Log("0th tag is wrong.")
|
|
t.FailNow()
|
|
}
|
|
|
|
if info.Tags[1] != "tag2" {
|
|
t.Log("1th tag is wrong.")
|
|
t.FailNow()
|
|
}
|
|
|
|
if info.Text[0] != "text1" {
|
|
t.Log("0th text is wrong.")
|
|
t.FailNow()
|
|
}
|
|
}
|