153 lines
4.3 KiB
Go
153 lines
4.3 KiB
Go
package openrouter
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
)
|
|
|
|
type Client struct {
|
|
BaseURL, APIKey, SiteURL, AppName string
|
|
HTTP *http.Client
|
|
}
|
|
|
|
func (c *Client) do(ctx context.Context, method, path string, body any) (*http.Response, error) {
|
|
var reader io.Reader
|
|
if body != nil {
|
|
b, err := json.Marshal(body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
reader = bytes.NewReader(b)
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(c.BaseURL, "/")+path, reader)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
|
if body != nil {
|
|
req.Header.Set("Content-Type", "application/json")
|
|
}
|
|
if c.SiteURL != "" {
|
|
req.Header.Set("HTTP-Referer", c.SiteURL)
|
|
}
|
|
if c.AppName != "" {
|
|
req.Header.Set("X-Title", c.AppName)
|
|
}
|
|
resp, err := c.HTTP.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
defer resp.Body.Close()
|
|
b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
|
return nil, fmt.Errorf("OpenRouter status %d: %s", resp.StatusCode, strings.TrimSpace(string(b)))
|
|
}
|
|
return resp, nil
|
|
}
|
|
|
|
type imageURL struct {
|
|
URL string `json:"url"`
|
|
}
|
|
type contentPart struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text,omitempty"`
|
|
ImageURL *imageURL `json:"image_url,omitempty"`
|
|
}
|
|
type message struct {
|
|
Role string `json:"role"`
|
|
Content any `json:"content"`
|
|
}
|
|
|
|
func (c *Client) Prompt(ctx context.Context, model, description, dataURL, system string) (string, error) {
|
|
body := map[string]any{
|
|
"model": model,
|
|
"messages": []message{
|
|
{Role: "system", Content: system},
|
|
{Role: "user", Content: []contentPart{
|
|
{Type: "text", Text: description},
|
|
{Type: "image_url", ImageURL: &imageURL{URL: dataURL}},
|
|
}},
|
|
},
|
|
"response_format": map[string]any{"type": "json_schema", "json_schema": map[string]any{"name": "video_prompt", "strict": true, "schema": map[string]any{"type": "object", "properties": map[string]any{"video_prompt": map[string]string{"type": "string"}}, "required": []string{"video_prompt"}, "additionalProperties": false}}},
|
|
}
|
|
resp, err := c.do(ctx, http.MethodPost, "/chat/completions", body)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer resp.Body.Close()
|
|
var out struct {
|
|
Choices []struct {
|
|
Message struct {
|
|
Content string `json:"content"`
|
|
} `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
if err = json.NewDecoder(io.LimitReader(resp.Body, 2<<20)).Decode(&out); err != nil {
|
|
return "", err
|
|
}
|
|
if len(out.Choices) == 0 {
|
|
return "", fmt.Errorf("empty model response")
|
|
}
|
|
var result struct {
|
|
VideoPrompt string `json:"video_prompt"`
|
|
}
|
|
if err = json.Unmarshal([]byte(out.Choices[0].Message.Content), &result); err != nil {
|
|
return "", err
|
|
}
|
|
if strings.TrimSpace(result.VideoPrompt) == "" {
|
|
return "", fmt.Errorf("empty video prompt")
|
|
}
|
|
return result.VideoPrompt, nil
|
|
}
|
|
|
|
type VideoRequest struct {
|
|
Model string `json:"model"`
|
|
Prompt string `json:"prompt"`
|
|
Duration int `json:"duration"`
|
|
Resolution string `json:"resolution"`
|
|
AspectRatio string `json:"aspect_ratio"`
|
|
GenerateAudio bool `json:"generate_audio"`
|
|
FrameImages []FrameImage `json:"frame_images"`
|
|
}
|
|
type FrameImage struct {
|
|
Type string `json:"type"`
|
|
ImageURL imageURL `json:"image_url"`
|
|
FrameType string `json:"frame_type"`
|
|
}
|
|
type VideoJob struct {
|
|
ID string `json:"id"`
|
|
Status string `json:"status"`
|
|
Error string `json:"error"`
|
|
UnsignedURLs []string `json:"unsigned_urls"`
|
|
}
|
|
|
|
func NewImageURL(raw string) imageURL { return imageURL{URL: raw} }
|
|
func (c *Client) Submit(ctx context.Context, v VideoRequest) (VideoJob, error) {
|
|
resp, err := c.do(ctx, http.MethodPost, "/videos", v)
|
|
if err != nil {
|
|
return VideoJob{}, err
|
|
}
|
|
defer resp.Body.Close()
|
|
var job VideoJob
|
|
err = json.NewDecoder(resp.Body).Decode(&job)
|
|
return job, err
|
|
}
|
|
func (c *Client) Poll(ctx context.Context, id string) (VideoJob, error) {
|
|
resp, err := c.do(ctx, http.MethodGet, "/videos/"+id, nil)
|
|
if err != nil {
|
|
return VideoJob{}, err
|
|
}
|
|
defer resp.Body.Close()
|
|
var job VideoJob
|
|
err = json.NewDecoder(resp.Body).Decode(&job)
|
|
return job, err
|
|
}
|
|
func (c *Client) Content(ctx context.Context, id string) (*http.Response, error) {
|
|
return c.do(ctx, http.MethodGet, "/videos/"+id+"/content?index=0", nil)
|
|
}
|