From d63d3bb7d2686da563cf59978ed27e69df113c53 Mon Sep 17 00:00:00 2001 From: tmimmanuel <14046872+tmimmanuel@users.noreply.github.com> Date: Tue, 12 May 2026 20:10:50 -1000 Subject: [PATCH] Go: implement provider: Novita.ai (#14850) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What problem does this PR solve? Add a Go driver for Novita.ai (https://novita.ai), one of the unchecked providers on the umbrella tracking issue #14736. Novita exposes an OpenAI-compatible REST API at `https://api.novita.ai/v3/openai` and proxies a large catalog of third-party models (DeepSeek, Llama, Qwen3, Kimi, Gemma, Mistral, MiniMax, GLM, etc.) behind a single OpenAI-shaped surface — 102 models live at the time of writing. Until this PR, a tenant who configured `novita` as a model provider in the Go layer fell through to the default branch of `internal/entity/models/factory.go` and got the dummy driver. ### What this PR includes - New `internal/entity/models/novita.go` with a `NovitaModel` implementing the `ModelDriver` interface (~520 lines). - New `conf/models/novita.json` with 7 representative chat models (DeepSeek-V4, Llama-3.3-70B, Qwen3-30B/235B reasoning, Kimi-K2, Gemma-3-27B, Mistral-Nemo). - `factory.go`: route `"novita"` to `NewNovitaModel`. - `internal/entity/models/novita_test.go`: 23 unit tests. ### Notable design point: `...` reasoning extraction Novita-routed reasoning models like `qwen3-*` and `deepseek-r1-*` embed their chain-of-thought **inline inside content as `...` tags**, rather than in a separate `reasoning_content` field. Verified live by probing `api.novita.ai`: ``` content head 200: Okay, let's see. I need to find 15% of 80. Hmm, percentages can sometimes be tricky, but I think content tail 100: h, that works. Alternatively, 0.15 × 80. If I move the decimal two places to the left for ``` Without handling, a tenant picking qwen3 via Novita would see raw `` tags in their UI answer — different from every other reasoning provider in the Go layer. The driver detects those tags and routes the inner text to `ChatResponse.ReasonContent` (non-stream) or the sender's second arg (stream), keeping the visible answer clean of tag clutter: - **`splitNovitaThink`** — scans a complete content string. Used by the non-streaming path. Handles multiple `` blocks, unclosed tags (the model got cut off mid-reasoning), pure-text content with no tags. - **`novitaThinkExtractor`** — stateful streaming version. Buffers trailing bytes that might be the start of a tag (e.g. ``), then emits segments in routing order so callers can pipe them to a UI. Tested with byte-level chunk boundaries and tag-spanning scenarios. ### Method coverage | Method | Behavior | |---|---| | `ChatWithMessages` | `POST /v3/openai/chat/completions`, `` extraction on response | | `ChatStreamlyWithSender` | SSE stream, stateful `` extraction across deltas | | `ListModels` / `CheckConnection` | `GET /v3/openai/models` (102 live) | | `Embed` / `Rerank` / `Balance` / `TranscribeAudio` / `AudioSpeech` / `OCRFile` | `"no such method"` — Novita's OpenAI-compatible surface does not expose any | No interface change. No new dependencies. ### How was this tested? **23 unit tests** in `internal/entity/models/novita_test.go` — all pass: ``` $ go test -vet=off -run "TestNovita|TestSplitNovita" -count=1 ./internal/entity/models/... ok ragflow/internal/entity/models 0.020s ``` Coverage: - `splitNovitaThink` (5 cases: pure text, single block, leading text, multiple blocks, unclosed tag) - `novitaThinkExtractor` (6 cases: single-chunk, opening tag span, closing tag span, byte-level chunking, no tags, lone `<` not as tag start) - `ChatWithMessages`: pure text, with `` tags, missing API key, empty messages, HTTP error - `ChatStreamlyWithSender`: tag-stripping with spanning deltas, pure content, sender-required, stream-true-required - `ListModels` / `CheckConnection` (happy paths) - All sentinel methods `go build ./internal/entity/models/...` exits 0 on go 1.25. **Live integration test** against `api.novita.ai/v3/openai`: ``` === RUN TestNovitaLiveSmoke [OK] Name() = "novita" [OK] CheckConnection [OK] ListModels: 102 models (showing first 6) [deepseek/deepseek-v4-pro deepseek/deepseek-v4-flash deepseek/deepseek-v3.2 xiaomimimo/mimo-v2.5-pro moonshotai/kimi-k2.6 zai-org/glm-5.1] [OK] Chat (llama-3.3) answer="ok" reason="" [OK] Chat (qwen3) answer len=0 head="" ReasonContent len=1657 head="Okay, so I need to figure out what 15% of 80 is. Hmm, percentages can sometimes trip me up, but let ..." [OK] Stream content: 0 chunks, 0 chars; reasoning: 600 chunks, 1667 chars [OK] Embed/Rerank/Balance/TranscribeAudio/AudioSpeech/OCRFile all return "novita, no such method" NOVITA LIVE SMOKE PASSED --- PASS: TestNovitaLiveSmoke (26.18s) ``` What the live run proves on the wire: - Auth (`Bearer `) accepted by `api.novita.ai`. - `/v3/openai/models` parser handles the real 102-model response. - Non-stream chat against `meta-llama/llama-3.3-70b-instruct`: clean string answer, empty ReasonContent (non-reasoning model, pure-text path). - Non-stream chat against `qwen/qwen3-30b-a3b-fp8`: 1657-char reasoning extracted from `...` and routed to `ChatResponse.ReasonContent`. Visible answer is 0 chars in this run because qwen3 spent its 600-token budget entirely on reasoning before reaching the answer phase — that's the model's behavior, not a driver bug. The important thing: **no `` tags leaked into the visible Answer field**. - Streaming against qwen3: 600 reasoning chunks (1667 chars) emitted via the sender's 2nd arg across SSE deltas; **no `` tag fragments leaked into the content channel** despite tag boundaries crossing chunk boundaries on the wire. - All 6 sentinel methods return the documented `"no such method"` strings. ### Type of change - [x] New Feature (non-breaking change which adds functionality) Tracking: #14736 --- conf/models/novita.json | 62 ++ internal/entity/models/factory.go | 2 + internal/entity/models/novita.go | 662 ++++++++++++++++++++++ internal/entity/models/novita_test.go | 776 ++++++++++++++++++++++++++ 4 files changed, 1502 insertions(+) create mode 100644 conf/models/novita.json create mode 100644 internal/entity/models/novita.go create mode 100644 internal/entity/models/novita_test.go diff --git a/conf/models/novita.json b/conf/models/novita.json new file mode 100644 index 0000000000..c5b3839e14 --- /dev/null +++ b/conf/models/novita.json @@ -0,0 +1,62 @@ +{ + "name": "Novita", + "url": { + "default": "https://api.novita.ai" + }, + "url_suffix": { + "chat": "openai/v1/chat/completions", + "models": "openai/v1/models" + }, + "class": "novita", + "models": [ + { + "name": "deepseek/deepseek-v4-pro", + "max_tokens": 65536, + "model_types": [ + "chat" + ] + }, + { + "name": "meta-llama/llama-3.3-70b-instruct", + "max_tokens": 131072, + "model_types": [ + "chat" + ] + }, + { + "name": "qwen/qwen3-30b-a3b-fp8", + "max_tokens": 32768, + "model_types": [ + "chat" + ] + }, + { + "name": "qwen/qwen3-235b-a22b-fp8", + "max_tokens": 32768, + "model_types": [ + "chat" + ] + }, + { + "name": "moonshotai/kimi-k2-instruct", + "max_tokens": 131072, + "model_types": [ + "chat" + ] + }, + { + "name": "google/gemma-3-27b-it", + "max_tokens": 131072, + "model_types": [ + "chat" + ] + }, + { + "name": "mistralai/mistral-nemo", + "max_tokens": 131072, + "model_types": [ + "chat" + ] + } + ] +} diff --git a/internal/entity/models/factory.go b/internal/entity/models/factory.go index 7540605d34..364af0c8b1 100644 --- a/internal/entity/models/factory.go +++ b/internal/entity/models/factory.go @@ -83,6 +83,8 @@ func (f *ModelFactory) CreateModelDriver(providerName string, baseURL map[string return NewBaichuanModel(baseURL, urlSuffix), nil case "jina": return NewJinaModel(baseURL, urlSuffix), nil + case "novita": + return NewNovitaModel(baseURL, urlSuffix), nil default: return NewDummyModel(baseURL, urlSuffix), nil } diff --git a/internal/entity/models/novita.go b/internal/entity/models/novita.go new file mode 100644 index 0000000000..50db75b727 --- /dev/null +++ b/internal/entity/models/novita.go @@ -0,0 +1,662 @@ +// +// 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 ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" +) + +// NovitaModel implements ModelDriver for Novita.ai +// (https://novita.ai/docs/api-reference/). +// +// Novita exposes an OpenAI-compatible REST API at +// https://api.novita.ai/v3/openai (chat completions at +// /chat/completions, list models at /models). It serves a large +// catalog of third-party models (DeepSeek, Llama, Qwen3, Kimi, +// Gemma, Mistral, etc.) behind a single OpenAI-shaped surface. +// +// The wire shape matches OpenAI standard with ONE notable +// difference: reasoning models like qwen3-* embed their +// chain-of-thought INLINE inside content as ... +// tags, rather than in a separate reasoning_content field. The +// driver detects those tags and routes the inner text to +// ChatResponse.ReasonContent (non-stream) or the sender's second +// arg (stream), keeping the answer clean of tag clutter. +type NovitaModel struct { + BaseURL map[string]string + URLSuffix URLSuffix + httpClient *http.Client +} + +// NewNovitaModel creates a new Novita model instance. +// +// Same transport convention as other Go drivers in this package: +// clone http.DefaultTransport, override the connection-pool fields, +// no client-level Timeout so SSE streams are not capped. +func NewNovitaModel(baseURL map[string]string, urlSuffix URLSuffix) *NovitaModel { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.MaxIdleConns = 100 + transport.MaxIdleConnsPerHost = 10 + transport.IdleConnTimeout = 90 * time.Second + transport.DisableCompression = false + transport.ResponseHeaderTimeout = 60 * time.Second + + return &NovitaModel{ + BaseURL: baseURL, + URLSuffix: urlSuffix, + httpClient: &http.Client{ + Transport: transport, + }, + } +} + +func (n *NovitaModel) NewInstance(baseURL map[string]string) ModelDriver { + return NewNovitaModel(baseURL, n.URLSuffix) +} + +func (n *NovitaModel) Name() string { + return "novita" +} + +func (n *NovitaModel) baseURLForRegion(region string) (string, error) { + base, ok := n.BaseURL[region] + if !ok || base == "" { + return "", fmt.Errorf("novita: no base URL configured for region %q", region) + } + // Strip a trailing "/" so callers can safely do + // fmt.Sprintf("%s/%s", base, suffix) without producing "//" in + // the path. The shipped config has no trailing slash, but a + // tenant can override the URL per-instance and may add one. + return strings.TrimSuffix(base, "/"), nil +} + +const ( + novitaThinkOpen = "" + novitaThinkClose = "" +) + +// splitNovitaThink walks a complete content string and returns the +// visible portion + the concatenated chain-of-thought from inside +// any ... blocks. Multiple think blocks are +// concatenated; tags themselves are stripped. Used by the +// non-streaming path where the whole content is available at once. +func splitNovitaThink(raw string) (visible, reasoning string) { + var v, r strings.Builder + inside := false + for { + var marker string + if inside { + marker = novitaThinkClose + } else { + marker = novitaThinkOpen + } + idx := strings.Index(raw, marker) + if idx < 0 { + if inside { + r.WriteString(raw) + } else { + v.WriteString(raw) + } + break + } + if inside { + r.WriteString(raw[:idx]) + } else { + v.WriteString(raw[:idx]) + } + raw = raw[idx+len(marker):] + inside = !inside + } + return v.String(), r.String() +} + +// novitaThinkExtractor maintains state across streaming chunks so +// that a ... block spanning multiple SSE events still +// gets split correctly between content and reasoning. The buffer +// preserves up to (len(closingMarker)-1) trailing bytes of each +// chunk in case the next chunk completes a partial tag. +type novitaThinkExtractor struct { + buf strings.Builder + inside bool +} + +// novitaThinkSegment is one routing decision: emit `content` via the +// sender's first arg, or emit `reasoning` via the sender's second arg. +// Exactly one of the two fields is non-empty. +type novitaThinkSegment struct { + content string + reasoning string +} + +// Feed appends an incoming chunk and returns any segments that are +// now safe to emit. Trailing bytes that could be the start of a tag +// are held back in the buffer until the next call. +func (e *novitaThinkExtractor) Feed(chunk string) []novitaThinkSegment { + e.buf.WriteString(chunk) + s := e.buf.String() + var out []novitaThinkSegment + for { + var marker, otherMarker string + if e.inside { + marker = novitaThinkClose + otherMarker = novitaThinkOpen + } else { + marker = novitaThinkOpen + otherMarker = novitaThinkClose + } + idx := strings.Index(s, marker) + if idx < 0 { + // No closing/opening marker yet. Emit everything except a + // possible partial-tag suffix at the very end. Reserve + // (max marker length - 1) trailing bytes so we don't + // emit "". + reserve := len(marker) - 1 + if len(otherMarker)-1 > reserve { + reserve = len(otherMarker) - 1 + } + safe := len(s) - reserve + if safe < 0 { + safe = 0 + } + // Don't reserve if the trailing bytes can't possibly be + // the start of a tag (no '<' suffix). + if safe < len(s) && !strings.Contains(s[safe:], "<") { + safe = len(s) + } + if safe > 0 { + if e.inside { + out = append(out, novitaThinkSegment{reasoning: s[:safe]}) + } else { + out = append(out, novitaThinkSegment{content: s[:safe]}) + } + s = s[safe:] + } + break + } + if idx > 0 { + if e.inside { + out = append(out, novitaThinkSegment{reasoning: s[:idx]}) + } else { + out = append(out, novitaThinkSegment{content: s[:idx]}) + } + } + s = s[idx+len(marker):] + e.inside = !e.inside + } + e.buf.Reset() + e.buf.WriteString(s) + return out +} + +// Flush returns the buffered tail when the stream ends. A stream that +// ends mid-tag would not normally happen with a well-behaved upstream, +// but if it does the partial bytes are emitted according to the +// current mode so nothing is silently lost. +func (e *novitaThinkExtractor) Flush() *novitaThinkSegment { + s := e.buf.String() + e.buf.Reset() + if s == "" { + return nil + } + if e.inside { + return &novitaThinkSegment{reasoning: s} + } + return &novitaThinkSegment{content: s} +} + +// ChatWithMessages sends multiple messages with roles and returns the response. +func (n *NovitaModel) 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") + } + if len(messages) == 0 { + return nil, fmt.Errorf("messages is empty") + } + + region := "default" + if apiConfig.Region != nil && *apiConfig.Region != "" { + region = *apiConfig.Region + } + + baseURL, err := n.baseURLForRegion(region) + if err != nil { + return nil, err + } + url := fmt.Sprintf("%s/%s", baseURL, n.URLSuffix.Chat) + + apiMessages := make([]map[string]interface{}, len(messages)) + for i, msg := range messages { + apiMessages[i] = map[string]interface{}{ + "role": msg.Role, + "content": msg.Content, + } + } + + reqBody := map[string]interface{}{ + "model": modelName, + "messages": apiMessages, + "stream": false, + } + + if chatModelConfig != nil { + if chatModelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *chatModelConfig.MaxTokens + } + if chatModelConfig.Temperature != nil { + reqBody["temperature"] = *chatModelConfig.Temperature + } + if chatModelConfig.TopP != nil { + reqBody["top_p"] = *chatModelConfig.TopP + } + if chatModelConfig.Stop != nil { + reqBody["stop"] = *chatModelConfig.Stop + } + // Map ChatConfig.Thinking -> Novita's `enable_thinking`. + // Per https://novita.ai/docs/api-reference/model-apis-llm-create-chat-completion, + // enable_thinking (boolean | null, default true) "controls the + // switches between thinking and non-thinking modes" for + // zai-org/glm-4.5, deepseek/deepseek-v3.1[-terminus|-exp]. For + // models outside that supported set Novita ignores the field, + // so it's safe to forward whenever the caller opts in. Tenants + // can now disable thinking mode at request time without having + // to use prompt-level hacks like "/no_think". + if chatModelConfig.Thinking != nil { + reqBody["enable_thinking"] = *chatModelConfig.Thinking + } + } + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 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 := n.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("API request failed with status %d: %s", resp.StatusCode, string(body)) + } + + var result map[string]interface{} + if err = json.Unmarshal(body, &result); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + choices, ok := result["choices"].([]interface{}) + if !ok || len(choices) == 0 { + return nil, fmt.Errorf("no choices in response") + } + + firstChoice, ok := choices[0].(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("invalid choice format") + } + + messageMap, ok := firstChoice["message"].(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("invalid message format") + } + + rawContent, ok := messageMap["content"].(string) + if !ok { + return nil, fmt.Errorf("invalid content format") + } + + // Novita emits chain-of-thought in two different shapes depending + // on the model and on enable_thinking: + // - qwen3-* and other inline-style models: chain-of-thought is + // embedded inside content as ... tags. + // - deepseek-v3.1 / glm-4.5 (and any model with separate + // reasoning enabled): chain-of-thought arrives in a separate + // `reasoning_content` field, with `content` already cleaned. + // Handle both so the visible Answer is always tag-free and any + // reasoning the upstream supplied is preserved. + visible, reasoning := splitNovitaThink(rawContent) + if r, ok := messageMap["reasoning_content"].(string); ok && r != "" { + if reasoning != "" { + reasoning += "\n" + r + } else { + reasoning = r + } + } + + return &ChatResponse{ + Answer: &visible, + ReasonContent: &reasoning, + }, nil +} + +// ChatStreamlyWithSender sends messages and streams the response via +// the sender. Handles both reasoning shapes Novita can emit: +// - delta.reasoning_content (deepseek-v3.1 / glm-4.5 / any model +// with separate reasoning): forwarded as-is to the second arg. +// - delta.content containing ... (qwen3-* and other +// inline-style models): a stateful extractor splits tag bytes +// across SSE chunk boundaries, then routes content/reasoning to +// the first/second sender arg respectively. +func (n *NovitaModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, sender func(*string, *string) error) error { + if sender == nil { + return fmt.Errorf("sender is required") + } + if len(messages) == 0 { + return fmt.Errorf("messages is empty") + } + if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" { + return fmt.Errorf("api key is required") + } + + region := "default" + if apiConfig.Region != nil && *apiConfig.Region != "" { + region = *apiConfig.Region + } + + baseURL, err := n.baseURLForRegion(region) + if err != nil { + return err + } + url := fmt.Sprintf("%s/%s", baseURL, n.URLSuffix.Chat) + + apiMessages := make([]map[string]interface{}, len(messages)) + for i, msg := range messages { + apiMessages[i] = map[string]interface{}{ + "role": msg.Role, + "content": msg.Content, + } + } + + reqBody := map[string]interface{}{ + "model": modelName, + "messages": apiMessages, + "stream": true, + } + + if chatModelConfig != nil { + if chatModelConfig.Stream != nil && !*chatModelConfig.Stream { + return fmt.Errorf("stream must be true in ChatStreamlyWithSender") + } + if chatModelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *chatModelConfig.MaxTokens + } + if chatModelConfig.Temperature != nil { + reqBody["temperature"] = *chatModelConfig.Temperature + } + if chatModelConfig.TopP != nil { + reqBody["top_p"] = *chatModelConfig.TopP + } + if chatModelConfig.Stop != nil { + reqBody["stop"] = *chatModelConfig.Stop + } + // See ChatWithMessages for why we forward this. + if chatModelConfig.Thinking != nil { + reqBody["enable_thinking"] = *chatModelConfig.Thinking + } + } + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(context.Background(), "POST", url, bytes.NewBuffer(jsonData)) + if err != nil { + return 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 := n.httpClient.Do(req) + if err != nil { + return fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) + } + + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 64*1024), 1024*1024) + extractor := &novitaThinkExtractor{} + sawTerminal := false + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data:") { + continue + } + data := strings.TrimSpace(line[5:]) + if data == "[DONE]" { + sawTerminal = true + break + } + var event map[string]interface{} + if err = json.Unmarshal([]byte(data), &event); err != nil { + continue + } + choices, ok := event["choices"].([]interface{}) + if !ok || len(choices) == 0 { + continue + } + firstChoice, ok := choices[0].(map[string]interface{}) + if !ok { + continue + } + delta, ok := firstChoice["delta"].(map[string]interface{}) + if !ok { + continue + } + // deepseek-v3.1 / glm-4.5 (and other models that emit reasoning + // separately) put chain-of-thought in delta.reasoning_content + // rather than inside content as .... Surface it + // before any content from the same chunk so callers piping to + // a UI render reasoning before the visible answer for that + // token, matching the wire ordering Novita emits. + if r, ok := delta["reasoning_content"].(string); ok && r != "" { + rr := r + if err := sender(nil, &rr); err != nil { + return err + } + } + if c, ok := delta["content"].(string); ok && c != "" { + for _, seg := range extractor.Feed(c) { + if seg.reasoning != "" { + r := seg.reasoning + if err := sender(nil, &r); err != nil { + return err + } + } + if seg.content != "" { + cc := seg.content + if err := sender(&cc, nil); err != nil { + return err + } + } + } + } + if finish, ok := firstChoice["finish_reason"].(string); ok && finish != "" { + sawTerminal = true + break + } + } + + // Flush any buffered tail (rare, but covers the case where the + // stream ends right after the last chunk without us seeing the + // closing tag). + if seg := extractor.Flush(); seg != nil { + if seg.reasoning != "" { + r := seg.reasoning + if err := sender(nil, &r); err != nil { + return err + } + } + if seg.content != "" { + cc := seg.content + if err := sender(&cc, nil); err != nil { + return err + } + } + } + + if err := scanner.Err(); err != nil { + return fmt.Errorf("failed to scan response body: %w", err) + } + if !sawTerminal { + return fmt.Errorf("novita: stream ended before [DONE] or finish_reason") + } + + endOfStream := "[DONE]" + if err := sender(&endOfStream, nil); err != nil { + return err + } + + return nil +} + +// ListModels returns the list of model ids visible to the API key. +func (n *NovitaModel) ListModels(apiConfig *APIConfig) ([]string, error) { + if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" { + return nil, fmt.Errorf("api key is required") + } + + region := "default" + if apiConfig.Region != nil && *apiConfig.Region != "" { + region = *apiConfig.Region + } + + baseURL, err := n.baseURLForRegion(region) + if err != nil { + return nil, err + } + url := fmt.Sprintf("%s/%s", baseURL, n.URLSuffix.Models) + + ctx, cancel := context.WithTimeout(context.Background(), nonStreamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, "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)) + + resp, err := n.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("API request failed with status %d: %s", resp.StatusCode, string(body)) + } + + var result map[string]interface{} + if err = json.Unmarshal(body, &result); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + data, ok := result["data"].([]interface{}) + if !ok { + return nil, fmt.Errorf("invalid models list format") + } + + models := make([]string, 0) + for _, model := range data { + modelMap, ok := model.(map[string]interface{}) + if !ok { + continue + } + modelName, ok := modelMap["id"].(string) + if !ok { + continue + } + models = append(models, modelName) + } + return models, nil +} + +// CheckConnection runs a lightweight ListModels call to verify the API key. +func (n *NovitaModel) CheckConnection(apiConfig *APIConfig) error { + _, err := n.ListModels(apiConfig) + return err +} + +// Embed is not exposed on Novita's OpenAI-compatible surface yet. +func (n *NovitaModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([]EmbeddingData, error) { + return nil, fmt.Errorf("%s, no such method", n.Name()) +} + +// Rerank is not exposed by the Novita API. +func (n *NovitaModel) Rerank(modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig) (*RerankResponse, error) { + return nil, fmt.Errorf("%s, no such method", n.Name()) +} + +// Balance is not exposed by the Novita API. +func (n *NovitaModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) { + return nil, fmt.Errorf("%s, no such method", n.Name()) +} + +func (n *NovitaModel) TranscribeAudio(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig) (*ASRResponse, error) { + return nil, fmt.Errorf("%s, no such method", n.Name()) +} + +func (n *NovitaModel) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, sender func(*string, *string) error) error { + return fmt.Errorf("%s, no such method", n.Name()) +} + +func (n *NovitaModel) AudioSpeech(modelName *string, audioContent *string, apiConfig *APIConfig, asrConfig *TTSConfig) (*TTSResponse, error) { + return nil, fmt.Errorf("%s, no such method", n.Name()) +} + +func (n *NovitaModel) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, sender func(*string, *string) error) error { + return fmt.Errorf("%s, no such method", n.Name()) +} + +func (n *NovitaModel) OCRFile(modelName *string, fileContent *string, apiConfig *APIConfig, ocrConfig *OCRConfig) (*OCRResponse, error) { + return nil, fmt.Errorf("%s, no such method", n.Name()) +} diff --git a/internal/entity/models/novita_test.go b/internal/entity/models/novita_test.go new file mode 100644 index 0000000000..0470918d3a --- /dev/null +++ b/internal/entity/models/novita_test.go @@ -0,0 +1,776 @@ +package models + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func newNovitaServer(t *testing.T, expectedPath string, handler func(t *testing.T, body map[string]interface{}, w http.ResponseWriter)) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != expectedPath { + t.Errorf("expected path=%s, got %s", expectedPath, r.URL.Path) + return + } + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Errorf("expected Authorization=Bearer test-key, got %q", got) + return + } + // Content-Type must declare JSON on BOTH POST chat (body is + // JSON) and GET ListModels (Novita platform expects callers to + // negotiate JSON content even though the body is empty — + // maintainer review explicitly flagged the missing header). + 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 + } + if r.Method == http.MethodPost { + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read body: %v", err) + return + } + var body map[string]interface{} + if err := json.Unmarshal(raw, &body); err != nil { + t.Errorf("unmarshal: %v\nraw=%s", err, string(raw)) + return + } + handler(t, body, w) + return + } + handler(t, nil, w) + })) +} + +func newNovitaForTest(baseURL string) *NovitaModel { + return NewNovitaModel( + map[string]string{"default": baseURL}, + URLSuffix{Chat: "openai/v1/chat/completions", Models: "openai/v1/models"}, + ) +} + +// newNovitaSSEServer asserts the SSE-chat wire contract (POST, path, +// Authorization, Content-Type) the same way newNovitaServer does for +// the JSON-chat path, then writes the supplied SSE payload. Closes +// the gap CodeRabbit flagged where streaming tests used +// httptest.NewServer directly and skipped the request-shape checks. +func newNovitaSSEServer(t *testing.T, expectedPath, ssePayload string) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + return + } + if r.URL.Path != expectedPath { + t.Errorf("expected path=%s, got %s", expectedPath, r.URL.Path) + return + } + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Errorf("expected Authorization=Bearer test-key, got %q", got) + return + } + 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 + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, ssePayload) + })) +} + +// ---- think-tag split helpers ---- + +func TestSplitNovitaThinkPureText(t *testing.T) { + v, r := splitNovitaThink("hello world") + if v != "hello world" || r != "" { + t.Errorf("got (%q,%q)", v, r) + } +} + +func TestSplitNovitaThinkSingleBlock(t *testing.T) { + v, r := splitNovitaThink("15% = 0.15. 0.15*80 = 12.The answer is 12.") + if v != "The answer is 12." { + t.Errorf("visible=%q", v) + } + if r != "15% = 0.15. 0.15*80 = 12." { + t.Errorf("reasoning=%q", r) + } +} + +func TestSplitNovitaThinkLeadingText(t *testing.T) { + v, r := splitNovitaThink("intro thoughttail") + if v != "intro tail" { + t.Errorf("visible=%q", v) + } + if r != "thought" { + t.Errorf("reasoning=%q", r) + } +} + +func TestSplitNovitaThinkMultipleBlocks(t *testing.T) { + v, r := splitNovitaThink("Apart1Bpart2") + if v != "part1part2" { + t.Errorf("visible=%q", v) + } + if r != "AB" { + t.Errorf("reasoning=%q", r) + } +} + +func TestSplitNovitaThinkUnclosedTag(t *testing.T) { + // Unclosed -> everything after the open tag is reasoning; + // content stops at the open tag. This matches a real upstream that + // got cut off mid-reasoning by max_tokens. + v, r := splitNovitaThink("visible still thinking when tokens ran out") + if v != "visible " { + t.Errorf("visible=%q", v) + } + if r != "still thinking when tokens ran out" { + t.Errorf("reasoning=%q", r) + } +} + +// ---- streaming extractor ---- + +// Helper to push multiple chunks through and concatenate all output by +// kind. Each chunk goes into Feed; the output is what's safe to emit. +func feedAll(e *novitaThinkExtractor, chunks []string) (content, reasoning string) { + var cb, rb strings.Builder + for _, c := range chunks { + for _, seg := range e.Feed(c) { + cb.WriteString(seg.content) + rb.WriteString(seg.reasoning) + } + } + if seg := e.Flush(); seg != nil { + cb.WriteString(seg.content) + rb.WriteString(seg.reasoning) + } + return cb.String(), rb.String() +} + +func TestNovitaThinkExtractorSingleChunk(t *testing.T) { + e := &novitaThinkExtractor{} + c, r := feedAll(e, []string{"hello thought world"}) + if c != "hello world" { + t.Errorf("content=%q", c) + } + if r != "thought" { + t.Errorf("reasoning=%q", r) + } +} + +func TestNovitaThinkExtractorTagSpansChunks(t *testing.T) { + // "" split across two SSE deltas: "" + e := &novitaThinkExtractor{} + c, r := feedAll(e, []string{"hello thoughttail"}) + if c != "hello tail" { + t.Errorf("content=%q", c) + } + if r != "thought" { + t.Errorf("reasoning=%q", r) + } +} + +func TestNovitaThinkExtractorClosingTagSpansChunks(t *testing.T) { + // "" split across two deltas + e := &novitaThinkExtractor{} + c, r := feedAll(e, []string{"reasoningvisible"}) + if c != "visible" { + t.Errorf("content=%q", c) + } + if r != "reasoning" { + t.Errorf("reasoning=%q", r) + } +} + +func TestNovitaThinkExtractorTokenBoundaries(t *testing.T) { + // Simulate the kind of chunking we saw on the wire for qwen3 — many + // small chunks, sometimes splitting tag bytes. + e := &novitaThinkExtractor{} + c, r := feedAll(e, []string{ + "<", "think>", "Ok", "ay, ", "compute. ", "12", "."}) + if c != "12." { + t.Errorf("content=%q", c) + } + if r != "Okay, compute. " { + t.Errorf("reasoning=%q", r) + } +} + +func TestNovitaThinkExtractorNoTags(t *testing.T) { + e := &novitaThinkExtractor{} + c, r := feedAll(e, []string{"plain ", "content ", "all ", "the way"}) + if c != "plain content all the way" { + t.Errorf("content=%q", c) + } + if r != "" { + t.Errorf("reasoning=%q", r) + } +} + +func TestNovitaThinkExtractorLessThanIsNotTagStart(t *testing.T) { + // "<10" or "... embedded in content. + // Driver must split it into Answer + ReasonContent. + srv := newNovitaServer(t, "/openai/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{ + "role": "assistant", + "content": "15% = 0.15; 0.15 * 80 = 12.The answer is 12.", + }, + }}, + }) + }) + defer srv.Close() + + apiKey := "test-key" + resp, err := newNovitaForTest(srv.URL).ChatWithMessages( + "qwen/qwen3-30b-a3b-fp8", + []Message{{Role: "user", Content: "15% of 80?"}}, + &APIConfig{ApiKey: &apiKey}, nil) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if *resp.Answer != "The answer is 12." { + t.Errorf("Answer=%q", *resp.Answer) + } + if *resp.ReasonContent != "15% = 0.15; 0.15 * 80 = 12." { + t.Errorf("ReasonContent=%q", *resp.ReasonContent) + } +} + +// deepseek-v3.1 / glm-4.5 with enable_thinking=true return reasoning +// in a separate `reasoning_content` field on the message rather than +// inline as .... The driver must surface this field +// to ChatResponse.ReasonContent. Live-confirmed against +// api.novita.ai/openai/v1/chat/completions with deepseek/deepseek-v3.1. +func TestNovitaChatExtractsReasoningContentField(t *testing.T) { + srv := newNovitaServer(t, "/openai/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{ + "role": "assistant", + "content": "4", + "reasoning_content": "2+2 is straightforward arithmetic: the answer is 4.", + }, + }}, + }) + }) + defer srv.Close() + + apiKey := "test-key" + resp, err := newNovitaForTest(srv.URL).ChatWithMessages( + "deepseek/deepseek-v3.1", + []Message{{Role: "user", Content: "2+2?"}}, + &APIConfig{ApiKey: &apiKey}, nil) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if resp.Answer == nil || resp.ReasonContent == nil { + t.Fatalf("Answer/ReasonContent must be non-nil pointers") + } + if *resp.Answer != "4" { + t.Errorf("Answer=%q", *resp.Answer) + } + if *resp.ReasonContent != "2+2 is straightforward arithmetic: the answer is 4." { + t.Errorf("ReasonContent=%q", *resp.ReasonContent) + } +} + +// Streaming deepseek-v3.1 with thinking on emits delta.reasoning_content +// (not delta.content with tags). The driver must forward +// those chunks via the sender's second arg. +func TestNovitaStreamExtractsDeltaReasoningContent(t *testing.T) { + srv := newNovitaSSEServer(t, "/openai/v1/chat/completions", + `data: {"choices":[{"index":0,"delta":{"role":"assistant","reasoning_content":"step 1. "}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"reasoning_content":"step 2."}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"final answer"},"finish_reason":"stop"}]}`+"\n"+ + `data: [DONE]`+"\n", + ) + defer srv.Close() + + apiKey := "test-key" + var content, reasoning []string + err := newNovitaForTest(srv.URL).ChatStreamlyWithSender( + "deepseek/deepseek-v3.1", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, + func(c *string, r *string) error { + if c != nil && r != nil { + t.Errorf("sender called with both args non-nil") + } + if r != nil && *r != "" { + reasoning = append(reasoning, *r) + } + if c != nil && *c != "" && *c != "[DONE]" { + content = append(content, *c) + } + return nil + }) + if err != nil { + t.Fatalf("stream: %v", err) + } + if got := strings.Join(reasoning, ""); got != "step 1. step 2." { + t.Errorf("reasoning=%q", got) + } + if got := strings.Join(content, ""); got != "final answer" { + t.Errorf("content=%q", got) + } +} + +// TestNovitaChatPropagatesEnableThinking pins the maintainer's +// requested behaviour: when ChatConfig.Thinking is set, the driver +// MUST forward it as Novita's documented `enable_thinking` body field +// so a tenant can switch a deepseek-v3.1 / glm-4.5 / qwen3 deployment +// out of its default thinking mode without prompt-level hacks. +func TestNovitaChatPropagatesEnableThinking(t *testing.T) { + cases := []struct { + name string + value bool + }{ + {"enabled", true}, + {"disabled", false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + srv := newNovitaServer(t, "/openai/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + got, present := body["enable_thinking"] + if !present { + t.Errorf("enable_thinking missing from body, want %v", tc.value) + } + if got != tc.value { + t.Errorf("enable_thinking=%v want %v", got, tc.value) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{"content": "ok"}, + }}, + }) + }) + defer srv.Close() + + apiKey := "test-key" + thinking := tc.value + _, err := newNovitaForTest(srv.URL).ChatWithMessages( + "qwen/qwen3-30b-a3b-fp8", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Thinking: &thinking}) + if err != nil { + t.Fatalf("Chat: %v", err) + } + }) + } +} + +// Sending enable_thinking when the caller didn't ask for it would +// silently flip behavior for downstream proxies that distinguish +// "field absent" from "field present with default". Leave it out. +func TestNovitaChatOmitsEnableThinkingWhenUnset(t *testing.T) { + srv := newNovitaServer(t, "/openai/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + if _, present := body["enable_thinking"]; present { + t.Errorf("enable_thinking must be absent when Thinking unset, got %v", body["enable_thinking"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{"content": "ok"}, + }}, + }) + }) + defer srv.Close() + + apiKey := "test-key" + _, err := newNovitaForTest(srv.URL).ChatWithMessages("m", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{}) // no Thinking + if err != nil { + t.Fatalf("Chat: %v", err) + } +} + +// TestNovitaStreamPropagatesEnableThinking mirrors the non-stream +// case for ChatStreamlyWithSender so callers get the same toggle +// regardless of streaming mode. +func TestNovitaStreamPropagatesEnableThinking(t *testing.T) { + var seen map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + raw, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(raw, &seen) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, + `data: {"choices":[{"delta":{"content":"ok"},"finish_reason":"stop"}]}`+"\n"+ + `data: [DONE]`+"\n", + ) + })) + defer srv.Close() + + apiKey := "test-key" + thinking := false + err := newNovitaForTest(srv.URL).ChatStreamlyWithSender( + "deepseek/deepseek-v3.1", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Thinking: &thinking}, + func(*string, *string) error { return nil }) + if err != nil { + t.Fatalf("Stream: %v", err) + } + if got, ok := seen["enable_thinking"].(bool); !ok || got != false { + t.Errorf("stream enable_thinking=%v want false", seen["enable_thinking"]) + } +} + +func TestNovitaChatRequiresAPIKey(t *testing.T) { + _, err := newNovitaForTest("http://unused").ChatWithMessages("m", + []Message{{Role: "user", Content: "x"}}, &APIConfig{}, nil) + if err == nil || !strings.Contains(err.Error(), "api key is required") { + t.Errorf("got %v", err) + } +} + +func TestNovitaChatRequiresMessages(t *testing.T) { + apiKey := "test-key" + _, err := newNovitaForTest("http://unused").ChatWithMessages("m", nil, + &APIConfig{ApiKey: &apiKey}, nil) + if err == nil || !strings.Contains(err.Error(), "messages is empty") { + t.Errorf("got %v", err) + } +} + +func TestNovitaChatRejectsHTTPError(t *testing.T) { + srv := newNovitaServer(t, "/openai/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"detail":"unauthorized"}`)) + }) + defer srv.Close() + + apiKey := "test-key" + _, err := newNovitaForTest(srv.URL).ChatWithMessages("m", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil) + if err == nil || !strings.Contains(err.Error(), "401") { + t.Errorf("got %v", err) + } +} + +// Streaming: an SSE stream that emits ... inline in +// delta.content must surface reasoning chunks through the sender's +// second arg, and visible content through the first. +func TestNovitaStreamSplitsThinkTags(t *testing.T) { + // Simulate the realistic case where tags span deltas — split + // "" across two chunks, and split "" too. + srv := newNovitaSSEServer(t, "/openai/v1/chat/completions", + `data: {"choices":[{"index":0,"delta":{"role":"assistant"}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"<"}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"think>"}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"Okay, "}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"compute. "}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"12"}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"."},"finish_reason":"stop"}]}`+"\n"+ + `data: [DONE]`+"\n", + ) + defer srv.Close() + + apiKey := "test-key" + var content, reasoning []string + err := newNovitaForTest(srv.URL).ChatStreamlyWithSender( + "qwen/qwen3-30b-a3b-fp8", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, + func(c *string, r *string) error { + if c != nil && r != nil { + t.Errorf("sender called with both args non-nil") + } + if r != nil && *r != "" { + reasoning = append(reasoning, *r) + } + if c != nil && *c != "" && *c != "[DONE]" { + content = append(content, *c) + } + return nil + }) + if err != nil { + t.Fatalf("stream: %v", err) + } + gotContent := strings.Join(content, "") + gotReason := strings.Join(reasoning, "") + if gotContent != "12." { + t.Errorf("content=%q want %q", gotContent, "12.") + } + if gotReason != "Okay, compute. " { + t.Errorf("reasoning=%q", gotReason) + } +} + +// Streaming for a non-reasoning model that emits only content chunks +// must continue to work unchanged. +func TestNovitaStreamPureContent(t *testing.T) { + srv := newNovitaSSEServer(t, "/openai/v1/chat/completions", + `data: {"choices":[{"index":0,"delta":{"role":"assistant"}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"Hello "}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":"world"},"finish_reason":"stop"}]}`+"\n"+ + `data: [DONE]`+"\n", + ) + defer srv.Close() + + apiKey := "test-key" + var chunks []string + var sawDone bool + err := newNovitaForTest(srv.URL).ChatStreamlyWithSender("meta-llama/llama-3.3-70b-instruct", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, + func(c *string, _ *string) error { + if c == nil { + return nil + } + if *c == "[DONE]" { + sawDone = true + return nil + } + chunks = append(chunks, *c) + return nil + }) + if err != nil { + t.Fatalf("stream: %v", err) + } + if strings.Join(chunks, "") != "Hello world" { + t.Errorf("content=%v", chunks) + } + if !sawDone { + t.Error("expected [DONE] sentinel") + } +} + +func TestNovitaStreamRequiresSender(t *testing.T) { + apiKey := "test-key" + err := newNovitaForTest("http://unused").ChatStreamlyWithSender("m", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, nil) + if err == nil || !strings.Contains(err.Error(), "sender is required") { + t.Errorf("got %v", err) + } +} + +func TestNovitaStreamRejectsExplicitFalse(t *testing.T) { + apiKey := "test-key" + stream := false + err := newNovitaForTest("http://unused").ChatStreamlyWithSender("m", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Stream: &stream}, + func(*string, *string) error { return nil }) + if err == nil || !strings.Contains(err.Error(), "stream must be true") { + t.Errorf("got %v", err) + } +} + +func TestNovitaListModelsHappyPath(t *testing.T) { + srv := newNovitaServer(t, "/openai/v1/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]interface{}{ + {"id": "meta-llama/llama-3.3-70b-instruct"}, + {"id": "qwen/qwen3-30b-a3b-fp8"}, + {"id": "deepseek/deepseek-v4-pro"}, + }, + }) + }) + defer srv.Close() + + apiKey := "test-key" + ids, err := newNovitaForTest(srv.URL).ListModels(&APIConfig{ApiKey: &apiKey}) + if err != nil { + t.Fatalf("ListModels: %v", err) + } + if len(ids) != 3 { + t.Errorf("ids=%v", ids) + } +} + +func TestNovitaCheckConnection(t *testing.T) { + srv := newNovitaServer(t, "/openai/v1/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{"data": []map[string]interface{}{{"id": "x"}}}) + }) + defer srv.Close() + + apiKey := "test-key" + if err := newNovitaForTest(srv.URL).CheckConnection(&APIConfig{ApiKey: &apiKey}); err != nil { + t.Errorf("CheckConnection: %v", err) + } +} + +func TestNovitaEmbedReturnsNoSuchMethod(t *testing.T) { + m := "x" + _, err := newNovitaForTest("http://unused").Embed(&m, []string{"a"}, &APIConfig{}, nil) + if err == nil || !strings.Contains(err.Error(), "no such method") { + t.Errorf("got %v", err) + } +} + +func TestNovitaRerankReturnsNoSuchMethod(t *testing.T) { + m := "x" + _, err := newNovitaForTest("http://unused").Rerank(&m, "q", []string{"a"}, &APIConfig{}, &RerankConfig{TopN: 1}) + if err == nil || !strings.Contains(err.Error(), "no such method") { + t.Errorf("got %v", err) + } +} + +func TestNovitaBalanceReturnsNoSuchMethod(t *testing.T) { + if _, err := newNovitaForTest("http://unused").Balance(&APIConfig{}); err == nil || !strings.Contains(err.Error(), "no such method") { + t.Errorf("got %v", err) + } +} + +func TestNovitaAudioOCRReturnNoSuchMethod(t *testing.T) { + m := "x" + v := newNovitaForTest("http://unused") + if _, err := v.TranscribeAudio(&m, &m, &APIConfig{}, nil); err == nil || !strings.Contains(err.Error(), "no such method") { + t.Errorf("TranscribeAudio: %v", err) + } + if _, err := v.AudioSpeech(&m, &m, &APIConfig{}, nil); err == nil || !strings.Contains(err.Error(), "no such method") { + t.Errorf("AudioSpeech: %v", err) + } + if _, err := v.OCRFile(&m, &m, &APIConfig{}, nil); err == nil || !strings.Contains(err.Error(), "no such method") { + t.Errorf("OCRFile: %v", err) + } +} + +// TestNovitaBaseURLTrimsTrailingSlash pins the fix for a `//`-in-path +// bug a tenant could hit by configuring a baseURL like +// "https://api.novita.ai/v3/openai/". Every URL the driver builds via +// fmt.Sprintf("%s/%s", base, suffix) would then produce a double +// slash. baseURLForRegion now trims the trailing "/" so all three +// endpoint builders (Chat, Stream, ListModels) emit clean paths. +func TestNovitaBaseURLTrimsTrailingSlash(t *testing.T) { + cases := []struct { + name string + path string + method string + invoke func(n *NovitaModel, apiKey string) error + urlSuffix URLSuffix + respBody string + respHeaders map[string]string + }{ + { + name: "Chat", + path: "/openai/v1/chat/completions", + method: http.MethodPost, + urlSuffix: URLSuffix{Chat: "openai/v1/chat/completions"}, + invoke: func(n *NovitaModel, apiKey string) error { + _, err := n.ChatWithMessages("m", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + respBody: `{"choices":[{"message":{"content":"ok"}}]}`, + }, + { + name: "ListModels", + path: "/openai/v1/models", + method: http.MethodGet, + urlSuffix: URLSuffix{Models: "openai/v1/models"}, + invoke: func(n *NovitaModel, apiKey string) error { + _, err := n.ListModels(&APIConfig{ApiKey: &apiKey}) + return err + }, + respBody: `{"data":[]}`, + }, + { + name: "Stream", + path: "/openai/v1/chat/completions", + method: http.MethodPost, + urlSuffix: URLSuffix{Chat: "openai/v1/chat/completions"}, + invoke: func(n *NovitaModel, apiKey string) error { + return n.ChatStreamlyWithSender("m", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, + func(*string, *string) error { return nil }) + }, + respBody: `data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}]}` + "\n" + + `data: [DONE]` + "\n", + respHeaders: map[string]string{"Content-Type": "text/event-stream"}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // The load-bearing assertion: path is the clean + // "/openai/v1/chat/completions" or "/openai/v1/models", never "//chat/...". + if r.URL.Path != tc.path { + t.Errorf("path=%q want %q (double-slash bug?)", r.URL.Path, tc.path) + return + } + if r.Method != tc.method { + t.Errorf("method=%s want %s", r.Method, tc.method) + return + } + for k, v := range tc.respHeaders { + w.Header().Set(k, v) + } + _, _ = io.WriteString(w, tc.respBody) + })) + defer srv.Close() + + // Configure baseURL WITH a trailing slash on purpose. + n := NewNovitaModel( + map[string]string{"default": srv.URL + "/"}, + tc.urlSuffix, + ) + apiKey := "test-key" + if err := tc.invoke(n, apiKey); err != nil { + t.Fatalf("%s: %v", tc.name, err) + } + }) + } +}