llmgw/internal/router/cachestream.go

package router

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
)

// A cached answer is kept as one JSON object whichever way it was asked
// for, so a client that streams and one that does not share it. A streamed
// answer is put back together into the object a model would have sent
// whole, and an answer is taken apart into chunks for a client that
// streams.

// streamChunk is one streamed chunk of a chat or a completion.
type streamChunk struct {
	ID                string `json:"id"`
	Object            string `json:"object"`
	Created           int64  `json:"created"`
	Model             string `json:"model"`
	SystemFingerprint string `json:"system_fingerprint"`
	Choices           []struct {
		Index int `json:"index"`
		// Delta is a chat chunk's, as fields: those not understood here
		// leave the answer unkept rather than kept without them.
		Delta        map[string]json.RawMessage `json:"delta"`
		Text         *string                    `json:"text"`
		FinishReason *string                    `json:"finish_reason"`
		Logprobs     json.RawMessage            `json:"logprobs"`
	} `json:"choices"`
}

// toolCallDelta is part of one tool call, as streamed.
type toolCallDelta struct {
	Index    int    `json:"index"`
	ID       string `json:"id"`
	Type     string `json:"type"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `json:"function"`
}

type toolCall struct {
	ID       string `json:"id"`
	Type     string `json:"type"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `json:"function"`
}

// assembledChoice is one choice being put back together.
type assembledChoice struct {
	role                     string
	content, reasoning, text strings.Builder
	hasContent, hasReasoning bool
	tools                    []*toolCall
	finish                   *string
}

// assembleStream puts a streamed chat or completion back together as the
// object the model would have answered without streaming, with usage as the
// engine reported it. ok is false for anything it does not fully
// understand — logprobs, audio, a chunk that is not JSON — which is then
// not kept: an answer kept with something missing would be served short.
func assembleStream(path string, sse []byte, usage *proxy.Usage) ([]byte, bool) {
	chat := path == "/v1/chat/completions"
	if !chat && path != "/v1/completions" {
		return nil, false
	}
	var head streamChunk
	choices := map[int]*assembledChoice{}
	maxIndex := -1
	done := false
	sc := bufio.NewScanner(bytes.NewReader(sse))
	sc.Buffer(make([]byte, 64<<10), len(sse)+1)
	for sc.Scan() {
		data, ok := strings.CutPrefix(sc.Text(), "data:")
		if !ok {
			continue
		}
		data = strings.TrimSpace(data)
		if data == "[DONE]" {
			done = true
			continue
		}
		var c streamChunk
		if json.Unmarshal([]byte(data), &c) != nil {
			return nil, false
		}
		if head.ID == "" {
			head.ID, head.Created, head.Model = c.ID, c.Created, c.Model
		}
		if c.SystemFingerprint != "" {
			head.SystemFingerprint = c.SystemFingerprint
		}
		for _, ch := range c.Choices {
			if len(ch.Logprobs) > 0 && string(ch.Logprobs) != "null" {
				return nil, false
			}
			a := choices[ch.Index]
			if a == nil {
				a = &assembledChoice{}
				choices[ch.Index] = a
				maxIndex = max(maxIndex, ch.Index)
			}
			if ch.FinishReason != nil {
				a.finish = ch.FinishReason
			}
			if ch.Text != nil {
				a.text.WriteString(*ch.Text)
			}
			if !addDelta(a, ch.Delta) {
				return nil, false
			}
		}
	}
	if sc.Err() != nil || !done || len(choices) == 0 || len(choices) != maxIndex+1 {
		// A stream that never said it was done may have been cut short.
		return nil, false
	}
	out := map[string]any{"id": head.ID, "created": head.Created, "model": head.Model}
	if head.SystemFingerprint != "" {
		out["system_fingerprint"] = head.SystemFingerprint
	}
	list := make([]map[string]any, 0, len(choices))
	for i := 0; i <= maxIndex; i++ {
		a := choices[i]
		if a.finish == nil {
			return nil, false
		}
		c := map[string]any{"index": i, "finish_reason": *a.finish, "logprobs": nil}
		if chat {
			role := a.role
			if role == "" {
				role = "assistant"
			}
			msg := map[string]any{"role": role, "content": nil}
			if a.hasContent || len(a.tools) == 0 {
				msg["content"] = a.content.String()
			}
			if a.hasReasoning {
				msg["reasoning_content"] = a.reasoning.String()
			}
			if len(a.tools) > 0 {
				msg["tool_calls"] = a.tools
			}
			c["message"] = msg
		} else {
			c["text"] = a.text.String()
		}
		list = append(list, c)
	}
	out["choices"] = list
	if chat {
		out["object"] = "chat.completion"
	} else {
		out["object"] = "text_completion"
	}
	if usage != nil {
		out["usage"] = map[string]int{"prompt_tokens": usage.PromptTokens,
			"completion_tokens": usage.CompletionTokens,
			"total_tokens":      usage.PromptTokens + usage.CompletionTokens}
	}
	b, err := json.Marshal(out)
	return b, err == nil
}

// addDelta adds one chat chunk's delta to its choice, and says whether it
// understood every field of it.
func addDelta(a *assembledChoice, delta map[string]json.RawMessage) bool {
	for k, raw := range delta {
		if string(raw) == "null" {
			continue
		}
		switch k {
		case "role":
			if json.Unmarshal(raw, &a.role) != nil {
				return false
			}
		case "content", "reasoning_content":
			var s string
			if json.Unmarshal(raw, &s) != nil {
				return false
			}
			// Empty says nothing: the first chunk carries the role with an
			// empty content even when the answer is a tool call.
			if s == "" {
				continue
			}
			if k == "content" {
				a.content.WriteString(s)
				a.hasContent = true
			} else {
				a.reasoning.WriteString(s)
				a.hasReasoning = true
			}
		case "refusal":
			// An empty refusal is a field some engines always send.
			var s string
			if json.Unmarshal(raw, &s) != nil || s != "" {
				return false
			}
		case "tool_calls":
			var parts []toolCallDelta
			if json.Unmarshal(raw, &parts) != nil {
				return false
			}
			for _, p := range parts {
				for len(a.tools) <= p.Index {
					a.tools = append(a.tools, &toolCall{})
				}
				t := a.tools[p.Index]
				if p.ID != "" {
					t.ID = p.ID
				}
				if p.Type != "" {
					t.Type = p.Type
				}
				t.Function.Name += p.Function.Name
				t.Function.Arguments += p.Function.Arguments
			}
		default:
			return false
		}
	}
	return true
}

// streamFromAnswer is a kept chat or completion as the chunks a model would
// have streamed: per choice, its whole text in one chunk and then its
// finish reason, and the usage last for a client that asks for it.
func streamFromAnswer(path string, body []byte, withUsage bool) ([]string, error) {
	var a struct {
		ID                string          `json:"id"`
		Created           int64           `json:"created"`
		Model             string          `json:"model"`
		SystemFingerprint string          `json:"system_fingerprint"`
		Usage             json.RawMessage `json:"usage"`
		Choices           []struct {
			Index        int             `json:"index"`
			FinishReason json.RawMessage `json:"finish_reason"`
			Logprobs     json.RawMessage `json:"logprobs"`
			Text         *string         `json:"text"`
			Message      *struct {
				Role             string            `json:"role"`
				Content          *string           `json:"content"`
				ReasoningContent *string           `json:"reasoning_content"`
				ToolCalls        []json.RawMessage `json:"tool_calls"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(body, &a); err != nil {
		return nil, err
	}
	if len(a.Choices) == 0 {
		return nil, errors.New("an answer with no choices")
	}
	chat := path == "/v1/chat/completions"
	object := "text_completion"
	if chat {
		object = "chat.completion.chunk"
	}
	chunk := func(choices []map[string]any) (string, error) {
		c := map[string]any{"id": a.ID, "object": object, "created": a.Created, "model": a.Model,
			"choices": choices}
		if a.SystemFingerprint != "" {
			c["system_fingerprint"] = a.SystemFingerprint
		}
		b, err := json.Marshal(c)
		return string(b), err
	}
	var events []string
	add := func(choices []map[string]any) error {
		e, err := chunk(choices)
		events = append(events, e)
		return err
	}
	for _, ch := range a.Choices {
		var logprobs any
		if len(ch.Logprobs) > 0 {
			logprobs = ch.Logprobs
		}
		finish := any(json.RawMessage("null"))
		if len(ch.FinishReason) > 0 {
			finish = ch.FinishReason
		}
		if !chat {
			if ch.Text == nil {
				return nil, errors.New("a completion with no text")
			}
			if err := add([]map[string]any{{"index": ch.Index, "text": *ch.Text, "logprobs": logprobs,
				"finish_reason": nil}}); err != nil {
				return nil, err
			}
			if err := add([]map[string]any{{"index": ch.Index, "text": "", "finish_reason": finish}}); err != nil {
				return nil, err
			}
			continue
		}
		m := ch.Message
		if m == nil {
			return nil, errors.New("a chat answer with no message")
		}
		role := m.Role
		if role == "" {
			role = "assistant"
		}
		if err := add([]map[string]any{{"index": ch.Index, "delta": map[string]any{"role": role, "content": ""},
			"finish_reason": nil}}); err != nil {
			return nil, err
		}
		if m.ReasoningContent != nil && *m.ReasoningContent != "" {
			if err := add([]map[string]any{{"index": ch.Index,
				"delta": map[string]any{"reasoning_content": *m.ReasoningContent}, "finish_reason": nil}}); err != nil {
				return nil, err
			}
		}
		if m.Content != nil && *m.Content != "" || logprobs != nil {
			content := ""
			if m.Content != nil {
				content = *m.Content
			}
			if err := add([]map[string]any{{"index": ch.Index, "delta": map[string]any{"content": content},
				"logprobs": logprobs, "finish_reason": nil}}); err != nil {
				return nil, err
			}
		}
		if len(m.ToolCalls) > 0 {
			calls := make([]map[string]any, 0, len(m.ToolCalls))
			for i, raw := range m.ToolCalls {
				var tc map[string]any
				if err := json.Unmarshal(raw, &tc); err != nil {
					return nil, err
				}
				tc["index"] = i
				calls = append(calls, tc)
			}
			if err := add([]map[string]any{{"index": ch.Index, "delta": map[string]any{"tool_calls": calls},
				"finish_reason": nil}}); err != nil {
				return nil, err
			}
		}
		if err := add([]map[string]any{{"index": ch.Index, "delta": map[string]any{},
			"finish_reason": finish}}); err != nil {
			return nil, err
		}
	}
	if withUsage && len(a.Usage) > 0 && string(a.Usage) != "null" {
		c := map[string]any{"id": a.ID, "object": object, "created": a.Created, "model": a.Model,
			"choices": []any{}, "usage": a.Usage}
		b, err := json.Marshal(c)
		if err != nil {
			return nil, err
		}
		events = append(events, string(b))
	}
	return events, nil
}