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
|
package agent
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
"github.com/crypto/byteclaw/internal/llm"
|
"github.com/crypto/byteclaw/internal/llm"
|
||||||
|
|
@ -11,9 +14,12 @@ import (
|
||||||
"github.com/crypto/byteclaw/internal/tools"
|
"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.
|
// Event is streamed to clients while the agent works.
|
||||||
type Event struct {
|
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"`
|
Text string `json:"text,omitempty"`
|
||||||
Tool string `json:"tool,omitempty"`
|
Tool string `json:"tool,omitempty"`
|
||||||
Args string `json:"args,omitempty"`
|
Args string `json:"args,omitempty"`
|
||||||
|
|
@ -21,16 +27,29 @@ type Event struct {
|
||||||
Session string `json:"session_id,omitempty"`
|
Session string `json:"session_id,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SummarizeConfig controls automatic summarisation.
|
||||||
|
type SummarizeConfig struct {
|
||||||
|
Enabled bool
|
||||||
|
MaxMessages int
|
||||||
|
MaxTokens int
|
||||||
|
KeepRecent int
|
||||||
|
}
|
||||||
|
|
||||||
type Agent struct {
|
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
|
Tools *tools.Registry
|
||||||
Store *store.Store
|
Store *store.Store
|
||||||
SystemPrompt string
|
SystemPrompt string
|
||||||
MaxIterations int
|
MaxIterations int
|
||||||
HistoryLimit int
|
HistoryLimit int
|
||||||
|
Summarize SummarizeConfig
|
||||||
|
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
locks map[string]*sync.Mutex
|
locks map[string]*sync.Mutex
|
||||||
|
runs map[string]context.CancelCauseFunc
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *Agent) sessionLock(id string) *sync.Mutex {
|
func (a *Agent) sessionLock(id string) *sync.Mutex {
|
||||||
|
|
@ -47,9 +66,55 @@ func (a *Agent) sessionLock(id string) *sync.Mutex {
|
||||||
return l
|
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
|
// Run appends the user message to the session and loops until the model
|
||||||
// answers without tool calls. Returns the final assistant text.
|
// answers without tool calls. Returns the final assistant text. If the run is
|
||||||
func (a *Agent) Run(ctx context.Context, sessionID, userText string, emit func(Event)) (string, error) {
|
// 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 {
|
if emit == nil {
|
||||||
emit = func(Event) {}
|
emit = func(Event) {}
|
||||||
}
|
}
|
||||||
|
|
@ -57,26 +122,57 @@ func (a *Agent) Run(ctx context.Context, sessionID, userText string, emit func(E
|
||||||
l.Lock()
|
l.Lock()
|
||||||
defer l.Unlock()
|
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
|
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
|
var specs []llm.ToolSpec
|
||||||
if a.Tools != nil {
|
if a.Tools != nil {
|
||||||
specs = a.Tools.Specs()
|
specs = a.Tools.Specs()
|
||||||
}
|
}
|
||||||
|
cancelled := func() bool { return errors.Is(context.Cause(ctx), ErrCancelled) }
|
||||||
for i := 0; i < a.MaxIterations; i++ {
|
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 {
|
if err != nil {
|
||||||
return "", err
|
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})
|
emit(Event{Type: "delta", Text: t})
|
||||||
})
|
})
|
||||||
if err != nil {
|
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
|
return "", err
|
||||||
}
|
}
|
||||||
asst := llm.Message{Role: llm.RoleAssistant, Content: resp.Content, ToolCalls: resp.ToolCalls}
|
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
|
return "", err
|
||||||
}
|
}
|
||||||
if len(resp.ToolCalls) == 0 {
|
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 {
|
for _, tc := range resp.ToolCalls {
|
||||||
emit(Event{Type: "tool_call", Tool: tc.Name, Args: tc.Arguments, CallID: tc.ID})
|
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})
|
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
|
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) {
|
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 {
|
if err != nil {
|
||||||
return nil, err
|
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 {
|
for len(stored) > 0 && stored[0].Role == llm.RoleTool {
|
||||||
stored = stored[1:]
|
stored = stored[1:]
|
||||||
}
|
}
|
||||||
msgs := make([]llm.Message, 0, len(stored)+1)
|
msgs := make([]llm.Message, 0, len(stored)+2)
|
||||||
if a.SystemPrompt != "" {
|
if a.SystemPrompt != "" {
|
||||||
msgs = append(msgs, llm.Message{Role: llm.RoleSystem, Content: 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 {
|
for _, m := range stored {
|
||||||
msgs = append(msgs, m.Message)
|
mm := m.Message
|
||||||
|
mm.Cancelled = false
|
||||||
|
msgs = append(msgs, mm)
|
||||||
}
|
}
|
||||||
return msgs, nil
|
return msgs, nil
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -2,6 +2,8 @@ package agent
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
@ -91,3 +93,116 @@ func TestHistoryDropsOrphanToolMessages(t *testing.T) {
|
||||||
t.Fatalf("%+v", h)
|
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
|
||||||
|
}
|
||||||
|
|
@ -20,12 +20,32 @@ type LLM struct {
|
||||||
Model string `yaml:"model" json:"model"`
|
Model string `yaml:"model" json:"model"`
|
||||||
TimeoutSec int `yaml:"timeout_sec" json:"timeout_sec"`
|
TimeoutSec int `yaml:"timeout_sec" json:"timeout_sec"`
|
||||||
Temperature float64 `yaml:"temperature" json:"temperature"`
|
Temperature float64 `yaml:"temperature" json:"temperature"`
|
||||||
|
// Providers holds optional per-preset settings (keys stay server-side) so the
|
||||||
|
// active provider can be switched at runtime, e.g. providers: {dahl: {api_key: ...}}.
|
||||||
|
Providers map[string]ProviderOverride `yaml:"providers" json:"providers"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type ProviderOverride struct {
|
||||||
|
BaseURL string `yaml:"base_url" json:"base_url"`
|
||||||
|
APIKey string `yaml:"api_key" json:"api_key"`
|
||||||
|
Model string `yaml:"model" json:"model"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type Agent struct {
|
type Agent struct {
|
||||||
SystemPrompt string `yaml:"system_prompt" json:"system_prompt"`
|
SystemPrompt string `yaml:"system_prompt" json:"system_prompt"`
|
||||||
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
|
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
|
||||||
HistoryLimit int `yaml:"history_limit" json:"history_limit"`
|
HistoryLimit int `yaml:"history_limit" json:"history_limit"`
|
||||||
|
Summarize Summarize `yaml:"summarize" json:"summarize"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Summarize controls automatic summarisation of long sessions.
|
||||||
|
type Summarize struct {
|
||||||
|
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||||
|
// Summarise when unsummarised messages exceed MaxMessages or ~MaxTokens (chars/4).
|
||||||
|
MaxMessages int `yaml:"max_messages" json:"max_messages"`
|
||||||
|
MaxTokens int `yaml:"max_tokens" json:"max_tokens"`
|
||||||
|
// KeepRecent newest messages always stay verbatim.
|
||||||
|
KeepRecent int `yaml:"keep_recent" json:"keep_recent"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type WebFetch struct {
|
type WebFetch struct {
|
||||||
|
|
@ -80,6 +100,7 @@ func Default() Config {
|
||||||
SystemPrompt: "You are ByteClaw, a concise and helpful self-hosted assistant. Use tools when they help.",
|
SystemPrompt: "You are ByteClaw, a concise and helpful self-hosted assistant. Use tools when they help.",
|
||||||
MaxIterations: 8,
|
MaxIterations: 8,
|
||||||
HistoryLimit: 50,
|
HistoryLimit: 50,
|
||||||
|
Summarize: Summarize{Enabled: true, MaxMessages: 40, MaxTokens: 12000, KeepRecent: 10},
|
||||||
},
|
},
|
||||||
Tools: Tools{
|
Tools: Tools{
|
||||||
WebFetch: WebFetch{Enabled: true, MaxBytes: 512 * 1024, TimeoutSec: 15},
|
WebFetch: WebFetch{Enabled: true, MaxBytes: 512 * 1024, TimeoutSec: 15},
|
||||||
|
|
|
||||||
|
|
@ -24,6 +24,8 @@ type Message struct {
|
||||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||||
Name string `json:"name,omitempty"`
|
Name string `json:"name,omitempty"`
|
||||||
|
// Cancelled marks a partial assistant answer that was stopped by the user.
|
||||||
|
Cancelled bool `json:"cancelled,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// ToolSpec describes a callable tool using JSON Schema parameters.
|
// ToolSpec describes a callable tool using JSON Schema parameters.
|
||||||
|
|
|
||||||
255
internal/providers/providers.go
Normal file
255
internal/providers/providers.go
Normal file
|
|
@ -0,0 +1,255 @@
|
||||||
|
// Package providers manages the set of configured LLM providers and which one
|
||||||
|
// is active (globally or per session). API keys never leave this package.
|
||||||
|
package providers
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/crypto/byteclaw/internal/config"
|
||||||
|
"github.com/crypto/byteclaw/internal/llm"
|
||||||
|
"github.com/crypto/byteclaw/internal/store"
|
||||||
|
)
|
||||||
|
|
||||||
|
var ErrUnknown = errors.New("unknown provider")
|
||||||
|
var ErrNotConfigured = errors.New("provider has no API key configured")
|
||||||
|
|
||||||
|
const settingKey = "active_provider"
|
||||||
|
|
||||||
|
type entry struct {
|
||||||
|
name, baseURL, model, apiKey string
|
||||||
|
keyless bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// Info is the public view of a provider (no secrets).
|
||||||
|
type Info struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
BaseURL string `json:"base_url"`
|
||||||
|
DefaultModel string `json:"default_model"`
|
||||||
|
Configured bool `json:"configured"`
|
||||||
|
Active bool `json:"active"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type Active struct {
|
||||||
|
Provider string `json:"provider"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type Manager struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
entries map[string]*entry
|
||||||
|
active Active
|
||||||
|
store *store.Store
|
||||||
|
timeout time.Duration
|
||||||
|
temperature float64
|
||||||
|
HTTP *http.Client
|
||||||
|
// New builds a provider; replaceable in tests.
|
||||||
|
New func(baseURL, apiKey, model string) llm.Provider
|
||||||
|
}
|
||||||
|
|
||||||
|
// New builds the manager from config: every preset (key from llm.providers.<name>
|
||||||
|
// or the preset's env var) plus the provider configured in llm.*.
|
||||||
|
func New(cfg config.LLM, getenv func(string) string, st *store.Store) *Manager {
|
||||||
|
m := &Manager{entries: map[string]*entry{}, store: st, timeout: time.Duration(cfg.TimeoutSec) * time.Second,
|
||||||
|
temperature: cfg.Temperature, HTTP: &http.Client{Timeout: 15 * time.Second}}
|
||||||
|
m.New = func(baseURL, apiKey, model string) llm.Provider {
|
||||||
|
return llm.NewOpenAI(baseURL, apiKey, model, m.temperature, m.timeout)
|
||||||
|
}
|
||||||
|
for name, p := range config.Presets {
|
||||||
|
e := &entry{name: name, baseURL: p.BaseURL, model: p.DefaultModel, keyless: p.KeyEnv == ""}
|
||||||
|
if p.KeyEnv != "" {
|
||||||
|
e.apiKey = getenv(p.KeyEnv)
|
||||||
|
}
|
||||||
|
m.entries[name] = e
|
||||||
|
}
|
||||||
|
for name, o := range cfg.Providers {
|
||||||
|
name = strings.ToLower(name)
|
||||||
|
e, ok := m.entries[name]
|
||||||
|
if !ok {
|
||||||
|
e = &entry{name: name}
|
||||||
|
m.entries[name] = e
|
||||||
|
}
|
||||||
|
if o.BaseURL != "" {
|
||||||
|
e.baseURL = o.BaseURL
|
||||||
|
}
|
||||||
|
if o.APIKey != "" {
|
||||||
|
e.apiKey = o.APIKey
|
||||||
|
}
|
||||||
|
if o.Model != "" {
|
||||||
|
e.model = o.Model
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// The primary llm.* settings win for the configured provider.
|
||||||
|
e, ok := m.entries[cfg.Provider]
|
||||||
|
if !ok {
|
||||||
|
e = &entry{name: cfg.Provider, keyless: true}
|
||||||
|
m.entries[cfg.Provider] = e
|
||||||
|
}
|
||||||
|
e.baseURL, e.model = cfg.BaseURL, cfg.Model
|
||||||
|
if cfg.APIKey != "" {
|
||||||
|
e.apiKey = cfg.APIKey
|
||||||
|
}
|
||||||
|
m.active = Active{Provider: cfg.Provider, Model: cfg.Model}
|
||||||
|
if st != nil {
|
||||||
|
if raw, err := st.GetSetting(context.Background(), settingKey); err == nil {
|
||||||
|
var a Active
|
||||||
|
if json.Unmarshal([]byte(raw), &a) == nil && m.usable(a.Provider) == nil {
|
||||||
|
m.active = a
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return m
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) usable(name string) error {
|
||||||
|
e, ok := m.entries[name]
|
||||||
|
if !ok || e.baseURL == "" {
|
||||||
|
return fmt.Errorf("%w: %q", ErrUnknown, name)
|
||||||
|
}
|
||||||
|
if e.apiKey == "" && !e.keyless {
|
||||||
|
return fmt.Errorf("%w: %q", ErrNotConfigured, name)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) List() []Info {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
out := make([]Info, 0, len(m.entries))
|
||||||
|
for _, e := range m.entries {
|
||||||
|
if e.baseURL == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, Info{Name: e.name, BaseURL: e.baseURL, DefaultModel: e.model,
|
||||||
|
Configured: m.usable(e.name) == nil, Active: e.name == m.active.Provider})
|
||||||
|
}
|
||||||
|
sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) Active() Active {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
return m.active
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) BaseURL(name string) string {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
if e, ok := m.entries[name]; ok {
|
||||||
|
return e.baseURL
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate checks provider/model and fills the default model.
|
||||||
|
func (m *Manager) Validate(name, model string) (Active, error) {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
name = strings.ToLower(strings.TrimSpace(name))
|
||||||
|
if err := m.usable(name); err != nil {
|
||||||
|
return Active{}, err
|
||||||
|
}
|
||||||
|
if model = strings.TrimSpace(model); model == "" {
|
||||||
|
model = m.entries[name].model
|
||||||
|
}
|
||||||
|
if model == "" {
|
||||||
|
return Active{}, errors.New("model required")
|
||||||
|
}
|
||||||
|
return Active{Provider: name, Model: model}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetActive switches the server-wide default and persists it.
|
||||||
|
func (m *Manager) SetActive(ctx context.Context, name, model string) (Active, error) {
|
||||||
|
a, err := m.Validate(name, model)
|
||||||
|
if err != nil {
|
||||||
|
return a, err
|
||||||
|
}
|
||||||
|
m.mu.Lock()
|
||||||
|
m.active = a
|
||||||
|
m.mu.Unlock()
|
||||||
|
if m.store != nil {
|
||||||
|
b, _ := json.Marshal(a)
|
||||||
|
if err := m.store.SetSetting(ctx, settingKey, string(b)); err != nil {
|
||||||
|
return a, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return a, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build returns a provider client for name/model.
|
||||||
|
func (m *Manager) Build(name, model string) (llm.Provider, error) {
|
||||||
|
a, err := m.Validate(name, model)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
m.mu.RLock()
|
||||||
|
e := *m.entries[a.Provider]
|
||||||
|
m.mu.RUnlock()
|
||||||
|
return m.New(e.baseURL, e.apiKey, a.Model), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ForSession resolves the provider for a session (session override, else active).
|
||||||
|
func (m *Manager) ForSession(ctx context.Context, sessionID string) (llm.Provider, Active, error) {
|
||||||
|
a := m.Active()
|
||||||
|
if m.store != nil && sessionID != "" {
|
||||||
|
if ss, err := m.store.GetSession(ctx, sessionID); err == nil && ss.Provider != "" {
|
||||||
|
a = Active{Provider: ss.Provider, Model: ss.Model}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
p, err := m.Build(a.Provider, a.Model)
|
||||||
|
return p, a, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Models lists model ids from the provider's OpenAI-compatible GET /models (key added server-side).
|
||||||
|
func (m *Manager) Models(ctx context.Context, name string) ([]string, error) {
|
||||||
|
m.mu.RLock()
|
||||||
|
e, ok := m.entries[strings.ToLower(name)]
|
||||||
|
var ec entry
|
||||||
|
if ok {
|
||||||
|
ec = *e
|
||||||
|
}
|
||||||
|
m.mu.RUnlock()
|
||||||
|
if !ok || ec.baseURL == "" {
|
||||||
|
return nil, fmt.Errorf("%w: %q", ErrUnknown, name)
|
||||||
|
}
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(ec.baseURL, "/")+"/models", nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if ec.apiKey != "" {
|
||||||
|
req.Header.Set("Authorization", "Bearer "+ec.apiKey)
|
||||||
|
}
|
||||||
|
resp, err := m.HTTP.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode/100 != 2 {
|
||||||
|
return nil, fmt.Errorf("models endpoint returned HTTP %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
var body struct {
|
||||||
|
Data []struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
} `json:"data"`
|
||||||
|
}
|
||||||
|
if err := json.NewDecoder(io.LimitReader(resp.Body, 4<<20)).Decode(&body); err != nil {
|
||||||
|
return nil, fmt.Errorf("decode models: %w", err)
|
||||||
|
}
|
||||||
|
ids := make([]string, 0, len(body.Data))
|
||||||
|
for _, d := range body.Data {
|
||||||
|
if d.ID != "" {
|
||||||
|
ids = append(ids, d.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
sort.Strings(ids)
|
||||||
|
return ids, nil
|
||||||
|
}
|
||||||
100
internal/providers/providers_test.go
Normal file
100
internal/providers/providers_test.go
Normal file
|
|
@ -0,0 +1,100 @@
|
||||||
|
package providers
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/crypto/byteclaw/internal/config"
|
||||||
|
"github.com/crypto/byteclaw/internal/llm"
|
||||||
|
"github.com/crypto/byteclaw/internal/store"
|
||||||
|
)
|
||||||
|
|
||||||
|
type fake struct{ base, key, model string }
|
||||||
|
|
||||||
|
func (fake) Chat(context.Context, llm.Request, llm.DeltaFunc) (llm.Response, error) {
|
||||||
|
return llm.Response{}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func newMgr(t *testing.T, env map[string]string) (*Manager, *store.Store) {
|
||||||
|
st, err := store.Open(filepath.Join(t.TempDir(), "p.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { st.Close() })
|
||||||
|
cfg := config.LLM{Provider: "openai", BaseURL: "https://api.openai.com/v1", Model: "gpt-4o-mini", APIKey: "sk-openai",
|
||||||
|
Providers: map[string]config.ProviderOverride{"gonkarouter": {APIKey: "sk-gonka"}}}
|
||||||
|
m := New(cfg, func(k string) string { return env[k] }, st)
|
||||||
|
m.New = func(b, k, mo string) llm.Provider { return fake{b, k, mo} }
|
||||||
|
return m, st
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestListAndSwitch(t *testing.T) {
|
||||||
|
m, st := newMgr(t, map[string]string{"DAHL_API_KEY": "dahl_k"})
|
||||||
|
got := map[string]Info{}
|
||||||
|
for _, i := range m.List() {
|
||||||
|
got[i.Name] = i
|
||||||
|
}
|
||||||
|
if !got["openai"].Active || !got["dahl"].Configured || !got["gonkarouter"].Configured || got["openrouter"].Configured || !got["ollama"].Configured {
|
||||||
|
t.Fatalf("%+v", got)
|
||||||
|
}
|
||||||
|
if _, err := m.SetActive(context.Background(), "openrouter", ""); !errors.Is(err, ErrNotConfigured) {
|
||||||
|
t.Fatalf("want not configured, got %v", err)
|
||||||
|
}
|
||||||
|
if _, err := m.SetActive(context.Background(), "nope", ""); !errors.Is(err, ErrUnknown) {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
a, err := m.SetActive(context.Background(), "Dahl", "")
|
||||||
|
if err != nil || a.Provider != "dahl" || a.Model != "MiniMaxAI/MiniMax-M2.7" {
|
||||||
|
t.Fatalf("%+v %v", a, err)
|
||||||
|
}
|
||||||
|
p, _ := m.Build("dahl", "zai-org/GLM-5.3-Flash")
|
||||||
|
if f := p.(fake); f.key != "dahl_k" || f.base != "https://inference.dahl.global/v1" || f.model != "zai-org/GLM-5.3-Flash" {
|
||||||
|
t.Fatalf("%+v", f)
|
||||||
|
}
|
||||||
|
// persisted across restarts
|
||||||
|
m2 := New(config.LLM{Provider: "openai", BaseURL: "x", Model: "y"}, func(k string) string { return map[string]string{"DAHL_API_KEY": "dahl_k"}[k] }, st)
|
||||||
|
if m2.Active().Provider != "dahl" {
|
||||||
|
t.Fatalf("not persisted: %+v", m2.Active())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSessionOverride(t *testing.T) {
|
||||||
|
m, st := newMgr(t, nil)
|
||||||
|
ctx := context.Background()
|
||||||
|
ss, _ := st.CreateSession(ctx, "")
|
||||||
|
st.SetSessionProvider(ctx, ss.ID, "gonkarouter", "MiniMaxAI/MiniMax-M2.7")
|
||||||
|
p, a, err := m.ForSession(ctx, ss.ID)
|
||||||
|
if err != nil || a.Provider != "gonkarouter" || p.(fake).key != "sk-gonka" {
|
||||||
|
t.Fatalf("%+v %v", a, err)
|
||||||
|
}
|
||||||
|
other, _ := st.CreateSession(ctx, "")
|
||||||
|
if _, a, _ := m.ForSession(ctx, other.ID); a.Provider != "openai" {
|
||||||
|
t.Fatalf("%+v", a)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelsUsesServerSideKey(t *testing.T) {
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.URL.Path != "/v1/models" || r.Header.Get("Authorization") != "Bearer sk-gonka" {
|
||||||
|
http.Error(w, "no", 401)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Fprint(w, `{"data":[{"id":"b-model"},{"id":"a-model"}]}`)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
m, _ := newMgr(t, nil)
|
||||||
|
m.entries["gonkarouter"].baseURL = srv.URL + "/v1"
|
||||||
|
ids, err := m.Models(context.Background(), "gonkarouter")
|
||||||
|
if err != nil || len(ids) != 2 || ids[0] != "a-model" {
|
||||||
|
t.Fatalf("%v %v", ids, err)
|
||||||
|
}
|
||||||
|
m.entries["gonkarouter"].apiKey = "wrong"
|
||||||
|
if _, err := m.Models(context.Background(), "gonkarouter"); err == nil {
|
||||||
|
t.Fatal("expected error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -18,8 +18,11 @@ import (
|
||||||
var ErrNotFound = errors.New("not found")
|
var ErrNotFound = errors.New("not found")
|
||||||
|
|
||||||
type Session struct {
|
type Session struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
Title string `json:"title"`
|
Title string `json:"title"`
|
||||||
|
// Provider/Model override the server-wide active provider for this session (empty = default).
|
||||||
|
Provider string `json:"provider,omitempty"`
|
||||||
|
Model string `json:"model,omitempty"`
|
||||||
CreatedAt time.Time `json:"created_at"`
|
CreatedAt time.Time `json:"created_at"`
|
||||||
UpdatedAt time.Time `json:"updated_at"`
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
}
|
}
|
||||||
|
|
@ -45,15 +48,46 @@ CREATE TABLE IF NOT EXISTS messages(
|
||||||
tool_calls TEXT NOT NULL DEFAULT '', tool_call_id TEXT NOT NULL DEFAULT '',
|
tool_calls TEXT NOT NULL DEFAULT '', tool_call_id TEXT NOT NULL DEFAULT '',
|
||||||
name TEXT NOT NULL DEFAULT '', created_at INTEGER NOT NULL);
|
name TEXT NOT NULL DEFAULT '', created_at INTEGER NOT NULL);
|
||||||
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, id);
|
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, id);
|
||||||
|
CREATE TABLE IF NOT EXISTS summaries(
|
||||||
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
|
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
||||||
|
upto_message_id INTEGER NOT NULL, content TEXT NOT NULL, created_at INTEGER NOT NULL);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_summaries_session ON summaries(session_id, id);
|
||||||
|
CREATE TABLE IF NOT EXISTS settings(key TEXT PRIMARY KEY, value TEXT NOT NULL);
|
||||||
`
|
`
|
||||||
|
|
||||||
|
// migrations add columns to databases created by older versions.
|
||||||
|
var migrations = []struct{ table, column, ddl string }{
|
||||||
|
{"sessions", "provider", "ALTER TABLE sessions ADD COLUMN provider TEXT NOT NULL DEFAULT ''"},
|
||||||
|
{"sessions", "model", "ALTER TABLE sessions ADD COLUMN model TEXT NOT NULL DEFAULT ''"},
|
||||||
|
{"messages", "cancelled", "ALTER TABLE messages ADD COLUMN cancelled INTEGER NOT NULL DEFAULT 0"},
|
||||||
|
}
|
||||||
|
|
||||||
|
func migrate(db *sql.DB) error {
|
||||||
|
if _, err := db.Exec(schema); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
for _, m := range migrations {
|
||||||
|
var n int
|
||||||
|
if err := db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info(?) WHERE name=?`, m.table, m.column).Scan(&n); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if n == 0 {
|
||||||
|
if _, err := db.Exec(m.ddl); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
func Open(path string) (*Store, error) {
|
func Open(path string) (*Store, error) {
|
||||||
db, err := sql.Open("sqlite", "file:"+path+"?_pragma=busy_timeout(5000)")
|
db, err := sql.Open("sqlite", "file:"+path+"?_pragma=busy_timeout(5000)")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
db.SetMaxOpenConns(1) // SQLite: serialize writers, avoids SQLITE_BUSY
|
db.SetMaxOpenConns(1) // SQLite: serialize writers, avoids SQLITE_BUSY
|
||||||
if _, err := db.Exec(schema); err != nil {
|
if err := migrate(db); err != nil {
|
||||||
db.Close()
|
db.Close()
|
||||||
return nil, fmt.Errorf("migrate: %w", err)
|
return nil, fmt.Errorf("migrate: %w", err)
|
||||||
}
|
}
|
||||||
|
|
@ -79,7 +113,7 @@ func (s *Store) CreateSession(ctx context.Context, title string) (Session, error
|
||||||
func (s *Store) GetSession(ctx context.Context, id string) (Session, error) {
|
func (s *Store) GetSession(ctx context.Context, id string) (Session, error) {
|
||||||
var ss Session
|
var ss Session
|
||||||
var c, u int64
|
var c, u int64
|
||||||
err := s.db.QueryRowContext(ctx, `SELECT id,title,created_at,updated_at FROM sessions WHERE id=?`, id).Scan(&ss.ID, &ss.Title, &c, &u)
|
err := s.db.QueryRowContext(ctx, `SELECT id,title,provider,model,created_at,updated_at FROM sessions WHERE id=?`, id).Scan(&ss.ID, &ss.Title, &ss.Provider, &ss.Model, &c, &u)
|
||||||
if errors.Is(err, sql.ErrNoRows) {
|
if errors.Is(err, sql.ErrNoRows) {
|
||||||
return ss, ErrNotFound
|
return ss, ErrNotFound
|
||||||
}
|
}
|
||||||
|
|
@ -88,7 +122,7 @@ func (s *Store) GetSession(ctx context.Context, id string) (Session, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *Store) ListSessions(ctx context.Context, limit int) ([]Session, error) {
|
func (s *Store) ListSessions(ctx context.Context, limit int) ([]Session, error) {
|
||||||
rows, err := s.db.QueryContext(ctx, `SELECT id,title,created_at,updated_at FROM sessions ORDER BY updated_at DESC LIMIT ?`, limit)
|
rows, err := s.db.QueryContext(ctx, `SELECT id,title,provider,model,created_at,updated_at FROM sessions ORDER BY updated_at DESC LIMIT ?`, limit)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
@ -97,7 +131,7 @@ func (s *Store) ListSessions(ctx context.Context, limit int) ([]Session, error)
|
||||||
for rows.Next() {
|
for rows.Next() {
|
||||||
var ss Session
|
var ss Session
|
||||||
var c, u int64
|
var c, u int64
|
||||||
if err := rows.Scan(&ss.ID, &ss.Title, &c, &u); err != nil {
|
if err := rows.Scan(&ss.ID, &ss.Title, &ss.Provider, &ss.Model, &c, &u); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
ss.CreatedAt, ss.UpdatedAt = time.UnixMilli(c).UTC(), time.UnixMilli(u).UTC()
|
ss.CreatedAt, ss.UpdatedAt = time.UnixMilli(c).UTC(), time.UnixMilli(u).UTC()
|
||||||
|
|
@ -138,8 +172,8 @@ func (s *Store) AppendMessage(ctx context.Context, sessionID string, m llm.Messa
|
||||||
return 0, err
|
return 0, err
|
||||||
}
|
}
|
||||||
defer tx.Rollback()
|
defer tx.Rollback()
|
||||||
res, err := tx.ExecContext(ctx, `INSERT INTO messages(session_id,role,content,tool_calls,tool_call_id,name,created_at) VALUES(?,?,?,?,?,?,?)`,
|
res, err := tx.ExecContext(ctx, `INSERT INTO messages(session_id,role,content,tool_calls,tool_call_id,name,cancelled,created_at) VALUES(?,?,?,?,?,?,?,?)`,
|
||||||
sessionID, string(m.Role), m.Content, tc, m.ToolCallID, m.Name, now)
|
sessionID, string(m.Role), m.Content, tc, m.ToolCallID, m.Name, m.Cancelled, now)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, err
|
return 0, err
|
||||||
}
|
}
|
||||||
|
|
@ -152,11 +186,16 @@ func (s *Store) AppendMessage(ctx context.Context, sessionID string, m llm.Messa
|
||||||
|
|
||||||
// Messages returns the last `limit` messages of a session in chronological order (limit<=0: all).
|
// Messages returns the last `limit` messages of a session in chronological order (limit<=0: all).
|
||||||
func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]StoredMessage, error) {
|
func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]StoredMessage, error) {
|
||||||
|
return s.MessagesAfter(ctx, sessionID, 0, limit)
|
||||||
|
}
|
||||||
|
|
||||||
|
// MessagesAfter returns the last `limit` messages with id > afterID (limit<=0: all), oldest first.
|
||||||
|
func (s *Store) MessagesAfter(ctx context.Context, sessionID string, afterID int64, limit int) ([]StoredMessage, error) {
|
||||||
if limit <= 0 {
|
if limit <= 0 {
|
||||||
limit = -1
|
limit = -1
|
||||||
}
|
}
|
||||||
rows, err := s.db.QueryContext(ctx, `SELECT * FROM (SELECT id,role,content,tool_calls,tool_call_id,name,created_at
|
rows, err := s.db.QueryContext(ctx, `SELECT * FROM (SELECT id,role,content,tool_calls,tool_call_id,name,cancelled,created_at
|
||||||
FROM messages WHERE session_id=? ORDER BY id DESC LIMIT ?) ORDER BY id ASC`, sessionID, limit)
|
FROM messages WHERE session_id=? AND id>? ORDER BY id DESC LIMIT ?) ORDER BY id ASC`, sessionID, afterID, limit)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
@ -166,7 +205,7 @@ func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]St
|
||||||
var m StoredMessage
|
var m StoredMessage
|
||||||
var role, tc string
|
var role, tc string
|
||||||
var c int64
|
var c int64
|
||||||
if err := rows.Scan(&m.ID, &role, &m.Content, &tc, &m.ToolCallID, &m.Name, &c); err != nil {
|
if err := rows.Scan(&m.ID, &role, &m.Content, &tc, &m.ToolCallID, &m.Name, &m.Cancelled, &c); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
m.Role, m.CreatedAt = llm.Role(role), time.UnixMilli(c).UTC()
|
m.Role, m.CreatedAt = llm.Role(role), time.UnixMilli(c).UTC()
|
||||||
|
|
@ -179,3 +218,58 @@ func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]St
|
||||||
}
|
}
|
||||||
return out, rows.Err()
|
return out, rows.Err()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (s *Store) SetSessionProvider(ctx context.Context, id, provider, model string) error {
|
||||||
|
res, err := s.db.ExecContext(ctx, `UPDATE sessions SET provider=?, model=? WHERE id=?`, provider, model, id)
|
||||||
|
return affected(res, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- summaries
|
||||||
|
|
||||||
|
type Summary struct {
|
||||||
|
ID int64 `json:"id"`
|
||||||
|
SessionID string `json:"session_id"`
|
||||||
|
UptoMessageID int64 `json:"upto_message_id"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Store) AddSummary(ctx context.Context, sessionID string, upto int64, content string) (Summary, error) {
|
||||||
|
now := time.Now().UTC()
|
||||||
|
res, err := s.db.ExecContext(ctx, `INSERT INTO summaries(session_id,upto_message_id,content,created_at) VALUES(?,?,?,?)`,
|
||||||
|
sessionID, upto, content, now.UnixMilli())
|
||||||
|
if err != nil {
|
||||||
|
return Summary{}, err
|
||||||
|
}
|
||||||
|
id, _ := res.LastInsertId()
|
||||||
|
return Summary{ID: id, SessionID: sessionID, UptoMessageID: upto, Content: content, CreatedAt: now}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// LatestSummary returns the newest summary of a session or ErrNotFound.
|
||||||
|
func (s *Store) LatestSummary(ctx context.Context, sessionID string) (Summary, error) {
|
||||||
|
var sm Summary
|
||||||
|
var c int64
|
||||||
|
err := s.db.QueryRowContext(ctx, `SELECT id,session_id,upto_message_id,content,created_at FROM summaries
|
||||||
|
WHERE session_id=? ORDER BY id DESC LIMIT 1`, sessionID).Scan(&sm.ID, &sm.SessionID, &sm.UptoMessageID, &sm.Content, &c)
|
||||||
|
if errors.Is(err, sql.ErrNoRows) {
|
||||||
|
return sm, ErrNotFound
|
||||||
|
}
|
||||||
|
sm.CreatedAt = time.UnixMilli(c).UTC()
|
||||||
|
return sm, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---- settings (key/value)
|
||||||
|
|
||||||
|
func (s *Store) GetSetting(ctx context.Context, key string) (string, error) {
|
||||||
|
var v string
|
||||||
|
err := s.db.QueryRowContext(ctx, `SELECT value FROM settings WHERE key=?`, key).Scan(&v)
|
||||||
|
if errors.Is(err, sql.ErrNoRows) {
|
||||||
|
return "", ErrNotFound
|
||||||
|
}
|
||||||
|
return v, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Store) SetSetting(ctx context.Context, key, value string) error {
|
||||||
|
_, err := s.db.ExecContext(ctx, `INSERT INTO settings(key,value) VALUES(?,?) ON CONFLICT(key) DO UPDATE SET value=excluded.value`, key, value)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -48,3 +48,43 @@ func TestStoreRoundTrip(t *testing.T) {
|
||||||
t.Fatal("cascade delete failed")
|
t.Fatal("cascade delete failed")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestSummariesSettingsProviderAndCancelled(t *testing.T) {
|
||||||
|
s, err := Open(filepath.Join(t.TempDir(), "t.db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer s.Close()
|
||||||
|
ctx := context.Background()
|
||||||
|
ss, _ := s.CreateSession(ctx, "x")
|
||||||
|
if _, err := s.LatestSummary(ctx, ss.ID); !errors.Is(err, ErrNotFound) {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
id1, _ := s.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleUser, Content: "a"})
|
||||||
|
s.AppendMessage(ctx, ss.ID, llm.Message{Role: llm.RoleAssistant, Content: "part", Cancelled: true})
|
||||||
|
s.AddSummary(ctx, ss.ID, id1, "old")
|
||||||
|
s.AddSummary(ctx, ss.ID, id1, "new")
|
||||||
|
sm, _ := s.LatestSummary(ctx, ss.ID)
|
||||||
|
if sm.Content != "new" || sm.UptoMessageID != id1 {
|
||||||
|
t.Fatalf("%+v", sm)
|
||||||
|
}
|
||||||
|
after, _ := s.MessagesAfter(ctx, ss.ID, id1, 0)
|
||||||
|
if len(after) != 1 || !after[0].Cancelled || after[0].Content != "part" {
|
||||||
|
t.Fatalf("%+v", after)
|
||||||
|
}
|
||||||
|
if all, _ := s.Messages(ctx, ss.ID, 0); len(all) != 2 {
|
||||||
|
t.Fatal("original messages must be kept")
|
||||||
|
}
|
||||||
|
if err := s.SetSessionProvider(ctx, ss.ID, "dahl", "m"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got, _ := s.GetSession(ctx, ss.ID)
|
||||||
|
if got.Provider != "dahl" || got.Model != "m" {
|
||||||
|
t.Fatalf("%+v", got)
|
||||||
|
}
|
||||||
|
s.SetSetting(ctx, "k", "v1")
|
||||||
|
s.SetSetting(ctx, "k", "v2")
|
||||||
|
if v, _ := s.GetSetting(ctx, "k"); v != "v2" {
|
||||||
|
t.Fatal(v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -107,14 +107,14 @@ func (s *Shell) Parse(line string) ([]step, error) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
var (
|
var (
|
||||||
steps []step
|
steps []step
|
||||||
cur []string // current command tokens
|
cur []string // current command tokens
|
||||||
pipe [][]string
|
pipe [][]string
|
||||||
tok strings.Builder
|
tok strings.Builder
|
||||||
inTok bool
|
inTok bool
|
||||||
quote rune
|
quote rune
|
||||||
nextOp = ""
|
nextOp = ""
|
||||||
runes = []rune(line)
|
runes = []rune(line)
|
||||||
)
|
)
|
||||||
flushTok := func() {
|
flushTok := func() {
|
||||||
if inTok {
|
if inTok {
|
||||||
|
|
@ -258,7 +258,9 @@ func (s *Shell) env() []string {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *Shell) Call(ctx context.Context, raw json.RawMessage) (string, error) {
|
func (s *Shell) Call(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||||
var a struct{ Command string `json:"command"` }
|
var a struct {
|
||||||
|
Command string `json:"command"`
|
||||||
|
}
|
||||||
if err := json.Unmarshal(raw, &a); err != nil {
|
if err := json.Unmarshal(raw, &a); err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue