diff --git a/frontend/app/api/queries/useGetSearchQuery.test.tsx b/frontend/app/api/queries/useGetSearchQuery.test.tsx index d29fadb5b..b79bb974e 100644 --- a/frontend/app/api/queries/useGetSearchQuery.test.tsx +++ b/frontend/app/api/queries/useGetSearchQuery.test.tsx @@ -5,6 +5,7 @@ import { server } from "@/test-utils/msw/server"; import { createQueryWrapper, renderHook, waitFor } from "@/test-utils/render"; import { type ChunkResult, + resolveScoreThreshold, type SearchPayload, useGetSearchQuery, } from "./useGetSearchQuery"; @@ -390,3 +391,23 @@ describe("useGetSearchQuery", () => { }); }); }); + +describe("resolveScoreThreshold", () => { + it("lowers the threshold to 1.0 for a short single token", () => { + expect(resolveScoreThreshold("docl", 1.25)).toBe(1.0); + expect(resolveScoreThreshold(" ab ", 3)).toBe(1.0); + }); + + it("never raises a threshold that is already below 1.0", () => { + expect(resolveScoreThreshold("docl", 0.5)).toBe(0.5); + }); + + it.each([ + ["docling", "longer than four characters"], + ["a b", "more than one token"], + ["*", "the wildcard"], + ["", "empty"], + ])("keeps the base threshold for %j (%s)", (query) => { + expect(resolveScoreThreshold(query, 1.25)).toBe(1.25); + }); +}); diff --git a/frontend/app/api/queries/useGetSearchQuery.ts b/frontend/app/api/queries/useGetSearchQuery.ts index da4a0dd2e..96266b4a6 100644 --- a/frontend/app/api/queries/useGetSearchQuery.ts +++ b/frontend/app/api/queries/useGetSearchQuery.ts @@ -117,6 +117,22 @@ const getFileIdentity = (chunk: ChunkResult): string => { return "Untitled source"; }; +/** Short single-token queries get a lower threshold so partial words still match. */ +export function resolveScoreThreshold( + query: string, + baseScoreThreshold: number, +): number { + const normalized = query.trim(); + const isShortSingleToken = + normalized !== "*" && + normalized.length > 0 && + normalized.length <= 4 && + !normalized.includes(" "); + return isShortSingleToken + ? Math.min(baseScoreThreshold, 1.0) + : baseScoreThreshold; +} + export const useGetSearchQuery = ( query: string, queryData?: ParsedQueryData | null, @@ -143,14 +159,10 @@ export const useGetSearchQuery = ( const baseScoreThreshold = queryData?.scoreThreshold ?? SEARCH_CONSTANTS.DEFAULT_SCORE_THRESHOLD; - const isShortSingleTokenQuery = - normalizedQuery !== "*" && - normalizedQuery.length > 0 && - normalizedQuery.length <= 4 && - !normalizedQuery.includes(" "); - const dynamicScoreThreshold = isShortSingleTokenQuery - ? Math.min(baseScoreThreshold, 1.0) - : baseScoreThreshold; + const dynamicScoreThreshold = resolveScoreThreshold( + normalizedQuery, + baseScoreThreshold, + ); const searchPayload: SearchPayload = { query: effectiveQuery, diff --git a/frontend/app/api/queries/useGetSearchSuggestions.test.tsx b/frontend/app/api/queries/useGetSearchSuggestions.test.tsx new file mode 100644 index 000000000..178677754 --- /dev/null +++ b/frontend/app/api/queries/useGetSearchSuggestions.test.tsx @@ -0,0 +1,236 @@ +import { HttpResponse, http } from "msw"; +import { describe, expect, it } from "vitest"; +import { server } from "@/test-utils/msw/server"; +import { + act, + createQueryWrapper, + renderHook, + waitFor, +} from "@/test-utils/render"; +import { + type SearchSuggestResult, + SUGGEST_KEYWORD_DEBOUNCE_MS, + SUGGEST_MAX_CHARS, + useSearchSuggestions, +} from "./useGetSearchSuggestions"; + +interface SuggestRequest { + query: string; + semantic: boolean; + filters: Record; + scoreThreshold?: number; +} + +function result( + phase: "keyword" | "semantic", + query: string, +): SearchSuggestResult { + return { + phase, + query, + suggestions: [{ text: `${query}-${phase}`, files: 1, chunks: 2 }], + files: [], + }; +} + +function serveSuggest() { + const requests: SuggestRequest[] = []; + server.use( + http.post("/api/search/suggest", async ({ request }) => { + const body = (await request.json()) as SuggestRequest; + requests.push(body); + return HttpResponse.json( + result(body.semantic ? "semantic" : "keyword", body.query), + ); + }), + ); + return requests; +} + +describe("useSearchSuggestions", () => { + it("sends a keyword request first, then upgrades to the semantic result", async () => { + const requests = serveSuggest(); + const { result: hook } = renderHook( + () => useSearchSuggestions("docling", undefined), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(hook.current.data?.phase).toBe("semantic")); + + expect(requests.map((r) => r.semantic)).toEqual( + expect.arrayContaining([false, true]), + ); + expect(hook.current.data?.suggestions[0].text).toBe("docling-semantic"); + expect(hook.current.isRefining).toBe(false); + }); + + it("sends nothing below the minimum length", async () => { + const requests = serveSuggest(); + const { result: hook } = renderHook( + () => useSearchSuggestions("d", undefined), + { wrapper: createQueryWrapper() }, + ); + + await new Promise((resolve) => + setTimeout(resolve, SUGGEST_KEYWORD_DEBOUNCE_MS + 100), + ); + + expect(requests).toHaveLength(0); + expect(hook.current.data).toBeUndefined(); + }); + + it("does not request anything while disabled", async () => { + const requests = serveSuggest(); + renderHook(() => useSearchSuggestions("docling", undefined, false), { + wrapper: createQueryWrapper(), + }); + + await new Promise((resolve) => + setTimeout(resolve, SUGGEST_KEYWORD_DEBOUNCE_MS + 100), + ); + + expect(requests).toHaveLength(0); + }); + + it("forwards active filters and drops the wildcard sentinel", async () => { + const requests = serveSuggest(); + renderHook( + () => + useSearchSuggestions("docling", { + data_sources: ["a.pdf"], + owners: ["*"], + }), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(requests.length).toBeGreaterThan(0)); + expect(requests[0].filters).toEqual({ data_sources: ["a.pdf"] }); + }); + + it("forwards the base score threshold on both phases", async () => { + const requests = serveSuggest(); + renderHook(() => useSearchSuggestions("docling", undefined, true, 2.5), { + wrapper: createQueryWrapper(), + }); + + await waitFor(() => expect(requests.some((r) => r.semantic)).toBe(true)); + expect(requests.map((r) => r.scoreThreshold)).toEqual( + expect.arrayContaining([2.5]), + ); + expect(requests.every((r) => r.scoreThreshold === 2.5)).toBe(true); + }); + + it("lowers the threshold for a short single token like the main search", async () => { + const requests = serveSuggest(); + renderHook(() => useSearchSuggestions("docl", undefined, true, 2.5), { + wrapper: createQueryWrapper(), + }); + + await waitFor(() => expect(requests.some((r) => r.semantic)).toBe(true)); + expect(requests.every((r) => r.scoreThreshold === 1.0)).toBe(true); + }); + + it("reports an error when the keyword request fails and nothing is cached", async () => { + server.use( + http.post("/api/search/suggest", () => + HttpResponse.json( + { error: "Suggestions unavailable" }, + { status: 500 }, + ), + ), + ); + const { result: hook } = renderHook( + () => useSearchSuggestions("docling", undefined), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(hook.current.isError).toBe(true)); + expect(hook.current.data).toBeUndefined(); + expect(hook.current.isRefining).toBe(false); + }); + + it("keeps showing keyword results when only the semantic request fails", async () => { + server.use( + http.post("/api/search/suggest", async ({ request }) => { + const body = (await request.json()) as SuggestRequest; + if (body.semantic) { + return HttpResponse.json({ error: "x" }, { status: 502 }); + } + return HttpResponse.json(result("keyword", body.query)); + }), + ); + const { result: hook } = renderHook( + () => useSearchSuggestions("docling", undefined), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(hook.current.data?.phase).toBe("keyword")); + await waitFor(() => expect(hook.current.isRefining).toBe(false)); + expect(hook.current.isError).toBe(false); + }); + + it("only debounces after typing pauses, collapsing a burst into one keyword request", async () => { + const requests = serveSuggest(); + const { rerender } = renderHook( + ({ q }) => useSearchSuggestions(q, undefined), + { wrapper: createQueryWrapper(), initialProps: { q: "" } }, + ); + + for (const q of ["do", "doc", "docl", "docli"]) { + await act(async () => { + rerender({ q }); + await new Promise((resolve) => setTimeout(resolve, 40)); + }); + } + + await waitFor(() => + expect(requests.some((r) => !r.semantic && r.query === "docli")).toBe( + true, + ), + ); + const keywordQueries = requests + .filter((r) => !r.semantic) + .map((r) => r.query); + expect(keywordQueries).toEqual(["docli"]); + }); + + it("truncates over-long input to the server's length cap", async () => { + const requests = serveSuggest(); + const { result: hook } = renderHook( + () => useSearchSuggestions("a".repeat(150), undefined), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => + expect(hook.current.data?.query).toBe("a".repeat(SUGGEST_MAX_CHARS)), + ); + expect(requests.every((r) => r.query.length === SUGGEST_MAX_CHARS)).toBe( + true, + ); + }); + + it("collapses internal whitespace so equivalent input shares one request", async () => { + const requests = serveSuggest(); + const { result: hook } = renderHook( + () => useSearchSuggestions(" open rag ", undefined), + { wrapper: createQueryWrapper() }, + ); + + await waitFor(() => expect(hook.current.data?.query).toBe("open rag")); + expect(requests.every((r) => r.query === "open rag")).toBe(true); + }); + + it("hides the previous query's suggestions while the new query debounces", async () => { + serveSuggest(); + const { result: hook, rerender } = renderHook( + ({ q }) => useSearchSuggestions(q, undefined), + { wrapper: createQueryWrapper(), initialProps: { q: "docling" } }, + ); + await waitFor(() => expect(hook.current.data?.query).toBe("docling")); + + rerender({ q: "docker" }); + expect(hook.current.data).toBeUndefined(); + + await waitFor(() => expect(hook.current.data?.query).toBe("docker")); + }); +}); diff --git a/frontend/app/api/queries/useGetSearchSuggestions.ts b/frontend/app/api/queries/useGetSearchSuggestions.ts new file mode 100644 index 000000000..0a7314d5f --- /dev/null +++ b/frontend/app/api/queries/useGetSearchSuggestions.ts @@ -0,0 +1,197 @@ +import { useQuery } from "@tanstack/react-query"; +import { resolveScoreThreshold } from "@/app/api/queries/useGetSearchQuery"; +import { useDebouncedValue } from "@/hooks/use-debounced-value"; +import { apiClient } from "@/lib/api-client"; +import { SEARCH_CONSTANTS } from "@/lib/constants"; +import { + buildSearchPayloadFilters, + type FilterInput, +} from "@/lib/filter-normalization"; + +/** Wait after the last keystroke before the cheap keyword request. */ +export const SUGGEST_KEYWORD_DEBOUNCE_MS = 180; +/** Wait before the embedding-backed request (a partial word has no meaning). */ +export const SUGGEST_SEMANTIC_DEBOUNCE_MS = 450; +export const SUGGEST_MIN_CHARS = 2; +/** Matches the server's MAX_QUERY_LENGTH; longer input is truncated there. */ +export const SUGGEST_MAX_CHARS = 100; +export const SUGGEST_SEMANTIC_MIN_CHARS = 3; +const SUGGEST_STALE_TIME_MS = 30_000; + +export interface SearchSuggestion { + text: string; + files: number; + chunks: number; +} + +export interface SuggestPreview { + /** One line; `keyword` previews wrap matched terms in ``. */ + text: string; + kind: "keyword" | "semantic"; +} + +export interface SuggestFile { + filename: string; + mimetype?: string | null; + connector_type?: string | null; + preview: SuggestPreview | null; +} + +export interface SearchSuggestResult { + phase: "keyword" | "semantic"; + query: string; + suggestions: SearchSuggestion[]; + files: SuggestFile[]; +} + +async function fetchSuggestions( + query: string, + semantic: boolean, + filters: FilterInput | undefined, + scoreThreshold: number, + signal: AbortSignal, +): Promise { + const response = await apiClient.post( + "/search/suggest", + { + query, + semantic, + filters: buildSearchPayloadFilters(filters) ?? {}, + scoreThreshold, + }, + { signal }, + ); + if (response.status < 200 || response.status >= 300) { + throw new Error(`Suggestions failed with status ${response.status}`); + } + return response.data as SearchSuggestResult; +} + +export interface UseSearchSuggestionsResult { + data: SearchSuggestResult | undefined; + /** True while the semantic refinement for the current text is still pending. */ + isRefining: boolean; + isFetching: boolean; + /** The keyword request failed and there is nothing to show. */ + isError: boolean; +} + +/** + * Search-as-you-type suggestions in two phases. + * + * Phase 1 (keyword) runs on a short debounce and never touches the embedding + * provider. Phase 2 (semantic) waits for a longer idle period, or fires at + * once on a word boundary (trailing space). The semantic result replaces the + * keyword one when it lands. Both are cached by TanStack Query (and again on + * the server), so retyping or backspacing costs no network call. + */ +export function useSearchSuggestions( + rawQuery: string, + filters: FilterInput | undefined, + enabled = true, + baseScoreThreshold: number = SEARCH_CONSTANTS.DEFAULT_SCORE_THRESHOLD, +): UseSearchSuggestionsResult { + // Mirror the server's normalization (whitespace collapse, length cap) so the + // returned query matches what was typed and equivalent inputs share a cache + // entry. Spread to count code points like the server does. + const trimmed = [...rawQuery.trim().replace(/\s+/g, " ")] + .slice(0, SUGGEST_MAX_CHARS) + .join(""); + const keywordQuery = useDebouncedValue(trimmed, SUGGEST_KEYWORD_DEBOUNCE_MS); + const atWordBoundary = rawQuery.endsWith(" "); + const semanticQuery = useDebouncedValue( + trimmed, + atWordBoundary ? 0 : SUGGEST_SEMANTIC_DEBOUNCE_MS, + ); + + const filtersKey = JSON.stringify(buildSearchPayloadFilters(filters) ?? {}); + + const keywordThreshold = resolveScoreThreshold( + keywordQuery, + baseScoreThreshold, + ); + const keyword = useQuery({ + queryKey: ["search-suggest", "keyword", keywordQuery, filtersKey], + queryFn: ({ signal }) => + fetchSuggestions(keywordQuery, false, filters, keywordThreshold, signal), + enabled: enabled && keywordQuery.length >= SUGGEST_MIN_CHARS, + staleTime: SUGGEST_STALE_TIME_MS, + retry: false, + }); + + const semanticEnabled = + enabled && semanticQuery.length >= SUGGEST_SEMANTIC_MIN_CHARS; + const scoreThreshold = resolveScoreThreshold( + semanticQuery, + baseScoreThreshold, + ); + const semantic = useQuery({ + queryKey: [ + "search-suggest", + "semantic", + semanticQuery, + filtersKey, + scoreThreshold, + ], + queryFn: ({ signal }) => + fetchSuggestions(semanticQuery, true, filters, scoreThreshold, signal), + enabled: semanticEnabled, + staleTime: SUGGEST_STALE_TIME_MS, + retry: false, + }); + + return { + ...mergePhases({ + trimmed, + keywordQuery, + semanticQuery, + semanticEnabled, + keyword, + semantic, + }), + isFetching: keyword.isFetching || semantic.isFetching, + }; +} + +interface PhaseState { + data: SearchSuggestResult | undefined; + isError: boolean; + isFetching: boolean; +} + +function mergePhases({ + trimmed, + keywordQuery, + semanticQuery, + semanticEnabled, + keyword, + semantic, +}: { + trimmed: string; + keywordQuery: string; + semanticQuery: string; + semanticEnabled: boolean; + keyword: PhaseState; + semantic: PhaseState; +}): Omit { + // Only trust results that match the text currently shown, so a late response + // or a still-debouncing query never displays another search's suggestions. + const semanticCurrent = + semantic.data && semanticQuery === trimmed ? semantic.data : undefined; + const keywordCurrent = + keyword.data && keywordQuery === trimmed ? keyword.data : undefined; + const data = semanticCurrent ?? keywordCurrent; + const visibleData = trimmed.length >= SUGGEST_MIN_CHARS ? data : undefined; + + const isRefining = + semanticEnabled && + !semanticCurrent && + !semantic.isError && + (semanticQuery !== trimmed || semantic.isFetching); + + return { + data: visibleData, + isRefining, + isError: !visibleData && (keyword.isError || semantic.isError), + }; +} diff --git a/frontend/app/knowledge/page.tsx b/frontend/app/knowledge/page.tsx index 624db4f1b..750eb9738 100644 --- a/frontend/app/knowledge/page.tsx +++ b/frontend/app/knowledge/page.tsx @@ -14,6 +14,11 @@ import { AgGridReact, type CustomCellRendererProps } from "ag-grid-react"; import { AlertTriangle, Cloud, FileIcon, Globe, RefreshCw } from "lucide-react"; import { useRouter } from "next/navigation"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { + BOLD_MATCH_CLASS, + HighlightedText, + PREVIEW_TEXT_CLASS, +} from "@/components/highlighted-text"; import { KnowledgeDropdown } from "@/components/knowledge-dropdown"; import { ProtectedRoute } from "@/components/protected-route"; import { Banner, BannerIcon, BannerTitle } from "@/components/ui/banner"; @@ -22,6 +27,7 @@ import { useOpenTaskMenu } from "@/contexts/console-status-context"; import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; import { useTask } from "@/contexts/task-context"; import { trackButton } from "@/lib/analytics"; +import { buildChunksUrl } from "@/lib/chunks-url"; import { isFileCancelled } from "@/lib/task-error-display"; import { EMPTY_SEARCH_RESULT, @@ -54,6 +60,11 @@ import { buildKnowledgeTableRows, getKnowledgeFileIdentity, } from "@/lib/knowledge-table-state"; +import { + computeTopPreviews, + type SearchPreview, + stripMarks, +} from "@/lib/search-preview"; import { parseTimestampMs } from "@/lib/time-utils"; import { cn } from "@/lib/utils"; import { @@ -254,17 +265,8 @@ export function getFileStatus(status?: File["status"]): string { return status || "active"; } -export function buildChunksUrl( - filename: string, - effectiveSearchText: string, -): string { - const params = new URLSearchParams({ filename }); - const trimmed = effectiveSearchText.trim(); - if (trimmed && trimmed !== "*") { - params.set("q", trimmed); - } - return `/knowledge/chunks?${params.toString()}`; -} +// Lives in lib/ so components can build the link without importing this page. +export { buildChunksUrl }; const AVG_SCORE_TOOLTIP = "Average relevance score across the matched chunks of this file. Higher means a stronger match — sort by this column to rank results by relevance."; @@ -663,6 +665,16 @@ function SearchPage() { const error = isWildcardQuery ? listFilesError : searchError; const isError = isWildcardQuery ? isListFilesError : isSearchError; + // One-line "why it matched" for the most relevant files across the whole + // result set (not just the visible page). + const topPreviews = useMemo( + () => + isWildcardQuery + ? new Map() + : computeTopPreviews(searchFiles, effectiveSearchText), + [isWildcardQuery, searchFiles, effectiveSearchText], + ); + const effectiveData: File[] = isWildcardQuery ? (listFilesData?.files ?? []) : searchFiles.slice( @@ -843,6 +855,7 @@ function SearchPage() { const isActive = status === "active"; const showOpenragSourceAnimation = isOpenragDocsRow(data) && hasOpenragRefreshCue; + const preview = topPreviews.get(getKnowledgeFileIdentity(data)); return (
{getSourceIcon(data?.connector_type)} - - + + + + + {filename} + + + + {filename} + + + {preview && ( - {filename} + {preview.kind === "keyword" ? ( + + ) : ( + preview.text + )} - - - {filename} - - + )} +
); @@ -895,9 +930,23 @@ function SearchPage() { hasOpenragRefreshCue, effectiveSearchText, router, + topPreviews, ], ); + // Rows that carry a preview line need extra height. The cloud grid already + // uses a 64px row, which fits two lines. + const getRowHeight = useCallback( + (params: { data?: File }) => + topPreviews.has(getKnowledgeFileIdentity(params.data)) ? 56 : undefined, + [topPreviews], + ); + + // biome-ignore lint/correctness/useExhaustiveDependencies: topPreviews changes which rows are tall, so heights must be recomputed. + useEffect(() => { + getGridApi()?.resetRowHeights(); + }, [topPreviews, getGridApi]); + const colSize: ColDef = { field: "size", headerName: "Size", @@ -1352,6 +1401,7 @@ function SearchPage() { getFileIdentity(params.data) } isRowSelectable={(params) => isDeletableKnowledgeRow(params.data)} + getRowHeight={getRowHeight} domLayout="normal" onGridReady={handleGridReady} onGridPreDestroyed={handleGridPreDestroyed} diff --git a/frontend/components/highlighted-text.test.tsx b/frontend/components/highlighted-text.test.tsx index bc8efb9c6..446acf8af 100644 --- a/frontend/components/highlighted-text.test.tsx +++ b/frontend/components/highlighted-text.test.tsx @@ -1,6 +1,38 @@ import { render, screen } from "@testing-library/react"; import { describe, expect, it } from "vitest"; -import { HighlightedText } from "./highlighted-text"; +import { + BOLD_MATCH_CLASS, + HighlightedText, + PREVIEW_TEXT_CLASS, +} from "./highlighted-text"; + +describe("preview highlight styles", () => { + it("bold match style is theme-token based with no hardcoded colours or yellow background", () => { + expect(BOLD_MATCH_CLASS).toContain("bg-transparent"); + expect(BOLD_MATCH_CLASS).toContain("text-foreground"); + expect(BOLD_MATCH_CLASS).toContain("font-semibold"); + expect(BOLD_MATCH_CLASS).toContain("underline"); + expect(BOLD_MATCH_CLASS).not.toMatch( + /yellow|#[0-9a-f]{3,6}|\bwhite\b|\bblack\b/i, + ); + expect(PREVIEW_TEXT_CLASS).not.toMatch( + /yellow|#[0-9a-f]{3,6}|\bwhite\b|\bblack\b/i, + ); + }); + + it("applies the bold style to matches via markClassName", () => { + const { container } = render( + match b"]} + fallbackText="" + markClassName={BOLD_MATCH_CLASS} + />, + ); + const mark = container.querySelector("mark"); + expect(mark?.className).toBe(BOLD_MATCH_CLASS); + expect(mark?.className).not.toContain("bg-yellow"); + }); +}); describe("HighlightedText", () => { it("renders fallback text when highlights array is empty", () => { diff --git a/frontend/components/highlighted-text.tsx b/frontend/components/highlighted-text.tsx index c4baf5d3e..a49a069f0 100644 --- a/frontend/components/highlighted-text.tsx +++ b/frontend/components/highlighted-text.tsx @@ -15,13 +15,34 @@ interface HighlightedTextProps { /** Plain text shown when highlights is empty. */ fallbackText: string; className?: string; + /** Replaces the default yellow-background style of matched terms. */ + markClassName?: string; } +const DEFAULT_MARK_CLASS = + "bg-yellow-200 dark:bg-yellow-800 text-foreground rounded-[2px] px-px not-italic"; + +/** + * Bold + underline match style for one-line previews (search suggestions and + * the table). `bg-transparent` is required: a bare gets the browser's + * yellow background, which is unreadable against theme foreground colours. + * Colours come from theme tokens so it follows light and dark mode. + */ +export const BOLD_MATCH_CLASS = + "bg-transparent text-foreground font-semibold underline decoration-1 underline-offset-2 not-italic"; + +/** Surrounding (non-matched) preview text: readable, but below the match. */ +export const PREVIEW_TEXT_CLASS = "text-foreground/70"; + /** * Parse a single fragment string like "foo bar baz" into an * array of React nodes, converting `` elements to styled spans. */ -function parseFragment(fragment: string, keyPrefix: string): React.ReactNode { +function parseFragment( + fragment: string, + keyPrefix: string, + markClassName: string, +): React.ReactNode { // Split on … pairs, keeping the captured groups. const parts = fragment.split(/(.*?<\/mark>)/g); let partId = 0; @@ -31,10 +52,7 @@ function parseFragment(fragment: string, keyPrefix: string): React.ReactNode { if (part.startsWith("") && part.endsWith("")) { const inner = part.slice(6, -7); // strip and return ( - + {inner} ); @@ -47,6 +65,7 @@ export function HighlightedText({ highlights, fallbackText, className, + markClassName = DEFAULT_MARK_CLASS, }: HighlightedTextProps) { if (!highlights || highlights.length === 0) { return {fallbackText}; @@ -65,7 +84,7 @@ export function HighlightedText({ {i > 0 && ( … )} - {parseFragment(fragment, key)} + {parseFragment(fragment, key, markClassName)} ); })} diff --git a/frontend/components/knowledge-search-bar.test.tsx b/frontend/components/knowledge-search-bar.test.tsx new file mode 100644 index 000000000..407e0e02e --- /dev/null +++ b/frontend/components/knowledge-search-bar.test.tsx @@ -0,0 +1,181 @@ +import { act, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { HttpResponse, http } from "msw"; +import { toast } from "sonner"; +import { describe, expect, it, vi } from "vitest"; +import { + type SearchSuggestResult, + SUGGEST_MAX_CHARS, +} from "@/app/api/queries/useGetSearchSuggestions"; +import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; +import { server } from "@/test-utils/msw/server"; +import { renderWithProviders } from "@/test-utils/render"; +import { mockRouter } from "@/test-utils/router"; +import { KnowledgeSearchBar } from "./knowledge-search-bar"; + +const PAYLOAD: SearchSuggestResult = { + phase: "semantic", + query: "docl", + suggestions: [{ text: "docling", files: 2, chunks: 61 }], + files: [ + { + filename: "docling.pdf", + mimetype: "application/pdf", + connector_type: "local", + preview: null, + }, + ], +}; + +function serveSuggest() { + const requests: unknown[] = []; + server.use( + http.post("/api/search/suggest", async ({ request }) => { + requests.push(await request.json()); + return HttpResponse.json(PAYLOAD); + }), + ); + return requests; +} + +/** KnowledgeDropdown, rendered inside the bar, loads these on mount. */ +function serveDropdownDefaults() { + server.use( + http.get("/api/upload_options", () => HttpResponse.json({})), + http.get("/api/connectors", () => HttpResponse.json({ connectors: {} })), + http.get("/api/connectors/:type/defaults", () => HttpResponse.json({})), + ); +} + +function Committed() { + const { queryOverride } = useKnowledgeFilter(); + return {queryOverride}; +} + +async function renderBar() { + serveDropdownDefaults(); + const user = userEvent.setup(); + renderWithProviders( + <> + + + , + { providers: ["knowledgeFilter", "task", "brand"] }, + ); + const input = await screen.findByPlaceholderText("Search knowledge"); + return { user, input }; +} + +describe("KnowledgeSearchBar suggestions", () => { + it("opens suggestions while typing", async () => { + serveSuggest(); + const { user, input } = await renderBar(); + + await user.click(input); + await user.type(input, "docl"); + + expect(input).toHaveAttribute("role", "combobox"); + expect( + await screen.findByRole("option", { name: /docling\.pdf/ }), + ).toBeInTheDocument(); + }); + + it("picking a suggestion fills the input and commits the search", async () => { + serveSuggest(); + const { user, input } = await renderBar(); + await user.click(input); + await user.type(input, "docl"); + + await user.click(await screen.findByRole("option", { name: /61 chunks/ })); + + expect(input).toHaveValue("docling"); + expect(screen.getByTestId("committed")).toHaveTextContent("docling"); + }); + + it("picking a file opens its chunks with the typed query", async () => { + serveSuggest(); + const { user, input } = await renderBar(); + await user.click(input); + await user.type(input, "docl"); + + await user.click( + await screen.findByRole("option", { name: /docling\.pdf/ }), + ); + + expect(mockRouter.push).toHaveBeenCalledWith( + "/knowledge/chunks?filename=docling.pdf&q=docl", + ); + }); + + it("submitting closes the dropdown and commits the trimmed text", async () => { + serveSuggest(); + const { user, input } = await renderBar(); + await user.click(input); + await user.type(input, "docl "); + await screen.findByRole("listbox"); + + await user.keyboard("{Enter}"); + + await waitFor(() => + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(), + ); + expect(screen.getByTestId("committed")).toHaveTextContent(/^docl$/); + }); + + it("clearing the input resets the committed search and stops suggesting", async () => { + serveSuggest(); + const { user, input } = await renderBar(); + await user.click(input); + await user.type(input, "docl"); + await screen.findByRole("listbox"); + + await user.click(screen.getByRole("button", { name: "Clear search" })); + + expect(input).toHaveValue(""); + expect(screen.getByTestId("committed")).toHaveTextContent(/^$/); + await waitFor(() => + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(), + ); + }); + + it("stops accepting characters at the cap and shows an error popup", async () => { + serveSuggest(); + const errorToast = vi.spyOn(toast, "error").mockReturnValue("toast"); + const { user, input } = await renderBar(); + await user.click(input); + + await user.type(input, "a".repeat(SUGGEST_MAX_CHARS)); + expect(errorToast).not.toHaveBeenCalled(); + + await user.type(input, "b"); + + expect(input).toHaveValue("a".repeat(SUGGEST_MAX_CHARS)); + expect(errorToast).toHaveBeenCalledWith( + `Search is limited to ${SUGGEST_MAX_CHARS} characters.`, + expect.anything(), + ); + }); + + it("truncates pasted text to the cap", async () => { + serveSuggest(); + vi.spyOn(toast, "error").mockReturnValue("toast"); + const { user, input } = await renderBar(); + await user.click(input); + + await user.paste("a".repeat(SUGGEST_MAX_CHARS + 25)); + + expect(input).toHaveValue("a".repeat(SUGGEST_MAX_CHARS)); + }); + + it("does not request suggestions until the input has focus", async () => { + const requests = serveSuggest(); + const { input } = await renderBar(); + + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 500)); + }); + + expect(input).toHaveValue(""); + expect(requests).toHaveLength(0); + }); +}); diff --git a/frontend/components/knowledge-search-bar.tsx b/frontend/components/knowledge-search-bar.tsx index ca8cc02c7..ee12cb0f9 100644 --- a/frontend/components/knowledge-search-bar.tsx +++ b/frontend/components/knowledge-search-bar.tsx @@ -1,4 +1,5 @@ import { ArrowRight, Download, RefreshCw, Search, X } from "lucide-react"; +import { useRouter } from "next/navigation"; import { type ChangeEvent, type FormEvent, useCallback, useState } from "react"; import { toast } from "sonner"; import { useRefreshOpenragDocs } from "@/app/api/mutations/useRefreshOpenragDocs"; @@ -10,9 +11,15 @@ import { import { RequirePermission } from "@/components/require-permission"; import { Button } from "@/components/ui/button"; import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; +import { buildChunksUrl } from "@/lib/chunks-url"; +import { limitSearchInput } from "@/lib/limit-search-input"; import { cn } from "@/lib/utils"; import { KnowledgeDropdown } from "./knowledge-dropdown"; import { filterAccentClasses } from "./knowledge-filter-panel"; +import { + SearchSuggestionsPopover, + useSearchSuggestionsBox, +} from "./search-suggestions"; import { SyncConfirmDialog } from "./sync-confirm-dialog"; export const KnowledgeSearchBar = () => { @@ -31,12 +38,28 @@ export const KnowledgeSearchBar = () => { setSearchQueryInput(queryOverride); } + const router = useRouter(); + const suggestions = useSearchSuggestionsBox({ + query: searchQueryInput, + filters: parsedFilterData?.filters, + scoreThreshold: parsedFilterData?.scoreThreshold, + enabled: true, + onPickSuggestion: (text) => { + setSearchQueryInput(text); + setQueryOverride(text); + }, + onPickFile: (filename) => + router.push(buildChunksUrl(filename, searchQueryInput)), + }); + const { dismiss: dismissSuggestions } = suggestions; + const handleSearch = useCallback( (e?: FormEvent) => { if (e) e.preventDefault(); + dismissSuggestions(); setQueryOverride(searchQueryInput.trim()); }, - [searchQueryInput, setQueryOverride], + [searchQueryInput, setQueryOverride, dismissSuggestions], ); const handleReset = useCallback(() => { @@ -93,120 +116,126 @@ export const KnowledgeSearchBar = () => { return (
-
- {!!selectedFilter?.name && ( -
- - {selectedFilter.name} - - -
- )} + + {selectedFilter.name} + + +
+ )} -
- -
+
+ +
-
- ) => - setSearchQueryInput(e.target.value) - } - className="h-full w-full bg-transparent text-sm text-foreground placeholder:text-[hsl(var(--placeholder))] focus:outline-hidden focus:ring-0" - /> - {searchQueryInput && ( - - )} - )} - type="submit" - > - - -
- - + +
- -
- + + + +
+ +
- + ({ + trackSearch: vi.fn(), + trackButton: vi.fn(), +})); + +const PAYLOAD: SearchSuggestResult = { + phase: "semantic", + query: "docl", + suggestions: [{ text: "docling", files: 2, chunks: 61 }], + files: [ + { + filename: "docling.pdf", + mimetype: "application/pdf", + connector_type: "local", + preview: null, + }, + ], +}; + +interface SuggestRequest { + query: string; + semantic: boolean; + scoreThreshold: number; + filters: Record; +} + +function serveSuggest() { + const requests: SuggestRequest[] = []; + server.use( + http.post("/api/search/suggest", async ({ request }) => { + requests.push((await request.json()) as SuggestRequest); + return HttpResponse.json(PAYLOAD); + }), + ); + return requests; +} + +function filterFixture(queryData: object): KnowledgeFilter { + return { + id: "f1", + name: "Reports", + description: "", + query_data: JSON.stringify(queryData), + owner: "u1", + created_at: "", + updated_at: "", + }; +} + +/** Selects a filter and exposes the committed query through the DOM. */ +function FilterSeed({ filter }: { filter?: KnowledgeFilter }) { + const { setSelectedFilter, queryOverride } = useKnowledgeFilter(); + // biome-ignore lint/correctness/useExhaustiveDependencies: seed once on mount + useEffect(() => { + if (filter) setSelectedFilter(filter); + }, []); + return {queryOverride}; +} + +function renderInput( + props: Parameters[0] = {}, + filter?: KnowledgeFilter, +) { + const user = userEvent.setup(); + renderWithProviders( + <> + + + , + { providers: ["knowledgeFilter"] }, + ); + return { + user, + input: screen.getByPlaceholderText("Search your documents..."), + }; +} + +describe("KnowledgeSearchInput suggestions", () => { + it("opens suggestions while typing in uncontrolled mode", async () => { + serveSuggest(); + const { user, input } = renderInput(); + + await user.click(input); + await user.type(input, "docl"); + + expect( + await screen.findByRole("option", { name: /docling\.pdf/ }), + ).toBeInTheDocument(); + }); + + it("picking a suggestion fills the input and commits the search", async () => { + serveSuggest(); + const { user, input } = renderInput(); + await user.click(input); + await user.type(input, "docl"); + + await user.click(await screen.findByRole("option", { name: /61 chunks/ })); + + expect(input).toHaveValue("docling"); + expect(screen.getByTestId("committed")).toHaveTextContent("docling"); + }); + + it("picking a file navigates to its chunks with the typed query", async () => { + serveSuggest(); + const { user, input } = renderInput(); + await user.click(input); + await user.type(input, "docl"); + + await user.click( + await screen.findByRole("option", { name: /docling\.pdf/ }), + ); + + expect(mockRouter.push).toHaveBeenCalledWith( + "/knowledge/chunks?filename=docling.pdf&q=docl", + ); + }); + + it("submitting closes the dropdown and commits the typed text", async () => { + serveSuggest(); + const { user, input } = renderInput(); + await user.click(input); + await user.type(input, "docl"); + await screen.findByRole("listbox"); + + await user.keyboard("{Enter}"); + + await waitFor(() => + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(), + ); + expect(screen.getByTestId("committed")).toHaveTextContent("docl"); + }); + + it("stops accepting characters at the cap and shows an error popup", async () => { + serveSuggest(); + const errorToast = vi.spyOn(toast, "error").mockReturnValue("toast"); + const { user, input } = renderInput(); + await user.click(input); + + await user.type(input, "a".repeat(SUGGEST_MAX_CHARS + 1)); + + expect(input).toHaveValue("a".repeat(SUGGEST_MAX_CHARS)); + expect(errorToast).toHaveBeenCalledWith( + `Search is limited to ${SUGGEST_MAX_CHARS} characters.`, + expect.anything(), + ); + }); + + it("does not suggest in controlled mode, where results filter live", async () => { + const requests = serveSuggest(); + const onSearch = vi.fn(); + const { user, input } = renderInput({ value: "", onSearch }); + + expect(input).not.toHaveAttribute("role", "combobox"); + await user.type(input, "docl"); + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 700)); + }); + + expect(onSearch).toHaveBeenCalled(); + expect(requests).toHaveLength(0); + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(); + }); + + it("sends the selected filter's facets and score threshold", async () => { + const requests = serveSuggest(); + const filter = filterFixture({ + query: "", + filters: { data_sources: ["a.pdf"] }, + scoreThreshold: 2.5, + }); + const { user, input } = renderInput({}, filter); + await screen.findByText("Reports"); + + await user.click(input); + await user.type(input, "docling"); + await waitFor(() => expect(requests.some((r) => r.semantic)).toBe(true)); + + expect(requests.every((r) => r.scoreThreshold === 2.5)).toBe(true); + expect(requests[0].filters).toEqual({ data_sources: ["a.pdf"] }); + }); +}); diff --git a/frontend/components/knowledge-search-input.tsx b/frontend/components/knowledge-search-input.tsx index 08133a6c6..73c0530af 100644 --- a/frontend/components/knowledge-search-input.tsx +++ b/frontend/components/knowledge-search-input.tsx @@ -1,10 +1,17 @@ import { ArrowRight, Search, X } from "lucide-react"; +import { useRouter } from "next/navigation"; import { type ChangeEvent, type FormEvent, useState } from "react"; import { Button } from "@/components/ui/button"; import { useKnowledgeFilter } from "@/contexts/knowledge-filter-context"; import { trackButton } from "@/lib/analytics"; +import { buildChunksUrl } from "@/lib/chunks-url"; +import { limitSearchInput } from "@/lib/limit-search-input"; import { cn } from "@/lib/utils"; import { filterAccentClasses } from "./knowledge-filter-panel"; +import { + SearchSuggestionsPopover, + useSearchSuggestionsBox, +} from "./search-suggestions"; export type KnowledgeSearchInputProps = { /** Controlled value — when set with onSearch, skips global queryOverride. */ @@ -40,6 +47,7 @@ export function KnowledgeSearchInput({ inputClassName, }: KnowledgeSearchInputProps = {}) { const controlled = onSearch != null; + const router = useRouter(); const { selectedFilter, setSelectedFilter, @@ -68,8 +76,23 @@ export function KnowledgeSearchInput({ } }; + // Suggestions are for the knowledge page's commit-on-submit search. Controlled + // inputs (file-scoped panels) already filter live as you type. + const suggestions = useSearchSuggestionsBox({ + query: inputValue, + filters: parsedFilterData?.filters, + scoreThreshold: parsedFilterData?.scoreThreshold, + enabled: !controlled, + onPickSuggestion: (text) => { + setDraft(text); + commitSearch(text); + }, + onPickFile: (filename) => router.push(buildChunksUrl(filename, inputValue)), + }); + const handleSubmit = (e: FormEvent) => { e.preventDefault(); + suggestions.dismiss(); commitSearch(inputValue); }; @@ -93,73 +116,77 @@ export function KnowledgeSearchInput({ )} onSubmit={handleSubmit} > -
- {!hideFilterChip && selectedFilter?.name && ( -
- - {selectedFilter.name} - - setSelectedFilter(null)} - /> -
- )} - - ) => { - const next = e.target.value; - if (controlled) { - // File-scoped panels filter as you type / paste. - onSearch(next); - } else { - setDraft(next); - } - }} - /> - {showClear && ( - - )} - {!hideSubmit && ( - - )} -
+ +
+ {!hideFilterChip && selectedFilter?.name && ( +
+ + {selectedFilter.name} + + setSelectedFilter(null)} + /> +
+ )} + + ) => { + const next = e.target.value; + if (controlled) { + // File-scoped panels filter as you type / paste. + onSearch(next); + } else { + setDraft(limitSearchInput(next)); + } + }} + /> + {showClear && ( + + )} + {!hideSubmit && ( + + )} +
+ ); } diff --git a/frontend/components/search-suggestions.test.tsx b/frontend/components/search-suggestions.test.tsx new file mode 100644 index 000000000..1d4a02cb2 --- /dev/null +++ b/frontend/components/search-suggestions.test.tsx @@ -0,0 +1,358 @@ +import { + act, + fireEvent, + render, + screen, + waitFor, +} from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { HttpResponse, http } from "msw"; +import { useState } from "react"; +import { describe, expect, it, vi } from "vitest"; +import type { SearchSuggestResult } from "@/app/api/queries/useGetSearchSuggestions"; +import { server } from "@/test-utils/msw/server"; +import { createQueryWrapper } from "@/test-utils/render"; +import { + SearchSuggestionsPopover, + useSearchSuggestionsBox, +} from "./search-suggestions"; + +const PAYLOAD: SearchSuggestResult = { + phase: "semantic", + query: "docl", + suggestions: [ + { text: "docling", files: 2, chunks: 61 }, + { text: "docling serve", files: 1, chunks: 1 }, + ], + files: [ + { + filename: "docling.pdf", + mimetype: "application/pdf", + connector_type: "local", + preview: { text: "Run docling serve", kind: "keyword" }, + }, + { + filename: "notes.md", + mimetype: "text/markdown", + connector_type: "local", + preview: { text: "Plain semantic sentence here.", kind: "semantic" }, + }, + ], +}; + +function serve(payload: SearchSuggestResult | { error: string }, status = 200) { + const requests: unknown[] = []; + server.use( + http.post("/api/search/suggest", async ({ request }) => { + requests.push(await request.json()); + return HttpResponse.json(payload, { status }); + }), + ); + return requests; +} + +interface HarnessProps { + initialQuery?: string; + enabled?: boolean; + onPickSuggestion?: (text: string) => void; + onPickFile?: (filename: string) => void; +} + +function Harness({ + initialQuery = "", + enabled = true, + onPickSuggestion = () => {}, + onPickFile = () => {}, +}: HarnessProps) { + const [query, setQuery] = useState(initialQuery); + const box = useSearchSuggestionsBox({ + query, + filters: undefined, + enabled, + onPickSuggestion: (text) => { + setQuery(text); + onPickSuggestion(text); + }, + onPickFile, + }); + return ( + +
+ setQuery(event.target.value)} + {...box.inputProps} + /> + +
+ + ); +} + +function setup(props: HarnessProps = {}) { + const Wrapper = createQueryWrapper(); + const user = userEvent.setup(); + render( + + + , + ); + return { user, input: screen.getByLabelText("search") }; +} + +async function openWithResults(user: ReturnType) { + const input = screen.getByLabelText("search"); + await user.click(input); + await user.type(input, "docl"); + await screen.findByRole("option", { name: /docling serve/ }); + return input; +} + +describe("search suggestions dropdown", () => { + it("stays closed until the input is focused and long enough", async () => { + serve(PAYLOAD); + const { user, input } = setup(); + + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(); + await user.click(input); + await user.type(input, "d"); + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(); + }); + + it("renders suggestions and files in labelled sections", async () => { + serve(PAYLOAD); + const { user } = setup(); + await openWithResults(user); + + expect(screen.getByText("Suggestions")).toBeInTheDocument(); + expect(screen.getByText("Files")).toBeInTheDocument(); + expect(screen.getAllByRole("option")).toHaveLength(4); + expect(screen.getByText("2 files · 61 chunks")).toBeInTheDocument(); + expect(screen.getByText("1 file · 1 chunk")).toBeInTheDocument(); + expect(screen.getByText("Plain semantic sentence here.")).toBeVisible(); + }); + + it("marks the keyword match inside a file preview", async () => { + serve(PAYLOAD); + const { user } = setup(); + await openWithResults(user); + + const mark = document.querySelector("mark"); + expect(mark?.textContent).toBe("docling"); + }); + + it("keeps non-option rows out of the listbox's option set", async () => { + serve(PAYLOAD); + const { user } = setup(); + await openWithResults(user); + + for (const label of ["Suggestions", "Files"]) { + expect(screen.getByText(label)).toHaveAttribute("role", "presentation"); + } + }); + + it("shows a loading row while the first response is pending", async () => { + server.use( + http.post("/api/search/suggest", async () => { + await new Promise((resolve) => setTimeout(resolve, 300)); + return HttpResponse.json(PAYLOAD); + }), + ); + const { user, input } = setup(); + await user.click(input); + await user.type(input, "docl"); + + expect(await screen.findByText("Searching…")).toBeInTheDocument(); + await screen.findByRole("option", { name: /docling serve/ }); + expect(screen.queryByText("Searching…")).not.toBeInTheDocument(); + }); + + it("shows an error row instead of searching forever when the request fails", async () => { + serve({ error: "Suggestions unavailable" }, 500); + const { user, input } = setup(); + await user.click(input); + await user.type(input, "docl"); + + expect( + await screen.findByText( + "Suggestions unavailable. Press Enter to search.", + ), + ).toBeInTheDocument(); + expect(screen.queryByText("Searching…")).not.toBeInTheDocument(); + }); + + it("shows an empty state when there is nothing to suggest", async () => { + serve({ ...PAYLOAD, suggestions: [], files: [] }); + const { user, input } = setup(); + await user.click(input); + await user.type(input, "docl"); + + expect( + await screen.findByText("No suggestions. Press Enter to search."), + ).toBeInTheDocument(); + }); + + it("never requests anything while disabled", async () => { + const requests = serve(PAYLOAD); + const { user, input } = setup({ enabled: false }); + await user.click(input); + await user.type(input, "docling"); + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 700)); + }); + + expect(requests).toHaveLength(0); + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(); + }); +}); + +describe("search suggestions keyboard and pointer", () => { + it("exposes combobox semantics and tracks the active option", async () => { + serve(PAYLOAD); + const { user } = setup(); + const input = await openWithResults(user); + + expect(input).toHaveAttribute("role", "combobox"); + expect(input).toHaveAttribute("aria-expanded", "true"); + expect(input).not.toHaveAttribute("aria-activedescendant"); + + await user.keyboard("{ArrowDown}"); + const [first] = screen.getAllByRole("option"); + expect(first).toHaveAttribute("aria-selected", "true"); + expect(input).toHaveAttribute("aria-activedescendant", first.id); + }); + + it("wraps from the last option back to none, and ArrowUp from none to the last", async () => { + serve(PAYLOAD); + const { user } = setup(); + const input = await openWithResults(user); + const options = screen.getAllByRole("option"); + + await user.keyboard("{ArrowUp}"); + expect(input).toHaveAttribute( + "aria-activedescendant", + options[options.length - 1].id, + ); + + await user.keyboard("{ArrowDown}"); + expect(input).not.toHaveAttribute("aria-activedescendant"); + }); + + it("Enter picks the highlighted suggestion", async () => { + serve(PAYLOAD); + const onPickSuggestion = vi.fn(); + const { user } = setup({ onPickSuggestion }); + await openWithResults(user); + + await user.keyboard("{ArrowDown}{ArrowDown}{Enter}"); + + expect(onPickSuggestion).toHaveBeenCalledWith("docling serve"); + await waitFor(() => + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(), + ); + }); + + it("Enter picks a highlighted file and passes its filename", async () => { + serve(PAYLOAD); + const onPickFile = vi.fn(); + const { user } = setup({ onPickFile }); + await openWithResults(user); + + await user.keyboard("{ArrowDown}{ArrowDown}{ArrowDown}{Enter}"); + + expect(onPickFile).toHaveBeenCalledWith("docling.pdf"); + }); + + it("Enter with no highlighted row is left alone so the form can submit", async () => { + serve(PAYLOAD); + const onPickSuggestion = vi.fn(); + const onPickFile = vi.fn(); + const { user } = setup({ onPickSuggestion, onPickFile }); + await openWithResults(user); + + await user.keyboard("{Enter}"); + + expect(onPickSuggestion).not.toHaveBeenCalled(); + expect(onPickFile).not.toHaveBeenCalled(); + }); + + it("Escape dismisses until the text changes", async () => { + serve(PAYLOAD); + const { user } = setup(); + const input = await openWithResults(user); + + await user.keyboard("{Escape}"); + await waitFor(() => + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(), + ); + expect(input).toHaveAttribute("aria-expanded", "false"); + + await user.type(input, "i"); + expect(await screen.findByRole("listbox")).toBeInTheDocument(); + }); + + it("the dismiss action closes the dropdown", async () => { + serve(PAYLOAD); + const { user } = setup(); + await openWithResults(user); + + await user.click(screen.getByRole("button", { name: "dismiss" })); + + await waitFor(() => + expect(screen.queryByRole("listbox")).not.toBeInTheDocument(), + ); + }); + + it("clicking a row picks it without taking focus from the input", async () => { + serve(PAYLOAD); + const onPickFile = vi.fn(); + const { user } = setup({ onPickFile }); + const input = await openWithResults(user); + + await user.click(screen.getByRole("option", { name: /notes\.md/ })); + + expect(onPickFile).toHaveBeenCalledWith("notes.md"); + expect(input).toHaveFocus(); + }); + + it("hovering a row makes it the active option", async () => { + serve(PAYLOAD); + const { user } = setup(); + await openWithResults(user); + + const option = screen.getAllByRole("option")[1]; + await user.hover(option); + + expect(option).toHaveAttribute("aria-selected", "true"); + }); + + it("ignores navigation keys while an IME composition is in progress", async () => { + serve(PAYLOAD); + const { user } = setup(); + const input = await openWithResults(user); + + fireEvent.keyDown(input, { key: "ArrowDown", isComposing: true }); + expect(input).not.toHaveAttribute("aria-activedescendant"); + + fireEvent.keyDown(input, { key: "ArrowDown", isComposing: false }); + expect(input).toHaveAttribute("aria-activedescendant"); + }); + + it("resets the active row when new results arrive", async () => { + serve(PAYLOAD); + const { user } = setup(); + const input = await openWithResults(user); + await user.keyboard("{ArrowDown}"); + expect(input).toHaveAttribute("aria-activedescendant"); + + serve({ ...PAYLOAD, suggestions: [PAYLOAD.suggestions[0]], files: [] }); + await user.type(input, "i"); + + await waitFor(() => + expect(input).not.toHaveAttribute("aria-activedescendant"), + ); + }); +}); diff --git a/frontend/components/search-suggestions.tsx b/frontend/components/search-suggestions.tsx new file mode 100644 index 000000000..864ffb52b --- /dev/null +++ b/frontend/components/search-suggestions.tsx @@ -0,0 +1,358 @@ +import { FileIcon, Info, Search } from "lucide-react"; +import { + type KeyboardEvent, + type ReactElement, + useId, + useMemo, + useState, +} from "react"; +import { + type SearchSuggestResult, + SUGGEST_MIN_CHARS, + type SuggestFile, + useSearchSuggestions, +} from "@/app/api/queries/useGetSearchSuggestions"; +import { + BOLD_MATCH_CLASS, + HighlightedText, + PREVIEW_TEXT_CLASS, +} from "@/components/highlighted-text"; +import { + Popover, + PopoverAnchor, + PopoverContent, +} from "@/components/ui/popover"; +import type { FilterInput } from "@/lib/filter-normalization"; +import { splitOnQuery } from "@/lib/search-preview"; +import { cn } from "@/lib/utils"; + +export type SuggestionItem = + | { + kind: "suggestion"; + key: string; + text: string; + files: number; + chunks: number; + } + | { kind: "file"; key: string; file: SuggestFile }; + +function buildSuggestionItems( + data: SearchSuggestResult | undefined, +): SuggestionItem[] { + if (!data) return []; + return [ + ...data.suggestions.map( + (suggestion): SuggestionItem => ({ + kind: "suggestion", + key: `s:${suggestion.text}`, + ...suggestion, + }), + ), + ...data.files.map( + (file): SuggestionItem => ({ + kind: "file", + key: `f:${file.filename}`, + file, + }), + ), + ]; +} + +interface UseSearchSuggestionsBoxOptions { + query: string; + filters: FilterInput | undefined; + enabled: boolean; + scoreThreshold?: number; + onPickSuggestion: (text: string) => void; + onPickFile: (filename: string) => void; +} + +/** + * State and keyboard handling for the search suggestions dropdown. + * + * Arrow keys move through the rows, Enter opens the highlighted row (and falls + * through to a normal search when none is highlighted), Escape dismisses. + */ +export function useSearchSuggestionsBox({ + query, + filters, + enabled, + scoreThreshold, + onPickSuggestion, + onPickFile, +}: UseSearchSuggestionsBoxOptions) { + const listId = useId(); + const [focused, setFocused] = useState(false); + const [dismissedQuery, setDismissedQuery] = useState(null); + const [activeIndex, setActiveIndex] = useState(-1); + + const { data, isRefining, isError } = useSearchSuggestions( + query, + filters, + enabled && focused, + scoreThreshold, + ); + const items = useMemo(() => buildSuggestionItems(data), [data]); + + const [prevItems, setPrevItems] = useState(items); + if (prevItems !== items) { + setPrevItems(items); + setActiveIndex(-1); + } + + // Compared trimmed: submitting commits the trimmed text, which would + // otherwise differ from the dismissed value and reopen the dropdown. + const trimmedQuery = query.trim(); + const open = + enabled && + focused && + trimmedQuery.length >= SUGGEST_MIN_CHARS && + dismissedQuery !== trimmedQuery; + + const pick = (item: SuggestionItem) => { + if (item.kind === "suggestion") { + setDismissedQuery(item.text); + onPickSuggestion(item.text); + } else { + setDismissedQuery(trimmedQuery); + onPickFile(item.file.filename); + } + }; + + const handleKeyDown = (event: KeyboardEvent) => { + if (!open || event.nativeEvent.isComposing) return; + const count = items.length; + + if (event.key === "ArrowDown" && count > 0) { + event.preventDefault(); + setActiveIndex((index) => (index + 1 >= count ? -1 : index + 1)); + } else if (event.key === "ArrowUp" && count > 0) { + event.preventDefault(); + setActiveIndex((index) => (index <= -1 ? count - 1 : index - 1)); + } else if (event.key === "Enter" && items[activeIndex]) { + event.preventDefault(); + pick(items[activeIndex]); + } else if (event.key === "Escape") { + event.preventDefault(); + setDismissedQuery(trimmedQuery); + } + }; + + const optionId = (index: number) => `${listId}-option-${index}`; + + return { + open, + dismiss: () => setDismissedQuery(trimmedQuery), + data, + items, + isRefining, + isError, + activeIndex, + setActiveIndex, + pick, + listId, + optionId, + onOpenChange: (next: boolean) => { + if (!next) setDismissedQuery(trimmedQuery); + }, + anchorProps: { "data-search-suggest-anchor": "" }, + inputProps: { + role: "combobox" as const, + "aria-autocomplete": "list" as const, + "aria-expanded": open, + "aria-controls": listId, + "aria-activedescendant": + open && activeIndex >= 0 ? optionId(activeIndex) : undefined, + autoComplete: "off", + onFocus: () => setFocused(true), + onBlur: () => setFocused(false), + onKeyDown: handleKeyDown, + }, + }; +} + +type SearchSuggestionsBox = ReturnType; + +function MatchedText({ text, query }: { text: string; query: string }) { + const { before, match, after } = splitOnQuery(text, query); + return ( + + {before} + {match && {match}} + {after} + + ); +} + +function countLabel(count: number, noun: string) { + return `${count} ${noun}${count === 1 ? "" : "s"}`; +} + +function SectionLabel({ children }: { children: string }) { + return ( +
+ {children} +
+ ); +} + +function SuggestionRows({ + box, + query, +}: { + box: SearchSuggestionsBox; + query: string; +}) { + const suggestionCount = box.items.filter( + (item) => item.kind === "suggestion", + ).length; + const rows = box.items.map((item, index) => { + const active = index === box.activeIndex; + const base = cn( + "flex w-full items-start gap-3 px-3 py-2 text-left text-sm", + active ? "bg-accent text-accent-foreground" : "hover:bg-accent/60", + ); + const handlers = { + id: box.optionId(index), + role: "option" as const, + "aria-selected": active, + type: "button" as const, + // Keep focus in the input so typing continues after a click. + onMouseDown: (event: React.MouseEvent) => event.preventDefault(), + onMouseEnter: () => box.setActiveIndex(index), + onClick: () => box.pick(item), + }; + + if (item.kind === "suggestion") { + return ( + + ); + } + + const { file } = item; + return ( + + ); + }); + + return ( + <> + {suggestionCount > 0 && Suggestions} + {rows.slice(0, suggestionCount)} + {rows.length > suggestionCount && Files} + {rows.slice(suggestionCount)} + + ); +} + +interface SearchSuggestionsPopoverProps { + box: SearchSuggestionsBox; + query: string; + /** The single element the dropdown anchors to (spread `box.anchorProps` on it). */ + children: ReactElement; +} + +export function SearchSuggestionsPopover({ + box, + query, + children, +}: SearchSuggestionsPopoverProps) { + const trimmed = query.trim(); + const empty = box.data && box.items.length === 0; + + return ( + + {children} + event.preventDefault()} + onCloseAutoFocus={(event) => event.preventDefault()} + // Radix reports clicks on the input as "outside"; ignore them so the + // dropdown does not flicker closed while the user is typing. + onInteractOutside={(event) => { + const target = event.target as HTMLElement | null; + if (target?.closest("[data-search-suggest-anchor]")) { + event.preventDefault(); + } + }} + > +
+ {!box.data && !box.isError && ( +
+ Searching… +
+ )} + {box.isError && ( +
+ Suggestions unavailable. Press Enter to search. +
+ )} + {empty && ( +
+ No suggestions. Press Enter to search. +
+ )} + +
+
+ + + Searching file names and content ·{" "} + {box.isRefining + ? "Refining with meaning…" + : "Hybrid (keyword + meaning)"} + +
+ + + ); +} diff --git a/frontend/hooks/use-debounced-value.test.ts b/frontend/hooks/use-debounced-value.test.ts new file mode 100644 index 000000000..6b29825cc --- /dev/null +++ b/frontend/hooks/use-debounced-value.test.ts @@ -0,0 +1,88 @@ +import { act, renderHook } from "@testing-library/react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { useDebouncedValue } from "./use-debounced-value"; + +describe("useDebouncedValue", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("returns the initial value immediately", () => { + const { result } = renderHook(() => useDebouncedValue("a", 100)); + expect(result.current).toBe("a"); + }); + + it("applies a change only after the delay has elapsed", () => { + const { result, rerender } = renderHook( + ({ value }) => useDebouncedValue(value, 100), + { initialProps: { value: "a" } }, + ); + + rerender({ value: "b" }); + expect(result.current).toBe("a"); + + act(() => { + vi.advanceTimersByTime(99); + }); + expect(result.current).toBe("a"); + + act(() => { + vi.advanceTimersByTime(1); + }); + expect(result.current).toBe("b"); + }); + + it("restarts the wait on every change so a burst applies only the last value", () => { + const { result, rerender } = renderHook( + ({ value }) => useDebouncedValue(value, 100), + { initialProps: { value: "a" } }, + ); + + for (const value of ["b", "c", "d"]) { + rerender({ value }); + act(() => { + vi.advanceTimersByTime(60); + }); + } + expect(result.current).toBe("a"); + + act(() => { + vi.advanceTimersByTime(100); + }); + expect(result.current).toBe("d"); + }); + + it("applies the value on the next tick when the delay is 0", () => { + const { result, rerender } = renderHook( + ({ value }) => useDebouncedValue(value, 0), + { initialProps: { value: "a" } }, + ); + + rerender({ value: "b" }); + act(() => { + vi.advanceTimersByTime(0); + }); + expect(result.current).toBe("b"); + }); + + it("does not update after unmount", () => { + const { result, rerender, unmount } = renderHook( + ({ value }) => useDebouncedValue(value, 100), + { initialProps: { value: "a" } }, + ); + + rerender({ value: "b" }); + unmount(); + + expect(() => { + act(() => { + vi.advanceTimersByTime(200); + }); + }).not.toThrow(); + expect(result.current).toBe("a"); + }); +}); diff --git a/frontend/hooks/use-debounced-value.ts b/frontend/hooks/use-debounced-value.ts new file mode 100644 index 000000000..12e809451 --- /dev/null +++ b/frontend/hooks/use-debounced-value.ts @@ -0,0 +1,20 @@ +import { useEffect, useState } from "react"; + +/** + * Returns `value` after it has stopped changing for `delayMs`. + * + * A debounce restarts on every change, so it fires on the first pause longer + * than `delayMs` rather than waiting for a whole word. A `delayMs` of 0 applies + * the value on the next tick, which callers use to skip the wait at word + * boundaries. + */ +export function useDebouncedValue(value: T, delayMs: number): T { + const [debounced, setDebounced] = useState(value); + + useEffect(() => { + const timer = setTimeout(() => setDebounced(value), delayMs); + return () => clearTimeout(timer); + }, [value, delayMs]); + + return debounced; +} diff --git a/frontend/lib/chunks-url.test.ts b/frontend/lib/chunks-url.test.ts new file mode 100644 index 000000000..aa1ab04aa --- /dev/null +++ b/frontend/lib/chunks-url.test.ts @@ -0,0 +1,29 @@ +import { describe, expect, it } from "vitest"; +import { buildChunksUrl } from "./chunks-url"; + +describe("buildChunksUrl", () => { + it("adds the query when there is one", () => { + expect(buildChunksUrl("report.pdf", "quarterly revenue")).toBe( + "/knowledge/chunks?filename=report.pdf&q=quarterly+revenue", + ); + }); + + it("trims the query before adding it", () => { + expect(buildChunksUrl("report.pdf", " docling ")).toBe( + "/knowledge/chunks?filename=report.pdf&q=docling", + ); + }); + + it.each(["", " ", "*", " * "])("omits q for %j", (query) => { + expect(buildChunksUrl("report.pdf", query)).toBe( + "/knowledge/chunks?filename=report.pdf", + ); + }); + + it("encodes characters that would break the URL", () => { + const url = buildChunksUrl("my doc & notes.pdf", "cost/benefit"); + const params = new URL(url, "http://localhost").searchParams; + expect(params.get("filename")).toBe("my doc & notes.pdf"); + expect(params.get("q")).toBe("cost/benefit"); + }); +}); diff --git a/frontend/lib/chunks-url.ts b/frontend/lib/chunks-url.ts new file mode 100644 index 000000000..c2d14393e --- /dev/null +++ b/frontend/lib/chunks-url.ts @@ -0,0 +1,11 @@ +export function buildChunksUrl( + filename: string, + effectiveSearchText: string, +): string { + const params = new URLSearchParams({ filename }); + const trimmed = effectiveSearchText.trim(); + if (trimmed && trimmed !== "*") { + params.set("q", trimmed); + } + return `/knowledge/chunks?${params.toString()}`; +} diff --git a/frontend/lib/limit-search-input.ts b/frontend/lib/limit-search-input.ts new file mode 100644 index 000000000..afec09fc3 --- /dev/null +++ b/frontend/lib/limit-search-input.ts @@ -0,0 +1,16 @@ +import { toast } from "sonner"; +import { SUGGEST_MAX_CHARS } from "@/app/api/queries/useGetSearchSuggestions"; + +/** + * Clamp typed or pasted search text to the suggestion length cap and tell the + * user when input was dropped. The fixed toast id keeps held keys from stacking + * popups. + */ +export function limitSearchInput(value: string): string { + const characters = [...value]; + if (characters.length <= SUGGEST_MAX_CHARS) return value; + toast.error(`Search is limited to ${SUGGEST_MAX_CHARS} characters.`, { + id: "search-query-limit", + }); + return characters.slice(0, SUGGEST_MAX_CHARS).join(""); +} diff --git a/frontend/lib/search-preview.test.ts b/frontend/lib/search-preview.test.ts new file mode 100644 index 000000000..d84e0889b --- /dev/null +++ b/frontend/lib/search-preview.test.ts @@ -0,0 +1,199 @@ +import { describe, expect, it } from "vitest"; +import type { ChunkResult, File } from "@/app/api/queries/useGetSearchQuery"; +import { + bestSentence, + buildFilePreview, + computeTopPreviews, + PREVIEW_MAX_CHARS, + splitOnQuery, + stripMarks, +} from "./search-preview"; + +function chunk(overrides: Partial = {}): ChunkResult { + return { + filename: "a.pdf", + mimetype: "application/pdf", + page: 1, + text: "body", + score: 1, + highlights: [], + ...overrides, + }; +} + +function file(name: string, chunks: ChunkResult[]): File { + return { + filename: name, + mimetype: "application/pdf", + source_url: "", + size: 0, + connector_type: "local", + chunks, + }; +} + +describe("buildFilePreview", () => { + it("uses the highlight fragment and marks where text was cut away", () => { + const preview = buildFilePreview( + file("a.pdf", [ + chunk({ + text: "Intro. The docling library parses PDFs. More follows.", + highlights: ["The docling library parses PDFs"], + }), + ]), + "docling", + ); + expect(preview?.kind).toBe("keyword"); + expect(preview?.text).toBe( + "…The docling library parses PDFs…", + ); + }); + + it("adds no ellipsis when the fragment is the whole chunk", () => { + const preview = buildFilePreview( + file("a.pdf", [ + chunk({ + text: "docling parses PDFs", + highlights: ["docling parses PDFs"], + }), + ]), + "docling", + ); + expect(preview?.text).toBe("docling parses PDFs"); + }); + + it("keeps the match and balanced marks when trimming a long fragment", () => { + const filler = "lorem ipsum ".repeat(40); + const fragment = `${filler}needle ${filler}`; + const preview = buildFilePreview( + file("a.pdf", [chunk({ text: fragment, highlights: [fragment] })]), + "needle", + ); + expect(preview?.text).toContain("needle"); + expect(stripMarks(preview?.text ?? "").length).toBeLessThanOrEqual( + PREVIEW_MAX_CHARS + 2, + ); + expect(preview?.text.match(//g)).toHaveLength(1); + expect(preview?.text.match(/<\/mark>/g)).toHaveLength(1); + }); + + it("prefers the highest-scoring chunk that has a literal match", () => { + const preview = buildFilePreview( + file("a.pdf", [ + chunk({ score: 1, highlights: ["low hit"] }), + chunk({ score: 9, highlights: ["high hit"] }), + ]), + "hit", + ); + expect(preview?.text).toContain("high"); + }); + + it("falls back to the best semantic sentence when nothing matched literally", () => { + const preview = buildFilePreview( + file("a.pdf", [ + chunk({ + text: "The weather was pleasant all week. Vector search ranks passages by similarity. Lunch was at noon daily.", + }), + ]), + "passages similarity", + ); + expect(preview).toEqual({ + kind: "semantic", + text: "Vector search ranks passages by similarity.", + }); + }); + + it("returns null when the file has no usable text", () => { + expect(buildFilePreview(file("a.pdf", []), "x")).toBeNull(); + expect( + buildFilePreview(file("a.pdf", [chunk({ text: " " })]), "x"), + ).toBeNull(); + }); +}); + +describe("bestSentence", () => { + it("picks the first sentence when none share words with the query", () => { + expect( + bestSentence( + "First sentence is right here. Second one follows it.", + "zzz", + ), + ).toBe("First sentence is right here."); + }); + + it("matches accented and non-Latin words the way the backend does", () => { + expect( + bestSentence( + "Nothing relevant lives in this first sentence. La cafétéria déjà ouverte sert du café.", + "café", + ), + ).toBe("La cafétéria déjà ouverte sert du café."); + expect( + bestSentence( + "Nothing relevant lives in this first sentence. 全文 検索 エンジン の 設計 と 実装 について 説明 する。", + "エンジン", + ), + ).toContain("エンジン"); + }); + + it("returns null for blank text", () => { + expect(bestSentence(" ", "q")).toBeNull(); + }); + + it("truncates a very long sentence with an ellipsis", () => { + const long = `${"word ".repeat(100)}end`; + const result = bestSentence(long, "word") ?? ""; + expect(result.endsWith("…")).toBe(true); + expect(result.length).toBeLessThanOrEqual(PREVIEW_MAX_CHARS + 1); + }); +}); + +describe("computeTopPreviews", () => { + const files = [ + file("low.pdf", [chunk({ score: 1, highlights: ["q low"] })]), + file("top.pdf", [chunk({ score: 9, highlights: ["q top"] })]), + file("mid.pdf", [chunk({ score: 5, highlights: ["q mid"] })]), + file("last.pdf", [chunk({ score: 0.5, highlights: ["q x"] })]), + ]; + + it("previews only the three most relevant files", () => { + const previews = computeTopPreviews(files, "q"); + expect(previews.size).toBe(3); + const texts = [...previews.values()].map((p) => p.text); + expect(texts.some((text) => text.includes("top"))).toBe(true); + expect(texts.some((text) => text.includes("mid"))).toBe(true); + expect(texts.some((text) => text.includes("low"))).toBe(true); + // The weakest file is the one left out. + expect(texts.some((text) => text.endsWith(" x"))).toBe(false); + }); + + it("does not mutate the input order", () => { + const names = files.map((f) => f.filename); + computeTopPreviews(files, "q"); + expect(files.map((f) => f.filename)).toEqual(names); + }); + + it("returns nothing for empty and wildcard queries", () => { + expect(computeTopPreviews(files, "").size).toBe(0); + expect(computeTopPreviews(files, " * ").size).toBe(0); + }); +}); + +describe("splitOnQuery", () => { + it("splits around the first case-insensitive match", () => { + expect(splitOnQuery("My Docling.pdf", "docling")).toEqual({ + before: "My ", + match: "Docling", + after: ".pdf", + }); + }); + + it("returns the whole text as `before` when there is no match", () => { + expect(splitOnQuery("report.pdf", "zzz")).toEqual({ + before: "report.pdf", + match: "", + after: "", + }); + expect(splitOnQuery("report.pdf", " ").before).toBe("report.pdf"); + }); +}); diff --git a/frontend/lib/search-preview.ts b/frontend/lib/search-preview.ts new file mode 100644 index 000000000..de12c25c1 --- /dev/null +++ b/frontend/lib/search-preview.ts @@ -0,0 +1,192 @@ +/** + * One-line "why does this file match?" previews for search results. + * + * A `keyword` preview is an OpenSearch highlight fragment (matched terms wrapped + * in ``), trimmed to one line with `…` where text was cut away. A + * `semantic` preview is the sentence of the best chunk that shares the most + * words with the query, shown when nothing matched literally. + */ +import type { ChunkResult, File } from "@/app/api/queries/useGetSearchQuery"; +import { getKnowledgeFileIdentity } from "@/lib/knowledge-table-state"; + +export const PREVIEW_MAX_CHARS = 180; +/** How many of the most relevant files get a preview line in the table. */ +const TABLE_PREVIEW_FILE_COUNT = 3; + +export interface SearchPreview { + /** One line; `keyword` previews wrap matched terms in ``. */ + text: string; + kind: "keyword" | "semantic"; +} + +const MARK_LENGTH = "".length; +const MARK_SPLIT = /(.*?<\/mark>)/; + +export function stripMarks(fragment: string): string { + return fragment.replace(/<\/?mark>/g, ""); +} + +function isMarkPart(part: string): boolean { + return part.startsWith("") && part.endsWith(""); +} + +/** Keep at most `limit` visible characters without splitting a ``. */ +function takePlain(source: string, limit: number): string { + let out = ""; + let used = 0; + for (const part of source.split(MARK_SPLIT)) { + if (!part) continue; + const plainLength = isMarkPart(part) + ? part.length - MARK_LENGTH + : part.length; + if (used + plainLength <= limit) { + out += part; + used += plainLength; + continue; + } + if (!isMarkPart(part)) out += part.slice(0, Math.max(limit - used, 0)); + break; + } + return out; +} + +function truncatePlain(text: string, limit: number): string { + if (text.length <= limit) return text; + const cut = text + .slice(0, limit) + .replace(/\s+\S*$/, "") + .replace(/[\s,;:-]+$/, ""); + return `${cut || text.slice(0, limit)}…`; +} + +/** Shrink an over-long fragment around its first ``, ending with `…`. */ +function clipFragment(fragment: string): string { + const flat = fragment.replace(/\s+/g, " ").trim(); + if (stripMarks(flat).length <= PREVIEW_MAX_CHARS) return flat; + + const first = /.*?<\/mark>/.exec(flat); + if (!first) return truncatePlain(flat, PREVIEW_MAX_CHARS); + + const before = flat.slice(0, first.index); + const match = first[0]; + const after = flat.slice(first.index + match.length); + + const budget = Math.max(PREVIEW_MAX_CHARS - (match.length - MARK_LENGTH), 0); + const keepBefore = Math.floor(budget / 3); + const keepAfter = budget - Math.min(keepBefore, before.length); + + const clippedBefore = + before.length > keepBefore + ? `…${before.slice(before.length - keepBefore).replace(/^\S*\s+/, "")}` + : before; + + const kept = takePlain(after, keepAfter); + const wasCut = stripMarks(kept).length < stripMarks(after).length; + const clippedAfter = wasCut + ? `${kept.replace(/\s+[^\s<>]*$/, "").replace(/[\s,;:-]+$/, "")}…` + : kept; + + return `${clippedBefore}${match}${clippedAfter}`; +} + +/** The sentence sharing the most words with the query (first one on ties). */ +export function bestSentence(text: string, query: string): string | null { + const cleaned = text.trim(); + if (!cleaned) return null; + + const sentences = cleaned + .split(/(?<=[.!?])\s+|\n{2,}/) + .map((sentence) => sentence.trim()) + .filter((sentence) => sentence.length >= 20); + const candidates = sentences.length > 0 ? sentences : [cleaned]; + + const words = (value: string) => + value.toLowerCase().match(/[\p{L}\p{N}_]+/gu) ?? []; + const queryWords = new Set(words(query)); + let best = candidates[0]; + let bestScore = -1; + for (const sentence of candidates) { + const score = words(sentence).filter((word) => queryWords.has(word)).length; + if (score > bestScore) { + best = sentence; + bestScore = score; + } + } + return truncatePlain(best.replace(/\s+/g, " "), PREVIEW_MAX_CHARS); +} + +function previewFromChunk(chunk: ChunkResult): SearchPreview | null { + const fragment = chunk.highlights?.find((item) => item.trim().length > 0); + if (!fragment) return null; + + let line = clipFragment(fragment); + const plain = stripMarks(fragment); + const start = chunk.text ? chunk.text.indexOf(plain) : -1; + if (start > 0 && !line.startsWith("…")) line = `…${line}`; + if ( + start >= 0 && + start + plain.length < chunk.text.length && + !line.endsWith("…") + ) { + line = `${line}…`; + } + return { text: line, kind: "keyword" }; +} + +function bestChunkScore(file: File): number { + return Math.max(...(file.chunks ?? []).map((chunk) => chunk.score), 0); +} + +/** Preview for one file: a literal match wins, else the best semantic sentence. */ +export function buildFilePreview( + file: File, + query: string, +): SearchPreview | null { + const chunks = [...(file.chunks ?? [])].sort((a, b) => b.score - a.score); + for (const chunk of chunks) { + const preview = previewFromChunk(chunk); + if (preview) return preview; + } + const top = chunks.find((chunk) => chunk.text?.trim()); + const sentence = top ? bestSentence(top.text, query) : null; + return sentence ? { text: sentence, kind: "semantic" } : null; +} + +/** + * Previews for the `count` most relevant files, keyed by file identity. + * Relevance is the best chunk score, so a file with one excellent chunk + * outranks one with many mediocre chunks. + */ +export function computeTopPreviews( + files: File[], + query: string, + count = TABLE_PREVIEW_FILE_COUNT, +): Map { + const previews = new Map(); + if (!query.trim() || query.trim() === "*") return previews; + + const ranked = files + .filter((file) => (file.chunks?.length ?? 0) > 0) + .sort((a, b) => bestChunkScore(b) - bestChunkScore(a)) + .slice(0, count); + for (const file of ranked) { + const preview = buildFilePreview(file, query); + if (preview) previews.set(getKnowledgeFileIdentity(file), preview); + } + return previews; +} + +/** Split `text` around the first case-insensitive occurrence of `query`. */ +export function splitOnQuery( + text: string, + query: string, +): { before: string; match: string; after: string } { + const needle = query.trim().toLowerCase(); + const index = needle ? text.toLowerCase().indexOf(needle) : -1; + if (index < 0) return { before: text, match: "", after: "" }; + return { + before: text.slice(0, index), + match: text.slice(index, index + needle.length), + after: text.slice(index + needle.length), + }; +} diff --git a/src/api/search.py b/src/api/search.py index 4ebdea764..2914e91fc 100644 --- a/src/api/search.py +++ b/src/api/search.py @@ -1,31 +1,56 @@ -from typing import Any, Dict +from typing import Any from fastapi import Depends -from pydantic import BaseModel, Field from fastapi.responses import JSONResponse -from utils.logging_config import get_logger -from utils.opensearch_utils import OpenSearchDiskSpaceError, DISK_SPACE_ERROR_MESSAGE +from pydantic import BaseModel, Field, field_validator from dependencies import ( get_search_service, get_session_manager, - get_current_user, require_permission, ) from session_manager import User +from utils.logging_config import get_logger +from utils.opensearch_utils import DISK_SPACE_ERROR_MESSAGE, OpenSearchDiskSpaceError logger = get_logger(__name__) +SUGGEST_MAX_FILTER_KEYS = 10 +SUGGEST_MAX_FILTER_VALUES = 100 + class SearchBody(BaseModel): query: str - filters: Dict[str, Any] = Field(default_factory=dict) + filters: dict[str, Any] = Field(default_factory=dict) limit: int = 10 scoreThreshold: float = Field(default=0, alias="scoreThreshold") model_config = {"populate_by_name": True} +class SuggestBody(BaseModel): + """Search-as-you-type request. Bounds are enforced server-side.""" + + query: str = Field(max_length=200) + filters: dict[str, Any] = Field(default_factory=dict) + # False: cheap keyword phase for every debounced keystroke. + # True: also run the semantic (embedding) search. + semantic: bool = False + scoreThreshold: float = Field(default=0, ge=0, le=100) + + model_config = {"populate_by_name": True} + + @field_validator("filters") + @classmethod + def _bound_filters(cls, filters: dict[str, Any]) -> dict[str, Any]: + if len(filters) > SUGGEST_MAX_FILTER_KEYS: + raise ValueError("too many filters") + for values in filters.values(): + if isinstance(values, list) and len(values) > SUGGEST_MAX_FILTER_VALUES: + raise ValueError("too many filter values") + return filters + + async def search( body: SearchBody, search_service=Depends(get_search_service), @@ -59,10 +84,34 @@ async def search( return JSONResponse({"error": DISK_SPACE_ERROR_MESSAGE}, status_code=507) except Exception as e: error_msg = str(e) - if ( - "AuthenticationException" in error_msg - or "access denied" in error_msg.lower() - ): + if "AuthenticationException" in error_msg or "access denied" in error_msg.lower(): return JSONResponse({"error": error_msg}, status_code=403) else: return JSONResponse({"error": error_msg}, status_code=500) + + +async def suggest( + body: SuggestBody, + search_service=Depends(get_search_service), + user: User = Depends(require_permission("search:use")), +): + """Query completions, matching files and per-file previews.""" + try: + result = await search_service.suggest( + body.query, + user_id=user.user_id, + jwt_token=user.jwt_token, + filters=body.filters, + semantic=body.semantic, + score_threshold=body.scoreThreshold, + ) + if result.get("error"): + logger.error("Search suggest failed", user_id=user.user_id, error=result["error"]) + return JSONResponse({"error": "Suggestions unavailable"}, status_code=502) + return JSONResponse(result, status_code=200) + except OpenSearchDiskSpaceError: + return JSONResponse({"error": DISK_SPACE_ERROR_MESSAGE}, status_code=507) + except Exception as e: + # Detailed error stays server-side; clients get a generic message. + logger.error("Search suggest failed", user_id=user.user_id, error=str(e)) + return JSONResponse({"error": "Suggestions unavailable"}, status_code=500) diff --git a/src/app/routes/internal.py b/src/app/routes/internal.py index efe5167d1..f6e67a0eb 100644 --- a/src/app/routes/internal.py +++ b/src/app/routes/internal.py @@ -157,6 +157,7 @@ def register_internal_routes(app: FastAPI): # Search endpoint app.add_api_route("/search", search.search, methods=["POST"], tags=["internal"]) + app.add_api_route("/search/suggest", search.suggest, methods=["POST"], tags=["internal"]) # File listing/search endpoints app.add_api_route("/files", files.list_files, methods=["GET"], tags=["internal"]) diff --git a/src/connectors/service.py b/src/connectors/service.py index a2be1e9cc..1e8d944d1 100644 --- a/src/connectors/service.py +++ b/src/connectors/service.py @@ -3,6 +3,7 @@ from opensearchpy.exceptions import NotFoundError from config.settings import get_index_name +from services.search_cache import invalidate_corpus_caches from utils.file_utils import clean_connector_filename from utils.logging_config import get_logger @@ -207,6 +208,7 @@ async def _update_connector_metadata( }, }, ) + invalidate_corpus_caches() logger.debug(f"Updated metadata for document {document.id}") except Exception as e: # A missing index means the chunks aren't where this write expects diff --git a/src/models/processors.py b/src/models/processors.py index cf1968d32..7c647b8f1 100644 --- a/src/models/processors.py +++ b/src/models/processors.py @@ -6,6 +6,7 @@ from config.settings import clients, get_embedding_model, get_index_name, get_openrag_config from observability.spans import trace_root, trace_span +from services.search_cache import invalidate_corpus_caches from session_manager import AnonymousUser from utils.document_processing import ( extract_relevant, @@ -1239,6 +1240,7 @@ async def _reconcile_shared_owner(self, filename: str, shared: bool) -> None: }, }, ) + invalidate_corpus_caches() except Exception as e: logger.warning( "Failed to reconcile owner fields for skipped duplicate", diff --git a/src/services/dls_principal_service.py b/src/services/dls_principal_service.py index 692de1c83..e3ddf8b87 100644 --- a/src/services/dls_principal_service.py +++ b/src/services/dls_principal_service.py @@ -7,6 +7,7 @@ from datetime import UTC, datetime from typing import Any +from services.search_cache import invalidate_user_spaces, invalidate_user_suggestions from session_manager import User from utils.group_acl import ( acl_principal_label, @@ -44,6 +45,7 @@ def __init__( self._ensure_lock = asyncio.Lock() self._index_checked = False self._cache: dict[str, tuple[float, list[str]]] = {} + self._last_principals: dict[str, frozenset[str]] = {} self._locks: dict[str, asyncio.Lock] = {} async def refresh_user_principals( @@ -146,6 +148,7 @@ async def _refresh_user_principals_uncached( refresh="wait_for", ) self._cache_principals(user, group_roles, principals) + self._drop_stale_suggestions(user, principals) except Exception as e: logger.warning( "Failed to refresh DLS principal lookup row", @@ -157,6 +160,15 @@ async def _refresh_user_principals_uncached( return principals + def _drop_stale_suggestions(self, user: User, principals: list[str]) -> None: + """Cached suggestions and spaces embed what the user could see; membership changes make them stale.""" + current = frozenset(principals) + previous = self._last_principals.get(user.user_id) + self._last_principals[user.user_id] = current + if previous is not None and previous != current: + invalidate_user_suggestions(user.user_id) + invalidate_user_spaces(user.user_id) + def invalidate_user(self, user: User | str) -> None: """Drop cached DLS principals for one auth user or OpenSearch user name.""" user_id = user.user_id if isinstance(user, User) else user diff --git a/src/services/document_index_writer.py b/src/services/document_index_writer.py index 0a25a991c..2f81af8f9 100644 --- a/src/services/document_index_writer.py +++ b/src/services/document_index_writer.py @@ -14,6 +14,7 @@ from dataclasses import dataclass, field from typing import Any +from services.search_cache import invalidate_corpus_caches from utils.embedding_fields import ensure_embedding_field_exists, get_embedding_space_id from utils.embeddings import create_index_body from utils.group_acl import unique_acl_principal_labels, unique_acl_principals @@ -156,10 +157,14 @@ async def index_chunks( ) ) - result = await self._bulk_with_retry(client, bulk_body, refresh=refresh) - self._raise_for_bulk_errors(result) - if final: - await self._refresh(index_name) + try: + result = await self._bulk_with_retry(client, bulk_body, refresh=refresh) + self._raise_for_bulk_errors(result) + if final: + await self._refresh(index_name) + finally: + # Even a failed batch may have written some chunks. + invalidate_corpus_caches() logger.info( "Indexed document chunks", @@ -198,6 +203,7 @@ async def delete_ingest_run(self, ingest_run_id: str, *, index_name: str | None conflicts="proceed", ) deleted = int(response.get("deleted", 0)) if isinstance(response, dict) else 0 + invalidate_corpus_caches() logger.info( "Deleted failed ingest run chunks", index_name=resolved_index, diff --git a/src/services/search_cache.py b/src/services/search_cache.py new file mode 100644 index 000000000..9022a97de --- /dev/null +++ b/src/services/search_cache.py @@ -0,0 +1,173 @@ +"""In-process caches for the search path. + +Typing in the search box fires many near-identical requests, and every search +used to repeat the same expensive work: discovering the corpus' embedding +spaces (paginated aggregations) and embedding the query through the provider. +These caches make a repeated or retyped query cheap. + +Three caches, each with different scoping rules: + +* ``QUERY_EMBEDDING_CACHE`` - query vectors keyed by ``(route_model, text)``. + A vector depends only on the text and the embedding route, never on the + caller, so it is safe to share across users. +* ``SPACES_CACHE`` - embedding spaces present in the corpus. The OpenSearch + client is per-user, so the key includes the user id and the filters. +* ``SUGGEST_CACHE`` - suggestion payloads containing file names and text + snippets. Always keyed by user id so one user can never be served another + user's results. + +Every cache is single-flight: concurrent identical misses share one +computation instead of each paying for it. +""" + +import asyncio +import json +from collections.abc import Awaitable, Callable, Hashable +from typing import Any, cast + +from cachetools import TTLCache + +QUERY_EMBEDDING_CACHE_SIZE = 512 +QUERY_EMBEDDING_TTL_SECONDS = 3600 +SPACES_CACHE_SIZE = 256 +SPACES_TTL_SECONDS = 60 +SUGGEST_CACHE_SIZE = 512 +SUGGEST_TTL_SECONDS = 45 + + +class SingleflightTTLCache: + """TTL + size-bounded cache whose misses are computed once per key.""" + + def __init__(self, maxsize: int, ttl: float) -> None: + self._cache: TTLCache[Hashable, Any] = TTLCache(maxsize=maxsize, ttl=ttl) + self._inflight: dict[Hashable, asyncio.Future[Any]] = {} + # Bumped by clear() so a computation that began before a clear cannot + # store its (now stale) result after it. + self._generation = 0 + + def get(self, key: Hashable) -> Any | None: + return self._cache.get(key) + + def set(self, key: Hashable, value: Any) -> None: + self._cache[key] = value + + def clear(self) -> None: + self._generation += 1 + self._cache.clear() + self._inflight.clear() + + def discard_where(self, predicate: Callable[[Hashable], bool]) -> None: + """Drop matching entries and keep in-flight computations from storing.""" + self._generation += 1 + for key in [key for key in self._cache if predicate(key)]: + del self._cache[key] + for key in [key for key in self._inflight if predicate(key)]: + del self._inflight[key] + + async def get_or_compute( + self, + key: Hashable, + compute: Callable[[], Awaitable[Any]], + *, + cacheable: Callable[[Any], bool] = lambda _value: True, + ) -> Any: + """Return the cached value, or run ``compute`` once for all callers. + + ``cacheable`` lets the caller refuse to store degraded results (for + example a response carrying provider warnings) so a transient failure + is not pinned for the whole TTL. Exceptions are never cached. + """ + while True: + hit = self._cache.get(key) + if hit is not None: + return hit + + pending = self._inflight.get(key) + if pending is None: + break + try: + # shield: a cancelled waiter must not cancel the shared future. + return await asyncio.shield(pending) + except asyncio.CancelledError: + # Only the leader's cancellation is recoverable; take over. + if not pending.cancelled(): + raise + + future: asyncio.Future[Any] = asyncio.get_running_loop().create_future() + self._inflight[key] = future + generation = self._generation + try: + value = await compute() + except asyncio.CancelledError: + future.cancel() + raise + except Exception as exc: + future.set_exception(exc) + # Mark retrieved so asyncio does not warn when nobody was waiting. + future.exception() + raise + else: + if generation == self._generation and cacheable(value): + self._cache[key] = value + future.set_result(value) + return value + finally: + if self._inflight.get(key) is future: + del self._inflight[key] + + +QUERY_EMBEDDING_CACHE = SingleflightTTLCache( + QUERY_EMBEDDING_CACHE_SIZE, QUERY_EMBEDDING_TTL_SECONDS +) +SPACES_CACHE = SingleflightTTLCache(SPACES_CACHE_SIZE, SPACES_TTL_SECONDS) +SUGGEST_CACHE = SingleflightTTLCache(SUGGEST_CACHE_SIZE, SUGGEST_TTL_SECONDS) + + +def clear_search_caches() -> None: + """Drop every cached entry (used by tests and after bulk index changes).""" + QUERY_EMBEDDING_CACHE.clear() + SPACES_CACHE.clear() + SUGGEST_CACHE.clear() + + +def normalize_embedding_text(query: str) -> str: + """Collapse whitespace only; case is preserved because it can change the vector.""" + return " ".join(query.split()) + + +def normalize_query(query: str) -> str: + """Canonical form used for result-cache keys (case and spacing insensitive).""" + return " ".join(query.split()).lower() + + +def fingerprint(value: Any) -> str: + """Stable, hashable representation of filter clauses / filter dicts.""" + return json.dumps(value, sort_keys=True, separators=(",", ":"), default=str) + + +def invalidate_user_suggestions(user_id: str) -> None: + """Drop one user's cached suggestions (keys start with the user id).""" + SUGGEST_CACHE.discard_where( + lambda key: ( + isinstance(key, tuple) and len(key) > 0 and cast(tuple[Any, ...], key)[0] == user_id + ) + ) + + +def invalidate_user_spaces(user_id: str) -> None: + """Drop one user's cached embedding spaces (keys start with the user id).""" + SPACES_CACHE.discard_where( + lambda key: ( + isinstance(key, tuple) and len(key) > 0 and cast(tuple[Any, ...], key)[0] == user_id + ) + ) + + +def invalidate_corpus_caches() -> None: + """Drop results that depend on what is indexed. + + Query embeddings are kept: a vector depends only on the model and the + text, and it is the one cache that saves a paid provider call. + """ + SPACES_CACHE.clear() + SUGGEST_CACHE.clear() diff --git a/src/services/search_service.py b/src/services/search_service.py index 56b64789b..649cab85f 100644 --- a/src/services/search_service.py +++ b/src/services/search_service.py @@ -14,6 +14,20 @@ from observability.spans import Span, trace_root, trace_span from services.llm_gateway import LlmGatewayError from services.llm_gateway import embeddings as gateway_embeddings +from services.search_cache import ( + QUERY_EMBEDDING_CACHE, + SPACES_CACHE, + SUGGEST_CACHE, + fingerprint, + normalize_embedding_text, + normalize_query, +) +from services.search_suggest import ( + build_preview, + escape_wildcard, + extract_candidates, + normalize_suggest_query, +) from utils.container_utils import transform_localhost_url from utils.embedding_fields import ( INDEXED_EMBEDDING_ROUTE_PREFIX, @@ -218,6 +232,79 @@ def _build_terms_agg(field: str, label_field: str | None = None) -> dict[str, An return chunks, aggregations +SUGGEST_HIT_LIMIT = 30 +SUGGEST_MAX_SUGGESTIONS = 4 +SUGGEST_MAX_FILES = 5 +SUGGEST_PREVIEW_FILES = 3 + +_SUGGEST_FILTER_FIELDS = { + "data_sources": "filename", + "document_types": "mimetype", + "owners": "owner", + "connector_types": "connector_type", +} + + +def _filters_to_clauses(filters: dict[str, Any]) -> list[dict[str, Any]]: + """Translate frontend facet filters into OpenSearch filter clauses. + + Unknown keys are ignored (an allowlist) so a caller cannot target + arbitrary index fields through the suggest endpoint. + """ + clauses: list[dict[str, Any]] = [] + for key, field_name in _SUGGEST_FILTER_FIELDS.items(): + values = filters.get(key) + if not isinstance(values, list): + continue + values = [v for v in values if isinstance(v, str)] + if not values: + # Empty list means "match nothing", mirroring search_tool. + clauses.append({"term": {field_name: "__IMPOSSIBLE_VALUE__"}}) + elif len(values) == 1: + clauses.append({"term": {field_name: values[0]}}) + else: + clauses.append({"terms": {field_name: values}}) + return clauses + + +def _group_suggest_files( + chunks: list[dict[str, Any]], query: str, *, allow_semantic: bool +) -> list[dict[str, Any]]: + """Collapse ranked chunks into per-file rows; preview only the top few.""" + grouped: dict[str, list[dict[str, Any]]] = {} + for chunk in chunks: + filename = chunk.get("filename") + if isinstance(filename, str) and filename: + grouped.setdefault(filename, []).append(chunk) + + files: list[dict[str, Any]] = [] + for rank, (filename, file_chunks) in enumerate(grouped.items()): + if rank >= SUGGEST_MAX_FILES: + break + preview = None + if rank < SUGGEST_PREVIEW_FILES: + candidates = ( + build_preview( + c.get("text"), c.get("highlights"), query, allow_semantic=allow_semantic + ) + for c in file_chunks + ) + previews = [p for p in candidates if p] + # A literal keyword match explains the hit better than a guess. + preview = next((p for p in previews if p["kind"] == "keyword"), None) or ( + previews[0] if previews else None + ) + files.append( + { + "filename": filename, + "mimetype": file_chunks[0].get("mimetype"), + "connector_type": file_chunks[0].get("connector_type"), + "preview": preview, + } + ) + return files + + def _record_search_response(span: Span, response: dict[str, Any]) -> None: """What a search returned, for its trace: counts, scores and ids, never text.""" results = response.get("results") or [] @@ -282,6 +369,48 @@ def _configure_provider_env(self): except Exception as e: logger.warning("[SEARCH] Could not configure Ollama endpoint from config", error=str(e)) + async def _discover_embedding_spaces( + self, opensearch_client, filter_clauses: list[dict[str, Any]] + ) -> list[EmbeddingSpace]: + """List the embedding spaces present in the (filtered) corpus.""" + available_spaces: list[EmbeddingSpace] = [] + seen_spaces: set[str] = set() + for legacy in (False, True): + after = None + while True: + agg_query = { + "size": 0, + "aggs": build_embedding_space_aggregation( + size=EMBEDDING_SPACE_PAGE_SIZE, + qualified_after=None if legacy else after, + legacy_after=after if legacy else None, + include_qualified=not legacy, + include_legacy=legacy, + ), + } + bool_query: dict[str, Any] = {} + if filter_clauses: + bool_query["filter"] = filter_clauses + if legacy: + bool_query["must_not"] = [{"exists": {"field": "embedding_space_id"}}] + if bool_query: + agg_query["query"] = {"bool": bool_query} + + agg_result = await opensearch_client.search( + index=get_index_name(), body=agg_query, params={"terminate_after": 0} + ) + for space in embedding_spaces_from_aggregation(agg_result): + if space.space_id not in seen_spaces: + seen_spaces.add(space.space_id) + available_spaces.append(space) + + qualified_after, legacy_after = embedding_space_after_keys(agg_result) + next_after = legacy_after if legacy else qualified_after + if not next_after or next_after == after: + break + after = next_after + return available_spaces + async def search_tool( self, query: str, embedding_model: str = None, *, source: str = "agent_tool" ) -> dict[str, Any]: @@ -293,8 +422,8 @@ async def search_tool( embedding_model (str): Optional override for embedding model. If not provided, uses the current embedding model from configuration. - source (str): Who is searching, for the trace: ``agent_tool`` or - ``api`` (see ``search``). + source (str): Who is searching, for the trace: ``agent_tool``, + ``api`` or ``suggest`` (see ``search``). Returns: dict (str, Any): {"results": [chunks]} on success @@ -402,41 +531,16 @@ async def _search(self, query: str, embedding_model: str) -> dict[str, Any]: filter_clauses.append({"terms": {field_name: values}}) try: - seen_spaces: set[str] = set() - for legacy in (False, True): - after = None - while True: - agg_query = { - "size": 0, - "aggs": build_embedding_space_aggregation( - size=EMBEDDING_SPACE_PAGE_SIZE, - qualified_after=None if legacy else after, - legacy_after=after if legacy else None, - include_qualified=not legacy, - include_legacy=legacy, - ), - } - bool_query: dict[str, Any] = {} - if filter_clauses: - bool_query["filter"] = filter_clauses - if legacy: - bool_query["must_not"] = [{"exists": {"field": "embedding_space_id"}}] - if bool_query: - agg_query["query"] = {"bool": bool_query} - - agg_result = await opensearch_client.search( - index=get_index_name(), body=agg_query, params={"terminate_after": 0} - ) - for space in embedding_spaces_from_aggregation(agg_result): - if space.space_id not in seen_spaces: - seen_spaces.add(space.space_id) - available_spaces.append(space) - - qualified_after, legacy_after = embedding_space_after_keys(agg_result) - next_after = legacy_after if legacy else qualified_after - if not next_after or next_after == after: - break - after = next_after + # Space discovery is several aggregation round-trips and the + # answer changes only when documents are (re)indexed, so reuse + # it briefly. Keyed per user because the client is per-user. + available_spaces = list( + await SPACES_CACHE.get_or_compute( + (user_id, fingerprint(filter_clauses)), + lambda: self._discover_embedding_spaces(opensearch_client, filter_clauses), + cacheable=bool, + ) + ) if not available_spaces: # Fallback to configured model if no documents indexed yet @@ -472,9 +576,7 @@ async def _search(self, query: str, embedding_model: str) -> dict[str, Any]: ] # Parallelize embedding generation for all models - async def embed_with_space( - space: EmbeddingSpace, - ) -> tuple[EmbeddingSpace, list[float]]: + async def _embed_uncached(space: EmbeddingSpace) -> list[float]: """Generate one query vector through its exact gateway route.""" delay = EMBED_RETRY_INITIAL_DELAY attempts = 0 @@ -495,7 +597,7 @@ async def embed_with_space( embedding = data[0].get("embedding") if data else None if embedding is None: raise RuntimeError("Embedding provider returned no vector") - return space, embedding + return embedding except LlmGatewayError: # Configuration/provenance errors are deterministic; # skip this space immediately and retain keyword search. @@ -528,6 +630,17 @@ async def embed_with_space( f"Failed to embed with model {space.space_id}" ) from last_exception + async def embed_with_space( + space: EmbeddingSpace, + ) -> tuple[EmbeddingSpace, list[float]]: + """Embed through the shared cache: retyped queries cost nothing, + and concurrent identical requests share one provider call.""" + embedding = await QUERY_EMBEDDING_CACHE.get_or_compute( + (space.route_model, normalize_embedding_text(query)), + lambda: _embed_uncached(space), + ) + return space, embedding + # Run all embeddings in parallel, tolerating per-model failures so # one broken model (e.g. provider credentials removed after ingest) # doesn't take down the entire search. If all models fail we fall @@ -918,12 +1031,15 @@ async def search( limit: int = 10, score_threshold: float = 0, embedding_model: str = None, + *, + source: str = "api", ) -> dict[str, Any]: """Public search method for API endpoints Args: embedding_model: Embedding model to use for search (defaults to the currently configured embedding model) + source: Trace tag distinguishing who is searching (see ``search_tool``) """ # Set auth context if provided (for direct API calls) from config.settings import is_no_auth_mode @@ -944,4 +1060,226 @@ async def search( set_search_limit(limit) set_score_threshold(score_threshold) - return await self.search_tool(query, embedding_model=embedding_model, source="api") + return await self.search_tool(query, embedding_model=embedding_model, source=source) + + # ------------------------------------------------------------------ + # Search-as-you-type suggestions + # ------------------------------------------------------------------ + + async def suggest( + self, + query: str, + user_id: str, + jwt_token: str | None = None, + filters: dict[str, Any] | None = None, + semantic: bool = False, + score_threshold: float = 0, + ) -> dict[str, Any]: + """Return query completions, matching files and one-line previews. + + ``semantic=False`` is the cheap keyword phase meant to run on every + (debounced) keystroke. ``semantic=True`` additionally runs the hybrid + search (query embedding + KNN) and merges its hits in. Results are + cached per user, and a cached semantic payload also answers keyword + requests because it is a superset. + """ + cleaned = normalize_suggest_query(query) + if not cleaned or not user_id: + return {"phase": "semantic" if semantic else "keyword", "query": cleaned, + "suggestions": [], "files": []} # fmt: skip + + filters = filters or {} + # Semantic hits depend on the score threshold, so it is part of the key. + base_key = (user_id, normalize_query(cleaned), fingerprint(filters), score_threshold) + if not semantic: + upgraded = SUGGEST_CACHE.get((*base_key, True)) + if upgraded is not None: + return upgraded + + def _cacheable(payload: dict[str, Any]) -> bool: + # Never pin a degraded (provider warning) or failed response. + return not payload.get("warnings") and not payload.get("error") + + return await SUGGEST_CACHE.get_or_compute( + (*base_key, semantic), + lambda: self._compute_suggestions( + cleaned, user_id, jwt_token, filters, semantic, score_threshold + ), + cacheable=_cacheable, + ) + + async def _compute_suggestions( + self, + query: str, + user_id: str, + jwt_token: str | None, + filters: dict[str, Any], + semantic: bool, + score_threshold: float, + ) -> dict[str, Any]: + opensearch_client = self.session_manager.get_user_opensearch_client(user_id, jwt_token) + filter_clauses = _filters_to_clauses(filters) + + warnings: list[dict[str, Any]] = [] + chunks: list[dict[str, Any]] + + if semantic: + # The separate keyword query cannot honor the hybrid score threshold, so + # it is skipped when one is set to avoid resurrecting rejected files. + keyword_task = self._suggest_keyword_chunks( + opensearch_client, query, filter_clauses, skip=score_threshold > 0 + ) + hybrid_task = self.search( + query, + user_id=user_id, + jwt_token=jwt_token, + filters=filters, + limit=SUGGEST_HIT_LIMIT, + score_threshold=score_threshold, + source="suggest", + ) + keyword_chunks, hybrid = await asyncio.gather(keyword_task, hybrid_task) + if hybrid.get("error"): + return {"error": hybrid["error"]} + warnings = hybrid.get("warnings", []) + # Ranked hybrid hits first, then keyword-only hits (e.g. partial + # filename matches) that the hybrid query did not return. + chunks = list(hybrid.get("results", [])) + seen_ids = {c.get("id") for c in chunks if c.get("id")} + chunks.extend(c for c in keyword_chunks if c.get("id") not in seen_ids) + else: + chunks = await self._suggest_keyword_chunks(opensearch_client, query, filter_clauses) + + candidates = extract_candidates([c.get("text") or "" for c in chunks], query) + suggestions = await self._count_suggestions(opensearch_client, candidates, filter_clauses) + + payload: dict[str, Any] = { + "phase": "semantic" if semantic else "keyword", + "query": query, + "suggestions": suggestions, + "files": _group_suggest_files(chunks, query, allow_semantic=semantic), + } + if warnings: + payload["warnings"] = warnings + return payload + + async def _suggest_keyword_chunks( + self, + opensearch_client, + query: str, + filter_clauses: list[dict[str, Any]], + *, + skip: bool = False, + ) -> list[dict[str, Any]]: + """Cheap lexical query: text (fuzzy + prefix) and partial filename.""" + if skip: + return [] + body: dict[str, Any] = { + "size": SUGGEST_HIT_LIMIT, + "_source": ["filename", "mimetype", "text", "connector_type"], + "query": { + "bool": { + "should": [ + { + "multi_match": { + "query": query, + "fields": ["text^2", "filename^1.5"], + "type": "best_fields", + "operator": "or", + "fuzziness": "AUTO:4,7", + } + }, + { + "match_phrase_prefix": { + "text": {"query": query, "max_expansions": 50, "boost": 0.5} + } + }, + { + "wildcard": { + "filename": { + "value": f"*{escape_wildcard(query.lower())}*", + "case_insensitive": True, + "boost": 2.0, + } + } + }, + ], + "minimum_should_match": 1, + "filter": filter_clauses, + } + }, + "highlight": { + "type": "unified", + "fields": { + "text": { + "fragment_size": 160, + "number_of_fragments": 1, + "pre_tags": [""], + "post_tags": [""], + } + }, + "require_field_match": False, + }, + } + result = await opensearch_client.search( + index=get_index_name(), body=body, params={"terminate_after": 0} + ) + chunks = [] + for hit in result.get("hits", {}).get("hits", []): + source = hit.get("_source", {}) + chunks.append( + { + "id": hit.get("_id"), + "filename": source.get("filename"), + "mimetype": source.get("mimetype"), + "connector_type": source.get("connector_type"), + "text": source.get("text"), + "highlights": hit.get("highlight", {}).get("text") or [], + } + ) + return chunks + + async def _count_suggestions( + self, + opensearch_client, + candidates: list[str], + filter_clauses: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Attach "N files / M chunks" to each candidate in a single request.""" + if not candidates: + return [] + body: dict[str, Any] = { + "size": 0, + "query": {"bool": {"filter": filter_clauses}} if filter_clauses else {"match_all": {}}, + "aggs": { + "candidates": { + "filters": { + "filters": [ + { + "match_phrase_prefix": { + "text": {"query": candidate, "max_expansions": 50} + } + } + for candidate in candidates + ] + }, + "aggs": {"files": {"cardinality": {"field": "filename"}}}, + } + }, + } + result = await opensearch_client.search( + index=get_index_name(), body=body, params={"terminate_after": 0} + ) + buckets = result.get("aggregations", {}).get("candidates", {}).get("buckets", []) + suggestions = [] + for candidate, bucket in zip(candidates, buckets, strict=False): + chunks = bucket.get("doc_count", 0) + if chunks > 0: + suggestions.append( + { + "text": candidate, + "files": bucket.get("files", {}).get("value", 0), + "chunks": chunks, + } + ) + return suggestions[:SUGGEST_MAX_SUGGESTIONS] diff --git a/src/services/search_suggest.py b/src/services/search_suggest.py new file mode 100644 index 000000000..aff5cf770 --- /dev/null +++ b/src/services/search_suggest.py @@ -0,0 +1,166 @@ +"""Pure helpers for the search-suggestion endpoint. + +Everything here is side-effect free so it can be unit-tested without +OpenSearch: building the one-line "why did this file match" preview, and +deriving query-completion candidates from the text of top-ranked chunks. +""" + +import re +from collections import Counter + +MIN_QUERY_LENGTH = 2 +MAX_QUERY_LENGTH = 100 +PREVIEW_MAX_CHARS = 180 + +PREVIEW_KEYWORD = "keyword" +PREVIEW_SEMANTIC = "semantic" + +_MARK_BLOCK = re.compile(r".*?", re.DOTALL) +_MARK_TAG = re.compile(r"") +_PARTIAL_TAG_AT_END = re.compile(r"<[^>]{0,6}$") +_SENTENCE_SPLIT = re.compile(r"(?<=[.!?])\s+|\n{2,}") +_TOKEN = re.compile(r"\w+", re.UNICODE) + +# Words that make a poor second word in a "query + next word" completion. +_STOPWORDS = frozenset( + { + "a", "an", "and", "are", "as", "at", "be", "by", "for", "from", "in", + "is", "it", "of", "on", "or", "that", "the", "this", "to", "was", "with", + } +) # fmt: skip + + +def normalize_suggest_query(query: str | None) -> str: + """Trim, collapse whitespace and cap length. Returns "" for unusable input.""" + if not isinstance(query, str): + return "" + cleaned = " ".join(query.split())[:MAX_QUERY_LENGTH] + if len(cleaned) < MIN_QUERY_LENGTH or cleaned == "*": + return "" + return cleaned + + +def strip_marks(fragment: str) -> str: + return _MARK_TAG.sub("", fragment) + + +def escape_wildcard(value: str) -> str: + """Escape OpenSearch wildcard metacharacters so user input is matched literally.""" + return value.replace("\\", "\\\\").replace("*", "\\*").replace("?", "\\?") + + +def _truncate_plain(text: str, max_chars: int) -> str: + if len(text) <= max_chars: + return text + cut = text[:max_chars].rsplit(" ", 1)[0].rstrip(" ,;:-") + return f"{cut or text[:max_chars]}…" + + +def _trim_fragment(fragment: str, max_chars: int) -> str: + """Shrink an over-long highlight fragment around its first .""" + if len(strip_marks(fragment)) <= max_chars: + return fragment + + match = _MARK_BLOCK.search(fragment) + if not match: + return _truncate_plain(fragment, max_chars) + + before, mark, after = fragment[: match.start()], match.group(0), fragment[match.end() :] + budget = max(max_chars - len(strip_marks(mark)), 0) + keep_before = budget // 3 + keep_after = budget - min(keep_before, len(before)) + + if len(before) > keep_before: + # Drop the (probably cut) leading word fragment. + before = "…" + before[-keep_before:].split(" ", 1)[-1] if keep_before else "…" + if len(strip_marks(after)) > keep_after: + after = after[:keep_after] + # Never leave a half-open or a partial tag behind. + after = _PARTIAL_TAG_AT_END.sub("", after) + if after.rfind("") > after.rfind(""): + after = after[: after.rfind("")] + after = after.rsplit(" ", 1)[0].rstrip(" ,;:-") + "…" + return before + mark + after + + +def preview_from_highlight(text: str | None, fragment: str) -> str: + """Turn an OpenSearch ```` fragment into a one-line preview. + + ``…`` marks the side(s) where the fragment was cut out of the chunk, so the + reader can tell the line is an excerpt. + """ + plain = strip_marks(fragment) + start = text.find(plain) if text else -1 + line = _trim_fragment(" ".join(fragment.split()), PREVIEW_MAX_CHARS) + if start > 0 and not line.startswith("…"): + line = f"…{line}" + if start >= 0 and text and start + len(plain) < len(text) and not line.endswith("…"): + line = f"{line}…" + return line + + +def best_sentence(text: str, query: str) -> str | None: + """Pick the sentence sharing the most words with the query (first on ties).""" + cleaned = text.strip() if isinstance(text, str) else "" + if not cleaned: + return None + + sentences = [s.strip() for s in _SENTENCE_SPLIT.split(cleaned) if len(s.strip()) >= 20] + if not sentences: + sentences = [cleaned] + + query_tokens = {t.lower() for t in _TOKEN.findall(query)} + + def overlap(sentence: str) -> int: + return len(query_tokens & {t.lower() for t in _TOKEN.findall(sentence)}) + + best = max(sentences, key=overlap) # max() keeps the first of equal scores + return _truncate_plain(" ".join(best.split()), PREVIEW_MAX_CHARS) + + +def build_preview( + text: str | None, highlights: list[str] | None, query: str, *, allow_semantic: bool +) -> dict[str, str] | None: + """Return ``{"text", "kind"}`` describing how a chunk relates to the query. + + ``kind`` is ``keyword`` when the line contains ````-wrapped matches and + ``semantic`` for a plain best-matching sentence. Returns None when there is + nothing meaningful to show. + """ + if highlights: + fragment = next((h for h in highlights if isinstance(h, str) and h.strip()), None) + if fragment: + return {"text": preview_from_highlight(text, fragment), "kind": PREVIEW_KEYWORD} + if allow_semantic and text: + sentence = best_sentence(text, query) + if sentence: + return {"text": sentence, "kind": PREVIEW_SEMANTIC} + return None + + +def extract_candidates(texts: list[str], query: str, limit: int = 6) -> list[str]: + """Derive completion candidates (``query`` + rest of word [+ next word]). + + For ``docl`` this yields ``docling`` and ``docling serve`` when the sampled + chunks contain them. Candidates are ranked by how often they occur. + """ + normalized = " ".join(query.lower().split()) + if not normalized: + return [] + + body = re.escape(normalized).replace(r"\ ", r"\s+") + pattern = re.compile(rf"(?= 2 and following not in _STOPWORDS: + counts[f"{completion} {following}"] += 1 + + ranked = sorted(counts.items(), key=lambda item: (-item[1], len(item[0]), item[0])) + return [candidate for candidate, _ in ranked[:limit]] diff --git a/src/utils/acl_utils.py b/src/utils/acl_utils.py index 07f76095c..bea858b8f 100644 --- a/src/utils/acl_utils.py +++ b/src/utils/acl_utils.py @@ -11,6 +11,7 @@ from typing import Any from config.settings import get_index_name +from services.search_cache import invalidate_corpus_caches from src.connectors.base import DocumentACL from utils.logging_config import get_logger @@ -157,6 +158,8 @@ async def update_document_acl( }, ) + # Visibility changed, so cached per-user suggestions may be wrong. + invalidate_corpus_caches() return {"status": "updated", "chunks_updated": response.get("updated", 0)} except Exception as e: @@ -224,6 +227,7 @@ async def batch_update_acls( try: results = await asyncio.gather(*update_tasks, return_exceptions=True) + invalidate_corpus_caches() total_chunks_updated = 0 errors = [] for result in results: diff --git a/src/utils/opensearch_delete.py b/src/utils/opensearch_delete.py index 4e9db031f..84ee34bce 100644 --- a/src/utils/opensearch_delete.py +++ b/src/utils/opensearch_delete.py @@ -2,6 +2,7 @@ from opensearchpy.exceptions import NotFoundError +from services.search_cache import invalidate_corpus_caches from utils.logging_config import get_logger logger = get_logger(__name__) @@ -59,17 +60,21 @@ async def delete_document_ids( a delete per ID. """ deleted_count = 0 - for document_id in document_ids: - try: - result = await opensearch_client.delete( - index=index, - id=document_id, - refresh=refresh, - ) - except NotFoundError: - continue + try: + for document_id in document_ids: + try: + result = await opensearch_client.delete( + index=index, + id=document_id, + refresh=refresh, + ) + except NotFoundError: + continue - if result.get("result") == "deleted": - deleted_count += 1 + if result.get("result") == "deleted": + deleted_count += 1 + finally: + if deleted_count: + invalidate_corpus_caches() return deleted_count diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index dae6d701e..6b84fa963 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -30,6 +30,19 @@ async def onboard_system(): yield +@pytest.fixture(autouse=True) +def _reset_search_caches(): + """Search caches are process-wide; keep tests from seeing each other's entries.""" + try: + from services.search_cache import clear_search_caches + except ImportError: + yield + return + clear_search_caches() + yield + clear_search_caches() + + @pytest.fixture(autouse=True) def _reset_db_engine_module_state(monkeypatch): """Defensive: per-test, ensure `db.engine`'s module-level singletons diff --git a/tests/unit/test_group_acl.py b/tests/unit/test_group_acl.py index 0b599e58d..07e5cbac6 100644 --- a/tests/unit/test_group_acl.py +++ b/tests/unit/test_group_acl.py @@ -490,6 +490,62 @@ async def index(self, **kwargs): assert service._locks == {} +@pytest.mark.asyncio +async def test_dls_principal_change_drops_that_users_cached_suggestions(): + from services import search_cache + from services.dls_principal_service import DLSPrincipalService + from session_manager import User + + current_principals = ["u:test:user-1", "g:test:engineering"] + + class ConnectionManager: + async def list_connections(self, user_id=None): + return [] + + def get_auth_user_principals(self, user): + return list(current_principals) + + class ConnectorService: + connection_manager = ConnectionManager() + + class Indices: + async def exists(self, index): + return True + + class OpenSearchClient: + indices = Indices() + + async def index(self, **kwargs): + return None + + service = DLSPrincipalService( + ConnectorService(), opensearch_client=OpenSearchClient(), refresh_ttl_seconds=0 + ) + user = User(user_id="user-1", email=None, name=None, provider="test") + other_key = ("user-2", "docling", "{}", 0, True) + user_key = ("user-1", "docling", "{}", 0, True) + other_space_key = ("user-2", "fp-2") + user_space_key = ("user-1", "fp-1") + + await service.refresh_user_principals(user) + search_cache.SUGGEST_CACHE.set(user_key, {"files": ["old.pdf"]}) + search_cache.SUGGEST_CACHE.set(other_key, {"files": ["other.pdf"]}) + search_cache.SPACES_CACHE.set(user_space_key, ["space-1"]) + search_cache.SPACES_CACHE.set(other_space_key, ["space-2"]) + + await service.refresh_user_principals(user) + assert search_cache.SUGGEST_CACHE.get(user_key) == {"files": ["old.pdf"]} + assert search_cache.SPACES_CACHE.get(user_space_key) == ["space-1"] + + current_principals.remove("g:test:engineering") + await service.refresh_user_principals(user) + + assert search_cache.SUGGEST_CACHE.get(user_key) is None + assert search_cache.SUGGEST_CACHE.get(other_key) == {"files": ["other.pdf"]} + assert search_cache.SPACES_CACHE.get(user_space_key) is None + assert search_cache.SPACES_CACHE.get(other_space_key) == ["space-2"] + + @pytest.mark.asyncio async def test_dls_principal_service_skips_connector_lookup_without_opensearch_client(monkeypatch): from services.dls_principal_service import DLSPrincipalService diff --git a/tests/unit/test_search_cache.py b/tests/unit/test_search_cache.py new file mode 100644 index 000000000..01e1b87e4 --- /dev/null +++ b/tests/unit/test_search_cache.py @@ -0,0 +1,250 @@ +"""Unit tests for services/search_cache.py.""" + +import asyncio + +import pytest + +from services import search_cache +from services.search_cache import ( + SingleflightTTLCache, + fingerprint, + invalidate_corpus_caches, + normalize_embedding_text, + normalize_query, +) + + +@pytest.mark.asyncio +async def test_second_call_is_served_from_cache(): + cache = SingleflightTTLCache(maxsize=8, ttl=60) + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + return "value" + + assert await cache.get_or_compute("k", compute) == "value" + assert await cache.get_or_compute("k", compute) == "value" + assert calls == 1 + + +@pytest.mark.asyncio +async def test_concurrent_misses_share_one_computation(): + cache = SingleflightTTLCache(maxsize=8, ttl=60) + calls = 0 + gate = asyncio.Event() + + async def compute(): + nonlocal calls + calls += 1 + await gate.wait() + return "value" + + tasks = [asyncio.create_task(cache.get_or_compute("k", compute)) for _ in range(5)] + await asyncio.sleep(0) + gate.set() + + assert await asyncio.gather(*tasks) == ["value"] * 5 + assert calls == 1 + + +@pytest.mark.asyncio +async def test_exceptions_are_not_cached_and_reach_all_waiters(): + cache = SingleflightTTLCache(maxsize=8, ttl=60) + calls = 0 + gate = asyncio.Event() + + async def failing(): + nonlocal calls + calls += 1 + await gate.wait() + raise RuntimeError("boom") + + tasks = [asyncio.create_task(cache.get_or_compute("k", failing)) for _ in range(3)] + await asyncio.sleep(0) + gate.set() + results = await asyncio.gather(*tasks, return_exceptions=True) + assert all(isinstance(r, RuntimeError) for r in results) + assert calls == 1 + + async def working(): + return "ok" + + assert await cache.get_or_compute("k", working) == "ok" + + +@pytest.mark.asyncio +async def test_uncacheable_results_are_recomputed(): + cache = SingleflightTTLCache(maxsize=8, ttl=60) + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + return {"warnings": ["degraded"]} + + def cacheable(value): + return not value["warnings"] + + await cache.get_or_compute("k", compute, cacheable=cacheable) + await cache.get_or_compute("k", compute, cacheable=cacheable) + assert calls == 2 + + +@pytest.mark.asyncio +async def test_cancelled_waiter_does_not_cancel_the_shared_computation(): + cache = SingleflightTTLCache(maxsize=8, ttl=60) + gate = asyncio.Event() + + async def compute(): + await gate.wait() + return "value" + + owner = asyncio.create_task(cache.get_or_compute("k", compute)) + await asyncio.sleep(0) + waiter = asyncio.create_task(cache.get_or_compute("k", compute)) + await asyncio.sleep(0) + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + + gate.set() + assert await owner == "value" + + +@pytest.mark.asyncio +async def test_cancelled_leader_does_not_cancel_waiters(): + cache = SingleflightTTLCache(maxsize=8, ttl=60) + started = asyncio.Event() + gate = asyncio.Event() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + started.set() + await gate.wait() + return "value" + + leader = asyncio.create_task(cache.get_or_compute("k", compute)) + await started.wait() + waiters = [asyncio.create_task(cache.get_or_compute("k", compute)) for _ in range(2)] + await asyncio.sleep(0) + + leader.cancel() + with pytest.raises(asyncio.CancelledError): + await leader + + gate.set() + assert await asyncio.gather(*waiters) == ["value", "value"] + assert calls == 2 + assert cache.get("k") == "value" + + +@pytest.mark.asyncio +async def test_least_recently_used_entry_is_evicted_when_full(): + cache = SingleflightTTLCache(maxsize=2, ttl=60) + + async def make(value): + return value + + for key in ("a", "b", "c"): + await cache.get_or_compute(key, lambda key=key: make(key)) + + assert cache.get("a") is None + assert cache.get("b") == "b" + assert cache.get("c") == "c" + + +def test_key_normalization(): + assert normalize_query(" Docling Serve ") == "docling serve" + # Embedding text keeps case: it can change the vector. + assert normalize_embedding_text(" Docling Serve ") == "Docling Serve" + assert fingerprint({"b": [1], "a": [2]}) == fingerprint({"a": [2], "b": [1]}) + + +@pytest.mark.asyncio +async def test_result_computed_across_a_clear_is_not_stored(): + cache = SingleflightTTLCache(maxsize=8, ttl=60) + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def slow_stale(): + nonlocal calls + calls += 1 + started.set() + await release.wait() + return "stale" + + async def fresh(): + nonlocal calls + calls += 1 + return "fresh" + + leader = asyncio.create_task(cache.get_or_compute("k", slow_stale)) + await started.wait() + cache.clear() + release.set() + assert await leader == "stale" + + assert await cache.get_or_compute("k", fresh) == "fresh" + assert calls == 2 + + +def test_invalidate_corpus_caches_keeps_query_embeddings(): + search_cache.QUERY_EMBEDDING_CACHE.set("q", [0.1]) + search_cache.SPACES_CACHE.set("s", ["space"]) + search_cache.SUGGEST_CACHE.set("g", {"files": []}) + + invalidate_corpus_caches() + + assert search_cache.QUERY_EMBEDDING_CACHE.get("q") == [0.1] + assert search_cache.SPACES_CACHE.get("s") is None + assert search_cache.SUGGEST_CACHE.get("g") is None + + +def test_invalidate_user_suggestions_only_drops_that_user(): + search_cache.SUGGEST_CACHE.set(("u1", "docling", "{}", 0, True), {"files": ["a"]}) + search_cache.SUGGEST_CACHE.set(("u2", "docling", "{}", 0, True), {"files": ["b"]}) + + search_cache.invalidate_user_suggestions("u1") + + assert search_cache.SUGGEST_CACHE.get(("u1", "docling", "{}", 0, True)) is None + assert search_cache.SUGGEST_CACHE.get(("u2", "docling", "{}", 0, True)) == {"files": ["b"]} + + +def test_invalidate_user_suggestions_ignores_non_tuple_and_empty_keys(): + search_cache.SUGGEST_CACHE.set(("u1", "docling", "{}", 0, True), {"files": ["a"]}) + search_cache.SUGGEST_CACHE.set("g", {"files": ["g"]}) + search_cache.SUGGEST_CACHE.set((), {"files": ["empty"]}) + search_cache.SUGGEST_CACHE.set(42, {"files": ["num"]}) + + # This should not raise IndexError, TypeError, etc. + search_cache.invalidate_user_suggestions("u1") + + assert search_cache.SUGGEST_CACHE.get(("u1", "docling", "{}", 0, True)) is None + assert search_cache.SUGGEST_CACHE.get("g") == {"files": ["g"]} + assert search_cache.SUGGEST_CACHE.get(()) == {"files": ["empty"]} + assert search_cache.SUGGEST_CACHE.get(42) == {"files": ["num"]} + + +@pytest.mark.asyncio +async def test_discard_where_keeps_an_inflight_computation_from_storing(): + cache = SingleflightTTLCache(maxsize=8, ttl=60) + release = asyncio.Event() + + async def compute(): + await release.wait() + return "stale" + + leader = asyncio.create_task(cache.get_or_compute(("u1", "q"), compute)) + await asyncio.sleep(0) + + cache.discard_where(lambda key: key[0] == "u1") + release.set() + assert await leader == "stale" + + assert cache.get(("u1", "q")) is None diff --git a/tests/unit/test_search_cache_invalidation.py b/tests/unit/test_search_cache_invalidation.py new file mode 100644 index 000000000..3fc2a74e4 --- /dev/null +++ b/tests/unit/test_search_cache_invalidation.py @@ -0,0 +1,156 @@ +"""Writes that change what is searchable must drop the corpus-dependent caches.""" + +import pytest + +from services import search_cache +from services.document_index_writer import ( + DocumentIndexChunk, + DocumentIndexContext, + DocumentIndexWriter, +) +from utils.opensearch_delete import delete_document_ids + + +def _prime_caches() -> None: + search_cache.SUGGEST_CACHE.set("g", {"files": ["old.pdf"]}) + search_cache.SPACES_CACHE.set("s", ["space"]) + search_cache.QUERY_EMBEDDING_CACHE.set("q", [0.1]) + + +def _assert_corpus_caches_dropped() -> None: + assert search_cache.SUGGEST_CACHE.get("g") is None + assert search_cache.SPACES_CACHE.get("s") is None + assert search_cache.QUERY_EMBEDDING_CACHE.get("q") == [0.1] + + +class _Indices: + async def exists(self, *, index): + return True + + async def get_mapping(self, *, index): + return {} + + async def refresh(self, *, index): + return None + + +class _WriteClient: + def __init__(self, bulk_errors=False): + self.indices = _Indices() + self.bulk_errors = bulk_errors + + async def bulk(self, *, body, refresh): + items = [ + {"index": {"_id": body[i]["index"]["_id"], "status": 500, "error": {"type": "boom"}}} + if self.bulk_errors + else {"index": {"_id": body[i]["index"]["_id"], "status": 200}} + for i in range(0, len(body), 2) + ] + return {"errors": self.bulk_errors, "items": items} + + async def delete_by_query(self, **_kwargs): + return {"deleted": 2} + + +def _writer(monkeypatch, client): + writer = DocumentIndexWriter(opensearch_client=client) + + async def no_op_ensure(*_args, **_kwargs): + return "embedding_field" + + monkeypatch.setattr(writer, "_ensure_index_and_embedding_field", no_op_ensure) + monkeypatch.setattr("config.settings.get_index_name", lambda: "documents") + return writer + + +def _context(): + return DocumentIndexContext( + document_id="doc", + filename="report.pdf", + mimetype="application/pdf", + embedding_model="test-model", + owner="user-a", + ) + + +_CHUNKS = [DocumentIndexChunk(chunk_id="doc_0", text="t", vector=[0.1, 0.2])] + + +@pytest.mark.asyncio +async def test_indexing_chunks_drops_corpus_caches(monkeypatch): + writer = _writer(monkeypatch, _WriteClient()) + _prime_caches() + + await writer.index_chunks(_context(), _CHUNKS) + + _assert_corpus_caches_dropped() + + +@pytest.mark.asyncio +async def test_a_failed_batch_still_drops_corpus_caches(monkeypatch): + writer = _writer(monkeypatch, _WriteClient(bulk_errors=True)) + monkeypatch.setattr("services.document_index_writer._BULK_RETRY_DELAY_SECONDS", 0) + _prime_caches() + + with pytest.raises(Exception): # noqa: B017 - any bulk failure type + await writer.index_chunks(_context(), _CHUNKS) + + _assert_corpus_caches_dropped() + + +@pytest.mark.asyncio +async def test_deleting_an_ingest_run_drops_corpus_caches(monkeypatch): + writer = _writer(monkeypatch, _WriteClient()) + _prime_caches() + + await writer.delete_ingest_run("run-1") + + _assert_corpus_caches_dropped() + + +class _DeleteClient: + def __init__(self, results): + self.results = results + + async def delete(self, *, index, id, refresh): + result = self.results[id] + if isinstance(result, Exception): + raise result + return {"result": result} + + +@pytest.mark.asyncio +async def test_deleting_documents_drops_corpus_caches(): + _prime_caches() + + deleted = await delete_document_ids( + _DeleteClient({"a": "deleted"}), index="documents", document_ids=["a"] + ) + + assert deleted == 1 + _assert_corpus_caches_dropped() + + +@pytest.mark.asyncio +async def test_a_delete_that_removed_nothing_keeps_the_caches(): + _prime_caches() + + await delete_document_ids( + _DeleteClient({"a": "not_found"}), index="documents", document_ids=["a"] + ) + + assert search_cache.SUGGEST_CACHE.get("g") is not None + + +@pytest.mark.asyncio +async def test_a_delete_that_fails_midway_still_drops_caches(): + _prime_caches() + + with pytest.raises(RuntimeError): + await delete_document_ids( + _DeleteClient({"a": "deleted", "b": RuntimeError("down")}), + index="documents", + document_ids=["a", "b"], + ) + + _assert_corpus_caches_dropped() diff --git a/tests/unit/test_search_service_suggest.py b/tests/unit/test_search_service_suggest.py new file mode 100644 index 000000000..2ef601e74 --- /dev/null +++ b/tests/unit/test_search_service_suggest.py @@ -0,0 +1,382 @@ +"""Unit tests for SearchService.suggest and the embedding/space caches.""" + +from types import SimpleNamespace + +import pytest + +from auth_context import set_auth_context, set_score_threshold, set_search_filters, set_search_limit +from services.search_cache import clear_search_caches +from services.search_service import SearchService, _filters_to_clauses, _group_suggest_files + + +@pytest.fixture(autouse=True) +def _fresh_caches(): + clear_search_caches() + yield + clear_search_caches() + + +class _OpenSearch: + """Routes the three request shapes the suggest path produces.""" + + def __init__(self, tag="u1"): + self.tag = tag + self.keyword_calls = 0 + self.count_calls = 0 + self.space_calls = 0 + self.bodies = [] + + async def search(self, *, index, body, params): + self.bodies.append(body) + aggs = body.get("aggs", {}) + if "candidates" in aggs: + self.count_calls += 1 + names = [ + f["match_phrase_prefix"]["text"]["query"] + for f in aggs["candidates"]["filters"]["filters"] + ] + return { + "aggregations": { + "candidates": { + "buckets": [ + {"doc_count": 61 if i == 0 else 0, "files": {"value": 2}} + for i, _ in enumerate(names) + ] + } + } + } + if body.get("size") == 0 and "embedding_spaces" in aggs: + self.space_calls += 1 + return { + "aggregations": { + "embedding_spaces": { + "buckets": [{"key": {"space_id": "openai:te3"}, "doc_count": 1}] + } + } + } + if body.get("size") == 0 and "legacy_embedding_models" in aggs: + return {"aggregations": {"legacy_embedding_models": {"buckets": []}}} + self.keyword_calls += 1 + return { + "hits": { + "hits": [ + { + "_id": f"{self.tag}-1", + "_source": { + "filename": f"{self.tag}-docling.pdf", + "mimetype": "application/pdf", + "connector_type": "local", + "text": "Run docling serve to start. Docling converts files.", + }, + "highlight": {"text": ["Run docling serve to start"]}, + }, + { + "_id": f"{self.tag}-2", + "_source": { + "filename": f"{self.tag}-notes.md", + "mimetype": "text/markdown", + "connector_type": "local", + "text": "Unrelated notes about lunch plans for the team offsite.", + }, + }, + ] + }, + "aggregations": {}, + } + + +def _service(opensearch_by_user, monkeypatch, embed_calls=None): + service = object.__new__(SearchService) + service.session_manager = SimpleNamespace( + get_user_opensearch_client=lambda user_id, jwt: opensearch_by_user[user_id] + ) + service.models_service = None + + async def create_embedding(body, *, interactive=False): + if embed_calls is not None: + embed_calls.append(body["input"][0]) + return {"data": [{"embedding": [0.1, 0.2]}]} + + monkeypatch.setattr("services.search_service.gateway_embeddings", create_embedding) + monkeypatch.setattr("services.search_service.get_index_name", lambda: "documents") + monkeypatch.setattr("services.search_service.get_embedding_model", lambda: "te3") + monkeypatch.setattr( + "services.search_service.get_openrag_config", + lambda: SimpleNamespace( + knowledge=SimpleNamespace( + embedding_provider="openai", legacy_embedding_provider_map={} + ), + providers=SimpleNamespace(ollama=SimpleNamespace(endpoint=None)), + ), + ) + return service + + +@pytest.mark.asyncio +async def test_keyword_phase_returns_suggestions_files_and_top_previews(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + + result = await service.suggest("docl", user_id="u1") + + assert result["phase"] == "keyword" + assert result["suggestions"][0] == {"text": "docling", "files": 2, "chunks": 61} + assert [f["filename"] for f in result["files"]] == ["u1-docling.pdf", "u1-notes.md"] + first = result["files"][0]["preview"] + assert first["kind"] == "keyword" + assert "docling" in first["text"] + # No keyword match and semantic is off: no guessed preview. + assert result["files"][1]["preview"] is None + + +@pytest.mark.asyncio +async def test_keyword_phase_never_embeds(monkeypatch): + embed_calls = [] + service = _service({"u1": _OpenSearch()}, monkeypatch, embed_calls) + + await service.suggest("docling", user_id="u1") + + assert embed_calls == [] + + +@pytest.mark.asyncio +async def test_repeat_request_is_served_from_cache(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + + first = await service.suggest("Docling", user_id="u1") + second = await service.suggest(" docling ", user_id="u1") + + assert first == second + assert opensearch.keyword_calls == 1 + assert opensearch.count_calls == 1 + + +@pytest.mark.asyncio +async def test_cache_is_isolated_between_users(monkeypatch): + u1, u2 = _OpenSearch("u1"), _OpenSearch("u2") + service = _service({"u1": u1, "u2": u2}, monkeypatch) + + r1 = await service.suggest("docling", user_id="u1") + r2 = await service.suggest("docling", user_id="u2") + + assert r1["files"][0]["filename"] == "u1-docling.pdf" + assert r2["files"][0]["filename"] == "u2-docling.pdf" + assert u2.keyword_calls == 1 + + +@pytest.mark.asyncio +async def test_semantic_phase_embeds_once_and_upgrades_keyword_cache(monkeypatch): + embed_calls = [] + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch, embed_calls) + set_auth_context("u1", "jwt") + set_search_filters({}) + set_search_limit(10) + set_score_threshold(0) + + semantic = await service.suggest("docling", user_id="u1", jwt_token="jwt", semantic=True) + assert semantic["phase"] == "semantic" + assert embed_calls == ["docling"] + # With semantic on, the unmatched file gets a best-sentence preview. + notes = next(f for f in semantic["files"] if f["filename"] == "u1-notes.md") + assert notes["preview"]["kind"] == "semantic" + + keyword_calls_before = opensearch.keyword_calls + again = await service.suggest("docling", user_id="u1", jwt_token="jwt", semantic=False) + assert again is semantic + assert opensearch.keyword_calls == keyword_calls_before + + +@pytest.mark.asyncio +async def test_search_tool_reuses_cached_embedding_and_spaces(monkeypatch): + embed_calls = [] + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch, embed_calls) + set_auth_context("u1", "jwt") + set_search_filters({}) + set_search_limit(10) + set_score_threshold(0) + + await service.search_tool("needle") + await service.search_tool("needle") + await service.search_tool("other") + + assert embed_calls == ["needle", "other"] + assert opensearch.space_calls == 1 + + +@pytest.mark.asyncio +async def test_too_short_query_returns_empty_without_touching_opensearch(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + + result = await service.suggest("d", user_id="u1") + + assert result["suggestions"] == [] and result["files"] == [] + assert opensearch.bodies == [] + + +@pytest.mark.asyncio +async def test_degraded_semantic_response_is_not_cached(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + calls = 0 + + async def degraded_search(query, **kwargs): + nonlocal calls + calls += 1 + return {"results": [], "warnings": [{"code": "embedding_unavailable"}]} + + monkeypatch.setattr(service, "search", degraded_search) + + await service.suggest("docling", user_id="u1", semantic=True) + await service.suggest("docling", user_id="u1", semantic=True) + + assert calls == 2 + + +@pytest.mark.asyncio +async def test_cache_is_keyed_by_score_threshold(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + thresholds = [] + + async def fake_search(query, **kwargs): + thresholds.append(kwargs["score_threshold"]) + return {"results": []} + + monkeypatch.setattr(service, "search", fake_search) + + await service.suggest("docling", user_id="u1", semantic=True, score_threshold=1.0) + await service.suggest("docling", user_id="u1", semantic=True, score_threshold=2.0) + await service.suggest("docling", user_id="u1", semantic=True, score_threshold=2.0) + + assert thresholds == [1.0, 2.0] + + +@pytest.mark.asyncio +async def test_semantic_suggest_tags_its_search_trace_as_suggest(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + sources = [] + + async def fake_search_tool(query, embedding_model=None, *, source): + sources.append(source) + return {"results": []} + + monkeypatch.setattr(service, "search_tool", fake_search_tool) + + await service.suggest("docling", user_id="u1", semantic=True) + + assert sources == ["suggest"] + + +@pytest.mark.asyncio +async def test_search_defaults_its_trace_source_to_api(monkeypatch): + service = _service({"u1": _OpenSearch()}, monkeypatch) + sources = [] + + async def fake_search_tool(query, embedding_model=None, *, source): + sources.append(source) + return {"results": []} + + monkeypatch.setattr(service, "search_tool", fake_search_tool) + + await service.search("docling", user_id="u1") + + assert sources == ["api"] + + +@pytest.mark.asyncio +async def test_keyword_request_is_not_answered_by_other_threshold_semantic_cache(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + + async def fake_search(query, **kwargs): + return {"results": []} + + monkeypatch.setattr(service, "search", fake_search) + + await service.suggest("docling", user_id="u1", semantic=True, score_threshold=1.0) + await service.suggest("docling", user_id="u1", semantic=False, score_threshold=2.0) + + assert opensearch.keyword_calls == 1 + + +@pytest.mark.asyncio +async def test_semantic_suggest_skips_keyword_query_when_threshold_is_set(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + + async def fake_search(query, **kwargs): + return {"results": []} + + monkeypatch.setattr(service, "search", fake_search) + + await service.suggest("docling", user_id="u1", semantic=True, score_threshold=1.0) + assert opensearch.keyword_calls == 0 + + await service.suggest("other", user_id="u1", semantic=True) + assert opensearch.keyword_calls == 1 + + +@pytest.mark.asyncio +async def test_suggest_without_jwt_still_resolves_the_user_client(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + + result = await service.suggest("docling", user_id="u1", jwt_token=None) + + assert result["files"] + assert opensearch.keyword_calls == 1 + + +@pytest.mark.asyncio +async def test_semantic_error_payload_is_returned_and_not_cached(monkeypatch): + opensearch = _OpenSearch() + service = _service({"u1": opensearch}, monkeypatch) + calls = 0 + + async def failing_search(query, **kwargs): + nonlocal calls + calls += 1 + return {"error": "Authentication required"} + + monkeypatch.setattr(service, "search", failing_search) + + first = await service.suggest("docling", user_id="u1", semantic=True) + await service.suggest("docling", user_id="u1", semantic=True) + + assert first == {"error": "Authentication required"} + assert calls == 2 + + +def test_filters_to_clauses_allowlists_fields(): + clauses = _filters_to_clauses( + { + "data_sources": ["a.pdf"], + "owners": ["x", "y"], + "connector_types": [], + "_index": ["evil"], + "document_types": "not-a-list", + } + ) + assert {"term": {"filename": "a.pdf"}} in clauses + assert {"terms": {"owner": ["x", "y"]}} in clauses + assert {"term": {"connector_type": "__IMPOSSIBLE_VALUE__"}} in clauses + assert len(clauses) == 3 + + +def test_group_suggest_files_limits_rows_and_previews(): + chunks = [ + { + "filename": f"f{i}.txt", + "text": "docling text here", + "highlights": ["docling text"], + } + for i in range(8) + ] + files = _group_suggest_files(chunks, "docling", allow_semantic=False) + assert len(files) == 5 + assert [f["preview"] is not None for f in files] == [True, True, True, False, False] diff --git a/tests/unit/test_search_suggest_endpoint.py b/tests/unit/test_search_suggest_endpoint.py new file mode 100644 index 000000000..7f18b2e34 --- /dev/null +++ b/tests/unit/test_search_suggest_endpoint.py @@ -0,0 +1,126 @@ +"""Endpoint-level tests for POST /search/suggest (api/search.py).""" + +import json +from types import SimpleNamespace + +import pytest + +from api.search import SUGGEST_MAX_FILTER_VALUES, SuggestBody, suggest +from utils.opensearch_utils import DISK_SPACE_ERROR_MESSAGE, OpenSearchDiskSpaceError + + +class _Service: + def __init__(self, result=None, exc=None): + self.result, self.exc, self.calls = result, exc, [] + + async def suggest(self, query, **kwargs): + self.calls.append((query, kwargs)) + if self.exc: + raise self.exc + return self.result + + +def _user(): + return SimpleNamespace(user_id="u1", jwt_token="jwt") + + +def _body(response): + return json.loads(response.body) + + +@pytest.mark.asyncio +async def test_forwards_request_and_returns_payload(): + payload = {"phase": "keyword", "query": "docling", "suggestions": [], "files": []} + service = _Service(payload) + + response = await suggest( + SuggestBody(query="docling", filters={"owners": ["x"]}, semantic=True, scoreThreshold=2), + search_service=service, + user=_user(), + ) + + assert response.status_code == 200 + assert _body(response) == payload + assert service.calls == [ + ( + "docling", + { + "user_id": "u1", + "jwt_token": "jwt", + "filters": {"owners": ["x"]}, + "semantic": True, + "score_threshold": 2, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_internal_errors_are_not_leaked_to_the_client(): + response = await suggest( + SuggestBody(query="docling"), + search_service=_Service(exc=RuntimeError("secret connection string")), + user=_user(), + ) + + assert response.status_code == 500 + assert _body(response) == {"error": "Suggestions unavailable"} + + +@pytest.mark.asyncio +async def test_service_error_payload_becomes_generic_502(): + response = await suggest( + SuggestBody(query="docling"), + search_service=_Service({"error": "embedding provider detail"}), + user=_user(), + ) + + assert response.status_code == 502 + assert "embedding provider detail" not in response.body.decode() + + +@pytest.mark.asyncio +async def test_disk_space_error_maps_to_507(): + response = await suggest( + SuggestBody(query="docling"), + search_service=_Service(exc=OpenSearchDiskSpaceError(DISK_SPACE_ERROR_MESSAGE)), + user=_user(), + ) + + assert response.status_code == 507 + assert _body(response) == {"error": DISK_SPACE_ERROR_MESSAGE} + + +@pytest.mark.asyncio +async def test_missing_jwt_is_forwarded_as_none(): + service = _Service({"phase": "keyword", "query": "docling", "suggestions": [], "files": []}) + + response = await suggest( + SuggestBody(query="docling"), + search_service=service, + user=SimpleNamespace(user_id="u1", jwt_token=None), + ) + + assert response.status_code == 200 + assert service.calls[0][1]["jwt_token"] is None + + +def test_request_bounds_are_enforced(): + with pytest.raises(ValueError): + SuggestBody(query="x" * 201) + with pytest.raises(ValueError): + SuggestBody(query="ok", scoreThreshold=-1) + with pytest.raises(ValueError): + SuggestBody(query="ok", scoreThreshold=101) + + +def test_filter_bounds_are_enforced(): + SuggestBody(query="ok", filters={"owners": ["x"] * SUGGEST_MAX_FILTER_VALUES}) + with pytest.raises(ValueError): + SuggestBody(query="ok", filters={"owners": ["x"] * (SUGGEST_MAX_FILTER_VALUES + 1)}) + with pytest.raises(ValueError): + SuggestBody(query="ok", filters={f"k{i}": [] for i in range(11)}) + + +def test_score_threshold_is_accepted_by_camel_case_name(): + assert SuggestBody.model_validate({"query": "ok", "scoreThreshold": 1.5}).scoreThreshold == 1.5 diff --git a/tests/unit/test_search_suggest_helpers.py b/tests/unit/test_search_suggest_helpers.py new file mode 100644 index 000000000..55ac9c6f8 --- /dev/null +++ b/tests/unit/test_search_suggest_helpers.py @@ -0,0 +1,122 @@ +"""Unit tests for services/search_suggest.py (pure helpers).""" + +from typing import Any, cast + +from services.search_suggest import ( + PREVIEW_KEYWORD, + PREVIEW_MAX_CHARS, + PREVIEW_SEMANTIC, + best_sentence, + build_preview, + escape_wildcard, + extract_candidates, + normalize_suggest_query, + preview_from_highlight, + strip_marks, +) + + +class TestNormalizeSuggestQuery: + def test_collapses_whitespace(self): + assert normalize_suggest_query(" docling serve ") == "docling serve" + + def test_rejects_short_wildcard_and_non_strings(self): + assert normalize_suggest_query("a") == "" + assert normalize_suggest_query("*") == "" + assert normalize_suggest_query(None) == "" + assert normalize_suggest_query(cast(Any, 123)) == "" + + def test_caps_length(self): + assert len(normalize_suggest_query("x" * 500)) == 100 + + +def test_escape_wildcard_neutralizes_metacharacters(): + assert escape_wildcard("a*b?c\\d") == "a\\*b\\?c\\\\d" + + +class TestPreviewFromHighlight: + def test_adds_ellipsis_when_fragment_is_mid_chunk(self): + text = "Intro text. The docling library parses PDFs. More text follows here." + fragment = "The docling library parses PDFs" + line = preview_from_highlight(text, fragment) + assert line.startswith("…") + assert line.endswith("…") + assert "docling" in line + + def test_no_ellipsis_when_fragment_is_the_whole_chunk(self): + text = "docling parses PDFs" + assert preview_from_highlight(text, "docling parses PDFs") == ( + "docling parses PDFs" + ) + + def test_long_fragment_is_trimmed_but_keeps_the_match_and_balanced_marks(self): + text = "x " * 400 + "needle " + "y " * 400 + fragment = "x " * 100 + "needle " + "y " * 100 + line = preview_from_highlight(text, fragment) + assert "needle" in line + assert len(strip_marks(line)) <= PREVIEW_MAX_CHARS + 4 + assert line.count("") == line.count("") + + +class TestBestSentence: + def test_picks_sentence_with_most_query_words(self): + text = ( + "The weather was pleasant all week long. " + "Vector search ranks passages by embedding similarity. " + "Lunch was served at noon every single day." + ) + assert best_sentence(text, "embedding similarity") == ( + "Vector search ranks passages by embedding similarity." + ) + + def test_falls_back_to_first_sentence_without_overlap(self): + text = "First sentence is right here. Second sentence follows it." + assert best_sentence(text, "zzz") == "First sentence is right here." + + def test_empty_text(self): + assert best_sentence(" ", "x") is None + + +class TestBuildPreview: + def test_prefers_keyword_highlight(self): + preview = build_preview("t", ["a b c"], "b", allow_semantic=True) + assert preview is not None + assert preview["kind"] == PREVIEW_KEYWORD + + def test_semantic_fallback_only_when_allowed(self): + text = "Retrieval augments generation with fetched context passages." + assert build_preview(text, [], "retrieval", allow_semantic=False) is None + preview = build_preview(text, [], "retrieval", allow_semantic=True) + assert preview is not None + assert preview["kind"] == PREVIEW_SEMANTIC + + def test_nothing_to_show(self): + assert build_preview(None, [], "x", allow_semantic=True) is None + + +class TestExtractCandidates: + def test_completes_partial_word_and_adds_next_word(self): + texts = [ + "Run docling serve to start the API.", + "docling serve exposes REST endpoints", + "The docling project converts documents.", + ] + candidates = extract_candidates(texts, "docl") + assert "docling" in candidates + assert "docling serve" in candidates + assert candidates.index("docling") < candidates.index("docling serve") + + def test_skips_stopwords_as_second_word(self): + assert extract_candidates(["docling and friends"], "docling") == ["docling"] + + def test_multi_word_query_is_matched_across_whitespace(self): + candidates = extract_candidates(["docling serve works"], "docling serve") + assert candidates[0] == "docling serve" + + def test_regex_metacharacters_in_query_are_literal(self): + assert extract_candidates(["a (b) c"], "(b") == ["(b"] + assert extract_candidates(["anything"], "[") == [] + + def test_ignores_non_string_texts_and_empty_query(self): + assert extract_candidates(cast(Any, [None, 3]), "docling") == [] + assert extract_candidates(["docling"], " ") == []