99 lines
2.6 KiB
Go
99 lines
2.6 KiB
Go
package llm
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"ai-operator/internal/config"
|
|
)
|
|
|
|
type OpenAIStreaming struct {
|
|
cfg config.Config
|
|
client *http.Client
|
|
}
|
|
|
|
func NewOpenAIStreaming(cfg config.Config) *OpenAIStreaming {
|
|
return &OpenAIStreaming{cfg: cfg, client: &http.Client{Timeout: cfg.LLM.Timeout}}
|
|
}
|
|
|
|
func (p *OpenAIStreaming) StreamGenerate(ctx context.Context, req GenerateRequest) (<-chan Event, error) {
|
|
if p.cfg.OpenAI.APIKey == "" {
|
|
return nil, errors.New("OPENAI_API_KEY is required for streaming LLM")
|
|
}
|
|
out := make(chan Event, 64)
|
|
go p.run(ctx, req, out)
|
|
return out, nil
|
|
}
|
|
|
|
func (p *OpenAIStreaming) run(ctx context.Context, req GenerateRequest, out chan<- Event) {
|
|
defer close(out)
|
|
payload := map[string]any{
|
|
"model": p.cfg.LLM.Model,
|
|
"temperature": p.cfg.LLM.Temperature,
|
|
"max_tokens": p.cfg.LLM.MaxOutputTokens,
|
|
"stream": true,
|
|
}
|
|
msgs := []map[string]string{{"role": "system", "content": req.SystemPrompt}}
|
|
for _, m := range req.Messages {
|
|
if strings.TrimSpace(m.Content) != "" {
|
|
msgs = append(msgs, map[string]string{"role": m.Role, "content": m.Content})
|
|
}
|
|
}
|
|
payload["messages"] = msgs
|
|
body, _ := json.Marshal(payload)
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://api.openai.com/v1/chat/completions", bytes.NewReader(body))
|
|
if err != nil {
|
|
out <- Event{Type: EventError, Error: err.Error()}
|
|
return
|
|
}
|
|
httpReq.Header.Set("Authorization", "Bearer "+p.cfg.OpenAI.APIKey)
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
resp, err := p.client.Do(httpReq)
|
|
if err != nil {
|
|
out <- Event{Type: EventError, Error: err.Error()}
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode > 299 {
|
|
out <- Event{Type: EventError, Error: "openai streaming llm returned non-2xx"}
|
|
return
|
|
}
|
|
sc := bufio.NewScanner(resp.Body)
|
|
for sc.Scan() {
|
|
line := strings.TrimSpace(sc.Text())
|
|
if !strings.HasPrefix(line, "data:") {
|
|
continue
|
|
}
|
|
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
|
if data == "[DONE]" {
|
|
out <- Event{Type: EventDone}
|
|
return
|
|
}
|
|
var chunk struct {
|
|
Choices []struct {
|
|
Delta struct {
|
|
Content string `json:"content"`
|
|
} `json:"delta"`
|
|
} `json:"choices"`
|
|
}
|
|
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
|
continue
|
|
}
|
|
for _, c := range chunk.Choices {
|
|
if c.Delta.Content != "" {
|
|
out <- Event{Type: EventTextDelta, Text: c.Delta.Content}
|
|
}
|
|
}
|
|
}
|
|
if err := sc.Err(); err != nil {
|
|
out <- Event{Type: EventError, Error: err.Error()}
|
|
return
|
|
}
|
|
out <- Event{Type: EventDone}
|
|
}
|