From afe73b08ff6432aca4749a024391c7f886a5bd2d Mon Sep 17 00:00:00 2001 From: feitianbubu Date: Tue, 20 Oct 2026 21:38:28 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E9=98=BF=E9=87=8C?= =?UTF-8?q?=E5=A4=9A=E6=A8=A1=E6=80=81=E5=90=91=E9=87=8F=E5=8C=96=E6=A8=A1?= =?UTF-8?q?=E5=9E=8Bqwen3-vl-embedding=20(#5179)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add ali MultimodalEmbedding qwen3-vl-embedding * fix: report image/video tokens in qwen3-vl-embedding usage for billing * fix(ali): serve every Bailian multimodal embedding model natively Route qwen2.5-vl-embedding, the tongyi-embedding-vision series and multimodal-embedding-v1, dated snapshots included, to the native multimodal-embedding endpoint like qwen3-vl-embedding. Text embedding models stay in compatible mode. Take tongyi-embedding-vision's text and image tokens from input_tokens_details so the image ratio applies to its image tokens, and set total_tokens to the billed input tokens: tongyi counts its unpriced output_tokens in total_tokens. A response without usage bills the request estimate, as compatible mode does. Return an error code carried in a 200 body as a 500, close the response body when reading it fails, and fold the input conversion into one function. Count the text, image and video parts of embedding input in the request estimate the way chat content is counted, so image calls are pre-charged and a response without usage is not free. With the default settings no image is downloaded for the estimate. * fix(gemini): reject embedding input that is not plain text Gemini embeds text only, so the converter kept the strings and dropped image, video and text parts. The embedding request estimate now counts those parts, and Gemini embeddings bill that estimate, so a dropped part would be charged without ever reaching Gemini. A part now fails the request with a 400 that skips retries, before anything is sent or charged. --------- Co-authored-by: CaIon --- relay/channel/ali/adaptor.go | 16 +- relay/channel/ali/adaptor_test.go | 151 ++++++++++++++ relay/channel/ali/constants.go | 1 + relay/channel/ali/dto.go | 3 + relay/channel/ali/embedding.go | 195 ++++++++++++++++++ relay/channel/gemini/adaptor.go | 15 ++ .../channel/gemini/adaptor_responses_test.go | 25 +++ relaykit/dto/embedding.go | 32 ++- relaykit/dto/token_count_meta_test.go | 17 ++ 9 files changed, 449 insertions(+), 6 deletions(-) create mode 100644 relay/channel/ali/embedding.go diff --git a/relay/channel/ali/adaptor.go b/relay/channel/ali/adaptor.go index 629151d8c..f3e797038 100644 --- a/relay/channel/ali/adaptor.go +++ b/relay/channel/ali/adaptor.go @@ -98,7 +98,11 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { default: switch info.RelayMode { case constant.RelayModeEmbeddings: - fullRequestURL = fmt.Sprintf("%s/compatible-mode/v1/embeddings", info.ChannelBaseUrl) + if isMultimodalEmbeddingModel(info.UpstreamModelName) { + fullRequestURL = fmt.Sprintf("%s/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding", info.ChannelBaseUrl) + } else { + fullRequestURL = fmt.Sprintf("%s/compatible-mode/v1/embeddings", info.ChannelBaseUrl) + } case constant.RelayModeRerank: fullRequestURL = fmt.Sprintf("%s/api/v1/services/rerank/text-rerank/text-rerank", info.ChannelBaseUrl) case constant.RelayModeResponses: @@ -170,6 +174,9 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt } func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) { + if isMultimodalEmbeddingModel(info.UpstreamModelName) { + return convertToMultimodalEmbeddingRequest(request) + } return request, nil } @@ -200,6 +207,13 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom switch info.RelayMode { case constant.RelayModeRerank: err, usage = RerankHandler(c, resp, info) + case constant.RelayModeEmbeddings: + if isMultimodalEmbeddingModel(info.UpstreamModelName) { + err, usage = MultimodalEmbeddingHandler(c, resp, info) + } else { + adaptor := openai.Adaptor{} + usage, err = adaptor.DoResponse(c, resp, info) + } default: adaptor := openai.Adaptor{} usage, err = adaptor.DoResponse(c, resp, info) diff --git a/relay/channel/ali/adaptor_test.go b/relay/channel/ali/adaptor_test.go index 1d45676d6..6e22c630d 100644 --- a/relay/channel/ali/adaptor_test.go +++ b/relay/channel/ali/adaptor_test.go @@ -3,8 +3,10 @@ package ali import ( "encoding/json" "fmt" + "io" "net/http" "net/http/httptest" + "strings" "testing" "github.com/QuantumNous/new-api/common" @@ -336,3 +338,152 @@ func TestAliResponsesWebSearchUsage(t *testing.T) { assert.Equal(t, map[string]int{aliSearchAgentKey: 2, dto.BuildInToolWebSearch: 0, dto.BuildInToolWebSearchPreview: 0}, counts, raw) } } + +// Bailian multimodal embedding models are served only by the native +// multimodal-embedding endpoint, which takes input.contents; text embedding +// models keep compatible mode and the OpenAI body. +func TestEmbeddingRequestByModel(t *testing.T) { + const ( + base = "https://dashscope.invalid" + nativePath = "/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding" + ) + tests := []struct { + model string + request string // OpenAI embedding request fields after model + wantPath string + wantBody string // "" means a conversion error + }{ + { + model: "qwen3-vl-embedding", + request: `"input":"a cat"`, + wantPath: nativePath, + wantBody: `{"model":"qwen3-vl-embedding","input":{"contents":[{"text":"a cat"}]}}`, + }, + { + model: "tongyi-embedding-vision-plus-2026-03-06", + request: `"input":["a cat","a dog"],"dimensions":1024`, + wantPath: nativePath, + wantBody: `{"model":"tongyi-embedding-vision-plus-2026-03-06","input":{"contents":[{"text":"a cat"},{"text":"a dog"}]},"parameters":{"dimension":1024}}`, + }, + { + model: "multimodal-embedding-v1", + request: `"input":[{"type":"text","text":"a cat"},` + + `{"type":"image_url","image_url":{"url":"data:image/png;base64,iVBORw0KGgo="}},` + + `{"type":"video_url","video_url":{"url":"https://example.com/cat.mp4"}}]`, + wantPath: nativePath, + wantBody: `{"model":"multimodal-embedding-v1","input":{"contents":[{"text":"a cat"},` + + `{"image":"data:image/png;base64,iVBORw0KGgo="},{"video":"https://example.com/cat.mp4"}]}}`, + }, + { + model: "qwen3-vl-embedding", + request: `"input":[{"type":"image_url","image_url":{"url":""}}]`, + wantPath: nativePath, + }, + { + model: "text-embedding-v4", + request: `"input":"a cat","dimensions":1024`, + wantPath: "/compatible-mode/v1/embeddings", + wantBody: `{"model":"text-embedding-v4","input":"a cat","dimensions":1024}`, + }, + } + for _, tt := range tests { + info := &relaycommon.RelayInfo{ + RelayMode: constant.RelayModeEmbeddings, + RelayFormat: types.RelayFormatEmbedding, + ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: base, UpstreamModelName: tt.model}, + } + url, err := (&Adaptor{}).GetRequestURL(info) + require.NoError(t, err, tt.request) + assert.Equal(t, base+tt.wantPath, url, tt.request) + + var request dto.EmbeddingRequest + require.NoError(t, common.UnmarshalJsonStr(`{"model":"`+tt.model+`",`+tt.request+`}`, &request)) + converted, err := (&Adaptor{}).ConvertEmbeddingRequest(nil, info, request) + if tt.wantBody == "" { + assert.Error(t, err, tt.request) + continue + } + require.NoError(t, err, tt.request) + wire, err := common.Marshal(converted) + require.NoError(t, err) + assert.JSONEq(t, tt.wantBody, string(wire), tt.request) + } +} + +// The native response becomes an OpenAI embedding list. prompt_tokens is all +// billable input and image_tokens its visual part, priced at the image ratio; +// output_tokens is not priced and never billed. A response without usage +// bills the request estimate. +func TestMultimodalEmbeddingResponse(t *testing.T) { + tests := []struct { + name string + model string + body string + estimate int + wantData []dto.EmbeddingResponseItem + wantUsage *dto.Usage // nil means an upstream error + }{ + { + name: "visual tokens reported apart", + model: "qwen3-vl-embedding", + body: `{"output":{"embeddings":[{"index":0,"embedding":[0.1,-0.2],"type":"text"},{"index":1,"embedding":[0.3,0.4],"type":"image"}]},` + + `"usage":{"input_tokens":43,"image_tokens":1247,"total_tokens":1290},"request_id":"r1"}`, + wantData: []dto.EmbeddingResponseItem{ + {Object: "embedding", Index: 0, Embedding: []float64{0.1, -0.2}}, + {Object: "embedding", Index: 1, Embedding: []float64{0.3, 0.4}}, + }, + wantUsage: &dto.Usage{PromptTokens: 1290, TotalTokens: 1290, + PromptTokensDetails: dto.InputTokenDetails{TextTokens: 43, ImageTokens: 1247}}, + }, + { + name: "input_tokens split in input_tokens_details", + model: "tongyi-embedding-vision-plus", + body: `{"output":{"embeddings":[{"index":0,"embedding":[0.5],"type":"fused"}]},` + + `"usage":{"input_tokens":903,"input_tokens_details":{"image_tokens":896,"text_tokens":7},"output_tokens":3,"total_tokens":906},"request_id":"r2"}`, + wantData: []dto.EmbeddingResponseItem{{Object: "embedding", Index: 0, Embedding: []float64{0.5}}}, + wantUsage: &dto.Usage{PromptTokens: 903, TotalTokens: 903, + PromptTokensDetails: dto.InputTokenDetails{TextTokens: 7, ImageTokens: 896}}, + }, + { + name: "no usage bills the estimate", + model: "qwen3-vl-embedding", + body: `{"output":{"embeddings":[{"index":0,"embedding":[0.5],"type":"text"}]},"request_id":"r3"}`, + estimate: 12, + wantData: []dto.EmbeddingResponseItem{{Object: "embedding", Index: 0, Embedding: []float64{0.5}}}, + wantUsage: &dto.Usage{PromptTokens: 12, TotalTokens: 12}, + }, + { + name: "error code in a 200 body", + model: "qwen3-vl-embedding", + body: `{"code":"InvalidParameter","message":"The input is invalid.","request_id":"r4"}`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + info := &relaycommon.RelayInfo{ + RelayMode: constant.RelayModeEmbeddings, + RelayFormat: types.RelayFormatEmbedding, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: tt.model}, + } + info.SetEstimatePromptTokens(tt.estimate) + resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(tt.body))} + + usage, apiErr := (&Adaptor{}).DoResponse(c, resp, info) + if tt.wantUsage == nil { + require.NotNil(t, apiErr) + assert.Equal(t, http.StatusInternalServerError, apiErr.StatusCode) + assert.Equal(t, "InvalidParameter", string(apiErr.GetErrorCode())) + assert.Empty(t, w.Body.String()) + return + } + require.Nil(t, apiErr) + assert.Equal(t, tt.wantUsage, usage) + + var got dto.EmbeddingResponse + require.NoError(t, common.Unmarshal(w.Body.Bytes(), &got)) + assert.Equal(t, dto.EmbeddingResponse{Object: "list", Data: tt.wantData, Model: tt.model, Usage: *tt.wantUsage}, got) + }) + } +} diff --git a/relay/channel/ali/constants.go b/relay/channel/ali/constants.go index df64439bc..e15503336 100644 --- a/relay/channel/ali/constants.go +++ b/relay/channel/ali/constants.go @@ -8,6 +8,7 @@ var ModelList = []string{ "qwq-32b", "qwen3-235b-a22b", "text-embedding-v1", + "qwen3-vl-embedding", "gte-rerank-v2", } diff --git a/relay/channel/ali/dto.go b/relay/channel/ali/dto.go index 671b88c56..cb0be9add 100644 --- a/relay/channel/ali/dto.go +++ b/relay/channel/ali/dto.go @@ -62,6 +62,9 @@ type AliUsage struct { InputTokens int `json:"input_tokens"` OutputTokens int `json:"output_tokens"` TotalTokens int `json:"total_tokens"` + ImageTokens int `json:"image_tokens,omitempty"` + // InputTokensDetails splits input_tokens on tongyi-embedding-vision. + InputTokensDetails *dto.InputTokenDetails `json:"input_tokens_details,omitempty"` } type AliRerankParameters struct { diff --git a/relay/channel/ali/embedding.go b/relay/channel/ali/embedding.go new file mode 100644 index 000000000..a4e7e4ffc --- /dev/null +++ b/relay/channel/ali/embedding.go @@ -0,0 +1,195 @@ +package ali + +import ( + "fmt" + "io" + "net/http" + "strings" + + "github.com/QuantumNous/new-api/common" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/service" + + "github.com/gin-gonic/gin" +) + +// Bailian multimodal embedding models (qwen3-vl-embedding, qwen2.5-vl-embedding, +// tongyi-embedding-vision-*, multimodal-embedding-v1) are not served in +// compatible mode; they only work on the native multimodal-embedding endpoint. +// Matching the name families also covers dated snapshots. +func isMultimodalEmbeddingModel(model string) bool { + return strings.Contains(model, "vl-embedding") || + strings.Contains(model, "embedding-vision") || + strings.HasPrefix(model, "multimodal-embedding") +} + +type AliMultimodalContent struct { + Text string `json:"text,omitempty"` + Image string `json:"image,omitempty"` + Video string `json:"video,omitempty"` +} + +type AliMultimodalEmbeddingRequest struct { + Model string `json:"model"` + Input struct { + Contents []AliMultimodalContent `json:"contents"` + } `json:"input"` + Parameters *AliMultimodalEmbeddingParameters `json:"parameters,omitempty"` +} + +type AliMultimodalEmbeddingParameters struct { + Dimension *int `json:"dimension,omitempty"` +} + +type AliMultimodalEmbeddingItem struct { + Index int `json:"index"` + Embedding []float64 `json:"embedding"` +} + +type AliMultimodalEmbeddingResponse struct { + Output struct { + Embeddings []AliMultimodalEmbeddingItem `json:"embeddings"` + } `json:"output"` + Usage AliUsage `json:"usage"` + RequestId string `json:"request_id"` + AliError +} + +func convertToMultimodalEmbeddingRequest(request dto.EmbeddingRequest) (*AliMultimodalEmbeddingRequest, error) { + var items []any + switch input := request.Input.(type) { + case nil: + return nil, fmt.Errorf("embedding input is empty") + case string: + items = []any{input} + case []any: + items = input + default: + return nil, fmt.Errorf("unsupported embedding input type: %T", request.Input) + } + if len(items) == 0 { + return nil, fmt.Errorf("embedding input is empty") + } + + aliReq := &AliMultimodalEmbeddingRequest{Model: request.Model} + aliReq.Input.Contents = make([]AliMultimodalContent, 0, len(items)) + for _, item := range items { + if text, ok := item.(string); ok { + aliReq.Input.Contents = append(aliReq.Input.Contents, AliMultimodalContent{Text: text}) + continue + } + part, ok := item.(map[string]any) + if !ok { + return nil, fmt.Errorf("unsupported embedding input item type: %T", item) + } + partType, _ := part["type"].(string) + var content AliMultimodalContent + switch partType { + case "": + return nil, fmt.Errorf("multimodal input item missing required 'type' field") + case dto.ContentTypeText: + content.Text, _ = part[dto.ContentTypeText].(string) + case dto.ContentTypeImageURL: + url, err := parseURLField(part, dto.ContentTypeImageURL) + if err != nil { + return nil, err + } + content.Image = url + case dto.ContentTypeVideoUrl: + url, err := parseURLField(part, dto.ContentTypeVideoUrl) + if err != nil { + return nil, err + } + content.Video = url + default: + return nil, fmt.Errorf("unsupported multimodal input type: %s", partType) + } + aliReq.Input.Contents = append(aliReq.Input.Contents, content) + } + if request.Dimensions != nil { + aliReq.Parameters = &AliMultimodalEmbeddingParameters{Dimension: request.Dimensions} + } + return aliReq, nil +} + +func parseURLField(m map[string]any, key string) (string, error) { + urlObj, ok := m[key].(map[string]any) + if !ok { + return "", fmt.Errorf("invalid %s format", key) + } + url, _ := urlObj["url"].(string) + if url == "" { + return "", fmt.Errorf("%s.url is required", key) + } + return url, nil +} + +func MultimodalEmbeddingHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.Usage) { + responseBody, err := io.ReadAll(resp.Body) + service.CloseResponseBodyGracefully(resp) + if err != nil { + return types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError), nil + } + + var aliResponse AliMultimodalEmbeddingResponse + if err = common.Unmarshal(responseBody, &aliResponse); err != nil { + return types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError), nil + } + // The native endpoint can answer 200 with an error code in the body; the + // client must not receive that failure as a 200. + if aliResponse.Code != "" { + return types.WithOpenAIError(types.OpenAIError{ + Message: aliResponse.Message, + Type: aliResponse.Code, + Param: aliResponse.RequestId, + Code: aliResponse.Code, + }, http.StatusInternalServerError), nil + } + + data := make([]dto.EmbeddingResponseItem, 0, len(aliResponse.Output.Embeddings)) + for _, embedding := range aliResponse.Output.Embeddings { + data = append(data, dto.EmbeddingResponseItem{ + Object: "embedding", + Index: embedding.Index, + Embedding: embedding.Embedding, + }) + } + + // The vl-embedding models and multimodal-embedding-v1 put text in + // input_tokens and image and video tokens in image_tokens; + // tongyi-embedding-vision counts all input in input_tokens and splits it in + // input_tokens_details. Only input is priced, so output_tokens (also part of + // total_tokens) is never billed. A response without usage falls back to the + // request estimate, as in compatible mode. + u := aliResponse.Usage + var usage dto.Usage + if details := u.InputTokensDetails; details != nil { + usage.PromptTokens = u.InputTokens + usage.PromptTokensDetails.TextTokens = details.TextTokens + usage.PromptTokensDetails.ImageTokens = details.ImageTokens + } else { + usage.PromptTokens = u.InputTokens + u.ImageTokens + usage.PromptTokensDetails.TextTokens = u.InputTokens + usage.PromptTokensDetails.ImageTokens = u.ImageTokens + } + if usage.PromptTokens == 0 { + usage.PromptTokens = info.GetEstimatePromptTokens() + } + usage.TotalTokens = usage.PromptTokens + + embeddingResponse := dto.EmbeddingResponse{ + Object: "list", + Data: data, + Model: info.UpstreamModelName, + Usage: usage, + } + + jsonResponse, err := common.Marshal(embeddingResponse) + if err != nil { + return types.NewError(err, types.ErrorCodeBadResponseBody), nil + } + service.IOCopyBytesGracefully(c, resp, jsonResponse) + return nil, &usage +} diff --git a/relay/channel/gemini/adaptor.go b/relay/channel/gemini/adaptor.go index e817896b1..e99c64714 100644 --- a/relay/channel/gemini/adaptor.go +++ b/relay/channel/gemini/adaptor.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "net/http" + "slices" "strings" "github.com/QuantumNous/new-api/relay/channel" @@ -186,6 +187,20 @@ func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.Rela if request.Input == nil { return nil, errors.New("input is required") } + // Gemini embeds text only. The request estimate billed for this channel + // counts image, video and text parts, so a part must fail the request + // instead of being dropped and charged without reaching Gemini. + if items, ok := request.Input.([]any); ok && slices.ContainsFunc(items, func(item any) bool { + _, isPart := item.(map[string]any) + return isPart + }) { + return nil, types.NewErrorWithStatusCode( + errors.New("gemini embedding models accept only plain text input; images, videos and other content parts are not supported"), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } inputs := request.ParseInput() if len(inputs) == 0 { diff --git a/relay/channel/gemini/adaptor_responses_test.go b/relay/channel/gemini/adaptor_responses_test.go index 2d5576f37..ecb53672d 100644 --- a/relay/channel/gemini/adaptor_responses_test.go +++ b/relay/channel/gemini/adaptor_responses_test.go @@ -1,11 +1,13 @@ package gemini import ( + "net/http" "testing" "github.com/QuantumNous/new-api/common" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relaykit/dto" + "github.com/QuantumNous/new-api/relaykit/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" @@ -191,3 +193,26 @@ func mustGeminiRawMessage(t *testing.T, value any) []byte { require.NoError(t, err) return raw } + +// Gemini embeds plain text only. A content part fails the request with a 400 +// that skips retries, because the request estimate billed for this channel +// counts parts Gemini would never receive. +func TestConvertEmbeddingRequestRejectsContentParts(t *testing.T) { + imagePart := map[string]any{"type": "image_url", "image_url": map[string]any{"url": "https://example.com/cat.png"}} + textPart := map[string]any{"type": "text", "text": "a cat"} + for _, input := range []any{[]any{"a cat", imagePart}, []any{textPart}} { + info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-embedding-001"}} + _, err := (&Adaptor{}).ConvertEmbeddingRequest(nil, info, dto.EmbeddingRequest{Model: "gemini-embedding-001", Input: input}) + var apiErr *types.NewAPIError + require.ErrorAs(t, err, &apiErr) + assert.Equal(t, http.StatusBadRequest, apiErr.StatusCode) + assert.True(t, types.IsSkipRetryError(apiErr)) + } + + info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-embedding-001"}} + got, err := (&Adaptor{}).ConvertEmbeddingRequest(nil, info, dto.EmbeddingRequest{Model: "gemini-embedding-001", Input: []any{"a cat", "a dog"}}) + require.NoError(t, err) + wire, err := common.Marshal(got) + require.NoError(t, err) + assert.JSONEq(t, `["a cat","a dog"]`, gjson.GetBytes(wire, "requests.#.content.parts.0.text").Raw) +} diff --git a/relaykit/dto/embedding.go b/relaykit/dto/embedding.go index 65a6b54b0..c709a4783 100644 --- a/relaykit/dto/embedding.go +++ b/relaykit/dto/embedding.go @@ -4,6 +4,7 @@ import ( "net/http" "strings" + kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil" "github.com/QuantumNous/new-api/relaykit/types" ) @@ -32,15 +33,36 @@ type EmbeddingRequest struct { } func (r *EmbeddingRequest) GetTokenCountMeta() *types.TokenCountMeta { - var texts = make([]string, 0) - - inputs := r.ParseInput() - for _, input := range inputs { - texts = append(texts, input) + texts := r.ParseInput() + var files []*types.FileMeta + // Chat-style parts count like chat message content: text parts as text, + // image and video parts as files. + items, _ := r.Input.([]any) + for _, item := range items { + if _, ok := item.(map[string]any); !ok { + continue + } + part, err := kitutil.Any2Type[MediaContent](item) + if err != nil { + continue + } + switch part.Type { + case ContentTypeText: + texts = append(texts, part.Text) + case ContentTypeImageURL: + if source := part.ToFileSource(); source != nil { + files = append(files, &types.FileMeta{FileType: types.FileTypeImage, Source: source, Detail: part.GetImageMedia().Detail}) + } + case ContentTypeVideoUrl: + if source := part.ToFileSource(); source != nil { + files = append(files, &types.FileMeta{FileType: types.FileTypeVideo, Source: source}) + } + } } return &types.TokenCountMeta{ CombineText: strings.Join(texts, "\n"), + Files: files, } } diff --git a/relaykit/dto/token_count_meta_test.go b/relaykit/dto/token_count_meta_test.go index 4aca6a054..4a63f8ca1 100644 --- a/relaykit/dto/token_count_meta_test.go +++ b/relaykit/dto/token_count_meta_test.go @@ -98,3 +98,20 @@ func TestTokenCountMetaCountsAgentLoopContent(t *testing.T) { }) } } + +// Embedding input can mix bare strings with chat-style parts; text parts count +// as text and image and video parts as files, like chat message content. +func TestEmbeddingTokenCountMetaCountsChatStyleParts(t *testing.T) { + const imageData = "data:image/png;base64,iVBORw0KGgo=" + var request EmbeddingRequest + require.NoError(t, kitutil.UnmarshalJsonStr(`{"model":"m","input":["a cat",{"type":"text","text":"on a mat"},`+ + `{"type":"image_url","image_url":{"url":"`+imageData+`","detail":"low"}},`+ + `{"type":"video_url","video_url":{"url":"https://example.com/cat.mp4"}}]}`, &request)) + + meta := request.GetTokenCountMeta() + assert.Equal(t, "a cat\non a mat", meta.CombineText) + assert.Equal(t, []*types.FileMeta{ + {FileType: types.FileTypeImage, Source: types.NewBase64FileSource(imageData, ""), Detail: "low"}, + {FileType: types.FileTypeVideo, Source: types.NewURLFileSource("https://example.com/cat.mp4")}, + }, meta.Files) +}