|
| 1 | +# SPDX-FileCopyrightText: Sudo Apt Holdings LLC |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +defmodule Trinity.Memory.Search do |
| 4 | + @moduledoc """ |
| 5 | + Full-text search over every message (slice 031). SQLite: the `messages_fts` FTS5 table, |
| 6 | + porter stemming, `bm25()` order, `snippet()`; Postgres: the generated `content_tsv` column, |
| 7 | + `plainto_tsquery('english')`, `ts_rank`, `ts_headline`. Same shape of hit either way, so the |
| 8 | + tool and the page do not know which database they run on. |
| 9 | +
|
| 10 | + The query text is bound as a parameter and never spliced into SQL. On SQLite every term is |
| 11 | + additionally quoted for FTS5, so a user's `"`, `*`, `-` or `OR` is text to find, not an |
| 12 | + operator; the search is "all these words", stemmed, in any order. Stemming is a suffix |
| 13 | + operation: "running" and "runs" meet at "run", and "ran" does not (NOTES.md, fact 2). |
| 14 | +
|
| 15 | + Schemaless queries on purpose: Memory depends on the core and LLM, never on Sessions |
| 16 | + (docs/01), so the tables are named here and the ids are cast through `Trinity.UUID`. |
| 17 | + """ |
| 18 | + |
| 19 | + import Ecto.Query |
| 20 | + |
| 21 | + alias Trinity.Repo |
| 22 | + |
| 23 | + @default_limit 20 |
| 24 | + @max_limit 100 |
| 25 | + |
| 26 | + # The adapter is fixed at compile time (config/config.exs, slice 010), so the two databases' |
| 27 | + # query shapes are compiled in, not chosen at run time; the postgres job proves the other. |
| 28 | + @adapter Application.compile_env(:trinity, :db_adapter, Ecto.Adapters.SQLite3) |
| 29 | + |
| 30 | + @type hit :: %{ |
| 31 | + message_id: String.t(), |
| 32 | + session_id: String.t(), |
| 33 | + session_title: String.t() | nil, |
| 34 | + seq: non_neg_integer(), |
| 35 | + role: String.t(), |
| 36 | + snippet: String.t(), |
| 37 | + inserted_at: DateTime.t(), |
| 38 | + rank: float() |
| 39 | + } |
| 40 | + |
| 41 | + @doc """ |
| 42 | + Ranked hits for `query`. Options: `limit:` (#{@default_limit}, at most #{@max_limit}), `role:` |
| 43 | + (a message role), `persona_id:`, `since:` and `until:` (`DateTime`, on the message's |
| 44 | + `inserted_at`). An empty or all-punctuation query is no hits. |
| 45 | + """ |
| 46 | + @spec messages(String.t(), keyword()) :: [hit()] |
| 47 | + def messages(query, opts \\ []) when is_binary(query) do |
| 48 | + case terms(query) do |
| 49 | + [] -> [] |
| 50 | + terms -> terms |> build(opts) |> Repo.all() |> Enum.map(&to_hit/1) |
| 51 | + end |
| 52 | + end |
| 53 | + |
| 54 | + @doc "The words of a query, punctuation dropped; the population FTS5 and tsquery both receive." |
| 55 | + @spec terms(String.t()) :: [String.t()] |
| 56 | + def terms(query) do |
| 57 | + query |
| 58 | + |> String.split(~r/[^\p{L}\p{N}_']+/u, trim: true) |
| 59 | + |> Enum.map(&String.replace(&1, "'", "")) |
| 60 | + |> Enum.reject(&(&1 == "")) |
| 61 | + |> Enum.take(32) |
| 62 | + end |
| 63 | + |
| 64 | + @doc "Rebuilds the index from `messages`; on Postgres the column is generated and this reports so." |
| 65 | + @spec reindex() :: {:ok, :rebuilt | :generated_column} |
| 66 | + if @adapter == Ecto.Adapters.SQLite3 do |
| 67 | + def reindex do |
| 68 | + Repo.transaction(fn -> |
| 69 | + Repo.query!("DELETE FROM messages_fts") |
| 70 | + |
| 71 | + Repo.query!( |
| 72 | + "INSERT INTO messages_fts(rowid, content, session_id, message_id) SELECT rowid, content, session_id, id FROM messages" |
| 73 | + ) |
| 74 | + |
| 75 | + Repo.query!("INSERT INTO messages_fts(messages_fts) VALUES('optimize')") |
| 76 | + end) |
| 77 | + |
| 78 | + {:ok, :rebuilt} |
| 79 | + end |
| 80 | + else |
| 81 | + def reindex, do: {:ok, :generated_column} |
| 82 | + end |
| 83 | + |
| 84 | + defp build(terms, opts) do |
| 85 | + limit = opts |> Keyword.get(:limit, @default_limit) |> min(@max_limit) |> max(1) |
| 86 | + |
| 87 | + base() |
| 88 | + |> match(terms) |
| 89 | + |> filter(:role, opts[:role]) |
| 90 | + |> filter(:persona_id, opts[:persona_id]) |
| 91 | + |> filter(:since, opts[:since]) |
| 92 | + |> filter(:until, opts[:until]) |
| 93 | + |> limit(^limit) |
| 94 | + end |
| 95 | + |
| 96 | + defp base do |
| 97 | + from(m in "messages", |
| 98 | + join: s in "sessions", |
| 99 | + on: s.id == m.session_id, |
| 100 | + select: %{ |
| 101 | + message_id: type(m.id, Trinity.UUID), |
| 102 | + session_id: type(m.session_id, Trinity.UUID), |
| 103 | + session_title: s.title, |
| 104 | + seq: m.seq, |
| 105 | + role: m.role, |
| 106 | + inserted_at: type(m.inserted_at, :utc_datetime_usec), |
| 107 | + persona_id: type(s.persona_id, Trinity.UUID) |
| 108 | + } |
| 109 | + ) |
| 110 | + end |
| 111 | + |
| 112 | + # SQLite: every term quoted for FTS5 (a double quote inside is doubled), joined by spaces, |
| 113 | + # which FTS5 reads as AND; the whole string is one bound parameter. |
| 114 | + if @adapter == Ecto.Adapters.SQLite3 do |
| 115 | + defp match(query, terms) do |
| 116 | + needle = Enum.map_join(terms, " ", &("\"" <> String.replace(&1, "\"", "\"\"") <> "\"")) |
| 117 | + |
| 118 | + from([m, s] in query, |
| 119 | + join: f in "messages_fts", |
| 120 | + on: f.rowid == m.rowid, |
| 121 | + where: fragment("messages_fts MATCH ?", ^needle), |
| 122 | + order_by: fragment("bm25(messages_fts)"), |
| 123 | + select_merge: %{ |
| 124 | + snippet: fragment("snippet(messages_fts, 0, '[', ']', '…', 12)"), |
| 125 | + rank: fragment("bm25(messages_fts)") |
| 126 | + } |
| 127 | + ) |
| 128 | + end |
| 129 | + else |
| 130 | + defp match(query, terms) do |
| 131 | + needle = Enum.join(terms, " ") |
| 132 | + |
| 133 | + from([m, s] in query, |
| 134 | + where: fragment("? @@ plainto_tsquery('english', ?)", m.content_tsv, ^needle), |
| 135 | + order_by: [ |
| 136 | + desc: fragment("ts_rank(?, plainto_tsquery('english', ?))", m.content_tsv, ^needle) |
| 137 | + ], |
| 138 | + select_merge: %{ |
| 139 | + snippet: |
| 140 | + fragment( |
| 141 | + "ts_headline('english', ?, plainto_tsquery('english', ?), 'StartSel=[, StopSel=], MaxWords=12, MinWords=6')", |
| 142 | + m.content, |
| 143 | + ^needle |
| 144 | + ), |
| 145 | + rank: fragment("ts_rank(?, plainto_tsquery('english', ?))", m.content_tsv, ^needle) |
| 146 | + } |
| 147 | + ) |
| 148 | + end |
| 149 | + end |
| 150 | + |
| 151 | + defp filter(query, _key, nil), do: query |
| 152 | + defp filter(query, :role, role), do: from([m, s] in query, where: m.role == ^role) |
| 153 | + |
| 154 | + defp filter(query, :persona_id, id), |
| 155 | + do: from([m, s] in query, where: s.persona_id == type(^id, Trinity.UUID)) |
| 156 | + |
| 157 | + defp filter(query, :since, at), do: from([m, s] in query, where: m.inserted_at >= ^at) |
| 158 | + defp filter(query, :until, at), do: from([m, s] in query, where: m.inserted_at <= ^at) |
| 159 | + |
| 160 | + defp to_hit(row) do |
| 161 | + row |
| 162 | + |> Map.delete(:persona_id) |
| 163 | + |> Map.update!(:rank, &(&1 * 1.0)) |
| 164 | + end |
| 165 | +end |
0 commit comments