mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-14 04:36:52 +08:00
327 lines
12 KiB
Go
327 lines
12 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package models
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"reflect"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
)
|
|
|
|
func newMWSTestDriver(serverURL string) (*MWSModel, *APIConfig) {
|
|
token := "token"
|
|
baseURL := serverURL + "/projects/test-project"
|
|
driver := NewMWSModel(map[string]string{"default": baseURL}, URLSuffix{})
|
|
return driver, &APIConfig{ApiKey: &token, BaseURL: &baseURL}
|
|
}
|
|
|
|
func decodeMWSRequest(t *testing.T, request *http.Request) map[string]any {
|
|
t.Helper()
|
|
if request.Method != http.MethodPost {
|
|
t.Fatalf("unexpected method: %s", request.Method)
|
|
}
|
|
if request.Header.Get("Authorization") != "Bearer token" {
|
|
t.Fatalf("unexpected authorization: %q", request.Header.Get("Authorization"))
|
|
}
|
|
if request.Header.Get("Content-Type") != "application/json" {
|
|
t.Fatalf("unexpected content type: %q", request.Header.Get("Content-Type"))
|
|
}
|
|
body, err := io.ReadAll(request.Body)
|
|
if err != nil {
|
|
t.Fatalf("read request body: %v", err)
|
|
}
|
|
var payload map[string]any
|
|
if err = json.Unmarshal(body, &payload); err != nil {
|
|
t.Fatalf("decode request body: %v", err)
|
|
}
|
|
return payload
|
|
}
|
|
|
|
func TestNormalizeMWSProjectURL(t *testing.T) {
|
|
got, err := normalizeMWSProjectURL("https://gpt.mwsapis.ru/projects/demo/")
|
|
if err != nil {
|
|
t.Fatalf("normalize URL: %v", err)
|
|
}
|
|
if got != "https://gpt.mwsapis.ru/projects/demo" {
|
|
t.Fatalf("unexpected normalized URL: %s", got)
|
|
}
|
|
if _, err = normalizeMWSProjectURL("https://gpt.mwsapis.ru/projects/demo/openai/v1"); err == nil {
|
|
t.Fatal("expected a non-root URL to be rejected")
|
|
}
|
|
}
|
|
|
|
func TestMWSListModelsUsesOpenAIEndpointAndFiltersTypes(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
|
if request.Method != http.MethodGet || request.URL.Path != "/projects/test-project/openai/v1/models" {
|
|
t.Fatalf("unexpected request: %s %s", request.Method, request.URL.Path)
|
|
}
|
|
if request.Header.Get("Authorization") != "Bearer token" {
|
|
t.Fatalf("unexpected authorization: %q", request.Header.Get("Authorization"))
|
|
}
|
|
body, _ := io.ReadAll(request.Body)
|
|
if len(body) != 0 {
|
|
t.Fatalf("GET models request must not have a body: %q", body)
|
|
}
|
|
response.Header().Set("Content-Type", "application/json")
|
|
_, _ = response.Write([]byte(`{"object":"list","data":[{"id":"bge-m3"},{"id":"bge-reranker-v2-m3"},{"id":"qwen3-32b"},{"id":"qwen-vl"}]}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
driver, apiConfig := newMWSTestDriver(server.URL)
|
|
models, err := driver.ListModels(context.Background(), apiConfig)
|
|
if err != nil {
|
|
t.Fatalf("list models: %v", err)
|
|
}
|
|
want := []ListModelResponse{
|
|
{Name: "bge-m3", ModelTypes: []string{"embedding"}},
|
|
{Name: "bge-reranker-v2-m3", ModelTypes: []string{"rerank"}},
|
|
{Name: "qwen3-32b", ModelTypes: []string{"chat"}},
|
|
}
|
|
if !reflect.DeepEqual(models, want) {
|
|
t.Fatalf("unexpected models: %#v", models)
|
|
}
|
|
}
|
|
|
|
func TestMWSChatSendsOnlyDocumentedFields(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
|
if request.URL.Path != "/projects/test-project/openai/v1/chat/completions" {
|
|
t.Fatalf("unexpected path: %s", request.URL.Path)
|
|
}
|
|
payload := decodeMWSRequest(t, request)
|
|
want := map[string]any{
|
|
"model": "qwen3-32b",
|
|
"messages": []any{
|
|
map[string]any{"role": "system", "content": "Be concise."},
|
|
map[string]any{"role": "user", "content": "Hello"},
|
|
},
|
|
"temperature": 0.25,
|
|
"max_completion_tokens": float64(128),
|
|
}
|
|
if !reflect.DeepEqual(payload, want) {
|
|
t.Fatalf("unexpected chat payload: %#v", payload)
|
|
}
|
|
response.Header().Set("Content-Type", "application/json")
|
|
_, _ = response.Write([]byte(`{"id":"chat-1","model":"qwen3-32b","choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"Hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":1,"total_tokens":5}}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
driver, apiConfig := newMWSTestDriver(server.URL)
|
|
temperature := 0.25
|
|
maxTokens := 128
|
|
topP := 0.9
|
|
stop := []string{"ignored"}
|
|
response, err := driver.ChatWithMessages(
|
|
context.Background(),
|
|
"qwen3-32b",
|
|
[]Message{
|
|
{Role: "system", Content: "Be concise."},
|
|
{Role: "user", Content: "Hello", ToolCallID: "ignored"},
|
|
},
|
|
apiConfig,
|
|
&ChatConfig{Temperature: &temperature, MaxTokens: &maxTokens, TopP: &topP, Stop: &stop, Tools: map[string]any{"ignored": true}},
|
|
nil,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("chat: %v", err)
|
|
}
|
|
if response.Answer == nil || *response.Answer != "Hi" || response.Usage == nil || response.Usage.TotalTokens != 5 {
|
|
t.Fatalf("unexpected chat response: %#v", response)
|
|
}
|
|
}
|
|
|
|
func TestMWSChatStreamingUsesDocumentedFields(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
|
if request.URL.Path != "/projects/test-project/openai/v1/chat/completions" {
|
|
t.Fatalf("unexpected path: %s", request.URL.Path)
|
|
}
|
|
payload := decodeMWSRequest(t, request)
|
|
want := map[string]any{
|
|
"model": "qwen3-32b",
|
|
"messages": []any{map[string]any{"role": "user", "content": "Hello"}},
|
|
"stream": true,
|
|
"stream_options": map[string]any{
|
|
"include_usage": true,
|
|
},
|
|
}
|
|
if !reflect.DeepEqual(payload, want) {
|
|
t.Fatalf("unexpected streaming chat payload: %#v", payload)
|
|
}
|
|
response.Header().Set("Content-Type", "text/event-stream")
|
|
_, _ = io.WriteString(response, "data: {\"model\":\"qwen3-32b\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hel\"}}]}\n\n")
|
|
_, _ = io.WriteString(response, "data: {\"model\":\"qwen3-32b\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"lo\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":4,\"completion_tokens\":1,\"total_tokens\":5}}\n\n")
|
|
_, _ = io.WriteString(response, "data: [DONE]\n\n")
|
|
}))
|
|
defer server.Close()
|
|
|
|
driver, apiConfig := newMWSTestDriver(server.URL)
|
|
stream := true
|
|
config := &ChatConfig{Stream: &stream}
|
|
var chunks []string
|
|
err := driver.ChatStreamlyWithSender(
|
|
context.Background(),
|
|
"qwen3-32b",
|
|
[]Message{{Role: "user", Content: "Hello"}},
|
|
apiConfig,
|
|
config,
|
|
nil,
|
|
func(content, _ *string) error {
|
|
if content != nil {
|
|
chunks = append(chunks, *content)
|
|
}
|
|
return nil
|
|
},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("stream chat: %v", err)
|
|
}
|
|
if !reflect.DeepEqual(chunks, []string{"Hel", "lo", "[DONE]"}) {
|
|
t.Fatalf("unexpected stream chunks: %#v", chunks)
|
|
}
|
|
if config.UsageResult == nil || config.UsageResult.TotalTokens != 5 {
|
|
t.Fatalf("unexpected stream usage: %#v", config.UsageResult)
|
|
}
|
|
}
|
|
|
|
func TestMWSEmbedSendsOnlyDocumentedFieldsAndOrdersVectors(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
|
if request.URL.Path != "/projects/test-project/openai/v1/embeddings" {
|
|
t.Fatalf("unexpected path: %s", request.URL.Path)
|
|
}
|
|
payload := decodeMWSRequest(t, request)
|
|
want := map[string]any{"model": "bge-m3", "input": []any{"first", "second"}}
|
|
if !reflect.DeepEqual(payload, want) {
|
|
t.Fatalf("unexpected embedding payload: %#v", payload)
|
|
}
|
|
response.Header().Set("Content-Type", "application/json")
|
|
_, _ = response.Write([]byte(`{"data":[{"index":1,"embedding":[0.3,0.4]},{"index":0,"embedding":[0.1,0.2]}],"usage":{"prompt_tokens":7,"total_tokens":7}}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
driver, apiConfig := newMWSTestDriver(server.URL)
|
|
modelName := "bge-m3"
|
|
embeddings, err := driver.Embed(context.Background(), &modelName, EmbedRequest{Texts: []string{"first", "second"}}, apiConfig, &EmbeddingConfig{Dimension: 256}, nil)
|
|
if err != nil {
|
|
t.Fatalf("embed: %v", err)
|
|
}
|
|
if len(embeddings) != 2 || embeddings[0].Index != 0 || embeddings[1].Index != 1 || embeddings[0].Embedding[0] != 0.1 || embeddings[1].Embedding[0] != 0.3 {
|
|
t.Fatalf("unexpected embeddings: %#v", embeddings)
|
|
}
|
|
}
|
|
|
|
func TestMWSRerankUsesCohereEndpointAndOriginalIndexOrder(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
|
if request.URL.Path != "/projects/test-project/cohere/v2/rerank" {
|
|
t.Fatalf("unexpected path: %s", request.URL.Path)
|
|
}
|
|
payload := decodeMWSRequest(t, request)
|
|
want := map[string]any{
|
|
"model": "bge-reranker-v2-m3",
|
|
"query": "query",
|
|
"documents": []any{"first", "second"},
|
|
"top_n": float64(2),
|
|
}
|
|
if !reflect.DeepEqual(payload, want) {
|
|
t.Fatalf("unexpected rerank payload: %#v", payload)
|
|
}
|
|
response.Header().Set("Content-Type", "application/json")
|
|
_, _ = response.Write([]byte(`{"id":"score-1","results":[{"index":1,"relevance_score":0.9},{"index":0,"relevance_score":0.2}],"meta":{"tokens":{"input_tokens":9}}}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
driver, apiConfig := newMWSTestDriver(server.URL)
|
|
modelName := "bge-reranker-v2-m3"
|
|
result, err := driver.Rerank(context.Background(), &modelName, RerankRequest{Query: "query", Documents: []string{"first", "second"}}, apiConfig, &RerankConfig{TopN: 1}, nil)
|
|
if err != nil {
|
|
t.Fatalf("rerank: %v", err)
|
|
}
|
|
if len(result.Data) != 2 || result.Data[0].Index != 0 || result.Data[0].RelevanceScore != 0.2 || result.Data[1].Index != 1 || result.Data[1].RelevanceScore != 0.9 {
|
|
t.Fatalf("unexpected rerank result: %#v", result.Data)
|
|
}
|
|
}
|
|
|
|
func TestMWSRejectsEmptyTokenWithoutRequest(t *testing.T) {
|
|
var calls atomic.Int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
|
calls.Add(1)
|
|
}))
|
|
defer server.Close()
|
|
|
|
driver, apiConfig := newMWSTestDriver(server.URL)
|
|
empty := " "
|
|
apiConfig.ApiKey = &empty
|
|
modelName := "bge-m3"
|
|
_, err := driver.Embed(context.Background(), &modelName, EmbedRequest{Texts: []string{"text"}}, apiConfig, nil, nil)
|
|
if err == nil || !strings.Contains(err.Error(), "api key is required") {
|
|
t.Fatalf("expected an API key error, got %v", err)
|
|
}
|
|
if calls.Load() != 0 {
|
|
t.Fatalf("unexpected HTTP calls: %d", calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestMWSEmptyInputsDoNotSendRequests(t *testing.T) {
|
|
var calls atomic.Int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
|
calls.Add(1)
|
|
}))
|
|
defer server.Close()
|
|
|
|
driver, apiConfig := newMWSTestDriver(server.URL)
|
|
modelName := "bge-m3"
|
|
embeddings, err := driver.Embed(context.Background(), &modelName, EmbedRequest{Texts: []string{}}, apiConfig, nil, nil)
|
|
if err != nil || len(embeddings) != 0 {
|
|
t.Fatalf("empty embedding input: data=%#v err=%v", embeddings, err)
|
|
}
|
|
ranked, err := driver.Rerank(context.Background(), &modelName, RerankRequest{Query: "query", Documents: []string{}}, apiConfig, nil, nil)
|
|
if err != nil || len(ranked.Data) != 0 {
|
|
t.Fatalf("empty rerank input: data=%#v err=%v", ranked, err)
|
|
}
|
|
if calls.Load() != 0 {
|
|
t.Fatalf("unexpected HTTP calls: %d", calls.Load())
|
|
}
|
|
}
|
|
|
|
func TestMWSErrorResponseIsReturned(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
|
|
http.Error(response, "MWS unavailable", http.StatusServiceUnavailable)
|
|
}))
|
|
defer server.Close()
|
|
|
|
driver, apiConfig := newMWSTestDriver(server.URL)
|
|
modelName := "bge-m3"
|
|
_, err := driver.Embed(context.Background(), &modelName, EmbedRequest{Texts: []string{"text"}}, apiConfig, nil, nil)
|
|
if err == nil || !strings.Contains(err.Error(), "status 503") || !strings.Contains(err.Error(), "MWS unavailable") {
|
|
t.Fatalf("unexpected MWS error: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestMWSFactoryRegistration(t *testing.T) {
|
|
driver, err := NewModelFactory().CreateModelDriver("MWS", map[string]string{"default": "https://gpt.mwsapis.ru/projects/demo"}, URLSuffix{})
|
|
if err != nil {
|
|
t.Fatalf("create driver: %v", err)
|
|
}
|
|
if driver.Name() != "MWS" {
|
|
t.Fatalf("unexpected driver: %s", driver.Name())
|
|
}
|
|
}
|