3e42475 · crush · git.iko.soy

git.iko.soy/crush

Improve summary to keep context (#159)

* improve summary to keep context * improve loop * remove debug msg

commit3e424754b48862fdd941f5d6434abda989caaa21 authorKujtim Hoxha <kujtimii.h@gmail.com> dateThu 15 May 2025, 15:59 (1 year ago) parent4f0c1c6

13 files changed+105−54

compare with main

unifiedsplit

internal/db/db.go+1−1 internal/db/files.sql.go+1−1 internal/db/messages.sql.go+1−1 internal/db/migrations/20250515105448_add_summary_message_id.sql+9 internal/db/models.go+2−1 internal/db/querier.go+1−1 internal/db/sessions.sql.go+19−10 internal/db/sql/sessions.sql+3 internal/llm/agent/agent.go+46−14 internal/session/session.go+7−1 internal/tui/components/chat/list.go+11 internal/tui/components/chat/message.go+4 internal/tui/tui.go−24

internal/db/db.go+1−1

@@ -1,6 +1,6 @@

11// Code generated by sqlc. DO NOT EDIT.

22// versions:

3−// sqlc v1.27.0

3+// sqlc v1.29.0

55package db

internal/db/files.sql.go+1−1

@@ -1,6 +1,6 @@

11// Code generated by sqlc. DO NOT EDIT.

22// versions:

3−// sqlc v1.27.0

3+// sqlc v1.29.0

44// source: files.sql

66package db

internal/db/messages.sql.go+1−1

@@ -1,6 +1,6 @@

11// Code generated by sqlc. DO NOT EDIT.

22// versions:

3−// sqlc v1.27.0

3+// sqlc v1.29.0

44// source: messages.sql

66package db

internal/db/migrations/20250515105448_add_summary_message_id.sqlnew file+9

@@ -0,0 +1,9 @@

1+-- +goose Up

2+-- +goose StatementBegin

3+ALTER TABLE sessions ADD COLUMN summary_message_id TEXT;

4+-- +goose StatementEnd

6+-- +goose Down

7+-- +goose StatementBegin

8+ALTER TABLE sessions DROP COLUMN summary_message_id;

9+-- +goose StatementEnd

internal/db/models.go+2−1

@@ -1,6 +1,6 @@

11// Code generated by sqlc. DO NOT EDIT.

22// versions:

3−// sqlc v1.27.0

3+// sqlc v1.29.0

55package db

@@ -39,4 +39,5 @@ type Session struct {

3939 Cost float64 `json:"cost"`

4040 UpdatedAt int64 `json:"updated_at"`

4141 CreatedAt int64 `json:"created_at"`

42+ SummaryMessageID sql.NullString `json:"summary_message_id"`

4243}

internal/db/querier.go+1−1

@@ -1,6 +1,6 @@

11// Code generated by sqlc. DO NOT EDIT.

22// versions:

3−// sqlc v1.27.0

3+// sqlc v1.29.0

55package db

internal/db/sessions.sql.go+19−10

@@ -1,6 +1,6 @@

11// Code generated by sqlc. DO NOT EDIT.

22// versions:

3−// sqlc v1.27.0

3+// sqlc v1.29.0

44// source: sessions.sql

66package db

@@ -19,6 +19,7 @@ INSERT INTO sessions (

1919 prompt_tokens,

2020 completion_tokens,

2121 cost,

22+ summary_message_id,

2223 updated_at,

2324 created_at

2425) VALUES (

@@ -29,9 +30,10 @@ INSERT INTO sessions (

2930 ?,

3031 ?,

3132 ?,

33+ null,

3234 strftime('%s', 'now'),

3335 strftime('%s', 'now')

34−) RETURNING id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at

36+) RETURNING id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id

3537`

3638

3739type CreateSessionParams struct {

@@ -65,6 +67,7 @@ func (q *Queries) CreateSession(ctx context.Context, arg CreateSessionParams) (S

6567 &i.Cost,

6668 &i.UpdatedAt,

6769 &i.CreatedAt,

70+ &i.SummaryMessageID,

6871 )

6972 return i, err

7073}

@@ -80,7 +83,7 @@ func (q *Queries) DeleteSession(ctx context.Context, id string) error {

8083}

8184

8285const getSessionByID = `-- name: GetSessionByID :one

83−SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at

86+SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id

8487FROM sessions

8588WHERE id = ? LIMIT 1

8689`

@@ -98,12 +101,13 @@ func (q *Queries) GetSessionByID(ctx context.Context, id string) (Session, error

98101 &i.Cost,

99102 &i.UpdatedAt,

100103 &i.CreatedAt,

104+ &i.SummaryMessageID,

101105 )

102106 return i, err

103107}

104108

105109const listSessions = `-- name: ListSessions :many

106−SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at

110+SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id

107111FROM sessions

108112WHERE parent_session_id is NULL

109113ORDER BY created_at DESC

@@ -128,6 +132,7 @@ func (q *Queries) ListSessions(ctx context.Context) ([]Session, error) {

128132 &i.Cost,

129133 &i.UpdatedAt,

130134 &i.CreatedAt,

135+ &i.SummaryMessageID,

131136 ); err != nil {

132137 return nil, err

133138 }

@@ -148,17 +153,19 @@ SET

148153 title = ?,

149154 prompt_tokens = ?,

150155 completion_tokens = ?,

156+ summary_message_id = ?,

151157 cost = ?

152158WHERE id = ?

153−RETURNING id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at

159+RETURNING id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id

154160`

155161

156162type UpdateSessionParams struct {

157− Title string `json:"title"`

158− PromptTokens int64 `json:"prompt_tokens"`

159− CompletionTokens int64 `json:"completion_tokens"`

160− Cost float64 `json:"cost"`

161− ID string `json:"id"`

163+ Title string `json:"title"`

164+ PromptTokens int64 `json:"prompt_tokens"`

165+ CompletionTokens int64 `json:"completion_tokens"`

166+ SummaryMessageID sql.NullString `json:"summary_message_id"`

167+ Cost float64 `json:"cost"`

168+ ID string `json:"id"`

162169}

163170

164171func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (Session, error) {

@@ -166,6 +173,7 @@ func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (S

166173 arg.Title,

167174 arg.PromptTokens,

168175 arg.CompletionTokens,

176+ arg.SummaryMessageID,

169177 arg.Cost,

170178 arg.ID,

171179 )

@@ -180,6 +188,7 @@ func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (S

180188 &i.Cost,

181189 &i.UpdatedAt,

182190 &i.CreatedAt,

191+ &i.SummaryMessageID,

183192 )

184193 return i, err

185194}

internal/db/sql/sessions.sql+3

@@ -7,6 +7,7 @@ INSERT INTO sessions (

77 prompt_tokens,

88 completion_tokens,

99 cost,

10+ summary_message_id,

1011 updated_at,

1112 created_at

1213) VALUES (

@@ -17,6 +18,7 @@ INSERT INTO sessions (

1718 ?,

1819 ?,

1920 ?,

21+ null,

2022 strftime('%s', 'now'),

2123 strftime('%s', 'now')

2224) RETURNING *;

@@ -38,6 +40,7 @@ SET

3840 title = ?,

3941 prompt_tokens = ?,

4042 completion_tokens = ?,

43+ summary_message_id = ?,

4144 cost = ?

4245WHERE id = ?

4346RETURNING *;

internal/llm/agent/agent.go+46−14

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

66 "fmt"

77 "strings"

88 "sync"

9+ "time"

910

1011 "github.com/opencode-ai/opencode/internal/config"

1112 "github.com/opencode-ai/opencode/internal/llm/models"

@@ -245,6 +246,23 @@ func (a *agent) processGeneration(ctx context.Context, sessionID, content string

245246 }

246247 }()

247248 }

249+ session, err := a.sessions.Get(ctx, sessionID)

250+ if err != nil {

251+ return a.err(fmt.Errorf("failed to get session: %w", err))

252+ }

253+ if session.SummaryMessageID != "" {

254+ summaryMsgInex := -1

255+ for i, msg := range msgs {

256+ if msg.ID == session.SummaryMessageID {

257+ summaryMsgInex = i

258+ break

259+ }

260+ }

261+ if summaryMsgInex != -1 {

262+ msgs = msgs[summaryMsgInex:]

263+ msgs[0].Role = message.User

264+ }

265+ }

248266

249267 userMsg, err := a.createUserMessage(ctx, sessionID, content, attachmentParts)

250268 if err != nil {

@@ -614,37 +632,51 @@ func (a *agent) Summarize(ctx context.Context, sessionID string) error {

614632 a.Publish(pubsub.CreatedEvent, event)

615633 return

616634 }

617− // Create a new session with the summary

618− newSession, err := a.sessions.Create(summarizeCtx, oldSession.Title+" - Continuation")

635+ // Create a message in the new session with the summary

636+ msg, err := a.messages.Create(summarizeCtx, oldSession.ID, message.CreateMessageParams{

637+ Role: message.Assistant,

638+ Parts: []message.ContentPart{

639+ message.TextContent{Text: summary},

640+ message.Finish{

641+ Reason: message.FinishReasonEndTurn,

642+ Time: time.Now().Unix(),

643+ },

644+ },

645+ Model: a.summarizeProvider.Model().ID,

646+ })

619647 if err != nil {

620648 event = AgentEvent{

621649 Type: AgentEventTypeError,

622− Error: fmt.Errorf("failed to create new session: %w", err),

650+ Error: fmt.Errorf("failed to create summary message: %w", err),

623651 Done: true,

624652 }

653+

625654 a.Publish(pubsub.CreatedEvent, event)

626655 return

627656 }

628−

629− // Create a message in the new session with the summary

630− _, err = a.messages.Create(summarizeCtx, newSession.ID, message.CreateMessageParams{

631− Role: message.Assistant,

632− Parts: []message.ContentPart{message.TextContent{Text: summary}},

633− Model: a.summarizeProvider.Model().ID,

634− })

657+ oldSession.SummaryMessageID = msg.ID

658+ oldSession.CompletionTokens = response.Usage.OutputTokens

659+ oldSession.PromptTokens = 0

660+ model := a.summarizeProvider.Model()

661+ usage := response.Usage

662+ cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +

663+ model.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +

664+ model.CostPer1MIn/1e6*float64(usage.InputTokens) +

665+ model.CostPer1MOut/1e6*float64(usage.OutputTokens)

666+ oldSession.Cost += cost

667+ _, err = a.sessions.Save(summarizeCtx, oldSession)

635668 if err != nil {

636669 event = AgentEvent{

637670 Type: AgentEventTypeError,

638− Error: fmt.Errorf("failed to create summary message: %w", err),

671+ Error: fmt.Errorf("failed to save session: %w", err),

639672 Done: true,

640673 }

641−

642674 a.Publish(pubsub.CreatedEvent, event)

643− return

644675 }

676+

645677 event = AgentEvent{

646678 Type: AgentEventTypeSummarize,

647− SessionID: newSession.ID,

679+ SessionID: oldSession.ID,

648680 Progress: "Summary complete",

649681 Done: true,

650682 }

internal/session/session.go+7−1

@@ -16,6 +16,7 @@ type Session struct {

1616 MessageCount int64

1717 PromptTokens int64

1818 CompletionTokens int64

19+ SummaryMessageID string

1920 Cost float64

2021 CreatedAt int64

2122 UpdatedAt int64

@@ -105,7 +106,11 @@ func (s *service) Save(ctx context.Context, session Session) (Session, error) {

105106 Title: session.Title,

106107 PromptTokens: session.PromptTokens,

107108 CompletionTokens: session.CompletionTokens,

108− Cost: session.Cost,

109+ SummaryMessageID: sql.NullString{

110+ String: session.SummaryMessageID,

111+ Valid: session.SummaryMessageID != "",

112+ },

113+ Cost: session.Cost,

109114 })

110115 if err != nil {

111116 return Session{}, err

@@ -135,6 +140,7 @@ func (s service) fromDBItem(item db.Session) Session {

135140 MessageCount: item.MessageCount,

136141 PromptTokens: item.PromptTokens,

137142 CompletionTokens: item.CompletionTokens,

143+ SummaryMessageID: item.SummaryMessageID.String,

138144 Cost: item.Cost,

139145 CreatedAt: item.CreatedAt,

140146 UpdatedAt: item.UpdatedAt,

internal/tui/components/chat/list.go+11

@@ -99,6 +99,14 @@ func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

9999 case renderFinishedMsg:

100100 m.rendering = false

101101 m.viewport.GotoBottom()

102+ case pubsub.Event[session.Session]:

103+ if msg.Type == pubsub.UpdatedEvent && msg.Payload.ID == m.session.ID {

104+ m.session = msg.Payload

105+ if m.session.SummaryMessageID == m.currentMsgID {

106+ delete(m.cachedContent, m.currentMsgID)

107+ m.renderView()

108+ }

109+ }

102110 case pubsub.Event[message.Message]:

103111 needsRerender := false

104112 if msg.Type == pubsub.CreatedEvent {

@@ -208,12 +216,15 @@ func (m *messagesCmp) renderView() {

208216 m.uiMessages = append(m.uiMessages, cache.content...)

209217 continue

210218 }

219+ isSummary := m.session.SummaryMessageID == msg.ID

220+

211221 assistantMessages := renderAssistantMessage(

212222 msg,

213223 inx,

214224 m.messages,

215225 m.app.Messages,

216226 m.currentMsgID,

227+ isSummary,

217228 m.width,

218229 pos,

219230 )

internal/tui/components/chat/message.go+4

@@ -120,6 +120,7 @@ func renderAssistantMessage(

120120 allMessages []message.Message, // we need this to get tool results and the user message

121121 messagesService message.Service, // We need this to get the task tool messages

122122 focusedUIMessageId string,

123+ isSummary bool,

123124 width int,

124125 position int,

125126) []uiMessage {

@@ -168,6 +169,9 @@ func renderAssistantMessage(

168169 if content == "" {

169170 content = "*Finished without output*"

170171 }

172+ if isSummary {

173+ info = append(info, baseStyle.Width(width-1).Foreground(t.TextMuted()).Render(" (summary)"))

174+ }

171175

172176 content = renderMessage(content, false, true, width, info...)

173177 messages = append(messages, uiMessage{

internal/tui/tui.go−24

@@ -331,30 +331,6 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

331331

332332 if payload.Done && payload.Type == agent.AgentEventTypeSummarize {

333333 a.isCompacting = false

334−

335− if payload.SessionID != "" {

336− // Switch to the new session

337− return a, func() tea.Msg {

338− sessions, err := a.app.Sessions.List(context.Background())

339− if err != nil {

340− return util.InfoMsg{

341− Type: util.InfoTypeError,

342− Msg: "Failed to list sessions: " + err.Error(),

343− }

344− }

345−

346− for _, s := range sessions {

347− if s.ID == payload.SessionID {

348− return dialog.SessionSelectedMsg{Session: s}

349− }

350− }

351−

352− return util.InfoMsg{

353− Type: util.InfoTypeError,

354− Msg: "Failed to find new session",

355− }

356− }

357− }

358334 return a, util.ReportInfo("Session summarization complete")

359335 } else if payload.Done && payload.Type == agent.AgentEventTypeResponse && a.selectedSession.ID != "" {

360336 model := a.app.CoderAgent.Model()