diff --git a/conf/models/ppio.json b/conf/models/ppio.json index 31a03913f9..24a40d8f5b 100644 --- a/conf/models/ppio.json +++ b/conf/models/ppio.json @@ -6,7 +6,9 @@ }, "url_suffix": { "chat": "chat/completions", - "models": "models" + "models": "models", + "embedding": "embeddings", + "rerank": "rerank" }, "class": "ppio", "models": [ @@ -219,6 +221,30 @@ "tools": { "support": true } + }, + { + "name": "qwen/qwen3-embedding-8b", + "model_types": [ + "embedding" + ] + }, + { + "name": "baai/bge-m3", + "model_types": [ + "embedding" + ] + }, + { + "name": "qwen/qwen3-embedding-0.6b", + "model_types": [ + "embedding" + ] + }, + { + "name": "baai/bge-reranker-v2-m3", + "model_types": [ + "rerank" + ] } ] } diff --git a/internal/entity/models/model_test.go b/internal/entity/models/model_test.go index fa5029d540..f74ddb3039 100644 --- a/internal/entity/models/model_test.go +++ b/internal/entity/models/model_test.go @@ -269,12 +269,12 @@ func TestPPIOProviderConfigLoadsIntoProviderManager(t *testing.T) { if provider.ModelDriver.Name() != "ppio" { t.Errorf("ModelDriver.Name()=%q", provider.ModelDriver.Name()) } - if len(provider.Models) != 21 { - t.Fatalf("PPIO model count=%d, want 21", len(provider.Models)) + if len(provider.Models) != 25 { + t.Fatalf("PPIO model count=%d, want 25", len(provider.Models)) } for _, model := range provider.Models { - if !model.ModelTypeMap["chat"] { - t.Errorf("model %q missing chat type map", model.Name) + if len(model.ModelTypes) == 0 { + t.Errorf("model %q missing model types", model.Name) } if model.Class == nil || *model.Class != "PPIO" { t.Errorf("model %q class=%v", model.Name, model.Class) @@ -285,8 +285,8 @@ func TestPPIOProviderConfigLoadsIntoProviderManager(t *testing.T) { if err != nil { t.Fatalf("ListModels: %v", err) } - if len(models) != 21 { - t.Errorf("ListModels count=%d, want 21", len(models)) + if len(models) != 25 { + t.Errorf("ListModels count=%d, want 25", len(models)) } model, err := pm.GetModelByName("ppio", "deepseek/deepseek-r1") @@ -310,6 +310,23 @@ func TestPPIOProviderConfigLoadsIntoProviderManager(t *testing.T) { if *model.MaxTokens != 1048576 { t.Errorf("deepseek/deepseek-v4-flash max_tokens=%d", *model.MaxTokens) } + if !model.ModelTypeMap["chat"] { + t.Errorf("deepseek/deepseek-v4-flash missing chat type map") + } + model, err = pm.GetModelByName("ppio", "qwen/qwen3-embedding-8b") + if err != nil { + t.Fatalf("GetModelByName qwen/qwen3-embedding-8b: %v", err) + } + if !model.ModelTypeMap["embedding"] { + t.Errorf("qwen/qwen3-embedding-8b missing embedding type map") + } + model, err = pm.GetModelByName("ppio", "baai/bge-reranker-v2-m3") + if err != nil { + t.Fatalf("GetModelByName baai/bge-reranker-v2-m3: %v", err) + } + if !model.ModelTypeMap["rerank"] { + t.Errorf("baai/bge-reranker-v2-m3 missing rerank type map") + } resp := pm.SearchByType("chat") if resp.Code != 0 { diff --git a/internal/entity/models/ppio.go b/internal/entity/models/ppio.go index 6096dd2e4d..cc2034c454 100644 --- a/internal/entity/models/ppio.go +++ b/internal/entity/models/ppio.go @@ -52,6 +52,23 @@ func (p *PPIOModel) Name() string { return "ppio" } +type PPIOChatResponse struct { + ID string `json:"id"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + } `json:"message"` + } `json:"choices"` + Created int `json:"created"` + Model string `json:"model"` + Object string `json:"object"` + Usage TokenUsage `json:"usage"` +} + func (p *PPIOModel) endpoint(apiConfig *APIConfig, suffix string) (string, error) { baseURL, err := p.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -64,7 +81,6 @@ func (p *PPIOModel) endpoint(apiConfig *APIConfig, suffix string) (string, error type ppioChatMessage struct { Content string `json:"content"` ReasoningContent string `json:"reasoning_content"` - Reasoning string `json:"reasoning"` } type ppioChatChoice struct { @@ -74,9 +90,15 @@ type ppioChatChoice struct { } type ppioChatResponse struct { - Choices []ppioChatChoice `json:"choices"` - Error interface{} `json:"error"` - FinishReason string `json:"finish_reason"` + ID string `json:"id"` + Choices []ppioChatChoice `json:"choices"` + Error interface{} `json:"error"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` + FinishReason string `json:"finish_reason"` } func (p *PPIOModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { @@ -124,27 +146,32 @@ func (p *PPIOModel) ChatWithMessages(ctx context.Context, modelName string, mess return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - var result ppioChatResponse - if err = json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to parse response: %w", err) - } - if result.Error != nil { - return nil, fmt.Errorf("ppio: upstream error: %v", result.Error) - } - if len(result.Choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result PPIOChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) + } + if len(result.Choices) == 0 { + var errResp struct { + Error interface{} `json:"error"` + } + if err := json.Unmarshal(body, &errResp); err == nil && errResp.Error != nil { + return chatResponseParts{}, fmt.Errorf("ppio: upstream error: %v", errResp.Error) + } + return chatResponseParts{}, fmt.Errorf("no choices in response") + } - content := result.Choices[0].Message.Content - reasonContent := result.Choices[0].Message.ReasoningContent - if reasonContent == "" { - reasonContent = result.Choices[0].Message.Reasoning - } + choice := &result.Choices[0] + content := choice.Message.Content + reasonContent := choice.Message.ReasoningContent - return &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - }, nil + return chatResponseParts{ + RequestID: result.ID, + Content: &content, + ReasonContent: &reasonContent, + Usage: &result.Usage, + }, nil + }) } func (p *PPIOModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -170,7 +197,10 @@ func (p *PPIOModel) ChatStreamlyWithSender(ctx context.Context, modelName string return err } - jsonData, err := json.Marshal(buildRequestBody(chatModelConfig, modelName, messages, true)) + reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} + + jsonData, err := json.Marshal(reqBody) if err != nil { return fmt.Errorf("failed to marshal request: %w", err) } @@ -202,15 +232,19 @@ func (p *PPIOModel) ChatStreamlyWithSender(ctx context.Context, modelName string if event.Error != nil { return fmt.Errorf("ppio: upstream stream error: %v", event.Error) } + if event.Usage.TotalTokens > 0 || event.Usage.PromptTokens > 0 || event.Usage.CompletionTokens > 0 { + applyStreamUsage(chatModelConfig, modelUsage, &TokenUsage{ + PromptTokens: event.Usage.PromptTokens, + CompletionTokens: event.Usage.CompletionTokens, + TotalTokens: event.Usage.TotalTokens, + }) + } if len(event.Choices) == 0 { return nil } choice := event.Choices[0] reasoning := choice.Delta.ReasoningContent - if reasoning == "" { - reasoning = choice.Delta.Reasoning - } if reasoning != "" { if err := sender(nil, &reasoning); err != nil { return err @@ -296,12 +330,207 @@ func (p *PPIOModel) CheckConnection(ctx context.Context, apiConfig *APIConfig) e return err } +type PPIOEmbeddingData struct { + EmbeddingData + Object string `json:"object"` +} + +type PPIOEmbeddingResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Data []PPIOEmbeddingData `json:"data"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + func (p *PPIOModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { - return nil, fmt.Errorf("%s, no such method", p.Name()) + if err := p.baseModel.APIConfigCheck(apiConfig); err != nil { + return nil, err + } + + if len(texts) == 0 { + return []EmbeddingData{}, nil + } + + if modelName == nil || *modelName == "" { + return nil, fmt.Errorf("model name is required") + } + + baseURL, err := p.baseModel.GetBaseURL(apiConfig) + if err != nil { + return nil, err + } + baseURL = strings.TrimSuffix(baseURL, "/") + url := fmt.Sprintf("%s/%s", baseURL, p.baseModel.URLSuffix.Embedding) + + reqBody := map[string]interface{}{ + "model": *modelName, + "input": texts, + "encoding_format": "float", + } + if embeddingConfig != nil && embeddingConfig.EncodingFormat != "" { + reqBody["encoding_format"] = embeddingConfig.EncodingFormat + } + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + + resp, err := p.baseModel.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("PPIO embeddings API error: %s, body: %s", resp.Status, string(body)) + } + + var parsed PPIOEmbeddingResponse + if err = json.Unmarshal(body, &parsed); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + embeddings := make([]EmbeddingData, len(texts)) + filled := make([]bool, len(texts)) + for _, item := range parsed.Data { + if item.Index < 0 || item.Index >= len(texts) { + return nil, fmt.Errorf("ppio: response index %d out of range for %d inputs", item.Index, len(texts)) + } + if filled[item.Index] { + return nil, fmt.Errorf("ppio: duplicate embedding index %d in response", item.Index) + } + embeddings[item.Index] = EmbeddingData{ + Embedding: item.Embedding, + Index: item.Index, + } + filled[item.Index] = true + } + for i, ok := range filled { + if !ok { + return nil, fmt.Errorf("ppio: missing embedding for input index %d", i) + } + } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "embedding") + + return embeddings, nil } func (p *PPIOModel) Rerank(ctx context.Context, modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig, modelUsage *common.ModelUsage) (*RerankResponse, error) { - return nil, fmt.Errorf("%s, no such method", p.Name()) + if err := p.baseModel.APIConfigCheck(apiConfig); err != nil { + return nil, err + } + + if len(documents) == 0 { + return &RerankResponse{}, nil + } + if modelName == nil || *modelName == "" { + return nil, fmt.Errorf("model name is required") + } + + resolvedBaseURL, err := p.baseModel.GetBaseURL(apiConfig) + if err != nil { + return nil, err + } + url := fmt.Sprintf("%s/%s", strings.TrimSuffix(resolvedBaseURL, "/"), p.baseModel.URLSuffix.Rerank) + + topN := len(documents) + if rerankConfig != nil && rerankConfig.TopN > 0 { + topN = rerankConfig.TopN + } + + reqBody := map[string]any{ + "model": *modelName, + "query": query, + "documents": documents, + "top_n": topN, + } + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + + resp, err := p.baseModel.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("PPIO rerank API error: status %d, body: %s", resp.StatusCode, string(body)) + } + + var rerankResp struct { + ID string `json:"id"` + Results []struct { + Index int `json:"index"` + RelevanceScore float64 `json:"relevance_score"` + } `json:"results"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` + } + + if err = json.Unmarshal(body, &rerankResp); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) + } + + var rerankResponse RerankResponse + for _, result := range rerankResp.Results { + rerankResult := RerankResult{ + Index: result.Index, + RelevanceScore: result.RelevanceScore, + } + rerankResponse.Data = append(rerankResponse.Data, rerankResult) + } + recordResponseUsage(modelUsage, rerankResp.ID, &TokenUsage{ + PromptTokens: rerankResp.Usage.PromptTokens, + CompletionTokens: rerankResp.Usage.CompletionTokens, + TotalTokens: rerankResp.Usage.TotalTokens, + }, "rerank") + + return &rerankResponse, nil } func (p *PPIOModel) Balance(ctx context.Context, apiConfig *APIConfig) (map[string]interface{}, error) { diff --git a/internal/entity/models/ppio_test.go b/internal/entity/models/ppio_test.go index 3dede3b64b..4b248abacd 100644 --- a/internal/entity/models/ppio_test.go +++ b/internal/entity/models/ppio_test.go @@ -5,6 +5,7 @@ import ( "io" "net/http" "net/http/httptest" + "ragflow/internal/common" "strings" "testing" ) @@ -40,7 +41,7 @@ func newPPIOServer(t *testing.T, handler func(t *testing.T, r *http.Request, bod func newPPIOForTest(baseURL string) *PPIOModel { return NewPPIOModel( map[string]string{"default": baseURL}, - URLSuffix{Chat: "chat/completions", Models: "models"}, + URLSuffix{Chat: "chat/completions", Models: "models", Embedding: "embeddings", Rerank: "rerank"}, ) } @@ -114,12 +115,18 @@ func TestPPIOChatHappyPath(t *testing.T) { t.Errorf("message=%#v", first) } _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "id": "chat-ppio", "choices": []map[string]interface{}{{ "message": map[string]interface{}{ "content": "pong", "reasoning_content": "thinking", }, }}, + "usage": map[string]interface{}{ + "prompt_tokens": 3, + "completion_tokens": 5, + "total_tokens": 8, + }, }) }) defer srv.Close() @@ -130,13 +137,14 @@ func TestPPIOChatHappyPath(t *testing.T) { topP := 0.9 stop := []string{"END"} effort := "high" + usage := &common.ModelUsage{} resp, err := newPPIOForTest(srv.URL).ChatWithMessages( ctx, "deepseek/deepseek-r1", []Message{{Role: "user", Content: "ping"}}, &APIConfig{ApiKey: &apiKey}, &ChatConfig{MaxTokens: &mt, Temperature: &temp, TopP: &topP, Stop: &stop, Effort: &effort}, - nil, + usage, ) if err != nil { t.Fatalf("ChatWithMessages: %v", err) @@ -147,37 +155,10 @@ func TestPPIOChatHappyPath(t *testing.T) { if *resp.ReasonContent != "thinking" { t.Errorf("ReasonContent=%q", *resp.ReasonContent) } -} - -func TestPPIOChatUsesReasoningFallback(t *testing.T) { - ctx := t.Context() - srv := newPPIOServer(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { - _ = json.NewEncoder(w).Encode(map[string]interface{}{ - "choices": []map[string]interface{}{{ - "message": map[string]interface{}{ - "content": "pong", - "reasoning": "fallback reasoning", - }, - }}, - }) - }) - defer srv.Close() - - apiKey := "test-key" - resp, err := newPPIOForTest(srv.URL).ChatWithMessages( - ctx, - "deepseek/deepseek-r1", - []Message{{Role: "user", Content: "ping"}}, - &APIConfig{ApiKey: &apiKey}, - nil, - nil, - ) - if err != nil { - t.Fatalf("ChatWithMessages: %v", err) - } - if *resp.ReasonContent != "fallback reasoning" { - t.Errorf("ReasonContent=%q", *resp.ReasonContent) + if resp.Usage == nil || resp.Usage.PromptTokens != 3 || resp.Usage.CompletionTokens != 5 || resp.Usage.TotalTokens != 8 { + t.Fatalf("Usage=%#v, want prompt=3 completion=5 total=8", resp.Usage) } + assertModelUsage(t, usage, 3, 5, 8) } func TestPPIOChatRequiresModelName(t *testing.T) { @@ -237,15 +218,19 @@ func TestPPIOStreamHappyPath(t *testing.T) { if body["stream"] != true { t.Errorf("stream=%v want true", body["stream"]) } + streamOptions, ok := body["stream_options"].(map[string]interface{}) + if !ok || streamOptions["include_usage"] != true { + t.Errorf("stream_options=%#v, want include_usage=true", body["stream_options"]) + } if got := r.Header.Get("Accept"); got != "text/event-stream" { t.Errorf("Accept=%q", got) } w.Header().Set("Content-Type", "text/event-stream") _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"reasoning_content":"think "}}]}`+"\n"+ - `data: {"choices":[{"delta":{"reasoning":"fallback "}}]}`+"\n"+ `data: {"choices":[{"delta":{"content":"Hello"}}]}`+"\n"+ - `data: {"choices":[{"delta":{"content":" world"},"finish_reason":"stop"}]}`+"\n", + `data: {"choices":[{"delta":{"content":" world"},"finish_reason":"stop"}]}`+"\n"+ + `data: {"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`+"\n", ) }) defer srv.Close() @@ -253,11 +238,13 @@ func TestPPIOStreamHappyPath(t *testing.T) { apiKey := "test-key" var content []string var reasoning []string + cfg := &ChatConfig{} + usage := &common.ModelUsage{} err := newPPIOForTest(srv.URL).ChatStreamlyWithSender( ctx, "deepseek/deepseek-r1", []Message{{Role: "user", Content: "hi"}}, - &APIConfig{ApiKey: &apiKey}, nil, nil, + &APIConfig{ApiKey: &apiKey}, cfg, usage, func(c *string, r *string) error { if c != nil { content = append(content, *c) @@ -274,12 +261,97 @@ func TestPPIOStreamHappyPath(t *testing.T) { if strings.Join(content, "") != "Hello world[DONE]" { t.Errorf("content=%q", strings.Join(content, "")) } - if strings.Join(reasoning, "") != "think fallback " { + if strings.Join(reasoning, "") != "think " { t.Errorf("reasoning=%q", strings.Join(reasoning, "")) } if len(content) == 0 || content[len(content)-1] != "[DONE]" { t.Errorf("final content sentinel missing: %#v", content) } + if cfg.UsageResult == nil || cfg.UsageResult.PromptTokens != 3 || cfg.UsageResult.CompletionTokens != 5 || cfg.UsageResult.TotalTokens != 8 { + t.Fatalf("UsageResult=%#v, want prompt=3 completion=5 total=8", cfg.UsageResult) + } + if usage.InputTokens != 3 || usage.OutputTokens != 5 || usage.TotalTokens != 8 { + t.Fatalf("stream usage=(%d,%d,%d), want (3,5,8)", usage.InputTokens, usage.OutputTokens, usage.TotalTokens) + } +} + +func TestPPIOEmbedRecordsUsage(t *testing.T) { + ctx := t.Context() + srv := newPPIOServer(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { + if r.URL.Path != "/embeddings" { + t.Errorf("path=%s", r.URL.Path) + } + if body["model"] != "embedding-model" { + t.Errorf("model=%v", body["model"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "id": "embed-ppio", + "data": []map[string]interface{}{{ + "embedding": []float64{0.1, 0.2}, + "index": 0, + }}, + "usage": map[string]interface{}{ + "prompt_tokens": 7, + "total_tokens": 7, + }, + }) + }) + defer srv.Close() + + apiKey := "test-key" + modelName := "embedding-model" + usage := &common.ModelUsage{} + embeddings, err := newPPIOForTest(srv.URL).Embed(ctx, &modelName, []string{"document"}, &APIConfig{ApiKey: &apiKey}, nil, usage) + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(embeddings) != 1 || len(embeddings[0].Embedding) != 2 { + t.Fatalf("embeddings=%#v", embeddings) + } + assertModelUsage(t, usage, 7, 0, 7) + if usage.Type != "embedding" { + t.Fatalf("usage type=%q, want embedding", usage.Type) + } +} + +func TestPPIORerankRecordsUsage(t *testing.T) { + ctx := t.Context() + srv := newPPIOServer(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { + if r.URL.Path != "/rerank" { + t.Errorf("path=%s", r.URL.Path) + } + if body["model"] != "rerank-model" { + t.Errorf("model=%v", body["model"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "id": "rerank-ppio", + "results": []map[string]interface{}{{ + "index": 0, + "relevance_score": 0.9, + }}, + "usage": map[string]interface{}{ + "prompt_tokens": 7, + "completion_tokens": 2, + "total_tokens": 9, + }, + }) + }) + defer srv.Close() + + apiKey := "test-key" + modelName := "rerank-model" + usage := &common.ModelUsage{} + reranked, err := newPPIOForTest(srv.URL).Rerank(ctx, &modelName, "query", []string{"document"}, &APIConfig{ApiKey: &apiKey}, &RerankConfig{TopN: 1}, usage) + if err != nil { + t.Fatalf("Rerank: %v", err) + } + if len(reranked.Data) != 1 || reranked.Data[0].Index != 0 { + t.Fatalf("reranked=%#v", reranked) + } + assertModelUsage(t, usage, 7, 2, 9) + if usage.Type != "rerank" { + t.Fatalf("usage type=%q, want rerank", usage.Type) + } } func TestPPIOStreamSurfacesHTTPError(t *testing.T) { @@ -498,12 +570,6 @@ func TestPPIOMissingRegionBaseURL(t *testing.T) { func TestPPIOUnsupportedMethods(t *testing.T) { ctx := t.Context() m := newPPIOForTest("http://unused") - if _, err := m.Embed(ctx, nil, nil, nil, nil, nil); err == nil || !strings.Contains(err.Error(), "no such method") { - t.Errorf("Embed error=%v", err) - } - if _, err := m.Rerank(ctx, nil, "", nil, nil, nil, nil); err == nil || !strings.Contains(err.Error(), "no such method") { - t.Errorf("Rerank error=%v", err) - } if _, err := m.Balance(ctx, nil); err == nil || !strings.Contains(err.Error(), "no such method") { t.Errorf("Balance error=%v", err) } diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index e5b4c09345..4f4d555c9d 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -191,7 +191,8 @@ type APIConfig struct { } type EmbeddingConfig struct { - Dimension int + Dimension int + EncodingFormat string } type RerankConfig struct {