BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
wangbo
2026-05-31 21:18:41 +08:00
parent 8ee7a7969e
commit 644a6f9d17
24 changed files with 1945 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/clients/clients_test.go
+101
View File
@@ -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) {
apps/api/internal/clients/openai.go
+26 -2
View File
@@ -7,6 +7,8 @@ import (
"net/http"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type OpenAIClient struct {
@@ -27,10 +29,10 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
body = NormalizeChatCompletionRequestBody(body)
}
body["model"] = upstreamModelName(request.Candidate)
stream := request.Stream || boolValue(body, "stream")
stream := openAIEndpointSupportsStream(request.Kind) && (request.Stream || boolValue(body, "stream"))
ensureOpenAIStreamUsage(body, request.Kind, stream)
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(request.Candidate.BaseURL, endpoint), bytes.NewReader(raw))
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(openAIBaseURL(request.Kind, request.Candidate), endpoint), bytes.NewReader(raw))
if err != nil {
return Response{}, err
}
@@ -81,6 +83,10 @@ func openAIEndpoint(kind string) string {
return "/chat/completions"
case "responses":
return "/responses"
case "embeddings":
return "/embeddings"
case "reranks":
return "/reranks"
case "images.generations":
return "/images/generations"
case "images.edits":
@@ -90,6 +96,24 @@ func openAIEndpoint(kind string) string {
}
}
func openAIEndpointSupportsStream(kind string) bool {
return kind == "chat.completions" || kind == "responses"
}
func openAIBaseURL(kind string, candidate store.RuntimeModelCandidate) string {
base := strings.TrimSpace(candidate.BaseURL)
if kind != "reranks" {
return base
}
if strings.Contains(base, "/compatible-mode/") && (strings.EqualFold(candidate.Provider, "aliyun-bailian-openai") || strings.Contains(base, "dashscope")) {
return strings.Replace(base, "/compatible-mode/", "/compatible-api/", 1)
}
if base == "" && strings.EqualFold(candidate.Provider, "aliyun-bailian-openai") {
return "https://dashscope.aliyuncs.com/compatible-api/v1"
}
return base
}
func cloneBody(body map[string]any) map[string]any {
out := map[string]any{}
for key, value := range body {
apps/api/internal/clients/simulation.go
+107
View File
@@ -131,6 +131,10 @@ func simulatedResult(request Request) map[string]any {
"output_text": fmt.Sprintf("simulation response from %s", request.Candidate.Provider),
"usage": map[string]any{"input_tokens": 12, "output_tokens": 8, "total_tokens": 20},
}
case "embeddings":
return simulatedEmbeddingResult(request)
case "reranks":
return simulatedRerankResult(request)
case "images.edits":
return map[string]any{
"id": "img-edit-simulated",
@@ -172,6 +176,106 @@ func simulatedResult(request Request) map[string]any {
}
}
func simulatedEmbeddingResult(request Request) map[string]any {
inputCount := simulatedEmbeddingInputCount(request.Body["input"])
dimensions := intValue(request.Body, "dimensions", 3)
if dimensions <= 0 {
dimensions = 3
}
if dimensions > 2048 {
dimensions = 2048
}
data := make([]any, 0, inputCount)
for index := 0; index < inputCount; index += 1 {
embedding := make([]any, 0, dimensions)
for dimension := 0; dimension < dimensions; dimension += 1 {
embedding = append(embedding, float64(index+1)/10+float64(dimension)/100)
}
data = append(data, map[string]any{
"object": "embedding",
"index": index,
"embedding": embedding,
})
}
usage := simulatedUsage(request)
return map[string]any{
"id": "embd-simulated",
"object": "list",
"model": request.Model,
"data": data,
"usage": map[string]any{"prompt_tokens": usage.InputTokens, "total_tokens": usage.TotalTokens},
}
}
func simulatedEmbeddingInputCount(value any) int {
switch typed := value.(type) {
case []any:
if len(typed) > 0 {
return len(typed)
}
case []string:
if len(typed) > 0 {
return len(typed)
}
}
return 1
}
func simulatedRerankResult(request Request) map[string]any {
documents := simulatedRerankDocuments(request.Body["documents"])
topN := intValue(request.Body, "top_n", len(documents))
if topN <= 0 || topN > len(documents) {
topN = len(documents)
}
results := make([]any, 0, topN)
for index := 0; index < topN; index += 1 {
score := 0.95 - float64(index)*0.1
if score < 0 {
score = 0
}
result := map[string]any{
"index": index,
"relevance_score": score,
}
if boolValue(request.Body, "return_documents") {
result["document"] = map[string]any{"text": documents[index]}
}
results = append(results, result)
}
usage := simulatedUsage(request)
return map[string]any{
"id": "rerank-simulated",
"object": "list",
"model": request.Model,
"results": results,
"usage": map[string]any{"total_tokens": usage.TotalTokens},
}
}
func simulatedRerankDocuments(value any) []string {
switch typed := value.(type) {
case []any:
out := make([]string, 0, len(typed))
for _, item := range typed {
text := stringValue(map[string]any{"value": item}, "value")
if text == "" {
if record, ok := item.(map[string]any); ok {
text = firstNonEmptyString(stringValue(record, "text"), stringValue(record, "content"))
}
}
out = append(out, text)
}
if len(out) > 0 {
return out
}
case []string:
if len(typed) > 0 {
return typed
}
}
return []string{"simulated document"}
}
func simulatedImageData(request Request, url string, fallbackPrompt string) []any {
count := simulatedOutputCount(request.Body)
items := make([]any, 0, count)
@@ -207,6 +311,9 @@ func simulatedUsage(request Request) Usage {
if request.ModelType == "chat" || request.ModelType == "text_generate" || request.Kind == "responses" {
return Usage{InputTokens: 12, OutputTokens: 8, TotalTokens: 20}
}
if request.ModelType == "text_embedding" || request.ModelType == "text_rerank" || request.Kind == "embeddings" || request.Kind == "reranks" {
return Usage{InputTokens: 16, TotalTokens: 16}
}
return Usage{}
}