diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 2c74ae8..e696a33 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -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 } diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go index 44d90fa..b56e8ab 100644 --- a/internal/agent/agent_test.go +++ b/internal/agent/agent_test.go @@ -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) + } +} diff --git a/internal/agent/summarize.go b/internal/agent/summarize.go new file mode 100644 index 0000000..5afd7c1 --- /dev/null +++ b/internal/agent/summarize.go @@ -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 +} diff --git a/internal/config/config.go b/internal/config/config.go index ca1e0f5..3ae6924 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -20,12 +20,32 @@ type LLM struct { Model string `yaml:"model" json:"model"` TimeoutSec int `yaml:"timeout_sec" json:"timeout_sec"` 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 { - SystemPrompt string `yaml:"system_prompt" json:"system_prompt"` - MaxIterations int `yaml:"max_iterations" json:"max_iterations"` - HistoryLimit int `yaml:"history_limit" json:"history_limit"` + SystemPrompt string `yaml:"system_prompt" json:"system_prompt"` + MaxIterations int `yaml:"max_iterations" json:"max_iterations"` + 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 { @@ -80,6 +100,7 @@ func Default() Config { SystemPrompt: "You are ByteClaw, a concise and helpful self-hosted assistant. Use tools when they help.", MaxIterations: 8, HistoryLimit: 50, + Summarize: Summarize{Enabled: true, MaxMessages: 40, MaxTokens: 12000, KeepRecent: 10}, }, Tools: Tools{ WebFetch: WebFetch{Enabled: true, MaxBytes: 512 * 1024, TimeoutSec: 15}, diff --git a/internal/llm/llm.go b/internal/llm/llm.go index 5c0772b..74b1baf 100644 --- a/internal/llm/llm.go +++ b/internal/llm/llm.go @@ -24,6 +24,8 @@ type Message struct { ToolCalls []ToolCall `json:"tool_calls,omitempty"` ToolCallID string `json:"tool_call_id,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. diff --git a/internal/providers/providers.go b/internal/providers/providers.go new file mode 100644 index 0000000..f8ce31f --- /dev/null +++ b/internal/providers/providers.go @@ -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. +// 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 +} diff --git a/internal/providers/providers_test.go b/internal/providers/providers_test.go new file mode 100644 index 0000000..9b71db2 --- /dev/null +++ b/internal/providers/providers_test.go @@ -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") + } +} diff --git a/internal/store/store.go b/internal/store/store.go index 0930090..5656a2b 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -18,8 +18,11 @@ import ( var ErrNotFound = errors.New("not found") type Session struct { - ID string `json:"id"` - Title string `json:"title"` + ID string `json:"id"` + 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"` 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 '', name TEXT NOT NULL DEFAULT '', created_at INTEGER NOT NULL); 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) { db, err := sql.Open("sqlite", "file:"+path+"?_pragma=busy_timeout(5000)") if err != nil { return nil, err } db.SetMaxOpenConns(1) // SQLite: serialize writers, avoids SQLITE_BUSY - if _, err := db.Exec(schema); err != nil { + if err := migrate(db); err != nil { db.Close() 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) { var ss Session 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) { 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) { - 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 { return nil, err } @@ -97,7 +131,7 @@ func (s *Store) ListSessions(ctx context.Context, limit int) ([]Session, error) for rows.Next() { var ss Session 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 } 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 } defer tx.Rollback() - res, err := tx.ExecContext(ctx, `INSERT INTO messages(session_id,role,content,tool_calls,tool_call_id,name,created_at) VALUES(?,?,?,?,?,?,?)`, - sessionID, string(m.Role), m.Content, tc, m.ToolCallID, m.Name, now) + 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, m.Cancelled, now) if err != nil { 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). 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 { limit = -1 } - rows, err := s.db.QueryContext(ctx, `SELECT * FROM (SELECT id,role,content,tool_calls,tool_call_id,name,created_at - FROM messages WHERE session_id=? ORDER BY id DESC LIMIT ?) ORDER BY id ASC`, sessionID, limit) + 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=? AND id>? ORDER BY id DESC LIMIT ?) ORDER BY id ASC`, sessionID, afterID, limit) if err != nil { return nil, err } @@ -166,7 +205,7 @@ func (s *Store) Messages(ctx context.Context, sessionID string, limit int) ([]St var m StoredMessage var role, tc string 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 } 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() } + +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 +} diff --git a/internal/store/store_test.go b/internal/store/store_test.go index 9c74829..7070bad 100644 --- a/internal/store/store_test.go +++ b/internal/store/store_test.go @@ -48,3 +48,43 @@ func TestStoreRoundTrip(t *testing.T) { 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) + } +} diff --git a/internal/tools/shell.go b/internal/tools/shell.go index 7433ef8..8f8eeef 100644 --- a/internal/tools/shell.go +++ b/internal/tools/shell.go @@ -107,14 +107,14 @@ func (s *Shell) Parse(line string) ([]step, error) { } } var ( - steps []step - cur []string // current command tokens - pipe [][]string - tok strings.Builder - inTok bool - quote rune - nextOp = "" - runes = []rune(line) + steps []step + cur []string // current command tokens + pipe [][]string + tok strings.Builder + inTok bool + quote rune + nextOp = "" + runes = []rune(line) ) flushTok := func() { if inTok { @@ -258,7 +258,9 @@ func (s *Shell) env() []string { } 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 { return "", err }