thread-master/apps/backend/internal/module/ai/stream.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
}