diff --git a/src/app/_layout.tsx b/src/app/_layout.tsx index c234166..5a37d8d 100644 --- a/src/app/_layout.tsx +++ b/src/app/_layout.tsx @@ -16,6 +16,7 @@ export default function RootLayout() { + diff --git a/src/app/courses.tsx b/src/app/courses.tsx index 97b8115..63da31e 100644 --- a/src/app/courses.tsx +++ b/src/app/courses.tsx @@ -6,19 +6,42 @@ import { ThemedText } from '@/components/themed-text'; import { ThemedView } from '@/components/themed-view'; import { MaxContentWidth, Spacing } from '@/constants/theme'; import { useTheme } from '@/hooks/use-theme'; +import { getRepo } from '@/lib/db'; import { useCourses } from '@/stores/coursesStore'; +/** Per-course rollup shown on each row. */ +interface CourseStat { + count: number; + hours: number; +} + export default function CoursesScreen() { const theme = useTheme(); const { items, refresh, createCourse, rename, remove } = useCourses(); const [name, setName] = useState(''); const [editing, setEditing] = useState(null); const [editName, setEditName] = useState(''); + const [stats, setStats] = useState>({}); + + // Load the courses list, then roll up lecture count + total hours per course + // in a single pass (one listByCourse() per course; fetched once on focus). + const reload = useCallback(async () => { + await refresh(); + const courses = useCourses.getState().items; + const entries = await Promise.all( + courses.map(async (c) => { + const metas = await getRepo().listByCourse(c.id); + const hours = metas.reduce((sum, m) => sum + m.durationSec, 0) / 3600; + return [c.id, { count: metas.length, hours }] as const; + }), + ); + setStats(Object.fromEntries(entries)); + }, [refresh]); useFocusEffect( useCallback(() => { - void refresh(); - }, [refresh]), + void reload(); + }, [reload]), ); const add = async () => { @@ -26,6 +49,7 @@ export default function CoursesScreen() { if (!n) return; await createCourse({ name: n }); setName(''); + await reload(); }; return ( @@ -74,7 +98,12 @@ export default function CoursesScreen() { ) : ( - {c.name} + + {c.name} + + {statLine(stats[c.id])} + + { @@ -84,7 +113,12 @@ export default function CoursesScreen() { hitSlop={8}> Rename - void remove(c.id)} hitSlop={8}> + { + await remove(c.id); + await reload(); + }} + hitSlop={8}> Delete @@ -97,6 +131,14 @@ export default function CoursesScreen() { ); } +/** "3 lectures · 4.2 h" — falls back to a zero state while stats load. */ +function statLine(stat: CourseStat | undefined): string { + const count = stat?.count ?? 0; + const hours = stat?.hours ?? 0; + const lectures = `${count} ${count === 1 ? 'lecture' : 'lectures'}`; + return `${lectures} · ${hours.toFixed(1)} h`; +} + const styles = StyleSheet.create({ fill: { flex: 1 }, content: { padding: Spacing.three, gap: Spacing.two, maxWidth: MaxContentWidth, width: '100%', alignSelf: 'center' }, diff --git a/src/app/index.tsx b/src/app/index.tsx index cfa01cc..8d64eb4 100644 --- a/src/app/index.tsx +++ b/src/app/index.tsx @@ -72,6 +72,11 @@ export default function LibraryScreen() { + + + Search + + Courses diff --git a/src/app/search.tsx b/src/app/search.tsx new file mode 100644 index 0000000..808d139 --- /dev/null +++ b/src/app/search.tsx @@ -0,0 +1,232 @@ +import { Stack, useFocusEffect, useRouter } from 'expo-router'; +import { useCallback, useEffect, useRef, useState } from 'react'; +import { ActivityIndicator, Pressable, ScrollView, StyleSheet, TextInput, View } from 'react-native'; + +import { ThemedText } from '@/components/themed-text'; +import { ThemedView } from '@/components/themed-view'; +import { MaxContentWidth, Spacing } from '@/constants/theme'; +import { useTheme } from '@/hooks/use-theme'; +import { formatClock } from '@/lib/format'; +import { searchLectures } from '@/lib/search/search'; +import type { SearchHit } from '@/lib/search/types'; +import { useCourses } from '@/stores/coursesStore'; +import { useEmbedding } from '@/stores/embeddingStore'; + +const ACCENT = '#3c87f7'; + +export default function SearchScreen() { + const theme = useTheme(); + const router = useRouter(); + + const courses = useCourses((s) => s.items); + const refreshCourses = useCourses((s) => s.refresh); + + const { status, progress, pending, refreshPending, buildIndex } = useEmbedding(); + + const [query, setQuery] = useState(''); + const [hits, setHits] = useState([]); + const [searching, setSearching] = useState(false); + // Distinguish "haven't searched yet" from "searched, got nothing". + const [searched, setSearched] = useState(false); + + // Refresh courses + pending count whenever the screen gains focus. + useFocusEffect( + useCallback(() => { + void refreshCourses(); + void refreshPending(); + }, [refreshCourses, refreshPending]), + ); + + // Debounced search on query change. The latest run wins (stale guard via seq). + const seq = useRef(0); + useEffect(() => { + const q = query.trim(); + if (!q) { + setHits([]); + setSearched(false); + setSearching(false); + return; + } + const mySeq = ++seq.current; + setSearching(true); + const handle = setTimeout(() => { + void searchLectures(q) + .then((res) => { + if (seq.current !== mySeq) return; // a newer query superseded us + setHits(res); + setSearched(true); + }) + .catch(() => { + if (seq.current !== mySeq) return; + setHits([]); + setSearched(true); + }) + .finally(() => { + if (seq.current !== mySeq) return; + setSearching(false); + }); + }, 250); + return () => clearTimeout(handle); + }, [query]); + + const runNow = useCallback(() => { + const q = query.trim(); + if (!q) return; + const mySeq = ++seq.current; + setSearching(true); + void searchLectures(q) + .then((res) => { + if (seq.current !== mySeq) return; + setHits(res); + setSearched(true); + }) + .catch(() => { + if (seq.current !== mySeq) return; + setHits([]); + setSearched(true); + }) + .finally(() => { + if (seq.current !== mySeq) return; + setSearching(false); + }); + }, [query]); + + const courseName = (cid: string | null) => + cid ? courses.find((c) => c.id === cid)?.name ?? 'Course' : 'Unsorted'; + + const openHit = (hit: SearchHit) => + router.push({ + pathname: '/transcript/[id]', + params: { id: hit.transcriptId, t: String(Math.floor(hit.start)) }, + }); + + const indexReady = status === 'ready' && pending === 0; + + return ( + + + + + Semantic search across your lectures — runs entirely on your device. + + + {pending > 0 && ( + + + Build search index ({pending} {pending === 1 ? 'lecture' : 'lectures'} pending) + + {status === 'indexing' ? ( + <> + + Indexing… {Math.round(progress * 100)}% + + + + ) : ( + void buildIndex()} + style={({ pressed }) => [styles.bannerBtn, { opacity: pressed ? 0.85 : 1 }]}> + Build index + + )} + + )} + + + + {searching ? ( + + ) : query.trim() === '' ? ( + + {indexReady + ? 'Type a question or topic to search across your lectures.' + : 'Build the index to search.'} + + ) : searched && hits.length === 0 ? ( + + No matches. + + ) : ( + hits.map((hit) => ( + openHit(hit)} + /> + )) + )} + + + ); +} + +function ResultRow({ + hit, + courseName, + onPress, +}: { + hit: SearchHit; + courseName: string; + onPress: () => void; +}) { + return ( + [pressed && styles.pressed]}> + + + {hit.text} + + + {courseName} · {formatClock(hit.start)} + + + + ); +} + +function ProgressBar({ value }: { value: number }) { + return ( + + + + ); +} + +const styles = StyleSheet.create({ + fill: { flex: 1 }, + content: { + padding: Spacing.three, + gap: Spacing.three, + maxWidth: MaxContentWidth, + width: '100%', + alignSelf: 'center', + }, + banner: { padding: Spacing.three, borderRadius: Spacing.three, gap: Spacing.two }, + bannerBtn: { + backgroundColor: ACCENT, + paddingVertical: Spacing.two, + borderRadius: Spacing.two, + alignItems: 'center', + }, + bannerBtnText: { color: '#fff', fontWeight: '700' }, + search: { + borderRadius: Spacing.two, + paddingHorizontal: Spacing.three, + paddingVertical: Spacing.two, + fontSize: 15, + }, + card: { padding: Spacing.three, borderRadius: Spacing.three, gap: Spacing.two }, + track: { height: 6, borderRadius: 3, backgroundColor: '#88888833', overflow: 'hidden' }, + bar: { height: 6, borderRadius: 3, backgroundColor: ACCENT }, + pad: { paddingVertical: Spacing.four, textAlign: 'center' }, + pressed: { opacity: 0.7 }, +}); diff --git a/src/app/transcript/[id].tsx b/src/app/transcript/[id].tsx index e17cbb6..7f16f42 100644 --- a/src/app/transcript/[id].tsx +++ b/src/app/transcript/[id].tsx @@ -23,7 +23,12 @@ import { useTranscribe } from '@/stores/transcribeStore'; export default function TranscriptScreen() { const theme = useTheme(); - const { id } = useLocalSearchParams<{ id: string }>(); + const { id, t } = useLocalSearchParams<{ id: string; t?: string }>(); + // Deep-link target time (seconds) from a search hit, if any. + const jumpTo = useMemo(() => { + const n = Number(t); + return Number.isFinite(n) && n >= 0 ? n : null; + }, [t]); const [transcript, setTranscript] = useState(undefined); const [title, setTitle] = useState(''); @@ -32,6 +37,12 @@ export default function TranscriptScreen() { const [saving, setSaving] = useState(false); const [currentTime, setCurrentTime] = useState(0); + const scrollRef = useRef(null); + // Y offset of each rendered segment row, captured via onLayout, for scroll-to. + const segmentYs = useRef>({}); + // Ensure we only auto-jump once per (id, t) deep-link. + const jumpedRef = useRef(false); + // Prefer the just-transcribed in-session audio; otherwise load the persisted // source media so playback/scrub works after a reload (ROADMAP Phase 0). const sessionAudioUrl = useTranscribe((s) => (s.lastTranscriptId === id ? s.audioUrl : undefined)); @@ -90,13 +101,48 @@ export default function TranscriptScreen() { [segments, currentTime], ); - const seek = (t: number) => { + const seek = (time: number) => { const el = audioRef.current; if (!el) return; - el.currentTime = t; + el.currentTime = time; void el.play(); }; + // Reset the one-shot jump guard whenever the deep-link target changes. + useEffect(() => { + jumpedRef.current = false; + }, [id, jumpTo]); + + // Deep-link jump: once the transcript is loaded and we have a finite `t`, + // seek the audio (if present), highlight the matching segment, and scroll it + // into view. Works even without audio (scroll/highlight by time only). + useEffect(() => { + if (jumpTo === null || jumpedRef.current) return; + if (!transcript || segments.length === 0) return; + + const target = segments.findIndex((s) => jumpTo >= s.start && jumpTo < s.end); + const idx = target >= 0 ? target : nearestIndex(segments, jumpTo); + if (idx < 0) return; + + jumpedRef.current = true; + // Drive the highlight by time even if audio isn't available. + setCurrentTime(jumpTo); + if (audioRef.current) seek(jumpTo); + + // Scroll once the row's Y offset has been measured (next frames). + const scrollToIdx = () => { + const y = segmentYs.current[idx]; + if (y === undefined) return false; + scrollRef.current?.scrollTo({ y: Math.max(0, y - Spacing.four), animated: true }); + return true; + }; + if (!scrollToIdx()) { + // Layout may not be measured yet; retry on the next ticks. + const timers = [requestAnimationFrame(() => { if (!scrollToIdx()) setTimeout(scrollToIdx, 120); })]; + return () => timers.forEach(cancelAnimationFrame); + } + }, [jumpTo, transcript, segments, audioUrl]); + const editSegment = (i: number, text: string) => { setSegments((prev) => prev.map((s, idx) => (idx === i ? { ...s, text } : s))); setDirty(true); @@ -130,11 +176,11 @@ export default function TranscriptScreen() { return ( - + { - setTitle(t); + onChangeText={(next) => { + setTitle(next); setDirty(true); }} style={[styles.title, { color: theme.text }]} @@ -160,7 +206,12 @@ export default function TranscriptScreen() { {segments.map((s, i) => ( - + { + segmentYs.current[i] = e.nativeEvent.layout.y; + }} + style={[styles.segRow, i === activeIndex && { backgroundColor: theme.backgroundSelected }]}> seek(s.start)} hitSlop={6}> {formatClock(s.start)} @@ -168,7 +219,7 @@ export default function TranscriptScreen() { editSegment(i, t)} + onChangeText={(next) => editSegment(i, next)} multiline style={[styles.segText, { color: theme.text }]} /> @@ -192,6 +243,24 @@ export default function TranscriptScreen() { ); } +/** Index of the segment whose start is closest to `time` (fallback for jumps + * that don't land inside any segment's [start, end) window). */ +function nearestIndex(segments: Segment[], time: number): number { + if (segments.length === 0) return -1; + let best = 0; + let bestDist = Infinity; + for (let i = 0; i < segments.length; i++) { + const seg = segments[i]; + if (!seg) continue; + const dist = Math.abs(seg.start - time); + if (dist < bestDist) { + bestDist = dist; + best = i; + } + } + return best; +} + function Centered({ children }: { children: React.ReactNode }) { return {children}; } diff --git a/src/lib/db/repo.native.ts b/src/lib/db/repo.native.ts index 8a5dbac..4e68ed3 100644 --- a/src/lib/db/repo.native.ts +++ b/src/lib/db/repo.native.ts @@ -30,7 +30,12 @@ import { type Course, type CourseDraft, } from './schema'; -import type { MediaInput, StorageRepo } from './repo'; +import type { + MediaInput, + StorageRepo, + SegmentVector, + VectorHit, +} from './repo'; // --------------------------------------------------------------------------- // Row shapes (as returned by getAllAsync). SQLite has no boolean/undefined: @@ -80,6 +85,20 @@ interface MediaRow { mime: string; } +// One stored segment embedding. `vector` is a SQLite BLOB; on read expo-sqlite +// hands it back as a Uint8Array whose underlying bytes we reinterpret as a +// Float32Array. start/end are REAL, courseId is nullable text. +interface SegVecRow { + transcriptId: string; + segmentId: string; + start: number; + end: number; + courseId: string | null; + text: string; + vector: Uint8Array; + model: string; +} + // --------------------------------------------------------------------------- // Pure derivation helpers // --------------------------------------------------------------------------- @@ -141,6 +160,33 @@ function rowToCourse(r: CourseRow): Course { return course; } +// ---- vector <-> BLOB codecs -------------------------------------------------- +// SQLite stores embeddings as raw little-endian float32 bytes. We bind a +// Uint8Array view of the Float32Array's buffer on write and reconstruct a +// Float32Array from the returned bytes on read. We copy on read so the result +// is 4-byte aligned (a Uint8Array sub-view may start at an unaligned offset, +// which `new Float32Array(buffer)` requires to be a multiple of 4). +function vectorToBlob(vector: Float32Array): Uint8Array { + return new Uint8Array(vector.buffer, vector.byteOffset, vector.byteLength); +} + +function blobToVector(blob: Uint8Array): Float32Array { + // Copy the bytes into a fresh, aligned ArrayBuffer, then view as float32. + const copy = new Uint8Array(blob.length); + copy.set(blob); + return new Float32Array(copy.buffer); +} + +/** Cosine similarity = dot product of two unit-normalized vectors. */ +function dot(a: Float32Array, b: Float32Array): number { + const n = Math.min(a.length, b.length); + let sum = 0; + for (let i = 0; i < n; i++) { + sum += a[i]! * b[i]!; + } + return sum; +} + // --------------------------------------------------------------------------- // Migration runner // --------------------------------------------------------------------------- @@ -151,7 +197,7 @@ function rowToCourse(r: CourseRow): Course { // idempotent (IF NOT EXISTS / pragma_table_info guards) so they also no-op on // DBs the original single-table code already created at version 0. -const TARGET_VERSION = 2; +const TARGET_VERSION = 3; type Migration = (db: SQLite.SQLiteDatabase) => Promise; @@ -246,6 +292,28 @@ const MIGRATIONS: Migration[] = [ ]); } }, + + // ---- v2 -> v3: segment-vector store for Phase 1 semantic search ------- + async (db) => { + // One row per (transcript, segment) embedding. courseId is denormalized + // from the owning transcript so course-scoped search needs no join; model + // tags the embedding model so a model change can be detected/backfilled. + await db.execAsync(` + CREATE TABLE IF NOT EXISTS segvecs ( + transcriptId TEXT NOT NULL, + segmentId TEXT NOT NULL, + start REAL NOT NULL, + end REAL NOT NULL, + courseId TEXT, + text TEXT NOT NULL, + vector BLOB NOT NULL, + model TEXT NOT NULL, + PRIMARY KEY (transcriptId, segmentId) + ); + CREATE INDEX IF NOT EXISTS idx_segvecs_model ON segvecs (model); + CREATE INDEX IF NOT EXISTS idx_segvecs_courseId ON segvecs (courseId); + `); + }, ]; async function runMigrations(db: SQLite.SQLiteDatabase): Promise { @@ -449,6 +517,11 @@ export const repo: StorageRepo = { id, ], ); + // Segments were patched => stored vectors are stale; drop them so the + // backfill re-embeds from the new text/timing. + if (patch.segments !== undefined) { + await db.runAsync(`DELETE FROM segvecs WHERE transcriptId = ?`, [id]); + } return updated; }, @@ -457,6 +530,8 @@ export const repo: StorageRepo = { // Delete persisted media first so we never orphan a file. await this.removeMedia(id); await db.runAsync(`DELETE FROM transcripts WHERE id = ?`, [id]); + // Cascade: drop this transcript's vectors too. + await db.runAsync(`DELETE FROM segvecs WHERE transcriptId = ?`, [id]); }, async search(query: string): Promise { @@ -514,10 +589,18 @@ export const repo: StorageRepo = { const transcript = JSON.parse(row.json) as Transcript; const updatedAt = Date.now(); const next: Transcript = { ...transcript, courseId, updatedAt }; - await db.runAsync( - `UPDATE transcripts SET courseId = ?, updatedAt = ?, json = ? WHERE id = ?`, - [courseId, updatedAt, JSON.stringify(next), transcriptId], - ); + await db.withTransactionAsync(async () => { + await db.runAsync( + `UPDATE transcripts SET courseId = ?, updatedAt = ?, json = ? WHERE id = ?`, + [courseId, updatedAt, JSON.stringify(next), transcriptId], + ); + // Keep the denormalized courseId on every vector row in sync so + // course-scoped search keeps matching after a move. + await db.runAsync(`UPDATE segvecs SET courseId = ? WHERE transcriptId = ?`, [ + courseId, + transcriptId, + ]); + }); }, // --- courses ------------------------------------------------------------- @@ -621,6 +704,11 @@ export const repo: StorageRepo = { `UPDATE transcripts SET courseId = NULL WHERE courseId = ?`, [id], ); + // Keep denormalized vector rows consistent: their transcripts are now + // Unsorted, so their courseId must follow. + await db.runAsync(`UPDATE segvecs SET courseId = NULL WHERE courseId = ?`, [ + id, + ]); await db.runAsync(`DELETE FROM courses WHERE id = ?`, [id]); }); }, @@ -695,4 +783,108 @@ export const repo: StorageRepo = { transcriptId, ]); }, + + // --- semantic search vectors (Phase 1) ---------------------------------- + async upsertVectors( + transcriptId: string, + model: string, + vectors: SegmentVector[], + ): Promise { + const db = await getDb(); + // Denormalize the transcript's current courseId onto each row so + // course-scoped search never has to join back to transcripts. + const owner = await db.getFirstAsync<{ courseId: string | null }>( + `SELECT courseId FROM transcripts WHERE id = ?`, + [transcriptId], + ); + const courseId = owner ? owner.courseId : null; + await db.withTransactionAsync(async () => { + // Replace, not merge: clear existing rows so a re-embed with fewer + // segments can't leave stale rows behind. + await db.runAsync(`DELETE FROM segvecs WHERE transcriptId = ?`, [ + transcriptId, + ]); + for (const v of vectors) { + await db.runAsync( + `INSERT INTO segvecs + (transcriptId, segmentId, start, end, courseId, text, vector, model) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + transcriptId, + v.segmentId, + v.start, + v.end, + courseId, + v.text, + vectorToBlob(v.vector), + model, + ], + ); + } + }); + }, + + async searchVectors( + query: Float32Array, + opts?: { courseId?: string | null; limit?: number }, + ): Promise { + const db = await getDb(); + // Scope: opts.courseId provided => filter (null => Unsorted via IS NULL); + // omitted => search every course. + let rows: SegVecRow[]; + if (opts && opts.courseId !== undefined) { + if (opts.courseId === null) { + rows = await db.getAllAsync( + `SELECT transcriptId, segmentId, start, end, courseId, text, vector, model + FROM segvecs WHERE courseId IS NULL`, + ); + } else { + rows = await db.getAllAsync( + `SELECT transcriptId, segmentId, start, end, courseId, text, vector, model + FROM segvecs WHERE courseId = ?`, + [opts.courseId], + ); + } + } else { + rows = await db.getAllAsync( + `SELECT transcriptId, segmentId, start, end, courseId, text, vector, model + FROM segvecs`, + ); + } + + const limit = opts?.limit ?? 30; + const hits: VectorHit[] = rows.map((r) => ({ + transcriptId: r.transcriptId, + segmentId: r.segmentId, + start: r.start, + end: r.end, + courseId: r.courseId, + text: r.text, + score: dot(query, blobToVector(r.vector)), + })); + hits.sort((a, b) => b.score - a.score); + return hits.slice(0, limit); + }, + + async clearVectors(transcriptId: string): Promise { + const db = await getDb(); + await db.runAsync(`DELETE FROM segvecs WHERE transcriptId = ?`, [ + transcriptId, + ]); + }, + + async unembeddedIds(model: string): Promise { + const db = await getDb(); + // Transcript ids with ZERO segvecs rows for this model: a LEFT-anti-join + // against the subset of segvecs tagged with the given model. + const rows = await db.getAllAsync<{ id: string }>( + `SELECT t.id AS id FROM transcripts t + WHERE NOT EXISTS ( + SELECT 1 FROM segvecs s + WHERE s.transcriptId = t.id AND s.model = ? + )`, + [model], + ); + return rows.map((r) => r.id); + }, }; diff --git a/src/lib/db/repo.test.ts b/src/lib/db/repo.test.ts index 5d62fd2..96aa466 100644 --- a/src/lib/db/repo.test.ts +++ b/src/lib/db/repo.test.ts @@ -18,6 +18,7 @@ import { parseDraft, type TranscriptDraft, type StorageRepo, + type SegmentVector, } from './index'; // Populated by the migration beforeAll (which seeds v1 then imports repo.web). @@ -322,3 +323,201 @@ describe('StorageRepo (Dexie web impl)', () => { await expect(repo.create(makeDraft({ durationSec: -1 }))).rejects.toThrow(); }); }); + +// --------------------------------------------------------------------------- +// Phase 1: semantic-search vector store +// --------------------------------------------------------------------------- +// +// We exercise the brute-force cosine store with tiny hand-made UNIT vectors so +// the math is checkable by eye. The repo must not hard-code 384 dims — these +// tests use 3-d vectors throughout. Cosine of two unit vectors is their dot +// product, so [1,0,0]·[1,0,0] = 1 and [1,0,0]·[0,1,0] = 0. + +const MODEL = 'Xenova/multilingual-e5-small'; + +// Three orthonormal basis vectors, used as stand-in segment embeddings. +const E0 = () => new Float32Array([1, 0, 0]); +const E1 = () => new Float32Array([0, 1, 0]); +const E2 = () => new Float32Array([0, 0, 1]); + +describe('vectors (semantic search store)', () => { + beforeEach(async () => { + // Wipe transcripts (cascades drop their vectors) + courses between tests. + const all = await repo.list(); + await Promise.all(all.map((m) => repo.remove(m.id))); + const courses = await repo.listCourses(); + await Promise.all(courses.map((c) => repo.deleteCourse(c.id))); + }); + + // Create a transcript whose segments carry stable ids, then return it. + async function makeTranscript(over: Partial = {}) { + return repo.create( + makeDraft({ + segments: [ + { start: 0, end: 1, text: 'alpha' }, + { start: 1, end: 2, text: 'beta' }, + { start: 2, end: 3, text: 'gamma' }, + ], + ...over, + }), + ); + } + + // Build SegmentVector rows pairing each segment id with one of the vectors. + function vecsFor( + t: { segments: { id?: string; start: number; end: number; text: string }[] }, + vectors: Float32Array[], + ): SegmentVector[] { + return t.segments.map((s, i) => ({ + segmentId: s.id!, + start: s.start, + end: s.end, + text: s.text, + vector: vectors[i]!, + })); + } + + it('upsertVectors + searchVectors ranks the matching segment first with score ~1', async () => { + const t = await makeTranscript(); + await repo.upsertVectors(t.id, MODEL, vecsFor(t, [E0(), E1(), E2()])); + + const hits = await repo.searchVectors(E0()); + expect(hits).toHaveLength(3); + // The [1,0,0] segment ("alpha") ranks first with cosine ~1. + expect(hits[0]!.segmentId).toBe(t.segments[0]!.id); + expect(hits[0]!.text).toBe('alpha'); + expect(hits[0]!.score).toBeCloseTo(1, 6); + // The orthogonal segments score ~0. + expect(hits[1]!.score).toBeCloseTo(0, 6); + expect(hits[2]!.score).toBeCloseTo(0, 6); + // Hits carry the segment-jump anchors + denormalized courseId. + expect(hits[0]!.transcriptId).toBe(t.id); + expect(hits[0]!.start).toBe(0); + expect(hits[0]!.end).toBe(1); + expect(hits[0]!.courseId).toBeNull(); + }); + + it('searchVectors honours the limit option', async () => { + const t = await makeTranscript(); + await repo.upsertVectors(t.id, MODEL, vecsFor(t, [E0(), E1(), E2()])); + + const hits = await repo.searchVectors(E0(), { limit: 1 }); + expect(hits).toHaveLength(1); + expect(hits[0]!.segmentId).toBe(t.segments[0]!.id); + }); + + it('upsertVectors replaces prior vectors for a transcript (no stale rows)', async () => { + const t = await makeTranscript(); + await repo.upsertVectors(t.id, MODEL, vecsFor(t, [E0(), E1(), E2()])); + // Re-embed with only the first two segments. + await repo.upsertVectors(t.id, MODEL, [ + { + segmentId: t.segments[0]!.id!, + start: 0, + end: 1, + text: 'alpha', + vector: E0(), + }, + { + segmentId: t.segments[1]!.id!, + start: 1, + end: 2, + text: 'beta', + vector: E1(), + }, + ]); + const hits = await repo.searchVectors(E0()); + expect(hits).toHaveLength(2); + }); + + it('courseId filter scopes search (course vs Unsorted vs all)', async () => { + const course = await repo.createCourse({ name: 'Physics' }); + const inCourse = await makeTranscript({ courseId: course.id }); + const unsorted = await makeTranscript(); + + await repo.upsertVectors(inCourse.id, MODEL, vecsFor(inCourse, [E0(), E1(), E2()])); + await repo.upsertVectors(unsorted.id, MODEL, vecsFor(unsorted, [E0(), E1(), E2()])); + + // Scoped to the course: only that transcript's rows are searched. + const courseHits = await repo.searchVectors(E0(), { courseId: course.id }); + expect(courseHits.every((h) => h.transcriptId === inCourse.id)).toBe(true); + expect(courseHits.every((h) => h.courseId === course.id)).toBe(true); + expect(courseHits).toHaveLength(3); + + // Scoped to Unsorted (null): only the unsorted transcript's rows. + const unsortedHits = await repo.searchVectors(E0(), { courseId: null }); + expect(unsortedHits.every((h) => h.transcriptId === unsorted.id)).toBe(true); + expect(unsortedHits.every((h) => h.courseId === null)).toBe(true); + expect(unsortedHits).toHaveLength(3); + + // No scope => everything (both transcripts). + const allHits = await repo.searchVectors(E0()); + expect(allHits).toHaveLength(6); + }); + + it('unembeddedIds excludes embedded transcripts and includes fresh ones', async () => { + const embedded = await makeTranscript(); + const fresh = await makeTranscript(); + await repo.upsertVectors(embedded.id, MODEL, vecsFor(embedded, [E0(), E1(), E2()])); + + const pending = await repo.unembeddedIds(MODEL); + expect(pending).not.toContain(embedded.id); + expect(pending).toContain(fresh.id); + + // A different model id has no vectors at all => both are unembedded. + const otherModel = await repo.unembeddedIds('some/other-model'); + expect(otherModel).toContain(embedded.id); + expect(otherModel).toContain(fresh.id); + }); + + it('clearVectors empties a transcript and re-adds it to unembeddedIds', async () => { + const t = await makeTranscript(); + await repo.upsertVectors(t.id, MODEL, vecsFor(t, [E0(), E1(), E2()])); + expect(await repo.searchVectors(E0())).toHaveLength(3); + + await repo.clearVectors(t.id); + expect(await repo.searchVectors(E0())).toHaveLength(0); + expect(await repo.unembeddedIds(MODEL)).toContain(t.id); + }); + + it('reassign updates the denormalized courseId on existing hits', async () => { + const course = await repo.createCourse({ name: 'Chemistry' }); + const t = await makeTranscript(); + await repo.upsertVectors(t.id, MODEL, vecsFor(t, [E0(), E1(), E2()])); + + // Initially Unsorted: hits carry courseId null and match the null scope. + expect((await repo.searchVectors(E0(), { courseId: null }))).toHaveLength(3); + + await repo.reassign(t.id, course.id); + + // Now the hits carry the new courseId and match the course scope... + const courseHits = await repo.searchVectors(E0(), { courseId: course.id }); + expect(courseHits).toHaveLength(3); + expect(courseHits.every((h) => h.courseId === course.id)).toBe(true); + // ...and no longer appear under Unsorted. + expect(await repo.searchVectors(E0(), { courseId: null })).toHaveLength(0); + }); + + it('updating segments clears stale vectors (transcript becomes unembedded again)', async () => { + const t = await makeTranscript(); + await repo.upsertVectors(t.id, MODEL, vecsFor(t, [E0(), E1(), E2()])); + expect(await repo.unembeddedIds(MODEL)).not.toContain(t.id); + + // Patching segments invalidates the embeddings => repo drops them. + await repo.update(t.id, { + segments: [{ start: 0, end: 1, text: 'rewritten' }], + }); + + expect(await repo.searchVectors(E0())).toHaveLength(0); + expect(await repo.unembeddedIds(MODEL)).toContain(t.id); + }); + + it('remove cascades to a transcript vectors', async () => { + const t = await makeTranscript(); + await repo.upsertVectors(t.id, MODEL, vecsFor(t, [E0(), E1(), E2()])); + expect(await repo.searchVectors(E0())).toHaveLength(3); + + await repo.remove(t.id); + expect(await repo.searchVectors(E0())).toHaveLength(0); + }); +}); diff --git a/src/lib/db/repo.ts b/src/lib/db/repo.ts index b405748..8d26a94 100644 --- a/src/lib/db/repo.ts +++ b/src/lib/db/repo.ts @@ -24,6 +24,27 @@ export interface MediaInput { mime: string; } +/** One stored embedding for a segment (vector is unit-normalized). */ +export interface SegmentVector { + segmentId: string; + start: number; + end: number; + text: string; + vector: Float32Array; +} + +/** A semantic-search hit at segment granularity. */ +export interface VectorHit { + transcriptId: string; + segmentId: string; + start: number; + end: number; + courseId: string | null; + text: string; + /** Cosine similarity (dot product of unit vectors). */ + score: number; +} + export interface StorageRepo { // --- transcripts --------------------------------------------------------- /** All transcript metadatas, newest first (by createdAt desc). */ @@ -84,4 +105,26 @@ export interface StorageRepo { /** Delete the persisted audio for a transcript. No-op if absent. */ removeMedia(transcriptId: string): Promise; + + // --- semantic search vectors (Phase 1) ---------------------------------- + /** + * Replace all stored vectors for a transcript with `vectors`, tagged with the + * embedding `model` id (so a model change can be detected and re-embedded). + */ + upsertVectors(transcriptId: string, model: string, vectors: SegmentVector[]): Promise; + + /** + * Brute-force cosine search over stored vectors (optionally scoped to a + * course; null = Unsorted), returning the top `limit` segment hits. + */ + searchVectors( + query: Float32Array, + opts?: { courseId?: string | null; limit?: number }, + ): Promise; + + /** Drop all vectors for a transcript. */ + clearVectors(transcriptId: string): Promise; + + /** Transcript ids that have NO vectors for `model` (need embedding/backfill). */ + unembeddedIds(model: string): Promise; } diff --git a/src/lib/db/repo.web.ts b/src/lib/db/repo.web.ts index be1e752..0254a96 100644 --- a/src/lib/db/repo.web.ts +++ b/src/lib/db/repo.web.ts @@ -24,7 +24,12 @@ import { type Course, type CourseDraft, } from './schema'; -import type { StorageRepo, MediaInput } from './repo'; +import type { + StorageRepo, + MediaInput, + SegmentVector, + VectorHit, +} from './repo'; // The shape persisted in the 'transcripts' store: a full Transcript plus the // derived searchText column. searchText is never returned to callers. @@ -40,6 +45,22 @@ interface StoredMedia { mime: string; } +// One row in the 'segvecs' store: a single segment's embedding plus the +// denormalized courseId of its owning transcript (so course-scoped search can +// filter without joining back to transcripts) and the embedding model tag. +// The compound key [transcriptId+segmentId] keeps one row per segment; the +// transcriptId / courseId / model indexes drive the cascade + filter paths. +interface StoredSegVec { + transcriptId: string; + segmentId: string; + start: number; + end: number; + courseId: string | null; + text: string; + vector: Float32Array; + model: string; +} + // --------------------------------------------------------------------------- // Pure derivation helpers (shared by create/update within this file) // --------------------------------------------------------------------------- @@ -68,6 +89,21 @@ function toMeta(row: StoredTranscript): TranscriptMeta { return meta; } +/** + * Cosine similarity between two vectors. Both the query and stored vectors are + * unit-normalized by contract, so this is just their dot product; we do not + * re-normalize here. Mismatched lengths fall back to the common prefix so a + * stale-dim row can never throw (it will simply score poorly). + */ +function dot(a: Float32Array, b: Float32Array): number { + const n = Math.min(a.length, b.length); + let sum = 0; + for (let i = 0; i < n; i++) { + sum += a[i]! * b[i]!; + } + return sum; +} + // --------------------------------------------------------------------------- // Dexie database // --------------------------------------------------------------------------- @@ -76,6 +112,9 @@ class WispDexie extends Dexie { transcripts!: Dexie.Table; courses!: Dexie.Table; media!: Dexie.Table; + // Keyed by the compound [transcriptId+segmentId]; secondary indexes on + // transcriptId (cascade), courseId (scoped search) and model (backfill). + segvecs!: Dexie.Table; constructor() { super('wisp'); @@ -102,6 +141,16 @@ class WispDexie extends Dexie { row.segments = withSegmentIds(row.segments); }); }); + // v3: add the segvecs store for Phase 1 on-device semantic search. The + // existing transcripts/courses/media stores are re-declared unchanged so + // Dexie keeps them; only the new store is added. No data backfill is needed + // — vectors are produced lazily by the embedding pipeline. + this.version(3).stores({ + transcripts: 'id, createdAt, courseId', + courses: 'id, name', + media: 'transcriptId', + segvecs: '[transcriptId+segmentId], transcriptId, courseId, model', + }); } } @@ -191,14 +240,20 @@ export const repo: StorageRepo = { updatedAt: Date.now(), }; await db.transcripts.put(updated); + // Segments were patched => any stored vectors are now stale (they were + // embedded from the old text/timing). Drop them so the backfill re-embeds. + if (patch.segments !== undefined) { + await db.segvecs.where('transcriptId').equals(id).delete(); + } return toTranscript(updated); }, async remove(id: string): Promise { - // Also delete any persisted media for this transcript. - await db.transaction('rw', db.transcripts, db.media, async () => { + // Also delete any persisted media + vectors for this transcript. + await db.transaction('rw', db.transcripts, db.media, db.segvecs, async () => { await db.transcripts.delete(id); await db.media.delete(id); + await db.segvecs.where('transcriptId').equals(id).delete(); }); }, @@ -226,10 +281,18 @@ export const repo: StorageRepo = { async reassign(transcriptId: string, courseId: string | null): Promise { const existing = await db.transcripts.get(transcriptId); if (!existing) return; - await db.transcripts.put({ - ...existing, - courseId, - updatedAt: Date.now(), + await db.transaction('rw', db.transcripts, db.segvecs, async () => { + await db.transcripts.put({ + ...existing, + courseId, + updatedAt: Date.now(), + }); + // Keep the denormalized courseId on every vector row in sync so + // course-scoped search keeps matching after a move. + await db.segvecs + .where('transcriptId') + .equals(transcriptId) + .modify({ courseId }); }); }, @@ -281,7 +344,7 @@ export const repo: StorageRepo = { async deleteCourse(id: string): Promise { // First reassign this course's transcripts to "Unsorted" (courseId=null), // THEN delete the course row, so no transcript is ever orphaned. - await db.transaction('rw', db.transcripts, db.courses, async () => { + await db.transaction('rw', db.transcripts, db.courses, db.segvecs, async () => { const now = Date.now(); const owned = await db.transcripts .filter((r) => r.courseId === id) @@ -291,6 +354,8 @@ export const repo: StorageRepo = { db.transcripts.put({ ...r, courseId: null, updatedAt: now }), ), ); + // Keep denormalized vector rows consistent with their now-Unsorted owners. + await db.segvecs.where('courseId').equals(id).modify({ courseId: null }); await db.courses.delete(id); }); }, @@ -326,4 +391,84 @@ export const repo: StorageRepo = { } }); }, + + // --- semantic search vectors (Phase 1) ---------------------------------- + async upsertVectors( + transcriptId: string, + model: string, + vectors: SegmentVector[], + ): Promise { + await db.transaction('rw', db.transcripts, db.segvecs, async () => { + // Denormalize the transcript's current courseId onto each row so + // course-scoped search never has to join back to transcripts. + const owner = await db.transcripts.get(transcriptId); + const courseId = owner ? owner.courseId ?? null : null; + // Replace, not merge: drop every existing row for this transcript first + // so a re-embed with fewer segments can't leave stale rows behind. + await db.segvecs.where('transcriptId').equals(transcriptId).delete(); + const rows: StoredSegVec[] = vectors.map((v) => ({ + transcriptId, + segmentId: v.segmentId, + start: v.start, + end: v.end, + courseId, + text: v.text, + vector: v.vector, + model, + })); + await db.segvecs.bulkPut(rows); + }); + }, + + async searchVectors( + query: Float32Array, + opts?: { courseId?: string | null; limit?: number }, + ): Promise { + // Scope: when opts.courseId is provided we filter (null => Unsorted, i.e. + // courseId == null); when it is omitted entirely we search every course. + let rows: StoredSegVec[]; + if (opts && opts.courseId !== undefined) { + const scope = opts.courseId; + if (scope === null) { + // IndexedDB cannot index null keys, so Unsorted is filtered in memory. + const all = await db.segvecs.toArray(); + rows = all.filter((r) => r.courseId === null); + } else { + rows = await db.segvecs.where('courseId').equals(scope).toArray(); + } + } else { + rows = await db.segvecs.toArray(); + } + + const limit = opts?.limit ?? 30; + const hits: VectorHit[] = rows.map((r) => ({ + transcriptId: r.transcriptId, + segmentId: r.segmentId, + start: r.start, + end: r.end, + courseId: r.courseId, + text: r.text, + score: dot(query, r.vector), + })); + hits.sort((a, b) => b.score - a.score); + return hits.slice(0, limit); + }, + + async clearVectors(transcriptId: string): Promise { + await db.segvecs.where('transcriptId').equals(transcriptId).delete(); + }, + + async unembeddedIds(model: string): Promise { + // Transcript ids with ZERO segvecs rows for this model. We collect the set + // of embedded ids for the model, then return every transcript id not in it. + const embedded = new Set(); + await db.segvecs + .where('model') + .equals(model) + .each((row) => { + embedded.add(row.transcriptId); + }); + const ids = await db.transcripts.toCollection().primaryKeys(); + return ids.filter((id) => !embedded.has(id)); + }, }; diff --git a/src/lib/embedding/engine.ts b/src/lib/embedding/engine.ts new file mode 100644 index 0000000..e22174a --- /dev/null +++ b/src/lib/embedding/engine.ts @@ -0,0 +1,28 @@ +// The embedding-engine interface: hides the platform-specific sentence-embedding +// backend (transformers.js feature-extraction on web; a follow-up native impl) +// behind one contract, mirroring TranscriptionEngine. Used for on-device +// semantic search (ROADMAP Phase 1) — vectors never leave the device. + +export interface EmbeddingEngine { + readonly platform: 'web' | 'native'; + /** Vector dimensionality (must match EMBED_DIM). */ + readonly dim: number; + + /** Synchronous check: is the model loaded in memory? */ + isLoaded(): boolean; + + /** Ensure the model is loaded. `onProgress` (0..1) fires during download. */ + loadModel(onProgress?: (p: number) => void): Promise; + + /** + * Embed texts into unit-normalized, mean-pooled vectors. `kind` adds the + * asymmetric prefix the model expects (e5: "query:" for search text, + * "passage:" for stored segments) so retrieval is well-calibrated. + */ + embed(texts: string[], kind: 'query' | 'passage'): Promise; +} + +/** Embedding model id — also stored alongside vectors to detect model changes. */ +export const EMBED_MODEL = 'Xenova/multilingual-e5-small'; +/** Output dimensionality of EMBED_MODEL. */ +export const EMBED_DIM = 384; diff --git a/src/lib/embedding/engineImpl.native.ts b/src/lib/embedding/engineImpl.native.ts new file mode 100644 index 0000000..6966ba8 --- /dev/null +++ b/src/lib/embedding/engineImpl.native.ts @@ -0,0 +1,28 @@ +// NATIVE-ONLY embedding engine stub. On-device sentence embeddings on native +// (a transformers/onnxruntime-react-native or whisper.rn-style binding) are a +// follow-up; until then this is a type-correct placeholder that throws clearly +// so the rest of the app stays platform-agnostic and never silently mis-embeds. +// +// Selected by Metro at build time for native targets; never imported by tests. + +import type { EmbeddingEngine } from './engine'; +import { EMBED_DIM } from './engine'; + +const NOT_AVAILABLE = 'On-device embeddings are not available on native yet'; + +export const engine: EmbeddingEngine = { + platform: 'native', + dim: EMBED_DIM, + + isLoaded(): boolean { + return false; + }, + + async loadModel(): Promise { + throw new Error(NOT_AVAILABLE); + }, + + async embed(): Promise { + throw new Error(NOT_AVAILABLE); + }, +}; diff --git a/src/lib/embedding/engineImpl.ts b/src/lib/embedding/engineImpl.ts new file mode 100644 index 0000000..0a03687 --- /dev/null +++ b/src/lib/embedding/engineImpl.ts @@ -0,0 +1,3 @@ +// Base resolver for TypeScript. Metro picks engineImpl.web.ts / engineImpl.native.ts +// by platform extension at build time; tsc resolves this file (defaults to web). +export { engine } from './engineImpl.web'; diff --git a/src/lib/embedding/engineImpl.web.ts b/src/lib/embedding/engineImpl.web.ts new file mode 100644 index 0000000..22acfc0 --- /dev/null +++ b/src/lib/embedding/engineImpl.web.ts @@ -0,0 +1,114 @@ +// WEB-ONLY embedding engine, backed by transformers.js (@huggingface/ +// transformers) running the multilingual-e5-small feature-extraction model in +// the browser (WebGPU when available, else WASM). Produces unit-normalized, +// mean-pooled 384-d vectors for on-device semantic search (Phase 1). +// +// WHY WE LOAD IT FROM A CDN AT RUNTIME (not a static import): +// transformers.js depends on onnxruntime-web, which uses a *computed* dynamic +// import (`import(/*webpackIgnore*/ a)`) and ships WASM — Metro (Expo's web +// bundler) cannot statically bundle either and fails the build. So we never let +// Metro see the package: we load the ESM build from a CDN at runtime via a +// dynamic import hidden behind `new Function` (so Metro's static analyzer can't +// trip over it). This mirrors src/lib/transcription/engineImpl.web.ts exactly. +// +// NOTE: the page must be cross-origin isolated (COOP + COEP) for multi-threaded +// WASM; see docker/nginx.conf. This module is web-only and is NEVER imported by +// any vitest test. + +import type { EmbeddingEngine } from './engine'; +import { EMBED_DIM, EMBED_MODEL } from './engine'; + +// Pin the transformers.js version we load at runtime (matches the ASR engine). +const TRANSFORMERS_CDN = 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.2.0'; + +// `new Function` hides the dynamic import() specifier from Metro's bundler so it +// never tries to resolve/transform transformers.js or onnxruntime-web. +const runtimeImport = new Function('u', 'return import(u)') as (u: string) => Promise; + +// Minimal structural types for the bits of transformers.js we use. +interface FeatureTensor { + /** Nested arrays: one row of floats per input text. */ + tolist(): number[][]; +} +type FeatureExtractor = ( + texts: string[], + opts: { pooling: 'mean'; normalize: boolean }, +) => Promise; +interface PipelineOptions { + device?: string; + dtype?: string; + progress_callback?: (e: { status?: string; progress?: number }) => void; +} +interface TransformersModule { + pipeline: (task: string, model: string, opts?: PipelineOptions) => Promise; + env: { allowLocalModels: boolean }; +} + +let libPromise: Promise | null = null; +async function lib(): Promise { + if (!libPromise) { + libPromise = runtimeImport(TRANSFORMERS_CDN).then((m) => { + // Never read models off the local filesystem in the browser. + m.env.allowLocalModels = false; + return m; + }); + } + return libPromise; +} + +/** The single cached feature-extraction pipeline (one model). */ +let extractor: FeatureExtractor | null = null; +let cachedWebGpu: boolean | undefined; + +async function detectWebGpu(): Promise { + if (cachedWebGpu !== undefined) return cachedWebGpu; + try { + if (typeof navigator === 'undefined' || !('gpu' in navigator)) { + cachedWebGpu = false; + return false; + } + const gpu = (navigator as { gpu?: { requestAdapter(): Promise } }).gpu; + const adapter = await gpu?.requestAdapter(); + cachedWebGpu = adapter != null; + } catch { + cachedWebGpu = false; + } + return cachedWebGpu; +} + +export const engine: EmbeddingEngine = { + platform: 'web', + dim: EMBED_DIM, + + isLoaded(): boolean { + return extractor != null; + }, + + async loadModel(onProgress?: (p: number) => void): Promise { + if (extractor != null) return; // idempotent + const { pipeline } = await lib(); + const webgpu = await detectWebGpu(); + extractor = await pipeline('feature-extraction', EMBED_MODEL, { + // WebGPU + fp16 when available; otherwise 8-bit weights on WASM, which + // stays small to download and runs acceptably on a plain CPU. + device: webgpu ? 'webgpu' : 'wasm', + dtype: webgpu ? 'fp16' : 'q8', + progress_callback: (e) => { + if (e.status === 'progress' && e.progress != null) onProgress?.(e.progress / 100); + }, + }); + }, + + async embed(texts: string[], kind: 'query' | 'passage'): Promise { + if (extractor == null) { + throw new Error('Embedding model is not loaded; call loadModel() first.'); + } + if (texts.length === 0) return []; + // e5 asymmetric convention: prefix with "query: " or "passage: ". + const prefix = `${kind}: `; + const prefixed = texts.map((t) => prefix + t); + const tensor = await extractor(prefixed, { pooling: 'mean', normalize: true }); + // tolist() yields one number[] row per input; convert each to Float32Array. + return tensor.tolist().map((row) => Float32Array.from(row)); + }, +}; diff --git a/src/lib/embedding/index.ts b/src/lib/embedding/index.ts new file mode 100644 index 0000000..c939e99 --- /dev/null +++ b/src/lib/embedding/index.ts @@ -0,0 +1,12 @@ +// Public entry point for the embedding engine. +// +// `./engineImpl` is resolved by Metro to engineImpl.web.ts or +// engineImpl.native.ts by platform extension; the base engineImpl.ts re-export +// (web) is what TypeScript resolves for typechecking. Consumers call +// getEmbeddingEngine() and stay platform-agnostic. +import { engine } from './engineImpl'; + +/** Return the platform-resolved embedding engine. */ +export const getEmbeddingEngine = () => engine; + +export * from './engine'; diff --git a/src/lib/search/lexical.test.ts b/src/lib/search/lexical.test.ts new file mode 100644 index 0000000..70b07a3 --- /dev/null +++ b/src/lib/search/lexical.test.ts @@ -0,0 +1,68 @@ +import { describe, it, expect } from 'vitest'; +import { lexicalRank } from './lexical'; + +describe('lexicalRank', () => { + it('returns [] for an empty query', () => { + const rows = [{ id: 'a', text: 'hello world' }]; + expect(lexicalRank(rows, '')).toEqual([]); + expect(lexicalRank(rows, ' ')).toEqual([]); + expect(lexicalRank(rows, '!!! ???')).toEqual([]); + }); + + it('returns [] when no row matches', () => { + const rows = [ + { id: 'a', text: 'the quick brown fox' }, + { id: 'b', text: 'lazy dog' }, + ]; + expect(lexicalRank(rows, 'elephant')).toEqual([]); + }); + + it('scores by total term-occurrence count', () => { + const rows = [ + { id: 'a', text: 'cat cat cat' }, + { id: 'b', text: 'cat dog' }, + ]; + const out = lexicalRank(rows, 'cat'); + expect(out).toEqual([ + { id: 'a', score: 3 }, + { id: 'b', score: 1 }, + ]); + }); + + it('is case-insensitive and tokenizes on non-word chars', () => { + const rows = [{ id: 'a', text: 'Neural-Networks, neural networks!' }]; + // "neural" appears twice, "networks" appears twice => 4. + expect(lexicalRank(rows, 'Neural networks')).toEqual([{ id: 'a', score: 4 }]); + }); + + it('sums distinct query terms and drops zero-score rows', () => { + const rows = [ + { id: 'a', text: 'gradient descent and gradient ascent' }, // gradient x2 + { id: 'b', text: 'gradient boosting' }, // gradient x1 + { id: 'c', text: 'random forest' }, // 0 -> dropped + ]; + const out = lexicalRank(rows, 'gradient descent'); + // a: gradient(2)+descent(1)=3 ; b: gradient(1)=1 ; c dropped + expect(out).toEqual([ + { id: 'a', score: 3 }, + { id: 'b', score: 1 }, + ]); + }); + + it('sorts by score descending', () => { + const rows = [ + { id: 'low', text: 'alpha' }, + { id: 'high', text: 'alpha alpha alpha' }, + { id: 'mid', text: 'alpha alpha' }, + ]; + expect(lexicalRank(rows, 'alpha').map((r) => r.id)).toEqual(['high', 'mid', 'low']); + }); + + it('skips rows with empty text', () => { + const rows = [ + { id: 'empty', text: '' }, + { id: 'a', text: 'match match' }, + ]; + expect(lexicalRank(rows, 'match')).toEqual([{ id: 'a', score: 2 }]); + }); +}); diff --git a/src/lib/search/lexical.ts b/src/lib/search/lexical.ts new file mode 100644 index 0000000..4fa60c5 --- /dev/null +++ b/src/lib/search/lexical.ts @@ -0,0 +1,44 @@ +// Pure lexical (keyword) ranking. Complements semantic search by catching exact +// term matches the embedding model might under-weight (names, acronyms, codes). +// No I/O, no platform deps — fully unit-testable. + +/** Lowercase + split on non-word chars, dropping empties. */ +function tokenize(s: string): string[] { + return s + .toLowerCase() + .split(/\W+/) + .filter((t) => t.length > 0); +} + +/** + * Rank `rows` by how often the query's terms occur in each row's text. + * + * Scoring: sum, over each query term, of the number of times that term appears + * as a token in the row's text. Rows scoring zero are dropped. Result is sorted + * by score descending (ties keep input order, since Array.sort is stable). + */ +export function lexicalRank( + rows: { id: string; text: string }[], + query: string, +): { id: string; score: number }[] { + const terms = tokenize(query); + if (terms.length === 0) return []; + + const out: { id: string; score: number }[] = []; + for (const row of rows) { + const tokens = tokenize(row.text); + if (tokens.length === 0) continue; + + // Count occurrences per token once, then sum the query terms' counts. + const counts = new Map(); + for (const tok of tokens) counts.set(tok, (counts.get(tok) ?? 0) + 1); + + let score = 0; + for (const term of terms) score += counts.get(term) ?? 0; + + if (score > 0) out.push({ id: row.id, score }); + } + + out.sort((a, b) => b.score - a.score); + return out; +} diff --git a/src/lib/search/rrf.test.ts b/src/lib/search/rrf.test.ts new file mode 100644 index 0000000..cab7cd8 --- /dev/null +++ b/src/lib/search/rrf.test.ts @@ -0,0 +1,55 @@ +import { describe, it, expect } from 'vitest'; +import { rrf } from './rrf'; + +describe('rrf', () => { + it('returns [] for no lists / empty lists', () => { + expect(rrf([])).toEqual([]); + expect(rrf([[], []])).toEqual([]); + }); + + it('ranks a single list by its existing order', () => { + const out = rrf([['a', 'b', 'c']]); + expect(out.map((r) => r.id)).toEqual(['a', 'b', 'c']); + // rank-1 score = 1/(60+1) + expect(out[0]!.score).toBeCloseTo(1 / 61, 10); + expect(out[1]!.score).toBeCloseTo(1 / 62, 10); + }); + + it('sums contributions for ids in multiple lists', () => { + // 'a' is rank 1 in list1 and rank 1 in list2 => 2/(k+1). + const out = rrf([ + ['a', 'b'], + ['a', 'c'], + ]); + const a = out.find((r) => r.id === 'a')!; + expect(a.score).toBeCloseTo(2 / 61, 10); + // 'a' should rank above 'b' and 'c' (each only 1/(k+2)). + expect(out[0]!.id).toBe('a'); + }); + + it('an item ranked high in both lists beats one ranked high in only one', () => { + const semantic = ['x', 'y', 'z']; + const lexical = ['y', 'x', 'w']; + const out = rrf([semantic, lexical]); + // y: 1/(60+2) + 1/(60+1) ; x: 1/(60+1) + 1/(60+2) -> tie, but both top. + expect(new Set([out[0]!.id, out[1]!.id])).toEqual(new Set(['x', 'y'])); + // z and w each appear once, so they rank lower. + expect(out.slice(2).map((r) => r.id).sort()).toEqual(['w', 'z']); + }); + + it('respects a custom k', () => { + const out = rrf([['a']], 0); + // rank 1 with k=0 => 1/1 = 1. + expect(out[0]!.score).toBeCloseTo(1, 10); + }); + + it('sorts by fused score descending', () => { + const out = rrf([ + ['a', 'b', 'c'], + ['a', 'b', 'c'], + ]); + expect(out.map((r) => r.id)).toEqual(['a', 'b', 'c']); + expect(out[0]!.score).toBeGreaterThan(out[1]!.score); + expect(out[1]!.score).toBeGreaterThan(out[2]!.score); + }); +}); diff --git a/src/lib/search/rrf.ts b/src/lib/search/rrf.ts new file mode 100644 index 0000000..ccd4bc1 --- /dev/null +++ b/src/lib/search/rrf.ts @@ -0,0 +1,30 @@ +// Pure Reciprocal Rank Fusion (RRF). Combines several ranked id lists into one +// ranking by summing 1/(k + rank) for each id across the lists it appears in +// (rank is 1-based; k dampens the contribution of low ranks). This is a robust, +// score-free way to fuse heterogeneous signals (semantic + lexical) — only the +// ORDER of each input list matters, not its raw scores. +// +// No I/O, no platform deps — fully unit-testable. + +/** + * Fuse ordered id lists into a single ranking, highest fused score first. + * + * @param lists ordered id lists (rank 1 = first element of each list). + * @param k RRF damping constant (default 60, the canonical value). + */ +export function rrf(lists: string[][], k = 60): { id: string; score: number }[] { + const scores = new Map(); + + for (const list of lists) { + for (let i = 0; i < list.length; i++) { + const id = list[i]; + if (id === undefined) continue; + const rank = i + 1; // 1-based + scores.set(id, (scores.get(id) ?? 0) + 1 / (k + rank)); + } + } + + return [...scores.entries()] + .map(([id, score]) => ({ id, score })) + .sort((a, b) => b.score - a.score); +} diff --git a/src/lib/search/search.ts b/src/lib/search/search.ts new file mode 100644 index 0000000..0284ee3 --- /dev/null +++ b/src/lib/search/search.ts @@ -0,0 +1,77 @@ +// Cross-lecture search orchestrator (Phase 1). Embeds the query on-device, +// pulls semantic candidates from the vector store, re-ranks them lexically, and +// fuses the two signals with Reciprocal Rank Fusion. Returns segment-level hits +// (with the matching signal tagged) that the UI can jump to. 100% on-device. +// +// IMPORTANT: relative imports only inside src/lib (vitest has no '@/*' alias). + +import { getRepo } from '../db'; +import type { VectorHit } from '../db/repo'; +import { getEmbeddingEngine } from '../embedding'; +import { lexicalRank } from './lexical'; +import { rrf } from './rrf'; +import type { SearchHit, SearchOptions } from './types'; + +/** How many semantic candidates to pull before re-ranking/fusing. */ +const CANDIDATE_LIMIT = 50; +/** Default number of hits returned to the caller. */ +const DEFAULT_LIMIT = 30; + +/** + * Search the user's lectures for `query`, returning fused semantic + lexical + * hits at segment granularity (best first). Empty/whitespace queries return []. + */ +export async function searchLectures( + query: string, + opts?: SearchOptions, +): Promise { + const q = query.trim(); + if (q.length === 0) return []; + + // 1) Embed the query on-device (lazy-load the model on first use). + const engine = getEmbeddingEngine(); + if (!engine.isLoaded()) await engine.loadModel(); + const [vec] = await engine.embed([q], 'query'); + if (!vec) return []; + + // 2) Semantic candidates from the vector store (cosine, top CANDIDATE_LIMIT). + const cand = await getRepo().searchVectors(vec, { + courseId: opts?.courseId, + limit: CANDIDATE_LIMIT, + }); + if (cand.length === 0) return []; + + // O(1) candidate lookup by segmentId for the fusion map-back step. + const bySegment = new Map(); + for (const c of cand) bySegment.set(c.segmentId, c); + + // 3) Two ranked id lists over the SAME candidate set. + // - semantic: candidates are already cosine-desc from the repo. + const semanticIds = cand.map((c) => c.segmentId); + const semanticSet = new Set(semanticIds); + // - lexical: keyword re-rank of just these candidates' texts. + const lexicalRanked = lexicalRank( + cand.map((c) => ({ id: c.segmentId, text: c.text })), + q, + ); + const lexicalIds = lexicalRanked.map((r) => r.id); + const lexicalSet = new Set(lexicalIds); + + // 4) Fuse the two orderings. + const fused = rrf([semanticIds, lexicalIds]); + + // 5) Map fused ids back to their VectorHit and tag the matching signal. + const limit = opts?.limit ?? DEFAULT_LIMIT; + const hits: SearchHit[] = []; + for (const { id } of fused) { + const hit = bySegment.get(id); + if (!hit) continue; + const inSemantic = semanticSet.has(id); + const inLexical = lexicalSet.has(id); + const via: SearchHit['via'] = inSemantic && inLexical ? 'both' : inLexical ? 'lexical' : 'semantic'; + hits.push({ ...hit, via }); + if (hits.length >= limit) break; + } + + return hits; +} diff --git a/src/lib/search/types.ts b/src/lib/search/types.ts new file mode 100644 index 0000000..4d7b155 --- /dev/null +++ b/src/lib/search/types.ts @@ -0,0 +1,14 @@ +// Shared types for cross-lecture search (Phase 1). +import type { VectorHit } from '../db/repo'; + +/** A search result at segment granularity, with which signal(s) matched. */ +export interface SearchHit extends VectorHit { + via: 'semantic' | 'lexical' | 'both'; +} + +export interface SearchOptions { + /** Scope to a course (null = Unsorted); omit for all. */ + courseId?: string | null; + /** Max hits to return. */ + limit?: number; +} diff --git a/src/stores/embeddingStore.ts b/src/stores/embeddingStore.ts new file mode 100644 index 0000000..d93c404 --- /dev/null +++ b/src/stores/embeddingStore.ts @@ -0,0 +1,118 @@ +// Drives on-device indexing for semantic search (Phase 1): loads the embedding +// model, embeds each transcript's segments into unit vectors, and persists them +// to the StorageRepo's vector store. UI/session state only — the repo is the +// record of truth for what's been embedded. +import { create } from 'zustand'; + +import { getRepo } from '@/lib/db'; +import { EMBED_MODEL, getEmbeddingEngine } from '@/lib/embedding'; + +type Status = 'idle' | 'indexing' | 'ready'; + +interface EmbeddingState { + status: Status; + /** Build progress in [0,1] (0.2 = model load, 0.8 = embedding). */ + progress: number; + /** Count of transcripts with no vectors for the active model. */ + pending: number; + + /** Refresh `pending` from the repo (cheap; no model load). */ + refreshPending: () => Promise; + /** Embed every un-embedded transcript. Loads the model first. */ + buildIndex: () => Promise; + /** Embed a single transcript (used right after a new transcription). */ + embedOne: (transcriptId: string) => Promise; +} + +/** + * Embed and persist vectors for one transcript. Replaces any existing vectors + * for it (upsertVectors semantics). No-op if the transcript or its segments are + * gone. Caller is responsible for model loading and status/progress bookkeeping. + */ +async function embedTranscript(transcriptId: string): Promise { + const repo = getRepo(); + const t = await repo.get(transcriptId); + if (!t) return; + if (t.segments.length === 0) { + // Nothing to embed, but record an (empty) vector set so it stops counting + // as "unembedded" for this model. + await repo.upsertVectors(transcriptId, EMBED_MODEL, []); + return; + } + const eng = getEmbeddingEngine(); + const vecs = await eng.embed( + t.segments.map((s) => s.text), + 'passage', + ); + await repo.upsertVectors( + transcriptId, + EMBED_MODEL, + t.segments.map((s, j) => ({ + segmentId: s.id ?? String(j), + start: s.start, + end: s.end, + text: s.text, + vector: vecs[j] ?? new Float32Array(eng.dim), + })), + ); +} + +export const useEmbedding = create((set, get) => ({ + status: 'idle', + progress: 0, + pending: 0, + + refreshPending: async () => { + const ids = await getRepo().unembeddedIds(EMBED_MODEL); + set({ pending: ids.length }); + }, + + buildIndex: async () => { + // Guard against concurrent runs. + if (get().status === 'indexing') return; + set({ status: 'indexing', progress: 0 }); + try { + const eng = getEmbeddingEngine(); + // Model load is the first 20% of the bar. + await eng.loadModel((p) => set({ progress: p * 0.2 })); + + const ids = await getRepo().unembeddedIds(EMBED_MODEL); + if (ids.length === 0) { + set({ status: 'ready', progress: 1, pending: 0 }); + return; + } + + for (let i = 0; i < ids.length; i++) { + const id = ids[i]; + if (id === undefined) continue; + await embedTranscript(id); + // Remaining 80% spread across the transcripts. + set({ progress: 0.2 + 0.8 * ((i + 1) / ids.length) }); + } + + set({ status: 'ready', progress: 1, pending: 0 }); + } catch (err) { + // Surface progress reset but don't crash callers; leave pending intact so + // a retry can pick up where it left off. + set({ status: 'idle', progress: 0 }); + throw err; + } + }, + + embedOne: async (transcriptId) => { + // Guard against clobbering an in-flight full build. + if (get().status === 'indexing') return; + set({ status: 'indexing' }); + try { + const eng = getEmbeddingEngine(); + // Lazy model load (no progress weighting for a single transcript). + if (!eng.isLoaded()) await eng.loadModel(); + await embedTranscript(transcriptId); + const pending = (await getRepo().unembeddedIds(EMBED_MODEL)).length; + set({ status: 'ready', pending }); + } catch (err) { + set({ status: 'idle' }); + throw err; + } + }, +})); diff --git a/src/stores/transcribeStore.ts b/src/stores/transcribeStore.ts index 90b58ce..c66afc6 100644 --- a/src/stores/transcribeStore.ts +++ b/src/stores/transcribeStore.ts @@ -10,6 +10,7 @@ import { DEFAULT_MODEL } from '@/lib/models/catalog'; import { getEngine } from '@/lib/transcription'; import { transcribe } from '@/lib/transcription/pipeline'; import type { ModelId, Segment } from '@/lib/types'; +import { useEmbedding } from './embeddingStore'; type Status = 'idle' | 'loading' | 'transcribing' | 'done' | 'error'; @@ -106,6 +107,11 @@ export const useTranscribe = create((set, get) => ({ console.warn('[wisp] could not persist source audio:', e); } + // Index the new transcript for semantic search in the background. Lazy + // model load happens inside embedOne; this must never block or fail the + // transcription itself, so it's fire-and-forget with a swallowed error. + void useEmbedding.getState().embedOne(saved.id).catch(() => {}); + set({ status: 'done', progress: 1, partial: segments, lastTranscriptId: saved.id }); return saved.id; } catch (err) {