| 1 | // Package chunk splits a document's text into retrievable passages. |
| 2 | // |
| 3 | // A chunk is the unit the index scores and the unit a citation points at, |
| 4 | // so its boundaries decide whether an answer can be checked. Two rules |
| 5 | // follow from that: |
| 6 | // |
| 7 | // - A chunk never straddles a heading. A passage whose first half |
| 8 | // belongs to one section and second half to another cites wrongly |
| 9 | // whichever way you label it. |
| 10 | // - Every chunk carries its heading trail. "1,350,000" is not |
| 11 | // retrievable; "FY2024 › Revenue › June: 1,350,000" is. The spec's |
| 12 | // §13 records the harder version of this problem for spreadsheets, |
| 13 | // which this package does not yet solve. |
| 14 | package chunk |
| 15 | |
| 16 | import ( |
| 17 | "strings" |
| 18 | "unicode" |
| 19 | ) |
| 20 | |
| 21 | // TargetWords is the size a chunk aims for. 400 words is roughly 500 |
| 22 | // tokens, which leaves eight chunks plus instructions inside a modest |
| 23 | // context window while staying long enough to carry an argument. |
| 24 | const TargetWords = 400 |
| 25 | |
| 26 | // MinWords stops a trailing fragment becoming its own chunk. A two-word |
| 27 | // chunk scores absurdly well on BM25 for those two words. |
| 28 | const MinWords = 40 |
| 29 | |
| 30 | // MaxWords is a hard ceiling. Without one, a document containing no blank |
| 31 | // line — a CSV, hard-wrapped prose, a log — becomes a single chunk of the |
| 32 | // whole file, and its one citation reads "passage 0" for everything. The |
| 33 | // form invites .csv uploads and a CSV never has a blank line, so this is |
| 34 | // the common case rather than the pathological one. |
| 35 | const MaxWords = 1000 |
| 36 | |
| 37 | // Chunk is one retrievable passage. |
| 38 | type Chunk struct { |
| 39 | // Ordinal is this chunk's position in the document, from 0. It is |
| 40 | // stored in plaintext: it is a coordinate, not content, exactly as |
| 41 | // Sheets keeps cell coordinates plaintext. |
| 42 | Ordinal int |
| 43 | // Trail is the heading path above this passage, outermost first. |
| 44 | Trail []string |
| 45 | // Text is the passage, with the trail prepended when Indexable is |
| 46 | // called. Stored sealed. |
| 47 | Text string |
| 48 | // Path is the document's ancestry in its SOURCE, outermost first — |
| 49 | // distinct from Trail, which is the heading path inside this one |
| 50 | // document. Empty for anything uploaded or pasted. |
| 51 | Path []string |
| 52 | } |
| 53 | |
| 54 | // Heading returns the trail joined for display and for indexing. |
| 55 | // |
| 56 | // Empty segments are dropped. The trail is padded when a document skips a |
| 57 | // level — one starting at "##" with no "#" above it — and an unfiltered |
| 58 | // join then renders "› Revenue" with a leading separator and nothing |
| 59 | // before it. Found by looking at the page; no test saw it, because every |
| 60 | // fixture started at "#". |
| 61 | func (c Chunk) Heading() string { return joinTrail(c.Trail) } |
| 62 | |
| 63 | // Indexable is what the index tokenises: the heading trail followed by |
| 64 | // the body. The trail is repeated into every chunk of a section |
| 65 | // deliberately — it is the only reason a query naming a section finds a |
| 66 | // passage that never repeats the section's name. |
| 67 | func (c Chunk) Indexable() string { |
| 68 | if len(c.Trail) == 0 { |
| 69 | return c.Text |
| 70 | } |
| 71 | return c.Heading() + "\n" + c.Text |
| 72 | } |
| 73 | |
| 74 | // Split divides plain text or Markdown into chunks. Markdown ATX |
| 75 | // headings ("## Revenue") set the trail; a blank line separates |
| 76 | // paragraphs; paragraphs accumulate until TargetWords is reached. |
| 77 | // |
| 78 | // It works line by line rather than paragraph by paragraph, which three |
| 79 | // bugs made necessary. A heading is a heading whether or not a blank line |
| 80 | // follows it — "## Revenue\nJune was strong." is the commonest Markdown |
| 81 | // style and the paragraph-based version filed it as body text with an |
| 82 | // empty trail. A "#" inside a fenced code block is a shell comment, not a |
| 83 | // heading, and was silently becoming one. And a document with no blank |
| 84 | // lines at all was becoming one chunk of the entire file. |
| 85 | func Split(text string) []Chunk { |
| 86 | var out []Chunk |
| 87 | var trail []string |
| 88 | var para []string |
| 89 | var buf []string |
| 90 | words := 0 |
| 91 | paraWords := 0 |
| 92 | inFence := false |
| 93 | fence := "" |
| 94 | |
| 95 | flushPara := func() { |
| 96 | if len(para) == 0 { |
| 97 | return |
| 98 | } |
| 99 | p := strings.Join(para, "\n") |
| 100 | para = para[:0] |
| 101 | n := len(strings.Fields(p)) |
| 102 | if words > 0 && words+n > TargetWords { |
| 103 | flushChunk(&out, &buf, &words, trail) |
| 104 | } |
| 105 | buf = append(buf, p) |
| 106 | words += n |
| 107 | // The hard ceiling: a single paragraph can exceed TargetWords on |
| 108 | // its own, so check after appending too. |
| 109 | if words >= MaxWords { |
| 110 | flushChunk(&out, &buf, &words, trail) |
| 111 | } |
| 112 | } |
| 113 | |
| 114 | for _, line := range strings.Split(normalise(text), "\n") { |
| 115 | trimmed := strings.TrimSpace(line) |
| 116 | |
| 117 | // Fences first: nothing inside one is structure. |
| 118 | if f := fenceToken(trimmed); f != "" { |
| 119 | if !inFence { |
| 120 | inFence, fence = true, f |
| 121 | } else if strings.HasPrefix(f, fence[:1]) && len(f) >= len(fence) { |
| 122 | inFence, fence = false, "" |
| 123 | } |
| 124 | para = append(para, line) |
| 125 | continue |
| 126 | } |
| 127 | if inFence { |
| 128 | para = append(para, line) |
| 129 | continue |
| 130 | } |
| 131 | |
| 132 | if trimmed == "" { |
| 133 | flushPara() |
| 134 | paraWords = 0 |
| 135 | continue |
| 136 | } |
| 137 | if depth, title, ok := heading(trimmed); ok { |
| 138 | // A heading closes what came before it, blank line or not. |
| 139 | flushPara() |
| 140 | paraWords = 0 |
| 141 | flushChunk(&out, &buf, &words, trail) |
| 142 | if depth-1 < len(trail) { |
| 143 | trail = trail[:depth-1] |
| 144 | } |
| 145 | for len(trail) < depth-1 { |
| 146 | trail = append(trail, "") |
| 147 | } |
| 148 | trail = append(trail, title) |
| 149 | continue |
| 150 | } |
| 151 | para = append(para, line) |
| 152 | // A paragraph is bounded too, not just a chunk. Without this a |
| 153 | // file with no blank line at all — a CSV, a log — accumulates |
| 154 | // every line into one paragraph and the ceiling below never gets |
| 155 | // a chance to bind, which is how a 4000-line CSV became a single |
| 156 | // passage. |
| 157 | if paraWords += len(strings.Fields(line)); paraWords >= MaxWords { |
| 158 | flushPara() |
| 159 | paraWords = 0 |
| 160 | } |
| 161 | } |
| 162 | flushPara() |
| 163 | flushChunk(&out, &buf, &words, trail) |
| 164 | |
| 165 | for i := range out { |
| 166 | out[i].Ordinal = i |
| 167 | } |
| 168 | return out |
| 169 | } |
| 170 | |
| 171 | // flushChunk emits the accumulated paragraphs as one chunk. |
| 172 | func flushChunk(out *[]Chunk, buf *[]string, words *int, trail []string) { |
| 173 | if len(*buf) == 0 { |
| 174 | return |
| 175 | } |
| 176 | body := strings.Join(*buf, "\n\n") |
| 177 | t := make([]string, len(trail)) |
| 178 | copy(t, trail) |
| 179 | cur := Chunk{Trail: t} |
| 180 | if *words < MinWords && len(*out) > 0 { |
| 181 | prev := &(*out)[len(*out)-1] |
| 182 | // Fold a runt into its predecessor, but only under the same |
| 183 | // heading — folding across one would file text under the wrong |
| 184 | // section. Compared through Heading() on both sides: an earlier |
| 185 | // version compared the filtered heading against an unfiltered |
| 186 | // join, so the fold never fired when the trail was padded. |
| 187 | if prev.Heading() == cur.Heading() { |
| 188 | prev.Text += "\n\n" + body |
| 189 | *buf, *words = nil, 0 |
| 190 | return |
| 191 | } |
| 192 | } |
| 193 | cur.Text = body |
| 194 | cur.Ordinal = len(*out) |
| 195 | *out = append(*out, cur) |
| 196 | *buf, *words = nil, 0 |
| 197 | } |
| 198 | |
| 199 | // fenceToken returns the fence marker if the line opens or closes a |
| 200 | // fenced code block. |
| 201 | func fenceToken(trimmed string) string { |
| 202 | for _, c := range []string{"```", "~~~"} { |
| 203 | if strings.HasPrefix(trimmed, c) { |
| 204 | n := 0 |
| 205 | for n < len(trimmed) && trimmed[n] == c[0] { |
| 206 | n++ |
| 207 | } |
| 208 | return trimmed[:n] |
| 209 | } |
| 210 | } |
| 211 | return "" |
| 212 | } |
| 213 | |
| 214 | // heading recognises a Markdown ATX heading and returns its depth. |
| 215 | func heading(line string) (depth int, title string, ok bool) { |
| 216 | if !strings.HasPrefix(line, "#") { |
| 217 | return 0, "", false |
| 218 | } |
| 219 | i := 0 |
| 220 | for i < len(line) && line[i] == '#' { |
| 221 | i++ |
| 222 | } |
| 223 | if i > 6 || i >= len(line) || !unicode.IsSpace(rune(line[i])) { |
| 224 | return 0, "", false |
| 225 | } |
| 226 | title = strings.TrimSpace(strings.TrimRight(line[i:], " #")) |
| 227 | if title == "" { |
| 228 | return 0, "", false |
| 229 | } |
| 230 | return i, title, true |
| 231 | } |
| 232 | |
| 233 | func normalise(s string) string { |
| 234 | s = strings.ReplaceAll(s, "\r\n", "\n") |
| 235 | s = strings.ReplaceAll(s, "\r", "\n") |
| 236 | // Collapse three or more newlines to exactly two so paragraph |
| 237 | // splitting is not defeated by generous spacing. |
| 238 | for strings.Contains(s, "\n\n\n") { |
| 239 | s = strings.ReplaceAll(s, "\n\n\n", "\n\n") |
| 240 | } |
| 241 | return s |
| 242 | } |
| 243 | |
| 244 | // Path is a document's ancestry in its source, outermost first, carried |
| 245 | // alongside the trail rather than merged into it. |
| 246 | // |
| 247 | // Kept separate because the two are scored differently and must be: a |
| 248 | // heading is written inside the document, while ancestry is inherited by |
| 249 | // every chunk of every document beneath it. index.PathWeight explains |
| 250 | // what merging them does to ranking. |
| 251 | func (c Chunk) PathString() string { return joinTrail(c.Path) } |
| 252 | |
| 253 | // SplitUnder is Split for a document that has ancestry. Every chunk |
| 254 | // carries the same path, because every passage of a nested page is |
| 255 | // equally nested. |
| 256 | func SplitUnder(path []string, text string) []Chunk { |
| 257 | out := Split(text) |
| 258 | if len(path) == 0 { |
| 259 | return out |
| 260 | } |
| 261 | p := make([]string, len(path)) |
| 262 | copy(p, path) |
| 263 | for i := range out { |
| 264 | out[i].Path = p |
| 265 | } |
| 266 | return out |
| 267 | } |
| 268 | |
| 269 | // joinTrail renders a trail for display and indexing, dropping empty |
| 270 | // segments. Shared by Heading and PathString so the two cannot drift — |
| 271 | // an earlier bug rendered "› Revenue" because one side filtered and the |
| 272 | // other did not. |
| 273 | func joinTrail(parts []string) string { |
| 274 | kept := make([]string, 0, len(parts)) |
| 275 | for _, t := range parts { |
| 276 | if strings.TrimSpace(t) != "" { |
| 277 | kept = append(kept, t) |
| 278 | } |
| 279 | } |
| 280 | return strings.Join(kept, " › ") |
| 281 | } |
| 282 | |
amadan