Luminary/internal/provider/stream.go

package provider

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"

	"github.com/rose_cat707/luminary/internal/model"
)

type StreamResult struct {
	StatusCode int
	TokensIn   int
	TokensOut  int
	Body       []byte // populated for non-2xx responses
}

func ForwardStreamHTTP(ctx context.Context, client *http.Client, method, url string, apiKey string, providerType model.ProviderType, body []byte, w io.Writer, flusher http.Flusher) (*StreamResult, error) {
	var reader io.Reader
	if len(body) > 0 {
		reader = bytes.NewReader(body)
	}
	req, err := http.NewRequestWithContext(ctx, method, url, reader)
	if err != nil {
		return nil, err
	}
	setAuthHeaders(req, apiKey, providerType)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream")
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		b, _ := io.ReadAll(resp.Body)
		if flusher != nil {
			w.Write(b)
			flusher.Flush()
		}
		return &StreamResult{StatusCode: resp.StatusCode, Body: b}, nil
	}

	result := &StreamResult{StatusCode: resp.StatusCode}
	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 64*1024), 1024*1024)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "data: ") {
			data := strings.TrimPrefix(line, "data: ")
			if data == "[DONE]" {
				continue
			}
			tokensIn, tokensOut := parseStreamChunk([]byte(data))
			result.TokensIn += tokensIn
			result.TokensOut += tokensOut
		}
		if _, err := w.Write([]byte(line + "\n")); err != nil {
			return result, err
		}
		if flusher != nil {
			flusher.Flush()
		}
	}
	return result, scanner.Err()
}

func parseStreamChunk(data []byte) (int, int) {
	var chunk struct {
		Usage struct {
			PromptTokens     int `json:"prompt_tokens"`
			CompletionTokens int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(data, &chunk); err == nil && chunk.Usage.PromptTokens+chunk.Usage.CompletionTokens > 0 {
		return chunk.Usage.PromptTokens, chunk.Usage.CompletionTokens
	}
	return 0, 0
}

func NewStreamHTTPClient() *http.Client {
	return &http.Client{} // no timeout for streams
}

func IsStreamRequest(body []byte) bool {
	var payload struct {
		Stream bool `json:"stream"`
	}
	_ = json.Unmarshal(body, &payload)
	return payload.Stream
}

func SupportsStreaming(endpoint EndpointType) bool {
	return endpoint == EndpointChatCompletion || endpoint == EndpointResponses
}