mirror of
https://github.com/coalaura/whiskr.git
synced 2025-09-09 01:09:54 +00:00
98 lines
2.9 KiB
Go
98 lines
2.9 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
|
|
"github.com/revrost/go-openrouter"
|
|
)
|
|
|
|
type Generation struct {
|
|
ID string `json:"id"`
|
|
TotalCost float64 `json:"total_cost"`
|
|
CreatedAt string `json:"created_at"`
|
|
Model string `json:"model"`
|
|
Origin string `json:"origin"`
|
|
Usage float64 `json:"usage"`
|
|
IsBYOK bool `json:"is_byok"`
|
|
UpstreamID *string `json:"upstream_id"`
|
|
CacheDiscount *float64 `json:"cache_discount"`
|
|
UpstreamInferenceCost *float64 `json:"upstream_inference_cost"`
|
|
AppID *int `json:"app_id"`
|
|
Streamed *bool `json:"streamed"`
|
|
Cancelled *bool `json:"cancelled"`
|
|
ProviderName *string `json:"provider_name"`
|
|
Latency *int `json:"latency"`
|
|
ModerationLatency *int `json:"moderation_latency"`
|
|
GenerationTime *int `json:"generation_time"`
|
|
FinishReason *string `json:"finish_reason"`
|
|
NativeFinishReason *string `json:"native_finish_reason"`
|
|
TokensPrompt *int `json:"tokens_prompt"`
|
|
TokensCompletion *int `json:"tokens_completion"`
|
|
NativeTokensPrompt *int `json:"native_tokens_prompt"`
|
|
NativeTokensCompletion *int `json:"native_tokens_completion"`
|
|
NativeTokensReasoning *int `json:"native_tokens_reasoning"`
|
|
NumMediaPrompt *int `json:"num_media_prompt"`
|
|
NumMediaCompletion *int `json:"num_media_completion"`
|
|
NumSearchResults *int `json:"num_search_results"`
|
|
}
|
|
|
|
func init() {
|
|
openrouter.DisableLogs()
|
|
}
|
|
|
|
func OpenRouterClient() *openrouter.Client {
|
|
return openrouter.NewClient(OpenRouterToken)
|
|
}
|
|
|
|
func OpenRouterStartStream(ctx context.Context, request openrouter.ChatCompletionRequest) (*openrouter.ChatCompletionStream, error) {
|
|
client := OpenRouterClient()
|
|
|
|
stream, err := client.CreateChatCompletionStream(ctx, request)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return stream, nil
|
|
}
|
|
|
|
func OpenRouterRun(ctx context.Context, request openrouter.ChatCompletionRequest) (openrouter.ChatCompletionResponse, error) {
|
|
client := OpenRouterClient()
|
|
|
|
return client.CreateChatCompletion(ctx, request)
|
|
}
|
|
|
|
func OpenRouterGetGeneration(ctx context.Context, id string) (*Generation, error) {
|
|
req, err := http.NewRequest("GET", fmt.Sprintf("https://openrouter.ai/api/v1/generation?id=%s", id), nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", OpenRouterToken))
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, errors.New(resp.Status)
|
|
}
|
|
|
|
var response struct {
|
|
Data Generation `json:"data"`
|
|
}
|
|
|
|
err = json.NewDecoder(resp.Body).Decode(&response)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &response.Data, nil
|
|
}
|