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