Skip to content

Commit

Permalink
add embeddings tests
Browse files Browse the repository at this point in the history
  • Loading branch information
sashabaranov committed Apr 8, 2023
1 parent 89219e3 commit 48481a8
Showing 1 changed file with 44 additions and 0 deletions.
44 changes: 44 additions & 0 deletions embeddings_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,14 @@ package openai_test

import (
. "github.com/sashabaranov/go-openai"
"github.com/sashabaranov/go-openai/internal/test"
"github.com/sashabaranov/go-openai/internal/test/checks"

"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"testing"
)

Expand Down Expand Up @@ -45,3 +49,43 @@ func TestEmbedding(t *testing.T) {
}
}
}

func TestEmbeddingModel(t *testing.T) {
var em EmbeddingModel
err := em.UnmarshalText([]byte("text-similarity-ada-001"))
checks.NoError(t, err, "Could not marshal embedding model")

if em != AdaSimilarity {
t.Errorf("Model is not equal to AdaSimilarity")
}

err = em.UnmarshalText([]byte("some-non-existent-model"))
checks.NoError(t, err, "Could not marshal embedding model")
if em != Unknown {
t.Errorf("Model is not equal to Unknown")
}
}

func TestEmbeddingEndpoint(t *testing.T) {
server := test.NewTestServer()
server.RegisterHandler(
"/v1/embeddings",
func(w http.ResponseWriter, r *http.Request) {
resBytes, _ := json.Marshal(EmbeddingResponse{})
fmt.Fprintln(w, string(resBytes))
},
)
// create the test server
var err error
ts := server.OpenAITestServer()
ts.Start()
defer ts.Close()

config := DefaultConfig(test.GetTestToken())
config.BaseURL = ts.URL + "/v1"
client := NewClientWithConfig(config)
ctx := context.Background()

_, err = client.CreateEmbeddings(ctx, EmbeddingRequest{})
checks.NoError(t, err, "CreateEmbeddings error")
}

0 comments on commit 48481a8

Please sign in to comment.