Add the code

Author: hsluoyz

.gitignore

@@ -19,3 +19,6 @@
 
 # Go workspace file
 go.work
+
+.idea/
+*.iml

chat.go

@@ -0,0 +1,48 @@
+package openrouter
+
+import (
+	"context"
+	"errors"
+	"net/http"
+)
+
+// Chat message role defined by the Sensa API.
+
+type ModelName string
+
+const (
+	ChatMessageRoleUser      = "user"
+	ChatMessageRoleSystem    = "system"
+	ChatMessageRoleAssistant = "assistant"
+)
+
+var (
+	ErrChatCompletionStreamNotSupported = errors.New("streaming is not supported with this method, please use CreateChatCompletionStream") //nolint:lll
+	ErrCompletionUnsupportedModel       = errors.New("this model is not supported with this method")                                       //nolint:lll
+)
+
+// CreateChatCompletion — API call to Create a completion for the chat message.
+func (c *Client) CreateChatCompletion(
+	ctx context.Context,
+	request *ChatCompletionRequest,
+) (response *ChatCompletionResponse, err error) {
+	if request.Stream {
+		err = ErrChatCompletionStreamNotSupported
+		return
+	}
+
+	urlSuffix := "/chat/completions"
+	request.Model = wrapperModels[request.Model]
+	if !checkSupportsModel(request.Model) {
+		err = ErrCompletionUnsupportedModel
+		return
+	}
+
+	req, err := c.requestBuilder.Build(ctx, http.MethodPost, c.fullURL(urlSuffix), request)
+	if err != nil {
+		return
+	}
+
+	err = c.sendRequest(req, &response)
+	return
+}

chat_stream.go

@@ -0,0 +1,51 @@
+package openrouter
+
+import (
+	"bufio"
+	"context"
+	utils "github.com/casibase/go-openrouter/internal"
+)
+
+type ChatCompletionStream struct {
+	streamReader
+}
+
+// CreateChatCompletionStream — API call to create a chat completion w/ streaming
+// support. It sets whether to stream back partial progress. If set, tokens will be
+// sent as data-only server-sent events as they become available, with the
+// stream terminated by a data: [DONE] message.
+func (c *Client) CreateChatCompletionStream(
+	ctx context.Context,
+	request *ChatCompletionRequest,
+) (stream *ChatCompletionStream, err error) {
+	urlSuffix := "/chat/completions"
+	request.Model = wrapperModels[request.Model]
+	if !checkSupportsModel(request.Model) {
+		err = ErrCompletionUnsupportedModel
+		return
+	}
+	request.Stream = true
+	req, err := c.newStreamRequest(ctx, "POST", urlSuffix, request)
+	if err != nil {
+		return
+	}
+
+	resp, err := c.config.HTTPClient.Do(req) //nolint:bodyclose // body is closed in stream.Close()
+	if err != nil {
+		return
+	}
+	if isFailureStatusCode(resp) {
+		return nil, c.handleErrorResp(resp)
+	}
+
+	stream = &ChatCompletionStream{
+		streamReader: streamReader{
+			emptyMessagesLimit: c.config.EmptyMessagesLimit,
+			reader:             bufio.NewReader(resp.Body),
+			response:           resp,
+			errAccumulator:     utils.NewErrorAccumulator(),
+			unmarshaler:        &utils.JSONUnmarshaler{},
+		},
+	}
+	return
+}

chat_test.go

@@ -0,0 +1,48 @@
+package openrouter
+
+import (
+	"context"
+	"testing"
+)
+
+func TestClient_CreateChatCompletion(t *testing.T) {
+	client, _ := NewClient("", "", "")
+
+	req := &ChatCompletionRequest{
+		Model: "claude-2",
+		Messages: []ChatCompletionMessage{
+			{
+				Role:    ChatMessageRoleSystem,
+				Content: "You are a helpful assistant.",
+			},
+			{
+				Role:    ChatMessageRoleUser,
+				Content: "what is today",
+			},
+		},
+		Stream:      false,
+		Temperature: nil,
+		TopP:        nil,
+	}
+
+	t.Log(client.CreateChatCompletion(context.Background(), req))
+	//
+	//r, err := client.CreateChatCompletionStream(context.Background(), req)
+	//if err != nil {
+	//	t.Error(err)
+	//}
+	//t.Log(r)
+	//for {
+	//	fmt.Println(1)
+	//	r, err := r.Recv()
+	//	if err != nil {
+	//		fmt.Println(err.Error())
+	//		if errors.Is(err, io.EOF) {
+	//			fmt.Println(1)
+	//			break
+	//		}
+	//		t.Error(err)
+	//	}
+	//	t.Logf("%#v", r.Choices)
+	//}
+}

client.go

@@ -0,0 +1,130 @@
+package openrouter
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	utils "github.com/casibase/go-openrouter/internal"
+	"io"
+	"net/http"
+)
+
+type Client struct {
+	config ClientConfig
+
+	requestBuilder utils.RequestBuilder
+}
+
+func NewClient(auth, xTitle, httpReferer string) (*Client, error) {
+	config, err := DefaultConfig(auth, xTitle, httpReferer)
+	if err != nil {
+		return nil, err
+	}
+	return NewClientWithConfig(config), nil
+}
+
+func NewClientWithConfig(config ClientConfig) *Client {
+	return &Client{
+		config:         config,
+		requestBuilder: utils.NewRequestBuilder(),
+	}
+}
+
+func (c *Client) sendRequest(req *http.Request, v any) error {
+	req.Header.Set("Accept", "application/json; charset=utf-8")
+
+	// Check whether Content-Type is already set, Upload Files API requires
+	// Content-Type == multipart/form-data
+	contentType := req.Header.Get("Content-Type")
+	if contentType == "" {
+		req.Header.Set("Content-Type", "application/json; charset=utf-8")
+	}
+
+	c.setCommonHeaders(req)
+
+	res, err := c.config.HTTPClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer res.Body.Close()
+
+	if isFailureStatusCode(res) {
+		return c.handleErrorResp(res)
+	}
+
+	return decodeResponse(res.Body, v)
+}
+
+func (c *Client) setCommonHeaders(req *http.Request) {
+	req.Header.Set("HTTP-Referer", c.config.HttpReferer)
+	req.Header.Set("X-Title", c.config.XTitle)
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.config.authToken))
+}
+
+func isFailureStatusCode(resp *http.Response) bool {
+	return resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusBadRequest
+}
+
+func decodeResponse(body io.Reader, v any) error {
+	if v == nil {
+		return nil
+	}
+
+	if result, ok := v.(*string); ok {
+		return decodeString(body, result)
+	}
+	return json.NewDecoder(body).Decode(v)
+}
+
+func decodeString(body io.Reader, output *string) error {
+	b, err := io.ReadAll(body)
+	if err != nil {
+		return err
+	}
+	*output = string(b)
+	return nil
+}
+
+// fullURL returns full URL for request.
+// args[0] is model name, if API type is Azure, model name is required to get deployment name.
+func (c *Client) fullURL(suffix string) string {
+	return fmt.Sprintf("%s%s", c.config.BaseURL, suffix)
+}
+
+func (c *Client) newStreamRequest(
+	ctx context.Context,
+	method string,
+	urlSuffix string,
+	body any) (*http.Request, error) {
+	req, err := c.requestBuilder.Build(ctx, method, c.fullURL(urlSuffix), body)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "text/event-stream")
+	req.Header.Set("Cache-Control", "no-cache")
+	req.Header.Set("Connection", "keep-alive")
+
+	c.setCommonHeaders(req)
+	return req, nil
+}
+
+func (c *Client) handleErrorResp(resp *http.Response) error {
+	var errRes ErrorResponse
+
+	err := json.NewDecoder(resp.Body).Decode(&errRes)
+	if err != nil || errRes.Error == nil {
+		reqErr := &RequestError{
+			HTTPStatusCode: resp.StatusCode,
+			Err:            err,
+		}
+		if errRes.Error != nil {
+			reqErr.Err = errRes.Error
+		}
+		return reqErr
+	}
+
+	errRes.Error.HTTPStatusCode = resp.StatusCode
+	return errRes.Error
+}

config.go

@@ -0,0 +1,37 @@
+package openrouter
+
+import (
+	"net/http"
+)
+
+const (
+	routerAPIURLv1                 = "https://openrouter.ai/api/v1"
+	defaultEmptyMessagesLimit uint = 300
+)
+
+// ClientConfig is a configuration of a client.
+// XTitle、HttpRefer your own site url
+type ClientConfig struct {
+	authToken          string
+	XTitle             string
+	HttpReferer        string
+	BaseURL            string
+	HTTPClient         *http.Client
+	EmptyMessagesLimit uint
+}
+
+func DefaultConfig(auth, xTitle, httpReferer string) (ClientConfig, error) {
+	return ClientConfig{
+		authToken:          auth,
+		HTTPClient:         &http.Client{},
+		XTitle:             xTitle,
+		HttpReferer:        httpReferer,
+		BaseURL:            routerAPIURLv1,
+		EmptyMessagesLimit: defaultEmptyMessagesLimit,
+	}, nil
+}
+
+func (c ClientConfig) WithHttpClientConfig(client *http.Client) ClientConfig {
+	c.HTTPClient = client
+	return c
+}

error.go

@@ -0,0 +1,68 @@
+package openrouter
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+// APIError provides error information returned by the OpenAI API.
+type APIError struct {
+	Code           any    `json:"code,omitempty"`
+	Message        string `json:"message"`
+	Details        any    `json:"details"`
+	HTTPStatusCode int    `json:"-"`
+}
+
+// RequestError provides informations about generic request errors.
+type RequestError struct {
+	HTTPStatusCode int
+	Err            error
+}
+
+type ErrorResponse struct {
+	Error *APIError `json:"error,omitempty"`
+}
+
+func (e *APIError) Error() string {
+	if e.HTTPStatusCode > 0 {
+		return fmt.Sprintf("error, status code: %d, message: %s", e.HTTPStatusCode, e.Message)
+	}
+
+	return e.Message
+}
+
+func (e *APIError) UnmarshalJSON(data []byte) (err error) {
+	var rawMap map[string]json.RawMessage
+	err = json.Unmarshal(data, &rawMap)
+	if err != nil {
+		return
+	}
+
+	err = json.Unmarshal(rawMap["message"], &e.Message)
+	if err != nil {
+		return
+	}
+
+	if _, ok := rawMap["code"]; !ok {
+		return nil
+	}
+
+	// if the api returned a number, we need to force an integer
+	// since the json package defaults to float64
+	var intCode int
+	err = json.Unmarshal(rawMap["code"], &intCode)
+	if err == nil {
+		e.Code = intCode
+		return nil
+	}
+
+	return json.Unmarshal(rawMap["code"], &e.Code)
+}
+
+func (e *RequestError) Error() string {
+	return fmt.Sprintf("error, status code: %d, message: %s", e.HTTPStatusCode, e.Err)
+}
+
+func (e *RequestError) Unwrap() error {
+	return e.Err
+}

go.mod

@@ -0,0 +1,3 @@
+module github.com/casibase/go-openrouter
+
+go 1.19