print.go RawHistoryFolder

1// Package modes implements zot's three run modes: print, json, interactive.
2package modes
3
4import (
5	"context"
6	"fmt"
7	"io"
8	"strings"
9
10	"github.com/patriceckhart/zot/packages/core"
11	"github.com/patriceckhart/zot/packages/provider"
12)
13
14// RunPrint runs the agent to completion and writes only the final
15// assistant text block to out. It returns usage for this invocation,
16// excluding any cumulative usage restored from a session.
17func RunPrint(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, out io.Writer) (provider.Usage, error) {
18	var finalText strings.Builder
19	var lastAssistant string
20	var usage provider.Usage
21	var haveUsage bool
22	var runErr error
23
24	sink := func(ev core.AgentEvent) {
25		switch e := ev.(type) {
26		case core.EvAssistantMessage:
27			// Keep the most recent assistant text block; by the end it's the final answer.
28			var sb strings.Builder
29			for _, c := range e.Message.Content {
30				if tb, ok := c.(provider.TextBlock); ok {
31					if sb.Len() > 0 {
32						sb.WriteString("\n")
33					}
34					sb.WriteString(tb.Text)
35				}
36			}
37			if sb.Len() > 0 {
38				lastAssistant = sb.String()
39			}
40		case core.EvUsage:
41			if haveUsage {
42				usage = usage.Add(e.Usage)
43			} else {
44				usage = e.Usage
45				haveUsage = true
46			}
47		case core.EvTurnEnd:
48			if e.Err != nil {
49				runErr = e.Err
50			}
51		}
52	}
53
54	if err := ag.Prompt(ctx, prompt, images, sink); err != nil {
55		return usage, err
56	}
57	if runErr != nil {
58		return usage, runErr
59	}
60
61	finalText.WriteString(lastAssistant)
62	if finalText.Len() > 0 && !strings.HasSuffix(finalText.String(), "\n") {
63		finalText.WriteString("\n")
64	}
65	_, err := fmt.Fprint(out, finalText.String())
66	return usage, err
67}