8e3a04664e
* added NoError check * corrected NoError * has error checks * replace more checks * Used checks test helper * Used checks test helper * remove duplicate import * fixed lint issues regarding length of messages --------- Co-authored-by: Rex Posadas <rposadas@redwoodlogistics.com>
48 lines
1.1 KiB
Go
48 lines
1.1 KiB
Go
package openai_test
|
|
|
|
import (
|
|
. "github.com/sashabaranov/go-openai"
|
|
"github.com/sashabaranov/go-openai/internal/test/checks"
|
|
|
|
"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)
|
|
checks.NoError(t, err, "Could not marshal embedding request")
|
|
if !bytes.Contains(marshaled, []byte(`"model":"`+model.String()+`"`)) {
|
|
t.Fatalf("Expected embedding request to contain model field")
|
|
}
|
|
}
|
|
}
|