Skip to content
Sign up

flagon-io/g1t

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