feat: provider manager, cancellation, session summaries (store+agent)
This commit is contained in:
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
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue