1package modes
2
3import (
4 "bytes"
5 "context"
6 "strings"
7 "sync/atomic"
8 "testing"
9
10 "github.com/bnema/zut/packages/core"
11 "github.com/bnema/zut/packages/provider"
12)
13
14// streamContinuationClient emits interim text with StopContinue on the
15// first call and the final answer with StopEnd on the second call.
16type streamContinuationClient struct {
17 calls int32
18}
19
20func (c *streamContinuationClient) Name() string { return "stream-continuation" }
21
22func (c *streamContinuationClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
23 call := atomic.AddInt32(&c.calls, 1)
24 out := make(chan provider.Event, 8)
25 go func() {
26 defer close(out)
27 text, stop := "final", provider.StopEnd
28 if call == 1 {
29 text, stop = "part one", provider.StopContinue
30 }
31 for _, chunk := range strings.Split(text, " ") {
32 select {
33 case <-ctx.Done():
34 return
35 case out <- provider.EventTextDelta{Delta: chunk + " "}:
36 }
37 }
38 select {
39 case <-ctx.Done():
40 return
41 case out <- provider.EventDone{Stop: stop, Message: provider.Message{
42 Role: provider.RoleAssistant,
43 Content: []provider.Content{provider.TextBlock{Text: text}},
44 }}:
45 }
46 }()
47 return out, nil
48}
49
50func TestRunStreamContinuation(t *testing.T) {
51 client := &streamContinuationClient{}
52 ag := core.NewAgent(client, "test-model", "", nil)
53 var out, diag bytes.Buffer
54 if err := RunStreamWithDiag(context.Background(), ag, "hi", nil, &out, &diag); err != nil {
55 t.Fatal(err)
56 }
57 // Two inferences happened: execution did not return after the
58 // interim continuation.
59 if got := atomic.LoadInt32(&client.calls); got != 2 {
60 t.Fatalf("provider calls = %d, want 2", got)
61 }
62 got := out.String()
63 if !strings.Contains(got, "final") {
64 t.Fatalf("stdout = %q, want it to include the final answer", got)
65 }
66 if !strings.Contains(got, "part one") {
67 t.Fatalf("stdout = %q, want it to include the interim text", got)
68 }
69}