package models import ( "bufio" "bytes" "encoding/json" "fmt" "io" "net/http" "ragflow/internal/common" "strings" "time" ) // OpenRouterModel implements ModelDriver for OpenRouter AI type OpenRouterModel struct { BaseURL map[string]string URLSuffix URLSuffix httpClient *http.Client } func (o *OpenRouterModel) ParseFile() { //TODO implement me panic("implement me") } // NewOpenRouterModel creates a new OpenRouter AI model instance func NewOpenRouterModel(baseURL map[string]string, urlSuffix URLSuffix) *OpenRouterModel { return &OpenRouterModel{ BaseURL: baseURL, URLSuffix: urlSuffix, httpClient: &http.Client{ Timeout: 120 * time.Second, Transport: &http.Transport{ MaxIdleConns: 10, MaxIdleConnsPerHost: 100, IdleConnTimeout: 90 * time.Second, DisableCompression: false, }, }, } } func (o *OpenRouterModel) NewInstance(baseURL map[string]string) ModelDriver { return &OpenRouterModel{ BaseURL: baseURL, URLSuffix: o.URLSuffix, httpClient: &http.Client{ Timeout: 120 * time.Second, Transport: &http.Transport{ MaxIdleConns: 10, MaxIdleConnsPerHost: 100, IdleConnTimeout: 90 * time.Second, DisableCompression: false, }, }, } } func (o *OpenRouterModel) Name() string { return "openrouter" } func (o *OpenRouterModel) 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 nil or empty") } if len(messages) == 0 { return nil, fmt.Errorf("messages is empty") } var region = "default" if apiConfig.Region != nil && *apiConfig.Region != "" { region = *apiConfig.Region } url := fmt.Sprintf("%s/%s", o.BaseURL[region], o.URLSuffix.Chat) // Convert messages to API format apiMessages := make([]map[string]interface{}, len(messages)) for i, msg := range messages { apiMessages[i] = map[string]interface{}{ "role": msg.Role, "content": msg.Content, } } // Build request body reqBody := map[string]interface{}{ "model": modelName, "messages": apiMessages, "stream": false, "temperature": 1, } if chatModelConfig != nil { if chatModelConfig.Temperature != nil { reqBody["temperature"] = *chatModelConfig.Temperature } if chatModelConfig.MaxTokens != nil { reqBody["max_tokens"] = *chatModelConfig.MaxTokens } if chatModelConfig.Stream != nil { reqBody["stream"] = *chatModelConfig.Stream } if chatModelConfig.TopP != nil { reqBody["top_p"] = *chatModelConfig.TopP } if chatModelConfig.DoSample != nil { reqBody["do_sample"] = *chatModelConfig.DoSample } reqBody["reasoning"] = map[string]interface{}{ "effort": "low", } } jsonData, err := json.Marshal(reqBody) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) } req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } req.Header.Add("Content-Type", "application/json") req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) resp, err := o.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 body: %w", err) } if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("failed to send request: %d %s", resp.StatusCode, string(body)) } // Parse response var result map[string]interface{} if err := json.Unmarshal(body, &result); err != nil { return nil, fmt.Errorf("failed to unmarshal response: %w", err) } choices, ok := result["choices"].([]interface{}) if !ok { return nil, fmt.Errorf("no choices in response") } firstChoice, ok := choices[0].(map[string]interface{}) if !ok { return nil, fmt.Errorf("no choices in response") } messageMap, ok := firstChoice["message"].(map[string]interface{}) if !ok { return nil, fmt.Errorf("no message in response") } content, ok := messageMap["content"].(string) if !ok { return nil, fmt.Errorf("no message in response") } var reasonContent string if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { reasonContent, ok = messageMap["reasoning"].(string) if !ok { return nil, fmt.Errorf("invalid content format") } if reasonContent != "" && reasonContent[0] == '\n' { reasonContent = reasonContent[1:] } } chatResponse := &ChatResponse{ Answer: &content, ReasonContent: &reasonContent, } return chatResponse, nil } func (o *OpenRouterModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error { if len(messages) == 0 { return fmt.Errorf("messages is empty") } var region = "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { region = *apiConfig.Region } url := fmt.Sprintf("%s/%s", o.BaseURL[region], o.URLSuffix.Chat) modelType := strings.Split(modelName, "_")[0] if modelType == "qwen" || modelType == "glm" { url = fmt.Sprintf("%s/%s", o.BaseURL[region], o.URLSuffix.AsyncChat) } // Convert messages to API format 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, "temperature": 1, } if modelConfig != nil { if modelConfig.Stream != nil { reqBody["stream"] = *modelConfig.Stream } if modelConfig.MaxTokens != nil { reqBody["max_tokens"] = *modelConfig.MaxTokens } if modelConfig.Temperature != nil { reqBody["temperature"] = *modelConfig.Temperature } if modelConfig.DoSample != nil { reqBody["do_sample"] = *modelConfig.DoSample } if modelConfig.TopP != nil { reqBody["top_p"] = *modelConfig.TopP } if modelConfig.Stop != nil { reqBody["stop"] = *modelConfig.Stop } 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) } req, err := http.NewRequest("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 := o.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("invalid status code: %d, body: %s", resp.StatusCode, string(body)) } // SSE parsing: read line by line 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:") { continue } // Extract JSON after "data:" data := strings.TrimSpace(line[5:]) // [DONE] marks the end of stream if data == "[DONE]" { break } // Parse the JSON event 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 } reasoningContent, ok := delta["reasoning"].(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 } } finishReason, ok := firstChoice["finish_reason"].(string) if ok && finishReason != "" { break } } // Send [DONE] marker for OpenAI compatibility endOfStream := "[DONE]" if err = sender(&endOfStream, nil); err != nil { return err } return scanner.Err() } type openrouterEmbeddingResponse struct { Data []openrouterEmbeddingData `json:"data"` Model string `json:"model"` Object string `json:"object"` Usage openrouterUsage `json:"usage"` } type openrouterEmbeddingData struct { Embedding []float64 `json:"embedding"` Object string `json:"object"` Index int `json:"index"` } type openrouterUsage struct { PromptTokens int `json:"prompt_tokens"` TotalTokens int `json:"total_tokens"` } func (o *OpenRouterModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([]EmbeddingData, error) { if len(texts) == 0 { return []EmbeddingData{}, nil } if modelName == nil || *modelName == "" { return nil, fmt.Errorf("model name is required") } var region = "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { region = *apiConfig.Region } url := fmt.Sprintf("%s/%s", o.BaseURL[region], o.URLSuffix.Embedding) reqBody := map[string]interface{}{ "model": *modelName, "input": texts, } 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) } req, err := http.NewRequest("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") if apiConfig != nil && apiConfig.ApiKey != nil && *apiConfig.ApiKey != "" { req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) } resp, err := o.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("OpenRouter embedding API error: status %d, body: %s", resp.StatusCode, string(body)) } var parsed openrouterEmbeddingResponse 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 } // OpenRouterRerankRequest OpenRouter official rerank request format type OpenRouterRerankRequest struct { Model string `json:"model"` Query string `json:"query"` Documents []string `json:"documents"` TopN int `json:"top_n,omitempty"` } // OpenRouterRerankResponse OpenRouter official rerank response format type OpenRouterRerankResponse struct { Model string `json:"model"` ID string `json:"id"` Results []struct { Index int `json:"index"` RelevanceScore float64 `json:"relevance_score"` Document *struct { Text string `json:"text"` } `json:"document,omitempty"` } `json:"results"` } func (o *OpenRouterModel) Rerank(modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig) (*RerankResponse, error) { if len(documents) == 0 { return &RerankResponse{}, nil } var region = "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { region = *apiConfig.Region } var topN = rerankConfig.TopN if rerankConfig.TopN == 0 { topN = len(documents) } reqBody := OpenRouterRerankRequest{ Model: *modelName, Query: query, Documents: documents, TopN: topN, } jsonData, err := json.Marshal(reqBody) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) } url := fmt.Sprintf("%s/%s", strings.TrimSuffix(o.BaseURL[region], "/"), o.URLSuffix.Rerank) req, err := http.NewRequest("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 := o.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("OpenRouter Rerank API error: %s, body: %s", resp.Status, string(body)) } var rerankResp OpenRouterRerankResponse if err = json.Unmarshal(body, &rerankResp); err != nil { return nil, fmt.Errorf("failed to decode response: %w", err) } var rerankResponse RerankResponse for _, result := range rerankResp.Results { rerankResult := RerankResult{ Index: result.Index, RelevanceScore: result.RelevanceScore, } rerankResponse.Data = append(rerankResponse.Data, rerankResult) } return &rerankResponse, nil } // TranscribeAudio transcribe audio func (o *OpenRouterModel) TranscribeAudio(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig) (*ASRResponse, error) { return nil, fmt.Errorf("%s, no such method", o.Name()) } func (z *OpenRouterModel) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, sender func(*string, *string) error) error { return fmt.Errorf("%s, no such method", z.Name()) } // AudioSpeech convert audio to text func (o *OpenRouterModel) AudioSpeech(modelName *string, audioContent *string, apiConfig *APIConfig, asrConfig *TTSConfig) (*TTSResponse, error) { return nil, fmt.Errorf("%s, no such method", o.Name()) } func (z *OpenRouterModel) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, sender func(*string, *string) error) error { return fmt.Errorf("%s, no such method", z.Name()) } // OCRFile OCR file func (m *OpenRouterModel) OCRFile(modelName *string, content []byte, url *string, apiConfig *APIConfig, ocrConfig *OCRConfig) (*OCRResponse, error) { return nil, fmt.Errorf("%s, no such method", m.Name()) } func (o *OpenRouterModel) ListModels(apiConfig *APIConfig) ([]string, error) { var region = "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { region = *apiConfig.Region } url := fmt.Sprintf("%s/%s", o.BaseURL[region], o.URLSuffix.Models) // Build request body reqBody := map[string]interface{}{} 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)) 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 := o.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 : %s", resp.StatusCode, string(body)) } // Parse response var result map[string]interface{} if err = json.Unmarshal(body, &result); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } // 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) } return models, nil } func (o *OpenRouterModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) { region := "default" if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { region = *apiConfig.Region } url := fmt.Sprintf("%s/%s", o.BaseURL[region], o.URLSuffix.Balance) 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)) resp, err := o.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 struct { Data struct { TotalCredits float64 `json:"total_credits"` TotalUsage float64 `json:"total_usage"` } `json:"data"` } if err := json.Unmarshal(body, &result); err != nil { return nil, fmt.Errorf("failed to parse balance response: %w", err) } remainingBalance := result.Data.TotalCredits - result.Data.TotalUsage return map[string]interface{}{ "total_credits": result.Data.TotalCredits, "total_usage": result.Data.TotalUsage, "balance": remainingBalance, "currency": "USD", }, nil } func (o *OpenRouterModel) CheckConnection(apiConfig *APIConfig) error { _, err := o.Balance(apiConfig) return err }