141 lines
3.0 KiB
Go
141 lines
3.0 KiB
Go
package openapi
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
)
|
|
|
|
type Client struct {
|
|
BaseURL string
|
|
APIKey string
|
|
HTTPClient *http.Client
|
|
}
|
|
|
|
func NewClient(baseURL, apiKey string) *Client {
|
|
if !strings.HasSuffix(baseURL, "/") {
|
|
baseURL += "/"
|
|
}
|
|
return &Client{
|
|
BaseURL: baseURL,
|
|
APIKey: apiKey,
|
|
HTTPClient: &http.Client{},
|
|
}
|
|
}
|
|
|
|
func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) {
|
|
req.Stream = false
|
|
data, err := json.Marshal(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal request: %w", err)
|
|
}
|
|
|
|
httpReq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"chat/completions", bytes.NewReader(data))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create request: %w", err)
|
|
}
|
|
|
|
c.setHeaders(httpReq)
|
|
|
|
resp, err := c.HTTPClient.Do(httpReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("do request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
return nil, fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
var result ChatCompletionResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
|
return nil, fmt.Errorf("decode response: %w", err)
|
|
}
|
|
|
|
return &result, nil
|
|
}
|
|
|
|
func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (<-chan ChatCompletionChunk, <-chan error) {
|
|
req.Stream = true
|
|
chunks := make(chan ChatCompletionChunk)
|
|
errs := make(chan error, 1)
|
|
|
|
go func() {
|
|
defer close(chunks)
|
|
defer close(errs)
|
|
|
|
data, err := json.Marshal(req)
|
|
if err != nil {
|
|
errs <- fmt.Errorf("marshal request: %w", err)
|
|
return
|
|
}
|
|
|
|
httpReq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"chat/completions", bytes.NewReader(data))
|
|
if err != nil {
|
|
errs <- fmt.Errorf("create request: %w", err)
|
|
return
|
|
}
|
|
|
|
c.setHeaders(httpReq)
|
|
|
|
resp, err := c.HTTPClient.Do(httpReq)
|
|
if err != nil {
|
|
errs <- fmt.Errorf("do request: %w", err)
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
errs <- fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(body))
|
|
return
|
|
}
|
|
|
|
scanner := bufio.NewScanner(resp.Body)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if line == "" {
|
|
continue
|
|
}
|
|
if !strings.HasPrefix(line, "data: ") {
|
|
continue
|
|
}
|
|
data := strings.TrimPrefix(line, "data: ")
|
|
if data == "[DONE]" {
|
|
break
|
|
}
|
|
|
|
var chunk ChatCompletionChunk
|
|
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
|
errs <- fmt.Errorf("unmarshal chunk: %w", err)
|
|
return
|
|
}
|
|
select {
|
|
case chunks <- chunk:
|
|
case <-ctx.Done():
|
|
errs <- ctx.Err()
|
|
return
|
|
}
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
errs <- fmt.Errorf("scanner error: %w", err)
|
|
}
|
|
}()
|
|
|
|
return chunks, errs
|
|
}
|
|
|
|
func (c *Client) setHeaders(req *http.Request) {
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if c.APIKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
|
}
|
|
}
|