add document embedding search on document view

This commit is contained in:
timothycarambat
2023-09-14 13:51:00 -07:00
parent 862416b1de
commit b4757763e2
10 changed files with 509 additions and 10 deletions
+38
View File
@@ -22,6 +22,9 @@ const { validEmbedding } = require("../../../utils/tokenizer");
const { documentDeletedJob } = require("../../../utils/jobs/documentDeleteJob");
const { cloneDocumentJob } = require("../../../utils/jobs/cloneDocumentJob");
const { selectConnector } = require("../../../utils/vectordatabases/providers");
const {
documentEmbeddingSearch,
} = require("../../../utils/search/documentEmbeddings");
process.env.NODE_ENV === "development"
? require("dotenv").config({ path: `.env.${process.env.NODE_ENV}` })
@@ -325,6 +328,41 @@ function documentEndpoints(app) {
}
}
);
app.get(
"/v1/documents/:documentId/search-embeddings",
[validSessionForUser],
async function (request, response) {
try {
const { documentId } = request.params;
const { method, q: query } = request.query;
const user = await userFromSession(request);
if (!user) {
response.sendStatus(403).end();
return;
}
const document = await WorkspaceDocument.get(`id = ${documentId}`);
if (!document) {
response.status(200).json({
fragments: [],
error: "No document found.",
});
return;
}
const { fragments, error } = await documentEmbeddingSearch(
document,
method,
query
);
response.status(200).json({ fragments, error });
} catch (e) {
console.log(e.message, e);
response.sendStatus(500).end();
}
}
);
}
module.exports = { documentEndpoints };
@@ -0,0 +1,44 @@
const { DocumentVectors } = require("../../../models/documentVectors");
const { WorkspaceDocument } = require("../../../models/workspaceDocument");
const { readJSON } = require("../../storage");
// Dirty, but works fast for most cases. Wont be perfect but also not something we should rely
// heavily on for exact text searching.
function fuzzyMatch(pattern, str) {
pattern = ".*" + pattern.split("").join(".*") + ".*";
const re = new RegExp(pattern);
return re.test(str);
}
async function findTextInDoc(wsDoc, query) {
try {
const fragmentIds = [];
const data = await readJSON(WorkspaceDocument.vectorFilepath(wsDoc));
for (const chunk of data) {
if (!chunk.hasOwnProperty("metadata")) continue;
for (const value of Object.values(chunk?.metadata)) {
const valid = fuzzyMatch(query, String(value));
if (valid) fragmentIds.push(chunk.vectorDbId);
}
}
return fragmentIds;
} catch (e) {
console.error(e);
return [];
}
}
async function exactTextSearch(document, query) {
const matchingVectorIds = await findTextInDoc(document, query);
if (matchingVectorIds.length === 0) return { fragments: [], error: null };
const queryString = matchingVectorIds.map((vid) => `'${vid}'`).join(",");
const fragments = await DocumentVectors.where(`vectorId IN (${queryString})`);
return { fragments, error: null };
}
module.exports = {
exactTextSearch,
};
@@ -0,0 +1,32 @@
const { Telemetry } = require("../../../models/telemetry");
const { exactTextSearch } = require("./exactText");
const { metadataSearch } = require("./metadata");
const { semanticSearch } = require("./semantic");
const { vectorIdSearch } = require("./vectorId");
const SEARCH_METHODS = {
semantic: semanticSearch,
exactText: exactTextSearch,
metadata: metadataSearch,
vectorId: vectorIdSearch,
};
function validSearchMethod(method) {
return Object.keys(SEARCH_METHODS).includes(method);
}
async function documentEmbeddingSearch(document, method, query) {
try {
if (!validSearchMethod(method))
throw new Error(`Invalid search method ${method}`);
await Telemetry.sendTelemetry("search_executed", { searchMethod: method });
return await SEARCH_METHODS[method](document, decodeURIComponent(query));
} catch (e) {
console.error("Workspace document search", e.message);
return { fragments: [], error: e.message };
}
}
module.exports = {
documentEmbeddingSearch,
};
@@ -0,0 +1,50 @@
const { DocumentVectors } = require("../../../models/documentVectors");
const { WorkspaceDocument } = require("../../../models/workspaceDocument");
const { readJSON } = require("../../storage");
// Dirty, but works fast for most cases. Wont be perfect but also not something we should rely
// heavily on for exact text searching.
function fuzzyMatch(pattern, str) {
pattern = ".*" + pattern.split("").join(".*") + ".*";
const re = new RegExp(pattern);
return re.test(str);
}
async function findKeyValueInDoc(wsDoc, query) {
try {
const fragmentIds = [];
const data = await readJSON(WorkspaceDocument.vectorFilepath(wsDoc));
const [keyToFind, valueToFind] = query.split(":");
for (const chunk of data) {
if (!chunk.hasOwnProperty("metadata")) continue;
for (const [key, value] of Object.entries(chunk?.metadata)) {
const validKey = fuzzyMatch(keyToFind, key);
if (!validKey) continue;
const match = fuzzyMatch(valueToFind, String(value));
if (match) fragmentIds.push(chunk.vectorDbId);
}
}
return fragmentIds;
} catch (e) {
console.error(e);
return [];
}
}
async function metadataSearch(document, query) {
const matchingVectorIds = await findKeyValueInDoc(document, query);
if (matchingVectorIds.length === 0) return { fragments: [], error: null };
const queryString = matchingVectorIds.map((vid) => `'${vid}'`).join(",");
const fragments = await DocumentVectors.where(
`vectorId IN (${queryString})`,
200
);
return { fragments, error: null };
}
module.exports = {
metadataSearch,
};
@@ -0,0 +1,53 @@
const { DocumentVectors } = require("../../../models/documentVectors");
const {
OrganizationConnection,
} = require("../../../models/organizationConnection");
const {
OrganizationWorkspace,
} = require("../../../models/organizationWorkspace");
const { SystemSettings } = require("../../../models/systemSettings");
const { WorkspaceDocument } = require("../../../models/workspaceDocument");
const { OpenAi } = require("../../openAi");
const { selectConnector } = require("../../vectordatabases/providers");
async function semanticSearch(document, query) {
const workspace = await OrganizationWorkspace.get(
`id = ${document.workspace_id}`
);
const connector = await OrganizationConnection.get(
`organization_id = ${document.organization_id}`
);
if (!connector)
return { fragments: [], error: "No connector found for org." };
const openAiKey = (await SystemSettings.get(`label = 'open_ai_api_key'`))
?.value;
if (!openAiKey)
return { fragments: [], error: "No OpenAI key available to embed query." };
const vectorDb = selectConnector(connector);
const openai = new OpenAi(openAiKey);
const queryVector = await openai.embedTextChunk(query);
if (!queryVector) return { fragments: [], error: "Failed to embed query." };
// Execute Similarity search for vector DB provider so we can find inferred documents.
const searchResults = await vectorDb.similarityResponse(
workspace.slug,
queryVector
);
// From similarity search we can find all document vector DB items to infer their associated
// document record.
const searchString = searchResults.vectorIds
.map((vid) => `'${vid}'`)
.join(",");
const fragments = await DocumentVectors.where(
`vectorId IN (${searchString})`
);
return { fragments, error: null };
}
module.exports = {
semanticSearch,
};
@@ -0,0 +1,12 @@
const { DocumentVectors } = require("../../../models/documentVectors");
async function vectorIdSearch(_document, query) {
const documentVector = await DocumentVectors.get(`vectorId = '${query}'`);
if (!documentVector)
return { fragments: [], error: "No document vector found with that id." };
return { fragments: [documentVector], error: null };
}
module.exports = {
vectorIdSearch,
};
+22
View File
@@ -1,3 +1,4 @@
import { ISearchTypes } from '../pages/DocumentView/FragmentList/SearchView';
import { API_BASE } from '../utils/constants';
import { baseHeaders } from '../utils/request';
@@ -106,6 +107,27 @@ const Document = {
return { success: false, error: e.message };
});
},
searchEmbeddings: async (
documentId: number,
method: ISearchTypes,
query: string
): Promise<{ documents: object[] }> => {
const searchEndpoint = new URL(
`${API_BASE}/v1/documents/${documentId}/search-embeddings`
);
searchEndpoint.searchParams.append('method', method);
searchEndpoint.searchParams.append('q', encodeURIComponent(query));
return await fetch(searchEndpoint, {
method: 'GET',
headers: baseHeaders(),
})
.then((res) => res.json())
.then((res) => res?.fragments || [])
.catch((e) => {
console.error(e.message);
return [];
});
},
};
export default Document;
@@ -0,0 +1,235 @@
import {
Dispatch,
SetStateAction,
SyntheticEvent,
useRef,
useState,
} from 'react';
import { ChevronDown, Search, Loader } from 'react-feather';
import Document from '../../../../models/document';
export type ISearchTypes = 'semantic' | 'exactText' | 'metadata' | 'vectorId';
const SEARCH_MODES = {
exactText: {
display: 'Fuzzy Text Search',
placeholder: 'Find embedding via a fuzzy text match on your query.',
},
semantic: {
display: 'Semantic Search',
placeholder:
'Search with natural language finding the most similar embedding by meaning. Use of this search will cost OpenAI credits to embed the query.',
},
metadata: {
display: 'Metadata',
placeholder:
'Find embedding by exact key:value pair. Formatted as key:value_to_look_for',
},
vectorId: {
display: 'Vector Id',
placeholder: 'Find by a specific vector ID',
},
};
export default function SearchView({
searchMode,
setSearchMode,
document,
FragmentItem,
canEdit,
}: {
searchMode: boolean;
document: object;
setSearchMode: Dispatch<SetStateAction<boolean>>;
FragmentItem: (props: any) => JSX.Element;
canEdit: boolean;
}) {
const formEl = useRef<HTMLFormElement>(null);
const [searching, setSearching] = useState(false);
const [showSearchMethods, setShowSearchMethods] = useState(false);
const [searchBy, setSearchBy] = useState<ISearchTypes>('exactText');
const [searchTerm, setSearchTerm] = useState<string>('');
const [fragments, setFragments] = useState([]);
const [sourceDoc, setSourceDoc] = useState(null);
const clearSearch = () => {
setSearchBy('exactText');
setSearchTerm('');
setFragments([]);
setSearching(false);
setSearchMode(false);
setSourceDoc(null);
(formEl.current as HTMLFormElement).reset();
};
const handleSearch = async (e: SyntheticEvent<HTMLElement, SubmitEvent>) => {
e.preventDefault();
setSearchMode(true);
const formData = new FormData(e.target as any);
const query = formData.get('query') as string;
setSearching(true);
setSearchTerm(query);
const matches = await Document.searchEmbeddings(
document.id,
searchBy,
query
);
const vectorIds = matches.map((fragment) => fragment.vectorId);
const metadataForIds = await Document.metadatas(document.id, vectorIds);
setSourceDoc(metadataForIds);
setFragments(matches);
setSearching(false);
};
return (
<div className="w-full flex-1 rounded-sm py-6">
<div className="flex items-center">
<form ref={formEl} onSubmit={handleSearch} className="w-full">
<div className="relative flex">
<button
onClick={() => setShowSearchMethods(!showSearchMethods)}
className="z-10 inline-flex flex-shrink-0 items-center rounded-l-lg border border-gray-300 bg-gray-100 px-4 py-2.5 text-center text-sm font-medium text-gray-900 hover:bg-gray-200 focus:outline-none focus:ring-4 focus:ring-gray-100 dark:border-gray-600 dark:bg-gray-700 dark:text-white dark:hover:bg-gray-600 dark:focus:ring-gray-700"
type="button"
>
{SEARCH_MODES[searchBy].display}
<ChevronDown size={18} />
</button>
<div
hidden={!showSearchMethods}
className="absolute left-0 top-12 z-99 w-44 divide-y divide-gray-100 rounded-lg bg-white shadow dark:bg-gray-700"
>
<ul
className="py-2 text-sm text-gray-700 dark:text-gray-200"
aria-labelledby="dropdown-button"
>
{Object.keys(SEARCH_MODES).map((_key, i) => {
const method = _key as ISearchTypes;
return (
<li key={i}>
<button
onClick={() => {
setSearchBy(method);
setShowSearchMethods(false);
setFragments([]);
}}
type="button"
className="inline-flex w-full px-4 py-2 hover:bg-gray-100 dark:hover:bg-gray-600 dark:hover:text-white"
>
{SEARCH_MODES[method].display}
</button>
</li>
);
})}
</ul>
</div>
<div className="relative w-full">
<input
type="search"
name="query"
className="z-20 block w-full rounded-r-lg border border-l-2 border-gray-300 border-l-gray-50 bg-gray-50 p-2.5 text-sm text-gray-900 focus:border-blue-500 focus:ring-blue-500 dark:border-gray-600 dark:border-l-gray-700 dark:bg-gray-700 dark:text-white dark:placeholder-gray-400 dark:focus:border-blue-500"
placeholder={SEARCH_MODES[searchBy].placeholder}
required
/>
<button
type="submit"
disabled={searching}
className="absolute right-0 top-0 h-full rounded-r-lg border border-blue-700 bg-blue-700 p-2.5 text-sm font-medium text-white hover:bg-blue-800 focus:outline-none focus:ring-4 focus:ring-blue-300 dark:bg-blue-600 dark:hover:bg-blue-700 dark:focus:ring-blue-800"
>
{searching ? (
<Loader size={18} className="animate-spin" />
) : (
<Search size={18} />
)}
<span className="sr-only">Search</span>
</button>
</div>
<button
onClick={clearSearch}
type="button"
className="ml-2 flex items-center rounded-lg px-4 py-2 text-center text-black hover:bg-gray-200"
>
X
</button>
</div>
</form>
</div>
<div hidden={!searchMode} className="h-auto w-auto">
{searching ? (
<div>
<div className="flex min-h-[40vh] w-full px-8">
<div className="flex flex h-auto w-full flex-col items-center justify-center gap-y-2 rounded-lg bg-slate-50">
<Loader size={15} className="animate-spin rounded-sm" />
<p className="text-sm">
Running {SEARCH_MODES[searchBy].display} for{' '}
<code className="bg-gray-200 px-2">"{searchTerm}"</code>
</p>
</div>
</div>
</div>
) : (
<>
{fragments.length > 0 ? (
<table className="w-full text-left text-sm text-gray-500 dark:text-gray-400">
<thead className="bg-gray-50 text-xs uppercase text-gray-700 dark:bg-gray-700 dark:text-gray-400">
<tr>
<th scope="col" className="px-6 py-3">
#
</th>
<th scope="col" className="px-6 py-3">
Vector DB Id
</th>
<th scope="col" className="px-6 py-3">
Text Chunk
</th>
<th scope="col" className="px-6 py-3">
Last Updated
</th>
<th scope="col" className="px-6 py-3">
Actions
</th>
</tr>
</thead>
<tbody>
{fragments.map((fragment) => {
return (
<FragmentItem
key={fragment.id}
fragment={fragment}
sourceDoc={sourceDoc}
canEdit={canEdit}
/>
);
})}
</tbody>
</table>
) : (
<>
<div>
<div className="flex min-h-[40vh] w-full px-8">
<div className="flex flex h-auto w-full flex-col items-center justify-center gap-y-2 rounded-lg bg-slate-50">
{!!searchTerm ? (
<p className="text-sm">
No results on {SEARCH_MODES[searchBy].display} for{' '}
<code className="bg-gray-200 px-2">
"{searchTerm}"
</code>
</p>
) : (
<p className="text-sm">
Type in a query to search for an embedding
</p>
)}
</div>
</div>
</div>
</>
)}
</>
)}
</div>
</div>
);
}
@@ -6,6 +6,7 @@ import moment from 'moment';
import { useParams } from 'react-router-dom';
import paths from '../../../utils/paths';
import DocumentListPagination from '../../../components/DocumentPaginator';
import SearchView from './SearchView';
const DeleteEmbeddingConfirmation = lazy(
() => import('./DeleteEmbeddingConfirmation')
);
@@ -23,6 +24,7 @@ export default function FragmentList({
}) {
const { slug, workspaceSlug } = useParams();
const [loading, setLoading] = useState(true);
const [searchMode, setSearchMode] = useState(false);
const [fragments, setFragments] = useState([]);
const [sourceDoc, setSourceDoc] = useState(null);
const [totalFragments, setTotalFragments] = useState(0);
@@ -99,7 +101,14 @@ export default function FragmentList({
</div>
</div>
<div className="px-6">
<SearchView
searchMode={searchMode}
setSearchMode={setSearchMode}
document={document}
FragmentItem={Fragment}
canEdit={canEdit}
/>
<div hidden={searchMode} className="px-6">
{loading ? (
<div>
<PreLoader />
@@ -140,17 +149,19 @@ export default function FragmentList({
</table>
)}
</div>
<DocumentListPagination
pageCount={totalPages}
currentPage={currentPage}
gotoPage={handlePageChange}
/>
{!searchMode && (
<DocumentListPagination
pageCount={totalPages}
currentPage={currentPage}
gotoPage={handlePageChange}
/>
)}
</div>
</>
);
}
const Fragment = ({
export const Fragment = ({
fragment,
sourceDoc,
canEdit,
@@ -242,7 +253,7 @@ const Fragment = ({
);
};
const FullTextWindow = memo(
export const FullTextWindow = memo(
({ data, fragment }: { data: any; fragment: any }) => {
return (
<dialog id={`${fragment.id}-text`} className="w-1/2 rounded-lg">
@@ -1,4 +1,4 @@
import { SyntheticEvent, useState } from 'react';
import { SyntheticEvent, useRef, useState } from 'react';
import { ChevronDown, FileText, Search, Loader } from 'react-feather';
import { CopyDocToModal } from '..';
import truncate from 'truncate';
@@ -42,17 +42,19 @@ export default function SearchView({
stopSearching: VoidFunction;
deleteDocument: (documentId: number) => void;
}) {
const formEl = useRef<HTMLFormElement>(null);
const [searching, setSearching] = useState(false);
const [showSearchMethods, setShowSearchMethods] = useState(false);
const [searchBy, setSearchBy] = useState<ISearchTypes>('exactText');
const [searchTerm, setSearchTerm] = useState<string>('');
const [documents, setDocuments] = useState([]);
const clearSearch = () => {
setSearchBy('semantic');
setSearchBy('exactText');
setSearchTerm('');
setDocuments([]);
setSearching(false);
stopSearching();
(formEl.current as HTMLFormElement).reset();
};
const handleSearch = async (e: SyntheticEvent<HTMLElement, SubmitEvent>) => {
e.preventDefault();