orchard

  1package tools
  2
  3// A local Wikipedia, served by kiwix off a ZIM on the bridge.
  4//
  5// The snapshot is the mini flavour, which carries every article's lead section
  6// and nothing below the first heading. That is the half that answers what and
  7// who and when, and it is the half that fits: a full article runs 18k to 50k
  8// tokens and one of them would spend most of the window on a single tool
  9// result. Depth stays a web_fetch away on the url this returns.
 10//
 11// It reaches a container on the bridge, so it costs no outbound request and
 12// keeps working while a search host has this address blocked.
 13
 14import (
 15	"context"
 16	"encoding/xml"
 17	"fmt"
 18	"html"
 19	"io"
 20	"net/http"
 21	"net/url"
 22	"regexp"
 23	"strings"
 24	"sync"
 25	"time"
 26)
 27
 28// WikiBase is the kiwix container. Set from main so a dev run can point at a
 29// server outside the compose network.
 30var WikiBase = "http://orchard-wiki:8000"
 31
 32// kiwix names a book after its file, so the deploy calls the ZIM wikipedia.zim
 33// and this matches. Getting it wrong answers 400 rather than an empty result.
 34const wikiBook = "wikipedia"
 35
 36// What one lead section is allowed to cost. Measured leads run well under this
 37// and the cap is for the occasional long one rather than the common case.
 38const wikiMaxChars = 4000
 39
 40var Wikipedia = Tool{
 41	Name: "wikipedia",
 42	Description: "Look a subject up in a local offline Wikipedia snapshot. It answers in " +
 43		"milliseconds, costs no web request, and returns the opening section of that article. " +
 44		"Use it first for background on any person, place, organisation, product, species, event " +
 45		"or technical term, before reaching for web_search. Pass the name of the thing and not " +
 46		"the question: for 'who is the leader of North Korea' pass 'North Korea', and for 'what " +
 47		"is a b-tree used for' pass 'b-tree'. It holds background rather than news: it knows " +
 48		"nothing after its snapshot date and carries each article's opening section only, so use " +
 49		"web_search or web_fetch for detail or for anything current.",
 50	Schema: obj(map[string]any{
 51		"query": str("the name of the thing to look up, like 'PostgreSQL' or 'North Korea', not a question"),
 52	}, "query"),
 53	Run: func(ctx context.Context, d *Deps, a map[string]any) (any, error) {
 54		q := strings.TrimSpace(argStr(a, "query"))
 55		if q == "" {
 56			return nil, fmt.Errorf("query is required")
 57		}
 58		if to := liveInstead(q); to != "" {
 59			return nil, fmt.Errorf("the snapshot cannot answer that, it holds background and not "+
 60				"what is happening. Call %s instead", to)
 61		}
 62		return wikiLookup(ctx, d, WikiBase, q)
 63	},
 64}
 65
 66// A subject whose head noun is something that changes by the hour, and the tool
 67// that does answer it. Asked for "Big news today" the snapshot returned the
 68// founding date of Universe Today, which is a real article and a useless answer,
 69// and the contract telling the model not to do this did not stop it.
 70//
 71// The head noun rather than any word in the phrase, so "News Corporation" and
 72// "Marketplace" still look up as the things they are.
 73func liveInstead(q string) string {
 74	f := strings.Fields(strings.ToLower(strings.Trim(q, " \t?.!,:;")))
 75	if len(f) == 0 {
 76		return ""
 77	}
 78	return liveHeadNoun[f[len(f)-1]]
 79}
 80
 81var liveHeadNoun = map[string]string{
 82	"news": News.Name, "headline": News.Name, "headlines": News.Name,
 83	"stories": News.Name, "story": News.Name, "events": News.Name,
 84	"weather": Weather.Name, "forecast": Weather.Name, "temperature": Weather.Name,
 85	"price": Markets.Name, "prices": Markets.Name, "stock": Markets.Name,
 86	"stocks": Markets.Name, "market": Markets.Name, "markets": Markets.Name,
 87	"score": SportsScores.Name, "scores": SportsScores.Name,
 88	"time": Now.Name, "date": Now.Name,
 89}
 90
 91func wikiLookup(ctx context.Context, d *Deps, base, q string) (any, error) {
 92	// An exact title beats the ranker. A one word lookup is nearly always the
 93	// article's own name, and full text search on it can rank a page that
 94	// merely mentions the word above the page about it.
 95	if guess, link, ok := wikiExact(ctx, d, base, q); ok {
 96		title, text, err := wikiArticle(ctx, d, base, link)
 97		if err == nil {
 98			// The page knows its own capitalisation and the guessed path does
 99			// not, so a lookup for "postgresql" comes back as PostgreSQL.
100			if title == "" {
101				title = guess
102			}
103			return wikiResult(ctx, d, base, title, text, nil), nil
104		}
105	}
106
107	hits, err := wikiSearch(ctx, d, base, q)
108	if err != nil {
109		return nil, err
110	}
111	if len(hits) == 0 {
112		return wikiMiss(nil), nil
113	}
114
115	// Titles of the runners up, so a wrong first guess can be corrected with a
116	// second call rather than a search.
117	var also []string
118	for _, h := range hits[1:] {
119		also = append(also, h.Title)
120	}
121
122	top := hits[0]
123	// The argument names a thing, so its article is named after it. Ranking
124	// lead sections across 19 million of them puts something unrelated on top
125	// often enough that handing the best hit over regardless is how the model
126	// gets told the leader of North Korea is a military history article.
127	if !wikiMatches(q, top.Title) {
128		return wikiMiss(append([]string{top.Title}, also...)), nil
129	}
130
131	_, text, err := wikiArticle(ctx, d, base, top.Link)
132	if err != nil {
133		return nil, err
134	}
135
136	return wikiResult(ctx, d, base, top.Title, text, also), nil
137}
138
139func wikiResult(ctx context.Context, d *Deps, base, title, text string, also []string) map[string]any {
140	date := wikiDate(ctx, d, base)
141	out := map[string]any{
142		"found":   true,
143		"title":   title,
144		"summary": text,
145		"url":     "https://en.wikipedia.org/wiki/" + strings.ReplaceAll(title, " ", "_"),
146		// Its own field as well as the sentence below, since anything reading
147		// this programmatically has to weigh the age without parsing prose.
148		"snapshot_date": date,
149		"note": "This is the opening section only, from an offline snapshot taken " + date +
150			". Anything after that date is not in it. For the rest of the article, or for anything " +
151			"current, call web_fetch on the url above or use web_search.",
152	}
153	if len(also) > 0 {
154		out["other_matches"] = also
155	}
156	return out
157}
158
159func wikiMiss(near []string) map[string]any {
160	out := map[string]any{
161		"found": false,
162		"note": "No article in the offline snapshot matches that. If you passed a question, call this " +
163			"again with just the name of the thing it is about. Otherwise use web_search, and do not " +
164			"treat the titles below as the answer, since they are what was rejected.",
165	}
166	if len(near) > 0 {
167		out["near_titles"] = near
168	}
169	return out
170}
171
172// wikiExact tries the query as an article title. Titles capitalise their first
173// letter, so a lowercase query needs the second try to hit.
174func wikiExact(ctx context.Context, d *Deps, base, q string) (title, link string, ok bool) {
175	t := strings.ReplaceAll(strings.TrimSpace(q), " ", "_")
176	if t == "" {
177		return "", "", false
178	}
179	for _, cand := range []string{t, strings.ToUpper(t[:1]) + t[1:]} {
180		link := "/content/" + wikiBook + "/" + url.PathEscape(cand)
181		if _, err := wikiGet(ctx, d, base+link); err == nil {
182			return strings.ReplaceAll(cand, "_", " "), link, true
183		}
184	}
185	return "", "", false
186}
187
188var wikiWord = regexp.MustCompile(`[a-z0-9]+`)
189
190// Words that carry no subject, so they neither count towards a match nor
191// against one. "The Beatles" has to match on Beatles alone.
192var wikiStop = map[string]bool{
193	"the": true, "a": true, "an": true, "of": true, "in": true, "on": true, "and": true,
194	"for": true, "to": true, "is": true, "are": true, "was": true, "were": true, "what": true,
195	"who": true, "when": true, "where": true, "why": true, "how": true, "does": true, "do": true,
196	"about": true, "me": true, "tell": true, "it": true, "its": true, "that": true, "this": true,
197}
198
199func wikiWords(s string) []string {
200	var out []string
201	for _, w := range wikiWord.FindAllString(strings.ToLower(s), -1) {
202		if !wikiStop[w] {
203			out = append(out, w)
204		}
205	}
206	return out
207}
208
209// wikiMatches asks whether the article is about the thing that was asked for,
210// rather than whether it mentions it. Half the title's own words have to appear
211// in the query, so "North Korea" matches and "Military history of Korea" does
212// not, which is the difference between an answer and a confident wrong one.
213func wikiMatches(q, title string) bool {
214	tw := wikiWords(title)
215	if len(tw) == 0 {
216		return false
217	}
218	asked := map[string]bool{}
219	for _, w := range wikiWords(q) {
220		asked[w] = true
221	}
222	hit := 0
223	for _, w := range tw {
224		if asked[w] {
225			hit++
226		}
227	}
228	return hit*2 >= len(tw)
229}
230
231type wikiHit struct {
232	Title string `xml:"title"`
233	Link  string `xml:"link"`
234}
235
236func wikiSearch(ctx context.Context, d *Deps, base, q string) ([]wikiHit, error) {
237	u := fmt.Sprintf("%s/search?books.name=%s&pattern=%s&format=xml&pageLength=5",
238		base, wikiBook, url.QueryEscape(q))
239	body, err := wikiGet(ctx, d, u)
240	if err != nil {
241		return nil, err
242	}
243	var feed struct {
244		Items []wikiHit `xml:"channel>item"`
245	}
246	if err := xml.Unmarshal(body, &feed); err != nil {
247		return nil, fmt.Errorf("the wikipedia snapshot sent a result that could not be read")
248	}
249	return feed.Items, nil
250}
251
252var (
253	wikiHead    = regexp.MustCompile(`(?is)<head[^>]*>.*?</head\s*>`)
254	wikiH1      = regexp.MustCompile(`(?is)<h1[^>]*>.*?</h1\s*>`)
255	wikiCutHead = regexp.MustCompile(`(?is)<h[23][^>]*>`)
256	// kiwix appends the Creative Commons notice to every article, so without
257	// this every single summary ends on the same two sentences of licence.
258	wikiFooter = regexp.MustCompile(`(?is)<div[^>]*class="[^"]*zim-footer[^"]*"`)
259	wikiTitle  = regexp.MustCompile(`(?is)<title[^>]*>(.*?)</title\s*>`)
260	// The infobox. Flattened to text it reads as a run of labels with no
261	// sentence in it, which is noise a small model has to wade through to reach
262	// the prose underneath.
263	wikiTable = regexp.MustCompile(`(?is)<table[^>]*>.*?</table\s*>`)
264)
265
266// stripTables runs to a fixed point because infoboxes nest, and RE2 has no
267// backreference to match an innermost table in one pass.
268func stripTables(h string) string {
269	for i := 0; i < 5; i++ {
270		out := wikiTable.ReplaceAllString(h, " ")
271		if out == h {
272			return h
273		}
274		h = out
275	}
276	return h
277}
278
279var (
280	wikiRow  = regexp.MustCompile(`(?is)<tr[^>]*>(.*?)</tr\s*>`)
281	wikiCell = regexp.MustCompile(`(?is)<t([hd])[^>]*>(.*?)</t[hd]\s*>`)
282	wikiCSS  = regexp.MustCompile(`(?is)<style[^>]*>.*?</style\s*>`)
283)
284
285// What the fact block is allowed to cost, since an infobox can run to forty
286// rows of styling and footnotes and the prose is what the reader came for.
287const (
288	wikiMaxFacts    = 8
289	wikiMaxFactLen  = 160
290	wikiMaxFactsLen = 600
291)
292
293// wikiInfobox pulls the labelled rows out of the first table.
294//
295// The whole table used to be thrown away, which read better and lost the one
296// line that answers who currently holds an office: "Prime Minister of Japan"
297// describes the office in its lead and names the incumbent only here. A row is
298// kept as "Label: value", and a row with no label is kept on its own, since
299// that is the shape the incumbent row comes in.
300// Captions on the images an infobox opens with. They are the only unlabelled
301// rows that are not facts, so they are named rather than guessed at.
302var captionStart = []string{"emblem of", "standard of", "flag of", "seal of", "logo of",
303	"coat of arms", "portrait of", "official portrait", "map of", "photograph of"}
304
305func wikiCaption(v string) bool {
306	l := strings.ToLower(strings.TrimSpace(v))
307	for _, p := range captionStart {
308		if strings.HasPrefix(l, p) {
309			return true
310		}
311	}
312	return false
313}
314
315func wikiInfobox(h string) []string {
316	m := wikiTable.FindString(h)
317	if m == "" {
318		return nil
319	}
320	var out []string
321	total := 0
322	for _, r := range wikiRow.FindAllStringSubmatch(m, -1) {
323		var label, value string
324		for _, c := range wikiCell.FindAllStringSubmatch(r[1], -1) {
325			t := strings.TrimSpace(Text(wikiCSS.ReplaceAllString(c[2], " ")))
326			t = strings.ReplaceAll(t, "\n", " ")
327			if c[1] == "h" && label == "" {
328				label = t
329				continue
330			}
331			if value == "" {
332				value = t
333			}
334		}
335		// A stylesheet that survived, which is what the class rules in a cell
336		// flatten to, and it is never a fact.
337		if strings.Contains(value, "mw-parser-output") || strings.Contains(label, "mw-parser-output") {
338			continue
339		}
340		// An unlabelled row is usually the caption under a picture, and the few
341		// that are not are the office rows worth keeping.
342		if label == "" && wikiCaption(value) {
343			continue
344		}
345		line := strings.TrimSpace(value)
346		if label != "" && value != "" {
347			line = label + ": " + value
348		} else if label != "" && value == "" {
349			continue
350		}
351		if line == "" || len(line) > wikiMaxFactLen {
352			continue
353		}
354		out = append(out, line)
355		total += len(line)
356		if len(out) >= wikiMaxFacts || total >= wikiMaxFactsLen {
357			break
358		}
359	}
360	return out
361}
362
363// wikiArticle returns the lead section as plain text. The mini snapshot has no
364// sections below the lead, and cutting at the first heading anyway means this
365// still returns a lead if the ZIM is ever swapped for a full flavour.
366func wikiArticle(ctx context.Context, d *Deps, base, link string) (title, text string, err error) {
367	if !strings.HasPrefix(link, "/") {
368		link = "/" + link
369	}
370	body, err := wikiGet(ctx, d, base+link)
371	if err != nil {
372		return "", "", err
373	}
374	h := string(body)
375	if m := wikiTitle.FindStringSubmatch(h); m != nil {
376		title = strings.TrimSpace(html.UnescapeString(m[1]))
377	}
378	h = wikiHead.ReplaceAllString(h, " ")
379	if loc := wikiCutHead.FindStringIndex(h); loc != nil {
380		h = h[:loc[0]]
381	}
382	h = wikiH1.ReplaceAllString(h, " ")
383	if loc := wikiFooter.FindStringIndex(h); loc != nil {
384		h = h[:loc[0]]
385	}
386	facts := wikiInfobox(h)
387	h = stripTables(h)
388
389	text = strings.TrimSpace(Text(h))
390	if len(facts) > 0 {
391		text = strings.Join(facts, "\n") + "\n\n" + text
392	}
393	if len(text) > wikiMaxChars {
394		// Cut on a sentence so the model is not handed half a clause.
395		cut := text[:wikiMaxChars]
396		if i := strings.LastIndex(cut, ". "); i > wikiMaxChars/2 {
397			cut = cut[:i+1]
398		}
399		text = cut
400	}
401	if text == "" {
402		return "", "", fmt.Errorf("that article is in the snapshot but its opening section is empty")
403	}
404	return title, text, nil
405}
406
407// wikiGet does not go through get(). The Guard and the budgets exist for third
408// party hosts that rate limit this address, and putting a container on the
409// bridge in the penalty box would take the snapshot out over a blip.
410func wikiGet(ctx context.Context, d *Deps, rawURL string) ([]byte, error) {
411	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
412	if err != nil {
413		return nil, err
414	}
415	resp, err := d.HTTP.Do(req)
416	if err != nil {
417		return nil, fmt.Errorf("the offline wikipedia is not answering: %w", err)
418	}
419	defer resp.Body.Close()
420	if resp.StatusCode >= 400 {
421		return nil, fmt.Errorf("the offline wikipedia answered %d", resp.StatusCode)
422	}
423	return io.ReadAll(io.LimitReader(resp.Body, 8<<20))
424}
425
426// The snapshot's own date, kept after the first call that answers. A model told
427// this is offline but not told how old it is will present a two month old fact
428// as current. Only a success latches, since caching the failure would leave a
429// container that started before kiwix reporting an unknown date for good.
430var (
431	wikiDateMu  sync.Mutex
432	wikiDateVal string
433)
434
435func wikiDate(ctx context.Context, d *Deps, base string) string {
436	wikiDateMu.Lock()
437	defer wikiDateMu.Unlock()
438	if wikiDateVal != "" {
439		return wikiDateVal
440	}
441	const unknown = "an unknown date"
442	body, err := wikiGet(ctx, d, base+"/catalog/v2/entries")
443	if err != nil {
444		return unknown
445	}
446	var feed struct {
447		Entries []struct {
448			Updated string `xml:"updated"`
449		} `xml:"entry"`
450	}
451	if err := xml.Unmarshal(body, &feed); err != nil || len(feed.Entries) == 0 {
452		return unknown
453	}
454	t, err := time.Parse(time.RFC3339, feed.Entries[0].Updated)
455	if err != nil {
456		return unknown
457	}
458	wikiDateVal = t.Format("January 2006")
459	return wikiDateVal
460}