49 lines
1.1 KiB
Go
49 lines
1.1 KiB
Go
package gogpt_test
|
|
|
|
import (
|
|
. "github.com/sashabaranov/go-openai"
|
|
|
|
"bytes"
|
|
"encoding/json"
|
|
"testing"
|
|
)
|
|
|
|
func TestEmbedding(t *testing.T) {
|
|
embeddedModels := []EmbeddingModel{
|
|
AdaSimilarity,
|
|
BabbageSimilarity,
|
|
CurieSimilarity,
|
|
DavinciSimilarity,
|
|
AdaSearchDocument,
|
|
AdaSearchQuery,
|
|
BabbageSearchDocument,
|
|
BabbageSearchQuery,
|
|
CurieSearchDocument,
|
|
CurieSearchQuery,
|
|
DavinciSearchDocument,
|
|
DavinciSearchQuery,
|
|
AdaCodeSearchCode,
|
|
AdaCodeSearchText,
|
|
BabbageCodeSearchCode,
|
|
BabbageCodeSearchText,
|
|
}
|
|
for _, model := range embeddedModels {
|
|
embeddingReq := EmbeddingRequest{
|
|
Input: []string{
|
|
"The food was delicious and the waiter",
|
|
"Other examples of embedding request",
|
|
},
|
|
Model: model,
|
|
}
|
|
// marshal embeddingReq to JSON and confirm that the model field equals
|
|
// the AdaSearchQuery type
|
|
marshaled, err := json.Marshal(embeddingReq)
|
|
if err != nil {
|
|
t.Fatalf("Could not marshal embedding request: %v", err)
|
|
}
|
|
if !bytes.Contains(marshaled, []byte(`"model":"`+model.String()+`"`)) {
|
|
t.Fatalf("Expected embedding request to contain model field")
|
|
}
|
|
}
|
|
}
|