// // 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 ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "net/url" "ragflow/internal/common" "sort" "strings" "time" ) const ( nvidiaHostedAPIHost = "integrate.api.nvidia.com" nvidiaCatalogURL = "https://api.ngc.nvidia.com/v2/search/catalog/resources/ENDPOINT" nvidiaCatalogPageSize = 500 ) // NvidiaModel implements ModelDriver for Nvidia type NvidiaModel struct { baseModel BaseModel catalogURL string hostedAPIHost string } // NewNvidiaModel creates a new Nvidia model instance func NewNvidiaModel(baseURL map[string]string, urlSuffix URLSuffix) *NvidiaModel { return &NvidiaModel{ baseModel: BaseModel{ BaseURL: baseURL, URLSuffix: urlSuffix, httpClient: NewDriverHTTPClient(), }, catalogURL: nvidiaCatalogURL, hostedAPIHost: nvidiaHostedAPIHost, } } func (n NvidiaModel) NewInstance(baseURL map[string]string) ModelDriver { return NewNvidiaModel(baseURL, n.baseModel.URLSuffix) } func (n NvidiaModel) Name() string { return "nvidia" } func (n *NvidiaModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := n.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } if len(messages) == 0 { return nil, fmt.Errorf("messages is empty") } resolvedBaseURL, err := n.baseModel.GetBaseURL(apiConfig) if err != nil { return nil, err } baseURL := resolvedBaseURL if baseURL == "" { baseURL = resolvedBaseURL } url := fmt.Sprintf("%s/%s", baseURL, n.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) if chatModelConfig != nil { if chatModelConfig.Thinking != nil { if *chatModelConfig.Thinking { reqBody["thinking"] = map[string]interface{}{"type": "enabled"} } else { reqBody["thinking"] = map[string]interface{}{"type": "disabled"} } } } jsonData, err := json.Marshal(reqBody) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) } ctx, cancel := context.WithTimeout(ctx, 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.baseModel.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") } content, ok := messageMap["content"].(string) toolCalls := extractToolCalls(messageMap) if !ok && len(toolCalls) == 0 { return nil, fmt.Errorf("invalid content format") } var reasonContent string if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { reasonContent, ok = messageMap["reasoning_content"].(string) if !ok { return nil, fmt.Errorf("invalid content format") } if reasonContent != "" && reasonContent[0] == '\n' { reasonContent = reasonContent[1:] } } chatResponse := &ChatResponse{ Answer: &content, ReasonContent: &reasonContent, ToolCalls: toolCalls, } return chatResponse, nil } func (n *NvidiaModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { if err := n.baseModel.APIConfigCheck(apiConfig); err != nil { return err } if sender == nil { return fmt.Errorf("sender is required") } if len(messages) == 0 { return fmt.Errorf("messages is empty") } resolvedBaseURL, err := n.baseModel.GetBaseURL(apiConfig) if err != nil { return err } baseURL := resolvedBaseURL if baseURL == "" { baseURL = resolvedBaseURL } url := fmt.Sprintf("%s/%s", baseURL, n.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(modelConfig, modelName, messages, true) if modelConfig != nil { if modelConfig.Thinking != nil { if *modelConfig.Thinking { reqBody["thinking"] = map[string]interface{}{"type": "enabled"} } else { reqBody["thinking"] = map[string]interface{}{"type": "disabled"} } } } jsonData, err := json.Marshal(reqBody) if err != nil { return fmt.Errorf("failed to marshal request: %w", err) } ctx, cancel := context.WithTimeout(ctx, streamCallTimeout) defer cancel() req, err := http.NewRequestWithContext(ctx, "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.baseModel.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)) } accumulatedToolCalls := make(map[int]map[string]any) if _, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { choices, ok := event["choices"].([]interface{}) if !ok || len(choices) == 0 { return nil } firstChoice, ok := choices[0].(map[string]interface{}) if !ok { return nil } delta, ok := firstChoice["delta"].(map[string]interface{}) if !ok { return nil } accumulateToolCallDeltas(delta, accumulatedToolCalls) reasoningContent, ok := delta["reasoning_content"].(string) if ok && reasoningContent != "" { if err := sender(nil, &reasoningContent); err != nil { return err } } content, ok := delta["content"].(string) if ok && content != "" { if err := sender(&content, nil); err != nil { return err } } return nil }); err != nil { return fmt.Errorf("failed to scan response body: %w", err) } setSortedToolCallsResult(modelConfig, accumulatedToolCalls) endOfStream := "[DONE]" if err = sender(&endOfStream, nil); err != nil { return err } return nil } type nvidiaEmbeddingResponse struct { Data []struct { Index int `json:"index"` Embedding []float64 `json:"embedding"` } `json:"data"` } func (n NvidiaModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { if err := n.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } if len(texts) == 0 { return []EmbeddingData{}, nil } if modelName == nil || *modelName == "" { return nil, fmt.Errorf("model name is required") } resolvedBaseURL, err := n.baseModel.GetBaseURL(apiConfig) if err != nil { return nil, err } baseURL := resolvedBaseURL if baseURL == "" { baseURL = resolvedBaseURL } url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), n.baseModel.URLSuffix.Embedding) reqBody := map[string]interface{}{ "model": *modelName, "input": texts, "input_type": "query", "encoding_format": "float", "truncate": "END", } if embeddingConfig != nil && embeddingConfig.Dimension > 0 { reqBody["dimensions"] = embeddingConfig.Dimension } jsonData, err := json.Marshal(reqBody) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) } ctx, cancel := context.WithTimeout(ctx, 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.baseModel.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("Nvidia embeddings API error: %s, body: %s", resp.Status, string(body)) } var parsed nvidiaEmbeddingResponse if err = json.Unmarshal(body, &parsed); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } var embeddings []EmbeddingData for _, dataElem := range parsed.Data { var embeddingData EmbeddingData embeddingData.Embedding = dataElem.Embedding embeddingData.Index = dataElem.Index embeddings = append(embeddings, embeddingData) } return embeddings, nil } // nvidiaRerankRequest mirrors the NIM /ranking request shape: // query is an object with a "text" field, passages is an array of // objects each with a "text" field. truncate=END matches the Python // NvidiaRerank reference at rag/llm/rerank_model.py. type nvidiaRerankRequest struct { Model string `json:"model"` Query nvidiaRerankText `json:"query"` Passages []nvidiaRerankText `json:"passages"` Truncate string `json:"truncate,omitempty"` TopN int `json:"top_n"` } type nvidiaRerankText struct { Text string `json:"text"` } // nvidiaRerankResponse maps the NIM rankings array. Each entry pairs // the original passage index with a logit score; the caller uses the // index to restore original input order. type nvidiaRerankResponse struct { Rankings []struct { Index int `json:"index"` Logit float64 `json:"logit"` } `json:"rankings"` } // Rerank scores documents against the query using an NVIDIA NIM // reranking model. Mirrors the Python NvidiaRerank class in // rag/llm/rerank_model.py for payload shape (passages/query/logit). // Defaults top_n to len(documents) so the API returns a score per // input; callers may shrink it via RerankConfig.TopN, in which case // only the top RerankConfig.TopN entries come back. Returned // RerankResult entries are in the API's ranking order; callers that // need original-input order should sort by Index. Same return-shape // contract as the Aliyun and ZhipuAI Rerank drivers. func (n NvidiaModel) Rerank(ctx context.Context, modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig, modelUsage *common.ModelUsage) (*RerankResponse, error) { if err := n.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } if len(documents) == 0 { return &RerankResponse{}, nil } if modelName == nil || *modelName == "" { return nil, fmt.Errorf("model name is required") } resolvedBaseURL, err := n.baseModel.GetBaseURL(apiConfig) if err != nil { return nil, err } baseURL := resolvedBaseURL if baseURL == "" { baseURL = resolvedBaseURL } url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), n.baseModel.URLSuffix.Rerank) topN := len(documents) if rerankConfig != nil && rerankConfig.TopN > 0 && rerankConfig.TopN < topN { topN = rerankConfig.TopN } passages := make([]nvidiaRerankText, len(documents)) for i, doc := range documents { passages[i] = nvidiaRerankText{Text: doc} } reqBody := nvidiaRerankRequest{ Model: *modelName, Query: nvidiaRerankText{Text: query}, Passages: passages, Truncate: "END", TopN: topN, } jsonData, err := json.Marshal(reqBody) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) } ctx, cancel := context.WithTimeout(ctx, 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.baseModel.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("Nvidia rerank API error: %s, body: %s", resp.Status, string(body)) } var parsed nvidiaRerankResponse if err = json.Unmarshal(body, &parsed); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } rerankResponse := RerankResponse{Data: make([]RerankResult, 0, len(parsed.Rankings))} for _, r := range parsed.Rankings { if r.Index < 0 || r.Index >= len(documents) { return nil, fmt.Errorf("unexpected rerank index %d for %d inputs", r.Index, len(documents)) } rerankResponse.Data = append(rerankResponse.Data, RerankResult{ Index: r.Index, RelevanceScore: r.Logit, }) } return &rerankResponse, nil } // TranscribeAudio transcribe audio func (n *NvidiaModel) TranscribeAudio(ctx context.Context, modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, modelUsage *common.ModelUsage) (*ASRResponse, error) { return nil, fmt.Errorf("%s, no such method", n.Name()) } func (n *NvidiaModel) TranscribeAudioWithSender(ctx context.Context, modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { return fmt.Errorf("%s, no such method", n.Name()) } // AudioSpeech convert text to audio func (n *NvidiaModel) AudioSpeech(ctx context.Context, modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, modelUsage *common.ModelUsage) (*TTSResponse, error) { return nil, fmt.Errorf("%s, no such method", n.Name()) } func (n *NvidiaModel) AudioSpeechWithSender(ctx context.Context, modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { return fmt.Errorf("%s, no such method", n.Name()) } // OCRFile OCR file func (n *NvidiaModel) OCRFile(ctx context.Context, modelName *string, content []byte, url *string, apiConfig *APIConfig, ocrConfig *OCRConfig, modelUsage *common.ModelUsage) (*OCRFileResponse, error) { return nil, fmt.Errorf("%s, no such method", n.Name()) } // ParseFile parse file func (n *NvidiaModel) ParseFile(ctx context.Context, modelName *string, content []byte, url *string, apiConfig *APIConfig, parseFileConfig *ParseFileConfig, modelUsage *common.ModelUsage) (*ParseFileResponse, error) { return nil, fmt.Errorf("%s, no such method", n.Name()) } // ListModels calls /v1/models on the configured NVIDIA NIM base URL // and returns the list of available model ids. The endpoint is // OpenAI-compatible, so the parsing follows the same shape used by // the moonshot, xai, and openai drivers. func (n NvidiaModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]ListModelResponse, error) { if err := n.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } resolvedBaseURL, err := n.baseModel.GetBaseURL(apiConfig) if err != nil { return nil, err } url := fmt.Sprintf("%s/%s", strings.TrimRight(resolvedBaseURL, "/"), strings.TrimLeft(n.baseModel.URLSuffix.Models, "/")) modelListCtx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) defer cancel() req, err := http.NewRequestWithContext(modelListCtx, "GET", url, 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("Accept", "application/json") resp, err := n.baseModel.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("Nvidia models API error: %s, body: %s", resp.Status, string(body)) } // Parse response var modelList ModelList if err = json.Unmarshal(body, &modelList); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } if modelList.Models == nil { return nil, fmt.Errorf("invalid models list format") } var provider *Provider if pm := GetProviderManager(); pm != nil { provider = pm.FindProvider("NVIDIA") } models := parseNvidiaModelList(modelList, provider) if n.usesHostedCatalog(resolvedBaseURL) { catalogCtx, catalogCancel := context.WithTimeout(ctx, nonStreamCallTimeout) catalog, catalogErr := n.fetchHostedCatalog(catalogCtx) catalogCancel() if catalogErr != nil { return nil, catalogErr } models = filterNvidiaHostedModels(models, catalog, time.Now().UTC()) } if len(models) == 0 { return nil, fmt.Errorf("Nvidia models API returned no usable models") } return models, nil } type nvidiaCatalogLabel struct { Key string `json:"key"` Values []string `json:"values"` UnresolvedValues []string `json:"unresolvedValues"` } type nvidiaCatalogAttribute struct { Key string `json:"key"` Value string `json:"value"` } type nvidiaCatalogResource struct { Name string `json:"name"` DisplayName string `json:"displayName"` Labels []nvidiaCatalogLabel `json:"labels"` Attributes []nvidiaCatalogAttribute `json:"attributes"` } type nvidiaCatalogGroup struct { GroupValue string `json:"groupValue"` TotalCount int `json:"totalCount"` Resources []nvidiaCatalogResource `json:"resources"` } type nvidiaCatalogResponse struct { Results []nvidiaCatalogGroup `json:"results"` } func (n NvidiaModel) usesHostedCatalog(baseURL string) bool { parsed, err := url.Parse(baseURL) return err == nil && strings.EqualFold(parsed.Hostname(), n.hostedAPIHost) } func (n NvidiaModel) fetchHostedCatalog(ctx context.Context) (*nvidiaCatalogResponse, error) { parsed, err := url.Parse(n.catalogURL) if err != nil { return nil, fmt.Errorf("failed to parse Nvidia endpoint catalog URL: %w", err) } var merged *nvidiaCatalogResponse endpointGroupIndex := -1 totalCount := -1 for page := 0; ; page++ { catalogQuery, err := json.Marshal(map[string]int{"page": page, "pageSize": nvidiaCatalogPageSize}) if err != nil { return nil, fmt.Errorf("failed to encode Nvidia endpoint catalog query: %w", err) } pageURL := *parsed query := pageURL.Query() query.Set("q", string(catalogQuery)) query.Set("group-labels-by-labelset", "true") pageURL.RawQuery = query.Encode() catalogPage, err := n.fetchHostedCatalogPage(ctx, pageURL.String()) if err != nil { return nil, err } pageGroupIndex := -1 for i := range catalogPage.Results { if catalogPage.Results[i].GroupValue == "ENDPOINT" { pageGroupIndex = i break } } if pageGroupIndex < 0 { return nil, fmt.Errorf("Nvidia endpoint catalog response is missing the ENDPOINT group") } pageGroup := catalogPage.Results[pageGroupIndex] if pageGroup.TotalCount < 0 { return nil, fmt.Errorf("Nvidia endpoint catalog returned invalid total count %d", pageGroup.TotalCount) } if merged == nil { merged = catalogPage endpointGroupIndex = pageGroupIndex totalCount = pageGroup.TotalCount merged.Results[endpointGroupIndex].Resources = nil } else if pageGroup.TotalCount != totalCount { return nil, fmt.Errorf("Nvidia endpoint catalog total count changed from %d to %d", totalCount, pageGroup.TotalCount) } merged.Results[endpointGroupIndex].Resources = append(merged.Results[endpointGroupIndex].Resources, pageGroup.Resources...) if len(pageGroup.Resources) < nvidiaCatalogPageSize { break } } if merged == nil || endpointGroupIndex < 0 { return nil, fmt.Errorf("Nvidia endpoint catalog returned no pages") } if len(merged.Results[endpointGroupIndex].Resources) < totalCount { return nil, fmt.Errorf( "Nvidia endpoint catalog returned %d of %d resources", len(merged.Results[endpointGroupIndex].Resources), totalCount, ) } return merged, nil } func (n NvidiaModel) fetchHostedCatalogPage(ctx context.Context, pageURL string) (*nvidiaCatalogResponse, error) { req, err := http.NewRequestWithContext(ctx, http.MethodGet, pageURL, nil) if err != nil { return nil, fmt.Errorf("failed to create Nvidia endpoint catalog request: %w", err) } req.Header.Set("Accept", "application/json") resp, err := n.baseModel.httpClient.Do(req) if err != nil { return nil, fmt.Errorf("failed to request Nvidia endpoint catalog: %w", err) } defer resp.Body.Close() body, err := io.ReadAll(resp.Body) if err != nil { return nil, fmt.Errorf("failed to read Nvidia endpoint catalog: %w", err) } if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("Nvidia endpoint catalog error: %s, body: %s", resp.Status, string(body)) } var catalog nvidiaCatalogResponse if err = json.Unmarshal(body, &catalog); err != nil { return nil, fmt.Errorf("failed to parse Nvidia endpoint catalog: %w", err) } return &catalog, nil } func filterNvidiaHostedModels(models []ListModelResponse, catalog *nvidiaCatalogResponse, now time.Time) []ListModelResponse { if catalog == nil { return nil } var resources []nvidiaCatalogResource for _, group := range catalog.Results { if group.GroupValue != "ENDPOINT" { continue } if group.TotalCount < 0 || len(group.Resources) < group.TotalCount { return nil } resources = group.Resources break } if resources == nil { return nil } activeEndpoints := make(map[string]struct{}, len(resources)) for _, resource := range resources { if !nvidiaCatalogResourceIsActive(resource, now) { continue } publishers := nvidiaCatalogLabelValues(resource, "publisher", true) displayName := resource.DisplayName if displayName == "" { displayName = resource.Name } if len(publishers) == 0 || displayName == "" { continue } activeEndpoints[nvidiaCatalogKey(publishers[0], displayName)] = struct{}{} } filtered := make([]ListModelResponse, 0, len(models)) for _, model := range models { publisher, modelName, found := strings.Cut(model.Name, "/") if !found { continue } if _, ok := activeEndpoints[nvidiaCatalogKey(publisher, modelName)]; ok { filtered = append(filtered, model) } } return filtered } func nvidiaCatalogResourceIsActive(resource nvidiaCatalogResource, now time.Time) bool { if !containsNvidiaCatalogValue(nvidiaCatalogLabelValues(resource, "nimType", false), "Free Endpoint") { return false } deprecation := "" for _, attribute := range resource.Attributes { if attribute.Key == "DEPRECATION" { deprecation = attribute.Value break } } if deprecation == "" { return true } cutoff, err := time.Parse("01/02/2006", deprecation) if err != nil { return false } today := time.Date(now.UTC().Year(), now.UTC().Month(), now.UTC().Day(), 0, 0, 0, 0, time.UTC) return cutoff.After(today) } func nvidiaCatalogLabelValues(resource nvidiaCatalogResource, key string, unresolved bool) []string { for _, label := range resource.Labels { if label.Key != key { continue } if unresolved { return label.UnresolvedValues } return label.Values } return nil } func nvidiaCatalogKey(publisher, modelName string) string { publisher = strings.ReplaceAll(strings.ToLower(strings.TrimSpace(publisher)), "_", "-") return publisher + "\x00" + modelName } func containsNvidiaCatalogValue(values []string, target string) bool { for _, value := range values { if value == target { return true } } return false } func parseNvidiaModelList(modelList ModelList, provider *Provider) []ListModelResponse { const defaultMaxTokens = 8192 models := make([]ListModelResponse, 0, len(modelList.Models)) seen := make(map[string]struct{}, len(modelList.Models)) for _, item := range modelList.Models { modelName := strings.TrimSpace(item.ID) if modelName == "" { continue } if _, ok := seen[modelName]; ok { continue } seen[modelName] = struct{}{} response := ListModelResponse{Name: modelName} var preset *Model if provider != nil { for _, model := range provider.Models { if strings.EqualFold(model.Name, modelName) { preset = model break } } } if preset != nil { response.MaxTokens = preset.MaxTokens response.ModelTypes = append([]string(nil), preset.ModelTypes...) response.Thinking = preset.Thinking response.MaxDimension = preset.MaxDimension response.Dimensions = append([]int(nil), preset.Dimensions...) } else { maxTokens := defaultMaxTokens response.MaxTokens = &maxTokens response.ModelTypes = InferModelTypes(modelName) } if len(response.ModelTypes) == 0 { response.ModelTypes = InferModelTypes(modelName) } models = append(models, response) } sort.Slice(models, func(i, j int) bool { return models[i].Name < models[j].Name }) return models } func (n NvidiaModel) Balance(ctx context.Context, apiConfig *APIConfig) (map[string]interface{}, error) { return nil, fmt.Errorf("no such method") } // CheckConnection verifies that the configured NVIDIA NIM base URL // is reachable and that the API key is accepted, by issuing a // lightweight ListModels call. Mirrors the pattern used by the xai, // moonshot, deepseek, aliyun, and gitee drivers. func (n NvidiaModel) CheckConnection(ctx context.Context, apiConfig *APIConfig) error { _, err := n.ListModels(ctx, apiConfig) return err } func (n *NvidiaModel) ListTasks(ctx context.Context, apiConfig *APIConfig) ([]ListTaskStatus, error) { return nil, fmt.Errorf("%s, no such method", n.Name()) } func (n *NvidiaModel) ShowTask(ctx context.Context, taskID string, apiConfig *APIConfig) (*TaskResponse, error) { return nil, fmt.Errorf("%s, no such method", n.Name()) }