Improve config lookup and OpenAI streaming decode

internal/openai/client.go

@@ -1,6 +1,7 @@
 package openai
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/json"
@@ -86,6 +87,9 @@ func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionReq
 	defer resp.Body.Close()
 
 	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
+		if req.Stream {
+			return decodeStream(resp.Body)
+		}
 		return decodeSuccess(resp.Body)
 	}
 
@@ -93,13 +97,121 @@ func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionReq
 }
 
 func decodeSuccess(r io.Reader) (*ChatCompletionResponse, error) {
-	var response ChatCompletionResponse
-	if err := json.NewDecoder(r).Decode(&response); err != nil {
-		return nil, fmt.Errorf("decode response: %w", err)
+	data, err := io.ReadAll(r)
+	if err != nil {
+		return nil, fmt.Errorf("read response: %w", err)
 	}
+
+	var response ChatCompletionResponse
+	if err := json.Unmarshal(data, &response); err != nil {
+		trimmed := bytes.TrimSpace(data)
+		if len(trimmed) == 0 {
+			return nil, fmt.Errorf("decode response: %w", err)
+		}
+		return &ChatCompletionResponse{
+			Choices: []ChatCompletionChoice{{
+				Message: ChatMessage{Role: "assistant", Content: string(trimmed)},
+			}},
+		}, nil
+	}
+
 	return &response, nil
 }
 
+func decodeStream(r io.Reader) (*ChatCompletionResponse, error) {
+	scanner := bufio.NewScanner(r)
+	var payloads []json.RawMessage
+
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.HasPrefix(line, "data: ") {
+			payload := strings.TrimPrefix(line, "data: ")
+			if payload == "[DONE]" {
+				break
+			}
+			payloads = append(payloads, json.RawMessage(payload))
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return nil, fmt.Errorf("read stream: %w", err)
+	}
+
+	if len(payloads) == 0 {
+		return nil, errors.New("empty stream response")
+	}
+
+	type streamChunk struct {
+		ID      string `json:"id"`
+		Object  string `json:"object"`
+		Choices []struct {
+			Index        int         `json:"index"`
+			Message      ChatMessage `json:"message"`
+			Delta        ChatMessage `json:"delta"`
+			FinishReason string      `json:"finish_reason"`
+		} `json:"choices"`
+		Usage Usage `json:"usage"`
+	}
+
+	var aggregated ChatCompletionResponse
+	var builder strings.Builder
+	role := "assistant"
+	finish := ""
+
+	for _, payload := range payloads {
+		var chunk streamChunk
+		if err := json.Unmarshal(payload, &chunk); err != nil {
+			return nil, fmt.Errorf("decode stream response: %w", err)
+		}
+		if aggregated.ID == "" {
+			aggregated.ID = chunk.ID
+		}
+		if aggregated.Object == "" {
+			aggregated.Object = chunk.Object
+		}
+		aggregated.Usage.PromptTokens += chunk.Usage.PromptTokens
+		aggregated.Usage.CompletionTokens += chunk.Usage.CompletionTokens
+		aggregated.Usage.TotalTokens += chunk.Usage.TotalTokens
+
+		if len(chunk.Choices) == 0 {
+			continue
+		}
+
+		choice := chunk.Choices[0]
+		if choice.Message.Role != "" {
+			role = choice.Message.Role
+		}
+		if choice.Delta.Role != "" {
+			role = choice.Delta.Role
+		}
+		if choice.Message.Content != "" {
+			builder.WriteString(choice.Message.Content)
+		}
+		if choice.Delta.Content != "" {
+			builder.WriteString(choice.Delta.Content)
+		}
+		if choice.FinishReason != "" {
+			finish = choice.FinishReason
+		}
+	}
+
+	content := strings.TrimSpace(builder.String())
+	if content == "" {
+		return nil, errors.New("stream response contained no content")
+	}
+
+	aggregated.Choices = []ChatCompletionChoice{{
+		Index:        0,
+		FinishReason: finish,
+		Message: ChatMessage{
+			Role:    role,
+			Content: content,
+		},
+	}}
+
+	return &aggregated, nil
+}
+
 func decodeError(r io.Reader, status int) error {
 	var apiErr ErrorResponse
 	if err := json.NewDecoder(r).Decode(&apiErr); err != nil {