160 lines
3.4 KiB
Go
160 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
|
|
tea "charm.land/bubbletea/v2"
|
|
)
|
|
|
|
type chatMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
type chatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []chatMessage `json:"messages"`
|
|
Stream bool `json:"stream"`
|
|
}
|
|
|
|
type sseChunk struct {
|
|
ID string `json:"id"`
|
|
Object string `json:"object"`
|
|
Created int64 `json:"created"`
|
|
Model string `json:"model"`
|
|
Choices []struct {
|
|
Index int `json:"index"`
|
|
Delta chatMessage `json:"delta"`
|
|
FinishReason *string `json:"finish_reason"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
type tagModel struct {
|
|
Name string `json:"name"`
|
|
}
|
|
|
|
type listTagsResp struct {
|
|
Models []tagModel `json:"models"`
|
|
}
|
|
|
|
type streamChunkMsg string
|
|
type streamDoneMsg struct{}
|
|
type streamErrMsg struct{ err error }
|
|
|
|
func fetchModels(baseURL string) ([]string, error) {
|
|
resp, err := http.Get(strings.TrimRight(baseURL, "/") + "/api/tags")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to fetch models: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
var r listTagsResp
|
|
if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
|
|
return nil, fmt.Errorf("failed to decode models: %w", err)
|
|
}
|
|
|
|
models := make([]string, len(r.Models))
|
|
for i, m := range r.Models {
|
|
models[i] = m.Name
|
|
}
|
|
return models, nil
|
|
}
|
|
|
|
func startStream(ctx context.Context, msgs []chatMessage, model, baseURL string) (<-chan tea.Msg, tea.Cmd) {
|
|
ch := make(chan tea.Msg, 16)
|
|
go runStream(ctx, ch, msgs, model, baseURL)
|
|
return ch, readNext(ch)
|
|
}
|
|
|
|
func readNext(ch <-chan tea.Msg) tea.Cmd {
|
|
return func() tea.Msg {
|
|
msg, ok := <-ch
|
|
if !ok {
|
|
return streamDoneMsg{}
|
|
}
|
|
return msg
|
|
}
|
|
}
|
|
|
|
func runStream(ctx context.Context, ch chan<- tea.Msg, msgs []chatMessage, model, baseURL string) {
|
|
defer close(ch)
|
|
|
|
reqBody := chatRequest{
|
|
Model: model,
|
|
Messages: msgs,
|
|
Stream: true,
|
|
}
|
|
|
|
body, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
ch <- streamErrMsg{err: fmt.Errorf("marshal request: %w", err)}
|
|
return
|
|
}
|
|
|
|
url := strings.TrimRight(baseURL, "/") + "/v1/chat/completions"
|
|
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
|
|
if err != nil {
|
|
ch <- streamErrMsg{err: fmt.Errorf("create request: %w", err)}
|
|
return
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "text/event-stream")
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return
|
|
}
|
|
ch <- streamErrMsg{err: fmt.Errorf("http request: %w", err)}
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
errBody, _ := io.ReadAll(resp.Body)
|
|
ch <- streamErrMsg{err: fmt.Errorf("api %d: %s", resp.StatusCode, string(errBody))}
|
|
return
|
|
}
|
|
|
|
scanner := bufio.NewScanner(resp.Body)
|
|
for scanner.Scan() {
|
|
if ctx.Err() != nil {
|
|
return
|
|
}
|
|
|
|
line := strings.TrimSpace(scanner.Text())
|
|
|
|
if line == "" || !strings.HasPrefix(line, "data:") {
|
|
continue
|
|
}
|
|
|
|
data := strings.TrimSpace(line[5:])
|
|
if data == "[DONE]" {
|
|
break
|
|
}
|
|
|
|
var chunk sseChunk
|
|
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
|
continue
|
|
}
|
|
|
|
if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" {
|
|
ch <- streamChunkMsg(chunk.Choices[0].Delta.Content)
|
|
}
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
if ctx.Err() != nil {
|
|
return
|
|
}
|
|
ch <- streamErrMsg{err: fmt.Errorf("scanner: %w", err)}
|
|
}
|
|
}
|