CompletionBatchingRequestSupport (#220)
* completionBatchingRequestSupport * lint fix * fix Run test fail * fix TestClientReturnsRequestBuilderErrors fail * fix Codecov check * ignore TestClientReturnsRequestBuilderErrors lint * fix lint again * lint again*2 * replace checkPromptType implementation * remove nil check --------- Co-authored-by: W <825708370@qq.com>
This commit is contained in:
@@ -9,6 +9,7 @@ import (
|
||||
var (
|
||||
ErrCompletionUnsupportedModel = errors.New("this model is not supported with this method, please use CreateChatCompletion client method instead") //nolint:lll
|
||||
ErrCompletionStreamNotSupported = errors.New("streaming is not supported with this method, please use CreateCompletionStream") //nolint:lll
|
||||
ErrCompletionRequestPromptTypeNotSupported = errors.New("the type of CompletionRequest.Promp only supports string and []string") //nolint:lll
|
||||
)
|
||||
|
||||
// GPT3 Defines the models provided by OpenAI to use when generating
|
||||
@@ -77,10 +78,16 @@ func checkEndpointSupportsModel(endpoint, model string) bool {
|
||||
return !disabledModelsForEndpoints[endpoint][model]
|
||||
}
|
||||
|
||||
func checkPromptType(prompt any) bool {
|
||||
_, isString := prompt.(string)
|
||||
_, isStringSlice := prompt.([]string)
|
||||
return isString || isStringSlice
|
||||
}
|
||||
|
||||
// CompletionRequest represents a request structure for completion API.
|
||||
type CompletionRequest struct {
|
||||
Model string `json:"model"`
|
||||
Prompt string `json:"prompt,omitempty"`
|
||||
Prompt any `json:"prompt,omitempty"`
|
||||
Suffix string `json:"suffix,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Temperature float32 `json:"temperature,omitempty"`
|
||||
@@ -143,6 +150,11 @@ func (c *Client) CreateCompletion(
|
||||
return
|
||||
}
|
||||
|
||||
if !checkPromptType(request.Prompt) {
|
||||
err = ErrCompletionRequestPromptTypeNotSupported
|
||||
return
|
||||
}
|
||||
|
||||
req, err := c.requestBuilder.build(ctx, http.MethodPost, c.fullURL(urlSuffix), request)
|
||||
if err != nil {
|
||||
return
|
||||
|
||||
@@ -98,14 +98,14 @@ func handleCompletionEndpoint(w http.ResponseWriter, r *http.Request) {
|
||||
// generate a random string of length completionReq.Length
|
||||
completionStr := strings.Repeat("a", completionReq.MaxTokens)
|
||||
if completionReq.Echo {
|
||||
completionStr = completionReq.Prompt + completionStr
|
||||
completionStr = completionReq.Prompt.(string) + completionStr
|
||||
}
|
||||
res.Choices = append(res.Choices, CompletionChoice{
|
||||
Text: completionStr,
|
||||
Index: i,
|
||||
})
|
||||
}
|
||||
inputTokens := numTokens(completionReq.Prompt) * completionReq.N
|
||||
inputTokens := numTokens(completionReq.Prompt.(string)) * completionReq.N
|
||||
completionTokens := completionReq.MaxTokens * completionReq.N
|
||||
res.Usage = Usage{
|
||||
PromptTokens: inputTokens,
|
||||
|
||||
@@ -51,7 +51,7 @@ func TestClientReturnsRequestBuilderErrors(t *testing.T) {
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err = client.CreateCompletion(ctx, CompletionRequest{})
|
||||
_, err = client.CreateCompletion(ctx, CompletionRequest{Prompt: "testing"})
|
||||
if !errors.Is(err, errTestRequestBuilderFailed) {
|
||||
t.Fatalf("Did not return error when request builder failed: %v", err)
|
||||
}
|
||||
@@ -146,3 +146,27 @@ func TestClientReturnsRequestBuilderErrors(t *testing.T) {
|
||||
t.Fatalf("Did not return error when request builder failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReturnsRequestBuilderErrorsAddtion(t *testing.T) {
|
||||
var err error
|
||||
ts := test.NewTestServer().OpenAITestServer()
|
||||
ts.Start()
|
||||
defer ts.Close()
|
||||
|
||||
config := DefaultConfig(test.GetTestToken())
|
||||
config.BaseURL = ts.URL + "/v1"
|
||||
client := NewClientWithConfig(config)
|
||||
client.requestBuilder = &failingRequestBuilder{}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err = client.CreateCompletion(ctx, CompletionRequest{Prompt: 1})
|
||||
if !errors.Is(err, ErrCompletionRequestPromptTypeNotSupported) {
|
||||
t.Fatalf("Did not return error when request builder failed: %v", err)
|
||||
}
|
||||
|
||||
_, err = client.CreateCompletionStream(ctx, CompletionRequest{Prompt: 1})
|
||||
if !errors.Is(err, ErrCompletionRequestPromptTypeNotSupported) {
|
||||
t.Fatalf("Did not return error when request builder failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user