initial commit
This commit is contained in:
@@ -0,0 +1,582 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user