Add OpenAI Mock Server (#31)
* add constants for completions, refactor usage, add test server Signed-off-by: Oleg <97077423+RobotSail@users.noreply.github.com> * append v1 endpoint to test Signed-off-by: Oleg <97077423+RobotSail@users.noreply.github.com> * add makefile for easy targets Signed-off-by: Oleg <97077423+RobotSail@users.noreply.github.com> * lint files & add linter Signed-off-by: Oleg <97077423+RobotSail@users.noreply.github.com> * disable real API tests in short mode Signed-off-by: Oleg <97077423+RobotSail@users.noreply.github.com> Signed-off-by: Oleg <97077423+RobotSail@users.noreply.github.com>
This commit is contained in:
Oleg
GitHub
235
api_test.go
235
api_test.go
@@ -1,14 +1,30 @@
|
||||
package gogpt
|
||||
package gogpt_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
. "github.com/sashabaranov/go-gpt3"
|
||||
)
|
||||
|
||||
const (
|
||||
testAPIToken = "this-is-my-secure-token-do-not-steal!!"
|
||||
)
|
||||
|
||||
func TestAPI(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("skipping test in short mode")
|
||||
}
|
||||
tokenBytes, err := ioutil.ReadFile(".openai-token")
|
||||
if err != nil {
|
||||
t.Fatalf("Could not load auth token from .openai-token file")
|
||||
@@ -38,16 +54,6 @@ func TestAPI(t *testing.T) {
|
||||
}
|
||||
} // else skip
|
||||
|
||||
req := CompletionRequest{
|
||||
MaxTokens: 5,
|
||||
Model: "ada",
|
||||
}
|
||||
req.Prompt = "Lorem ipsum"
|
||||
_, err = c.CreateCompletion(ctx, req)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateCompletion error: %v", err)
|
||||
}
|
||||
|
||||
searchReq := SearchRequest{
|
||||
Documents: []string{"White House", "hospital", "school"},
|
||||
Query: "the president",
|
||||
@@ -70,6 +76,60 @@ func TestAPI(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestCompletions Tests the completions endpoint of the API using the mocked server.
|
||||
func TestCompletions(t *testing.T) {
|
||||
// create the test server
|
||||
var err error
|
||||
ts := OpenAITestServer()
|
||||
ts.Start()
|
||||
defer ts.Close()
|
||||
|
||||
client := NewClient(testAPIToken)
|
||||
ctx := context.Background()
|
||||
client.BaseURL = ts.URL + "/v1"
|
||||
|
||||
req := CompletionRequest{
|
||||
MaxTokens: 5,
|
||||
Model: "ada",
|
||||
}
|
||||
req.Prompt = "Lorem ipsum"
|
||||
_, err = client.CreateCompletion(ctx, req)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateCompletion error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEdits Tests the edits endpoint of the API using the mocked server.
|
||||
func TestEdits(t *testing.T) {
|
||||
// create the test server
|
||||
var err error
|
||||
ts := OpenAITestServer()
|
||||
ts.Start()
|
||||
defer ts.Close()
|
||||
|
||||
client := NewClient(testAPIToken)
|
||||
ctx := context.Background()
|
||||
client.BaseURL = ts.URL + "/v1"
|
||||
|
||||
// create an edit request
|
||||
model := "ada"
|
||||
editReq := EditsRequest{
|
||||
Model: &model,
|
||||
Input: "Lorem ipsum dolor sit amet, consectetur adipiscing elit, " +
|
||||
"sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim" +
|
||||
" ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip" +
|
||||
" ex ea commodo consequat. Duis aute irure dolor in reprehe",
|
||||
Instruction: "test instruction",
|
||||
N: 3,
|
||||
}
|
||||
response, err := client.Edits(ctx, editReq)
|
||||
if err != nil {
|
||||
t.Fatalf("Edits error: %v", err)
|
||||
}
|
||||
if len(response.Choices) != editReq.N {
|
||||
t.Fatalf("edits does not properly return the correct number of choices")
|
||||
}
|
||||
}
|
||||
func TestEmbedding(t *testing.T) {
|
||||
embeddedModels := []EmbeddingModel{
|
||||
AdaSimilarity,
|
||||
@@ -108,3 +168,156 @@ func TestEmbedding(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getEditBody Returns the body of the request to create an edit.
|
||||
func getEditBody(r *http.Request) (EditsRequest, error) {
|
||||
edit := EditsRequest{}
|
||||
// read the request body
|
||||
reqBody, err := ioutil.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
return EditsRequest{}, err
|
||||
}
|
||||
err = json.Unmarshal(reqBody, &edit)
|
||||
if err != nil {
|
||||
return EditsRequest{}, err
|
||||
}
|
||||
return edit, nil
|
||||
}
|
||||
|
||||
// handleEditEndpoint Handles the edit endpoint by the test server.
|
||||
func handleEditEndpoint(w http.ResponseWriter, r *http.Request) {
|
||||
var err error
|
||||
var resBytes []byte
|
||||
|
||||
// edits only accepts POST requests
|
||||
if r.Method != "POST" {
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
}
|
||||
var editReq EditsRequest
|
||||
editReq, err = getEditBody(r)
|
||||
if err != nil {
|
||||
http.Error(w, "could not read request", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
// create a response
|
||||
res := EditsResponse{
|
||||
Object: "test-object",
|
||||
Created: uint64(time.Now().Unix()),
|
||||
}
|
||||
// edit and calculate token usage
|
||||
editString := "edited by mocked OpenAI server :)"
|
||||
inputTokens := numTokens(editReq.Input+editReq.Instruction) * editReq.N
|
||||
completionTokens := int(float32(len(editString))/4) * editReq.N
|
||||
for i := 0; i < editReq.N; i++ {
|
||||
// instruction will be hidden and only seen by OpenAI
|
||||
res.Choices = append(res.Choices, EditsChoice{
|
||||
Text: editReq.Input + editString,
|
||||
Index: i,
|
||||
})
|
||||
}
|
||||
res.Usage = Usage{
|
||||
PromptTokens: inputTokens,
|
||||
CompletionTokens: completionTokens,
|
||||
TotalTokens: inputTokens + completionTokens,
|
||||
}
|
||||
resBytes, _ = json.Marshal(res)
|
||||
fmt.Fprint(w, string(resBytes))
|
||||
}
|
||||
|
||||
// handleCompletionEndpoint Handles the completion endpoint by the test server.
|
||||
func handleCompletionEndpoint(w http.ResponseWriter, r *http.Request) {
|
||||
var err error
|
||||
var resBytes []byte
|
||||
|
||||
// completions only accepts POST requests
|
||||
if r.Method != "POST" {
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
}
|
||||
var completionReq CompletionRequest
|
||||
if completionReq, err = getCompletionBody(r); err != nil {
|
||||
http.Error(w, "could not read request", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
res := CompletionResponse{
|
||||
ID: strconv.Itoa(int(time.Now().Unix())),
|
||||
Object: "test-object",
|
||||
Created: uint64(time.Now().Unix()),
|
||||
// would be nice to validate Model during testing, but
|
||||
// this may not be possible with how much upkeep
|
||||
// would be required / wouldn't make much sense
|
||||
Model: completionReq.Model,
|
||||
}
|
||||
// create completions
|
||||
for i := 0; i < completionReq.N; i++ {
|
||||
// generate a random string of length completionReq.Length
|
||||
completionStr := strings.Repeat("a", completionReq.MaxTokens)
|
||||
if completionReq.Echo {
|
||||
completionStr = completionReq.Prompt + completionStr
|
||||
}
|
||||
res.Choices = append(res.Choices, CompletionChoice{
|
||||
Text: completionStr,
|
||||
Index: i,
|
||||
})
|
||||
}
|
||||
inputTokens := numTokens(completionReq.Prompt) * completionReq.N
|
||||
completionTokens := completionReq.MaxTokens * completionReq.N
|
||||
res.Usage = Usage{
|
||||
PromptTokens: inputTokens,
|
||||
CompletionTokens: completionTokens,
|
||||
TotalTokens: inputTokens + completionTokens,
|
||||
}
|
||||
resBytes, _ = json.Marshal(res)
|
||||
fmt.Fprintln(w, string(resBytes))
|
||||
}
|
||||
|
||||
// getCompletionBody Returns the body of the request to create a completion.
|
||||
func getCompletionBody(r *http.Request) (CompletionRequest, error) {
|
||||
completion := CompletionRequest{}
|
||||
// read the request body
|
||||
reqBody, err := ioutil.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
return CompletionRequest{}, err
|
||||
}
|
||||
err = json.Unmarshal(reqBody, &completion)
|
||||
if err != nil {
|
||||
return CompletionRequest{}, err
|
||||
}
|
||||
return completion, nil
|
||||
}
|
||||
|
||||
// numTokens Returns the number of GPT-3 encoded tokens in the given text.
|
||||
// This function approximates based on the rule of thumb stated by OpenAI:
|
||||
// https://beta.openai.com/tokenizer
|
||||
//
|
||||
// TODO: implement an actual tokenizer for GPT-3 and Codex (once available)
|
||||
func numTokens(s string) int {
|
||||
return int(float32(len(s)) / 4)
|
||||
}
|
||||
|
||||
// OpenAITestServer Creates a mocked OpenAI server which can pretend to handle requests during testing.
|
||||
func OpenAITestServer() *httptest.Server {
|
||||
return httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
log.Printf("received request at path %q\n", r.URL.Path)
|
||||
|
||||
// check auth
|
||||
if r.Header.Get("Authorization") != "Bearer "+testAPIToken {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
|
||||
// OPTIMIZE: create separate handler functions for these
|
||||
switch r.URL.Path {
|
||||
case "/v1/edits":
|
||||
handleEditEndpoint(w, r)
|
||||
return
|
||||
case "/v1/completions":
|
||||
handleCompletionEndpoint(w, r)
|
||||
return
|
||||
// TODO: implement the other endpoints
|
||||
default:
|
||||
// the endpoint doesn't exist
|
||||
http.Error(w, "the resource path doesn't exist", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user