rework llm

1 files changed, 10 insertions, 0 deletions

diff --git a/internal/db/db.go b/internal/db/db.go
index e882106c4..75f626013 100644
--- a/internal/db/db.go
+++ b/internal/db/db.go
@@ -51,6 +51,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) {
 	if q.listSessionsStmt, err = db.PrepareContext(ctx, listSessions); err != nil {
 		return nil, fmt.Errorf("error preparing query ListSessions: %w", err)
 	}
+	if q.updateMessageStmt, err = db.PrepareContext(ctx, updateMessage); err != nil {
+		return nil, fmt.Errorf("error preparing query UpdateMessage: %w", err)
+	}
 	if q.updateSessionStmt, err = db.PrepareContext(ctx, updateSession); err != nil {
 		return nil, fmt.Errorf("error preparing query UpdateSession: %w", err)
 	}
@@ -104,6 +107,11 @@ func (q *Queries) Close() error {
 			err = fmt.Errorf("error closing listSessionsStmt: %w", cerr)
 		}
 	}
+	if q.updateMessageStmt != nil {
+		if cerr := q.updateMessageStmt.Close(); cerr != nil {
+			err = fmt.Errorf("error closing updateMessageStmt: %w", cerr)
+		}
+	}
 	if q.updateSessionStmt != nil {
 		if cerr := q.updateSessionStmt.Close(); cerr != nil {
 			err = fmt.Errorf("error closing updateSessionStmt: %w", cerr)
@@ -157,6 +165,7 @@ type Queries struct {
 	getSessionByIDStmt        *sql.Stmt
 	listMessagesBySessionStmt *sql.Stmt
 	listSessionsStmt          *sql.Stmt
+	updateMessageStmt         *sql.Stmt
 	updateSessionStmt         *sql.Stmt
 }
 
@@ -173,6 +182,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries {
 		getSessionByIDStmt:        q.getSessionByIDStmt,
 		listMessagesBySessionStmt: q.listMessagesBySessionStmt,
 		listSessionsStmt:          q.listSessionsStmt,
+		updateMessageStmt:         q.updateMessageStmt,
 		updateSessionStmt:         q.updateSessionStmt,
 	}
 }