mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: Add full text search over chat messages (#27126)
Closes CODAGT-721 Closes CODAGT-722 Closes CODAGT-723 Closes CODAGT-724 Closes CODAGT-725 This PR adds the database and API pieces necessary to support full-text chat message search. - Adds required chat schema for full-text search - Adds dbpurge job to populate search_tsv in the background - Adds `search` parameter to GetChats query - Adds `search` filter to `searchquery.Chats` - Wires chat search filter into chats API > Implemented by Coder Agents, reviewed and tested by a human.
This commit is contained in:
Cian Johnston
Generated
+103
-10
@@ -6032,6 +6032,36 @@ func (q *sqlQuerier) AutoArchiveInactiveChats(ctx context.Context, arg AutoArchi
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const backfillChatMessagesSearchTsv = `-- name: BackfillChatMessagesSearchTsv :execrows
|
||||
WITH batch AS (
|
||||
SELECT id FROM chat_messages
|
||||
WHERE search_tsv IS NULL
|
||||
AND deleted = false
|
||||
AND visibility IN ('user', 'both')
|
||||
AND role IN ('user', 'assistant')
|
||||
ORDER BY id DESC
|
||||
LIMIT $1::int
|
||||
)
|
||||
UPDATE chat_messages cm
|
||||
SET search_tsv = COALESCE(
|
||||
to_tsvector('simple', chat_message_search_text(cm.content)),
|
||||
''::tsvector)
|
||||
FROM batch WHERE cm.id = batch.id
|
||||
`
|
||||
|
||||
// Backfills chat_messages.search_tsv for pending rows, newest first.
|
||||
// The WHERE clause must match the predicate of
|
||||
// idx_chat_messages_search_tsv_pending exactly so the partial index
|
||||
// serves this query.
|
||||
// NULL means "pending", empty tsvector means "backfilled, no text".
|
||||
func (q *sqlQuerier) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
|
||||
result, err := q.db.ExecContext(ctx, backfillChatMessagesSearchTsv, batchSize)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
const backoffChatDiffStatus = `-- name: BackoffChatDiffStatus :exec
|
||||
UPDATE
|
||||
chat_diff_statuses
|
||||
@@ -6096,6 +6126,19 @@ func (q *sqlQuerier) BatchUpsertChatHeartbeats(ctx context.Context, arg BatchUps
|
||||
return err
|
||||
}
|
||||
|
||||
const chatSearchQueryIsEmpty = `-- name: ChatSearchQueryIsEmpty :one
|
||||
SELECT numnode(websearch_to_tsquery('simple', $1::text)) = 0 AS is_empty
|
||||
`
|
||||
|
||||
// Reports whether search text tokenizes to an empty tsquery (e.g. '!!!').
|
||||
// Used to reject input that would silently match nothing.
|
||||
func (q *sqlQuerier) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
|
||||
row := q.db.QueryRowContext(ctx, chatSearchQueryIsEmpty, search)
|
||||
var is_empty bool
|
||||
err := row.Scan(&is_empty)
|
||||
return is_empty, err
|
||||
}
|
||||
|
||||
const countChatQueuedMessages = `-- name: CountChatQueuedMessages :one
|
||||
SELECT COUNT(*)::bigint AS count
|
||||
FROM chat_queued_messages
|
||||
@@ -7400,7 +7443,7 @@ func (q *sqlQuerier) GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatP
|
||||
|
||||
const getChatMessageByID = `-- name: GetChatMessageByID :one
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7436,6 +7479,7 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -7525,7 +7569,7 @@ func (q *sqlQuerier) GetChatMessageSummariesPerChat(ctx context.Context, created
|
||||
|
||||
const getChatMessagesByChatID = `-- name: GetChatMessagesByChatID :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7576,6 +7620,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7592,7 +7637,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
|
||||
|
||||
const getChatMessagesByChatIDAscPaginated = `-- name: GetChatMessagesByChatIDAscPaginated :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7646,6 +7691,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7662,7 +7708,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar
|
||||
|
||||
const getChatMessagesByChatIDDescPaginated = `-- name: GetChatMessagesByChatIDDescPaginated :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7729,6 +7775,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7745,7 +7792,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a
|
||||
|
||||
const getChatMessagesByRevisionForStream = `-- name: GetChatMessagesByRevisionForStream :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7795,6 +7842,7 @@ func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7827,7 +7875,7 @@ WITH latest_compressed_summary AS (
|
||||
1
|
||||
)
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7902,6 +7950,7 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -8588,6 +8637,46 @@ WHERE
|
||||
)
|
||||
ELSE true
|
||||
END
|
||||
-- websearch_to_tsquery accepts quoted phrases, OR, and -negation;
|
||||
-- the 'simple' config folds case and skips stemming.
|
||||
AND CASE
|
||||
WHEN $16::text != '' THEN (
|
||||
-- Served by idx_chats_title_fts.
|
||||
to_tsvector('simple', chats_expanded.title) @@ websearch_to_tsquery('simple', $16)
|
||||
-- Served by idx_chat_diff_statuses_pr_title_fts.
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_diff_statuses cds
|
||||
WHERE cds.chat_id = chats_expanded.id
|
||||
AND to_tsvector('simple', cds.pull_request_title) @@ websearch_to_tsquery('simple', $16)
|
||||
)
|
||||
-- The WHERE clause must repeat the predicate of the partial index
|
||||
-- idx_chat_messages_search_tsv so the planner can use it. Additional
|
||||
-- filters should still be fine.
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_messages cm
|
||||
WHERE cm.chat_id = chats_expanded.id
|
||||
AND cm.search_tsv IS NOT NULL
|
||||
AND cm.deleted = false
|
||||
AND cm.visibility IN ('user', 'both')
|
||||
AND cm.role IN ('user', 'assistant')
|
||||
AND cm.search_tsv @@ websearch_to_tsquery('simple', $16)
|
||||
)
|
||||
-- Skip an explicit pr_number lookup unless the search is a valid bigint.
|
||||
OR CASE
|
||||
WHEN $16 ~ '^[0-9]{1,18}$' THEN EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_diff_statuses cds
|
||||
WHERE cds.chat_id = chats_expanded.id
|
||||
AND cds.pr_number IS NOT NULL
|
||||
AND cds.pr_number = $16::bigint
|
||||
)
|
||||
ELSE false
|
||||
END
|
||||
)
|
||||
ELSE true
|
||||
END
|
||||
-- Paginate over root chats only. Children are fetched
|
||||
-- separately via GetChildChatsByParentIDs and embedded under
|
||||
-- each parent. Other callers that need the full set should
|
||||
@@ -8604,11 +8693,11 @@ ORDER BY
|
||||
-chats_expanded.pin_order DESC,
|
||||
chats_expanded.updated_at DESC,
|
||||
chats_expanded.id DESC
|
||||
OFFSET $16
|
||||
OFFSET $17
|
||||
LIMIT
|
||||
-- The chat list is unbounded and expected to grow large.
|
||||
-- Default to 50 to prevent accidental excessively large queries.
|
||||
COALESCE(NULLIF($17 :: int, 0), 50)
|
||||
COALESCE(NULLIF($18 :: int, 0), 50)
|
||||
`
|
||||
|
||||
type GetChatsParams struct {
|
||||
@@ -8627,6 +8716,7 @@ type GetChatsParams struct {
|
||||
PrNumber int32 `db:"pr_number" json:"pr_number"`
|
||||
RepoQuery string `db:"repo_query" json:"repo_query"`
|
||||
PrTitleQuery string `db:"pr_title_query" json:"pr_title_query"`
|
||||
Search string `db:"search" json:"search"`
|
||||
OffsetOpt int32 `db:"offset_opt" json:"offset_opt"`
|
||||
LimitOpt int32 `db:"limit_opt" json:"limit_opt"`
|
||||
}
|
||||
@@ -8653,6 +8743,7 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]GetCha
|
||||
arg.PrNumber,
|
||||
arg.RepoQuery,
|
||||
arg.PrTitleQuery,
|
||||
arg.Search,
|
||||
arg.OffsetOpt,
|
||||
arg.LimitOpt,
|
||||
)
|
||||
@@ -9147,7 +9238,7 @@ func (q *sqlQuerier) GetDatabaseNow(ctx context.Context) (time.Time, error) {
|
||||
|
||||
const getLastChatMessageByRole = `-- name: GetLastChatMessageByRole :one
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -9193,6 +9284,7 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -9702,7 +9794,7 @@ SELECT
|
||||
NULLIF(UNNEST($18::bigint[]), 0),
|
||||
NULLIF(UNNEST($19::bigint[]), 0)
|
||||
RETURNING
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
`
|
||||
|
||||
type InsertChatMessagesParams struct {
|
||||
@@ -9781,6 +9873,7 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user