odidere/internal/llm/client_completions.go

package llm

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

// ErrMaxIterations is returned when the agent loop exceeds the configured
// maximum number of iterations.
var ErrMaxIterations = fmt.Errorf("max iterations exceeded")

// CompletionsResult bundles the messages generated during a completions
// request along with usage and timings from the final API response.
type CompletionsResult struct {
	Messages []Message
	Usage    *Usage
	Timings  *Timings
}

// CompletionsParams holds parameters for Completions and CompletionsStream.
type CompletionsParams struct {
	Model                string
	SystemMessage        string
	MaxIterations        int
	Messages             []Message
	ThinkingBudgetTokens *int
}

// Completions sends messages to the LLM using the specified model.
// If systemMessage is non-empty, it overrides the configured system message.
// maxIterations caps the number of agent loop iterations.
// A value of 0 means unlimited (no cap).
// Returns all messages generated during the completions request, including
// tool calls and tool results, along with usage and timings from the final response.
func (c *Client) Completions(
	ctx context.Context, p CompletionsParams,
) (*CompletionsResult, error) {
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	model, systemMessage, maxIterations, messages, thinkingBudgetTokens :=
		p.Model, p.SystemMessage, p.MaxIterations, p.Messages, p.ThinkingBudgetTokens

	// Use request system message if provided, config default otherwise.
	smsg := c.systemMessage
	if systemMessage != "" {
		smsg = systemMessage
	}
	// Prepend system message, if configured and not already present.
	if smsg != "" &&
		(len(messages) == 0 || messages[0].Role != RoleSystem) {
		messages = append(
			[]Message{{
				Role: RoleSystem,
				ContentParts: []ContentPart{{
					Type: ContentTypeText,
					Text: smsg,
				}},
			}},
			messages...,
		)
	}

	// Remember the starting length so we can return only the messages
	// generated during this call (the suffix of the conversation history).
	start := len(messages)
	// Loop for tool calls. maxIterations == 0 means unlimited.
	var finalUsage *Usage
	var finalTimings *Timings
	for i := 0; maxIterations == 0 || i < maxIterations; i++ {
		cr, err := c.completions(ctx, model, thinkingBudgetTokens, messages)
		if err != nil {
			return nil, fmt.Errorf("completions: %w", err)
		}
		// Capture usage and timings from this iteration.
		if cr.Usage != nil {
			finalUsage = cr.Usage
		}
		if cr.Timings != nil {
			finalTimings = cr.Timings
		}
		// Use first choice for tool call loop.
		message := &cr.Choices[0].Message
		messages = append(messages, *message)
		if len(message.ToolCalls) == 0 {
			return &CompletionsResult{
				Messages: messages[start:],
				Usage:    finalUsage,
				Timings:  finalTimings,
			}, nil
		}
		toolResults := c.callTools(ctx, *message)
		messages = append(messages, toolResults...)
	}

	return &CompletionsResult{
		Messages: messages[start:],
		Usage:    finalUsage,
		Timings:  finalTimings,
	}, ErrMaxIterations
}

// completions sends a single non-streaming request and returns the
// full API response.
func (c *Client) completions(
	ctx context.Context, model string, thinkingBudgetTokens *int, messages []Message,
) (*ChatResponse, error) {
	kwargs := map[string]any{}
	if thinkingBudgetTokens != nil {
		kwargs["enable_thinking"] = *thinkingBudgetTokens != 0
	}
	body, err := json.Marshal(ChatRequest{
		Model:                model,
		Messages:             messages,
		Tools:                c.tools,
		ThinkingBudgetTokens: thinkingBudgetTokens,
		ChatTemplateKWARGS:   kwargs,
		ParallelToolCalls:    new(true),
	})
	if err != nil {
		return nil, fmt.Errorf(
			"completions: marshal request body: %w", err,
		)
	}
	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		c.baseURL+"/chat/completions",
		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 *struct {
				Type    string `json:"type"`
				Code    string `json:"code"`
				Message string `json:"message"`
			} `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: %s: %s",
				apiError.Error.Code,
				apiError.Error.Message,
			)
		}
		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.NewDecoder(res.Body).Decode(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
}

// StreamEvent wraps a message or delta produced during streaming.
type StreamEvent struct {
	// Type identifies the event kind: delta, reasoning_delta, done, or message.
	Type string
	// Delta is a text fragment from the assistant response.
	Delta string
	// ReasoningDelta is a reasoning text fragment from the assistant response.
	ReasoningDelta string
	// Message is the complete chat completion message.
	Message Message
	// Usage is token usage info from the final chunk, sent on the done event.
	Usage *Usage
	// Timings is llama.cpp timing info from the final chunk, sent on the done event.
	Timings *Timings
}

const (
	streamEventDelta          = "delta"
	streamEventReasoningDelta = "reasoning_delta"
	streamEventDone           = "done"
	streamEventMessage        = "message"
)

// CompletionsStream sends messages to the LLM using the specified model and
// streams results. Each complete message (assistant reply, tool call,
// tool result) is sent to the events channel as it becomes available.
// The channel is closed before returning.
// If systemMessage is non-empty, it overrides the configured system message.
// maxIterations caps the number of agent loop iterations; a value of 0
// means unlimited (no cap).
// Returns ErrMaxIterations if the iteration limit is exceeded.
func (c *Client) CompletionsStream(
	ctx context.Context, p CompletionsParams, events chan<- StreamEvent,
) error {
	defer close(events)
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	model, systemMessage, maxIterations, messages, thinkingBudgetTokens :=
		p.Model, p.SystemMessage, p.MaxIterations, p.Messages, p.ThinkingBudgetTokens

	// Use request system message if provided, config default otherwise.
	smsg := c.systemMessage
	if systemMessage != "" {
		smsg = systemMessage
	}

	// Prepend system message, if configured and not already present.
	if smsg != "" &&
		(len(messages) == 0 || messages[0].Role != RoleSystem) {
		messages = append(
			[]Message{{
				Role: RoleSystem,
				ContentParts: []ContentPart{
					{Type: ContentTypeText, Text: smsg},
				},
			}},
			messages...,
		)
	}

	// Loop for tool calls. maxIterations == 0 means unlimited.
	for i := 0; maxIterations == 0 || i < maxIterations; i++ {
		choices, err := c.completionsStream(ctx, messages, model, thinkingBudgetTokens, events)
		if err != nil {
			return fmt.Errorf("completions stream: %w", err)
		}
		// Use first choice for tool call loop.
		message := choices[0]
		if len(message.ToolCalls) == 0 {
			return nil
		}
		toolResults := c.callTools(ctx, message)
		for _, tr := range toolResults {
			events <- StreamEvent{
				Type:    streamEventMessage,
				Message: tr,
			}
		}
		messages = append(messages, choices...)
		messages = append(messages, toolResults...)
	}

	return ErrMaxIterations
}

// streamAccum accumulates deltas for a single choice in a streaming response.
type streamAccum struct {
	content   strings.Builder
	reasoning strings.Builder
	refusal   strings.Builder
	role      string
	toolCalls []ToolCall
}

// completionsStream sends a single streaming request, accumulates the
// response, and returns all assistant messages from the response choices.
// Sends each message to the events channel. Usage and timings are included
// on the done event for the last message.
func (c *Client) completionsStream(ctx context.Context, messages []Message, model string, thinkingBudgetTokens *int, events chan<- StreamEvent) ([]Message, error) {
	kwargs := map[string]any{}
	if thinkingBudgetTokens != nil {
		kwargs["enable_thinking"] = *thinkingBudgetTokens != 0
	}
	body, err := json.Marshal(ChatRequest{
		Model:                model,
		Messages:             messages,
		Stream:               true,
		Tools:                c.tools,
		ThinkingBudgetTokens: thinkingBudgetTokens,
		ChatTemplateKWARGS:   kwargs,
		ParallelToolCalls:    new(true),
		StreamOptions:        &StreamOptions{IncludeUsage: true},
	})
	if err != nil {
		return nil, fmt.Errorf(
			"completions stream: marshal request body: %w", err,
		)
	}
	req, err := http.NewRequestWithContext(
		ctx, http.MethodPost,
		c.baseURL+"/chat/completions",
		bytes.NewReader(body),
	)
	if err != nil {
		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 {
		return nil, fmt.Errorf("completions stream: %w", err)
	}

	if res.StatusCode != http.StatusOK {
		body, err := io.ReadAll(res.Body)
		res.Body.Close()
		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 *struct {
				Type    string `json:"type"`
				Code    string `json:"code"`
				Message string `json:"message"`
			} `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: %s: %s",
				apiError.Error.Code,
				apiError.Error.Message,
			)
		}
		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),
		)
	}

	// Accumulate the streamed response per choice index.
	accums := make(map[int]*streamAccum)
	var usage *Usage
	var timings *Timings

	scanner := bufio.NewScanner(res.Body)
	for scanner.Scan() {
		line := scanner.Text()
		// Skip empty lines and comments.
		if line == "" || line == ":" {
			continue
		}
		// Parse data lines.
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := strings.TrimPrefix(line, "data: ")
		// [DONE] signals end of stream.
		if data == "[DONE]" {
			break
		}
		var chunk StreamChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			continue
		}
		for _, ch := range chunk.Choices {
			acc, ok := accums[ch.Index]
			if !ok {
				acc = &streamAccum{}
				accums[ch.Index] = acc
			}
			delta := ch.Delta
			if delta.Role != "" {
				acc.role = delta.Role
			}
			if delta.Content != "" {
				acc.content.WriteString(delta.Content)
				events <- StreamEvent{
					Type:  streamEventDelta,
					Delta: delta.Content,
				}
			}
			if delta.Refusal != "" {
				acc.refusal.WriteString(delta.Refusal)
			}
			if delta.ReasoningContent != "" {
				acc.reasoning.WriteString(
					delta.ReasoningContent,
				)
				events <- StreamEvent{
					Type:           streamEventReasoningDelta,
					ReasoningDelta: delta.ReasoningContent,
				}
			}
			if delta.Reasoning != "" {
				acc.reasoning.WriteString(delta.Reasoning)
				events <- StreamEvent{
					Type:           streamEventReasoningDelta,
					ReasoningDelta: delta.Reasoning,
				}
			}
			// Accumulate tool call deltas by index.
			for _, tc := range delta.ToolCalls {
				idx := 0
				if tc.Index != nil {
					idx = *tc.Index
				}
				// Grow the slice as needed.
				for len(acc.toolCalls) <= idx {
					acc.toolCalls = append(
						acc.toolCalls,
						ToolCall{Type: "function"},
					)
				}
				if tc.ID != "" {
					acc.toolCalls[idx].ID = tc.ID
				}
				if tc.Function.Name != "" {
					acc.toolCalls[idx].Function.Name +=
						tc.Function.Name
				}
				if tc.Function.Arguments != "" {
					acc.toolCalls[idx].Function.Arguments +=
						tc.Function.Arguments
				}
			}
		}
		// Capture usage from the final chunk (has null choices).
		if chunk.Usage != nil {
			usage = chunk.Usage
		}
		if chunk.Timings != nil {
			timings = chunk.Timings
		}
	}
	res.Body.Close()

	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf(
			"completions stream: read error: %w", err,
		)
	}

	// Build messages in index order.
	indices := make([]int, 0, len(accums))
	for idx := range accums {
		indices = append(indices, idx)
	}
	slices.Sort(indices)

	msgs := make([]Message, 0, len(accums))
	for i, idx := range indices {
		acc := accums[idx]

		// Content and refusal are mutually exclusive in the API.
		var contentParts []ContentPart
		if refusalStr := acc.refusal.String(); refusalStr != "" {
			contentParts = []ContentPart{{
				Type:    ContentTypeRefusal,
				Refusal: refusalStr,
			}}
		} else {
			contentParts = []ContentPart{{
				Type: ContentTypeText,
				Text: acc.content.String(),
			}}
		}

		message := Message{
			Role:             acc.role,
			ContentParts:     contentParts,
			Refusal:          acc.refusal.String(),
			ReasoningContent: acc.reasoning.String(),
			ToolCalls:        acc.toolCalls,
		}

		// Attach usage and timings to the done event of the last message.
		var evUsage *Usage
		var evTimings *Timings
		if i == len(indices)-1 {
			evUsage = usage
			evTimings = timings
		}
		events <- StreamEvent{
			Type:    streamEventDone,
			Message: message,
			Usage:   evUsage,
			Timings: evTimings,
		}
		msgs = append(msgs, message)
	}

	return msgs, nil
}