1package modes
2
3import (
4 "context"
5 "fmt"
6 "io"
7 "os"
8 "strings"
9
10 "github.com/patriceckhart/zot/packages/core"
11 "github.com/patriceckhart/zot/packages/provider"
12)
13
14// RunStream runs the agent and writes assistant text to out as it
15// arrives (EvTextDelta). Tool activity is written to stderr so stdout
16// stays the live assistant transcript. When a provider does not emit
17// deltas, the full assistant text is written when the message completes.
18func RunStream(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, out io.Writer) error {
19 return RunStreamWithDiag(ctx, ag, prompt, images, out, os.Stderr)
20}
21
22// RunStreamWithDiag is RunStream with an explicit diagnostics writer
23// for tool activity (tests pass a buffer).
24func RunStreamWithDiag(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, out, diag io.Writer) error {
25 if diag == nil {
26 diag = io.Discard
27 }
28 var (
29 runErr error
30 streamed bool
31 wroteText bool
32 lastWasNL bool
33 )
34 flush := func() {
35 if f, ok := out.(interface{ Sync() error }); ok {
36 _ = f.Sync()
37 }
38 }
39 writeText := func(s string) {
40 if s == "" {
41 return
42 }
43 _, _ = fmt.Fprint(out, s)
44 flush()
45 wroteText = true
46 lastWasNL = strings.HasSuffix(s, "\n")
47 }
48 assistantText := func(m provider.Message) string {
49 var sb strings.Builder
50 for _, c := range m.Content {
51 if tb, ok := c.(provider.TextBlock); ok {
52 if sb.Len() > 0 {
53 sb.WriteString("\n")
54 }
55 sb.WriteString(tb.Text)
56 }
57 }
58 return sb.String()
59 }
60
61 sink := func(ev core.AgentEvent) {
62 switch e := ev.(type) {
63 case core.EvAssistantStart:
64 streamed = false
65 case core.EvTextDelta:
66 streamed = true
67 writeText(e.Delta)
68 case core.EvAssistantMessage:
69 if !streamed {
70 writeText(assistantText(e.Message))
71 }
72 case core.EvToolCall:
73 fmt.Fprintf(diag, "\n[tool] %s\n", e.Name)
74 if len(e.Args) > 0 {
75 fmt.Fprintf(diag, "%s\n", string(e.Args))
76 }
77 case core.EvToolProgress:
78 if e.Text != "" {
79 fmt.Fprint(diag, e.Text)
80 }
81 case core.EvToolResult:
82 for _, c := range e.Result.Content {
83 if tb, ok := c.(provider.TextBlock); ok && tb.Text != "" {
84 fmt.Fprint(diag, tb.Text)
85 if !strings.HasSuffix(tb.Text, "\n") {
86 fmt.Fprint(diag, "\n")
87 }
88 }
89 }
90 if e.Result.IsError {
91 fmt.Fprintln(diag, "[tool error]")
92 }
93 case core.EvTurnEnd:
94 if e.Err != nil {
95 runErr = e.Err
96 }
97 case core.EvError:
98 if e.Err != nil {
99 runErr = e.Err
100 }
101 }
102 }
103
104 if err := ag.Prompt(ctx, prompt, images, sink); err != nil {
105 return err
106 }
107 if runErr != nil {
108 return runErr
109 }
110 if wroteText && !lastWasNL {
111 _, _ = fmt.Fprint(out, "\n")
112 flush()
113 }
114 return nil
115}