odidere/internal/llm/client_completions.go

package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json/v2"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
)

// CompletionsParams contains the inputs to a chat completion.
type CompletionsParams struct {
	// Messages is the conversation presented to the model, in order.
	Messages []Message
	// Model identifies the model according to the configured endpoint.
	Model string
	// ThinkingBudgetTokens is sent as the endpoint's reasoning-token budget.
	// Nil omits both the budget and enable_thinking; zero sets enable_thinking
	// to false; any other value sets it to true.
	ThinkingBudgetTokens *int
	// Tools describes the functions available for the model to call.
	Tools []APITool
}

// CompletionResult contains one generated Message and its response metadata.
type CompletionResult struct {
	// Message is the generated message.
	Message Message
	// Metadata describes this inference's selected choice and endpoint,
	// separately from the model-visible message.
	Metadata Metadata
}

// Completions requests a non-streaming chat completion and returns the first
// choice. It returns an error if the response contains no choices.
func (c *Client) Completions(
	ctx context.Context, p CompletionsParams,
) (*CompletionResult, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	cr, err := c.completions(ctx, p)
	if err != nil {
		return nil, err
	}
	choice := cr.Choices[0]
	return &CompletionResult{
		Message: choice.Message,
		Metadata: Metadata{
			ChoiceIndex:       choice.Index,
			Created:           cr.Created,
			FinishReason:      choice.FinishReason,
			ID:                cr.ID,
			Logprobs:          choice.Logprobs,
			Model:             cr.Model,
			Provider:          new(c.provider),
			RequestedModel:    &p.Model,
			ServiceTier:       cr.ServiceTier,
			SystemFingerprint: cr.SystemFingerprint,
			Timings:           cr.Timings,
			Usage:             cr.Usage,
		},
	}, nil
}

// completions sends a single non-streaming request and returns the
// full API response.
func (c *Client) completions(
	ctx context.Context, p CompletionsParams,
) (*ChatResponse, error) {
	kwargs := map[string]any{}
	if p.ThinkingBudgetTokens != nil {
		kwargs[chatTemplateEnableThinking] = *p.ThinkingBudgetTokens != 0
	}
	body, err := json.Marshal(ChatRequest{
		Model:                p.Model,
		Messages:             p.Messages,
		Tools:                p.Tools,
		ThinkingBudgetTokens: p.ThinkingBudgetTokens,
		ChatTemplateKWARGS:   kwargs,
		ParallelToolCalls:    new(true),
	}, json.FormatNilSliceAsNull(true))
	if err != nil {
		return nil, fmt.Errorf(
			"completions: marshal request body: %w", err,
		)
	}
	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		c.baseURL+chatCompletionsPath,
		bytes.NewReader(body),
	)
	if err != nil {
		return nil, fmt.Errorf("completions: create request: %w", err)
	}
	if c.key != "" {
		req.Header.Set("Authorization", "Bearer "+c.key)
	}
	req.Header.Set("Content-Type", "application/json")

	res, err := c.httpc.Do(req)
	if err != nil {
		return nil, fmt.Errorf("completions: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		body, err := io.ReadAll(res.Body)
		if err != nil {
			c.log.ErrorContext(
				ctx,
				"completions: failed to read error response body",
				slog.Any("error", err),
			)
			return nil, fmt.Errorf(
				"completions: status code %d", res.StatusCode,
			)
		}
		var apiError struct {
			Error *APIError `json:"error"`
		}
		if err := json.Unmarshal(
			body, &apiError,
		); err == nil && apiError.Error != nil {
			c.log.ErrorContext(
				ctx,
				"completions: request error",
				slog.Int("status", res.StatusCode),
				slog.String("code", apiError.Error.Code),
				slog.String(
					"message", apiError.Error.Message,
				),
				slog.String("type", apiError.Error.Type),
			)
			return nil, fmt.Errorf("completions: %w", apiError.Error)
		}
		c.log.ErrorContext(
			ctx,
			"completions: API error",
			slog.Int("status", res.StatusCode),
			slog.String("body", string(body)),
		)
		return nil, fmt.Errorf(
			"completions: %d: %s",
			res.StatusCode, string(body),
		)
	}

	var cr = &ChatResponse{}
	if err := json.UnmarshalRead(res.Body, cr); err != nil {
		return nil, fmt.Errorf("completions: decode response: %w", err)
	}
	if len(cr.Choices) == 0 {
		return nil, fmt.Errorf(
			"completions: no response choices returned",
		)
	}
	return cr, nil
}

// StreamAccumulator reconstructs a Message from the ordered deltas for one
// streaming choice. The zero value is ready for use.
type StreamAccumulator struct {
	// Content contains the accumulated message text.
	Content strings.Builder
	// Reasoning combines fragments received under either supported reasoning
	// field name.
	Reasoning strings.Builder
	// Refusal contains the accumulated refusal text.
	Refusal strings.Builder
	// Role is the latest non-empty role received.
	Role Role
	// ToolCalls contains calls assembled by tool-call index.
	ToolCalls []ToolCall
}

// Add incorporates the next delta. Content, refusal, reasoning, and tool-call
// fragments are appended in arrival order. Tool-call fragments are grouped by
// Index; a missing index denotes the first call. A non-empty Role replaces the
// previously accumulated role.
func (acc *StreamAccumulator) Add(delta Delta) {
	if delta.Role != "" {
		acc.Role = delta.Role
	}
	acc.Content.WriteString(delta.Content)
	acc.Refusal.WriteString(delta.Refusal)
	acc.Reasoning.WriteString(delta.ReasoningContent)
	acc.Reasoning.WriteString(delta.Reasoning)

	for _, tc := range delta.ToolCalls {
		idx := 0
		if tc.Index != nil {
			idx = *tc.Index
		}
		for len(acc.ToolCalls) <= idx {
			acc.ToolCalls = append(
				acc.ToolCalls,
				ToolCall{Type: ToolTypeFunction},
			)
		}
		if tc.ID != "" {
			acc.ToolCalls[idx].ID = tc.ID
		}
		if tc.Type != "" {
			acc.ToolCalls[idx].Type = tc.Type
		}
		acc.ToolCalls[idx].Function.Name += tc.Function.Name
		acc.ToolCalls[idx].Function.Arguments += tc.Function.Arguments
	}
}

// Message returns the accumulated state without resetting the accumulator. If
// the stream contains a refusal, the returned content consists of that refusal
// rather than the accumulated text.
func (acc *StreamAccumulator) Message() Message {
	var contentParts []ContentPart
	if refusal := acc.Refusal.String(); refusal != "" {
		contentParts = []ContentPart{{
			Type:    ContentTypeRefusal,
			Refusal: refusal,
		}}
	} else {
		contentParts = []ContentPart{{
			Type: ContentTypeText,
			Text: acc.Content.String(),
		}}
	}

	return Message{
		ContentParts:     contentParts,
		ReasoningContent: acc.Reasoning.String(),
		Refusal:          acc.Refusal.String(),
		Role:             acc.Role,
		ToolCalls:        acc.ToolCalls,
	}
}

// CompletionStream is a caller-driven chat-completion stream. A stream owns
// its response body and timeout context until it terminates or is closed.
type CompletionStream struct {
	ctx    context.Context
	reader *bufio.Reader
	body   io.ReadCloser
	cancel context.CancelFunc

	acc     StreamAccumulator
	current StreamChunk
	// metadata retains reported scalars and the latest usage/timing
	// snapshots.
	// Logprob sequences contain every reported token in arrival order.
	// Nested metadata values are shared with Current and Result.
	metadata      Metadata
	err           error
	cleanupErr    error
	terminated    bool
	resourcesDone bool
}

// CompletionsStream opens a streaming chat completion. Request setup,
// transport, and non-success response errors are returned before a stream is
// created.
func (c *Client) CompletionsStream(
	ctx context.Context,
	p CompletionsParams,
) (*CompletionStream, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)

	kwargs := map[string]any{}
	if p.ThinkingBudgetTokens != nil {
		kwargs[chatTemplateEnableThinking] = *p.ThinkingBudgetTokens != 0
	}
	body, err := json.Marshal(ChatRequest{
		ChatTemplateKWARGS:   kwargs,
		Messages:             p.Messages,
		Model:                p.Model,
		ParallelToolCalls:    new(true),
		Stream:               true,
		StreamOptions:        &StreamOptions{IncludeUsage: true},
		ThinkingBudgetTokens: p.ThinkingBudgetTokens,
		Tools:                p.Tools,
	}, json.FormatNilSliceAsNull(true))
	if err != nil {
		cancel()
		return nil, fmt.Errorf(
			"completions stream: marshal request body: %w", err,
		)
	}

	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		c.baseURL+chatCompletionsPath,
		bytes.NewReader(body),
	)
	if err != nil {
		cancel()
		return nil, fmt.Errorf(
			"completions stream: create request: %w", err,
		)
	}
	if c.key != "" {
		req.Header.Set("Authorization", "Bearer "+c.key)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream")

	res, err := c.httpc.Do(req)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("completions stream: %w", err)
	}

	if res.StatusCode != http.StatusOK {
		defer cancel()
		defer res.Body.Close()
		body, err := io.ReadAll(res.Body)
		if err != nil {
			c.log.ErrorContext(
				ctx,
				"completions stream: failed to read error response body",
				slog.Any("error", err),
			)
			return nil, fmt.Errorf(
				"completions stream: status code %d",
				res.StatusCode,
			)
		}
		var apiError struct {
			Error *APIError `json:"error"`
		}
		if err := json.Unmarshal(
			body, &apiError,
		); err == nil && apiError.Error != nil {
			c.log.ErrorContext(
				ctx,
				"completions stream: request error",
				slog.Int("status", res.StatusCode),
				slog.String("code", apiError.Error.Code),
				slog.String("message", apiError.Error.Message),
				slog.String("type", apiError.Error.Type),
			)
			return nil, fmt.Errorf(
				"completions stream: %w", apiError.Error,
			)
		}
		c.log.ErrorContext(
			ctx,
			"completions stream: API error",
			slog.Int("status", res.StatusCode),
			slog.String("body", string(body)),
		)
		return nil, fmt.Errorf(
			"completions stream: %d: %s",
			res.StatusCode,
			string(body),
		)
	}

	return &CompletionStream{
		ctx:      ctx,
		reader:   bufio.NewReader(res.Body),
		body:     res.Body,
		cancel:   cancel,
		metadata: Metadata{Provider: new(c.provider), RequestedModel: &p.Model},
	}, nil
}

// Next advances to the next JSON data event. It returns false after [DONE], an
// error, or Close. On an error, Err reports the cause and Result remains
// available with all chunks accumulated before the error.
func (s *CompletionStream) Next() bool {
	if s.terminated {
		return false
	}

	var data []string
	for {
		line, err := s.reader.ReadString('\n')
		if err != nil && err != io.EOF {
			s.fail(fmt.Errorf("completions stream: read response: %w", err))
			return false
		}
		if len(line) > 0 {
			line = strings.TrimSuffix(line, "\n")
			line = strings.TrimSuffix(line, "\r")
			if line == "" {
				if len(data) > 0 {
					break
				}
			} else if line[0] != ':' {
				field, value, found := strings.Cut(line, ":")
				if !found {
					field = line
					value = ""
				}
				value = strings.TrimPrefix(value, " ")
				if field == "data" {
					data = append(data, value)
				}
			}
		}
		if err == io.EOF {
			s.fail(io.ErrUnexpectedEOF)
			return false
		}
	}

	payload := strings.Join(data, "\n")
	if payload == sseDone {
		s.terminated = true
		s.release()
		return false
	}

	var chunk StreamChunk
	if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
		s.fail(fmt.Errorf("completions stream: decode chunk: %w", err))
		return false
	}
	for _, choice := range chunk.Choices {
		if choice.Index != nil && *choice.Index != 0 {
			s.fail(fmt.Errorf(
				"completions stream: unexpected choice index %d",
				*choice.Index,
			))
			return false
		}
		if choice.Index != nil {
			s.metadata.ChoiceIndex = choice.Index
		}
		if choice.FinishReason != nil {
			s.metadata.FinishReason = choice.FinishReason
		}
		if choice.Logprobs != nil {
			if s.metadata.Logprobs == nil {
				s.metadata.Logprobs = &Logprobs{}
			}
			if choice.Logprobs.Content != nil &&
				s.metadata.Logprobs.Content == nil {
				s.metadata.Logprobs.Content = []TokenLogprob{}
			}
			if choice.Logprobs.Refusal != nil &&
				s.metadata.Logprobs.Refusal == nil {
				s.metadata.Logprobs.Refusal = []TokenLogprob{}
			}
			s.metadata.Logprobs.Content = append(
				s.metadata.Logprobs.Content,
				choice.Logprobs.Content...,
			)
			s.metadata.Logprobs.Refusal = append(
				s.metadata.Logprobs.Refusal,
				choice.Logprobs.Refusal...,
			)
		}
		s.acc.Add(choice.Delta)
	}
	if chunk.Created != nil {
		s.metadata.Created = chunk.Created
	}
	if chunk.ID != nil {
		s.metadata.ID = chunk.ID
	}
	if chunk.Model != nil {
		s.metadata.Model = chunk.Model
	}
	if chunk.ServiceTier != nil {
		s.metadata.ServiceTier = chunk.ServiceTier
	}
	if chunk.SystemFingerprint != nil {
		s.metadata.SystemFingerprint = chunk.SystemFingerprint
	}
	if chunk.Usage != nil {
		s.metadata.Usage = chunk.Usage
	}
	if chunk.Timings != nil {
		s.metadata.Timings = chunk.Timings
	}
	s.current = chunk
	return true
}

// Current returns the most recently decoded chunk. The value is replaced by
// the next successful call to Next. Previously returned chunks remain valid.
// Reported metadata shares pointer targets and nested logprob values with
// the stream and Result. Mutations through those references affect both.
func (s *CompletionStream) Current() StreamChunk {
	return s.current
}

// Result returns the message and metadata accumulated so far. Advancing the
// stream does not change previously returned metadata values. Metadata pointer
// targets and nested logprob slices are shared with the stream and other
// results; mutations through those references are visible through the aliases.
// Err reports failures. Nil Err alone does not establish completion: the
// stream completes successfully when Next consumes the DONE marker.
func (s *CompletionStream) Result() CompletionResult {
	message := s.acc.Message()
	message.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
	metadata := s.metadata
	if metadata.Logprobs != nil {
		logprobs := &Logprobs{}
		if metadata.Logprobs.Content != nil {
			logprobs.Content = make(
				[]TokenLogprob, len(metadata.Logprobs.Content),
			)
			copy(logprobs.Content, metadata.Logprobs.Content)
		}
		if metadata.Logprobs.Refusal != nil {
			logprobs.Refusal = make(
				[]TokenLogprob, len(metadata.Logprobs.Refusal),
			)
			copy(logprobs.Refusal, metadata.Logprobs.Refusal)
		}
		metadata.Logprobs = logprobs
	}
	return CompletionResult{Message: message, Metadata: metadata}
}

// Err returns the stream failure, if any. Closing a stream does not create or
// replace a stream failure.
func (s *CompletionStream) Err() error {
	return s.err
}

// Close releases the response body and timeout context. It is safe to call
// repeatedly and returns the first response-body cleanup error.
func (s *CompletionStream) Close() error {
	s.terminated = true
	s.release()
	return s.cleanupErr
}

func (s *CompletionStream) fail(err error) {
	if ctxErr := s.ctx.Err(); ctxErr != nil {
		err = ctxErr
	}
	s.err = err
	s.terminated = true
	s.release()
}

func (s *CompletionStream) release() {
	if s.resourcesDone {
		return
	}
	s.resourcesDone = true
	if err := s.body.Close(); err != nil {
		s.cleanupErr = err
	}
	s.cancel()
}