Repository: sashabaranov/go-openai
Pull request #417 (2023-06-29 03:41:38)

Allow embeddings requests to be tokens or strings
Addresses #240 

There may be a cleaner way to do this, but this is the best way I could think of quickly to not cause any breaking changes and to maintain some semblance of type checking for the user (instead of just changing the Input field type in EmbeddingRequest from string to interface{}).

Also modified the test case for it.

Thanks for all of the work you guys do on this project!

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/sashabaranov/go-openai/pull/417?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#417](https://app.codecov.io/gh/sashabaranov/go-openai/pull/417?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (0df26ac) into [master](https://app.codecov.io/gh/sashabaranov/go-openai/commit/9c99f3626f1d80382e187df8adc38f7e8e929a75?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (9c99f36) will **decrease** coverage by `0.80%`.
> The diff coverage is `60.00%`.

```diff
@@            Coverage Diff             @@
##           master     #417      +/-   ##
==========================================
- Coverage   96.95%   96.15%   -0.80%     
==========================================
  Files          17       17              
  Lines         690      703      +13     
==========================================
+ Hits          669      676       +7     
- Misses         15       21       +6     
  Partials        6        6              
```


| [Impacted Files](https://app.codecov.io/gh/sashabaranov/go-openai/pull/417?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [embeddings.go](https://app.codecov.io/gh/sashabaranov/go-openai/pull/417?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-ZW1iZWRkaW5ncy5nbw==) | `81.81% <60.00%> (-18.19%)` | :arrow_down: |



Patch:
diff --git a/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md b/.github/PULL_REQUEST_TEMPLATE.md
similarity index 77%
rename from .github/PULL_REQUEST_TEMPLATE/pull_request_template.md
rename to .github/PULL_REQUEST_TEMPLATE.md
index b078d19..44bf697 100644
--- a/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -1,15 +1,8 @@
----
-name: Pull Request
-about: Propose changes to the codebase
-title: ''
-labels: ''
-assignees: ''
-
----
-
 A similar PR may already be submitted!
 Please search among the [Pull request](https://github.com/sashabaranov/go-openai/pulls) before creating one.
 
+If your changes introduce breaking changes, please prefix the title of your pull request with "[BREAKING_CHANGES]". This allows for clear identification of such changes in the 'What's Changed' section on the release page, making it developer-friendly.
+
 Thanks for submitting a pull request! Please provide enough information so that others can review your pull request.
 
 **Describe the change**
diff --git a/embeddings.go b/embeddings.go
index ba327ce..41af50b 100644
--- a/embeddings.go
+++ b/embeddings.go
@@ -113,10 +113,25 @@ type EmbeddingResponse struct {
 	Usage  Usage          `json:"usage"`
 }
 
-// EmbeddingRequest is the input to a Create embeddings request.
+type EmbeddingRequestConverter interface {
+	// Needs to be of type EmbeddingRequestStrings or EmbeddingRequestTokens
+	Convert() EmbeddingRequest
+}
+
 type EmbeddingRequest struct {
+	Input any            `json:"input"`
+	Model EmbeddingModel `json:"model"`
+	User  string         `json:"user"`
+}
+
+func (r EmbeddingRequest) Convert() EmbeddingRequest {
+	return r
+}
+
+// EmbeddingRequestStrings is the input to a create embeddings request with a slice of strings.
+type EmbeddingRequestStrings struct {
 	// Input is a slice of strings for which you want to generate an Embedding vector.
-	// Each input must not exceed 2048 tokens in length.
+	// Each input must not exceed 8192 tokens in length.
 	// OpenAPI suggests replacing newlines (\n) in your input with a single space, as they
 	// have observed inferior results when newlines are present.
 	// E.g.
@@ -129,15 +144,50 @@ type EmbeddingRequest struct {
 	User string `json:"user"`
 }
 
-// CreateEmbeddings returns an EmbeddingResponse which will contain an Embedding for every item in |request.Input|.
+func (r EmbeddingRequestStrings) Convert() EmbeddingRequest {
+	return EmbeddingRequest{
+		Input: r.Input,
+		Model: r.Model,
+		User:  r.User,
+	}
+}
+
+type EmbeddingRequestTokens struct {
+	// Input is a slice of slices of ints ([][]int) for which you want to generate an Embedding vector.
+	// Each input must not exceed 8192 tokens in length.
+	// OpenAPI suggests replacing newlines (\n) in your input with a single space, as they
+	// have observed inferior results when newlines are present.
+	// E.g.
+	//	"The food was delicious and the waiter..."
+	Input [][]int `json:"input"`
+	// ID of the model to use. You can use the List models API to see all of your available models,
+	// or see our Model overview for descriptions of them.
+	Model EmbeddingModel `json:"model"`
+	// A unique identifier representing your end-user, which will help OpenAI to monitor and detect abuse.
+	User string `json:"user"`
+}
+
+func (r EmbeddingRequestTokens) Convert() EmbeddingRequest {
+	return EmbeddingRequest{
+		Input: r.Input,
+		Model: r.Model,
+		User:  r.User,
+	}
+}
+
+// CreateEmbeddings returns an EmbeddingResponse which will contain an Embedding for every item in |body.Input|.
 // https://beta.openai.com/docs/api-reference/embeddings/create
-func (c *Client) CreateEmbeddings(ctx context.Context, request EmbeddingRequest) (resp EmbeddingResponse, err error) {
-	req, err := c.newRequest(ctx, http.MethodPost, c.fullURL("/embeddings", request.Model.String()), withBody(request))
+//
+// Body should be of type EmbeddingRequestStrings for embedding strings or EmbeddingRequestTokens
+// for embedding groups of text already converted to tokens.
+func (c *Client) CreateEmbeddings(ctx context.Context, conv EmbeddingRequestConverter) (res EmbeddingResponse, err error) { //nolint:lll
+	baseReq := conv.Convert()
+	req, err := c.newRequest(ctx, http.MethodPost, c.fullURL("/embeddings", baseReq.Model.String()), withBody(baseReq))
 	if err != nil {
 		return
 	}
 
-	err = c.sendRequest(req, &resp)
+	err = c.sendRequest(req, &res)
 
 	return
 }
diff --git a/jsonschema/json.go b/jsonschema/json.go
index 24af858..c02d250 100644
--- a/jsonschema/json.go
+++ b/jsonschema/json.go
@@ -27,7 +27,7 @@ type Definition struct {
 	// one element, where each element is unique. You will probably only use this with strings.
 	Enum []string `json:"enum,omitempty"`
 	// Properties describes the properties of an object, if the schema type is Object.
-	Properties map[string]Definition `json:"properties,omitempty"`
+	Properties map[string]Definition `json:"properties"`
 	// Required specifies which properties are required, if the schema type is Object.
 	Required []string `json:"required,omitempty"`
 	// Items specifies which data type an array contains, if the schema type is Array.


Test patch:
diff --git a/embeddings_test.go b/embeddings_test.go
index d7892cd..47c4f51 100644
--- a/embeddings_test.go
+++ b/embeddings_test.go
@@ -32,6 +32,7 @@ func TestEmbedding(t *testing.T) {
 		BabbageCodeSearchText,
 	}
 	for _, model := range embeddedModels {
+		// test embedding request with strings (simple embedding request)
 		embeddingReq := EmbeddingRequest{
 			Input: []string{
 				"The food was delicious and the waiter",
@@ -46,6 +47,34 @@ func TestEmbedding(t *testing.T) {
 		if !bytes.Contains(marshaled, []byte(`"model":"`+model.String()+`"`)) {
 			t.Fatalf("Expected embedding request to contain model field")
 		}
+
+		// test embedding request with strings
+		embeddingReqStrings := EmbeddingRequestStrings{
+			Input: []string{
+				"The food was delicious and the waiter",
+				"Other examples of embedding request",
+			},
+			Model: model,
+		}
+		marshaled, err = json.Marshal(embeddingReqStrings)
+		checks.NoError(t, err, "Could not marshal embedding request")
+		if !bytes.Contains(marshaled, []byte(`"model":"`+model.String()+`"`)) {
+			t.Fatalf("Expected embedding request to contain model field")
+		}
+
+		// test embedding request with tokens
+		embeddingReqTokens := EmbeddingRequestTokens{
+			Input: [][]int{
+				{464, 2057, 373, 12625, 290, 262, 46612},
+				{6395, 6096, 286, 11525, 12083, 2581},
+			},
+			Model: model,
+		}
+		marshaled, err = json.Marshal(embeddingReqTokens)
+		checks.NoError(t, err, "Could not marshal embedding request")
+		if !bytes.Contains(marshaled, []byte(`"model":"`+model.String()+`"`)) {
+			t.Fatalf("Expected embedding request to contain model field")
+		}
 	}
 }
 
@@ -75,6 +104,15 @@ func TestEmbeddingEndpoint(t *testing.T) {
 			fmt.Fprintln(w, string(resBytes))
 		},
 	)
+	// test create embeddings with strings (simple embedding request)
 	_, err := client.CreateEmbeddings(context.Background(), EmbeddingRequest{})
 	checks.NoError(t, err, "CreateEmbeddings error")
+
+	// test create embeddings with strings
+	_, err = client.CreateEmbeddings(context.Background(), EmbeddingRequestStrings{})
+	checks.NoError(t, err, "CreateEmbeddings strings error")
+
+	// test create embeddings with tokens
+	_, err = client.CreateEmbeddings(context.Background(), EmbeddingRequestTokens{})
+	checks.NoError(t, err, "CreateEmbeddings tokens error")
 }