go-deepseek/completion.go

package deepseek

import (
	"context"
	"errors"
	"io"
	"net/http"
	"strings"
)

// CompletionReq configures a chat completion request.
type CompletionReq struct {
	ChatSessionID   string    `json:"chat_session_id"`
	ParentMessageID *uint64   `json:"parent_message_id"`
	ModelType       ModelType `json:"model_type"`
	promptBuilder   strings.Builder
	Prompt          string   `json:"prompt"`
	RefFileIDs      []string `json:"ref_file_ids"`
	ThinkingEnabled bool     `json:"thinking_enabled"`
	SearchEnabled   bool     `json:"search_enabled"`
}

// NewCompletionReq creates a completion request for chatID using the default
// model type.
func NewCompletionReq(chatID string) *CompletionReq {
	return &CompletionReq{
		ChatSessionID: chatID, ModelType: ModelTypeDefault,
	}
}

// SetParentMessageID sets the parent message for a continued conversation.
func (req *CompletionReq) SetParentMessageID(id uint64) *CompletionReq {
	req.ParentMessageID = new(id)
	return req
}

// SetModelType selects the model used for the completion.
func (req *CompletionReq) SetModelType(t ModelType) *CompletionReq {
	req.ModelType = t
	return req
}

// AddPrompt appends s to the prompt under construction.
func (req *CompletionReq) AddPrompt(s string) *CompletionReq {
	_, _ = req.promptBuilder.WriteString(s)
	return req
}

// BuildPrompt stores the accumulated prompt and resets the internal builder.
func (req *CompletionReq) BuildPrompt() *CompletionReq {
	req.Prompt = req.promptBuilder.String()
	req.promptBuilder = strings.Builder{}
	return req
}

// SetPrompt replaces the request prompt.
func (req *CompletionReq) SetPrompt(prompt string) *CompletionReq {
	req.Prompt = prompt
	return req
}

// SetThinkingEnabled enables or disables model thinking output.
func (req *CompletionReq) SetThinkingEnabled(b bool) *CompletionReq {
	req.ThinkingEnabled = b
	return req
}

// SetSearchEnabled enables or disables web search for the completion.
func (req *CompletionReq) SetSearchEnabled(b bool) *CompletionReq {
	req.SearchEnabled = b
	return req
}

// ModelType identifies a DeepSeek chat model mode.
type ModelType string

const (
	// ModelTypeDefault selects the default model mode.
	ModelTypeDefault ModelType = "default"
	// ModelTypeExpert selects the expert model mode.
	ModelTypeExpert ModelType = "expert"
)

// Completion starts a streaming completion using a background context.
// The caller owns and must either close the response body or pass the response
// to ReadStream or ReadStreamAsString.
func (api *Client) Completion(body CompletionReq) (*http.Response, error) {
	return api.CompletionWithContext(context.Background(), body)
}

// CompletionWithContext starts a streaming completion and honors ctx
// cancellation. The caller owns and must either close the response body or
// pass the response to ReadStream or ReadStreamAsString.
func (api *Client) CompletionWithContext(ctx context.Context, body CompletionReq) (*http.Response, error) {
	req := NewRequest[CreatePowChallengeRes]("POST", "chat/completion", body)
	err := req.SolvePow(ctx, api)
	if err != nil {
		return nil, err
	}

	resp, err := req.DoWithContext(ctx, api)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// ReadStream consumes and closes resp.Body, updates stream state, and calls
// onText for newly appended response text. A nil callback is allowed.
func ReadStream(ctx context.Context, resp *http.Response, onText func(string)) (*StreamState, error) {
	defer resp.Body.Close()

	reader := NewSSEReader(resp.Body)
	state := &StreamState{}

	var previousLength int

	for {
		event, err := reader.Next(ctx)
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return nil, err
		}

		if err := state.Apply(event); err != nil {
			return nil, err
		}

		current := state.Content.String()

		if len(current) > previousLength {
			chunk := current[previousLength:]

			if onText != nil {
				onText(chunk)
			}

			previousLength = len(current)
		}

		if event.Event == "close" {
			break
		}
	}

	return state, nil
}

// ReadStreamAsString consumes and closes resp.Body and returns the final
// response content.
func ReadStreamAsString(ctx context.Context, resp *http.Response) (string, error) {
	state, err := ReadStream(ctx, resp, nil)
	if err != nil {
		return "", err
	}
	return state.Content.String(), nil
}