152 lines
3.9 KiB
Go
152 lines
3.9 KiB
Go
package ai
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"strings"
|
||
)
|
||
|
||
// CompleteStream streams chat completion deltas via onDelta; returns full text.
|
||
// onDelta may be called many times; return error to abort.
|
||
func (p *OpenAICompatible) CompleteStream(ctx context.Context, apiKey, model, prompt string, onDelta func(chunk string) error) (string, error) {
|
||
if strings.TrimSpace(apiKey) == "" {
|
||
return "", fmt.Errorf("missing API key for %s", p.ID)
|
||
}
|
||
if model == "" {
|
||
model = "default"
|
||
}
|
||
sys := SystemPromptForLanguage(ResponseLanguageFrom(ctx))
|
||
maxTokens := MaxTokensFrom(ctx)
|
||
if maxTokens <= 0 {
|
||
maxTokens = 2048
|
||
if p.ID == ProviderOpenCodeGo {
|
||
// 預設 stream 預算；靈感會再透過 WithMaxTokens 壓更小
|
||
maxTokens = 1536
|
||
}
|
||
}
|
||
temp := 0.9
|
||
if t := TemperatureFrom(ctx); t > 0 {
|
||
temp = t
|
||
}
|
||
|
||
body := map[string]any{
|
||
"model": model,
|
||
"max_tokens": maxTokens,
|
||
"temperature": temp,
|
||
"stream": true,
|
||
"messages": []map[string]string{
|
||
{"role": "system", "content": sys},
|
||
{"role": "user", "content": prompt},
|
||
},
|
||
}
|
||
raw, err := json.Marshal(body)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.BaseURL+"/chat/completions", bytes.NewReader(raw))
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("Accept", "text/event-stream")
|
||
req.Header.Set("User-Agent", "HarborDesk/1.0 (+threads-tool; OpenAI-compatible)")
|
||
|
||
res, err := p.Client.Do(req)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
defer res.Body.Close()
|
||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||
data, _ := io.ReadAll(io.LimitReader(res.Body, 64<<10))
|
||
return "", fmt.Errorf("%s stream failed: HTTP %d %s", p.ID, res.StatusCode, truncateRunes(string(data), 240))
|
||
}
|
||
|
||
var full strings.Builder
|
||
sc := bufio.NewScanner(res.Body)
|
||
// SSE 行可能較長
|
||
buf := make([]byte, 0, 64*1024)
|
||
sc.Buffer(buf, 1024*1024)
|
||
|
||
for sc.Scan() {
|
||
line := sc.Text()
|
||
if line == "" {
|
||
continue
|
||
}
|
||
if !strings.HasPrefix(line, "data:") {
|
||
continue
|
||
}
|
||
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||
if payload == "[DONE]" {
|
||
break
|
||
}
|
||
chunk, ok := extractStreamDelta(payload)
|
||
if !ok || chunk == "" {
|
||
continue
|
||
}
|
||
full.WriteString(chunk)
|
||
if onDelta != nil {
|
||
if err := onDelta(chunk); err != nil {
|
||
return full.String(), err
|
||
}
|
||
}
|
||
}
|
||
if err := sc.Err(); err != nil && ctx.Err() == nil {
|
||
return full.String(), err
|
||
}
|
||
text := strings.TrimSpace(full.String())
|
||
if text == "" {
|
||
// 部分供應商 stream 只吐 reasoning、content 在非 stream 才有；
|
||
// 或 max_tokens 太小 content=null。降級 Complete（內含 length 再試）。
|
||
return p.Complete(ctx, apiKey, model, prompt)
|
||
}
|
||
return text, nil
|
||
}
|
||
|
||
func extractStreamDelta(payload string) (string, bool) {
|
||
var obj struct {
|
||
Choices []struct {
|
||
Delta struct {
|
||
Content json.RawMessage `json:"content"`
|
||
ReasoningContent string `json:"reasoning_content"`
|
||
} `json:"delta"`
|
||
// 少數 gateway 用 message
|
||
Message struct {
|
||
Content json.RawMessage `json:"content"`
|
||
} `json:"message"`
|
||
} `json:"choices"`
|
||
}
|
||
if err := json.Unmarshal([]byte(payload), &obj); err != nil {
|
||
return "", false
|
||
}
|
||
if len(obj.Choices) == 0 {
|
||
return "", false
|
||
}
|
||
d := obj.Choices[0].Delta
|
||
if t := decodeMessageContent(d.Content); t != "" {
|
||
return t, true
|
||
}
|
||
// 不把 reasoning 當正文 stream 出去（避免滿屏思考）
|
||
if t := decodeMessageContent(obj.Choices[0].Message.Content); t != "" {
|
||
return t, true
|
||
}
|
||
return "", false
|
||
}
|
||
|
||
// CompleteStream on FakeClient — 一次吐出（測試）
|
||
func (f *FakeClient) CompleteStream(ctx context.Context, apiKey, model, prompt string, onDelta func(string) error) (string, error) {
|
||
text, err := f.Complete(ctx, apiKey, model, prompt)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
if onDelta != nil {
|
||
_ = onDelta(text)
|
||
}
|
||
return text, nil
|
||
}
|