Skip to content

Commit a5929ab

Browse files
committed
extend openai embedding api
Signed-off-by: yxia216 <yxia216@bloomberg.net>
1 parent 3a56916 commit a5929ab

12 files changed

Lines changed: 903 additions & 364 deletions

File tree

internal/apischema/gcp/gcp.go

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -110,35 +110,49 @@ type PredictResponse struct {
110110

111111
// EmbedContentRequest is the request body for the embedContent endpoint used by newer embedding models
112112
// (e.g. gemini-embedding-2-*). All input texts are packed as parts in a single Content object.
113-
// https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api
113+
//
114+
// The REST API reference documents deprecated top-level fields (taskType, outputDimensionality, etc.):
115+
//
116+
// https://docs.cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.publishers.models/embedContent
117+
//
118+
// However, Vertex AI also accepts "embedContentConfig" as a nested config object (undocumented in REST
119+
// reference but used by the genai SDK v1.54+ (https://github.com/googleapis/go-genai/blob/v1.54.0/models.go#L727)
120+
// and verified experimentally). We use "embedContentConfig" to follow the SDK convention and avoid
121+
// the deprecated fields.
114122
type EmbedContentRequest struct {
115-
Content genai.Content `json:"content"`
123+
Content genai.Content `json:"content"`
124+
Config *EmbedContentConfig `json:"embedContentConfig,omitempty"`
125+
}
126+
127+
// EmbedContentConfig contains optional parameters for the embedContent method.
128+
// https://github.com/googleapis/go-genai/blob/v1.54.0/models.go#L727
129+
type EmbedContentConfig struct {
116130
TaskType openai.EmbeddingTaskType `json:"taskType,omitempty"`
117131
Title string `json:"title,omitempty"`
118132
OutputDimensionality int `json:"outputDimensionality,omitempty"`
119133
AutoTruncate *bool `json:"autoTruncate,omitempty"`
120134
}
121135

122136
// EmbedContentResponse is the response from the embedContent endpoint.
123-
// Each part in the request produces a corresponding embedding in the response.
137+
// The REST API returns a single embedding (not an array), plus usage metadata.
138+
// https://docs.cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.publishers.models/embedContent
124139
type EmbedContentResponse struct {
125-
Embeddings []*EmbedContentEmbedding `json:"embeddings,omitempty"`
140+
// The embedding generated from the input content (singular object, not an array).
141+
Embedding *EmbedContentEmbedding `json:"embedding,omitempty"`
142+
// Usage metadata about the response.
143+
UsageMetadata *EmbedContentUsageMetadata `json:"usageMetadata,omitempty"`
144+
// Whether the input content was truncated before generating the embedding.
145+
Truncated bool `json:"truncated,omitempty"`
126146
}
127147

128-
// EmbedContentEmbedding represents a single embedding from the embedContent response.
148+
// EmbedContentEmbedding represents the embedding values from the embedContent response.
129149
type EmbedContentEmbedding struct {
130150
// The embedding values.
131151
Values []float32 `json:"values,omitempty"`
132-
// Vertex API only. Statistics of the input text associated with this embedding.
133-
Statistics *EmbedContentEmbeddingStatistics `json:"statistics,omitempty"`
134152
}
135153

136-
// EmbedContentEmbeddingStatistics contains per-embedding statistics from the embedContent endpoint.
137-
// Note: Unlike the predict endpoint which uses snake_case (token_count), the embedContent endpoint
138-
// uses camelCase (tokenCount), matching the genai SDK's ContentEmbeddingStatistics type.
139-
type EmbedContentEmbeddingStatistics struct {
140-
// Number of tokens of the input text.
141-
TokenCount float32 `json:"tokenCount,omitempty"`
142-
// Whether the input text was truncated.
143-
Truncated bool `json:"truncated,omitempty"`
154+
// EmbedContentUsageMetadata contains token usage from the embedContent response.
155+
type EmbedContentUsageMetadata struct {
156+
PromptTokenCount int `json:"promptTokenCount,omitempty"`
157+
TotalTokenCount int `json:"totalTokenCount,omitempty"`
144158
}

internal/apischema/openai/openai.go

Lines changed: 82 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1602,15 +1602,8 @@ type Model struct {
16021602
OwnedBy string `json:"owned_by"`
16031603
}
16041604

1605-
// EmbeddingRequest represents a request structure for embeddings API.
1606-
type EmbeddingRequest struct {
1607-
// Input: Input text to embed, encoded as a string or array of tokens.
1608-
// To embed multiple inputs in a single request, pass an array of strings or array of token arrays.
1609-
// The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002),
1610-
// cannot be an empty string, and any array must be 2048 dimensions or less.
1611-
// Docs: https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-input
1612-
Input EmbeddingRequestInput `json:"input"`
1613-
1605+
// EmbeddingBaseRequest holds fields shared by both embedding request variants.
1606+
type EmbeddingBaseRequest struct {
16141607
// Model: ID of the model to use.
16151608
// Docs: https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-model
16161609
Model string `json:"model"`
@@ -1632,6 +1625,76 @@ type EmbeddingRequest struct {
16321625
*GCPVertexAIEmbeddingVendorFields `json:",inline,omitempty"`
16331626
}
16341627

1628+
// EmbeddingCompletionRequest is the text-only embedding request (classic OpenAI style).
1629+
// https://developers.openai.com/api/reference/resources/embeddings/methods/create
1630+
type EmbeddingCompletionRequest struct {
1631+
EmbeddingBaseRequest
1632+
// Input: Input text to embed, encoded as a string or array of tokens.
1633+
// To embed multiple inputs in a single request, pass an array of strings or array of token arrays.
1634+
// The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002),
1635+
// cannot be an empty string, and any array must be 2048 dimensions or less.
1636+
Input EmbeddingRequestInput `json:"input"`
1637+
}
1638+
1639+
// EmbeddingChatRequest is the chat-style embedding request (supports multimodal).
1640+
// Following vLLM convention: https://github.com/vllm-project/vllm/blob/2a16ece2d342c0c154a4949ad317b521f8c04ec4/vllm/entrypoints/pooling/embed/protocol.py#L83
1641+
type EmbeddingChatRequest struct {
1642+
EmbeddingBaseRequest
1643+
// Messages: Chat messages for multimodal embedding.
1644+
// Uses the same chat message format as /v1/chat/completions.
1645+
// Only user messages are processed; system/assistant/tool messages are ignored.
1646+
Messages []ChatCompletionMessageParamUnion `json:"messages"`
1647+
}
1648+
1649+
// EmbeddingRequest is a discriminated union: exactly one of OfCompletion or OfChat is set.
1650+
// Discrimination is by presence of "input" (→ completion) vs "messages" (→ chat) in JSON.
1651+
type EmbeddingRequest struct {
1652+
EmbeddingBaseRequest // promoted shared fields
1653+
OfCompletion *EmbeddingCompletionRequest // set when JSON has "input"
1654+
OfChat *EmbeddingChatRequest // set when JSON has "messages"
1655+
}
1656+
1657+
// UnmarshalJSON discriminates between completion and chat embedding requests.
1658+
func (r *EmbeddingRequest) UnmarshalJSON(data []byte) error {
1659+
hasInput := gjson.GetBytes(data, "input").Exists()
1660+
hasMessages := gjson.GetBytes(data, "messages").Exists()
1661+
1662+
if hasInput && hasMessages {
1663+
return fmt.Errorf("embedding request must have either 'input' or 'messages', not both")
1664+
}
1665+
1666+
if hasMessages {
1667+
var chat EmbeddingChatRequest
1668+
if err := json.Unmarshal(data, &chat); err != nil {
1669+
return err
1670+
}
1671+
r.EmbeddingBaseRequest = chat.EmbeddingBaseRequest
1672+
r.OfChat = &chat
1673+
return nil
1674+
}
1675+
1676+
// Default to completion (input-based) — this handles both explicit "input" and legacy cases.
1677+
var comp EmbeddingCompletionRequest
1678+
if err := json.Unmarshal(data, &comp); err != nil {
1679+
return err
1680+
}
1681+
r.EmbeddingBaseRequest = comp.EmbeddingBaseRequest
1682+
r.OfCompletion = &comp
1683+
return nil
1684+
}
1685+
1686+
// MarshalJSON delegates to the active variant.
1687+
func (r EmbeddingRequest) MarshalJSON() ([]byte, error) {
1688+
if r.OfChat != nil {
1689+
return json.Marshal(r.OfChat)
1690+
}
1691+
if r.OfCompletion != nil {
1692+
return json.Marshal(r.OfCompletion)
1693+
}
1694+
// Fallback: marshal just the base fields.
1695+
return json.Marshal(r.EmbeddingBaseRequest)
1696+
}
1697+
16351698
type EmbeddingTaskType string
16361699

16371700
const (
@@ -1645,16 +1708,20 @@ const (
16451708
EmbeddingTaskTypeCodeRetrievalQuery EmbeddingTaskType = "CODE_RETRIEVAL_QUERY"
16461709
)
16471710

1648-
// GCPVertexAIEmbeddingVendorFields contains GCP Vertex AI (Gemini) vendor-specific fields for embeddings.
1711+
// GCPVertexAIEmbeddingVendorFields contains GCP Vertex AI vendor-specific fields for embeddings.
1712+
// The translator maps these to the appropriate wire format per endpoint:
1713+
// - predict: parameters.auto_truncate, instances[].task_type, instances[].title
1714+
// - embedContent: embedContentConfig.autoTruncate, embedContentConfig.taskType, embedContentConfig.title
16491715
type GCPVertexAIEmbeddingVendorFields struct {
1650-
// When set to true, input text will be truncated. When set to false, an error is returned if the input text is longer than the maximum length supported by the model. Defaults to true.
1651-
// https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#parameter-list
1652-
1716+
// Auto-truncate input text if it exceeds the model's max length. Defaults to true.
16531717
AutoTruncate bool `json:"auto_truncate,omitempty"`
16541718

1655-
// This is global task_type set, which is convenient for users. If left blank, the default used is RETRIEVAL_QUERY.
1656-
// For more information about task types, see https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types
1719+
// Global task type for the request. Defaults to RETRIEVAL_QUERY if unset.
1720+
// https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types
16571721
TaskType EmbeddingTaskType `json:"task_type,omitempty"`
1722+
1723+
// Title for the embedding content. Helps the model produce better embeddings.
1724+
Title string `json:"title,omitempty"`
16581725
}
16591726

16601727
// EmbeddingResponse represents a response from /v1/embeddings.

internal/endpointspec/endpointspec_test.go

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,8 +143,14 @@ func TestEmbeddingsEndpointSpec_ParseBody(t *testing.T) {
143143
require.ErrorContains(t, err, "malformed request")
144144
})
145145

146-
t.Run("success", func(t *testing.T) {
147-
req := openai.EmbeddingRequest{Model: "text-embedding-3-large", Input: openai.EmbeddingRequestInput{Value: "input"}}
146+
t.Run("success with input", func(t *testing.T) {
147+
req := openai.EmbeddingRequest{
148+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-large"},
149+
OfCompletion: &openai.EmbeddingCompletionRequest{
150+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-large"},
151+
Input: openai.EmbeddingRequestInput{Value: "input"},
152+
},
153+
}
148154
body, err := json.Marshal(req)
149155
require.NoError(t, err)
150156

@@ -155,6 +161,17 @@ func TestEmbeddingsEndpointSpec_ParseBody(t *testing.T) {
155161
require.NotNil(t, parsed)
156162
require.Nil(t, mutated)
157163
})
164+
165+
t.Run("success with messages", func(t *testing.T) {
166+
body := []byte(`{"model":"gemini-embedding-2","messages":[{"role":"user","content":"embed this"}]}`)
167+
168+
model, parsed, stream, mutated, err := spec.ParseBody(body, false)
169+
require.NoError(t, err)
170+
require.Equal(t, "gemini-embedding-2", model)
171+
require.False(t, stream)
172+
require.NotNil(t, parsed)
173+
require.Nil(t, mutated)
174+
})
158175
}
159176

160177
func TestEmbeddingsEndpointSpec_GetTranslator(t *testing.T) {

internal/tracing/openinference/openai/embeddings_test.go

Lines changed: 48 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,23 +24,46 @@ import (
2424
// Test data.
2525
var (
2626
basicEmbeddingReq = &openai.EmbeddingRequest{
27-
Model: "text-embedding-3-small",
28-
Input: openai.EmbeddingRequestInput{Value: "How do I reset my password?"},
27+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
28+
OfCompletion: &openai.EmbeddingCompletionRequest{
29+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
30+
Input: openai.EmbeddingRequestInput{Value: "How do I reset my password?"},
31+
},
2932
}
3033
basicEmbeddingReqBody = []byte(`{"model":"text-embedding-3-small","input":"How do I reset my password?"}`)
3134

3235
multiInputEmbeddingReq = &openai.EmbeddingRequest{
33-
Model: "text-embedding-3-small",
34-
Input: openai.EmbeddingRequestInput{Value: []string{"How", "do", "I", "reset", "my", "password?"}},
36+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
37+
OfCompletion: &openai.EmbeddingCompletionRequest{
38+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
39+
Input: openai.EmbeddingRequestInput{Value: []string{"How", "do", "I", "reset", "my", "password?"}},
40+
},
3541
}
3642
multiInputEmbeddingReqBody = []byte(`{"model":"text-embedding-3-small","input":["How","do","I","reset","my","password?"]}`)
3743

3844
tokenInputEmbeddingReq = &openai.EmbeddingRequest{
39-
Model: "text-embedding-3-small",
40-
Input: openai.EmbeddingRequestInput{Value: []int{14438, 656, 358, 7738, 856, 3636, 30}},
45+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
46+
OfCompletion: &openai.EmbeddingCompletionRequest{
47+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
48+
Input: openai.EmbeddingRequestInput{Value: []int64{4438, 656, 358, 7738, 856, 3636, 30}},
49+
},
4150
}
4251
tokenInputEmbeddingReqBody = []byte(`{"model":"text-embedding-3-small","input":[4438,656,358,7738,856,3636,30]}`)
4352

53+
chatEmbeddingReq = &openai.EmbeddingRequest{
54+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "gemini-embedding-2"},
55+
OfChat: &openai.EmbeddingChatRequest{
56+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "gemini-embedding-2"},
57+
Messages: []openai.ChatCompletionMessageParamUnion{
58+
{OfUser: &openai.ChatCompletionUserMessageParam{
59+
Role: "user",
60+
Content: openai.StringOrUserRoleContentUnion{Value: "embed this text"},
61+
}},
62+
},
63+
},
64+
}
65+
chatEmbeddingReqBody = []byte(`{"model":"gemini-embedding-2","messages":[{"role":"user","content":"embed this text"}]}`)
66+
4467
basicEmbeddingResp = &openai.EmbeddingResponse{
4568
Model: "text-embedding-3-small",
4669
Usage: openai.EmbeddingUsage{
@@ -95,6 +118,12 @@ func TestEmbeddingsRecorder_StartParams(t *testing.T) {
95118
reqBody: multiInputEmbeddingReqBody,
96119
expectedSpanName: "CreateEmbeddings",
97120
},
121+
{
122+
name: "chat embedding request",
123+
req: chatEmbeddingReq,
124+
reqBody: chatEmbeddingReqBody,
125+
expectedSpanName: "CreateEmbeddings",
126+
},
98127
}
99128

100129
for _, tt := range tests {
@@ -172,6 +201,19 @@ func TestEmbeddingsRecorder_RecordRequest(t *testing.T) {
172201
attribute.String(openinference.EmbeddingInvocationParameters, `{"model":"text-embedding-3-small"}`),
173202
},
174203
},
204+
{
205+
name: "chat embedding request (no embedding text attributes)",
206+
req: chatEmbeddingReq,
207+
reqBody: chatEmbeddingReqBody,
208+
config: &openinference.TraceConfig{},
209+
expectedAttrs: []attribute.KeyValue{
210+
attribute.String(openinference.SpanKind, openinference.SpanKindEmbedding),
211+
attribute.String(openinference.InputValue, string(chatEmbeddingReqBody)),
212+
attribute.String(openinference.InputMimeType, openinference.MimeTypeJSON),
213+
attribute.String(openinference.EmbeddingInvocationParameters, `{"model":"gemini-embedding-2"}`),
214+
// No embedding.text attributes for chat-style requests.
215+
},
216+
},
175217
{
176218
name: "hidden invocation parameters",
177219
req: basicEmbeddingReq,

internal/tracing/openinference/openai/request_attrs.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -255,8 +255,8 @@ func buildEmbeddingsRequestAttributes(embRequest *openai.EmbeddingRequest, body
255255
// 3. It would require model-specific tokenizer libraries (tiktoken, sentencepiece, etc.)
256256
// 4. Azure deployments don't affect this (they only host OpenAI models with cl100k_base)
257257
// Following OpenInference spec guidance to only record human-readable text.
258-
if !config.HideInputs && !config.HideEmbeddingsText {
259-
switch input := embRequest.Input.Value.(type) {
258+
if !config.HideInputs && !config.HideEmbeddingsText && embRequest.OfCompletion != nil {
259+
switch input := embRequest.OfCompletion.Input.Value.(type) {
260260
case string:
261261
attrs = append(attrs, attribute.String(openinference.EmbeddingTextAttribute(0), input))
262262
case []string:

internal/tracing/tracing_test.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -775,9 +775,12 @@ func TestNewTracingFromEnv_Embeddings_Redaction(t *testing.T) {
775775

776776
// Create a test request with sensitive data.
777777
req := &openai.EmbeddingRequest{
778-
Model: "text-embedding-3-small",
779-
Input: openai.EmbeddingRequestInput{
780-
Value: "Sensitive embedding text",
778+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
779+
OfCompletion: &openai.EmbeddingCompletionRequest{
780+
EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
781+
Input: openai.EmbeddingRequestInput{
782+
Value: "Sensitive embedding text",
783+
},
781784
},
782785
}
783786
reqBody := []byte(`{"input":"Sensitive embedding text","model":"text-embedding-3-small"}`)

internal/translator/openai_awsbedrock_embeddings.go

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,12 @@ func (o *openAIToAWSBedrockTranslatorV1Embedding) RequestBody(_ []byte, req *ope
4141
}
4242
o.requestModel = model
4343

44+
if req.OfCompletion == nil {
45+
return nil, nil, fmt.Errorf("%w: AWS Bedrock Titan requires an input-based embedding request (messages not supported)", internalapi.ErrInvalidRequestBody)
46+
}
47+
4448
var inputText string
45-
switch v := req.Input.Value.(type) {
49+
switch v := req.OfCompletion.Input.Value.(type) {
4650
case string:
4751
inputText = v
4852
case []string:
@@ -52,7 +56,7 @@ func (o *openAIToAWSBedrockTranslatorV1Embedding) RequestBody(_ []byte, req *ope
5256
}
5357
inputText = v[0]
5458
default:
55-
return nil, nil, fmt.Errorf("%w: unsupported input type %T", internalapi.ErrInvalidRequestBody, req.Input.Value)
59+
return nil, nil, fmt.Errorf("%w: unsupported input type %T", internalapi.ErrInvalidRequestBody, req.OfCompletion.Input.Value)
5660
}
5761

5862
bedrockReq := awsbedrock.TitanEmbeddingRequest{

0 commit comments

Comments
 (0)