diff --git a/internal/entity/models/302ai.go b/internal/entity/models/302ai.go index a093681a68..c2dad5aad3 100644 --- a/internal/entity/models/302ai.go +++ b/internal/entity/models/302ai.go @@ -10,9 +10,9 @@ import ( "io" "mime/multipart" "net/http" + "net/url" "os" "path/filepath" - "ragflow/internal/common" "strconv" "strings" "time" @@ -62,9 +62,36 @@ func (a *AI302Model) Name() string { return "302ai" } +func validateAI302APIKey(apiConfig *APIConfig) (string, error) { + if apiConfig == nil || apiConfig.ApiKey == nil || strings.TrimSpace(*apiConfig.ApiKey) == "" { + return "", fmt.Errorf("api key is required") + } + return strings.TrimSpace(*apiConfig.ApiKey), nil +} + +func validateAI302ModelName(modelName *string) (string, error) { + if modelName == nil || strings.TrimSpace(*modelName) == "" { + return "", fmt.Errorf("model name is required") + } + return strings.TrimSpace(*modelName), nil +} + +func validateAI302DocumentURL(rawURL string) (string, error) { + documentURL := strings.TrimSpace(rawURL) + parsedURL, err := url.Parse(documentURL) + if err != nil || parsedURL.Host == "" || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") { + return "", fmt.Errorf("invalid document URL") + } + return documentURL, nil +} + func (a *AI302Model) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) { - if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" { - return nil, fmt.Errorf("api key is required") + apiKey, err := validateAI302APIKey(apiConfig) + if err != nil { + return nil, err + } + if strings.TrimSpace(modelName) == "" { + return nil, fmt.Errorf("model name is required") } if len(messages) == 0 { return nil, fmt.Errorf("messages is empty") @@ -88,17 +115,13 @@ func (a *AI302Model) ChatWithMessages(modelName string, messages []Message, apiC // Build request body reqBody := map[string]interface{}{ - "model": modelName, + "model": strings.TrimSpace(modelName), "messages": apiMessages, "stream": false, "temperature": 1, } if chatModelConfig != nil { - if chatModelConfig.Stream != nil { - reqBody["stream"] = *chatModelConfig.Stream - } - if chatModelConfig.MaxTokens != nil { reqBody["max_tokens"] = *chatModelConfig.MaxTokens } @@ -145,7 +168,7 @@ func (a *AI302Model) ChatWithMessages(modelName string, messages []Message, apiC } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) req.Header.Set("Accept", "application/json") resp, err := a.httpClient.Do(req) @@ -210,6 +233,16 @@ func (a *AI302Model) ChatWithMessages(modelName string, messages []Message, apiC } func (a *AI302Model) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error { + apiKey, err := validateAI302APIKey(apiConfig) + if err != nil { + return err + } + if strings.TrimSpace(modelName) == "" { + return fmt.Errorf("model name is required") + } + if sender == nil { + return fmt.Errorf("sender is required") + } if len(messages) == 0 { return fmt.Errorf("messages is empty") } @@ -232,17 +265,13 @@ func (a *AI302Model) ChatStreamlyWithSender(modelName string, messages []Message // Build request body with streaming enabled reqBody := map[string]interface{}{ - "model": modelName, + "model": strings.TrimSpace(modelName), "messages": apiMessages, "stream": true, "temperature": 1, } if modelConfig != nil { - if modelConfig.Stream != nil { - reqBody["stream"] = *modelConfig.Stream - } - if modelConfig.MaxTokens != nil { reqBody["max_tokens"] = *modelConfig.MaxTokens } @@ -293,7 +322,8 @@ func (a *AI302Model) ChatStreamlyWithSender(modelName string, messages []Message } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) + req.Header.Set("Accept", "text/event-stream") resp, err := a.httpClient.Do(req) if err != nil { @@ -310,7 +340,6 @@ func (a *AI302Model) ChatStreamlyWithSender(modelName string, messages []Message scanner := bufio.NewScanner(resp.Body) for scanner.Scan() { line := scanner.Text() - common.Info(line) // SSE data line starts with "data:" if !strings.HasPrefix(line, "data:") { @@ -379,6 +408,14 @@ func (a *AI302Model) Embed(modelName *string, texts []string, apiConfig *APIConf if len(texts) == 0 { return []EmbeddingData{}, nil } + model, err := validateAI302ModelName(modelName) + if err != nil { + return nil, err + } + apiKey, err := validateAI302APIKey(apiConfig) + if err != nil { + return nil, err + } var region = "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { @@ -388,7 +425,7 @@ func (a *AI302Model) Embed(modelName *string, texts []string, apiConfig *APIConf url := fmt.Sprintf("%s/%s", a.BaseURL[region], a.URLSuffix.Embedding) reqBody := map[string]interface{}{ - "model": *modelName, + "model": model, "input": texts, } @@ -403,7 +440,7 @@ func (a *AI302Model) Embed(modelName *string, texts []string, apiConfig *APIConf } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) resp, err := a.httpClient.Do(req) if err != nil { @@ -450,6 +487,17 @@ func (a *AI302Model) Rerank(modelName *string, query string, documents []string, if len(documents) == 0 { return &RerankResponse{}, nil } + if strings.TrimSpace(query) == "" { + return nil, fmt.Errorf("query is required") + } + model, err := validateAI302ModelName(modelName) + if err != nil { + return nil, err + } + apiKey, err := validateAI302APIKey(apiConfig) + if err != nil { + return nil, err + } var region = "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { @@ -458,14 +506,14 @@ func (a *AI302Model) Rerank(modelName *string, query string, documents []string, url := fmt.Sprintf("%s/%s", a.BaseURL[region], a.URLSuffix.Rerank) - var topN = rerankConfig.TopN - if rerankConfig.TopN != 0 { + var topN int + if rerankConfig != nil && rerankConfig.TopN != 0 { topN = rerankConfig.TopN } reqBody := map[string]interface{}{ - "model": *modelName, - "query": query, + "model": model, + "query": strings.TrimSpace(query), "documents": documents, "top_n": topN, } @@ -481,7 +529,7 @@ func (a *AI302Model) Rerank(modelName *string, query string, documents []string, } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) resp, err := a.httpClient.Do(req) if err != nil { @@ -522,9 +570,17 @@ func (a *AI302Model) Rerank(modelName *string, query string, documents []string, } func (a *AI302Model) TranscribeAudio(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig) (*ASRResponse, error) { - if file == nil || *file == "" { + if file == nil || strings.TrimSpace(*file) == "" { return nil, fmt.Errorf("file is missing") } + model, err := validateAI302ModelName(modelName) + if err != nil { + return nil, err + } + apiKey, err := validateAI302APIKey(apiConfig) + if err != nil { + return nil, err + } region := "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { @@ -538,7 +594,7 @@ func (a *AI302Model) TranscribeAudio(modelName *string, file *string, apiConfig writer := multipart.NewWriter(&body) // open audio file - audioFile, err := os.Open(*file) + audioFile, err := os.Open(strings.TrimSpace(*file)) if err != nil { return nil, fmt.Errorf("failed to open audio file: %w", err) } @@ -547,7 +603,7 @@ func (a *AI302Model) TranscribeAudio(modelName *string, file *string, apiConfig // create multipart file field part, err := writer.CreateFormFile( "file", - filepath.Base(*file), + filepath.Base(strings.TrimSpace(*file)), ) if err != nil { return nil, fmt.Errorf("failed to create multipart file: %w", err) @@ -559,7 +615,7 @@ func (a *AI302Model) TranscribeAudio(modelName *string, file *string, apiConfig } // model field - if err := writer.WriteField("model", *modelName); err != nil { + if err := writer.WriteField("model", model); err != nil { return nil, fmt.Errorf("failed to write model field: %w", err) } @@ -602,7 +658,7 @@ func (a *AI302Model) TranscribeAudio(modelName *string, file *string, apiConfig return nil, fmt.Errorf("failed to create request: %w", err) } - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) req.Header.Set("Content-Type", writer.FormDataContentType()) req.Header.Set("Accept", "application/json") @@ -648,9 +704,17 @@ func (a *AI302Model) AudioSpeechWithSender(modelName *string, audioContent *stri } func (a *AI302Model) OCRFile(modelName *string, content []byte, urls *string, apiConfig *APIConfig, ocrConfig *OCRConfig) (*OCRFileResponse, error) { - if (urls == nil || *urls == "") && (content == nil || len(content) == 0) { + if (urls == nil || strings.TrimSpace(*urls) == "") && (content == nil || len(content) == 0) { return nil, fmt.Errorf("file url or content is required") } + model, err := validateAI302ModelName(modelName) + if err != nil { + return nil, err + } + apiKey, err := validateAI302APIKey(apiConfig) + if err != nil { + return nil, err + } region := "default" if apiConfig.Region != nil && *apiConfig.Region != "" { @@ -660,8 +724,11 @@ func (a *AI302Model) OCRFile(modelName *string, content []byte, urls *string, ap url := fmt.Sprintf("%s/%s", a.BaseURL[region], a.URLSuffix.OCR) var docURL string - if urls != nil && *urls != "" { - docURL = *urls + if urls != nil && strings.TrimSpace(*urls) != "" { + docURL, err = validateAI302DocumentURL(*urls) + if err != nil { + return nil, err + } } else { mimeType := http.DetectContentType(content) base64Str := base64.StdEncoding.EncodeToString(content) @@ -669,7 +736,7 @@ func (a *AI302Model) OCRFile(modelName *string, content []byte, urls *string, ap } reqData := map[string]interface{}{ - "model": *modelName, + "model": model, "document": map[string]interface{}{ "type": "document_url", "document_url": docURL, @@ -690,7 +757,7 @@ func (a *AI302Model) OCRFile(modelName *string, content []byte, urls *string, ap } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) resp, err := a.httpClient.Do(req) if err != nil { @@ -732,11 +799,15 @@ func (a *AI302Model) OCRFile(modelName *string, content []byte, urls *string, ap } func (a *AI302Model) ParseFile(modelName *string, content []byte, documentURL *string, apiConfig *APIConfig, parseFileConfig *ParseFileConfig) (*ParseFileResponse, error) { - if documentURL == nil || *documentURL == "" { + if documentURL == nil || strings.TrimSpace(*documentURL) == "" { return nil, fmt.Errorf("302.ai API requires a valid public document URL; direct file upload is not supported") } + docURL, err := validateAI302DocumentURL(*documentURL) + if err != nil { + return nil, err + } - if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" { + if apiConfig == nil || apiConfig.ApiKey == nil || strings.TrimSpace(*apiConfig.ApiKey) == "" { return nil, fmt.Errorf("api key is required") } @@ -748,11 +819,11 @@ func (a *AI302Model) ParseFile(modelName *string, content []byte, documentURL *s apiURL := fmt.Sprintf("%s/%s", a.BaseURL[region], a.URLSuffix.DocumentParse) reqBody := map[string]interface{}{ - "url": *documentURL, + "url": docURL, } - if modelName != nil && *modelName != "" { - reqBody["model_version"] = *modelName + if modelName != nil && strings.TrimSpace(*modelName) != "" { + reqBody["model_version"] = strings.TrimSpace(*modelName) } jsonData, err := json.Marshal(reqBody) @@ -766,7 +837,7 @@ func (a *AI302Model) ParseFile(modelName *string, content []byte, documentURL *s } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", strings.TrimSpace(*apiConfig.ApiKey))) resp, err := a.httpClient.Do(req) if err != nil { @@ -798,6 +869,10 @@ func (a *AI302Model) ParseFile(modelName *string, content []byte, documentURL *s } func (a *AI302Model) ListModels(apiConfig *APIConfig) ([]string, error) { + apiKey, err := validateAI302APIKey(apiConfig) + if err != nil { + return nil, err + } var region = "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { region = *apiConfig.Region @@ -805,20 +880,13 @@ func (a *AI302Model) ListModels(apiConfig *APIConfig) ([]string, error) { url := fmt.Sprintf("%s/%s", a.BaseURL[region], a.URLSuffix.Models) - reqBody := map[string]string{} - - jsonData, err := json.Marshal(reqBody) - if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) - } - - req, err := http.NewRequest("GET", url, bytes.NewBuffer(jsonData)) + req, err := http.NewRequest("GET", url, nil) 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)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) + req.Header.Set("Accept", "application/json") resp, err := a.httpClient.Do(req) if err != nil { @@ -835,18 +903,24 @@ func (a *AI302Model) ListModels(apiConfig *APIConfig) ([]string, error) { return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - // Parse response - var result map[string]interface{} + var result struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } if err = json.Unmarshal(body, &result); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } + if result.Data == nil { + return nil, fmt.Errorf("models response missing data") + } - // convert result["data"] to []map[string]interface{} - models := make([]string, 0) - for _, model := range result["data"].([]interface{}) { - modelMap := model.(map[string]interface{}) - modelName := modelMap["id"].(string) - models = append(models, modelName) + models := make([]string, 0, len(result.Data)) + for _, model := range result.Data { + if strings.TrimSpace(model.ID) == "" { + return nil, fmt.Errorf("models response contains empty id") + } + models = append(models, strings.TrimSpace(model.ID)) } return models, nil @@ -866,7 +940,10 @@ func (a *AI302Model) ListTasks(apiConfig *APIConfig) ([]ListTaskStatus, error) { } func (a *AI302Model) ShowTask(taskID string, apiConfig *APIConfig) (*TaskResponse, error) { - if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" { + if strings.TrimSpace(taskID) == "" { + return nil, fmt.Errorf("task id is required") + } + if apiConfig == nil || apiConfig.ApiKey == nil || strings.TrimSpace(*apiConfig.ApiKey) == "" { return nil, fmt.Errorf("api key is required") } @@ -876,14 +953,14 @@ func (a *AI302Model) ShowTask(taskID string, apiConfig *APIConfig) (*TaskRespons } // URL: https://mineru.net/api/v4/extract/task/{task_id} - apiURL := fmt.Sprintf("%s/%s/%s", a.BaseURL[region], a.URLSuffix.DocumentParse, taskID) + apiURL := fmt.Sprintf("%s/%s/%s", a.BaseURL[region], a.URLSuffix.DocumentParse, url.PathEscape(strings.TrimSpace(taskID))) req, err := http.NewRequest("GET", apiURL, nil) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", strings.TrimSpace(*apiConfig.ApiKey))) resp, err := a.httpClient.Do(req) if err != nil { diff --git a/internal/entity/models/302ai_test.go b/internal/entity/models/302ai_test.go new file mode 100644 index 0000000000..472415f5da --- /dev/null +++ b/internal/entity/models/302ai_test.go @@ -0,0 +1,388 @@ +package models + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func newAI302Server(t *testing.T, handler func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter)) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Errorf("expected Authorization=Bearer test-key, got %q", got) + return + } + + var body map[string]interface{} + if r.Method == http.MethodPost { + if got := r.Header.Get("Content-Type"); !strings.HasPrefix(got, "application/json") { + t.Errorf("expected Content-Type to start with application/json, got %q", got) + return + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read body: %v", err) + return + } + if err := json.Unmarshal(raw, &body); err != nil { + t.Errorf("unmarshal: %v\nraw=%s", err, string(raw)) + return + } + } else { + if r.ContentLength > 0 { + t.Errorf("expected %s request without body, ContentLength=%d", r.Method, r.ContentLength) + return + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read body: %v", err) + return + } + if len(raw) != 0 { + t.Errorf("expected %s request without body, got %q", r.Method, string(raw)) + return + } + } + + handler(t, r, body, w) + })) +} + +func newAI302ForTest(baseURL string) *AI302Model { + return NewAI302Model( + map[string]string{"default": baseURL}, + URLSuffix{ + Chat: "v1/chat/completions", + Embedding: "jina/v1/embeddings", + Rerank: "jina/v1/rerank", + Models: "v1/models", + ASR: "v1/audio/transcriptions", + OCR: "mistral/v1/ocr", + DocumentParse: "mineru/api/v4/extract/task", + }, + ) +} + +func TestAI302ChatForcesNonStreaming(t *testing.T) { + srv := newAI302Server(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { + if r.Method != http.MethodPost { + t.Errorf("method=%s, want POST", r.Method) + } + if r.URL.Path != "/v1/chat/completions" { + t.Errorf("path=%s, want /v1/chat/completions", r.URL.Path) + } + if body["stream"] != false { + t.Errorf("stream=%v, want false", body["stream"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{ + "content": "pong", + "reasoning_content": "\nthought", + }, + }}, + }) + }) + defer srv.Close() + + apiKey := "test-key" + stream := true + thinking := true + resp, err := newAI302ForTest(srv.URL).ChatWithMessages( + "gpt-5", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Stream: &stream, Thinking: &thinking}, + ) + if err != nil { + t.Fatalf("ChatWithMessages: %v", err) + } + if resp.Answer == nil || *resp.Answer != "pong" { + t.Errorf("Answer=%v, want pong", resp.Answer) + } + if resp.ReasonContent == nil || *resp.ReasonContent != "thought" { + t.Errorf("ReasonContent=%v, want thought", resp.ReasonContent) + } +} + +func TestAI302StreamForcesStreaming(t *testing.T) { + srv := newAI302Server(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { + if r.URL.Path != "/v1/chat/completions" { + t.Errorf("path=%s, want /v1/chat/completions", r.URL.Path) + } + if body["stream"] != true { + t.Errorf("stream=%v, want true", body["stream"]) + } + if got := r.Header.Get("Accept"); got != "text/event-stream" { + t.Errorf("Accept=%q, want text/event-stream", got) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, strings.Join([]string{ + `data: {"choices":[{"delta":{"reasoning_content":"thinking"}}]}`, + `data: {"choices":[{"delta":{"content":"hello"}}]}`, + `data: [DONE]`, + ``, + }, "\n")) + }) + defer srv.Close() + + apiKey := "test-key" + stream := false + var content, reasoning []string + err := newAI302ForTest(srv.URL).ChatStreamlyWithSender( + "gpt-5", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Stream: &stream}, + func(answer, reason *string) error { + if answer != nil { + content = append(content, *answer) + } + if reason != nil { + reasoning = append(reasoning, *reason) + } + return nil + }, + ) + if err != nil { + t.Fatalf("ChatStreamlyWithSender: %v", err) + } + if got := strings.Join(content, ""); got != "hello[DONE]" { + t.Errorf("content=%q, want hello[DONE]", got) + } + if got := strings.Join(reasoning, ""); got != "thinking" { + t.Errorf("reasoning=%q, want thinking", got) + } +} + +func TestAI302ListModelsHappyPath(t *testing.T) { + srv := newAI302Server(t, func(t *testing.T, r *http.Request, _ map[string]interface{}, w http.ResponseWriter) { + if r.Method != http.MethodGet { + t.Errorf("method=%s, want GET", r.Method) + } + if r.URL.Path != "/v1/models" { + t.Errorf("path=%s, want /v1/models", r.URL.Path) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]string{ + {"id": "gpt-5"}, + {"id": " jina-embeddings-v3 "}, + }, + }) + }) + defer srv.Close() + + apiKey := "test-key" + models, err := newAI302ForTest(srv.URL).ListModels(&APIConfig{ApiKey: &apiKey}) + if err != nil { + t.Fatalf("ListModels: %v", err) + } + if got := strings.Join(models, ","); got != "gpt-5,jina-embeddings-v3" { + t.Errorf("models=%q", got) + } +} + +func TestAI302ListModelsRejectsMalformedResponse(t *testing.T) { + apiKey := "test-key" + for name, response := range map[string]interface{}{ + "missing data": map[string]interface{}{"object": "list"}, + "empty id": map[string]interface{}{"data": []map[string]string{{"id": ""}}}, + } { + t.Run(name, func(t *testing.T) { + srv := newAI302Server(t, func(t *testing.T, _ *http.Request, _ map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(response) + }) + defer srv.Close() + + if _, err := newAI302ForTest(srv.URL).ListModels(&APIConfig{ApiKey: &apiKey}); err == nil { + t.Fatal("expected malformed response error") + } + }) + } +} + +func TestAI302ShowTaskEscapesTaskID(t *testing.T) { + srv := newAI302Server(t, func(t *testing.T, r *http.Request, _ map[string]interface{}, w http.ResponseWriter) { + if r.Method != http.MethodGet { + t.Errorf("method=%s, want GET", r.Method) + } + want := "/mineru/api/v4/extract/task/task%2Fwith%3Fquery%23fragment" + if r.RequestURI != want { + t.Errorf("RequestURI=%q, want %q", r.RequestURI, want) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "code": 0, + "data": map[string]interface{}{ + "state": "done", + "full_zip_url": "https://example.com/result.zip", + }, + }) + }) + defer srv.Close() + + apiKey := "test-key" + resp, err := newAI302ForTest(srv.URL).ShowTask(" task/with?query#fragment ", &APIConfig{ApiKey: &apiKey}) + if err != nil { + t.Fatalf("ShowTask: %v", err) + } + if len(resp.Segments) != 1 || resp.Segments[0].Content != "https://example.com/result.zip" { + t.Fatalf("Segments=%v", resp.Segments) + } +} + +func TestAI302ValidatesInputs(t *testing.T) { + apiKey := "test-key" + emptyKey := " " + model := "gpt-5" + file := " " + blankURL := " " + docURL := "https://example.com/doc.pdf" + invalidURL := "ftp://example.com/doc.pdf" + send := func(*string, *string) error { return nil } + + tests := []struct { + name string + run func() error + want string + }{ + { + name: "chat api key", + run: func() error { + _, err := newAI302ForTest("http://unused").ChatWithMessages("gpt-5", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &emptyKey}, nil) + return err + }, + want: "api key is required", + }, + { + name: "chat model", + run: func() error { + _, err := newAI302ForTest("http://unused").ChatWithMessages(" ", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "model name is required", + }, + { + name: "stream api key", + run: func() error { + return newAI302ForTest("http://unused").ChatStreamlyWithSender("gpt-5", []Message{{Role: "user", Content: "x"}}, nil, nil, send) + }, + want: "api key is required", + }, + { + name: "stream sender", + run: func() error { + return newAI302ForTest("http://unused").ChatStreamlyWithSender("gpt-5", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &apiKey}, nil, nil) + }, + want: "sender is required", + }, + { + name: "embed model", + run: func() error { + _, err := newAI302ForTest("http://unused").Embed(nil, []string{"x"}, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "model name is required", + }, + { + name: "rerank api key", + run: func() error { + _, err := newAI302ForTest("http://unused").Rerank(&model, "q", []string{"doc"}, nil, nil) + return err + }, + want: "api key is required", + }, + { + name: "rerank query", + run: func() error { + _, err := newAI302ForTest("http://unused").Rerank(&model, " ", []string{"doc"}, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "query is required", + }, + { + name: "asr model", + run: func() error { + _, err := newAI302ForTest("http://unused").TranscribeAudio(nil, &docURL, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "model name is required", + }, + { + name: "asr file", + run: func() error { + _, err := newAI302ForTest("http://unused").TranscribeAudio(&model, &file, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "file is missing", + }, + { + name: "ocr api key", + run: func() error { + _, err := newAI302ForTest("http://unused").OCRFile(&model, nil, &docURL, nil, nil) + return err + }, + want: "api key is required", + }, + { + name: "ocr input", + run: func() error { + _, err := newAI302ForTest("http://unused").OCRFile(&model, nil, &blankURL, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "file url or content is required", + }, + { + name: "ocr invalid url", + run: func() error { + _, err := newAI302ForTest("http://unused").OCRFile(&model, nil, &invalidURL, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "invalid document URL", + }, + { + name: "parse file url", + run: func() error { + _, err := newAI302ForTest("http://unused").ParseFile(&model, nil, &blankURL, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "valid public document URL", + }, + { + name: "parse file invalid url", + run: func() error { + _, err := newAI302ForTest("http://unused").ParseFile(&model, nil, &invalidURL, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "invalid document URL", + }, + { + name: "models api key", + run: func() error { + _, err := newAI302ForTest("http://unused").ListModels(&APIConfig{}) + return err + }, + want: "api key is required", + }, + { + name: "show task id", + run: func() error { + _, err := newAI302ForTest("http://unused").ShowTask(" ", &APIConfig{ApiKey: &apiKey}) + return err + }, + want: "task id is required", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.run() + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("expected %q error, got %v", tt.want, err) + } + }) + } +}