Skip to content
Sign up

flagon-io/g1t

public
0
156 linesCodeBlameRaw
1/**
2 * Recall's choices, apart from where the passages come from: how the
3 * vector query is filtered, which matches are close enough, and how
4 * passages are picked (an agent's required reading first, at most two per
5 * document, words only to fill). Pure; src/index.ts `recallForAgent` and
6 * hybrid search run it over Vectorize and D1.
7 */
8
9/**
10 * The least cosine similarity a passage needs to count as being about the
11 * query. Measured on g1t's own docs folder (606 passages, chunked and
12 * embedded as here, bge-base-en-v1.5 with Workers AI's mean pooling): the
13 * best passage for 16 questions the docs answer scored 0.72 to 0.84, while
14 * the best for 16 they don't ("hello", a recipe, Postgres tuning, a
15 * vacation policy, SSO with Okta) scored 0.54 to 0.67. At 0.6, 11 of those
16 * 16 pulled in passages (up to 89 above it); 0.7 keeps every answer and
17 * keeps recall quiet when the docs say nothing about the question.
18 */
19export const MEANING_FLOOR = 0.7;
20/** The score a passage found only by its words carries: below the floor, so callers can tell. */
21export const WORDS_SCORE = 0.5;
22/** Nearest passages asked of the index. */
23export const TOP_K = 24;
24/** Asked when the index can't filter by space and recall filters after: more, so enough are left. */
25export const TOP_K_UNFILTERED = 50;
26/** The most space ids put in one `$in` filter; past it, filter after the query (a Vectorize filter is at most 2 KB of JSON). */
27export const MAX_IN_FILTER = 40;
28/** Passages from one page or file at most. */
29export const PER_DOC = 2;
30export const DEFAULT_LIMIT = 5;
31export const MAX_LIMIT = 10;
32
33export function recallLimit(limit: unknown): number {
34 const n = Math.floor(Number(limit));
35 if (!Number.isFinite(n) || n < 1) return DEFAULT_LIMIT;
36 return Math.min(n, MAX_LIMIT);
37}
38
39/**
40 * How to ask the index: by the allowed spaces when there are few enough
41 * to name, otherwise by workspace alone, more of them, filtered after.
42 * Null when nothing may be read.
43 */
44export function vectorQueryPlan(workspaceId: string, allowed: string[]): { topK: number; filter: { workspace_id: string; space_ids?: string[] }; filterAfter: boolean } | null {
45 const ids = [...new Set(allowed)];
46 if (!ids.length) return null;
47 if (ids.length <= MAX_IN_FILTER) return { topK: TOP_K, filter: { workspace_id: workspaceId, space_ids: ids }, filterAfter: false };
48 return { topK: TOP_K_UNFILTERED, filter: { workspace_id: workspaceId }, filterAfter: true };
49}
50
51/**
52 * The spaces recall looks in first: those asked for (an agent's required
53 * reading) that may be read. Never wider than what may be read.
54 */
55export function requiredSpaces(allowed: string[], asked: unknown): string[] {
56 if (!Array.isArray(asked)) return [];
57 const may = new Set(allowed);
58 return [...new Set(asked.map(String))].filter((id) => may.has(id));
59}
60
61export type Candidate = {
62 /** The passage's id: `<doc>:<seq>`. */
63 id: string;
64 /** Its page's or file's id. */
65 doc_id: string;
66 space_id: string;
67 score: number;
68 /** Found by meaning (the index) or by its words (full text). */
69 by: "meaning" | "words";
70};
71
72/**
73 * The passages to hand over, best first: by meaning above the floor,
74 * required spaces first, then the rest; then, if that is fewer than
75 * `limit`, by words, required spaces first. At most PER_DOC from one
76 * document, each passage once, only from `allowed`.
77 */
78export function pickPassages<T extends Candidate>(candidates: T[], options: { allowed: Set<string>; required?: string[]; limit: number; floor?: number }): T[] {
79 const floor = options.floor ?? MEANING_FLOOR;
80 const required = new Set(options.required ?? []);
81 const usable = candidates.filter((c) => options.allowed.has(c.space_id));
82 const meaning = usable.filter((c) => c.by === "meaning" && c.score >= floor).sort((a, b) => b.score - a.score);
83 const words = usable.filter((c) => c.by === "words");
84 const out: T[] = [];
85 const taken = new Set<string>();
86 const perDoc = new Map<string, number>();
87 const take = (list: T[]) => {
88 for (const c of list) {
89 if (out.length >= options.limit) return;
90 if (taken.has(c.id)) continue;
91 const n = perDoc.get(c.doc_id) ?? 0;
92 if (n >= PER_DOC) continue;
93 taken.add(c.id);
94 perDoc.set(c.doc_id, n + 1);
95 out.push(c);
96 }
97 };
98 take(meaning.filter((c) => required.has(c.space_id)));
99 take(meaning.filter((c) => !required.has(c.space_id)));
100 take(words.filter((c) => required.has(c.space_id)));
101 take(words.filter((c) => !required.has(c.space_id)));
102 return out;
103}
104
105/**
106 * Hybrid search's order for people: each document's place in the word
107 * results and in the meaning results, fused by reciprocal rank (k = 60),
108 * so a page both find comes first and either alone still counts.
109 */
110export function fuseRanks(words: string[], meaning: string[], k = 60): string[] {
111 const score = new Map<string, number>();
112 const add = (ids: string[]) =>
113 [...new Set(ids)].forEach((id, rank) => {
114 score.set(id, (score.get(id) ?? 0) + 1 / (k + rank + 1));
115 });
116 add(words);
117 add(meaning);
118 return [...score.entries()].sort((a, b) => b[1] - a[1]).map(([id]) => id);
119}
120
121/**
122 * A query's embeddings, kept a minute per isolate: an agent asked the same
123 * thing again in a session, or a search page reloaded, embeds once.
124 */
125export class QueryCache {
126 private readonly entries = new Map<string, { at: number; vector: number[] }>();
127 private readonly ttlMs: number;
128 private readonly max: number;
129 constructor(ttlMs = 60_000, max = 200) {
130 this.ttlMs = ttlMs;
131 this.max = max;
132 }
133
134 get(key: string, now = Date.now()): number[] | null {
135 const hit = this.entries.get(key);
136 if (!hit) return null;
137 if (now - hit.at > this.ttlMs) {
138 this.entries.delete(key);
139 return null;
140 }
141 return hit.vector;
142 }
143
144 set(key: string, vector: number[], now = Date.now()): void {
145 if (this.entries.size >= this.max) {
146 for (const [k, v] of this.entries) if (now - v.at > this.ttlMs) this.entries.delete(k);
147 while (this.entries.size >= this.max) this.entries.delete(this.entries.keys().next().value!);
148 }
149 this.entries.set(key, { at: now, vector });
150 }
151}
152
153/** A query as the cache keys it: the same words, however spaced or cased. */
154export function queryKey(query: string): string {
155 return String(query ?? "").replace(/\s+/g, " ").trim().toLowerCase().slice(0, 2000);
156}