b550c69 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat: steering (#605)

b550c69Kujtim Hoxha1 year ago

internal/llm/agent/agent.go+61−1

compare with main

unifiedsplit

@@ -60,6 +60,8 @@ type Service interface {

60 IsBusy() bool

61 Summarize(ctx context.Context, sessionID string) error

62 UpdateModel() error

63+ QueuedPrompts(sessionID string) int

64+ ClearQueue(sessionID string)

63}

65}

65type agent struct {

67type agent struct {

@@ -79,6 +81,8 @@ type agent struct {

79 summarizeProviderID string

81 summarizeProviderID string

81 activeRequests *csync.Map[string, context.CancelFunc]

83 activeRequests *csync.Map[string, context.CancelFunc]

84+

85+ promptQueue *csync.Map[string, []string]

82}

86}

84var agentPromptMap = map[string]prompt.PromptID{

88var agentPromptMap = map[string]prompt.PromptID{

@@ -228,6 +232,7 @@ func NewAgent(

228 summarizeProviderID: string(providerCfg.ID),

232 summarizeProviderID: string(providerCfg.ID),

229 activeRequests: csync.NewMap[string, context.CancelFunc](),

233 activeRequests: csync.NewMap[string, context.CancelFunc](),

230 tools: csync.NewLazySlice(toolFn),

234 tools: csync.NewLazySlice(toolFn),

235+ promptQueue: csync.NewMap[string, []string](),

231 }, nil

236 }, nil

232}

237}

233

238

@@ -247,6 +252,11 @@ func (a *agent) Cancel(sessionID string) {

247 slog.Info("Summarize cancellation initiated", "session_id", sessionID)

252 slog.Info("Summarize cancellation initiated", "session_id", sessionID)

248 cancel()

253 cancel()

249 }

254 }

255+

256+ if a.QueuedPrompts(sessionID) > 0 {

257+ slog.Info("Clearing queued prompts", "session_id", sessionID)

258+ a.promptQueue.Del(sessionID)

259+ }

250}

260}

251

261

252func (a *agent) IsBusy() bool {

262func (a *agent) IsBusy() bool {

@@ -265,6 +275,14 @@ func (a *agent) IsSessionBusy(sessionID string) bool {

265 return busy

275 return busy

266}

276}

267

277

278+func (a *agent) QueuedPrompts(sessionID string) int {

279+ l, ok := a.promptQueue.Get(sessionID)

280+ if !ok {

281+ return 0

282+ }

283+ return len(l)

284+}

285+

268func (a *agent) generateTitle(ctx context.Context, sessionID string, content string) error {

286func (a *agent) generateTitle(ctx context.Context, sessionID string, content string) error {

269 if content == "" {

287 if content == "" {

270 return nil

288 return nil

@@ -327,7 +345,13 @@ func (a *agent) Run(ctx context.Context, sessionID string, content string, attac

327 }

345 }

328 events := make(chan AgentEvent)

346 events := make(chan AgentEvent)

329 if a.IsSessionBusy(sessionID) {

347 if a.IsSessionBusy(sessionID) {

348+ existing, ok := a.promptQueue.Get(sessionID)

349+ if !ok {

350+ existing = []string{}

351+ }

352+ existing = append(existing, content)

353+ a.promptQueue.Set(sessionID, existing)

330− return nil, ErrSessionBusy

354+ return nil, nil

331 }

355 }

332

356

333 genCtx, cancel := context.WithCancel(ctx)

357 genCtx, cancel := context.WithCancel(ctx)

@@ -422,7 +446,36 @@ func (a *agent) processGeneration(ctx context.Context, sessionID, content string

422 if (agentMessage.FinishReason() == message.FinishReasonToolUse) && toolResults != nil {

446 if (agentMessage.FinishReason() == message.FinishReasonToolUse) && toolResults != nil {

423 // We are not done, we need to respond with the tool response

447 // We are not done, we need to respond with the tool response

424 msgHistory = append(msgHistory, agentMessage, *toolResults)

448 msgHistory = append(msgHistory, agentMessage, *toolResults)

449+ // If there are queued prompts, process the next one

450+ nextPrompt, ok := a.promptQueue.Take(sessionID)

451+ if ok {

452+ for _, prompt := range nextPrompt {

453+ // Create a new user message for the queued prompt

454+ userMsg, err := a.createUserMessage(ctx, sessionID, prompt, nil)

455+ if err != nil {

456+ return a.err(fmt.Errorf("failed to create user message for queued prompt: %w", err))

457+ }

458+ // Append the new user message to the conversation history

459+ msgHistory = append(msgHistory, userMsg)

460+ }

461+ }

462+

425 continue

463 continue

464+ } else if agentMessage.FinishReason() == message.FinishReasonEndTurn {

465+ queuePrompts, ok := a.promptQueue.Take(sessionID)

466+ if ok {

467+ for _, prompt := range queuePrompts {

468+ if prompt == "" {

469+ continue

470+ }

471+ userMsg, err := a.createUserMessage(ctx, sessionID, prompt, nil)

472+ if err != nil {

473+ return a.err(fmt.Errorf("failed to create user message for queued prompt: %w", err))

474+ }

475+ msgHistory = append(msgHistory, userMsg)

476+ }

477+ continue

478+ }

426 }

479 }

427 if agentMessage.FinishReason() == "" {

480 if agentMessage.FinishReason() == "" {

428 // Kujtim: could not track down where this is happening but this means its cancelled

481 // Kujtim: could not track down where this is happening but this means its cancelled

@@ -852,6 +905,13 @@ func (a *agent) Summarize(ctx context.Context, sessionID string) error {

852 return nil

905 return nil

853}

906}

854

907

908+func (a *agent) ClearQueue(sessionID string) {

909+ if a.QueuedPrompts(sessionID) > 0 {

910+ slog.Info("Clearing queued prompts", "session_id", sessionID)

911+ a.promptQueue.Del(sessionID)

912+ }

913+}

914+

855func (a *agent) CancelAll() {

915func (a *agent) CancelAll() {

856 if !a.IsBusy() {

916 if !a.IsBusy() {

857 return

917 return

internal/tui/components/chat/chat.go+18−1

@@ -18,6 +18,7 @@ import (

18 "github.com/charmbracelet/crush/internal/tui/exp/list"

19 "github.com/charmbracelet/crush/internal/tui/styles"

20 "github.com/charmbracelet/crush/internal/tui/util"

21+ "github.com/charmbracelet/lipgloss/v2"

21)

22)

23type SendMsg struct {

24type SendMsg struct {

@@ -198,13 +199,29 @@ func (m *messageListCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

198// View renders the message list or an initial screen if empty.

199// View renders the message list or an initial screen if empty.

199func (m *messageListCmp) View() string {

200func (m *messageListCmp) View() string {

200 t := styles.CurrentTheme()

201 t := styles.CurrentTheme()

201− return t.S().Base.

202+ listView := t.S().Base.

202 Padding(1, 1, 0, 1).

203 Padding(1, 1, 0, 1).

203 Width(m.width).

204 Width(m.width).

204 Height(m.height).

205 Height(m.height).

205 Render(

206 Render(

206 m.listCmp.View(),

207 m.listCmp.View(),

207 )

208 )

209+

210+ if m.app.CoderAgent != nil && m.app.CoderAgent.QueuedPrompts(m.session.ID) > 0 {

211+ queue := m.app.CoderAgent.QueuedPrompts(m.session.ID)

212+ queuePill := queuePill(queue, t)

213+ layers := []*lipgloss.Layer{

214+ lipgloss.NewLayer(listView),

215+ lipgloss.NewLayer(

216+ queuePill,

217+ ).X(4).Y(m.height - 3),

218+ }

219+ canvas := lipgloss.NewCanvas(

220+ layers...,

221+ )

222+ return canvas.Render()

223+ }

224+ return listView

208}

225}

209

226

210func (m *messageListCmp) handlePermissionRequest(permission permission.PermissionNotification) tea.Cmd {

227func (m *messageListCmp) handlePermissionRequest(permission permission.PermissionNotification) tea.Cmd {

internal/tui/components/chat/editor/editor.go−7

@@ -138,13 +138,6 @@ func (m *editorCmp) Init() tea.Cmd {

138}

139

140func (m *editorCmp) send() tea.Cmd {

141− if m.app.CoderAgent == nil {

142− return util.ReportError(fmt.Errorf("coder agent is not initialized"))

143− }

144− if m.app.CoderAgent.IsSessionBusy(m.session.ID) {

145− return util.ReportWarn("Agent is working, please wait...")

146− }

147−

148 value := m.textarea.Value()

141 value := m.textarea.Value()

149 value = strings.TrimSpace(value)

142 value = strings.TrimSpace(value)

150

143

internal/tui/components/chat/queue.gonew file+28

@@ -0,0 +1,28 @@

1+package chat

3+import (

4+ "fmt"

5+ "strings"

7+ "github.com/charmbracelet/crush/internal/tui/styles"

8+ "github.com/charmbracelet/lipgloss/v2"

9+)

10+

11+func queuePill(queue int, t *styles.Theme) string {

12+ if queue <= 0 {

13+ return ""

14+ }

15+ triangles := styles.ForegroundGrad("▶▶▶▶▶▶▶▶▶", false, t.RedDark, t.Accent)

16+ if queue < 10 {

17+ triangles = triangles[:queue]

18+ }

19+

20+ allTriangles := strings.Join(triangles, "")

21+

22+ return t.S().Base.

23+ BorderStyle(lipgloss.RoundedBorder()).

24+ BorderForeground(t.BgOverlay).

25+ PaddingLeft(1).

26+ PaddingRight(1).

27+ Render(fmt.Sprintf("%s %d Queued", allTriangles, queue))

28+}

internal/tui/page/chat/chat.go+10

@@ -653,6 +653,10 @@ func (p *chatPage) cancel() tea.Cmd {

653 return nil

654 }

655

656+ if p.app.CoderAgent != nil && p.app.CoderAgent.QueuedPrompts(p.session.ID) > 0 {

657+ p.app.CoderAgent.ClearQueue(p.session.ID)

658+ return nil

659+ }

656 p.isCanceling = true

660 p.isCanceling = true

657 return cancelTimerCmd()

661 return cancelTimerCmd()

658}

662}

@@ -828,6 +832,12 @@ func (p *chatPage) Help() help.KeyMap {

828 key.WithHelp("esc", "press again to cancel"),

832 key.WithHelp("esc", "press again to cancel"),

829 )

833 )

830 }

834 }

835+ if p.app.CoderAgent.QueuedPrompts(p.session.ID) > 0 {

836+ cancelBinding = key.NewBinding(

837+ key.WithKeys("esc"),

838+ key.WithHelp("esc", "clear queue"),

839+ )

840+ }

831 shortList = append(shortList, cancelBinding)

841 shortList = append(shortList, cancelBinding)

832 fullList = append(fullList,

842 fullList = append(fullList,

833 []key.Binding{

843 []key.Binding{

internal/tui/styles/theme.go+29−25

@@ -591,18 +591,18 @@ func Lighten(c color.Color, percent float64) color.Color {

591 }

592}

593

594−// ApplyForegroundGrad renders a given string with a horizontal gradient

595−// foreground.

596−func ApplyForegroundGrad(input string, color1, color2 color.Color) string {

594+func ForegroundGrad(input string, bold bool, color1, color2 color.Color) []string {

597 if input == "" {

595 if input == "" {

598− return ""

596+ return []string{""}

599 }

597 }

600−

598+ t := CurrentTheme()

601− var o strings.Builder

602 if len(input) == 1 {

599 if len(input) == 1 {

603− return lipgloss.NewStyle().Foreground(color1).Render(input)

600+ style := t.S().Base.Foreground(color1)

601+ if bold {

602+ style.Bold(true)

603+ }

604+ return []string{style.Render(input)}

604 }

605 }

605−

606 var clusters []string

607 gr := uniseg.NewGraphemes(input)

608 for gr.Next() {

@@ -611,9 +611,26 @@ func ApplyForegroundGrad(input string, color1, color2 color.Color) string {

611

612 ramp := blendColors(len(clusters), color1, color2)

613 for i, c := range ramp {

614− fmt.Fprint(&o, CurrentTheme().S().Base.Foreground(c).Render(clusters[i]))

614+ style := t.S().Base.Foreground(c)

615+ if bold {

616+ style.Bold(true)

617+ }

618+ clusters[i] = style.Render(clusters[i])

615 }

619 }

620+ return clusters

621+}

616

622

623+// ApplyForegroundGrad renders a given string with a horizontal gradient

624+// foreground.

625+func ApplyForegroundGrad(input string, color1, color2 color.Color) string {

626+ if input == "" {

627+ return ""

628+ }

629+ var o strings.Builder

630+ clusters := ForegroundGrad(input, false, color1, color2)

631+ for _, c := range clusters {

632+ fmt.Fprint(&o, c)

633+ }

617 return o.String()

634 return o.String()

618}

635}

619

636

@@ -623,24 +640,11 @@ func ApplyBoldForegroundGrad(input string, color1, color2 color.Color) string {

623 if input == "" {

640 if input == "" {

624 return ""

641 return ""

625 }

642 }

626− t := CurrentTheme()

627−

628 var o strings.Builder

643 var o strings.Builder

629− if len(input) == 1 {

644+ clusters := ForegroundGrad(input, true, color1, color2)

630− return t.S().Base.Bold(true).Foreground(color1).Render(input)

631− }

632−

633− var clusters []string

634− gr := uniseg.NewGraphemes(input)

635− for gr.Next() {

636− clusters = append(clusters, string(gr.Runes()))

637− }

638−

639− ramp := blendColors(len(clusters), color1, color2)

640− for i, c := range ramp {

645+ for _, c := range clusters {

641− fmt.Fprint(&o, t.S().Base.Bold(true).Foreground(c).Render(clusters[i]))

646+ fmt.Fprint(&o, c)

642 }

647 }

643−

644 return o.String()

648 return o.String()

645}

649}

646

650