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