diff --git a/backend/package-lock.json b/backend/package-lock.json index 4462ada..86118c4 100644 --- a/backend/package-lock.json +++ b/backend/package-lock.json @@ -21,6 +21,7 @@ "dotenv": "^16.5.0", "express": "^5.1.0", "express-rate-limit": "^7.5.1", + "gpt-tokenizer": "^3.4.0", "jsonwebtoken": "^9.0.2", "mammoth": "^1.8.0", "mongodb": "^6.17.0", @@ -6266,6 +6267,12 @@ "url": "https://github.com/sponsors/ljharb" } }, + "node_modules/gpt-tokenizer": { + "version": "3.4.0", + "resolved": "https://registry.npmjs.org/gpt-tokenizer/-/gpt-tokenizer-3.4.0.tgz", + "integrity": "sha512-wxFLnhIXTDjYebd9A9pGl3e31ZpSypbpIJSOswbgop5jLte/AsZVDvjlbEuVFlsqZixVKqbcoNmRlFDf6pz/UQ==", + "license": "MIT" + }, "node_modules/graceful-fs": { "version": "4.2.11", "resolved": "https://registry.npmjs.org/graceful-fs/-/graceful-fs-4.2.11.tgz", diff --git a/backend/package.json b/backend/package.json index c0779e4..02161d7 100644 --- a/backend/package.json +++ b/backend/package.json @@ -59,10 +59,11 @@ "dotenv": "^16.5.0", "express": "^5.1.0", "express-rate-limit": "^7.5.1", + "gpt-tokenizer": "^3.4.0", "jsonwebtoken": "^9.0.2", + "mammoth": "^1.8.0", "mongodb": "^6.17.0", "mongoose": "^8.16.0", - "mammoth": "^1.8.0", "multer": "^2.0.0", "nodemailer": "^7.0.4", "pdf-parse": "^1.1.1", diff --git a/backend/src/app.ts b/backend/src/app.ts index 07184cb..40e600f 100644 --- a/backend/src/app.ts +++ b/backend/src/app.ts @@ -14,6 +14,7 @@ import connectDB from "./db/mongo.js"; import fileRouter from "./routes/file.js"; import { requestContextMiddleware } from "./middleware/requestContext.js"; import { csrfMiddleware } from "./middleware/csrf.js"; +import { warmupWeaviate, startWeaviateKeepalive } from "./db/weaviate_client.js"; const app = express(); const PORT = process.env.PORT || 4000; @@ -38,6 +39,11 @@ app.use(csrfMiddleware); await connectDB() await setupPubSub(); +// Pre-warm Weaviate so the first user query doesn't pay TCP+TLS handshake cost. +// Don't block startup if it fails β€” search will retry on the first real query. +warmupWeaviate().catch(() => { /* logged inside */ }); +startWeaviateKeepalive(); + app.get("/", (req, res) => { res.send("SmartDrive backend running πŸš€"); diff --git a/backend/src/db/weaviate_client.ts b/backend/src/db/weaviate_client.ts index 2eea01d..89a7e78 100644 --- a/backend/src/db/weaviate_client.ts +++ b/backend/src/db/weaviate_client.ts @@ -1,29 +1,90 @@ -import weaviate from 'weaviate-client' +import weaviate, { WeaviateClient } from 'weaviate-client'; import logger from '../logger.js'; const WEAVIATE_URL = process.env.WEAVIATE_URL as string; const WEAVIATE_API_KEY = process.env.WEAVIATE_API_KEY as string; -const getWeaviateClient = async () => { +// ---------- Singleton connection ---------- +// Previously, every call to getWeaviateClient() created a NEW connection, +// adding 200-500ms TCP+TLS handshake to every search and chat query. +// Now we reuse one client across the process lifetime. +let _client: WeaviateClient | null = null; +let _connecting: Promise | null = null; + +const _connect = async (): Promise => { if (!WEAVIATE_URL || !WEAVIATE_API_KEY) { - logger.error("WEAVIATE_CLUSTER_URL or WEAVIATE_API_KEY environment variables not set.") - return; + logger.error("WEAVIATE_URL or WEAVIATE_API_KEY environment variables not set."); + return null; } try { - const client = await weaviate.connectToWeaviateCloud( - WEAVIATE_URL, { - authCredentials: new weaviate.ApiKey(WEAVIATE_API_KEY), - } - ) - logger.info("Successfully connected to Weaviate.") - - return client + const c = await weaviate.connectToWeaviateCloud( + WEAVIATE_URL, + { authCredentials: new weaviate.ApiKey(WEAVIATE_API_KEY) }, + ); + logger.info("Successfully connected to Weaviate (cached for process lifetime)."); + return c; + } catch (e) { + logger.error(`Failed to connect to Weaviate: ${e}`); + return null; } - catch (e) { - logger.error(`Failed to connect to Weaviate: ${e}`) - return +}; + +const getWeaviateClient = async (): Promise => { + if (_client) return _client; + + // If a connection is already being established, await the same promise + // instead of starting a second handshake (avoids thundering-herd on cold start). + if (_connecting) return _connecting; + + _connecting = _connect().then((c) => { + _client = c; + _connecting = null; + return c; + }); + return _connecting; +}; + +// ---------- Pre-warm on boot ---------- +// Called from app.ts at startup so the first request doesn't pay connection cost. + +export const warmupWeaviate = async (): Promise => { + const start = Date.now(); + const c = await getWeaviateClient(); + if (c) { + try { + await c.isLive(); + logger.info(`Weaviate warmup OK in ${Date.now() - start}ms`); + } catch (e) { + logger.warn(`Weaviate warmup failed: ${e}`); + } } -} +}; + +// ---------- Keepalive ---------- +// Cloud Run keeps idle TCP for some time, but Weaviate Cloud may also drop +// connections after periods of inactivity. A periodic isLive() ping keeps +// the connection warm and lets us detect dropped connections early. + +const KEEPALIVE_MS = 5 * 60 * 1000; // 5 minutes +let _keepaliveTimer: NodeJS.Timeout | null = null; + +export const startWeaviateKeepalive = (): void => { + if (_keepaliveTimer) return; // already running + _keepaliveTimer = setInterval(async () => { + try { + const c = await getWeaviateClient(); + if (c) await c.isLive(); + } catch (e) { + // If the keepalive fails, reset the client so the next real query + // re-establishes. Better to fail fast than serve from a dead conn. + logger.warn(`Weaviate keepalive failed, resetting client: ${e}`); + _client = null; + } + }, KEEPALIVE_MS); + // Don't keep the process alive just for this timer. + if (typeof _keepaliveTimer.unref === "function") _keepaliveTimer.unref(); + logger.info(`Weaviate keepalive started (interval ${KEEPALIVE_MS / 1000}s)`); +}; -export default getWeaviateClient; \ No newline at end of file +export default getWeaviateClient; diff --git a/backend/src/handlers/fileHandler.ts b/backend/src/handlers/fileHandler.ts index b6a1261..3f8e25d 100644 --- a/backend/src/handlers/fileHandler.ts +++ b/backend/src/handlers/fileHandler.ts @@ -21,6 +21,14 @@ const capHistory = (history: ChatTurn[] | undefined): ChatTurn[] => { return valid.slice(-HISTORY_TURN_CAP); }; +// Personalization: bump access count + recency on file interaction. Fire-and-forget. +const touchFileAccess = (fileId: string): void => { + UserFile.updateOne( + { _id: fileId }, + { $inc: { accessCount: 1 }, $set: { lastAccessedAt: new Date() } }, + ).catch((err) => logger.warn(`touchFileAccess fileId=${fileId} failed: ${err}`)); +}; + const getUserFile = (fileRecord: UserFileType | null) => { const filePath = `${fileRecord?.userId}/${fileRecord?.fileHash}`; return bucket.file(filePath); @@ -90,6 +98,7 @@ const generateFileSignedUrl = async (req: AuthenticatedRequest, res: Response): } const [url] = await file.getSignedUrl(options); + touchFileAccess(fileId); // R6 personalization signal res.status(200).json({ url }); return @@ -269,6 +278,7 @@ const prepareChat = async (req: AuthenticatedRequest, res: Response): Promise => { const { userQuery, queryCollection } = req.query; @@ -31,5 +33,45 @@ const queryHandler = async (req: AuthenticatedRequest, res: Response): Promise => { + try { + const userId = req.user?._id; + if (!userId) { + res.status(401).json({ error: 'Not authorized.' }); + return; + } + const { query, fileId, rank } = req.body as { + query?: string; + fileId?: string; + rank?: number; + }; + if (!query || !fileId || typeof rank !== 'number' || rank < 1) { + res.status(400).json({ error: 'Missing query/fileId/rank.' }); + return; + } + if (!mongoose.Types.ObjectId.isValid(fileId)) { + res.status(400).json({ error: 'Invalid fileId.' }); + return; + } + const day = new Date().toISOString().slice(0, 10); + // Fire and forget β€” return 200 immediately, don't block on the write. + SearchClick.create({ + userId, + query: query.trim().toLowerCase().slice(0, 500), + fileId: new mongoose.Types.ObjectId(fileId), + rank: Math.min(rank, 100), + day, + }).catch((err) => logger.warn(`searchClick insert failed: ${err}`)); + res.status(200).json({ ok: true }); + } catch (error) { + logger.error('logSearchClick failed:', error); + res.status(500).json({ error: 'Server Error' }); + } +}; +export { logSearchClick }; export default queryHandler; \ No newline at end of file diff --git a/backend/src/handlers/uploadHandler.ts b/backend/src/handlers/uploadHandler.ts index 776f9db..3d7257d 100644 --- a/backend/src/handlers/uploadHandler.ts +++ b/backend/src/handlers/uploadHandler.ts @@ -162,6 +162,9 @@ const getUploads = async (req: AuthenticatedRequest, res: Response): Promise("SearchClick", schema); +export default SearchClick; diff --git a/backend/src/models/userFileModel.ts b/backend/src/models/userFileModel.ts index f30e41d..dd27e23 100644 --- a/backend/src/models/userFileModel.ts +++ b/backend/src/models/userFileModel.ts @@ -18,6 +18,13 @@ export interface UserFileType extends mongoose.Document { * no body embedding, no chunks. Only filename + metadata are indexed so the * user can still find the file by name. Chat is disabled for private files. */ isPrivate?: boolean; + /** Personalization signal β€” files the user interacts with rank higher in + * search. Touched on chat-prep, chat-stream, view, download. */ + lastAccessedAt?: Date | null; + accessCount?: number; + /** Extraction progress for UI ("extracting page 3 of 12"). Worker + * updates this periodically during processing. Cleared when status='done'. */ + extractionProgress?: { current: number; total: number; stage: string } | null; createdAt: Date; updatedAt: Date; } @@ -64,6 +71,13 @@ const userFileSchema = new mongoose.Schema( default: false, index: true, }, + lastAccessedAt: { type: Date, default: null }, + accessCount: { type: Number, default: 0 }, + extractionProgress: { + current: { type: Number, default: 0 }, + total: { type: Number, default: 0 }, + stage: { type: String, default: "" }, + }, }, { timestamps: true, diff --git a/backend/src/routes/query.ts b/backend/src/routes/query.ts index f812c19..608ef7f 100644 --- a/backend/src/routes/query.ts +++ b/backend/src/routes/query.ts @@ -1,9 +1,11 @@ import {Router} from "express"; -import queryHandler from "../handlers/queryHandler.js"; +import queryHandler, { logSearchClick } from "../handlers/queryHandler.js"; import { verifyToken } from "../middleware/auth.js"; const queryRouter = Router(); queryRouter.get('/', verifyToken, queryHandler); +// R10 β€” click logging +queryRouter.post('/click', verifyToken, logSearchClick); export default queryRouter; \ No newline at end of file diff --git a/backend/src/services/chunking.ts b/backend/src/services/chunking.ts index f2d462c..02f0a8e 100644 --- a/backend/src/services/chunking.ts +++ b/backend/src/services/chunking.ts @@ -16,7 +16,18 @@ * used by the retrieval boost for "totals / count / how many" queries. */ -const CHARS_PER_TOKEN = 4; +// Real tokenizer. Previously used CHARS_PER_TOKEN = 4 which is off by 2-4x +// for code or multilingual content (off in opposite directions). gpt-tokenizer +// uses the cl100k_base BPE β€” same family as Gemini's tokenizer for our purposes. +// Token-count diff between gpt-tokenizer and gemini-embedding-001 is <5% on +// English content; close enough for chunk-size budgeting. +import { encode as gptEncode } from "gpt-tokenizer"; + +// Used for char-budget conversion when we need approximate char counts in +// splitters (since slicing strings character-wise is O(1) vs token-encode +// being O(n)). Tuned to median English content; chunk boundaries are +// refined later by re-counting with the real tokenizer. +const APPROX_CHARS_PER_TOKEN = 4; export type ChildChunk = { index: number; @@ -25,6 +36,9 @@ export type ChildChunk = { parent_index: number; has_table: boolean; tokens: number; + /** C8 β€” Full heading path for this chunk's section, e.g. ["Intro", "Background"]. + * Used for citations in chat answers and as a future filter signal. */ + heading_path?: string[]; }; type Block = { @@ -32,9 +46,54 @@ type Block = { text: string; atomic: boolean; sticky: boolean; + /** C8 β€” Heading depth (1-6) if kind === "heading", else undefined. */ + level?: number; }; -const tokens = (s: string): number => Math.max(1, Math.ceil(s.length / CHARS_PER_TOKEN)); +// ---------- C4: Boilerplate detection ---------- +// Repeated lines that appear on every page (headers, footers, disclaimers, +// "Confidential" stamps, "Page N of M") pollute every chunk with the same +// terms β€” BM25 thinks they're content keywords. Detect lines that appear +// many times and strip them before chunking. + +const BOILERPLATE_MIN_OCCURRENCES = 3; +const BOILERPLATE_MIN_LENGTH = 4; +const BOILERPLATE_MAX_LENGTH = 120; + +const detectBoilerplate = (text: string): Set => { + const lineCounts = new Map(); + for (const raw of text.split("\n")) { + const line = raw.trim(); + if (line.length < BOILERPLATE_MIN_LENGTH || line.length > BOILERPLATE_MAX_LENGTH) continue; + lineCounts.set(line, (lineCounts.get(line) ?? 0) + 1); + } + const boilerplate = new Set(); + for (const [line, count] of lineCounts) { + if (count >= BOILERPLATE_MIN_OCCURRENCES) boilerplate.add(line); + } + return boilerplate; +}; + +const stripBoilerplate = (text: string, boilerplate: Set): string => { + if (boilerplate.size === 0) return text; + const cleaned: string[] = []; + for (const raw of text.split("\n")) { + if (!boilerplate.has(raw.trim())) cleaned.push(raw); + } + return cleaned.join("\n"); +}; + +// Accurate token count for budgeting and metadata. +const tokens = (s: string): number => { + if (!s) return 1; + try { + return Math.max(1, gptEncode(s).length); + } catch { + // gpt-tokenizer rarely throws (e.g. on certain control characters); + // fall back to char approximation so chunking never breaks. + return Math.max(1, Math.ceil(s.length / APPROX_CHARS_PER_TOKEN)); + } +}; // ---------- block parsing ---------- @@ -61,7 +120,14 @@ const parseBlocks = (text: string): Block[] => { const m = line.match(HEADING_RE); if (m) { - blocks.push({ kind: "heading", text: line.trim(), atomic: false, sticky: true }); + // C8 β€” capture heading level (# = 1, ## = 2, etc.) + blocks.push({ + kind: "heading", + text: line.trim(), + atomic: false, + sticky: true, + level: m[1].length, + }); i++; continue; } @@ -95,18 +161,69 @@ const parseBlocks = (text: string): Block[] => { return blocks; }; +// C3 β€” Better sentence segmentation. +// The old regex `/(?<=[.!?])\s+(?=[A-Z("'])/` broke on: +// β€’ Abbreviations: "Dr. Smith said hello" β†’ split between "Dr." and "Smith" +// β€’ Decimals: "$3.14 trillion" β†’ split between "3." and "14" +// β€’ Initials: "J. K. Rowling" β†’ split multiple times +// β€’ URLs: "see www.foo.com for" β†’ split +// +// We use a regex pass + a context-aware filter that re-joins false-positive +// splits. Not full NLP-grade (that'd need spaCy), but catches the common cases +// without adding a 50MB dependency. + +const COMMON_ABBREVIATIONS = new Set([ + "Dr", "Mr", "Mrs", "Ms", "Prof", "Sr", "Jr", "Hon", + "Inc", "Ltd", "Co", "Corp", "Llc", "Plc", + "vs", "etc", "ie", "eg", "Eg", "Ie", "Etc", "Vs", + "U.S", "U.K", "U.S.A", "E.U", "St", "Ave", "Blvd", + "Mt", "Ft", "No", "Vol", "Op", "Cit", "Ed", + "Jan", "Feb", "Mar", "Apr", "Jun", "Jul", "Aug", "Sep", "Sept", "Oct", "Nov", "Dec", +]); + +const looksLikeAbbreviation = (lastWord: string): boolean => { + // Strip trailing period + const w = lastWord.replace(/\.$/, ""); + if (COMMON_ABBREVIATIONS.has(w)) return true; + // Single capital letter (initials like "J.") + if (/^[A-Z]$/.test(w)) return true; + // Numbers that look like decimals (e.g. "$3.14" β†’ previous part "$3" stays) + if (/^\$?\d+$/.test(w)) return true; + return false; +}; + const splitBySentences = (text: string, maxChars: number): string[] => { if (text.length <= maxChars) return [text]; - const sentences = text.split(/(?<=[.!?])\s+(?=[A-Z(\"'])/); + + // First pass: candidate breaks after .!? followed by whitespace + capital/quote/paren. + const candidates = text.split(/(?<=[.!?])\s+(?=[A-Z("'])/); + + // Second pass: rejoin splits that came after a known abbreviation or single-letter initial. + const sentences: string[] = []; + for (const piece of candidates) { + if (sentences.length > 0) { + const last = sentences[sentences.length - 1]; + const lastWord = last.trimEnd().split(/\s+/).pop() ?? ""; + if (looksLikeAbbreviation(lastWord)) { + sentences[sentences.length - 1] = last + " " + piece; + continue; + } + } + sentences.push(piece); + } + + // Third pass: pack into max-size buckets. const out: string[] = []; let buf: string[] = []; let buflen = 0; for (const s of sentences) { if (buflen + s.length + 1 > maxChars && buf.length > 0) { out.push(buf.join(" ")); - buf = [s]; buflen = s.length; + buf = [s]; + buflen = s.length; } else { - buf.push(s); buflen += s.length + 1; + buf.push(s); + buflen += s.length + 1; } } if (buf.length > 0) out.push(buf.join(" ")); @@ -129,17 +246,42 @@ const splitAtomicByLines = (text: string, maxChars: number): string[] => { return out.map((g) => g.join("\n")).filter((s) => s.trim().length > 0); }; -type Section = { blocks: Block[]; has_table: boolean }; +type Section = { + blocks: Block[]; + has_table: boolean; + /** C8 β€” Full heading path leading INTO this section. + * e.g. ["Introduction", "Background"] for a level-3 section under those. */ + heading_path: string[]; +}; const groupIntoSections = (blocks: Block[]): Section[] => { const out: Section[] = []; - let cur: Section = { blocks: [], has_table: false }; + // C8 β€” track heading path as we walk blocks. headingStack[i] = current + // heading at level i+1. When we see a deeper heading, we extend; when + // we see a same-or-shallower heading, we pop back. + const headingStack: string[] = []; // index = level-1 + let cur: Section = { blocks: [], has_table: false, heading_path: [] }; + for (const b of blocks) { if (b.kind === "heading") { - const level = (b.text.match(/^(#+)/)?.[1].length ?? 1); + const level = b.level ?? 1; + // Cleanly trim heading text: "## My Heading" β†’ "My Heading" + const headingText = b.text.replace(/^#+\s*/, "").trim(); + // Truncate stack to level-1 (this heading replaces everything at and below its level) + headingStack.length = Math.max(0, level - 1); + headingStack.push(headingText); + + // Start a new section when we hit a top-level heading (h1/h2). if (level <= 2 && cur.blocks.length > 0) { out.push(cur); - cur = { blocks: [], has_table: false }; + cur = { + blocks: [], + has_table: false, + heading_path: [...headingStack], + }; + } else { + // Otherwise update the heading path of the current section. + cur.heading_path = [...headingStack]; } } cur.blocks.push(b); @@ -156,7 +298,7 @@ const packSectionIntoParents = (section: Section, parentTargetChars: number): Se if (sectionText(section).length <= parentTargetChars) return [section]; const out: Section[] = []; - let cur: Section = { blocks: [], has_table: false }; + let cur: Section = { blocks: [], has_table: false, heading_path: [...section.heading_path] }; let curLen = 0; for (const b of section.blocks) { const pieces: Block[] = b.atomic @@ -165,7 +307,7 @@ const packSectionIntoParents = (section: Section, parentTargetChars: number): Se for (const piece of pieces) { if (curLen + piece.text.length + 2 > parentTargetChars && cur.blocks.length > 0) { out.push(cur); - cur = { blocks: [], has_table: false }; + cur = { blocks: [], has_table: false, heading_path: [...section.heading_path] }; curLen = 0; } cur.blocks.push(piece); @@ -201,6 +343,9 @@ const splitParentIntoChildren = ( parent_index: parentIndex, has_table: curHasTable || parent.has_table, tokens: tokens(body), + // C8 β€” attach the heading_path so chat answers can cite the + // exact section ("Section: Introduction > Background"). + heading_path: parent.heading_path.length > 0 ? [...parent.heading_path] : undefined, }); const tail = body.length > overlapChars ? body.slice(-overlapChars) : body; cur = tail ? [tail] : []; @@ -208,7 +353,24 @@ const splitParentIntoChildren = ( curHasTable = false; }; + // C6 (light semantic chunking) β€” document structure is a natural semantic + // boundary. When we encounter a heading or a kind-transition (paragraphβ†’table, + // paragraphβ†’code) AND the current chunk is at least 35% full, flush first. + // This avoids fusing unrelated topics into one chunk without needing + // expensive embedding-based break detection. + const semanticFlushThreshold = Math.floor(childTargetChars * 0.35); + let prevBlockKind: Block["kind"] | null = null; + for (const block of parent.blocks) { + const isSemanticBoundary = + block.kind === "heading" || + (prevBlockKind !== null && prevBlockKind !== block.kind && + (block.kind === "table" || block.kind === "code" || + prevBlockKind === "table" || prevBlockKind === "code")); + if (isSemanticBoundary && curLen >= semanticFlushThreshold) { + flush(); + } + const pieces = block.atomic ? splitAtomicByLines(block.text, childTargetChars) : splitBySentences(block.text, childTargetChars); @@ -218,6 +380,7 @@ const splitParentIntoChildren = ( curLen += piece.length + 2; if (block.kind === "table") curHasTable = true; } + prevBlockKind = block.kind; } flush(); if (out.length === 0) { @@ -228,6 +391,7 @@ const splitParentIntoChildren = ( parent_index: parentIndex, has_table: parent.has_table, tokens: tokens(parentText), + heading_path: parent.heading_path.length > 0 ? [...parent.heading_path] : undefined, }); } return out; @@ -240,15 +404,37 @@ export type ChunkOptions = { }; export const chunkMarkdown = (text: string, opts: ChunkOptions = {}): ChildChunk[] => { - const clean = (text ?? "").trim(); - if (!clean) return []; - const parentTargetTokens = opts.parentTargetTokens ?? 1500; - const childTargetTokens = opts.childTargetTokens ?? 400; - const overlapTokens = opts.overlapTokens ?? 60; - - const parentTargetChars = parentTargetTokens * CHARS_PER_TOKEN; - const childTargetChars = childTargetTokens * CHARS_PER_TOKEN; - const overlapChars = overlapTokens * CHARS_PER_TOKEN; + const raw = (text ?? "").trim(); + if (!raw) return []; + + // C4 β€” strip boilerplate (repeated headers/footers/disclaimers) BEFORE + // chunking. Otherwise every chunk contains the same boilerplate text + // and BM25 thinks "Confidential" or "Page 1 of 50" are key terms. + const boilerplate = detectBoilerplate(raw); + const clean = boilerplate.size > 0 ? stripBoilerplate(raw, boilerplate) : raw; + + // C7 β€” adaptive chunk sizing by document type. + // Code-heavy docs benefit from larger chunks (functions are atomic units). + // Very short docs use smaller chunks for retrieval granularity. + // Caller-provided opts override the heuristic. + const docTokens = tokens(clean); + const isCodeHeavy = (clean.match(/```/g) || []).length >= 4 || /^\s{4,}/m.test(clean); + const defaultParent = isCodeHeavy ? 2000 : docTokens < 500 ? 800 : 1500; + const defaultChild = isCodeHeavy ? 600 : docTokens < 500 ? 250 : 400; + + const parentTargetTokens = opts.parentTargetTokens ?? defaultParent; + const childTargetTokens = opts.childTargetTokens ?? defaultChild; + // C2: bumped from 60 β†’ 120. Cross-chunk antecedent context (sentences + // that reference "this approach" or "the previous result") needs more + // headroom than 60 tokens to survive a chunk boundary. + const overlapTokens = opts.overlapTokens ?? 120; + + // We need char budgets for splitter functions (which slice strings by + // character). Approximation here is intentional β€” chunks get tokenized + // accurately at the end via `tokens()` for the stored token count. + const parentTargetChars = parentTargetTokens * APPROX_CHARS_PER_TOKEN; + const childTargetChars = childTargetTokens * APPROX_CHARS_PER_TOKEN; + const overlapChars = overlapTokens * APPROX_CHARS_PER_TOKEN; const blocks = parseBlocks(clean); const sections = groupIntoSections(blocks); diff --git a/backend/src/services/queryWeaviate.ts b/backend/src/services/queryWeaviate.ts index 1b5240a..be53c6a 100644 --- a/backend/src/services/queryWeaviate.ts +++ b/backend/src/services/queryWeaviate.ts @@ -2,7 +2,7 @@ import getWeaviateClient from "../db/weaviate_client.js"; import { Filters } from 'weaviate-client'; import logger from "../logger.js"; import generateQueryEmbedding from "../utils/getQueryEmbedding.js"; -import { runSearchPipeline } from "./searchPipeline.js"; +import UserFile from "../models/userFileModel.js"; type SmartDriveSchema = { user_id: string; @@ -24,10 +24,245 @@ const chunkCollections: Record = { Media: "SmartDriveMediaChunks", }; +// ---------- RRF fusion ---------- +// Reciprocal Rank Fusion. Each signal contributes weight Β· (1 / (k + rank)) +// per file. Multi-signal matches naturally score higher than single-signal +// wins. Calibration-free: doesn't matter that BM25 scores 0-30 while hybrid +// scores 0-1 β€” we fuse on RANK, not raw score. +const RRF_K = 60; + +// Per-signal weights. Higher = signal is more trusted. +const SIGNAL_WEIGHTS: Record = { + "filename:full": 3.0, // explicit filename paste β†’ strongest signal + "filename:tokenized": 2.0, // descriptive query β†’ filename match + "chunk:bm25": 2.0, // body lexical match (chunks) + "summary:bm25": 2.0, // summary lexical match + "chunk:hybrid": 1.5, // semantic chunk match (paraphrase) + "raw_text:bm25": 1.5, // body lexical fallback for unchatted files +}; + +type SignalEntry = { rank: number; properties: Record; matched_chunk?: string }; +type FileSignalMap = Map>; // file_id β†’ signal_name β†’ entry + +// ---------- Recency boost ---------- +// Tiny exponential decay so recent files break ties for ambiguous queries +// without overriding genuine relevance. Capped at 8% of base RRF score. +const recencyBoost = (createdAt: unknown, baseScore: number): number => { + if (!createdAt || (typeof createdAt !== "string" && !(createdAt instanceof Date))) return 0; + const t = createdAt instanceof Date ? createdAt.getTime() : new Date(createdAt).getTime(); + if (!Number.isFinite(t)) return 0; + const ageDays = Math.max(0, (Date.now() - t) / (1000 * 60 * 60 * 24)); + // 30-day half-life. After 30 days the boost halves; after 90 days it's negligible. + return baseScore * 0.08 * Math.exp(-ageDays / 30); +}; + +// ---------- R6: Personalization boost ---------- +// Files the user actually interacts with (chat, view, download) rank higher +// for ambiguous queries. Two signals combined: +// - access count: log-scaled so frequent files boost mildly, not unboundedly +// - last-access recency: exp decay with 14-day half-life +// Capped at 12% of base score. Stronger than the upload-recency boost +// because explicit interaction is a stronger relevance signal. +const personalizationBoost = ( + accessCount: unknown, + lastAccessedAt: unknown, + baseScore: number, +): number => { + const count = typeof accessCount === "number" ? accessCount : 0; + if (count <= 0) return 0; + let recencyFactor = 0.5; // default if no lastAccessedAt + if (lastAccessedAt && (typeof lastAccessedAt === "string" || lastAccessedAt instanceof Date)) { + const t = lastAccessedAt instanceof Date + ? lastAccessedAt.getTime() + : new Date(lastAccessedAt as string).getTime(); + if (Number.isFinite(t)) { + const ageDays = Math.max(0, (Date.now() - t) / (1000 * 60 * 60 * 24)); + recencyFactor = Math.exp(-ageDays / 14); + } + } + // log(1+count) maps: 1 β†’ 0.69, 5 β†’ 1.79, 20 β†’ 3.04, 100 β†’ 4.62 + // Multiplied by recencyFactor (0-1), then scaled to max 0.12 of baseScore. + const intensity = Math.log(1 + count) * recencyFactor; + return baseScore * 0.12 * Math.min(1, intensity / 3); +}; + +// ---------- Result LRU cache ---------- +// Repeat queries within a session (typing/refinement, common terms) skip +// the entire pipeline. 5-min TTL means the cache stays warm without serving +// stale data when files change frequently. +const RESULT_CACHE_TTL_MS = 5 * 60 * 1000; +const RESULT_CACHE_MAX = 200; +type CacheEntry = { results: Record[]; expiresAt: number }; +const resultCache = new Map(); + +const resultCacheKey = (userId: string, query: string, collection: string): string => { + return `${userId}|${collection}|${query.trim().toLowerCase()}`; +}; + +const resultCacheGet = (key: string): Record[] | null => { + const entry = resultCache.get(key); + if (!entry) return null; + if (Date.now() > entry.expiresAt) { + resultCache.delete(key); + return null; + } + // LRU touch + resultCache.delete(key); + resultCache.set(key, entry); + return entry.results; +}; + +const resultCacheSet = (key: string, results: Record[]): void => { + resultCache.set(key, { results, expiresAt: Date.now() + RESULT_CACHE_TTL_MS }); + if (resultCache.size > RESULT_CACHE_MAX) { + const oldest = resultCache.keys().next().value; + if (oldest !== undefined) resultCache.delete(oldest); + } +}; + +// Exported so chat / delete-file paths can invalidate when content changes. +// Cheap to call: O(n) scan but n is bounded at RESULT_CACHE_MAX. +export const invalidateUserSearchCache = (userId: string): void => { + const prefix = `${userId}|`; + for (const k of resultCache.keys()) { + if (k.startsWith(prefix)) resultCache.delete(k); + } +}; + +// ---------- R7: MMR diversity ---------- +// Maximal Marginal Relevance β€” pick the next result by trading off "high +// relevance" against "low similarity to already-picked results". Prevents +// 5 near-duplicate file versions from dominating top 10. Uses filename +// token overlap as a cheap similarity proxy (avoiding an extra embed call). +const filenameTokens = (filename: string): Set => { + return new Set( + (filename ?? "") + .toLowerCase() + .replace(/([a-z])([A-Z])/g, "$1 $2") + .split(/[\s_\-.]+/) + .filter((t) => t.length > 1), + ); +}; + +const tokenSetSimilarity = (a: Set, b: Set): number => { + if (a.size === 0 || b.size === 0) return 0; + let inter = 0; + for (const t of a) if (b.has(t)) inter++; + return inter / Math.min(a.size, b.size); // 0..1 +}; + +const applyMMR = ( + candidates: Record[], + lambda: number = 0.7, +): Record[] => { + if (candidates.length <= 2) return candidates; + const picked: Record[] = []; + const remaining = [...candidates]; + // Always pick the top-scored first. + picked.push(remaining.shift()!); + while (remaining.length > 0) { + let bestIdx = 0; + let bestScore = -Infinity; + const pickedTokens = picked.map((p) => filenameTokens(String(p.filename ?? ""))); + for (let i = 0; i < remaining.length; i++) { + const cand = remaining[i]; + const candTokens = filenameTokens(String(cand.filename ?? "")); + // max similarity to any already-picked + let maxSim = 0; + for (const pt of pickedTokens) { + const sim = tokenSetSimilarity(candTokens, pt); + if (sim > maxSim) maxSim = sim; + } + const mmr = lambda * (cand.score as number) - (1 - lambda) * maxSim; + if (mmr > bestScore) { + bestScore = mmr; + bestIdx = i; + } + } + picked.push(remaining.splice(bestIdx, 1)[0]); + } + return picked; +}; + +// Common English stopwords + search-intent verbs/pronouns. Dropped from +// queries before tokenized filename matching so the user's wrapper words +// ("I am searching for…") don't dilute meaningful tokens. +const STOPWORDS = new Set([ + "i", "me", "my", "mine", "we", "our", "ours", + "a", "an", "the", "is", "are", "was", "were", "be", "been", "being", + "have", "has", "had", "do", "does", "did", + "for", "of", "to", "from", "with", "in", "on", "at", "by", "about", + "and", "or", "but", "if", "as", "this", "that", "these", "those", + "find", "show", "get", "search", "looking", "want", "need", "remember", + "stored", "saved", "uploaded", "file", "files", "where", "which", +]); + +/** Extract meaningful tokens from a conversational query. + * Keeps tokens of length β‰₯ 3 that aren't stopwords. Lowercased. */ +const extractKeywords = (query: string): string[] => { + const tokens = (query.toLowerCase().match(/\b\w+\b/g) ?? []) + .filter((t) => t.length >= 3 && !STOPWORDS.has(t)); + // Dedupe while preserving order. + const seen = new Set(); + const out: string[] = []; + for (const t of tokens) { + if (!seen.has(t)) { + seen.add(t); + out.push(t); + } + } + return out; +}; + +/** + * Search pipeline β€” four signals, max-fusion. + * + * (1) Chunk hybrid (BM25 + dense, score > 0.4) β€” for files that have + * chunks built (chatted with). Catches semantic + lexical body matches. + * (2) Filename substring (`like *q*`) β€” guarantees exact-filename hits even + * when BM25 tokenization would split the query. Fixed score 0.95. + * (3) Summary hybrid restricted to `summary` property ONLY β€” catches matches + * against the LLM-generated summary text (e.g. "IIT Guwahati" mentioned + * in a resume's summary). Strict 0.5 threshold so weak semantic matches + * don't pad results. Critical: queryProperties restriction prevents + * Weaviate from also searching raw_text (which caused "same results for + * every query" because short docs gamed BM25 length norm). + * (4) raw_text pure BM25 β€” strict lexical body fallback. Returns nothing + * for queries with no actual term match. Catches body content for + * files without chunks. No-ops if the property isn't index_searchable. + * + * Fuse via max(score). Drop nameless results, top 30. + */ const queryWeaviate = async (userId: string, userQuery: string, queryCollection: string) => { try { + // Check result cache first β€” same query within 5 minutes is free. + const cacheKey = resultCacheKey(userId, userQuery, queryCollection); + const cached = resultCacheGet(cacheKey); + if (cached) { + logger.info(`search: cache HIT for "${userQuery}" (${cached.length} results)`); + return cached; + } + + const client = await getWeaviateClient(); + if (!client) return []; + + // CRITICAL: build a cleaned BM25 query that strips conversational + // stopwords. Raw user queries like "I am searching for a person who + // studies at IIT" have 11 tokens, only 3 of which carry meaning. BM25 + // is OR-by-token: noise tokens pull in noise files. With cleaned + // ["person", "studies", "iit"], the rare token "iit" (high IDF) + // dominates and the resume wins reliably. + // + // Dense embedding still uses the original query β€” sentence embeddings + // handle stopwords gracefully and conversational context can help. + const keywords = extractKeywords(userQuery); + const bm25Query = keywords.length > 0 ? keywords.join(" ") : userQuery; + logger.info(`search: raw="${userQuery}" bm25="${bm25Query}" tokens=${keywords.length}`); + + const queryVector = await generateQueryEmbedding(userQuery); + let pairs: { summary: string; chunks: string }[] = []; - if (queryCollection === "SmartDrive") { + if (queryCollection === "SmartDrive" || queryCollection === "all") { pairs = Object.keys(summaryCollections).map((k) => ({ summary: summaryCollections[k], chunks: chunkCollections[k], @@ -42,7 +277,303 @@ const queryWeaviate = async (userId: string, userQuery: string, queryCollection: return []; } - return await runSearchPipeline(userId, userQuery, pairs); + // RRF data structure: file_id β†’ signal_name β†’ {rank, properties, matched_chunk} + // Multiple signals per file lets RRF reward multi-signal matches. + const fileSignals: FileSignalMap = new Map(); + + const recordSignal = ( + signal: string, + file_id: string, + rank: number, + properties: Record, + matched_chunk?: string, + ) => { + if (!file_id) return; + let sigs = fileSignals.get(file_id); + if (!sigs) { + sigs = new Map(); + fileSignals.set(file_id, sigs); + } + // If the same signal already recorded a hit (from a different + // collection), keep the better-ranked one. + const existing = sigs.get(signal); + if (!existing || rank < existing.rank) { + sigs.set(signal, { rank, properties, matched_chunk }); + } + }; + + for (const { summary: summaryName, chunks: chunksName } of pairs) { + // (1a) Chunk pure BM25 β€” only docs that actually contain the term. + // No dense, no noise. For rare-term queries this is what + // reliably ranks the right file at the top. + if (await client.collections.exists(chunksName)) { + try { + const chunkCol = client.collections.get(chunksName); + const bm25Hits = await chunkCol.query.bm25(bm25Query, { + queryProperties: ["chunk_text"], + limit: 20, + filters: chunkCol.filter.byProperty("user_id").equal(userId), + returnMetadata: ['score'], + }); + // Dedupe chunks β†’ files: keep the highest-ranking chunk per file. + const seenFiles = new Set(); + let rank = 0; + for (const obj of bm25Hits.objects) { + const raw = obj.metadata?.score ?? 0; + if (raw <= 0) continue; + const fid = String((obj.properties as { file_id?: string }).file_id ?? ""); + if (!fid || seenFiles.has(fid)) continue; + seenFiles.add(fid); + rank++; + recordSignal( + "chunk:bm25", + fid, + rank, + obj.properties as Record, + String((obj.properties as { chunk_text?: string }).chunk_text ?? "") || undefined, + ); + } + } catch (err) { + logger.warn(`chunk bm25 on ${chunksName} failed: ${err}`); + } + } + + // (1b) Chunk hybrid β€” adds semantic recall for paraphrase queries. + // Strict 0.55 threshold so weak dense matches don't pad. + if (await client.collections.exists(chunksName)) { + try { + const chunkCol = client.collections.get(chunksName); + const chunkHits = await chunkCol.query.hybrid(userQuery, { + vector: queryVector, + alpha: 0.5, + limit: 20, + filters: chunkCol.filter.byProperty("user_id").equal(userId), + returnMetadata: ['score'], + }); + const seenFiles = new Set(); + let rank = 0; + for (const obj of chunkHits.objects) { + const score = obj.metadata?.score ?? 0; + if (score <= 0.55) continue; + const fid = String((obj.properties as { file_id?: string }).file_id ?? ""); + if (!fid || seenFiles.has(fid)) continue; + seenFiles.add(fid); + rank++; + recordSignal( + "chunk:hybrid", + fid, + rank, + obj.properties as Record, + String((obj.properties as { chunk_text?: string }).chunk_text ?? "") || undefined, + ); + } + } catch (err) { + logger.warn(`chunk hybrid on ${chunksName} failed: ${err}`); + } + } + + if (await client.collections.exists(summaryName)) { + const sumCol = client.collections.get(summaryName); + + // (2a) Filename FULL-string substring β€” for explicit-filename queries. + try { + const filenameHits = await sumCol.query.fetchObjects({ + limit: 10, + filters: Filters.and( + sumCol.filter.byProperty("user_id").equal(userId), + sumCol.filter.byProperty("filename").like(`*${userQuery}*`), + ), + }); + let rank = 0; + for (const obj of filenameHits.objects) { + const fid = String((obj.properties as { file_id?: string }).file_id ?? ""); + if (!fid) continue; + rank++; + recordSignal("filename:full", fid, rank, obj.properties as Record); + } + } catch (err) { + logger.warn(`filename match on ${summaryName} failed: ${err}`); + } + + // (2b) Filename TOKENIZED match β€” for descriptive queries like + // "I'm looking for my undergrad transcripts" β†’ tokens + // [undergrad, transcripts] β†’ OR-match each against filename. + // Catches Undergrad_Transcripts.pdf even when the file has + // no summary or chunks. Score reflects # of tokens matched. + if (keywords.length > 0) { + try { + const orParts = keywords.map((kw) => + sumCol.filter.byProperty("filename").like(`*${kw}*`), + ); + const tokenHits = await sumCol.query.fetchObjects({ + limit: 20, + filters: Filters.and( + sumCol.filter.byProperty("user_id").equal(userId), + orParts.length === 1 ? orParts[0] : Filters.or(...orParts), + ), + }); + // Sort by # of matched keywords so the most-relevant filename ranks first. + const scored: { obj: typeof tokenHits.objects[number]; matched: number }[] = []; + for (const obj of tokenHits.objects) { + const props = obj.properties as { file_id?: string; filename?: string }; + const fid = String(props.file_id ?? ""); + if (!fid) continue; + const filenameLower = String(props.filename ?? "").toLowerCase(); + const matchedCount = keywords.filter((kw) => filenameLower.includes(kw)).length; + if (matchedCount === 0) continue; + scored.push({ obj, matched: matchedCount }); + } + scored.sort((a, b) => b.matched - a.matched); + let rank = 0; + for (const { obj } of scored) { + const fid = String((obj.properties as { file_id?: string }).file_id ?? ""); + rank++; + recordSignal("filename:tokenized", fid, rank, obj.properties as Record); + } + } catch (err) { + logger.warn(`filename token match on ${summaryName} failed: ${err}`); + } + } + + // (3) Summary PURE BM25 β€” only docs that actually contain the + // term in the summary. NOT hybrid: hybrid includes dense-only + // matches and Weaviate's score normalization can rank a + // dense-noise match higher than a real BM25 win. For a query + // like "Indraneel" (rare proper noun, only in the resume's + // summary), pure BM25 returns ONLY the resume β€” no padding. + try { + const summaryHits = await sumCol.query.bm25(bm25Query, { + queryProperties: ["summary"], + limit: 20, + filters: sumCol.filter.byProperty("user_id").equal(userId), + returnMetadata: ['score'], + }); + let rank = 0; + for (const obj of summaryHits.objects) { + const raw = obj.metadata?.score ?? 0; + if (raw <= 0) continue; + const fid = String((obj.properties as { file_id?: string }).file_id ?? ""); + if (!fid) continue; + rank++; + recordSignal("summary:bm25", fid, rank, obj.properties as Record); + } + } catch (err) { + logger.warn(`summary bm25 on ${summaryName} failed: ${err}`); + } + + // (4) raw_text pure BM25 β€” strict lexical match in document body. + // This is the fallback for files without chunks (most files, + // since chunks are lazy). No dense vector β€” purely + // "does the doc actually contain these words". The signal + // returns nothing for queries that have no term match. + // Will silently no-op if raw_text isn't index_searchable. + try { + // R3: prefer the new `body_text` property (index_searchable=true) + // over `raw_text` (storage-only). New ingests populate both; + // older rows still have raw_text only. + const rawHits = await sumCol.query.bm25(bm25Query, { + queryProperties: ["body_text", "raw_text"], + limit: 20, + filters: sumCol.filter.byProperty("user_id").equal(userId), + returnMetadata: ['score'], + }); + let rank = 0; + for (const obj of rawHits.objects) { + const raw = obj.metadata?.score ?? 0; + if (raw <= 0) continue; + const fid = String((obj.properties as { file_id?: string }).file_id ?? ""); + if (!fid) continue; + rank++; + recordSignal("raw_text:bm25", fid, rank, obj.properties as Record); + } + } catch (err) { + logger.warn(`raw_text bm25 on ${summaryName} failed (property may not be searchable): ${err}`); + } + } + } + + if (fileSignals.size === 0) { + logger.info(`search: 0 results for "${userQuery}"`); + resultCacheSet(cacheKey, []); + return []; + } + + // ---------- RRF fusion ---------- + // For each file, sum weighted 1/(k+rank) across every signal it appears in. + // This rewards files matched by multiple signals β€” exactly the "the + // file we're looking for shows up in chunks AND filename AND summary" + // case that max-fusion missed. + // R6 β€” fetch personalization data (accessCount, lastAccessedAt) for + // candidate files in one Mongo round-trip. Skipped if no signals fired. + const candidateIds = [...fileSignals.keys()]; + const personalizationMap = new Map(); + if (candidateIds.length > 0) { + try { + const personalizationDocs = await UserFile.find( + { _id: { $in: candidateIds }, userId }, + { _id: 1, accessCount: 1, lastAccessedAt: 1 }, + ).lean(); + for (const doc of personalizationDocs) { + personalizationMap.set(doc._id.toString(), { + accessCount: (doc as { accessCount?: number }).accessCount ?? 0, + lastAccessedAt: (doc as { lastAccessedAt?: Date }).lastAccessedAt ?? null, + }); + } + } catch (err) { + logger.warn(`personalization lookup failed (continuing without boost): ${err}`); + } + } + + const merged: Record[] = []; + for (const [fid, sigs] of fileSignals.entries()) { + let score = 0; + let bestProps: Record = {}; + let bestChunk: string | undefined; + const matchedIn: string[] = []; + for (const [signalName, entry] of sigs.entries()) { + const weight = SIGNAL_WEIGHTS[signalName] ?? 1.0; + score += weight * (1 / (RRF_K + entry.rank)); + matchedIn.push(signalName); + if (signalName.startsWith("filename:") || signalName.startsWith("summary:") || signalName.startsWith("raw_text:")) { + bestProps = entry.properties; + } else if (Object.keys(bestProps).length === 0) { + bestProps = entry.properties; + } + if (!bestChunk && entry.matched_chunk) bestChunk = entry.matched_chunk; + } + + // Recency boost (upload age) β€” capped 8%. + const recBoost = recencyBoost(bestProps.created_at, score); + // R6 β€” personalization boost (user interaction history) β€” capped 12%. + const personalization = personalizationMap.get(fid); + const personBoost = personalization + ? personalizationBoost(personalization.accessCount, personalization.lastAccessedAt, score) + : 0; + + const row: Record = { + ...bestProps, + file_id: fid, + score: score + recBoost + personBoost, + matched_chunk: bestChunk, + matched_in: matchedIn, + }; + if (typeof row.filename === "string" && (row.filename as string).length > 0) { + merged.push(row); + } + } + merged.sort((a, b) => (b.score as number) - (a.score as number)); + // R7 β€” MMR diversification on the top 30. Filename-token similarity + // is the cheap proxy. lambda=0.7 keeps the relevance bias high while + // breaking up near-duplicates (file v1, v2, v3 etc). + const rerankPool = merged.slice(0, 30); + const diversified = applyMMR(rerankPool, 0.7); + + const topScore = diversified.length > 0 ? (diversified[0].score as number).toFixed(4) : "n/a"; + const topFile = diversified.length > 0 ? String(diversified[0].filename ?? "?") : "n/a"; + logger.info(`search: ${diversified.length} results for "${userQuery}" (top=${topFile} @ ${topScore})`); + + resultCacheSet(cacheKey, diversified); + return diversified; } catch (error) { logger.error('queryWeaviate failed:', error); return { status: 500, error: 'Search failed due to an internal error.' }; @@ -124,6 +655,10 @@ const deleteWeaviateFile = async ( await deleteFrom(summaryCollection); if (chunkCollectionName) await deleteFrom(chunkCollectionName); + // Invalidate cached search results so deleted files don't keep + // showing up in cached query responses for the next 5 minutes. + invalidateUserSearchCache(userId); + return true; } catch (error) { logger.error(`deleteWeaviateFile fileId=${fileId} failed:`, error); diff --git a/backend/src/services/searchPipeline.ts b/backend/src/services/searchPipeline.ts deleted file mode 100644 index 1cea4cf..0000000 --- a/backend/src/services/searchPipeline.ts +++ /dev/null @@ -1,567 +0,0 @@ -/** - * Search pipeline β€” multi-signal retrieval + RRF fusion + optional rerank. - * - * Architecture: - * 1. analyzeQuery β€” detects filename vs content vs recency intent, extracts phrases - * 2. Parallel signals: - * (a) chunk hybrid (BM25 + vector) β†’ covers semantic + body text - * (b) filename BM25 over tokenized name β†’ covers "find file named …" - * (c) summary hybrid β†’ covers paraphrased queries against the LLM summary - * (d) entity / date / doc_id structured β†’ high-precision facts - * 3. RRF fusion with intent-based per-signal weights (k=60) - * 4. Recency boost (small) on top of fused rank - * 5. Optional cross-encoder / LLM rerank for top-30 (off by default; latency cost) - */ - -import { Filters } from "weaviate-client"; -import getWeaviateClient from "../db/weaviate_client.js"; -import generateQueryEmbedding from "../utils/getQueryEmbedding.js"; -import { extractQueryEntities } from "./queryEntityExtractor.js"; -import logger from "../logger.js"; - -// ---------- types ---------- - -type Intent = { - filename: boolean; - recency: boolean; - content: boolean; - exactPhrases: string[]; -}; - -export type QueryAnalysis = { - raw: string; - cleaned: string; - intent: Intent; -}; - -type SignalHit = { - file_id: string; - rank: number; // 1-based - matched_chunk?: string; -}; - -type FusedHit = { - file_id: string; - score: number; - matched_in: string[]; - matched_chunk?: string; - matched_entities?: string[]; - matched_dates?: string[]; - matched_doc_ids?: string[]; -}; - -export type SearchResult = FusedHit & { - [key: string]: unknown; -}; - -// ---------- 1. Query analyzer ---------- - -const STOPWORDS = new Set([ - "find", "show", "me", "get", "search", "the", "a", "an", "please", - "for", "of", "to", "with", "about", "i", "we", "want", "need", -]); - -const FILENAME_HINTS = /\.\w{2,5}(\s|$)|[_\-]|[a-z][A-Z]|[A-Z]{2,}/; - -const RECENCY_HINTS = /\b(recent|latest|today|yesterday|this\s+week|last\s+week|new(est)?|just\s+uploaded)\b/i; - -export const analyzeQuery = (raw: string): QueryAnalysis => { - const trimmed = raw.trim(); - - // 1. Pull quoted phrases out β€” these must match exactly in some signal. - const exactPhrases: string[] = []; - const phraseRe = /"([^"]+)"/g; - let m: RegExpExecArray | null; - while ((m = phraseRe.exec(trimmed))) exactPhrases.push(m[1].trim()); - const unquoted = trimmed.replace(/"[^"]+"/g, " ").replace(/\s+/g, " ").trim(); - - // 2. Strip conversational stopwords for retrieval, but keep raw form for filename/exact match. - const cleaned = unquoted - .split(/\s+/) - .filter((w) => w.length > 0 && !STOPWORDS.has(w.toLowerCase())) - .join(" ") - .trim() || unquoted; - - // 3. Intent classification. - const filename = FILENAME_HINTS.test(unquoted); - const recency = RECENCY_HINTS.test(unquoted); - const content = !filename || cleaned.split(/\s+/).length > 1; - - return { - raw: trimmed, - cleaned, - intent: { filename, recency, content, exactPhrases }, - }; -}; - -// ---------- 2. Tokenize filenames for BM25 ---------- - -/** Tokenize `DS_Cheat-Sheet_v2.pdf` β†’ ['ds','cheat','sheet','v2','pdf']. */ -export const tokenizeFilename = (name: string): string[] => { - if (!name) return []; - return name - .replace(/([a-z])([A-Z])/g, "$1 $2") // camelCase β†’ camel Case - .split(/[\s_\-.]+/) - .map((t) => t.toLowerCase()) - .filter((t) => t.length > 1); -}; - -// ---------- 3. RRF fusion ---------- - -const RRF_K = 60; - -/** - * Reciprocal Rank Fusion. Each signal contributes weightΒ·(1 / (k + rank)) per file. - * Per-signal weights let us bias the fusion by detected intent - * (e.g. boost filename signal 2x when the query looks like a filename). - */ -export const rrfFuse = ( - signals: { name: string; hits: SignalHit[]; weight: number }[], - k: number = RRF_K, -): Map; bestChunk?: string }> => { - const acc = new Map; bestChunk?: string }>(); - for (const { name, hits, weight } of signals) { - for (const hit of hits) { - const cur = acc.get(hit.file_id) ?? { - score: 0, - matched_in: new Set(), - bestChunk: undefined, - }; - cur.score += weight * (1 / (k + hit.rank)); - cur.matched_in.add(name); - if (!cur.bestChunk && hit.matched_chunk) cur.bestChunk = hit.matched_chunk; - acc.set(hit.file_id, cur); - } - } - return acc; -}; - -// ---------- 4. Recency boost ---------- - -/** Small bump (max ~5% of RRF top score) for files uploaded recently. Never dominates relevance. */ -const recencyBoost = (createdAt: unknown, baseScore: number): number => { - if (!createdAt || typeof createdAt !== "string") return 0; - const t = new Date(createdAt).getTime(); - if (!Number.isFinite(t)) return 0; - const ageDays = Math.max(0, (Date.now() - t) / (1000 * 60 * 60 * 24)); - // Half-life of 30 days. Boost = baseScore Β· 0.05 Β· exp(-age/30) - return baseScore * 0.05 * Math.exp(-ageDays / 30); -}; - -// ---------- 5. Parallel signal collection ---------- - -type WClient = Awaited>; - -/** Pick alpha based on query length. Short queries (1-3 tokens) are usually - * keyword-driven β€” bias toward BM25. Longer queries are natural language β€” - * let dense pull more weight. */ -const pickAlpha = (query: string): number => { - const tokens = query.trim().split(/\s+/).filter(Boolean); - if (tokens.length <= 3) return 0.25; // 75% BM25, 25% dense - if (tokens.length <= 6) return 0.5; // balanced - return 0.7; // 70% dense, 30% BM25 -}; - -const chunkHybridSignal = async ( - client: NonNullable, - chunksName: string, - userId: string, - query: string, - vector: number[], -): Promise => { - if (!(await client.collections.exists(chunksName))) return []; - const col = client.collections.get(chunksName); - const res = await col.query.hybrid(query, { - vector, - alpha: pickAlpha(query), - limit: 50, - filters: col.filter.byProperty("user_id").equal(userId), - returnMetadata: ["score"], - }); - // Collapse to one rank per file β€” keep the best-ranked chunk per file. - const seen = new Map(); - let rank = 1; - for (const obj of res.objects) { - const p = obj.properties as { file_id?: string; chunk_text?: string }; - const fid = String(p.file_id ?? ""); - if (!fid || seen.has(fid)) { rank++; continue; } - seen.set(fid, { file_id: fid, rank: rank++, matched_chunk: p.chunk_text }); - } - return [...seen.values()]; -}; - -/** Pure BM25 over chunk text. Critical for rare-term queries like proper nouns - * ("Guwahati", "Acme Corp") β€” hybrid's dense half dilutes exact matches, - * but pure BM25 ranks the unique-term file at #1 reliably. RRF fuses this - * with the hybrid signal so common-word queries still benefit from semantic. */ -const chunkBM25Signal = async ( - client: NonNullable, - chunksName: string, - userId: string, - query: string, -): Promise => { - if (!(await client.collections.exists(chunksName))) return []; - const col = client.collections.get(chunksName); - try { - const res = await col.query.bm25(query, { - queryProperties: ["chunk_text"], - limit: 50, - filters: col.filter.byProperty("user_id").equal(userId), - returnMetadata: ["score"], - }); - const seen = new Map(); - let rank = 1; - for (const obj of res.objects) { - const p = obj.properties as { file_id?: string; chunk_text?: string }; - const fid = String(p.file_id ?? ""); - if (!fid || seen.has(fid)) { rank++; continue; } - seen.set(fid, { file_id: fid, rank: rank++, matched_chunk: p.chunk_text }); - } - return [...seen.values()]; - } catch (e) { - logger.warn(`chunkBM25Signal on ${chunksName} failed: ${e}`); - return []; - } -}; - -const filenameBM25Signal = async ( - client: NonNullable, - summaryName: string, - userId: string, - query: string, -): Promise => { - if (!(await client.collections.exists(summaryName))) return []; - const col = client.collections.get(summaryName); - try { - const res = await col.query.bm25(query, { - queryProperties: ["filename"], - limit: 50, - filters: col.filter.byProperty("user_id").equal(userId), - returnMetadata: ["score"], - }); - let rank = 1; - return res.objects.map((obj) => { - const p = obj.properties as { file_id?: string }; - return { file_id: String(p.file_id ?? ""), rank: rank++ }; - }).filter((h) => h.file_id); - } catch (e) { - logger.warn(`filenameBM25Signal on ${summaryName} failed (collection may lack BM25 index on filename): ${e}`); - return []; - } -}; - -const summaryHybridSignal = async ( - client: NonNullable, - summaryName: string, - userId: string, - query: string, - vector: number[], -): Promise => { - if (!(await client.collections.exists(summaryName))) return []; - const col = client.collections.get(summaryName); - try { - const res = await col.query.hybrid(query, { - vector, - alpha: 0.5, - limit: 50, - filters: col.filter.byProperty("user_id").equal(userId), - returnMetadata: ["score"], - }); - let rank = 1; - return res.objects.map((obj) => { - const p = obj.properties as { file_id?: string }; - return { file_id: String(p.file_id ?? ""), rank: rank++ }; - }).filter((h) => h.file_id); - } catch { - // Older summary collections may not have a summary embedding β€” that's fine. - return []; - } -}; - -const entitySignal = async ( - client: NonNullable, - summaryName: string, - userId: string, - qe: { entities: string[]; dates: string[]; doc_ids: string[]; topics: string[] }, -): Promise<{ - hits: SignalHit[]; - matches: Map; -}> => { - const matches = new Map(); - if (!(await client.collections.exists(summaryName))) return { hits: [], matches }; - const col = client.collections.get(summaryName); - const orParts = []; - if (qe.entities.length) orParts.push(col.filter.byProperty("entities").containsAny(qe.entities)); - if (qe.dates.length) orParts.push(col.filter.byProperty("dates").containsAny(qe.dates)); - if (qe.doc_ids.length) orParts.push(col.filter.byProperty("doc_ids").containsAny(qe.doc_ids)); - if (qe.topics.length) orParts.push(col.filter.byProperty("topics").containsAny(qe.topics)); - if (orParts.length === 0) return { hits: [], matches }; - try { - const res = await col.query.fetchObjects({ - limit: 50, - filters: Filters.and( - col.filter.byProperty("user_id").equal(userId), - Filters.or(...orParts), - ), - }); - // Rank by # of matched signals (more matches = higher rank). - const scored = res.objects - .map((obj) => { - const p = obj.properties as { - file_id?: string; - entities?: string[]; - dates?: string[]; - doc_ids?: string[]; - topics?: string[]; - }; - const fid = String(p.file_id ?? ""); - if (!fid) return null; - const me = qe.entities.filter((e) => (p.entities ?? []).includes(e)); - const md = qe.dates.filter((d) => (p.dates ?? []).includes(d)); - const mi = qe.doc_ids.filter((d) => (p.doc_ids ?? []).includes(d)); - const mt = qe.topics.filter((t) => (p.topics ?? []).includes(t)); - return { fid, total: me.length + md.length + mi.length + mt.length, me, md, mi }; - }) - .filter((x): x is NonNullable => x !== null) - .sort((a, b) => b.total - a.total); - const hits: SignalHit[] = scored.map((s, i) => ({ file_id: s.fid, rank: i + 1 })); - for (const s of scored) { - matches.set(s.fid, { entities: s.me, dates: s.md, doc_ids: s.mi }); - } - return { hits, matches }; - } catch (e) { - // Older collections may not have the array fields β€” degrade gracefully. - logger.warn(`entitySignal on ${summaryName} failed: ${e}`); - return { hits: [], matches }; - } -}; - -// ---------- 6. Orchestrator ---------- - -const intentWeights = (intent: Intent): { chunk: number; chunkBM25: number; filename: number; summary: number; entity: number } => { - // Defaults are balanced. Filename intent boosts filename + summary (where the name lives). - // Pure content intent boosts chunk/body. - // chunkBM25 is the pure-lexical companion to chunk hybrid β€” weight it strongly - // because exact-term matches (proper nouns, IDs, rare words) are what users - // most often expect to win. - if (intent.filename) return { chunk: 1.0, chunkBM25: 1.5, filename: 2.5, summary: 1.5, entity: 1.5 }; - if (intent.exactPhrases.length > 0) return { chunk: 1.5, chunkBM25: 2.5, filename: 1.0, summary: 1.5, entity: 2.0 }; - return { chunk: 1.5, chunkBM25: 2.0, filename: 1.0, summary: 1.0, entity: 1.5 }; -}; - -/** Exact-substring match post-RRF bonus. If the query token literally appears - * in the file's filename, summary, or matched chunk, give a meaningful score - * bump. Catches the "one file uniquely contains this term β€” it should be #1" - * failure mode that RRF alone misses. - * - * Scaled to ~0.05 per match, which dominates the typical RRF score - * (range 0.01-0.10) without overwhelming intentional ranking. */ -const exactMatchBoost = ( - query: string, - properties: Record, - matchedChunk?: string, -): number => { - const tokens = query.toLowerCase().split(/\s+/).filter((t) => t.length >= 3); - if (tokens.length === 0) return 0; - const filename = String(properties.filename ?? "").toLowerCase(); - const summary = String(properties.summary ?? "").toLowerCase(); - const chunk = (matchedChunk ?? "").toLowerCase(); - let bonus = 0; - for (const tok of tokens) { - // Filename match is the strongest signal β€” user named the file with this word. - if (filename.includes(tok)) bonus += 0.08; - if (summary.includes(tok)) bonus += 0.04; - if (chunk.includes(tok)) bonus += 0.03; - } - return bonus; -}; - -export const runSearchPipeline = async ( - userId: string, - rawQuery: string, - pairs: { summary: string; chunks: string }[], -): Promise => { - const analysis = analyzeQuery(rawQuery); - const queryForRetrieval = analysis.cleaned || analysis.raw; - - const client = await getWeaviateClient(); - if (!client) return []; - - // Embedding + entity extraction run in parallel β€” both are slow network calls. - const [queryVector, qe] = await Promise.all([ - generateQueryEmbedding(queryForRetrieval), - extractQueryEntities(queryForRetrieval).catch(() => null), - ]); - - const weights = intentWeights(analysis.intent); - logger.info( - `search: q="${rawQuery}" intent=` + - `${analysis.intent.filename ? "FILENAME " : ""}` + - `${analysis.intent.recency ? "RECENT " : ""}` + - `${analysis.intent.content ? "CONTENT " : ""}` + - `phrases=${analysis.intent.exactPhrases.length} ` + - `weights=${JSON.stringify(weights)}`, - ); - - // Run all signals across all collection pairs in parallel. - type SignalBlock = { name: string; hits: SignalHit[]; weight: number }; - const allSignals: SignalBlock[] = []; - const entityMatchByFile = new Map(); - - await Promise.all(pairs.flatMap(({ summary: summaryName, chunks: chunksName }) => [ - chunkHybridSignal(client, chunksName, userId, queryForRetrieval, queryVector) - .then((hits) => allSignals.push({ name: "content", hits, weight: weights.chunk })), - chunkBM25Signal(client, chunksName, userId, queryForRetrieval) - .then((hits) => allSignals.push({ name: "content-bm25", hits, weight: weights.chunkBM25 })), - filenameBM25Signal(client, summaryName, userId, queryForRetrieval) - .then((hits) => allSignals.push({ name: "filename", hits, weight: weights.filename })), - summaryHybridSignal(client, summaryName, userId, queryForRetrieval, queryVector) - .then((hits) => allSignals.push({ name: "summary", hits, weight: weights.summary })), - (async () => { - if (!qe || (!qe.entities.length && !qe.dates.length && !qe.doc_ids.length && !qe.topics.length)) { - return; - } - const { hits, matches } = await entitySignal(client, summaryName, userId, qe); - allSignals.push({ name: "entities", hits, weight: weights.entity }); - for (const [fid, m] of matches) entityMatchByFile.set(fid, m); - })(), - ])); - - // Fuse. - const fused = rrfFuse(allSignals); - if (fused.size === 0) { - logger.info(`search: 0 results for "${rawQuery}"`); - return []; - } - - // Enrich with parent summary properties so the UI gets filename, summary, created_at, etc. - const fileIds = [...fused.keys()]; - const summariesByFile = new Map>(); - await Promise.all(pairs.map(async ({ summary: summaryName }) => { - if (!(await client.collections.exists(summaryName))) return; - const col = client.collections.get(summaryName); - const res = await col.query.fetchObjects({ - limit: fileIds.length, - filters: Filters.and( - col.filter.byProperty("user_id").equal(userId), - col.filter.byProperty("file_id").containsAny(fileIds), - ), - }); - for (const obj of res.objects) { - const fid = String((obj.properties as { file_id?: string }).file_id ?? ""); - if (fid && !summariesByFile.has(fid)) { - summariesByFile.set(fid, obj.properties as Record); - } - } - })); - - // Final scoring: fused RRF + recency boost + exact-match boost. - // The exact-match boost is the critical safety net for rare-term queries - // ("Guwahati", "INV-12345"). Without it, RRF can demote the one file that - // uniquely contains the term in favor of files that rank middling-but-everywhere. - const recencyMultiplier = analysis.intent.recency ? 3 : 1; - const results: SearchResult[] = fileIds.map((fid) => { - const f = fused.get(fid)!; - const parent = summariesByFile.get(fid) ?? {}; - const recBoost = recencyMultiplier * recencyBoost(parent.created_at, f.score); - const exactBoost = exactMatchBoost(queryForRetrieval, parent, f.bestChunk); - const entityMatch = entityMatchByFile.get(fid); - return { - ...parent, - file_id: fid, - score: f.score + recBoost + exactBoost, - matched_in: [...f.matched_in], - matched_chunk: f.bestChunk, - matched_entities: entityMatch?.entities.length ? entityMatch.entities : undefined, - matched_dates: entityMatch?.dates.length ? entityMatch.dates : undefined, - matched_doc_ids: entityMatch?.doc_ids.length ? entityMatch.doc_ids : undefined, - }; - }); - - // Drop nameless results. These come from file_ids that the parent fetch - // couldn't find a matching summary row for (race condition between delete - // and re-extract, or orphaned chunk rows). The UI can't render them and - // they'd crash on .split() / .localeCompare(). - const cleaned = results.filter((r) => typeof r.filename === "string" && r.filename.length > 0); - cleaned.sort((a, b) => b.score - a.score); - if (cleaned.length < results.length) { - logger.warn(`search: dropped ${results.length - cleaned.length} orphan result(s) without filename`); - } - logger.info(`search: ${cleaned.length} results for "${rawQuery}" (top score=${cleaned[0]?.score.toFixed(4)})`); - - // Cap to a reasonable top-K β€” UI doesn't need 200 results. - return cleaned.slice(0, 30); -}; - -// ---------- 7. Optional LLM rerank for the top results ---------- - -/** - * Re-rank top candidates by feeding (query, filename + summary snippet) to an LLM. - * Adds 500-1500ms of latency. Off by default; opt in via SEARCH_RERANK=on env var. - * Caller is responsible for invoking this only when latency budget allows. - */ -export const llmRerankSearchResults = async ( - query: string, - results: SearchResult[], - keep: number, - geminiJSON: ( - prompt: string, - schema: Record, - maxTokens: number, - label: string, - ) => Promise, -): Promise => { - if (results.length <= keep) return results; - const pool = results.slice(0, Math.min(20, results.length)); - const formatted = pool.map((r, i) => { - const fn = String(r.filename ?? ""); - const sum = String(r.summary ?? "").slice(0, 400); - const chunk = r.matched_chunk ? String(r.matched_chunk).slice(0, 200) : ""; - return `[${i}] filename: ${fn}\nsummary: ${sum}${chunk ? `\nmatched: ${chunk}` : ""}`; - }).join("\n\n---\n\n"); - - const SCHEMA = { - type: "OBJECT", - properties: { - scores: { - type: "ARRAY", - items: { - type: "OBJECT", - properties: { - idx: { type: "INTEGER" }, - score: { type: "INTEGER", description: "0-10 relevance to the query." }, - }, - required: ["idx", "score"], - }, - }, - }, - required: ["scores"], - }; - - try { - const out = await geminiJSON<{ scores: { idx: number; score: number }[] }>( - `Rate each file's relevance to the search query on a 0-10 scale. -Query: "${query}" - -Files: -${formatted} - -Output one score per file in the same order, referencing each by its idx.`, - SCHEMA, - 1024, - "search:rerank", - ); - const map = new Map(); - for (const s of out?.scores ?? []) map.set(s.idx, s.score); - const reranked = pool - .map((r, i) => ({ r, s: map.get(i) ?? -1 })) - .filter((x) => x.s >= 0) - .sort((a, b) => b.s - a.s) - .map((x) => x.r); - // Tail (positions 21+) keeps RRF order β€” we didn't rerank them. - return [...reranked.slice(0, keep), ...results.slice(pool.length)]; - } catch (e) { - logger.warn(`llmRerankSearchResults failed: ${e}`); - return results; - } -}; diff --git a/backend/src/utils/getQueryEmbedding.ts b/backend/src/utils/getQueryEmbedding.ts index 553f966..f3a7a2d 100644 --- a/backend/src/utils/getQueryEmbedding.ts +++ b/backend/src/utils/getQueryEmbedding.ts @@ -2,33 +2,69 @@ import logger from '../logger.js'; import { GoogleGenAI } from '@google/genai'; -// 1. Initialize the Client with your API Key const genai = new GoogleGenAI({ apiKey: process.env.GOOGLE_API_KEY }); -const generateQueryEmbedding = async (query: string) => { - try { +// ---------- LRU cache for query embeddings ---------- +// Real search traffic has 30-50% repeat queries (autosuggest, refinements, +// common search terms). Caching makes the embedding LLM call free on hits. +// Max 500 entries, ~3 MB at 768 floats Γ— 4 bytes Γ— 500 β€” negligible. - const response = await genai.models.embedContent({ - model: 'gemini-embedding-001', // The current stable standard - contents: query, // Pass the string directly - config: { - taskType: 'RETRIEVAL_QUERY', - outputDimensionality: 768 - } - }); - - // 3. Access the embedding values from the response - // result.embeddings is an array; the first element contains your values - const embedding = response.embeddings; - - if (!embedding || embedding.length === 0 || !embedding[0].values) { - logger.error('API call succeeded but returned no embedding values.'); - throw new Error('Embedding generation resulted in empty values.'); - } +const MAX_CACHE_ENTRIES = 500; +const embeddingCache = new Map(); + +const cacheKey = (query: string): string => query.trim().toLowerCase(); + +const cacheGet = (key: string): number[] | undefined => { + const hit = embeddingCache.get(key); + if (!hit) return undefined; + // LRU: re-insert to mark as most recently used. + embeddingCache.delete(key); + embeddingCache.set(key, hit); + return hit; +}; + +const cacheSet = (key: string, value: number[]): void => { + embeddingCache.set(key, value); + if (embeddingCache.size > MAX_CACHE_ENTRIES) { + // Evict least-recently-used (first entry in insertion order). + const oldest = embeddingCache.keys().next().value; + if (oldest !== undefined) embeddingCache.delete(oldest); + } +}; - logger.info('Generated embedding successfully'); +const _generateRaw = async (query: string): Promise => { + const response = await genai.models.embedContent({ + model: 'gemini-embedding-001', + contents: query, + config: { + taskType: 'RETRIEVAL_QUERY', + outputDimensionality: 768, + }, + }); - return embedding[0].values; + const embedding = response.embeddings; + if (!embedding || embedding.length === 0 || !embedding[0].values) { + logger.error('API call succeeded but returned no embedding values.'); + throw new Error('Embedding generation resulted in empty values.'); + } + return embedding[0].values; +}; + +const generateQueryEmbedding = async (query: string): Promise => { + const key = cacheKey(query); + if (key) { + const cached = cacheGet(key); + if (cached) { + logger.info(`Embedding cache HIT (size=${embeddingCache.size})`); + return cached; + } + } + + try { + const vec = await _generateRaw(query); + if (key) cacheSet(key, vec); + logger.info(`Embedding cache MISS β€” generated (size=${embeddingCache.size})`); + return vec; } catch (error) { logger.error('Failed to generate query embedding:', error); throw new Error('Could not generate embedding for query.'); diff --git a/smartdrive-extractor/main.py b/smartdrive-extractor/main.py index 5670470..168df38 100644 --- a/smartdrive-extractor/main.py +++ b/smartdrive-extractor/main.py @@ -4,6 +4,7 @@ from app.app import create_app from app.environment import Environment from utils.docling import get_converter +from smartdrive_core.mongo_status import sweep_orphaned_files env = Environment.from_env() @@ -24,6 +25,16 @@ except Exception as e: logger.warning(f"Docling warm-up failed (will retry lazily): {e}") +# Sweep for orphaned files on boot. Any file stuck in `processing` for >10m +# is a victim of a previous worker death (OOM, deploy, crash). Reset to +# `pending` so the next Pub/Sub redelivery re-extracts them. +try: + n = sweep_orphaned_files(stale_minutes=10) + if n > 0: + logger.warning(f"Reset {n} orphaned file(s) on boot β€” re-extraction queued") +except Exception as e: + logger.warning(f"Orphan sweep failed (non-fatal): {e}") + if __name__ == "__main__": app = create_app(env) diff --git a/smartdrive-extractor/pyproject.toml b/smartdrive-extractor/pyproject.toml index 299e4a4..893f6ee 100644 --- a/smartdrive-extractor/pyproject.toml +++ b/smartdrive-extractor/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "weaviate-client>=4.15.3", "docling[easyocr]", "google-genai", + "pypdf>=5.0", "smartdrive-core", ] diff --git a/smartdrive-extractor/utils/docling.py b/smartdrive-extractor/utils/docling.py index 31ed4f8..683809e 100644 --- a/smartdrive-extractor/utils/docling.py +++ b/smartdrive-extractor/utils/docling.py @@ -1,4 +1,7 @@ import logging +import os +import re + from docling.document_converter import DocumentConverter, ImageFormatOption, PdfFormatOption from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import PdfPipelineOptions, EasyOcrOptions, TableStructureOptions @@ -8,69 +11,264 @@ _converter = None + def get_converter(): global _converter if _converter is None: logger.info("Initializing Docling DocumentConverter (EasyOCR forced)...") - opts = PdfPipelineOptions() opts.do_ocr = True - - # FORCE EasyOCR only (no auto) opts.ocr_options = EasyOcrOptions(lang=["en"], use_gpu=False) - - # Tables (optional) opts.do_table_structure = True opts.table_structure_options = TableStructureOptions(do_cell_matching=True) - _converter = DocumentConverter( format_options={ InputFormat.IMAGE: ImageFormatOption(pipeline_options=opts), InputFormat.PDF: PdfFormatOption(pipeline_options=opts), - # other formats use defaults } ) - return _converter -def extract_content(file_path: str) -> dict: + +# ============================================================================ +# E2 β€” Early checks: encrypted/protected PDFs and obvious junk +# ============================================================================ + +def _is_pdf(file_path: str) -> bool: + return os.path.splitext(file_path)[1].lower() == ".pdf" + + +def _pdf_is_encrypted(file_path: str) -> bool: + """Read first 4KB of a PDF and check for /Encrypt marker. Costs ~1ms vs + 60s + OOM risk of letting docling try to parse an encrypted PDF.""" + try: + with open(file_path, "rb") as f: + head = f.read(4096) + # PDF spec: /Encrypt entry in trailer or document catalog. + # Heuristic: if /Encrypt appears in the first 4KB header, it's encrypted. + return b"/Encrypt" in head + except Exception as e: + logger.warning(f"_pdf_is_encrypted check failed: {e}") + return False + + +def _pdf_is_valid(file_path: str) -> bool: + """Sanity check: file starts with %PDF magic bytes.""" + try: + with open(file_path, "rb") as f: + return f.read(5) == b"%PDF-" + except Exception: + return False + + +# ============================================================================ +# E3 β€” Quality check on extracted text +# ============================================================================ + +_NON_ALNUM_RE = re.compile(r"[^a-zA-Z0-9\s]") + + +def _looks_like_garbage(text: str) -> tuple[bool, str]: + """Return (is_garbage, reason). Catches OCR garbage and binary echoes + before we waste an LLM call summarizing them.""" + if not text: + return True, "empty" + if len(text) < 50: + return True, f"too_short ({len(text)} chars)" + + # Non-alphanumeric ratio: garbage OCR produces lots of symbols. + sample = text[:5000] + non_alnum = len(_NON_ALNUM_RE.findall(sample)) + ratio = non_alnum / max(1, len(sample)) + if ratio > 0.4: + return True, f"high_symbol_ratio ({ratio:.2f})" + + # Repetitiveness: if 80% of the text is one repeated chunk, + # it's probably a header/footer that wasn't stripped, or OCR loop. + words = text.split() + if len(words) >= 30: + unique_ratio = len(set(w.lower() for w in words[:200])) / min(200, len(words)) + if unique_ratio < 0.15: + return True, f"low_unique_ratio ({unique_ratio:.2f})" + + return False, "" + + +# ============================================================================ +# E3 fallback: PyPDF for text-extraction when docling fails +# ============================================================================ + +def _pypdf_fallback(file_path: str) -> dict: + """Last-resort PDF text extraction. No layout awareness, no OCR, but + very fast and never OOMs. For text-PDF cases where docling crashes, + this gets us *something* indexable. """ - Docs-only extraction (no OCR). - Returns both: - - markdown: for LLM + chunking - - metadata: for page/table/heading provenance (store in Mongo) + try: + from pypdf import PdfReader # pypdf >= 4.0 ships with the project's deps + except ImportError: + try: + from PyPDF2 import PdfReader # type: ignore + except ImportError: + logger.warning("pypdf/PyPDF2 not installed β€” fallback unavailable") + return {"created": False, "markdown": "", "metadata": None} + + try: + reader = PdfReader(file_path) + if reader.is_encrypted: + return {"created": False, "markdown": "", "metadata": {"error": "encrypted"}} + pages = [] + for i, page in enumerate(reader.pages): + try: + pages.append((page.extract_text() or "").strip()) + except Exception as e: + logger.warning(f"pypdf page {i} extract failed: {e}") + text = "\n\n".join(p for p in pages if p).strip() + if not text: + return {"created": False, "markdown": "", "metadata": None} + logger.info(f"pypdf fallback extracted {len(text)} chars from {len(pages)} pages") + return { + "created": True, + "markdown": text, + "metadata": {"fallback": "pypdf", "page_count": len(pages)}, + } + except Exception as e: + logger.warning(f"pypdf fallback failed: {e}") + return {"created": False, "markdown": "", "metadata": None} + + +# ============================================================================ +# Main entry β€” docling primary with layered guards and fallbacks +# ============================================================================ + +def extract_content(file_path: str) -> dict: + """Extraction with guards: + 1. If PDF: validate magic bytes + encryption check (fail fast) + 2. Try docling (primary) + 3. If docling fails OR returns garbage: try pypdf fallback for PDFs + 4. If still nothing: plain-text fallback for known extensions """ - # logger.info(f"πŸ“„ Docling extracting from: {file_path}") + # E2 β€” Early PDF validation + if _is_pdf(file_path): + if not _pdf_is_valid(file_path): + logger.warning(f"PDF magic bytes missing for {file_path} β€” not a valid PDF") + return {"created": False, "markdown": "", "metadata": {"error": "invalid_pdf"}} + if _pdf_is_encrypted(file_path): + logger.warning(f"PDF is encrypted, skipping: {file_path}") + return { + "created": False, + "markdown": "", + "metadata": {"error": "encrypted"}, + "error_kind": "encrypted_pdf", + } + + # Primary: docling + docling_result = _run_docling(file_path) + + # E3 β€” Quality check + if docling_result.get("created"): + text = docling_result.get("markdown", "") + is_garbage, reason = _looks_like_garbage(text) + if is_garbage: + logger.warning(f"Docling output failed quality check ({reason}) β€” trying fallback") + docling_result = {"created": False, "markdown": "", "metadata": {"quality_fail": reason}} + + if docling_result.get("created"): + return docling_result + + # Fallback chain + if _is_pdf(file_path): + pypdf_result = _pypdf_fallback(file_path) + if pypdf_result.get("created"): + # Quality check the fallback too + text = pypdf_result.get("markdown", "") + is_garbage, reason = _looks_like_garbage(text) + if not is_garbage: + return pypdf_result + logger.warning(f"pypdf fallback also failed quality ({reason})") + # Last resort: plain-text for known extensions + return extract_plain_text(file_path) + + +def _is_image_only_pdf(file_path: str) -> bool: + """E7 β€” heuristic detection of image-only (scanned) PDFs. + Image-only PDFs have low/no extractable text but high page count or + large file size. These are OCR-heavy and need more memory + time. + Costs ~50ms via pypdf vs minutes via docling-then-fail.""" try: + from pypdf import PdfReader + except ImportError: + return False + try: + reader = PdfReader(file_path) + if not reader.pages: + return False + # Sample first 3 pages β€” if they have <50 chars of extractable text + # but the PDF is >100KB, it's likely image-only (scanned). + sample_text = "" + for page in reader.pages[:3]: + try: + sample_text += (page.extract_text() or "") + except Exception: + pass + file_size = os.path.getsize(file_path) + if len(sample_text.strip()) < 50 and file_size > 100_000: + return True + return False + except Exception: + return False + + +def _run_docling(file_path: str) -> dict: + try: + # E7 β€” log image-only route for memory observability. + if _is_pdf(file_path) and _is_image_only_pdf(file_path): + file_size_mb = os.path.getsize(file_path) / (1024 * 1024) + logger.info( + f"πŸ“· Image-only PDF detected ({file_size_mb:.1f}MB) β€” " + f"OCR-heavy path; expect higher memory + slower extraction" + ) + result = get_converter().convert(file_path) doc = result.document - markdown = (doc.export_to_markdown() or "").strip() - text = (doc.export_to_text() or "").strip() + text = (doc.export_to_text() or "").strip() metadata = doc.export_to_dict() + # E6 β€” per-page observability. Pulls page count from metadata if + # present so we can compare extracted_chars / page_count and spot + # half-extracted PDFs. + page_count = 0 + try: + page_count = len((metadata or {}).get("pages", []) or []) + except Exception: + pass + if markdown: - logger.info(f"βœ… Extracted text preview: {markdown[:50]}...") - elif not markdown and text: - logger.info(f"βœ… Extracted text (no markdown) preview: {text[:50]}...") + logger.info( + f"βœ… Docling extracted {len(markdown)} chars across {page_count or '?'} pages " + f"β€” preview: {markdown[:50]}..." + ) + if page_count > 0: + chars_per_page = len(markdown) / page_count + if chars_per_page < 50: + logger.warning( + f"⚠️ Low chars/page ratio ({chars_per_page:.0f}) β€” " + f"PDF may be partially extracted (image-only pages, OCR misfires)" + ) + elif text: + logger.info(f"βœ… Docling extracted {len(text)} chars (text-only fallback)") markdown = text else: - logger.warning("⚠️ Docling returned empty markdown.") + logger.warning(f"⚠️ Docling returned empty content (page_count={page_count})") + return {"created": False, "markdown": "", "metadata": metadata} - return { - "created": True, - "markdown": markdown, - "metadata": metadata, # rich structure for citations later - } + return {"created": True, "markdown": markdown, "metadata": metadata} except Exception as e: msg = str(e) - - # βœ… fallback only on β€œnot allowed” / unsupported formats if "File format not allowed" in msg or "does not match any" in msg: - logger.warning(f"Docling unsupported format, using plain-text fallback: {msg}") - return extract_plain_text(file_path) - - logger.exception(f"❌ Docling extraction failed: {e}") - return {"created": False, "markdown": "", "metadata": None} \ No newline at end of file + logger.warning(f"Docling unsupported format: {msg}") + else: + logger.exception(f"❌ Docling extraction crashed: {e}") + return {"created": False, "markdown": "", "metadata": None} diff --git a/smartdrive-extractor/utils/document_extractor.py b/smartdrive-extractor/utils/document_extractor.py index 82ec8a8..93f90e1 100644 --- a/smartdrive-extractor/utils/document_extractor.py +++ b/smartdrive-extractor/utils/document_extractor.py @@ -3,6 +3,7 @@ from .docling import extract_content from smartdrive_core.llm import LLM_doc_summarize, get_embedding from smartdrive_core.metrics import stage_timer +from smartdrive_core.mongo_status import update_progress from .weaviate_utils import save_doc, save_doc_private logger = logging.getLogger(__name__) @@ -32,6 +33,7 @@ def process_document(file_path: str, data: dict) -> dict: } # ---- extract ---- + update_progress(file_id, "Extracting text and tables", 1, 4) with stage_timer("extract", file_id=file_id): res = extract_content(file_path) if not res or not res.get("created"): @@ -42,12 +44,14 @@ def process_document(file_path: str, data: dict) -> dict: return {"message": f"No text extracted for {filename}", "created": False, "error_kind": "no_content"} # ---- summarise ---- + update_progress(file_id, "Summarizing with AI", 2, 4) with stage_timer("summarize", file_id=file_id, chars=len(markdown)): summary, index_json = LLM_doc_summarize(markdown) if not summary: return {"message": "Failed to generate summary", "created": False, "error_kind": "llm_failed"} # ---- embed summary only (file-level vector for cross-file search) ---- + update_progress(file_id, "Embedding for search", 3, 4) summary_embed_text = f"{summary}\n\nKeywords: {index_json}" with stage_timer("embed_summary", file_id=file_id): summary_vector = get_embedding(summary_embed_text) @@ -55,6 +59,7 @@ def process_document(file_path: str, data: dict) -> dict: return {"message": "Failed to embed summary", "created": False, "error_kind": "embedding_failed"} # ---- save summary row with raw_text (used later for lazy chat prep) ---- + update_progress(file_id, "Indexing", 4, 4) with stage_timer("save_summary", file_id=file_id): save_doc( data, diff --git a/smartdrive-extractor/utils/weaviate_utils.py b/smartdrive-extractor/utils/weaviate_utils.py index f3e51b1..fd14ecb 100644 --- a/smartdrive-extractor/utils/weaviate_utils.py +++ b/smartdrive-extractor/utils/weaviate_utils.py @@ -45,6 +45,11 @@ def _norm_list(values) -> list[str]: wvc.config.Property(name="user_id", data_type=wvc.config.DataType.TEXT), wvc.config.Property(name="summary", data_type=wvc.config.DataType.TEXT), wvc.config.Property(name="raw_text", data_type=wvc.config.DataType.TEXT, index_searchable=False), + # R3 β€” searchable body text. raw_text was marked index_searchable=False + # (storage-only), which silently broke BM25 over body content for files + # without chunks. We add a searchable mirror; backend uses this for + # body-text BM25 queries. Auto-populated alongside raw_text below. + wvc.config.Property(name="body_text", data_type=wvc.config.DataType.TEXT, index_searchable=True), wvc.config.Property(name="filename", data_type=wvc.config.DataType.TEXT), wvc.config.Property(name="filetype", data_type=wvc.config.DataType.TEXT), wvc.config.Property(name="created_at", data_type=wvc.config.DataType.DATE, index_filterable=True), @@ -77,6 +82,8 @@ def _norm_list(values) -> list[str]: wvc.config.Property(name="user_id", data_type=wvc.config.DataType.TEXT), wvc.config.Property(name="summary", data_type=wvc.config.DataType.TEXT), wvc.config.Property(name="raw_text", data_type=wvc.config.DataType.TEXT, index_searchable=False), + # R3 β€” searchable body mirror (see DOC_PROPERTIES for explanation). + wvc.config.Property(name="body_text", data_type=wvc.config.DataType.TEXT, index_searchable=True), wvc.config.Property(name="filename", data_type=wvc.config.DataType.TEXT), wvc.config.Property(name="filetype", data_type=wvc.config.DataType.TEXT), wvc.config.Property(name="created_at", data_type=wvc.config.DataType.DATE, index_filterable=True), @@ -114,6 +121,7 @@ def save_doc_private(data): "user_id": data["userId"], "summary": _PRIVATE_PLACEHOLDER_SUMMARY, "raw_text": "", + "body_text": "", # R3 β€” searchable mirror (empty for private/stub files) "index_json": "{}", "filetype": data.get("fileType"), "created_at": data.get("uploadedAt"), @@ -136,6 +144,7 @@ def save_doc(data, summary, index_json, embedding, raw_text: str = "", chunk_cou "user_id": data["userId"], "summary": summary, "raw_text": raw_text or "", + "body_text": raw_text or "", # R3 β€” searchable mirror "index_json": json.dumps(idx, ensure_ascii=False), "filetype": data.get("fileType"), "created_at": data.get("uploadedAt"), @@ -186,6 +195,7 @@ def save_image_private(data): "user_id": data["userId"], "summary": _PRIVATE_PLACEHOLDER_SUMMARY, "raw_text": "", + "body_text": "", # R3 β€” searchable mirror (empty for private/stub files) "filetype": data.get("fileType"), "created_at": data.get("uploadedAt"), "processing_type": "private", @@ -202,6 +212,7 @@ def save_image(data, summary, embedding, processing_type, raw_text: str = "", ch "user_id": data["userId"], "summary": summary, "raw_text": raw_text or "", + "body_text": raw_text or "", # R3 β€” searchable mirror "filetype": data.get("fileType"), "created_at": data.get("uploadedAt"), "processing_type": processing_type, diff --git a/smartdrive-frontend/src/components/FileListWithDrawer.tsx b/smartdrive-frontend/src/components/FileListWithDrawer.tsx index 5c9a2f2..4ce687b 100644 --- a/smartdrive-frontend/src/components/FileListWithDrawer.tsx +++ b/smartdrive-frontend/src/components/FileListWithDrawer.tsx @@ -80,6 +80,9 @@ export interface UploadItem { extraction_error?: string; index_json?: IndexJson; is_private?: boolean; + /** Live extraction progress reported by the worker. Present only while + * status is pending/processing. e.g. { current: 2, total: 4, stage: "Summarizing with AI" } */ + extraction_progress?: { current?: number; total?: number; stage?: string } | null; /** Set by the search endpoint: the chunk text that scored highest against * the user's query. Surfaced on the card so users know *why* a file matched. */ @@ -207,7 +210,13 @@ function typeAccent(filetype: string) { }; } -function ExtractionBadge({ status }: { status: ExtractionStatus | undefined }) { +function ExtractionBadge({ + status, + progress, +}: { + status: ExtractionStatus | undefined; + progress?: { current?: number; total?: number; stage?: string } | null; +}) { const s = status ?? 'done'; if (s === 'done') { return ( @@ -225,11 +234,21 @@ function ExtractionBadge({ status }: { status: ExtractionStatus | undefined }) { ); } - // pending or processing + // pending or processing β€” show live stage if the worker is reporting it + const stage = progress?.stage?.trim(); + const cur = progress?.current ?? 0; + const total = progress?.total ?? 0; + const showStage = s === 'processing' && stage && total > 0; return ( - - - {s === 'pending' ? 'Queued' : 'Processing'} + + + + {showStage ? `${stage} (${cur}/${total})` : (s === 'pending' ? 'Queued' : 'Processing')} + ); } @@ -580,7 +599,10 @@ function FileCard({ {/* Status row */}
- +
{/* Hero summary β€” 3 lines, soft fade-out for overflow */} @@ -1235,7 +1257,10 @@ export function FileListWithDrawer({ {fileTypeLabel(file.filetype)} Β· {formatRelative(file.created_at)} - +