Files
ragflow/internal/entity/models/openrouter.go
Dexterity 04aa8d04e8 fix(go-models): raise SSE scanner buffer so large stream chunks are not dropped (#15382)
### Summary

Closes #15381 

Every provider in `internal/entity/models/` reads its streaming response
with `bufio.NewScanner(resp.Body)` and iterates over `scanner.Scan()`.
The default `bufio.Scanner` maximum token size is 64KB, so when an
upstream sends a single SSE `data:` line larger than 64KB (long content
deltas, large tool or function call argument blobs, bundled
`reasoning_content`, or providers that emit a whole message in one
event) `scanner.Scan()` returns `false` and `scanner.Err()` returns
`bufio.ErrTooLong`. Streaming chat then ends with an error partway
through the response.

This change adds `scanner.Buffer(make([]byte, 64*1024), 1024*1024)`
immediately after every SSE scanner that was still bare, raising the cap
to 1MB. 1MB is the value already used for streaming chat in `openai.go`,
`modelscope.go`, `groq.go`, `mistral.go`, `xai.go` and the other already
patched providers (the 8MB cap in the repo is reserved for TTS and
embedding paths), so this simply converges the remaining providers onto
the established pattern. Nothing else changes: line parsing, `data:`
prefix handling, `[DONE]` detection, JSON unmarshalling, error handling,
and the existing `scanner.Err()` checks all stay the same.

Providers covered (23 scanners across 22 files): 302ai, aliyun,
baichuan, baidu, cohere, deepinfra, deepseek, gitee, huggingface,
lmstudio, minimax (the chat scanner, whose TTS scanner was already
bumped), moonshot, nvidia, ollama, openrouter, orcarouter, paddleocr,
siliconflow, tokenhub, vllm, volcengine, xunfei, zhipu-ai. `jiekouai.go`
is excluded because it is covered by the in flight #15337.

A table driven regression test (`sse_scanner_buffer_test.go`) streams a
single 128KB `data:` content delta followed by `data: [DONE]` through an
`httptest` server and asserts that `ChatStreamlyWithSender` delivers the
full content with no error across a representative subset of providers.
Without the buffer fix the test fails with `bufio.Scanner: token too
long`.

This PR also removes three duplicate declarations of the package level
`roundTripperFunc` test helper that several recently merged provider PRs
each added independently, which had left the `internal/entity/models`
test package unable to compile. The helper now lives in a single place
and is shared.

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
2026-05-29 19:34:00 +08:00

837 lines
22 KiB
Go

package models
import (
"bufio"
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"ragflow/internal/common"
"strings"
"time"
)
// OpenRouterModel implements ModelDriver for OpenRouter AI
type OpenRouterModel struct {
BaseURL map[string]string
URLSuffix URLSuffix
httpClient *http.Client
}
// 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
}
if chatModelConfig.Effort != nil {
reqBody["reasoning"] = map[string]interface{}{
"effort": chatModelConfig.Effort,
}
}
}
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)
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
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
}
type openRouterTranscriptionResponse struct {
Text string `json:"text"`
}
func openRouterAudioFormat(file string, asrConfig *ASRConfig) string {
if asrConfig != nil && asrConfig.Params != nil {
if format, ok := asrConfig.Params["format"]; ok && format != nil {
if value := strings.TrimPrefix(fmt.Sprint(format), "."); value != "" {
return value
}
}
}
ext := strings.TrimPrefix(strings.ToLower(filepath.Ext(file)), ".")
if ext == "" {
return "wav"
}
return ext
}
// TranscribeAudio transcribe audio
func (o *OpenRouterModel) TranscribeAudio(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig) (*ASRResponse, error) {
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
return nil, fmt.Errorf("OpenRouter API key is missing")
}
if modelName == nil || *modelName == "" {
return nil, fmt.Errorf("model name is required")
}
if file == nil || *file == "" {
return nil, fmt.Errorf("file is missing")
}
if o.URLSuffix.ASR == "" {
return nil, fmt.Errorf("OpenRouter ASR url suffix is missing")
}
audio, err := os.ReadFile(*file)
if err != nil {
return nil, fmt.Errorf("failed to read audio file: %w", err)
}
region := "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
}
reqBody := map[string]interface{}{
"model": *modelName,
"input_audio": map[string]interface{}{
"data": base64.StdEncoding.EncodeToString(audio),
"format": openRouterAudioFormat(*file, asrConfig),
},
}
if asrConfig != nil && asrConfig.Params != nil {
for key, value := range asrConfig.Params {
switch key {
case "format", "model", "input_audio":
continue
}
reqBody[key] = value
}
}
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.ASR)
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 body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("OpenRouter ASR API error: %s, body: %s", resp.Status, string(body))
}
var result openRouterTranscriptionResponse
if err = json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse transcription response: %w", err)
}
return &ASRResponse{Text: result.Text}, nil
}
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 text to audio
func (o *OpenRouterModel) AudioSpeech(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig) (*TTSResponse, error) {
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
return nil, fmt.Errorf("OpenRouter API key is missing")
}
if audioContent == nil || *audioContent == "" {
return nil, fmt.Errorf("text content is empty")
}
var region = "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
}
url := fmt.Sprintf("%s/%s", o.BaseURL[region], o.URLSuffix.TTS)
// OpenRouter:response Audio bytes stream
reqBody := map[string]interface{}{
"model": modelName,
"input": audioContent,
}
if ttsConfig != nil && ttsConfig.Params != nil {
for key, value := range ttsConfig.Params {
reqBody[key] = value
}
}
if ttsConfig != nil && ttsConfig.Format != "" {
reqBody["response_format"] = ttsConfig.Format
}
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")
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 body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("OpenRouter API error: %s, body: %s", resp.Status, string(body))
}
return &TTSResponse{Audio: body}, nil
}
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) (*OCRFileResponse, error) {
return nil, fmt.Errorf("%s, no such method", m.Name())
}
// ParseFile parse file
func (z *OpenRouterModel) ParseFile(modelName *string, content []byte, url *string, apiConfig *APIConfig, parseFileConfig *ParseFileConfig) (*ParseFileResponse, error) {
return nil, fmt.Errorf("%s, no such method", z.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", 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
}
func (z *OpenRouterModel) ListTasks(apiConfig *APIConfig) ([]ListTaskStatus, error) {
return nil, fmt.Errorf("%s, no such method", z.Name())
}
func (z *OpenRouterModel) ShowTask(taskID string, apiConfig *APIConfig) (*TaskResponse, error) {
return nil, fmt.Errorf("%s, no such method", z.Name())
}