x

Programs, configuration and documentation that don't fit anywhere else
Log | Files | Refs | README | LICENSE

openai.go (3160B)

      1 package openai
      2 
      3 import (
      4 	"bytes"
      5 	"encoding/json"
      6 	"fmt"
      7 	"net/http"
      8 	"strings"
      9 	"time"
     10 )
     11 
     12 type Role string
     13 
     14 const (
     15 	RoleSystem    = "system"
     16 	RoleUser      = "user"
     17 	RoleAssistant = "assistant"
     18 )
     19 
     20 type Message struct {
     21 	Role    string `json:"role"`
     22 	Content string `json:"content"`
     23 }
     24 
     25 type Chat struct {
     26 	Messages       []Message `json:"messages"`
     27 	Model          string    `json:"model"`
     28 	ResponseFormat *struct {
     29 		Type string `json:"type"`
     30 	} `json:"response_format,omitempty"`
     31 }
     32 
     33 type Model struct {
     34 	ID               string
     35 	Name             string
     36 	Description      string
     37 	Aliases          []string
     38 	MaxContextLength int
     39 	Deprecation      time.Time
     40 }
     41 
     42 type Client struct {
     43 	*http.Client
     44 	Token   string
     45 	BaseURL string
     46 }
     47 
     48 type apiError struct {
     49 	Message struct {
     50 		Detail []struct {
     51 			Msg string
     52 		}
     53 	}
     54 	Type string
     55 }
     56 
     57 func (e apiError) Error() string {
     58 	messages := make([]string, len(e.Message.Detail))
     59 	for i := range e.Message.Detail {
     60 		messages[i] = e.Message.Detail[i].Msg
     61 	}
     62 	return fmt.Sprintf("%s: %s", e.Type, strings.Join(messages, ", "))
     63 }
     64 
     65 func (c *Client) do(req *http.Request) (*http.Response, error) {
     66 	if c.Client == nil {
     67 		c.Client = http.DefaultClient
     68 	}
     69 	if c.Token != "" {
     70 		req.Header.Set("Authorization", "Bearer "+c.Token)
     71 	}
     72 	req.Header.Set("Accept", "application/json")
     73 	if req.Body != nil {
     74 		req.Header.Set("Content-Type", "application/json")
     75 	}
     76 	return c.Do(req)
     77 }
     78 
     79 type completeResponse struct {
     80 	Choices []struct {
     81 		Message Message
     82 	}
     83 }
     84 
     85 func (c *Client) Complete(chat *Chat) (*Message, error) {
     86 	b, err := json.Marshal(chat)
     87 	if err != nil {
     88 		return nil, fmt.Errorf("encode messages: %w", err)
     89 	}
     90 	u := c.BaseURL + "/v1/chat/completions"
     91 	req, err := http.NewRequest(http.MethodPost, u, bytes.NewReader(b))
     92 	if err != nil {
     93 		return nil, err
     94 	}
     95 	resp, err := c.do(req)
     96 	if err != nil {
     97 		return nil, err
     98 	}
     99 	defer resp.Body.Close()
    100 	if resp.StatusCode == http.StatusUnauthorized {
    101 		return nil, fmt.Errorf("unauthorised")
    102 	} else if resp.StatusCode >= 400 && resp.StatusCode <= 499 {
    103 		var aerr apiError
    104 		if err := json.NewDecoder(resp.Body).Decode(&aerr); err != nil {
    105 			return nil, fmt.Errorf(resp.Status)
    106 		}
    107 		return nil, aerr
    108 	} else if resp.StatusCode >= 500 {
    109 		return nil, fmt.Errorf(resp.Status)
    110 	}
    111 
    112 	var cresp completeResponse
    113 	if err := json.NewDecoder(resp.Body).Decode(&cresp); err != nil {
    114 		return nil, fmt.Errorf("decode response: %w", err)
    115 	}
    116 	if len(cresp.Choices) == 0 {
    117 		return nil, fmt.Errorf("no completions in response")
    118 	}
    119 	return &cresp.Choices[0].Message, nil
    120 }
    121 
    122 func (c *Client) Models() ([]Model, error) {
    123 	u := c.BaseURL + "/v1/models"
    124 	req, err := http.NewRequest(http.MethodGet, u, nil)
    125 	if err != nil {
    126 		return nil, err
    127 	}
    128 	resp, err := c.do(req)
    129 	if err != nil {
    130 		return nil, err
    131 	}
    132 	defer resp.Body.Close()
    133 	if resp.StatusCode != http.StatusOK {
    134 		var aerr apiError
    135 		if err := json.NewDecoder(resp.Body).Decode(&aerr); err != nil {
    136 			return nil, fmt.Errorf(resp.Status)
    137 		}
    138 		return nil, aerr
    139 	}
    140 	v := struct {
    141 		Data []Model
    142 	}{}
    143 	if err := json.NewDecoder(resp.Body).Decode(&v); err != nil {
    144 		return nil, fmt.Errorf("decode response: %w", err)
    145 	}
    146 	return v.Data, nil
    147 }