feat: provider manager, cancellation, session summaries (store+agent)
This commit is contained in:
OpenClaw Dev
parent
406a523d54
commit
f1b2093b22
10 changed files with 883 additions and 37 deletions
|
|
@ -1,9 +1,12 @@
|
|||
// Package agent implements the LLM tool-calling loop.
|
||||
// Package agent implements the LLM tool-calling loop, cancellation and
|
||||
// summarisation of long sessions.
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/crypto/byteclaw/internal/llm"
|
||||
|
|
@ -11,9 +14,12 @@ import (
|
|||
"github.com/crypto/byteclaw/internal/tools"
|
||||
)
|
||||
|
||||
// ErrCancelled is returned by Run when the user stopped the answer.
|
||||
var ErrCancelled = errors.New("cancelled by user")
|
||||
|
||||
// Event is streamed to clients while the agent works.
|
||||
type Event struct {
|
||||
Type string `json:"type"` // delta | tool_call | tool_result | done | error
|
||||
Type string `json:"type"` // session | delta | tool_call | tool_result | summarized | done | cancelled | error
|
||||
Text string `json:"text,omitempty"`
|
||||
Tool string `json:"tool,omitempty"`
|
||||
Args string `json:"args,omitempty"`
|
||||
|
|
@ -21,16 +27,29 @@ type Event struct {
|
|||
Session string `json:"session_id,omitempty"`
|
||||
}
|
||||
|
||||
// SummarizeConfig controls automatic summarisation.
|
||||
type SummarizeConfig struct {
|
||||
Enabled bool
|
||||
MaxMessages int
|
||||
MaxTokens int
|
||||
KeepRecent int
|
||||
}
|
||||
|
||||
type Agent struct {
|
||||
LLM llm.Provider
|
||||
// LLM is used when Resolve is nil.
|
||||
LLM llm.Provider
|
||||
// Resolve picks the provider for a session (e.g. per-session override).
|
||||
Resolve func(ctx context.Context, sessionID string) (llm.Provider, error)
|
||||
Tools *tools.Registry
|
||||
Store *store.Store
|
||||
SystemPrompt string
|
||||
MaxIterations int
|
||||
HistoryLimit int
|
||||
Summarize SummarizeConfig
|
||||
|
||||
mu sync.Mutex
|
||||
locks map[string]*sync.Mutex
|
||||
runs map[string]context.CancelCauseFunc
|
||||
}
|
||||
|
||||
func (a *Agent) sessionLock(id string) *sync.Mutex {
|
||||
|
|
@ -47,9 +66,55 @@ func (a *Agent) sessionLock(id string) *sync.Mutex {
|
|||
return l
|
||||
}
|
||||
|
||||
func (a *Agent) provider(ctx context.Context, sessionID string) (llm.Provider, error) {
|
||||
if a.Resolve != nil {
|
||||
return a.Resolve(ctx, sessionID)
|
||||
}
|
||||
if a.LLM == nil {
|
||||
return nil, errors.New("no LLM provider configured")
|
||||
}
|
||||
return a.LLM, nil
|
||||
}
|
||||
|
||||
// Cancel stops the running answer of a session. Returns false if none is running.
|
||||
func (a *Agent) Cancel(sessionID string) bool {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if c, ok := a.runs[sessionID]; ok {
|
||||
c(ErrCancelled)
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Running reports whether an answer is in progress for the session.
|
||||
func (a *Agent) Running(sessionID string) bool {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
_, ok := a.runs[sessionID]
|
||||
return ok
|
||||
}
|
||||
|
||||
func (a *Agent) register(sessionID string, c context.CancelCauseFunc) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if a.runs == nil {
|
||||
a.runs = map[string]context.CancelCauseFunc{}
|
||||
}
|
||||
a.runs[sessionID] = c
|
||||
}
|
||||
|
||||
func (a *Agent) unregister(sessionID string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
delete(a.runs, sessionID)
|
||||
}
|
||||
|
||||
// Run appends the user message to the session and loops until the model
|
||||
// answers without tool calls. Returns the final assistant text.
|
||||
func (a *Agent) Run(ctx context.Context, sessionID, userText string, emit func(Event)) (string, error) {
|
||||
// answers without tool calls. Returns the final assistant text. If the run is
|
||||
// cancelled via Cancel, the partial answer is stored (marked cancelled) and
|
||||
// ErrCancelled is returned together with the partial text.
|
||||
func (a *Agent) Run(parent context.Context, sessionID, userText string, emit func(Event)) (string, error) {
|
||||
if emit == nil {
|
||||
emit = func(Event) {}
|
||||
}
|
||||
|
|
@ -57,26 +122,57 @@ func (a *Agent) Run(ctx context.Context, sessionID, userText string, emit func(E
|
|||
l.Lock()
|
||||
defer l.Unlock()
|
||||
|
||||
if _, err := a.Store.AppendMessage(ctx, sessionID, llm.Message{Role: llm.RoleUser, Content: userText}); err != nil {
|
||||
ctx, cancel := context.WithCancelCause(parent)
|
||||
defer cancel(nil)
|
||||
a.register(sessionID, cancel)
|
||||
defer a.unregister(sessionID)
|
||||
// Writes after a cancellation must still succeed.
|
||||
saveCtx := context.WithoutCancel(ctx)
|
||||
|
||||
if _, err := a.Store.AppendMessage(saveCtx, sessionID, llm.Message{Role: llm.RoleUser, Content: userText}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
prov, err := a.provider(ctx, sessionID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if a.Summarize.Enabled {
|
||||
if sm, err := a.summarizeLocked(ctx, prov, sessionID, false); err == nil && sm != nil {
|
||||
emit(Event{Type: "summarized", Text: sm.Content})
|
||||
}
|
||||
}
|
||||
var specs []llm.ToolSpec
|
||||
if a.Tools != nil {
|
||||
specs = a.Tools.Specs()
|
||||
}
|
||||
cancelled := func() bool { return errors.Is(context.Cause(ctx), ErrCancelled) }
|
||||
for i := 0; i < a.MaxIterations; i++ {
|
||||
msgs, err := a.history(ctx, sessionID)
|
||||
if cancelled() {
|
||||
return "", ErrCancelled
|
||||
}
|
||||
msgs, err := a.history(saveCtx, sessionID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
resp, err := a.LLM.Chat(ctx, llm.Request{Messages: msgs, Tools: specs}, func(t string) {
|
||||
var partial strings.Builder
|
||||
resp, err := prov.Chat(ctx, llm.Request{Messages: msgs, Tools: specs}, func(t string) {
|
||||
partial.WriteString(t)
|
||||
emit(Event{Type: "delta", Text: t})
|
||||
})
|
||||
if err != nil {
|
||||
if cancelled() {
|
||||
text := partial.String()
|
||||
if text != "" {
|
||||
if _, serr := a.Store.AppendMessage(saveCtx, sessionID, llm.Message{Role: llm.RoleAssistant, Content: text, Cancelled: true}); serr != nil {
|
||||
return text, serr
|
||||
}
|
||||
}
|
||||
return text, ErrCancelled
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
asst := llm.Message{Role: llm.RoleAssistant, Content: resp.Content, ToolCalls: resp.ToolCalls}
|
||||
if _, err := a.Store.AppendMessage(ctx, sessionID, asst); err != nil {
|
||||
if _, err := a.Store.AppendMessage(saveCtx, sessionID, asst); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(resp.ToolCalls) == 0 {
|
||||
|
|
@ -84,9 +180,13 @@ func (a *Agent) Run(ctx context.Context, sessionID, userText string, emit func(E
|
|||
}
|
||||
for _, tc := range resp.ToolCalls {
|
||||
emit(Event{Type: "tool_call", Tool: tc.Name, Args: tc.Arguments, CallID: tc.ID})
|
||||
out := a.Tools.Call(ctx, tc.Name, tc.Arguments)
|
||||
out := "error: cancelled by user"
|
||||
if !cancelled() {
|
||||
out = a.Tools.Call(ctx, tc.Name, tc.Arguments)
|
||||
}
|
||||
emit(Event{Type: "tool_result", Tool: tc.Name, Text: clip(out, 2000), CallID: tc.ID})
|
||||
if _, err := a.Store.AppendMessage(ctx, sessionID, llm.Message{Role: llm.RoleTool, ToolCallID: tc.ID, Name: tc.Name, Content: out}); err != nil {
|
||||
// Every tool call gets a result so the history stays valid for the provider.
|
||||
if _, err := a.Store.AppendMessage(saveCtx, sessionID, llm.Message{Role: llm.RoleTool, ToolCallID: tc.ID, Name: tc.Name, Content: out}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
|
@ -95,7 +195,14 @@ func (a *Agent) Run(ctx context.Context, sessionID, userText string, emit func(E
|
|||
}
|
||||
|
||||
func (a *Agent) history(ctx context.Context, sessionID string) ([]llm.Message, error) {
|
||||
stored, err := a.Store.Messages(ctx, sessionID, a.HistoryLimit)
|
||||
var after int64
|
||||
summary := ""
|
||||
if sm, err := a.Store.LatestSummary(ctx, sessionID); err == nil {
|
||||
after, summary = sm.UptoMessageID, sm.Content
|
||||
} else if !errors.Is(err, store.ErrNotFound) {
|
||||
return nil, err
|
||||
}
|
||||
stored, err := a.Store.MessagesAfter(ctx, sessionID, after, a.HistoryLimit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -103,12 +210,17 @@ func (a *Agent) history(ctx context.Context, sessionID string) ([]llm.Message, e
|
|||
for len(stored) > 0 && stored[0].Role == llm.RoleTool {
|
||||
stored = stored[1:]
|
||||
}
|
||||
msgs := make([]llm.Message, 0, len(stored)+1)
|
||||
msgs := make([]llm.Message, 0, len(stored)+2)
|
||||
if a.SystemPrompt != "" {
|
||||
msgs = append(msgs, llm.Message{Role: llm.RoleSystem, Content: a.SystemPrompt})
|
||||
}
|
||||
if summary != "" {
|
||||
msgs = append(msgs, llm.Message{Role: llm.RoleSystem, Content: "Summary of the earlier conversation:\n" + summary})
|
||||
}
|
||||
for _, m := range stored {
|
||||
msgs = append(msgs, m.Message)
|
||||
mm := m.Message
|
||||
mm.Cancelled = false
|
||||
msgs = append(msgs, mm)
|
||||
}
|
||||
return msgs, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@ package agent
|
|||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
|
@ -91,3 +93,116 @@ func TestHistoryDropsOrphanToolMessages(t *testing.T) {
|
|||
t.Fatalf("%+v", h)
|
||||
}
|
||||
}
|
||||
|
||||
// blocking streams one delta then waits for cancellation, like a slow provider.
|
||||
type blocking struct{ started chan struct{} }
|
||||
|
||||
func (b *blocking) Chat(ctx context.Context, _ llm.Request, d llm.DeltaFunc) (llm.Response, error) {
|
||||
d("Teil")
|
||||
d("antwort")
|
||||
close(b.started)
|
||||
<-ctx.Done()
|
||||
return llm.Response{}, ctx.Err()
|
||||
}
|
||||
|
||||
func TestCancelStoresPartialAnswer(t *testing.T) {
|
||||
st := newStore(t)
|
||||
ss, _ := st.CreateSession(context.Background(), "")
|
||||
b := &blocking{started: make(chan struct{})}
|
||||
a := &Agent{LLM: b, Store: st, Tools: tools.NewRegistry(), MaxIterations: 3}
|
||||
if a.Cancel(ss.ID) {
|
||||
t.Fatal("nothing should be running")
|
||||
}
|
||||
go func() {
|
||||
<-b.started
|
||||
if !a.Cancel(ss.ID) {
|
||||
t.Error("cancel should find the run")
|
||||
}
|
||||
}()
|
||||
text, err := a.Run(context.Background(), ss.ID, "erzähl was", nil)
|
||||
if !errors.Is(err, ErrCancelled) || text != "Teilantwort" {
|
||||
t.Fatalf("%q %v", text, err)
|
||||
}
|
||||
msgs, _ := st.Messages(context.Background(), ss.ID, 0)
|
||||
if len(msgs) != 2 || !msgs[1].Cancelled || msgs[1].Content != "Teilantwort" {
|
||||
t.Fatalf("%+v", msgs)
|
||||
}
|
||||
if a.Running(ss.ID) {
|
||||
t.Fatal("run must be unregistered")
|
||||
}
|
||||
}
|
||||
|
||||
// recorder answers summary requests with "SUMMARY" and normal turns with "ok".
|
||||
type recorder struct{ reqs []llm.Request }
|
||||
|
||||
func (r *recorder) Chat(_ context.Context, req llm.Request, _ llm.DeltaFunc) (llm.Response, error) {
|
||||
r.reqs = append(r.reqs, req)
|
||||
if strings.Contains(req.Messages[0].Content, "compress chat histories") {
|
||||
return llm.Response{Content: "SUMMARY"}, nil
|
||||
}
|
||||
return llm.Response{Content: "ok"}, nil
|
||||
}
|
||||
|
||||
func TestAutoSummarize(t *testing.T) {
|
||||
st := newStore(t)
|
||||
ctx := context.Background()
|
||||
ss, _ := st.CreateSession(ctx, "")
|
||||
p := &recorder{}
|
||||
a := &Agent{LLM: p, Store: st, Tools: tools.NewRegistry(), MaxIterations: 2, SystemPrompt: "sys",
|
||||
Summarize: SummarizeConfig{Enabled: true, MaxMessages: 6, KeepRecent: 2}}
|
||||
for i := 0; i < 3; i++ { // 6 messages -> at threshold, no summary yet
|
||||
a.Run(ctx, ss.ID, fmt.Sprintf("frage %d", i), nil)
|
||||
}
|
||||
if _, err := st.LatestSummary(ctx, ss.ID); err == nil {
|
||||
t.Fatal("summary created too early")
|
||||
}
|
||||
var events []string
|
||||
a.Run(ctx, ss.ID, "frage 3", func(e Event) { events = append(events, e.Type) })
|
||||
sm, err := st.LatestSummary(ctx, ss.ID)
|
||||
if err != nil || sm.Content != "SUMMARY" {
|
||||
t.Fatalf("%+v %v", sm, err)
|
||||
}
|
||||
if events[0] != "summarized" {
|
||||
t.Fatalf("%v", events)
|
||||
}
|
||||
last := p.reqs[len(p.reqs)-1].Messages
|
||||
if last[1].Role != llm.RoleSystem || !strings.Contains(last[1].Content, "SUMMARY") {
|
||||
t.Fatalf("summary not in context: %+v", last)
|
||||
}
|
||||
// the verbatim tail starts at a user message and older messages are not sent
|
||||
if last[2].Role != llm.RoleUser || strings.Contains(fmt.Sprint(last), "frage 0") {
|
||||
t.Fatalf("bad window: %+v", last)
|
||||
}
|
||||
if all, _ := st.Messages(ctx, ss.ID, 0); len(all) != 8 {
|
||||
t.Fatalf("originals must be kept, got %d", len(all))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutoSummarizeByTokens(t *testing.T) {
|
||||
st := newStore(t)
|
||||
ctx := context.Background()
|
||||
ss, _ := st.CreateSession(ctx, "")
|
||||
a := &Agent{LLM: &recorder{}, Store: st, Tools: tools.NewRegistry(), MaxIterations: 2,
|
||||
Summarize: SummarizeConfig{Enabled: true, MaxMessages: 1000, MaxTokens: 50, KeepRecent: 1}}
|
||||
a.Run(ctx, ss.ID, strings.Repeat("lang ", 100), nil)
|
||||
a.Run(ctx, ss.ID, "kurz", nil)
|
||||
if _, err := st.LatestSummary(ctx, ss.ID); err != nil {
|
||||
t.Fatal("token threshold should trigger summary")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizeNowManual(t *testing.T) {
|
||||
st := newStore(t)
|
||||
ctx := context.Background()
|
||||
ss, _ := st.CreateSession(ctx, "")
|
||||
a := &Agent{LLM: &recorder{}, Store: st, Tools: tools.NewRegistry(), MaxIterations: 2, Summarize: SummarizeConfig{KeepRecent: 2}}
|
||||
if sm, err := a.SummarizeNow(ctx, ss.ID); err != nil || sm != nil {
|
||||
t.Fatalf("empty session: %v %v", sm, err)
|
||||
}
|
||||
a.Run(ctx, ss.ID, "eins", nil)
|
||||
a.Run(ctx, ss.ID, "zwei", nil)
|
||||
sm, err := a.SummarizeNow(ctx, ss.ID)
|
||||
if err != nil || sm == nil || sm.Content != "SUMMARY" {
|
||||
t.Fatalf("%v %v", sm, err)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
105
internal/agent/summarize.go
Normal file
105
internal/agent/summarize.go
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/crypto/byteclaw/internal/llm"
|
||||
"github.com/crypto/byteclaw/internal/store"
|
||||
)
|
||||
|
||||
const summarizePrompt = `You compress chat histories. Write a concise summary of the conversation below so an assistant can continue it.
|
||||
Keep facts, user preferences, decisions, results of tool calls and open tasks. Use the language of the conversation. No preamble.`
|
||||
|
||||
// estimateTokens is a cheap approximation (≈4 characters per token).
|
||||
func estimateTokens(msgs []store.StoredMessage) int {
|
||||
n := 0
|
||||
for _, m := range msgs {
|
||||
n += len(m.Content) / 4
|
||||
for _, tc := range m.ToolCalls {
|
||||
n += len(tc.Arguments) / 4
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// SummarizeNow summarises the session regardless of thresholds (manual trigger).
|
||||
// Returns nil summary if there is nothing old enough to summarise.
|
||||
func (a *Agent) SummarizeNow(ctx context.Context, sessionID string) (*store.Summary, error) {
|
||||
l := a.sessionLock(sessionID)
|
||||
l.Lock()
|
||||
defer l.Unlock()
|
||||
prov, err := a.provider(ctx, sessionID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return a.summarizeLocked(ctx, prov, sessionID, true)
|
||||
}
|
||||
|
||||
func (a *Agent) summarizeLocked(ctx context.Context, prov llm.Provider, sessionID string, force bool) (*store.Summary, error) {
|
||||
var after int64
|
||||
prev := ""
|
||||
if sm, err := a.Store.LatestSummary(ctx, sessionID); err == nil {
|
||||
after, prev = sm.UptoMessageID, sm.Content
|
||||
} else if !errors.Is(err, store.ErrNotFound) {
|
||||
return nil, err
|
||||
}
|
||||
msgs, err := a.Store.MessagesAfter(ctx, sessionID, after, 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cfg := a.Summarize
|
||||
if !force && len(msgs) <= cfg.MaxMessages && (cfg.MaxTokens <= 0 || estimateTokens(msgs) <= cfg.MaxTokens) {
|
||||
return nil, nil
|
||||
}
|
||||
keep := cfg.KeepRecent
|
||||
if keep < 1 {
|
||||
keep = 1
|
||||
}
|
||||
cut := len(msgs) - keep
|
||||
// The verbatim tail must start at a user message so tool call/result pairs stay together.
|
||||
for cut > 0 && cut < len(msgs) && msgs[cut].Role != llm.RoleUser {
|
||||
cut--
|
||||
}
|
||||
if cut <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var b strings.Builder
|
||||
if prev != "" {
|
||||
b.WriteString("Previous summary:\n" + prev + "\n\n")
|
||||
}
|
||||
b.WriteString("Conversation:\n")
|
||||
for _, m := range msgs[:cut] {
|
||||
switch {
|
||||
case m.Role == llm.RoleTool:
|
||||
fmt.Fprintf(&b, "tool %s result: %s\n", m.Name, clip(m.Content, 500))
|
||||
case len(m.ToolCalls) > 0:
|
||||
for _, tc := range m.ToolCalls {
|
||||
fmt.Fprintf(&b, "assistant called %s(%s)\n", tc.Name, clip(tc.Arguments, 200))
|
||||
}
|
||||
if m.Content != "" {
|
||||
fmt.Fprintf(&b, "assistant: %s\n", m.Content)
|
||||
}
|
||||
default:
|
||||
fmt.Fprintf(&b, "%s: %s\n", m.Role, m.Content)
|
||||
}
|
||||
}
|
||||
resp, err := prov.Chat(ctx, llm.Request{Messages: []llm.Message{
|
||||
{Role: llm.RoleSystem, Content: summarizePrompt},
|
||||
{Role: llm.RoleUser, Content: b.String()},
|
||||
}}, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("summarize: %w", err)
|
||||
}
|
||||
text := strings.TrimSpace(resp.Content)
|
||||
if text == "" {
|
||||
return nil, errors.New("summarize: empty summary")
|
||||
}
|
||||
sm, err := a.Store.AddSummary(context.WithoutCancel(ctx), sessionID, msgs[cut-1].ID, text)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &sm, nil
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue