go-deepseek/sse.go

package deepseek

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"strings"
)

// SSEEvent is one event decoded from a server-sent event stream.
type SSEEvent struct {
	ID    string
	Event string
	Data  []byte
	Retry string
}

// SSEReader parses server-sent events from an input stream.
type SSEReader struct {
	reader *bufio.Reader
}

// NewSSEReader creates an SSE reader for r.
func NewSSEReader(r io.Reader) *SSEReader {
	return &SSEReader{
		reader: bufio.NewReaderSize(r, 64*1024),
	}
}

// Next reads and returns the next SSE event. It returns io.EOF when the input
// ends without another event and checks ctx between blocking reads.
func (r *SSEReader) Next(ctx context.Context) (SSEEvent, error) {
	var event SSEEvent
	var dataLines []string

	for {
		if err := ctx.Err(); err != nil {
			return SSEEvent{}, err
		}

		line, err := r.reader.ReadString('\n')
		if err != nil && !errors.Is(err, io.EOF) {
			return SSEEvent{}, fmt.Errorf(
				"read SSE stream: %w",
				err,
			)
		}

		line = strings.TrimSuffix(line, "\n")
		line = strings.TrimSuffix(line, "\r")

		if line == "" {
			if len(dataLines) == 0 &&
				event.ID == "" &&
				event.Event == "" &&
				event.Retry == "" {

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

				continue
			}

			event.Data = []byte(
				strings.Join(dataLines, "\n"),
			)

			return event, nil
		}

		// Строки-комментарии SSE.
		if strings.HasPrefix(line, ":") {
			if errors.Is(err, io.EOF) {
				return SSEEvent{}, io.EOF
			}

			continue
		}

		field, value, found := strings.Cut(line, ":")
		if found && strings.HasPrefix(value, " ") {
			value = value[1:]
		}

		switch field {
		case "event":
			event.Event = value

		case "data":
			dataLines = append(dataLines, value)

		case "id":
			event.ID = value

		case "retry":
			event.Retry = value
		}

		if errors.Is(err, io.EOF) {
			if len(dataLines) > 0 ||
				event.ID != "" ||
				event.Event != "" ||
				event.Retry != "" {

				event.Data = []byte(
					strings.Join(dataLines, "\n"),
				)

				return event, nil
			}

			return SSEEvent{}, io.EOF
		}
	}
}

// ReadyEvent identifies the request and response messages for a stream.
type ReadyEvent struct {
	RequestMessageID  int64  `json:"request_message_id"`
	ResponseMessageID int64  `json:"response_message_id"`
	ModelType         string `json:"model_type"`
}

// UpdateSessionEvent reports the latest session update timestamp.
type UpdateSessionEvent struct {
	UpdatedAt float64 `json:"updated_at"`
}

// TitleEvent contains a generated chat title.
type TitleEvent struct {
	Content string `json:"content"`
}

// CloseEvent describes how a completed stream should be closed or resumed.
type CloseEvent struct {
	ClickBehavior string `json:"click_behavior"`
	AutoResume    bool   `json:"auto_resume"`
}

// PatchEvent describes a compact patch in the DeepSeek stream protocol.
//
// Pointer fields distinguish:
//
//	{"v":"text"}
//
// from:
//
//	{"p":"","o":"","v":"text"}
//
// This distinction is required to inherit the previous path and operation.
type PatchEvent struct {
	Path      *string         `json:"p,omitempty"`
	Operation *string         `json:"o,omitempty"`
	Value     json.RawMessage `json:"v"`
}

// ResponseFragment is one typed content fragment in a response snapshot.
type ResponseFragment struct {
	ID         int64             `json:"id"`
	Type       string            `json:"type"`
	Content    string            `json:"content"`
	References []json.RawMessage `json:"references"`
	StageID    int64             `json:"stage_id"`
}

// ResponseSnapshot is a complete response state delivered by the stream.
type ResponseSnapshot struct {
	MessageID             int64              `json:"message_id"`
	ParentID              int64              `json:"parent_id"`
	Model                 string             `json:"model"`
	Role                  string             `json:"role"`
	ThinkingEnabled       bool               `json:"thinking_enabled"`
	Status                string             `json:"status"`
	QuasiStatus           string             `json:"quasi_status"`
	AccumulatedTokenUsage int64              `json:"accumulated_token_usage"`
	Content               string             `json:"content"`
	ThinkingContent       *string            `json:"thinking_content"`
	Fragments             []ResponseFragment `json:"fragments"`
}

// SnapshotEvent wraps a complete response snapshot.
type SnapshotEvent struct {
	Response *ResponseSnapshot `json:"response"`
}

// StreamState is the accumulated state of a completion stream.
// Content and ThinkingContent must not be copied after their first use.
type StreamState struct {
	RequestMessageID  int64
	ResponseMessageID int64

	Content         strings.Builder
	ThinkingContent strings.Builder

	Status         string
	TokenUsage     int64
	Title          string
	SessionUpdated float64

	Finished bool
	Closed   bool

	lastPath string
	lastOp   string
}

// Apply incorporates event into s.
func (s *StreamState) Apply(event SSEEvent) error {
	switch event.Event {
	case "ready":
		return s.applyReady(event.Data)

	case "update_session":
		return s.applyUpdateSession(event.Data)

	case "title":
		return s.applyTitle(event.Data)

	case "finish":
		s.Finished = true
		return nil

	case "close":
		s.Closed = true
		return nil
	}

	if len(bytes.TrimSpace(event.Data)) == 0 {
		return nil
	}

	var patch PatchEvent
	if err := json.Unmarshal(event.Data, &patch); err != nil {
		return fmt.Errorf(
			"decode stream patch: %w; data=%s",
			err,
			event.Data,
		)
	}

	if len(patch.Value) == 0 {
		return nil
	}

	/*
		Snapshot имеет вид:

		{
		  "v": {
		    "response": {
		      ...
		    }
		  }
		}

		У него отсутствуют и p, и o. Проверять snapshot нужно до
		наследования lastPath/lastOp, иначе он может быть принят
		за продолжение предыдущего patch.
	*/
	if patch.Path == nil && patch.Operation == nil {
		isSnapshot, err := s.tryApplySnapshot(patch.Value)
		if err != nil {
			return err
		}

		if isSnapshot {
			return nil
		}
	}

	path, operation := s.resolvePatchLocation(patch)

	return s.applyPatch(
		path,
		operation,
		patch.Value,
	)
}

func (s *StreamState) applyReady(data []byte) error {
	var value ReadyEvent

	if err := json.Unmarshal(data, &value); err != nil {
		return fmt.Errorf("decode ready: %w", err)
	}

	s.RequestMessageID = value.RequestMessageID
	s.ResponseMessageID = value.ResponseMessageID

	return nil
}

func (s *StreamState) applyUpdateSession(data []byte) error {
	var value UpdateSessionEvent

	if err := json.Unmarshal(data, &value); err != nil {
		return fmt.Errorf(
			"decode update_session: %w",
			err,
		)
	}

	s.SessionUpdated = value.UpdatedAt

	return nil
}

func (s *StreamState) applyTitle(data []byte) error {
	var value TitleEvent

	if err := json.Unmarshal(data, &value); err != nil {
		return fmt.Errorf("decode title: %w", err)
	}

	s.Title = value.Content

	return nil
}

func (s *StreamState) resolvePatchLocation(
	patch PatchEvent,
) (path string, operation string) {
	switch {
	/*
		Пришёл новый путь.

		Если операция отсутствует, она не должна наследовать старый
		APPEND. Для этого сохраняем lastOp как пустую строку.
	*/
	case patch.Path != nil:
		path = *patch.Path
		s.lastPath = path

		if patch.Operation != nil {
			operation = *patch.Operation
		} else {
			operation = ""
		}

		s.lastOp = operation

	/*
		Операция изменилась, а путь остался прежним.
	*/
	case patch.Operation != nil:
		path = s.lastPath
		operation = *patch.Operation
		s.lastOp = operation

	/*
		Сокращённое событие:

			{"v":"текст"}

		Наследует и путь, и операцию.
	*/
	default:
		path = s.lastPath
		operation = s.lastOp
	}

	return path, operation
}

func (s *StreamState) tryApplySnapshot(
	value json.RawMessage,
) (bool, error) {
	var snapshot SnapshotEvent

	if err := json.Unmarshal(value, &snapshot); err != nil {
		// Значение может быть обычной строкой, числом и т. п.
		// В таком случае это не snapshot.
		return false, nil
	}

	if snapshot.Response == nil {
		return false, nil
	}

	s.applyResponseSnapshot(*snapshot.Response)

	return true, nil
}

func (s *StreamState) applyResponseSnapshot(
	response ResponseSnapshot,
) {
	if response.MessageID != 0 {
		s.ResponseMessageID = response.MessageID
	}

	s.Status = response.Status
	s.TokenUsage = response.AccumulatedTokenUsage

	s.Content.Reset()
	s.ThinkingContent.Reset()

	if response.Content != "" {
		s.Content.WriteString(response.Content)
	}

	if response.ThinkingContent != nil {
		s.ThinkingContent.WriteString(
			*response.ThinkingContent,
		)
	}

	for _, fragment := range response.Fragments {
		switch fragment.Type {
		case "RESPONSE":
			s.Content.WriteString(fragment.Content)

		case "THINKING":
			s.ThinkingContent.WriteString(fragment.Content)
		}
	}

	if response.QuasiStatus != "" {
		s.Status = response.QuasiStatus
	}
}

func (s *StreamState) applyPatch(
	path string,
	operation string,
	value json.RawMessage,
) error {
	switch path {
	case "response/content",
		"response/fragments/-1/content":

		return applyStringPatch(
			&s.Content,
			operation,
			value,
		)

	case "response/thinking_content",
		"response/fragments/-1/thinking_content":

		return applyStringPatch(
			&s.ThinkingContent,
			operation,
			value,
		)

	case "response/accumulated_token_usage":
		if err := json.Unmarshal(value, &s.TokenUsage); err != nil {
			return fmt.Errorf(
				"decode token usage: %w",
				err,
			)
		}

	case "response/status",
		"response/quasi_status":

		if err := json.Unmarshal(value, &s.Status); err != nil {
			return fmt.Errorf(
				"decode response status: %w",
				err,
			)
		}

		if s.Status == "FINISHED" {
			s.Finished = true
		}

	case "response":
		if operation != "BATCH" {
			log.Printf(
				"unknown response operation: op=%q value=%s",
				operation,
				value,
			)

			return nil
		}

		if err := s.applyBatch(value); err != nil {
			return err
		}

	default:
		log.Printf(
			"unknown patch: path=%q op=%q value=%s",
			path,
			operation,
			value,
		)
	}

	return nil
}

func (s *StreamState) applyBatch(
	value json.RawMessage,
) error {
	var batch []PatchEvent

	if err := json.Unmarshal(value, &batch); err != nil {
		return fmt.Errorf("decode batch: %w", err)
	}

	for _, item := range batch {
		if len(item.Value) == 0 {
			continue
		}

		itemPath := ""
		if item.Path != nil {
			itemPath = *item.Path
		}

		itemOperation := ""
		if item.Operation != nil {
			itemOperation = *item.Operation
		}

		/*
			Внутри response/BATCH пути относительные:

				accumulated_token_usage
				quasi_status

			Преобразуем их в абсолютные.
		*/
		if itemPath != "" &&
			!strings.HasPrefix(itemPath, "response/") {

			itemPath = "response/" + itemPath
		}

		if err := s.applyPatch(
			itemPath,
			itemOperation,
			item.Value,
		); err != nil {
			return err
		}
	}

	return nil
}

func applyStringPatch(
	builder *strings.Builder,
	operation string,
	value json.RawMessage,
) error {
	var text string

	if err := json.Unmarshal(value, &text); err != nil {
		return fmt.Errorf(
			"decode string patch: %w",
			err,
		)
	}

	switch operation {
	case "", "APPEND":
		builder.WriteString(text)

	case "SET":
		builder.Reset()
		builder.WriteString(text)

	default:
		return fmt.Errorf(
			"unsupported string patch operation %q",
			operation,
		)
	}

	return nil
}