完善文档页文本向量与重排序调用支持

2026-05-31 21:18:41 +08:00
parent 8ee7a7969e
commit 644a6f9d17
24 changed files with 1945 additions and 71 deletions
@@ -151,6 +151,107 @@ func TestOpenAIClientChatContract(t *testing.T) {
 	}
 }

+func TestOpenAIClientEmbeddingsContract(t *testing.T) {
+	var gotPath string
+	var gotModel string
+	var gotDimensions float64
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotPath = r.URL.Path
+		var body map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			t.Fatalf("decode request: %v", err)
+		}
+		gotModel, _ = body["model"].(string)
+		gotDimensions, _ = body["dimensions"].(float64)
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"id":     "embd-test",
+			"object": "list",
+			"model":  gotModel,
+			"data": []any{map[string]any{
+				"object":    "embedding",
+				"index":     0,
+				"embedding": []any{0.1, 0.2, 0.3},
+			}},
+			"usage": map[string]any{"prompt_tokens": 3, "total_tokens": 3},
+		})
+	}))
+	defer server.Close()
+
+	response, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
+		Kind:  "embeddings",
+		Model: "aliyun-bailian-openai:text-embedding-v4",
+		Body: map[string]any{
+			"model":      "aliyun-bailian-openai:text-embedding-v4",
+			"input":      []any{"hello"},
+			"dimensions": 3,
+		},
+		Candidate: store.RuntimeModelCandidate{
+			BaseURL:           server.URL,
+			ProviderModelName: "text-embedding-v4",
+			Credentials:       map[string]any{"apiKey": "test-key"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("run embeddings client: %v", err)
+	}
+	if gotPath != "/embeddings" || gotModel != "text-embedding-v4" || gotDimensions != 3 {
+		t.Fatalf("unexpected embeddings request path=%s model=%s dimensions=%v", gotPath, gotModel, gotDimensions)
+	}
+	if response.Usage.InputTokens != 3 || response.Usage.TotalTokens != 3 || response.Result["id"] != "embd-test" {
+		t.Fatalf("unexpected embeddings response: %+v", response)
+	}
+}
+
+func TestOpenAIClientAliyunRerankUsesCompatibleAPIBase(t *testing.T) {
+	var gotPath string
+	var gotModel string
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotPath = r.URL.Path
+		var body map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			t.Fatalf("decode request: %v", err)
+		}
+		gotModel, _ = body["model"].(string)
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"id":     "rerank-test",
+			"object": "list",
+			"model":  gotModel,
+			"results": []any{
+				map[string]any{"index": 0, "relevance_score": 0.93},
+				map[string]any{"index": 2, "relevance_score": 0.34},
+			},
+			"usage": map[string]any{"total_tokens": 9},
+		})
+	}))
+	defer server.Close()
+
+	response, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
+		Kind:  "reranks",
+		Model: "aliyun-bailian-openai:qwen3-rerank",
+		Body: map[string]any{
+			"model":     "aliyun-bailian-openai:qwen3-rerank",
+			"query":     "what is rerank",
+			"documents": []any{"rerank sorts documents", "unrelated"},
+			"top_n":     2,
+		},
+		Candidate: store.RuntimeModelCandidate{
+			Provider:          "aliyun-bailian-openai",
+			BaseURL:           server.URL + "/compatible-mode/v1",
+			ProviderModelName: "qwen3-rerank",
+			Credentials:       map[string]any{"apiKey": "test-key"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("run rerank client: %v", err)
+	}
+	if gotPath != "/compatible-api/v1/reranks" || gotModel != "qwen3-rerank" {
+		t.Fatalf("unexpected rerank request path=%s model=%s", gotPath, gotModel)
+	}
+	if response.Usage.TotalTokens != 9 || response.Result["id"] != "rerank-test" {
+		t.Fatalf("unexpected rerank response: %+v", response)
+	}
+}
+
 func TestOpenAIClientChatRequestNormalizesToolContext(t *testing.T) {
 	var captured map[string]any
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {