nub/internal/agent/consume.go
Tom a97013d876 initial commit
- v 0.1.0 siehe CHANGELOG.md
2026-07-25 11:08:01 +02:00

86 lines
2.1 KiB
Go

package agent
import (
"encoding/json"
"iter"
"strings"
"nub/internal/llm"
"nub/internal/tool"
)
type blockAcc struct {
kind llm.BlockKind
id string
name string
text strings.Builder
args strings.Builder
}
// consume liest den Event-Stream eines Turns, emittiert Deltas als UIEvents
// und akkumuliert parallel die vollständigen Blöcke der Assistant-Message
// (Ablauf pro Turn, Schritt 2).
func (l *Loop) consume(events iter.Seq2[llm.Event, error], out chan<- tool.UIEvent) (llm.Message, llm.StopReason, llm.Usage, error) {
blocks := map[int]*blockAcc{}
var order []int
stop := llm.StopEnd
var usage llm.Usage
var streamErr error
events(func(ev llm.Event, err error) bool {
if err != nil {
streamErr = err
return false
}
switch e := ev.(type) {
case llm.BlockStart:
b := &blockAcc{kind: e.Block.Kind, id: e.Block.ID, name: e.Block.Name}
blocks[e.Index] = b
order = append(order, e.Index)
if e.Block.Kind == llm.KindToolUse {
out <- tool.ToolCallStart{ID: e.Block.ID, Name: e.Block.Name}
}
case llm.BlockDelta:
b := blocks[e.Index]
if b == nil {
return true
}
switch b.kind {
case llm.KindText:
b.text.WriteString(e.Text)
out <- tool.TextDelta{Text: e.Text}
case llm.KindThinking:
b.text.WriteString(e.Text)
out <- tool.ThinkingDelta{Text: e.Text}
case llm.KindToolUse:
b.args.WriteString(e.PartialJSON)
}
case llm.Done:
stop = e.Stop
usage = e.Usage
}
return true
})
if streamErr != nil {
return llm.Message{}, "", llm.Usage{}, streamErr
}
msg := llm.Message{Role: llm.RoleAssistant}
for _, idx := range order {
b := blocks[idx]
switch b.kind {
case llm.KindText, llm.KindThinking:
if b.text.Len() > 0 {
msg.Content = append(msg.Content, llm.Block{Kind: b.kind, Text: b.text.String()})
}
case llm.KindToolUse:
input := json.RawMessage(b.args.String())
if len(input) == 0 {
input = json.RawMessage("{}")
}
msg.Content = append(msg.Content, llm.Block{Kind: llm.KindToolUse, ID: b.id, Name: b.name, Input: input})
}
}
return msg, stop, usage, nil
}