Go: implement embed, rerank for PPIO provider (#17486)

### Summary As title #17284 #### verified from CLI ``` RAGFlow(api/default)> embed text 'walkerwhat' 'jumperwho' with 'qwen/qwen3-embedding-0.6b@test@ppio' dimension 16 +-----------+-------+ | dimension | index | +-----------+-------+ | 1024 | 0 | | 1024 | 1 | +-----------+-------+ RAGFlow(api/default)> rerank query 'what is rag' document 'rag is retrieval augment generation' 'rag need llm' 'famous rag project includes ragflow' with 'baai/bge-reranker-v2-m3@test@ppio' top 3 +-------+-----------------+ | index | relevance_score | +-------+-----------------+ | 0 | 0.9830034 | | 2 | 0.06399203 | | 1 | 0.04665664 | +-------+-----------------+ ```
2026-07-29 12:09:31 +08:00 · 2026-07-28 19:16:31 +08:00
parent e5c038a411
commit e0ad4f8339
5 changed files with 418 additions and 79 deletions
--- 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"
+      ]
    }
  ]
 }
--- 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 {
--- 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) {
--- 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)
 	}
--- 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 {