Repository: sashabaranov/go-openai
Pull request #882 (2024-10-17 03:11:15)

add chatcompletion stream delta refusal and logprobs
**Describe the change**
Like #844, In ChatCompletionStream.delta the refusal and logprobs field is missing

**Provide OpenAI documentation link**
[Provide a relevant API doc from https://platform.openai.com/docs/api-reference](https://platform.openai.com/docs/api-reference/chat/streaming#chat/streaming-choices)

**Describe your solution**
Add Refusal and Logprobs Fields

**Tests**
n/a

**Additional context**
n/a

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/sashabaranov/go-openai/pull/882?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Alex+Baranov) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.04%. Comparing base [(`774fc9d`)](https://app.codecov.io/gh/sashabaranov/go-openai/commit/774fc9dd12ed60c10a9f9f03319ddb9cd5f8780c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Alex+Baranov) to head [(`520c829`)](https://app.codecov.io/gh/sashabaranov/go-openai/commit/520c829db6341ed0f4e5ef0d5eadcf9710b79ea4?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Alex+Baranov).
> Report is 64 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #882      +/-   ##
==========================================
+ Coverage   98.46%   99.04%   +0.58%     
==========================================
  Files          24       26       +2     
  Lines        1364     1462      +98     
==========================================
+ Hits         1343     1448     +105     
+ Misses         15        8       -7     
  Partials        6        6              
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/sashabaranov/go-openai/pull/882?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Alex+Baranov).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Alex+Baranov).


Patch:
diff --git a/assistant.go b/assistant.go
index 4c89c1b..8aab5bc 100644
--- a/assistant.go
+++ b/assistant.go
@@ -14,17 +14,20 @@ const (
 )
 
 type Assistant struct {
-	ID            string                 `json:"id"`
-	Object        string                 `json:"object"`
-	CreatedAt     int64                  `json:"created_at"`
-	Name          *string                `json:"name,omitempty"`
-	Description   *string                `json:"description,omitempty"`
-	Model         string                 `json:"model"`
-	Instructions  *string                `json:"instructions,omitempty"`
-	Tools         []AssistantTool        `json:"tools"`
-	FileIDs       []string               `json:"file_ids,omitempty"`
-	Metadata      map[string]any         `json:"metadata,omitempty"`
-	ToolResources *AssistantToolResource `json:"tool_resources,omitempty"`
+	ID             string                 `json:"id"`
+	Object         string                 `json:"object"`
+	CreatedAt      int64                  `json:"created_at"`
+	Name           *string                `json:"name,omitempty"`
+	Description    *string                `json:"description,omitempty"`
+	Model          string                 `json:"model"`
+	Instructions   *string                `json:"instructions,omitempty"`
+	Tools          []AssistantTool        `json:"tools"`
+	ToolResources  *AssistantToolResource `json:"tool_resources,omitempty"`
+	FileIDs        []string               `json:"file_ids,omitempty"` // Deprecated in v2
+	Metadata       map[string]any         `json:"metadata,omitempty"`
+	Temperature    *float32               `json:"temperature,omitempty"`
+	TopP           *float32               `json:"top_p,omitempty"`
+	ResponseFormat any                    `json:"response_format,omitempty"`
 
 	httpHeader
 }
diff --git a/chat_stream.go b/chat_stream.go
index f43d018..58b2651 100644
--- a/chat_stream.go
+++ b/chat_stream.go
@@ -10,13 +10,33 @@ type ChatCompletionStreamChoiceDelta struct {
 	Role         string        `json:"role,omitempty"`
 	FunctionCall *FunctionCall `json:"function_call,omitempty"`
 	ToolCalls    []ToolCall    `json:"tool_calls,omitempty"`
+	Refusal      string        `json:"refusal,omitempty"`
+}
+
+type ChatCompletionStreamChoiceLogprobs struct {
+	Content []ChatCompletionTokenLogprob `json:"content,omitempty"`
+	Refusal []ChatCompletionTokenLogprob `json:"refusal,omitempty"`
+}
+
+type ChatCompletionTokenLogprob struct {
+	Token       string                                 `json:"token"`
+	Bytes       []int64                                `json:"bytes,omitempty"`
+	Logprob     float64                                `json:"logprob,omitempty"`
+	TopLogprobs []ChatCompletionTokenLogprobTopLogprob `json:"top_logprobs"`
+}
+
+type ChatCompletionTokenLogprobTopLogprob struct {
+	Token   string  `json:"token"`
+	Bytes   []int64 `json:"bytes"`
+	Logprob float64 `json:"logprob"`
 }
 
 type ChatCompletionStreamChoice struct {
-	Index                int                             `json:"index"`
-	Delta                ChatCompletionStreamChoiceDelta `json:"delta"`
-	FinishReason         FinishReason                    `json:"finish_reason"`
-	ContentFilterResults ContentFilterResults            `json:"content_filter_results,omitempty"`
+	Index                int                                 `json:"index"`
+	Delta                ChatCompletionStreamChoiceDelta     `json:"delta"`
+	Logprobs             *ChatCompletionStreamChoiceLogprobs `json:"logprobs,omitempty"`
+	FinishReason         FinishReason                        `json:"finish_reason"`
+	ContentFilterResults ContentFilterResults                `json:"content_filter_results,omitempty"`
 }
 
 type PromptFilterResult struct {
diff --git a/client.go b/client.go
index 1e228a0..ed8595e 100644
--- a/client.go
+++ b/client.go
@@ -289,9 +289,6 @@ func (c *Client) handleErrorResp(resp *http.Response) error {
 	if err != nil {
 		return fmt.Errorf("error, reading response body: %w", err)
 	}
-	if !strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json") {
-		return fmt.Errorf("error, status code: %d, status: %s, body: %s", resp.StatusCode, resp.Status, body)
-	}
 	var errRes ErrorResponse
 	err = json.Unmarshal(body, &errRes)
 	if err != nil || errRes.Error == nil {
diff --git a/completion.go b/completion.go
index 84ef2ad..77ea8c3 100644
--- a/completion.go
+++ b/completion.go
@@ -161,7 +161,23 @@ func checkEndpointSupportsModel(endpoint, model string) bool {
 func checkPromptType(prompt any) bool {
 	_, isString := prompt.(string)
 	_, isStringSlice := prompt.([]string)
-	return isString || isStringSlice
+	if isString || isStringSlice {
+		return true
+	}
+
+	// check if it is prompt is []string hidden under []any
+	slice, isSlice := prompt.([]any)
+	if !isSlice {
+		return false
+	}
+
+	for _, item := range slice {
+		_, itemIsString := item.(string)
+		if !itemIsString {
+			return false
+		}
+	}
+	return true // all items in the slice are string, so it is []string
 }
 
 var unsupportedToolsForO1Models = map[ToolType]struct{}{
diff --git a/error.go b/error.go
index fc9e7cd..8a74bd5 100644
--- a/error.go
+++ b/error.go
@@ -104,7 +104,10 @@ func (e *APIError) UnmarshalJSON(data []byte) (err error) {
 }
 
 func (e *RequestError) Error() string {
-	return fmt.Sprintf("error, status code: %d, status: %s, message: %s", e.HTTPStatusCode, e.HTTPStatus, e.Err)
+	return fmt.Sprintf(
+		"error, status code: %d, status: %s, message: %s, body: %s",
+		e.HTTPStatusCode, e.HTTPStatus, e.Err, e.Body,
+	)
 }
 
 func (e *RequestError) Unwrap() error {


Test patch:
diff --git a/chat_stream_test.go b/chat_stream_test.go
index 2e7c99b..1468414 100644
--- a/chat_stream_test.go
+++ b/chat_stream_test.go
@@ -358,6 +358,271 @@ func TestCreateChatCompletionStreamRateLimitError(t *testing.T) {
 	t.Logf("%+v\n", apiErr)
 }
 
+func TestCreateChatCompletionStreamWithRefusal(t *testing.T) {
+	client, server, teardown := setupOpenAITestServer()
+	defer teardown()
+	server.RegisterHandler("/v1/chat/completions", func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+
+		dataBytes := []byte{}
+
+		//nolint:lll
+		dataBytes = append(dataBytes, []byte(`data: {"id":"1","object":"chat.completion.chunk","created":1729585728,"model":"gpt-4o-mini-2024-07-18","system_fingerprint":"fp_d9767fc5b9","choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"finish_reason":null}]}`)...)
+		dataBytes = append(dataBytes, []byte("\n\n")...)
+
+		//nolint:lll
+		dataBytes = append(dataBytes, []byte(`data: {"id":"2","object":"chat.completion.chunk","created":1729585728,"model":"gpt-4o-mini-2024-07-18","system_fingerprint":"fp_d9767fc5b9","choices":[{"index":0,"delta":{"refusal":"Hello"},"finish_reason":null}]}`)...)
+		dataBytes = append(dataBytes, []byte("\n\n")...)
+
+		//nolint:lll
+		dataBytes = append(dataBytes, []byte(`data: {"id":"3","object":"chat.completion.chunk","created":1729585728,"model":"gpt-4o-mini-2024-07-18","system_fingerprint":"fp_d9767fc5b9","choices":[{"index":0,"delta":{"refusal":" World"},"finish_reason":null}]}`)...)
+		dataBytes = append(dataBytes, []byte("\n\n")...)
+
+		//nolint:lll
+		dataBytes = append(dataBytes, []byte(`data: {"id":"4","object":"chat.completion.chunk","created":1729585728,"model":"gpt-4o-mini-2024-07-18","system_fingerprint":"fp_d9767fc5b9","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`)...)
+		dataBytes = append(dataBytes, []byte("\n\n")...)
+
+		dataBytes = append(dataBytes, []byte("data: [DONE]\n\n")...)
+
+		_, err := w.Write(dataBytes)
+		checks.NoError(t, err, "Write error")
+	})
+
+	stream, err := client.CreateChatCompletionStream(context.Background(), openai.ChatCompletionRequest{
+		MaxTokens: 2000,
+		Model:     openai.GPT4oMini20240718,
+		Messages: []openai.ChatCompletionMessage{
+			{
+				Role:    openai.ChatMessageRoleUser,
+				Content: "Hello!",
+			},
+		},
+		Stream: true,
+	})
+	checks.NoError(t, err, "CreateCompletionStream returned error")
+	defer stream.Close()
+
+	expectedResponses := []openai.ChatCompletionStreamResponse{
+		{
+			ID:                "1",
+			Object:            "chat.completion.chunk",
+			Created:           1729585728,
+			Model:             openai.GPT4oMini20240718,
+			SystemFingerprint: "fp_d9767fc5b9",
+			Choices: []openai.ChatCompletionStreamChoice{
+				{
+					Index: 0,
+					Delta: openai.ChatCompletionStreamChoiceDelta{},
+				},
+			},
+		},
+		{
+			ID:                "2",
+			Object:            "chat.completion.chunk",
+			Created:           1729585728,
+			Model:             openai.GPT4oMini20240718,
+			SystemFingerprint: "fp_d9767fc5b9",
+			Choices: []openai.ChatCompletionStreamChoice{
+				{
+					Index: 0,
+					Delta: openai.ChatCompletionStreamChoiceDelta{
+						Refusal: "Hello",
+					},
+				},
+			},
+		},
+		{
+			ID:                "3",
+			Object:            "chat.completion.chunk",
+			Created:           1729585728,
+			Model:             openai.GPT4oMini20240718,
+			SystemFingerprint: "fp_d9767fc5b9",
+			Choices: []openai.ChatCompletionStreamChoice{
+				{
+					Index: 0,
+					Delta: openai.ChatCompletionStreamChoiceDelta{
+						Refusal: " World",
+					},
+				},
+			},
+		},
+		{
+			ID:                "4",
+			Object:            "chat.completion.chunk",
+			Created:           1729585728,
+			Model:             openai.GPT4oMini20240718,
+			SystemFingerprint: "fp_d9767fc5b9",
+			Choices: []openai.ChatCompletionStreamChoice{
+				{
+					Index:        0,
+					FinishReason: "stop",
+				},
+			},
+		},
+	}
+
+	for ix, expectedResponse := range expectedResponses {
+		b, _ := json.Marshal(expectedResponse)
+		t.Logf("%d: %s", ix, string(b))
+
+		receivedResponse, streamErr := stream.Recv()
+		checks.NoError(t, streamErr, "stream.Recv() failed")
+		if !compareChatResponses(expectedResponse, receivedResponse) {
+			t.Errorf("Stream response %v is %v, expected %v", ix, receivedResponse, expectedResponse)
+		}
+	}
+
+	_, streamErr := stream.Recv()
+	if !errors.Is(streamErr, io.EOF) {
+		t.Errorf("stream.Recv() did not return EOF in the end: %v", streamErr)
+	}
+}
+
+func TestCreateChatCompletionStreamWithLogprobs(t *testing.T) {
+	client, server, teardown := setupOpenAITestServer()
+	defer teardown()
+	server.RegisterHandler("/v1/chat/completions", func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+
+		// Send test responses
+		dataBytes := []byte{}
+
+		//nolint:lll
+		dataBytes = append(dataBytes, []byte(`data: {"id":"1","object":"chat.completion.chunk","created":1729585728,"model":"gpt-4o-mini-2024-07-18","system_fingerprint":"fp_d9767fc5b9","choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"logprobs":{"content":[],"refusal":null},"finish_reason":null}]}`)...)
+		dataBytes = append(dataBytes, []byte("\n\n")...)
+
+		//nolint:lll
+		dataBytes = append(dataBytes, []byte(`data: {"id":"2","object":"chat.completion.chunk","created":1729585728,"model":"gpt-4o-mini-2024-07-18","system_fingerprint":"fp_d9767fc5b9","choices":[{"index":0,"delta":{"content":"Hello"},"logprobs":{"content":[{"token":"Hello","logprob":-0.000020458236,"bytes":[72,101,108,108,111],"top_logprobs":[]}],"refusal":null},"finish_reason":null}]}`)...)
+		dataBytes = append(dataBytes, []byte("\n\n")...)
+
+		//nolint:lll
+		dataBytes = append(dataBytes, []byte(`data: {"id":"3","object":"chat.completion.chunk","created":1729585728,"model":"gpt-4o-mini-2024-07-18","system_fingerprint":"fp_d9767fc5b9","choices":[{"index":0,"delta":{"content":" World"},"logprobs":{"content":[{"token":" World","logprob":-0.00055303273,"bytes":[32,87,111,114,108,100],"top_logprobs":[]}],"refusal":null},"finish_reason":null}]}`)...)
+		dataBytes = append(dataBytes, []byte("\n\n")...)
+
+		//nolint:lll
+		dataBytes = append(dataBytes, []byte(`data: {"id":"4","object":"chat.completion.chunk","created":1729585728,"model":"gpt-4o-mini-2024-07-18","system_fingerprint":"fp_d9767fc5b9","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}]}`)...)
+		dataBytes = append(dataBytes, []byte("\n\n")...)
+
+		dataBytes = append(dataBytes, []byte("data: [DONE]\n\n")...)
+
+		_, err := w.Write(dataBytes)
+		checks.NoError(t, err, "Write error")
+	})
+
+	stream, err := client.CreateChatCompletionStream(context.Background(), openai.ChatCompletionRequest{
+		MaxTokens: 2000,
+		Model:     openai.GPT3Dot5Turbo,
+		Messages: []openai.ChatCompletionMessage{
+			{
+				Role:    openai.ChatMessageRoleUser,
+				Content: "Hello!",
+			},
+		},
+		Stream: true,
+	})
+	checks.NoError(t, err, "CreateCompletionStream returned error")
+	defer stream.Close()
+
+	expectedResponses := []openai.ChatCompletionStreamResponse{
+		{
+			ID:                "1",
+			Object:            "chat.completion.chunk",
+			Created:           1729585728,
+			Model:             openai.GPT4oMini20240718,
+			SystemFingerprint: "fp_d9767fc5b9",
+			Choices: []openai.ChatCompletionStreamChoice{
+				{
+					Index: 0,
+					Delta: openai.ChatCompletionStreamChoiceDelta{},
+					Logprobs: &openai.ChatCompletionStreamChoiceLogprobs{
+						Content: []openai.ChatCompletionTokenLogprob{},
+					},
+				},
+			},
+		},
+		{
+			ID:                "2",
+			Object:            "chat.completion.chunk",
+			Created:           1729585728,
+			Model:             openai.GPT4oMini20240718,
+			SystemFingerprint: "fp_d9767fc5b9",
+			Choices: []openai.ChatCompletionStreamChoice{
+				{
+					Index: 0,
+					Delta: openai.ChatCompletionStreamChoiceDelta{
+						Content: "Hello",
+					},
+					Logprobs: &openai.ChatCompletionStreamChoiceLogprobs{
+						Content: []openai.ChatCompletionTokenLogprob{
+							{
+								Token:       "Hello",
+								Logprob:     -0.000020458236,
+								Bytes:       []int64{72, 101, 108, 108, 111},
+								TopLogprobs: []openai.ChatCompletionTokenLogprobTopLogprob{},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			ID:                "3",
+			Object:            "chat.completion.chunk",
+			Created:           1729585728,
+			Model:             openai.GPT4oMini20240718,
+			SystemFingerprint: "fp_d9767fc5b9",
+			Choices: []openai.ChatCompletionStreamChoice{
+				{
+					Index: 0,
+					Delta: openai.ChatCompletionStreamChoiceDelta{
+						Content: " World",
+					},
+					Logprobs: &openai.ChatCompletionStreamChoiceLogprobs{
+						Content: []openai.ChatCompletionTokenLogprob{
+							{
+								Token:       " World",
+								Logprob:     -0.00055303273,
+								Bytes:       []int64{32, 87, 111, 114, 108, 100},
+								TopLogprobs: []openai.ChatCompletionTokenLogprobTopLogprob{},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			ID:                "4",
+			Object:            "chat.completion.chunk",
+			Created:           1729585728,
+			Model:             openai.GPT4oMini20240718,
+			SystemFingerprint: "fp_d9767fc5b9",
+			Choices: []openai.ChatCompletionStreamChoice{
+				{
+					Index:        0,
+					Delta:        openai.ChatCompletionStreamChoiceDelta{},
+					FinishReason: "stop",
+				},
+			},
+		},
+	}
+
+	for ix, expectedResponse := range expectedResponses {
+		b, _ := json.Marshal(expectedResponse)
+		t.Logf("%d: %s", ix, string(b))
+
+		receivedResponse, streamErr := stream.Recv()
+		checks.NoError(t, streamErr, "stream.Recv() failed")
+		if !compareChatResponses(expectedResponse, receivedResponse) {
+			t.Errorf("Stream response %v is %v, expected %v", ix, receivedResponse, expectedResponse)
+		}
+	}
+
+	_, streamErr := stream.Recv()
+	if !errors.Is(streamErr, io.EOF) {
+		t.Errorf("stream.Recv() did not return EOF in the end: %v", streamErr)
+	}
+}
+
 func TestAzureCreateChatCompletionStreamRateLimitError(t *testing.T) {
 	wantCode := "429"
 	wantMessage := "Requests to the Creates a completion for the chat message Operation under Azure OpenAI API " +
diff --git a/client_test.go b/client_test.go
index 18da787..354a6b3 100644
--- a/client_test.go
+++ b/client_test.go
@@ -194,26 +194,31 @@ func TestHandleErrorResp(t *testing.T) {
 				{
 					"error":{}
 				}`)),
-			expected: "error, status code: 503, status: , message: ",
+			expected: `error, status code: 503, status: , message: , body: 
+				{
+					"error":{}
+				}`,
 		},
 		{
 			name:        "413 Request Entity Too Large",
 			httpCode:    http.StatusRequestEntityTooLarge,
 			contentType: "text/html",
-			body: bytes.NewReader([]byte(`<html>
-<head><title>413 Request Entity Too Large</title></head>
-<body>
-<center><h1>413 Request Entity Too Large</h1></center>
-<hr><center>nginx</center>
-</body>
-</html>`)),
-			expected: `error, status code: 413, status: , body: <html>
-<head><title>413 Request Entity Too Large</title></head>
-<body>
-<center><h1>413 Request Entity Too Large</h1></center>
-<hr><center>nginx</center>
-</body>
-</html>`,
+			body: bytes.NewReader([]byte(`
+	<html>
+	<head><title>413 Request Entity Too Large</title></head>
+	<body>
+	<center><h1>413 Request Entity Too Large</h1></center>
+	<hr><center>nginx</center>
+	</body>
+	</html>`)),
+			expected: `error, status code: 413, status: , message: invalid character '<' looking for beginning of value, body: 
+	<html>
+	<head><title>413 Request Entity Too Large</title></head>
+	<body>
+	<center><h1>413 Request Entity Too Large</h1></center>
+	<hr><center>nginx</center>
+	</body>
+	</html>`,
 		},
 		{
 			name:        "errorReader",
diff --git a/completion_test.go b/completion_test.go
index 89950bf..935bbe8 100644
--- a/completion_test.go
+++ b/completion_test.go
@@ -59,6 +59,38 @@ func TestCompletions(t *testing.T) {
 	checks.NoError(t, err, "CreateCompletion error")
 }
 
+// TestMultiplePromptsCompletionsWrong Tests the completions endpoint of the API using the mocked server
+// where the completions requests has a list of prompts with wrong type.
+func TestMultiplePromptsCompletionsWrong(t *testing.T) {
+	client, server, teardown := setupOpenAITestServer()
+	defer teardown()
+	server.RegisterHandler("/v1/completions", handleCompletionEndpoint)
+	req := openai.CompletionRequest{
+		MaxTokens: 5,
+		Model:     "ada",
+		Prompt:    []interface{}{"Lorem ipsum", 9},
+	}
+	_, err := client.CreateCompletion(context.Background(), req)
+	if !errors.Is(err, openai.ErrCompletionRequestPromptTypeNotSupported) {
+		t.Fatalf("CreateCompletion should return ErrCompletionRequestPromptTypeNotSupported, but returned: %v", err)
+	}
+}
+
+// TestMultiplePromptsCompletions Tests the completions endpoint of the API using the mocked server
+// where the completions requests has a list of prompts.
+func TestMultiplePromptsCompletions(t *testing.T) {
+	client, server, teardown := setupOpenAITestServer()
+	defer teardown()
+	server.RegisterHandler("/v1/completions", handleCompletionEndpoint)
+	req := openai.CompletionRequest{
+		MaxTokens: 5,
+		Model:     "ada",
+		Prompt:    []interface{}{"Lorem ipsum", "Lorem ipsum"},
+	}
+	_, err := client.CreateCompletion(context.Background(), req)
+	checks.NoError(t, err, "CreateCompletion error")
+}
+
 // handleCompletionEndpoint Handles the completion endpoint by the test server.
 func handleCompletionEndpoint(w http.ResponseWriter, r *http.Request) {
 	var err error
@@ -87,24 +119,50 @@ func handleCompletionEndpoint(w http.ResponseWriter, r *http.Request) {
 	if n == 0 {
 		n = 1
 	}
+	// Handle different types of prompts: single string or list of strings
+	prompts := []string{}
+	switch v := completionReq.Prompt.(type) {
+	case string:
+		prompts = append(prompts, v)
+	case []interface{}:
+		for _, item := range v {
+			if str, ok := item.(string); ok {
+				prompts = append(prompts, str)
+			}
+		}
+	default:
+		http.Error(w, "Invalid prompt type", http.StatusBadRequest)
+		return
+	}
+
 	for i := 0; i < n; i++ {
-		// generate a random string of length completionReq.Length
-		completionStr := strings.Repeat("a", completionReq.MaxTokens)
-		if completionReq.Echo {
-			completionStr = completionReq.Prompt.(string) + completionStr
+		for _, prompt := range prompts {
+			// Generate a random string of length completionReq.MaxTokens
+			completionStr := strings.Repeat("a", completionReq.MaxTokens)
+			if completionReq.Echo {
+				completionStr = prompt + completionStr
+			}
+
+			res.Choices = append(res.Choices, openai.CompletionChoice{
+				Text:  completionStr,
+				Index: len(res.Choices),
+			})
 		}
-		res.Choices = append(res.Choices, openai.CompletionChoice{
-			Text:  completionStr,
-			Index: i,
-		})
 	}
-	inputTokens := numTokens(completionReq.Prompt.(string)) * n
-	completionTokens := completionReq.MaxTokens * n
+
+	inputTokens := 0
+	for _, prompt := range prompts {
+		inputTokens += numTokens(prompt)
+	}
+	inputTokens *= n
+	completionTokens := completionReq.MaxTokens * len(prompts) * n
 	res.Usage = openai.Usage{
 		PromptTokens:     inputTokens,
 		CompletionTokens: completionTokens,
 		TotalTokens:      inputTokens + completionTokens,
 	}
+
+	// Serialize the response and send it back
 	resBytes, _ = json.Marshal(res)
 	fmt.Fprintln(w, string(resBytes))
 }