184 lines
7.1 KiB
Go
184 lines
7.1 KiB
Go
// This file implements streamAssistantResponse (US-003): it shapes the context
|
|
// into a provider request, resolves the API key dynamically, drives the
|
|
// provider stream, and back-fills the partial assistant message into the
|
|
// context while emitting message_start / message_update / message_end events.
|
|
package runtime
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/smallnest/pigo/internal/agentcore"
|
|
"github.com/smallnest/pigo/internal/compaction"
|
|
"github.com/smallnest/pigo/internal/provider"
|
|
)
|
|
|
|
// LoopConfig holds the pluggable behavior of the agent loop. Every hook is
|
|
// optional (nil = use the default). The pointer/func-field pattern mirrors pi's
|
|
// optional callbacks.
|
|
type LoopConfig struct {
|
|
// Model is the model id passed to StreamFn.
|
|
Model string
|
|
// APIKey is the static fallback key when GetAPIKey is nil or returns "".
|
|
APIKey string
|
|
// ThinkingLevel is the reasoning effort for requests.
|
|
ThinkingLevel agentcore.ThinkingLevel
|
|
// Stream produces the provider stream. Required (defaults are wired by
|
|
// callers/tests, e.g. a fake provider).
|
|
Stream provider.StreamFn
|
|
|
|
// TransformContext optionally rewrites the message list before conversion
|
|
// (context trimming/injection). Contract: must not error; on failure return
|
|
// a safe fallback. Runs first.
|
|
TransformContext func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList
|
|
// ConvertToLlm optionally filters UI-only messages. Defaults to identity.
|
|
// Contract: must not error.
|
|
ConvertToLlm func(msgs agentcore.MessageList) agentcore.MessageList
|
|
// GetAPIKey optionally resolves a fresh key per request (handles short-lived
|
|
// token expiry). Falls back to APIKey when nil or empty.
|
|
GetAPIKey func(ctx context.Context, provider string) string
|
|
// Provider is the provider name passed to GetAPIKey.
|
|
Provider string
|
|
|
|
// ContextWindow is the model's total context-token budget, used to decide
|
|
// automatic compaction. When <= 0 the window is unknown and auto-compaction
|
|
// is disabled (ShouldCompact returns false), so the loop behaves exactly as
|
|
// before for callers that do not plumb it through.
|
|
ContextWindow int
|
|
// Compaction holds the thresholds/retention knobs for auto-compaction. Its
|
|
// Enabled flag gates the feature independently of ContextWindow.
|
|
Compaction compaction.CompactionSettings
|
|
// SummaryStream produces the provider stream used to generate compaction
|
|
// summaries. Defaults to Stream when nil.
|
|
SummaryStream provider.StreamFn
|
|
// SummaryModel is the model used for summarization. When zero, a model is
|
|
// synthesized from Model/ContextWindow.
|
|
SummaryModel provider.Model
|
|
|
|
// Extra is forwarded to StreamConfig.Extra.
|
|
Extra map[string]any
|
|
}
|
|
|
|
// streamAssistantResponse runs one assistant turn: it builds the request from
|
|
// agentCtx, streams the provider response, back-fills the partial into
|
|
// agentCtx.Messages, and returns the final assistant message. The sequence
|
|
// (transformContext → convertToLlm → resolve key → stream → drain) is kept
|
|
// identical to pi. It never returns an error for a request failure — such
|
|
// failures arrive as a terminal assistant message with stopReason error/aborted.
|
|
func streamAssistantResponse(ctx context.Context, agentCtx *agentcore.AgentContext, cfg LoopConfig, emit agentcore.EmitFunc) (agentcore.AssistantMessage, error) {
|
|
// 1. transformContext (optional, must not error).
|
|
msgs := agentCtx.Messages
|
|
if cfg.TransformContext != nil {
|
|
msgs = cfg.TransformContext(ctx, msgs)
|
|
}
|
|
// 2. convertToLlm (filter UI-only; default identity).
|
|
if cfg.ConvertToLlm != nil {
|
|
msgs = cfg.ConvertToLlm(msgs)
|
|
}
|
|
// 3. shape the LLM context.
|
|
llm := provider.LlmContext{
|
|
SystemPrompt: agentCtx.SystemPrompt,
|
|
Messages: msgs,
|
|
Tools: agentCtx.Tools,
|
|
}
|
|
// 4. resolve API key dynamically, fall back to static.
|
|
key := cfg.APIKey
|
|
if cfg.GetAPIKey != nil {
|
|
if dyn := cfg.GetAPIKey(ctx, cfg.Provider); dyn != "" {
|
|
key = dyn
|
|
}
|
|
}
|
|
// 5. build the provider stream.
|
|
stream, err := cfg.Stream(ctx, cfg.Model, llm, provider.StreamConfig{
|
|
APIKey: key,
|
|
ThinkingLevel: cfg.ThinkingLevel,
|
|
Extra: cfg.Extra,
|
|
})
|
|
if err != nil {
|
|
// Early "cannot build stream" failure: synthesize a terminal message so
|
|
// the loop has a uniform assistant message to record.
|
|
return newErrorAssistantMessage(cfg, err), nil
|
|
}
|
|
|
|
// 6. drain the stream, back-filling the partial into the context.
|
|
addedPartial := false
|
|
backfill := func(partial agentcore.AssistantMessage) {
|
|
if !addedPartial {
|
|
agentCtx.Messages = append(agentCtx.Messages, partial)
|
|
addedPartial = true
|
|
} else {
|
|
agentCtx.Messages[len(agentCtx.Messages)-1] = partial
|
|
}
|
|
}
|
|
|
|
for ev := range stream.Events() {
|
|
switch e := ev.(type) {
|
|
case provider.StreamStartEvent:
|
|
backfill(e.Partial)
|
|
if err := emit(ctx, agentcore.MessageStartEvent{Message: e.Partial}); err != nil {
|
|
return agentcore.AssistantMessage{}, err
|
|
}
|
|
case provider.StreamTextEvent:
|
|
backfill(e.Partial)
|
|
if err := emit(ctx, agentcore.MessageUpdateEvent{Message: e.Partial, AssistantMessageEvent: e}); err != nil {
|
|
return agentcore.AssistantMessage{}, err
|
|
}
|
|
case provider.StreamThinkingEvent:
|
|
backfill(e.Partial)
|
|
if err := emit(ctx, agentcore.MessageUpdateEvent{Message: e.Partial, AssistantMessageEvent: e}); err != nil {
|
|
return agentcore.AssistantMessage{}, err
|
|
}
|
|
case provider.StreamToolCallEvent:
|
|
backfill(e.Partial)
|
|
if err := emit(ctx, agentcore.MessageUpdateEvent{Message: e.Partial, AssistantMessageEvent: e}); err != nil {
|
|
return agentcore.AssistantMessage{}, err
|
|
}
|
|
case provider.StreamDoneEvent:
|
|
finalizeMessage(agentCtx, e.Message, &addedPartial)
|
|
if err := emit(ctx, agentcore.MessageEndEvent{Message: e.Message}); err != nil {
|
|
return agentcore.AssistantMessage{}, err
|
|
}
|
|
return e.Message, nil
|
|
case provider.StreamErrorEvent:
|
|
finalizeMessage(agentCtx, e.Message, &addedPartial)
|
|
if err := emit(ctx, agentcore.MessageEndEvent{Message: e.Message}); err != nil {
|
|
return agentcore.AssistantMessage{}, err
|
|
}
|
|
return e.Message, nil
|
|
}
|
|
}
|
|
|
|
// 7. stream ended without done/error: fall back to the stream result.
|
|
final, resErr := stream.Result(ctx)
|
|
if resErr != nil {
|
|
return newErrorAssistantMessage(cfg, resErr), nil
|
|
}
|
|
finalizeMessage(agentCtx, final, &addedPartial)
|
|
if err := emit(ctx, agentcore.MessageEndEvent{Message: final}); err != nil {
|
|
return agentcore.AssistantMessage{}, err
|
|
}
|
|
return final, nil
|
|
}
|
|
|
|
// finalizeMessage replaces the placeholder partial with the final message, or
|
|
// appends it if the provider sent done/error without a prior start.
|
|
func finalizeMessage(agentCtx *agentcore.AgentContext, final agentcore.AssistantMessage, addedPartial *bool) {
|
|
if *addedPartial {
|
|
agentCtx.Messages[len(agentCtx.Messages)-1] = final
|
|
} else {
|
|
agentCtx.Messages = append(agentCtx.Messages, final)
|
|
*addedPartial = true
|
|
}
|
|
}
|
|
|
|
// newErrorAssistantMessage builds a terminal assistant message for an early
|
|
// failure that never produced a provider stream.
|
|
func newErrorAssistantMessage(cfg LoopConfig, err error) agentcore.AssistantMessage {
|
|
return agentcore.AssistantMessage{
|
|
RoleField: agentcore.RoleAssistant,
|
|
Model: cfg.Model,
|
|
Provider: cfg.Provider,
|
|
StopReason: agentcore.StopReasonError,
|
|
ErrorMessage: err.Error(),
|
|
}
|
|
}
|