880 lines
23 KiB
Go
880 lines
23 KiB
Go
package protocol
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
const maxStreamBytes = 16 << 20
|
|
const maxStreamBlocks = 4096
|
|
const maxStreamUsageBytes = 16 << 10
|
|
const maxPassthroughBytes = 64 << 20
|
|
|
|
type streamBlock struct {
|
|
part
|
|
content, identity, functionName strings.Builder
|
|
index int
|
|
closed bool
|
|
emitted bool
|
|
sent int
|
|
}
|
|
type streamState struct {
|
|
thinking map[int]*thinkingBlock
|
|
ctx context.Context
|
|
w io.Writer
|
|
from, to Kind
|
|
r response
|
|
blocks []*streamBlock
|
|
source map[int]*streamBlock
|
|
reasoning map[int]*reasoningItem
|
|
started, finished, terminal bool
|
|
sequence int
|
|
retained, blockCount int
|
|
byteLimit, blockLimit int
|
|
}
|
|
|
|
func (s *streamState) retain(n int) error {
|
|
if n > s.byteLimit-s.retained {
|
|
return fmt.Errorf("stream retained content exceeds %d bytes", s.byteLimit)
|
|
}
|
|
s.retained += n
|
|
return nil
|
|
}
|
|
func (s *streamState) reserveBlock() error {
|
|
if s.blockCount >= s.blockLimit {
|
|
return fmt.Errorf("stream exceeds %d blocks", s.blockLimit)
|
|
}
|
|
s.blockCount++
|
|
return nil
|
|
}
|
|
|
|
func writeFlush(w io.Writer, b []byte) error {
|
|
n, e := w.Write(b)
|
|
if e != nil {
|
|
return e
|
|
}
|
|
if n != len(b) {
|
|
return io.ErrShortWrite
|
|
}
|
|
if f, ok := w.(interface{ Flush() error }); ok {
|
|
return f.Flush()
|
|
}
|
|
if f, ok := w.(interface{ Flush() }); ok {
|
|
f.Flush()
|
|
}
|
|
return nil
|
|
}
|
|
func (s *streamState) event(name string, m object) error {
|
|
if e := s.ctx.Err(); e != nil {
|
|
return e
|
|
}
|
|
if s.to == Responses {
|
|
m["sequence_number"] = s.sequence
|
|
s.sequence++
|
|
}
|
|
b, e := json.Marshal(m)
|
|
if e != nil {
|
|
return e
|
|
}
|
|
prefix := ""
|
|
if name != "" {
|
|
prefix = "event: " + name + "\n"
|
|
}
|
|
return writeFlush(s.w, []byte(prefix+"data: "+string(b)+"\n\n"))
|
|
}
|
|
func (s *streamState) typed(name string, m object) error { m["type"] = name; return s.event(name, m) }
|
|
func (s *streamState) chat(delta object, finish any) error {
|
|
return s.event("", object{"id": s.r.id, "object": "chat.completion.chunk", "created": 0, "model": s.r.model, "choices": []any{object{"index": 0, "delta": delta, "finish_reason": finish}}})
|
|
}
|
|
func (s *streamState) start() error {
|
|
if s.started {
|
|
return nil
|
|
}
|
|
s.started = true
|
|
if s.r.id == "" {
|
|
s.r.id = newID()
|
|
}
|
|
switch s.to {
|
|
case Chat:
|
|
return s.chat(object{"role": "assistant", "content": ""}, nil)
|
|
case Anthropic:
|
|
m := object{"id": s.r.id, "type": "message", "role": "assistant", "model": s.r.model, "content": []any{}, "stop_reason": nil, "stop_sequence": nil}
|
|
m["usage"] = usageJSON(Anthropic, s.r.usage)
|
|
return s.typed("message_start", object{"message": m})
|
|
case Responses:
|
|
r := object{"id": s.r.id, "object": "response", "created_at": time.Now().Unix(), "model": s.r.model, "status": "in_progress", "output": []any{}, "error": nil, "incomplete_details": nil, "usage": nil}
|
|
if e := s.typed("response.created", object{"response": r}); e != nil {
|
|
return e
|
|
}
|
|
return s.typed("response.in_progress", object{"response": r})
|
|
}
|
|
return nil
|
|
}
|
|
func (s *streamState) itemID(b *streamBlock) string {
|
|
if b.typ == "call" {
|
|
return fmt.Sprintf("fc_%s_%d", s.r.id, b.index)
|
|
}
|
|
return fmt.Sprintf("msg_%s_%d", s.r.id, b.index)
|
|
}
|
|
func (s *streamState) item(b *streamBlock, done bool) object {
|
|
status := "in_progress"
|
|
if done {
|
|
status = "completed"
|
|
}
|
|
if b.typ == "call" {
|
|
args := ""
|
|
if done {
|
|
args = b.args
|
|
}
|
|
return object{"id": s.itemID(b), "type": "function_call", "call_id": b.id, "name": b.name, "arguments": args, "status": status}
|
|
}
|
|
content := []any{}
|
|
if done {
|
|
content = append(content, object{"type": "output_text", "text": b.text, "annotations": []any{}})
|
|
}
|
|
return object{"id": s.itemID(b), "type": "message", "role": "assistant", "status": status, "content": content}
|
|
}
|
|
func (s *streamState) open(b *streamBlock) error {
|
|
if b.emitted {
|
|
return nil
|
|
}
|
|
if b.typ == "call" && (b.id == "" || b.name == "") {
|
|
return nil
|
|
}
|
|
if e := s.start(); e != nil {
|
|
return e
|
|
}
|
|
b.emitted = true
|
|
switch s.to {
|
|
case Chat:
|
|
if b.typ == "call" {
|
|
return s.chat(object{"tool_calls": []any{object{"index": s.toolIndex(b), "id": b.id, "type": "function", "function": object{"name": b.name, "arguments": ""}}}}, nil)
|
|
}
|
|
case Anthropic:
|
|
c := object{"type": "text", "text": ""}
|
|
if b.typ == "call" {
|
|
c = object{"type": "tool_use", "id": b.id, "name": b.name, "input": object{}}
|
|
}
|
|
return s.typed("content_block_start", object{"index": b.index, "content_block": c})
|
|
case Responses:
|
|
if e := s.typed("response.output_item.added", object{"output_index": b.index, "item": s.item(b, false)}); e != nil {
|
|
return e
|
|
}
|
|
if b.typ == "text" {
|
|
return s.typed("response.content_part.added", object{"output_index": b.index, "item_id": s.itemID(b), "content_index": 0, "part": object{"type": "output_text", "text": "", "annotations": []any{}}})
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
func (s *streamState) toolIndex(b *streamBlock) int {
|
|
i := 0
|
|
for _, x := range s.blocks {
|
|
if x == b {
|
|
return i
|
|
}
|
|
if x.typ == "call" {
|
|
i++
|
|
}
|
|
}
|
|
return i
|
|
}
|
|
func (s *streamState) add(key int, p part) (*streamBlock, error) {
|
|
if s.finished {
|
|
return nil, fmt.Errorf("content after finish")
|
|
}
|
|
if _, ok := s.source[key]; ok {
|
|
return nil, fmt.Errorf("duplicate content index %d", key)
|
|
}
|
|
if e := s.reserveBlock(); e != nil {
|
|
return nil, e
|
|
}
|
|
if e := s.retain(len(p.text) + len(p.args) + len(p.id) + len(p.name)); e != nil {
|
|
return nil, e
|
|
}
|
|
b := &streamBlock{part: p, index: len(s.blocks)}
|
|
b.content.WriteString(p.text)
|
|
b.content.WriteString(p.args)
|
|
b.identity.WriteString(p.id)
|
|
b.functionName.WriteString(p.name)
|
|
s.blocks = append(s.blocks, b)
|
|
s.source[key] = b
|
|
return b, s.open(b)
|
|
}
|
|
func (s *streamState) delta(b *streamBlock, text string) error {
|
|
if b == nil || b.closed || s.finished {
|
|
return fmt.Errorf("delta outside active content block")
|
|
}
|
|
if e := s.retain(len(text)); e != nil {
|
|
return e
|
|
}
|
|
b.content.WriteString(text)
|
|
// Builder.String shares storage without copying accumulated content.
|
|
if b.typ == "text" {
|
|
b.text = b.content.String()
|
|
} else {
|
|
b.args = b.content.String()
|
|
}
|
|
if e := s.open(b); e != nil {
|
|
return e
|
|
}
|
|
if !b.emitted {
|
|
return nil
|
|
}
|
|
whole := b.text
|
|
if b.typ == "call" {
|
|
whole = b.args
|
|
}
|
|
text = whole[b.sent:]
|
|
b.sent = len(whole)
|
|
if text == "" {
|
|
return nil
|
|
}
|
|
switch s.to {
|
|
case Chat:
|
|
if b.typ == "text" {
|
|
return s.chat(object{"content": text}, nil)
|
|
}
|
|
return s.chat(object{"tool_calls": []any{object{"index": s.toolIndex(b), "function": object{"arguments": text}}}}, nil)
|
|
case Anthropic:
|
|
d := object{"type": "text_delta", "text": text}
|
|
if b.typ == "call" {
|
|
d = object{"type": "input_json_delta", "partial_json": text}
|
|
}
|
|
return s.typed("content_block_delta", object{"index": b.index, "delta": d})
|
|
case Responses:
|
|
name := "response.output_text.delta"
|
|
m := object{"output_index": b.index, "item_id": s.itemID(b), "delta": text}
|
|
if b.typ == "call" {
|
|
name = "response.function_call_arguments.delta"
|
|
} else {
|
|
m["content_index"] = 0
|
|
}
|
|
return s.typed(name, m)
|
|
}
|
|
return nil
|
|
}
|
|
func (s *streamState) close(b *streamBlock) error {
|
|
if b == nil {
|
|
return fmt.Errorf("unknown content block")
|
|
}
|
|
if b.closed {
|
|
return nil
|
|
}
|
|
if b.typ == "call" {
|
|
if b.id == "" || b.name == "" {
|
|
return fmt.Errorf("incomplete tool identity")
|
|
}
|
|
if b.args == "" {
|
|
if e := s.delta(b, "{}"); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
if _, e := decode([]byte(b.args)); e != nil {
|
|
return fmt.Errorf("invalid streamed tool arguments: %w", e)
|
|
}
|
|
}
|
|
if e := s.open(b); e != nil {
|
|
return e
|
|
}
|
|
if e := s.delta(b, ""); e != nil {
|
|
return e
|
|
}
|
|
b.closed = true
|
|
switch s.to {
|
|
case Anthropic:
|
|
return s.typed("content_block_stop", object{"index": b.index})
|
|
case Responses:
|
|
m := object{"output_index": b.index, "item_id": s.itemID(b)}
|
|
name := "response.function_call_arguments.done"
|
|
if b.typ == "call" {
|
|
m["arguments"] = b.args
|
|
} else {
|
|
name = "response.output_text.done"
|
|
m["text"] = b.text
|
|
m["content_index"] = 0
|
|
}
|
|
if e := s.typed(name, m); e != nil {
|
|
return e
|
|
}
|
|
if b.typ == "text" {
|
|
if e := s.typed("response.content_part.done", object{"output_index": b.index, "item_id": s.itemID(b), "content_index": 0, "part": object{"type": "output_text", "text": b.text, "annotations": []any{}}}); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
return s.typed("response.output_item.done", object{"output_index": b.index, "item": s.item(b, true)})
|
|
}
|
|
return nil
|
|
}
|
|
func (s *streamState) finish(reason string) error {
|
|
if s.finished {
|
|
return fmt.Errorf("duplicate finish")
|
|
}
|
|
if e := s.start(); e != nil {
|
|
return e
|
|
}
|
|
for _, b := range s.blocks {
|
|
if e := s.close(b); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
s.r.finish = reason
|
|
s.finished = true
|
|
return nil
|
|
}
|
|
func (s *streamState) end() error {
|
|
if !s.finished {
|
|
return fmt.Errorf("stream ended without finish reason")
|
|
}
|
|
switch s.to {
|
|
case Chat:
|
|
if e := s.chat(object{}, wireFinish(Chat, s.r.finish)); e != nil {
|
|
return e
|
|
}
|
|
if s.r.usage.Known {
|
|
if e := s.event("", object{"id": s.r.id, "object": "chat.completion.chunk", "created": 0, "model": s.r.model, "choices": []any{}, "usage": usageJSON(Chat, s.r.usage)}); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
return writeFlush(s.w, []byte("data: [DONE]\n\n"))
|
|
case Anthropic:
|
|
m := object{"delta": object{"stop_reason": wireFinish(Anthropic, s.r.finish), "stop_sequence": nil}}
|
|
if s.r.usage.Known {
|
|
m["usage"] = usageJSON(Anthropic, s.r.usage)
|
|
}
|
|
if e := s.typed("message_delta", m); e != nil {
|
|
return e
|
|
}
|
|
return s.typed("message_stop", object{})
|
|
case Responses:
|
|
r := s.r
|
|
for _, b := range s.blocks {
|
|
r.parts = append(r.parts, b.part)
|
|
}
|
|
m, e := writeResponse(Responses, r)
|
|
if e != nil {
|
|
return e
|
|
}
|
|
items := []any{}
|
|
for _, b := range s.blocks {
|
|
items = append(items, s.item(b, true))
|
|
}
|
|
m["output"] = items
|
|
name := "response.completed"
|
|
if r.finish == "length" {
|
|
name = "response.incomplete"
|
|
}
|
|
return s.typed(name, object{"response": m})
|
|
}
|
|
return nil
|
|
}
|
|
func mergeUsage(k Kind, old Usage, m object) (Usage, error) {
|
|
if r := obj(m["response"]); r != nil {
|
|
m = r
|
|
}
|
|
if r := obj(m["message"]); r != nil {
|
|
m = r
|
|
}
|
|
u := obj(m["usage"])
|
|
if u == nil {
|
|
return old, nil
|
|
}
|
|
raw, err := json.Marshal(u)
|
|
if err != nil {
|
|
return old, err
|
|
}
|
|
if len(raw) > maxStreamUsageBytes {
|
|
return old, fmt.Errorf("stream usage exceeds %d bytes", maxStreamUsageBytes)
|
|
}
|
|
if len(old.Raw) > 0 {
|
|
prior, e := decode(old.Raw)
|
|
if e != nil {
|
|
return old, e
|
|
}
|
|
for key, v := range u {
|
|
prior[key] = v
|
|
}
|
|
u = prior
|
|
}
|
|
n := usage(k, u)
|
|
if len(n.Raw) > maxStreamUsageBytes {
|
|
return old, fmt.Errorf("stream usage exceeds %d bytes", maxStreamUsageBytes)
|
|
}
|
|
return n, nil
|
|
}
|
|
func streamError(m object) error {
|
|
if m["error"] != nil {
|
|
return fmt.Errorf("upstream stream error: %v", m["error"])
|
|
}
|
|
switch str(m["type"]) {
|
|
case "error", "response.failed":
|
|
return fmt.Errorf("upstream stream error: %v", m)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *streamState) consume(m object, event string) error {
|
|
if e := streamError(m); e != nil {
|
|
return e
|
|
}
|
|
u, err := mergeUsage(s.from, s.r.usage, m)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
s.r.usage = u
|
|
if s.from == Chat {
|
|
if id := str(m["id"]); id != "" && !s.started {
|
|
s.r.id = id
|
|
}
|
|
if s.r.model == "" {
|
|
s.r.model = str(m["model"])
|
|
}
|
|
cs := arr(m["choices"])
|
|
if len(cs) > 1 {
|
|
return fmt.Errorf("multiple streaming choices unsupported")
|
|
}
|
|
for _, v := range cs {
|
|
c := obj(v)
|
|
if num(c["index"]) != 0 {
|
|
return fmt.Errorf("nonzero choice index unsupported")
|
|
}
|
|
if c["logprobs"] != nil {
|
|
return fmt.Errorf("logprobs unsupported")
|
|
}
|
|
d, e := chatOutput(obj(c["delta"]))
|
|
if e != nil {
|
|
return e
|
|
}
|
|
if e := fields(d, "role", "content", "tool_calls"); e != nil {
|
|
return e
|
|
}
|
|
if role := str(d["role"]); role != "" && role != "assistant" {
|
|
return fmt.Errorf("invalid delta role")
|
|
}
|
|
if e := s.start(); e != nil {
|
|
return e
|
|
}
|
|
if text, ok := d["content"]; ok && text != nil {
|
|
t, ok := text.(string)
|
|
if !ok {
|
|
return fmt.Errorf("non-text streamed content")
|
|
}
|
|
b := s.source[-1]
|
|
if b == nil {
|
|
var e error
|
|
b, e = s.add(-1, part{typ: "text"})
|
|
if e != nil {
|
|
return e
|
|
}
|
|
}
|
|
if e := s.delta(b, t); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
for _, v := range arr(d["tool_calls"]) {
|
|
t := obj(v)
|
|
if e := fields(t, "index", "id", "type", "function"); e != nil {
|
|
return e
|
|
}
|
|
idx, ok := t["index"]
|
|
if !ok || num(idx) < 0 {
|
|
return fmt.Errorf("invalid tool index")
|
|
}
|
|
if typ := str(t["type"]); typ != "" && typ != "function" {
|
|
return fmt.Errorf("unsupported streamed tool")
|
|
}
|
|
f := obj(t["function"])
|
|
if e := fields(f, "name", "arguments"); e != nil {
|
|
return e
|
|
}
|
|
key := int(num(idx))
|
|
b := s.source[key]
|
|
if b == nil {
|
|
var e error
|
|
b, e = s.add(key, part{typ: "call"})
|
|
if e != nil {
|
|
return e
|
|
}
|
|
}
|
|
if b.emitted && (str(t["id"]) != "" || str(f["name"]) != "") {
|
|
return fmt.Errorf("fragmented tool identity after arguments unsupported")
|
|
}
|
|
if e := s.retain(len(str(t["id"])) + len(str(f["name"]))); e != nil {
|
|
return e
|
|
}
|
|
b.identity.WriteString(str(t["id"]))
|
|
b.functionName.WriteString(str(f["name"]))
|
|
b.id = b.identity.String()
|
|
b.name = b.functionName.String()
|
|
if e := s.delta(b, str(f["arguments"])); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
if reason := str(c["finish_reason"]); reason != "" {
|
|
f, e := finishReason(Chat, reason)
|
|
if e != nil {
|
|
return e
|
|
}
|
|
if e = s.finish(f); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
typ := str(m["type"])
|
|
if typ == "" {
|
|
typ = event
|
|
}
|
|
if event != "" && typ != event {
|
|
return fmt.Errorf("SSE event/type mismatch")
|
|
}
|
|
// Discarded blocks still retain lifecycle bookkeeping in separate maps.
|
|
if (s.from == Anthropic && typ == "content_block_start" && (str(obj(m["content_block"])["type"]) == "thinking" || str(obj(m["content_block"])["type"]) == "redacted_thinking")) ||
|
|
(s.from == Responses && typ == "response.output_item.added" && str(obj(m["item"])["type"]) == "reasoning") ||
|
|
(s.from == Responses && (typ == "response.content_part.added" || typ == "response.reasoning_summary_part.added") && s.reasoning[int(num(m["output_index"]))] != nil) {
|
|
if e := s.reserveBlock(); e != nil {
|
|
return e
|
|
}
|
|
if typ == "response.output_item.added" {
|
|
if e := s.retain(len(str(obj(m["item"])["id"]))); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
}
|
|
if s.from == Anthropic {
|
|
if handled, err := s.thinkingEvent(typ, m); handled || err != nil {
|
|
return err
|
|
}
|
|
switch typ {
|
|
case "ping":
|
|
return nil
|
|
case "message_start":
|
|
if s.started {
|
|
return fmt.Errorf("duplicate message_start")
|
|
}
|
|
x := obj(m["message"])
|
|
s.r.id = str(x["id"])
|
|
if s.r.model == "" {
|
|
s.r.model = str(x["model"])
|
|
}
|
|
if len(arr(x["content"])) != 0 {
|
|
return fmt.Errorf("nonempty message_start content")
|
|
}
|
|
return s.start()
|
|
case "content_block_start":
|
|
x := obj(m["content_block"])
|
|
p := part{typ: "text"}
|
|
initial := ""
|
|
switch str(x["type"]) {
|
|
case "text":
|
|
if e := fields(x, "type", "text"); e != nil {
|
|
return e
|
|
}
|
|
initial = str(x["text"])
|
|
case "tool_use":
|
|
if e := fields(x, "type", "id", "name", "input"); e != nil {
|
|
return e
|
|
}
|
|
p = part{typ: "call", id: str(x["id"]), name: str(x["name"])}
|
|
if obj(x["input"]) == nil {
|
|
return fmt.Errorf("tool input must be an object")
|
|
}
|
|
if len(obj(x["input"])) > 0 {
|
|
b, e := json.Marshal(x["input"])
|
|
if e != nil {
|
|
return e
|
|
}
|
|
initial = string(b)
|
|
}
|
|
default:
|
|
return fmt.Errorf("unsupported streamed block %q", str(x["type"]))
|
|
}
|
|
b, e := s.add(int(num(m["index"])), p)
|
|
if e != nil {
|
|
return e
|
|
}
|
|
return s.delta(b, initial)
|
|
case "content_block_delta":
|
|
b := s.source[int(num(m["index"]))]
|
|
d := obj(m["delta"])
|
|
switch str(d["type"]) {
|
|
case "text_delta":
|
|
if b == nil || b.typ != "text" {
|
|
return fmt.Errorf("text delta block mismatch")
|
|
}
|
|
return s.delta(b, str(d["text"]))
|
|
case "input_json_delta":
|
|
if b == nil || b.typ != "call" {
|
|
return fmt.Errorf("tool delta block mismatch")
|
|
}
|
|
return s.delta(b, str(d["partial_json"]))
|
|
default:
|
|
return fmt.Errorf("unsupported delta %q", str(d["type"]))
|
|
}
|
|
case "content_block_stop":
|
|
return s.close(s.source[int(num(m["index"]))])
|
|
case "message_delta":
|
|
d := obj(m["delta"])
|
|
if e := fields(d, "stop_reason", "stop_sequence"); e != nil {
|
|
return e
|
|
}
|
|
if d["stop_sequence"] != nil {
|
|
return fmt.Errorf("stop sequence cannot be preserved")
|
|
}
|
|
f, e := finishReason(Anthropic, str(d["stop_reason"]))
|
|
if e != nil {
|
|
return e
|
|
}
|
|
return s.finish(f)
|
|
case "message_stop":
|
|
s.terminal = true
|
|
return nil
|
|
default:
|
|
return fmt.Errorf("unsupported Anthropic event %q", typ)
|
|
}
|
|
}
|
|
if handled, err := s.reasoningEvent(typ, m); handled || err != nil {
|
|
return err
|
|
}
|
|
switch typ {
|
|
case "response.created", "response.in_progress":
|
|
x := obj(m["response"])
|
|
if !s.started {
|
|
s.r.id = str(x["id"])
|
|
if s.r.model == "" {
|
|
s.r.model = str(x["model"])
|
|
}
|
|
}
|
|
return s.start()
|
|
case "response.output_item.added":
|
|
x := obj(m["item"])
|
|
key := int(num(m["output_index"]))
|
|
p := part{typ: "text"}
|
|
switch str(x["type"]) {
|
|
case "message":
|
|
if len(arr(x["content"])) > 0 {
|
|
return fmt.Errorf("nonempty initial response content")
|
|
}
|
|
case "function_call":
|
|
if _, ok := x["arguments"].(string); !ok {
|
|
return fmt.Errorf("function arguments must be a string")
|
|
}
|
|
p = part{typ: "call", id: str(x["call_id"]), name: str(x["name"])}
|
|
default:
|
|
return fmt.Errorf("unsupported output item %q", str(x["type"]))
|
|
}
|
|
b, e := s.add(key, p)
|
|
if e != nil {
|
|
return e
|
|
}
|
|
if p.typ == "call" {
|
|
return s.delta(b, str(x["arguments"]))
|
|
}
|
|
return nil
|
|
case "response.content_part.added":
|
|
p := obj(m["part"])
|
|
if str(p["type"]) != "output_text" || num(m["content_index"]) != 0 {
|
|
return fmt.Errorf("unsupported response content part")
|
|
}
|
|
if len(arr(p["annotations"])) > 0 {
|
|
return fmt.Errorf("annotations unsupported")
|
|
}
|
|
return s.delta(s.source[int(num(m["output_index"]))], str(p["text"]))
|
|
case "response.output_text.delta", "response.function_call_arguments.delta":
|
|
b := s.source[int(num(m["output_index"]))]
|
|
if b == nil {
|
|
return fmt.Errorf("delta before output item")
|
|
}
|
|
if (typ == "response.output_text.delta") != (b.typ == "text") || num(m["content_index"]) != 0 {
|
|
return fmt.Errorf("response delta block mismatch")
|
|
}
|
|
return s.delta(b, str(m["delta"]))
|
|
case "response.output_text.done", "response.function_call_arguments.done":
|
|
b := s.source[int(num(m["output_index"]))]
|
|
if b == nil {
|
|
return fmt.Errorf("done before output item")
|
|
}
|
|
full := str(m["text"])
|
|
got := b.text
|
|
if b.typ == "call" {
|
|
full = str(m["arguments"])
|
|
got = b.args
|
|
}
|
|
if full != got {
|
|
return fmt.Errorf("final content disagrees with deltas")
|
|
}
|
|
return nil
|
|
case "response.content_part.done":
|
|
p := obj(m["part"])
|
|
if str(p["type"]) != "output_text" || len(arr(p["annotations"])) > 0 || num(m["content_index"]) != 0 {
|
|
return fmt.Errorf("unsupported content part")
|
|
}
|
|
return nil
|
|
case "response.output_item.done":
|
|
return s.close(s.source[int(num(m["output_index"]))])
|
|
case "response.completed", "response.incomplete":
|
|
r, e := readResponse(Responses, obj(m["response"]))
|
|
if e != nil {
|
|
return e
|
|
}
|
|
if len(r.parts) != len(s.blocks) {
|
|
return fmt.Errorf("final response disagrees with streamed items")
|
|
}
|
|
for i, p := range r.parts {
|
|
b := s.blocks[i]
|
|
if p.typ != b.typ || p.text != b.text || p.args != b.args || p.id != b.id || p.name != b.name {
|
|
return fmt.Errorf("final response disagrees with streamed content")
|
|
}
|
|
}
|
|
if e = s.finish(r.finish); e != nil {
|
|
return e
|
|
}
|
|
s.terminal = true
|
|
return nil
|
|
default:
|
|
return fmt.Errorf("unsupported Responses event %q", typ)
|
|
}
|
|
}
|
|
|
|
// Stream does not own r. A caller must cancel or close a blocked upstream read.
|
|
func Stream(ctx context.Context, w io.Writer, r io.Reader, from, to Kind, model string) (Usage, error) {
|
|
return streamWithLimits(ctx, w, r, from, to, model, maxStreamBytes, maxStreamBlocks, maxPassthroughBytes)
|
|
}
|
|
|
|
func readSSELine(ctx context.Context, br *bufio.Reader, remaining int) (string, error) {
|
|
var line strings.Builder
|
|
for {
|
|
if err := ctx.Err(); err != nil {
|
|
return "", err
|
|
}
|
|
fragment, err := br.ReadSlice('\n')
|
|
if len(fragment) > remaining-line.Len() {
|
|
return "", fmt.Errorf("SSE line/event exceeds size limit")
|
|
}
|
|
line.Write(fragment)
|
|
if err != bufio.ErrBufferFull {
|
|
return line.String(), err
|
|
}
|
|
}
|
|
}
|
|
|
|
func streamWithLimits(ctx context.Context, w io.Writer, r io.Reader, from, to Kind, model string, byteLimit, blockLimit, totalLimit int) (Usage, error) {
|
|
if !valid(from) || !valid(to) {
|
|
return Usage{}, fmt.Errorf("unknown protocol")
|
|
}
|
|
s := &streamState{ctx: ctx, w: w, from: from, to: to, r: response{model: model}, source: map[int]*streamBlock{}, byteLimit: byteLimit, blockLimit: blockLimit}
|
|
br := bufio.NewReader(r)
|
|
data := []string{}
|
|
event := ""
|
|
size := 0
|
|
total := 0
|
|
dispatch := func() error {
|
|
if len(data) == 0 {
|
|
event = ""
|
|
return nil
|
|
}
|
|
payload := strings.Join(data, "\n")
|
|
data = nil
|
|
ev := event
|
|
event = ""
|
|
size = 0
|
|
if payload == "[DONE]" {
|
|
if from != Chat {
|
|
return fmt.Errorf("unexpected DONE marker")
|
|
}
|
|
s.terminal = true
|
|
return nil
|
|
}
|
|
m, e := decode([]byte(payload))
|
|
if e != nil {
|
|
return fmt.Errorf("invalid SSE JSON: %w", e)
|
|
}
|
|
if e = streamError(m); e != nil {
|
|
return e
|
|
}
|
|
if from == to {
|
|
u, err := mergeUsage(from, s.r.usage, m)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
s.r.usage = u
|
|
switch str(m["type"]) {
|
|
case "message_stop", "response.completed", "response.incomplete":
|
|
s.terminal = true
|
|
}
|
|
return nil
|
|
}
|
|
if s.terminal {
|
|
return fmt.Errorf("data after terminal event")
|
|
}
|
|
return s.consume(m, ev)
|
|
}
|
|
for {
|
|
if e := ctx.Err(); e != nil {
|
|
return s.r.usage, e
|
|
}
|
|
line, e := readSSELine(ctx, br, byteLimit-size)
|
|
if e != nil && e != io.EOF {
|
|
return s.r.usage, e
|
|
}
|
|
if len(line) > 0 {
|
|
if from == to {
|
|
if total+len(line) > totalLimit {
|
|
return s.r.usage, fmt.Errorf("stream passthrough exceeds %d bytes", totalLimit)
|
|
}
|
|
total += len(line)
|
|
if err := writeFlush(w, []byte(line)); err != nil {
|
|
return s.r.usage, err
|
|
}
|
|
}
|
|
size += len(line)
|
|
line = strings.TrimSuffix(strings.TrimSuffix(line, "\n"), "\r")
|
|
if line == "" {
|
|
if err := dispatch(); err != nil {
|
|
return s.r.usage, err
|
|
}
|
|
size = 0
|
|
if s.terminal {
|
|
break
|
|
}
|
|
} else if !strings.HasPrefix(line, ":") {
|
|
key, val, _ := strings.Cut(line, ":")
|
|
val = strings.TrimPrefix(val, " ")
|
|
switch key {
|
|
case "data":
|
|
data = append(data, val)
|
|
case "event":
|
|
event = val
|
|
}
|
|
}
|
|
}
|
|
if e != nil {
|
|
if e != io.EOF {
|
|
return s.r.usage, e
|
|
}
|
|
if err := dispatch(); err != nil {
|
|
return s.r.usage, err
|
|
}
|
|
break
|
|
}
|
|
}
|
|
if e := ctx.Err(); e != nil {
|
|
return s.r.usage, e
|
|
}
|
|
if from == to {
|
|
return s.r.usage, nil
|
|
}
|
|
if !s.terminal {
|
|
return s.r.usage, fmt.Errorf("truncated SSE stream: missing terminal event")
|
|
}
|
|
if e := s.end(); e != nil {
|
|
return s.r.usage, e
|
|
}
|
|
return s.r.usage, nil
|
|
}
|