268 lines
9.0 KiB
Go
268 lines
9.0 KiB
Go
package protocol
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRequestMatrix(t *testing.T) {
|
|
chat := []byte(`{"model":"old","max_tokens":100,"stream":true,"messages":[{"role":"system","content":"help"},{"role":"user","content":[{"type":"text","text":"look"},{"type":"image_url","image_url":{"url":"data:image/png;base64,YQ=="}}]},{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"read","arguments":"{\"path\":\"a\"}"}}]},{"role":"tool","tool_call_id":"c1","content":"ok"},{"role":"user","content":"continue"}],"tools":[{"type":"function","function":{"name":"read","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}}],"tool_choice":"auto"}`)
|
|
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
|
src := chat
|
|
var e error
|
|
if from != Chat {
|
|
src, e = ConvertRequest(Chat, from, chat, "new")
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
}
|
|
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
|
t.Run(string(from)+"/"+string(to), func(t *testing.T) {
|
|
b, e := ConvertRequest(from, to, src, "target")
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
m, e := decode(b)
|
|
if e != nil || m["model"] != "target" {
|
|
t.Fatalf("%s %v", b, e)
|
|
}
|
|
r, e := readRequest(to, m)
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
var calls, results, images int
|
|
for _, m := range r.messages {
|
|
for _, p := range m.parts {
|
|
switch p.typ {
|
|
case "call":
|
|
calls++
|
|
if p.id != "c1" {
|
|
t.Fatal(p)
|
|
}
|
|
case "result":
|
|
results++
|
|
case "image":
|
|
images++
|
|
}
|
|
}
|
|
}
|
|
if calls != 1 || results != 1 || images != 1 {
|
|
t.Fatalf("lost content: %s", b)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
func TestStrictRejection(t *testing.T) {
|
|
cases := []struct {
|
|
k Kind
|
|
b string
|
|
to Kind
|
|
}{
|
|
{Chat, `{"messages":[],"response_format":{"type":"json_object"}}`, Responses},
|
|
{Chat, `{"messages":[]}`, Anthropic},
|
|
{Responses, `{"input":[],"previous_response_id":"r"}`, Chat},
|
|
{Anthropic, `{"max_tokens":10,"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"x"}]}]}`, Chat},
|
|
{Chat, `{"messages":[],"n":2}`, Responses},
|
|
{Chat, `{"messages":[],"stop":["x"]}`, Responses},
|
|
{Chat, `{"max_tokens":10,"messages":[],"tools":[{"type":"function","function":{"name":"x","parameters":{},"strict":true}}]}`, Anthropic},
|
|
}
|
|
for _, c := range cases {
|
|
if b, e := ConvertRequest(c.k, c.to, []byte(c.b), "m"); e == nil {
|
|
t.Fatalf("accepted %s => %s", c.b, b)
|
|
}
|
|
}
|
|
}
|
|
func TestResponseMatrix(t *testing.T) {
|
|
chat := []byte(`{"id":"r","model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"working","tool_calls":[{"id":"c","type":"function","function":{"name":"read","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":4,"prompt_tokens_details":{"cached_tokens":3}}}`)
|
|
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
|
src, e := ConvertResponse(Chat, from, chat)
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
|
b, e := ConvertResponse(from, to, src)
|
|
if e != nil {
|
|
t.Fatal(from, to, e)
|
|
}
|
|
m, _ := decode(b)
|
|
r, e := readResponse(to, m)
|
|
if e != nil || len(r.parts) != 2 || r.finish != "tool" {
|
|
t.Fatalf("%s %v", b, e)
|
|
}
|
|
u := ExtractUsage(to, b)
|
|
if u.Input != 12 || u.Output != 4 || u.CacheRead != 3 || !u.Known {
|
|
t.Fatal(u)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
type flushBuffer struct {
|
|
bytes.Buffer
|
|
flushes int
|
|
}
|
|
|
|
func (b *flushBuffer) Flush() { b.flushes++ }
|
|
func chatStream() string {
|
|
return strings.Join([]string{
|
|
`data: {"id":"r","model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}`,
|
|
`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"read","arguments":"{\"x\":"}},{"index":1,"id":"c2","type":"function","function":{"name":"write","arguments":"{"}}]}}]}`,
|
|
`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"}"}},{"index":0,"function":{"arguments":"1}"}}]}}]}`,
|
|
`data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`,
|
|
`data: {"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":7}}`,
|
|
`data: [DONE]`,
|
|
}, "\n\n") + "\n\n"
|
|
}
|
|
func TestStreamMatrix(t *testing.T) {
|
|
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
|
src := chatStream()
|
|
if from != Chat {
|
|
var b bytes.Buffer
|
|
if _, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, from, "m"); e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
src = b.String()
|
|
}
|
|
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
|
t.Run(string(from)+"/"+string(to), func(t *testing.T) {
|
|
var b flushBuffer
|
|
u, e := Stream(context.Background(), &b, strings.NewReader(src), from, to, "m")
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
if u.Input != 10 || u.Output != 7 || !u.Known {
|
|
t.Fatal(u)
|
|
}
|
|
if b.flushes < 3 {
|
|
t.Fatal("not flushed")
|
|
}
|
|
if from == to && b.String() != src {
|
|
t.Fatal("passthrough changed bytes")
|
|
}
|
|
if to == Responses {
|
|
for _, event := range []string{"response.output_item.added", "response.function_call_arguments.delta", "response.function_call_arguments.done", "response.output_item.done", "response.completed"} {
|
|
if !strings.Contains(b.String(), event) {
|
|
t.Fatal("missing", event)
|
|
}
|
|
}
|
|
}
|
|
if to == Anthropic && !strings.Contains(b.String(), "message_stop") {
|
|
t.Fatal("missing stop")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
func TestUsageCumulative(t *testing.T) {
|
|
src := "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":5,\"output_tokens\":1,\"cache_read_input_tokens\":8,\"cache_creation_input_tokens\":2}}}\n\n" + "data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":9}}\n\n"
|
|
var b bytes.Buffer
|
|
u, e := Stream(context.Background(), &b, strings.NewReader(src), Anthropic, Anthropic, "")
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
if u.Input != 15 || u.Output != 9 || u.CacheRead != 8 || u.CacheWrite != 2 {
|
|
t.Fatal(u)
|
|
}
|
|
raw, e := decode(u.Raw)
|
|
if e != nil || num(raw["input_tokens"]) != 5 || num(raw["output_tokens"]) != 9 || num(raw["cache_read_input_tokens"]) != 8 {
|
|
t.Fatal(string(u.Raw), e)
|
|
}
|
|
}
|
|
func TestRealtimeBeforeEOF(t *testing.T) {
|
|
r, w := io.Pipe()
|
|
defer r.Close()
|
|
defer w.Close()
|
|
done := make(chan error, 1)
|
|
out := &signalWriter{signal: make(chan struct{}, 1)}
|
|
go func() { _, e := Stream(context.Background(), out, r, Chat, Responses, "m"); done <- e }()
|
|
first, rest, _ := strings.Cut(chatStream(), "\n\n")
|
|
if _, e := io.WriteString(w, first+"\n\n"); e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
select {
|
|
case <-out.signal:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("first delta buffered until EOF")
|
|
}
|
|
if _, e := io.WriteString(w, rest); e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
select {
|
|
case e := <-done:
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
case <-time.After(time.Second):
|
|
t.Fatal("terminal marker waits for EOF")
|
|
}
|
|
}
|
|
|
|
type signalWriter struct{ signal chan struct{} }
|
|
|
|
func (w *signalWriter) Write(p []byte) (int, error) {
|
|
select {
|
|
case w.signal <- struct{}{}:
|
|
default:
|
|
}
|
|
return len(p), nil
|
|
}
|
|
func TestUnknownUsageAndUniqueID(t *testing.T) {
|
|
src := "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"a\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"
|
|
var previous string
|
|
for i := 0; i < 2; i++ {
|
|
var b bytes.Buffer
|
|
u, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Anthropic, "m")
|
|
if e != nil || u.Known {
|
|
t.Fatal(u, e)
|
|
}
|
|
line := strings.Split(b.String(), "\n")[1]
|
|
m, e := decode([]byte(strings.TrimPrefix(line, "data: ")))
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
msg := obj(m["message"])
|
|
if obj(msg["usage"]) == nil {
|
|
t.Fatal("missing initial usage")
|
|
}
|
|
id := str(msg["id"])
|
|
if id == "" || id == previous {
|
|
t.Fatal("nonunique id")
|
|
}
|
|
previous = id
|
|
}
|
|
}
|
|
func TestStreamErrors(t *testing.T) {
|
|
for _, src := range []string{"data: nope\n\n", "data: {\"error\":{\"message\":\"bad\"}}\n\n", strings.Replace(chatStream(), "data: [DONE]\n\n", "", 1), strings.Replace(chatStream(), `"arguments":"1}"`, `"arguments":"oops}"`, 1)} {
|
|
var b bytes.Buffer
|
|
if _, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Responses, "m"); e == nil {
|
|
t.Fatal("accepted broken stream", src)
|
|
}
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if _, e := Stream(ctx, io.Discard, strings.NewReader(chatStream()), Chat, Chat, ""); !errors.Is(e, context.Canceled) {
|
|
t.Fatal(e)
|
|
}
|
|
if _, e := Stream(context.Background(), badWriter{}, strings.NewReader(chatStream()), Chat, Anthropic, ""); !errors.Is(e, io.ErrClosedPipe) {
|
|
t.Fatal(e)
|
|
}
|
|
}
|
|
|
|
type badWriter struct{}
|
|
|
|
func (badWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe }
|
|
func TestSSEMultilineAndCRLF(t *testing.T) {
|
|
src := ": comment\r\n\r\ndata: {\r\ndata: \"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3}}\r\n\r\ndata: [DONE]\r\n\r\n"
|
|
var b flushBuffer
|
|
u, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Chat, "")
|
|
if e != nil || u.Input != 2 || u.Output != 3 || b.String() != src {
|
|
t.Fatal(u, e, b.String())
|
|
}
|
|
}
|