agentu/internal/llm/openai.go

package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

const maxChatAttempts = 3

const retryDelay = 100 * time.Millisecond

type OpenAICompatibleClient struct {
	baseURL    string
	apiKey     string
	httpClient *http.Client
}

func NewOpenAICompatibleClient(baseURL, apiKey string, httpClient *http.Client) *OpenAICompatibleClient {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	return &OpenAICompatibleClient{
		baseURL:    strings.TrimRight(baseURL, "/"),
		apiKey:     apiKey,
		httpClient: httpClient,
	}
}

func (c *OpenAICompatibleClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error {
	req.Stream = true
	body, err := json.Marshal(req)
	if err != nil {
		return fmt.Errorf("marshal chat request: %w", err)
	}

	var lastErr error
	for attempt := 1; attempt <= maxChatAttempts; attempt++ {
		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
		if err != nil {
			return fmt.Errorf("create chat request: %w", err)
		}
		httpReq.Header.Set("Content-Type", "application/json")
		httpReq.Header.Set("Accept", "text/event-stream")
		if c.apiKey != "" {
			httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
		}

		resp, err := c.httpClient.Do(httpReq)
		if err != nil {
			lastErr = fmt.Errorf("send chat request: %w", err)
			if !shouldRetryRequest(ctx, attempt, 0) {
				return lastErr
			}
			if waitErr := waitBeforeRetry(ctx); waitErr != nil {
				return waitErr
			}
			continue
		}

		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
			limit := io.LimitReader(resp.Body, 4096)
			data, _ := io.ReadAll(limit)
			_ = resp.Body.Close()
			lastErr = fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
			if !shouldRetryRequest(ctx, attempt, resp.StatusCode) {
				return lastErr
			}
			if waitErr := waitBeforeRetry(ctx); waitErr != nil {
				return waitErr
			}
			continue
		}

		defer resp.Body.Close()
		return readSSE(resp.Body, emit)
	}
	return lastErr
}

func shouldRetryRequest(ctx context.Context, attempt int, status int) bool {
	if ctx.Err() != nil || attempt >= maxChatAttempts {
		return false
	}
	if status == 0 {
		return true
	}
	return status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500
}

func waitBeforeRetry(ctx context.Context) error {
	timer := time.NewTimer(retryDelay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func readSSE(r io.Reader, emit func(StreamEvent) error) error {
	reader := bufio.NewReader(r)
	for {
		line, err := reader.ReadString('\n')
		if err != nil && !errors.Is(err, io.EOF) {
			return fmt.Errorf("read stream: %w", err)
		}

		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "data:") {
			payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if payload == "[DONE]" {
				return nil
			}
			event, parseErr := parseStreamPayload(payload)
			if parseErr != nil {
				return parseErr
			}
			if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" {
				if emitErr := emit(event); emitErr != nil {
					return emitErr
				}
			}
		}

		if errors.Is(err, io.EOF) {
			return nil
		}
	}
}

type streamPayload struct {
	Choices []struct {
		Delta struct {
			Content   string              `json:"content"`
			ToolCalls []toolCallDeltaJSON `json:"tool_calls"`
		} `json:"delta"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
		Type    string `json:"type"`
	} `json:"error,omitempty"`
}

type toolCallDeltaJSON struct {
	Index    int    `json:"index"`
	ID       string `json:"id"`
	Type     string `json:"type"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `json:"function"`
}

func parseStreamPayload(payload string) (StreamEvent, error) {
	var decoded streamPayload
	if err := json.Unmarshal([]byte(payload), &decoded); err != nil {
		return StreamEvent{}, fmt.Errorf("parse stream payload: %w", err)
	}
	if decoded.Error != nil {
		if decoded.Error.Type != "" {
			return StreamEvent{}, fmt.Errorf("provider error: %s: %s", decoded.Error.Type, decoded.Error.Message)
		}
		return StreamEvent{}, fmt.Errorf("provider error: %s", decoded.Error.Message)
	}
	if len(decoded.Choices) == 0 {
		return StreamEvent{}, nil
	}

	choice := decoded.Choices[0]
	event := StreamEvent{
		Content:      choice.Delta.Content,
		FinishReason: choice.FinishReason,
	}
	for _, tc := range choice.Delta.ToolCalls {
		event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
			Index:     tc.Index,
			ID:        tc.ID,
			Type:      tc.Type,
			Name:      tc.Function.Name,
			Arguments: tc.Function.Arguments,
		})
	}
	return event, nil
}