svs-services-server/completion/openrouter/openrouter.go

97 lines
2.3 KiB
Go

package openrouter
import (
"bytes"
"context"
"encoding/json"
"fmt"
"log"
"net/http"
"time"
)
type OpenRouterProvider struct {
// Endpoint string
Token string
Model string // "openai/gpt-oss-20b:free"
}
type Message struct {
Role string `json:"role"` // "system" | "user" | "assistant"
Content string `json:"content"`
}
type ChatCompletionRequest struct {
Model string `json:"model"`
Messages []Message `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
MaxTokens *int `json:"max_tokens,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Stop json.RawMessage `json:"stop,omitempty"` // string or []string; keep flexible
Provider struct {
Sort string `json:"sort,omitempty"`
} `json:"provider,omitempty"`
}
type ChatCompletionResponse struct {
ID string `json:"id"`
Choices []struct {
Index int `json:"index"`
Message struct {
Role string `json:"role"`
Content string `json:"content"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
}
func (p OpenRouterProvider) Complete(ctx context.Context, prompt string) (string, error) {
req := ChatCompletionRequest{
Model: p.Model,
Messages: []Message{
{
Role: "user",
Content: prompt,
},
},
Provider: struct {
Sort string `json:"sort,omitempty"`
}{
Sort: "throughput",
},
}
httpClient := http.Client{Timeout: 10 * time.Second}
body, err := json.Marshal(req)
if err != nil {
return "", err
}
httpReq, err := http.NewRequestWithContext(ctx, "POST", "https://openrouter.ai/api/v1/chat/completions", bytes.NewReader(body))
if err != nil {
return "", err
}
httpReq.Header.Set("Authorization", fmt.Sprintf("Bearer %v", p.Token))
httpReq.Header.Set("Content-Type", "application/json")
resp, err := httpClient.Do(httpReq)
if err != nil {
return "", err
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
// You may want to decode OpenRouter's error JSON here for better messages
return "", fmt.Errorf("openrouter status %d", resp.StatusCode)
}
var out ChatCompletionResponse
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
log.Println(err)
log.Println(out)
return "", err
}
return out.Choices[0].Message.Content, nil
}