Files
go-openai/moderation_test.go
sashabaranov 1eb5d625f8 Better configuration (#79)
* Configurable Transport (#75)

* new functions to allow HTTPClient configuration

* updated go.mod for testing from remote

* updated go.mod for remote testing

* revert go.mod replace directives

* Fixed NewOrgClientWithTransport comment

* Make client fully configurable

* make empty messages limit configurable #70 #71

* make auth token private in config

* add docs

* lint

---------

Co-authored-by: Michael Fox <m.will.fox@gmail.com>
2023-02-21 00:16:44 +04:00

103 lines
2.8 KiB
Go

package gogpt_test
import (
. "github.com/sashabaranov/go-gpt3"
"github.com/sashabaranov/go-gpt3/internal/test"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"testing"
"time"
)
// TestModeration Tests the moderations endpoint of the API using the mocked server.
func TestModerations(t *testing.T) {
server := test.NewTestServer()
server.RegisterHandler("/v1/moderations", handleModerationEndpoint)
// 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()
// create an edit request
model := "text-moderation-stable"
moderationReq := ModerationRequest{
Model: &model,
Input: "I want to kill them.",
}
_, err = client.Moderations(ctx, moderationReq)
if err != nil {
t.Fatalf("Moderation error: %v", err)
}
}
// handleModerationEndpoint Handles the moderation endpoint by the test server.
func handleModerationEndpoint(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 moderationReq ModerationRequest
if moderationReq, err = getModerationBody(r); err != nil {
http.Error(w, "could not read request", http.StatusInternalServerError)
return
}
resCat := ResultCategories{}
resCatScore := ResultCategoryScores{}
switch {
case strings.Contains(moderationReq.Input, "kill"):
resCat = ResultCategories{Violence: true}
resCatScore = ResultCategoryScores{Violence: 1}
case strings.Contains(moderationReq.Input, "hate"):
resCat = ResultCategories{Hate: true}
resCatScore = ResultCategoryScores{Hate: 1}
case strings.Contains(moderationReq.Input, "suicide"):
resCat = ResultCategories{SelfHarm: true}
resCatScore = ResultCategoryScores{SelfHarm: 1}
case strings.Contains(moderationReq.Input, "porn"):
resCat = ResultCategories{Sexual: true}
resCatScore = ResultCategoryScores{Sexual: 1}
}
result := Result{Categories: resCat, CategoryScores: resCatScore, Flagged: true}
res := ModerationResponse{
ID: strconv.Itoa(int(time.Now().Unix())),
Model: *moderationReq.Model,
}
res.Results = append(res.Results, result)
resBytes, _ = json.Marshal(res)
fmt.Fprintln(w, string(resBytes))
}
// getModerationBody Returns the body of the request to do a moderation.
func getModerationBody(r *http.Request) (ModerationRequest, error) {
moderation := ModerationRequest{}
// read the request body
reqBody, err := io.ReadAll(r.Body)
if err != nil {
return ModerationRequest{}, err
}
err = json.Unmarshal(reqBody, &moderation)
if err != nil {
return ModerationRequest{}, err
}
return moderation, nil
}