The Simple Client for OpenRouter with Go

openrouter.go

package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"

	"github.com/go-resty/resty/v2"
)

type ChatCompletion struct {
	ID      string   `json:"id"`
	Error   RespErr  `json:"error"`
	Model   string   `json:"model"`
	Object  string   `json:"object"`
	Created int      `json:"created"`
	Choices []Choice `json:"choices"`
	Usage   Usage    `json:"usage"`
}

type RespErr struct {
	Message string `json:"message"`
	Code    int    `json:"code"`
}

type Choice struct {
	Logprobs     interface{} `json:"logprobs"`
	FinishReason string      `json:"finish_reason"`
	Index        int         `json:"index"`
	Message      Message     `json:"message"`
	Delta        Delta       `json:"delta"` // for streaming
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type Delta struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type Option struct {
	APIKey  string
	BaseURL string
}

type Client struct {
	client *resty.Client
	opt    Option
}

type ChatCompletionParams struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Temperature float64   `json:"temperature"`
}

func NewClient(opt Option) *Client {
	client := resty.New()
	client.SetHeader("Content-Type", "application/json")
	client.SetHeader("Authorization", "Bearer "+opt.APIKey)

	return &Client{
		client: client,
		opt:    opt,
	}
}

func (c *Client) ChatCompletion(params ChatCompletionParams) (string, error) {
	resp, err := c.client.R().
		SetBody(map[string]interface{}{
			"model":       params.Model,
			"messages":    params.Messages,
			"temperature": params.Temperature,
		}).
		Post(c.opt.BaseURL + "/chat/completions")

	if err != nil {
		return "", err
	}

	var chatCompletion ChatCompletion
	err = json.Unmarshal(resp.Body(), &chatCompletion)
	if err != nil {
		return "", err
	}

	if chatCompletion.Error.Message != "" {
		return "", errors.New(chatCompletion.Error.Message)
	}

	return chatCompletion.Choices[0].Message.Content, nil
}

type ChatCompletionStream struct {
	scanner *bufio.Scanner
	err     error
	current string
}

func (c *Client) CreateChatCompletionStream(params ChatCompletionParams) (*ChatCompletionStream, error) {
	resp, err := c.client.R().
		SetBody(map[string]interface{}{
			"model":       params.Model,
			"messages":    params.Messages,
			"stream":      true,
			"temperature": params.Temperature,
		}).
		SetDoNotParseResponse(true).
		Post(c.opt.BaseURL + "/chat/completions")

	if err != nil {
		return nil, err
	}

	return &ChatCompletionStream{
		scanner: bufio.NewScanner(resp.RawResponse.Body),
	}, nil
}

func (s *ChatCompletionStream) Next() bool {
	for s.scanner.Scan() {
		_res := s.scanner.Text()
		if _res == "" || _res == "OPENROUTER PROCESSING" || _res == "data: [DONE]" {
			continue
		}
		_res = strings.TrimPrefix(_res, "data: ")
		_res = strings.TrimSpace(_res)

		var chatCompletion ChatCompletion
		err := json.Unmarshal([]byte(_res), &chatCompletion)
		// ignore json parse error
		if err != nil {
			s.current = ""
			return true
		}

		s.current = chatCompletion.Choices[0].Delta.Content
		return true
	}

	if err := s.scanner.Err(); err != nil {
		s.err = err
	}
	return false
}

func (s *ChatCompletionStream) Current() string {
	return s.current
}

func (s *ChatCompletionStream) Err() error {
	return s.err
}

func main() {
	client := NewClient(Option{
		APIKey:  "Your key",
		BaseURL: "https://openrouter.ai/api/v1",
	})

	stream, err := client.CreateChatCompletionStream(ChatCompletionParams{
		Model: "google/gemini-flash-1.5",
		Messages: []Message{
			{
				Role:    "user",
				Content: "who are you?",
			},
		},
		Temperature: 0.7,
	})

	if err != nil {
		log.Println(err)
		return
	}

	for stream.Next() {
		chunk := stream.Current()
		fmt.Print(chunk)
	}

	if err := stream.Err(); err != nil {
		log.Println("流读取错误:", err)
	}
}