diff --git a/backend/endpoints/v1/documents/index.js b/backend/endpoints/v1/documents/index.js index faa762d..601a863 100644 --- a/backend/endpoints/v1/documents/index.js +++ b/backend/endpoints/v1/documents/index.js @@ -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 }; diff --git a/backend/utils/search/documentEmbeddings/exactText.js b/backend/utils/search/documentEmbeddings/exactText.js new file mode 100644 index 0000000..0ee1a88 --- /dev/null +++ b/backend/utils/search/documentEmbeddings/exactText.js @@ -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, +}; diff --git a/backend/utils/search/documentEmbeddings/index.js b/backend/utils/search/documentEmbeddings/index.js new file mode 100644 index 0000000..b053b8d --- /dev/null +++ b/backend/utils/search/documentEmbeddings/index.js @@ -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, +}; diff --git a/backend/utils/search/documentEmbeddings/metadata.js b/backend/utils/search/documentEmbeddings/metadata.js new file mode 100644 index 0000000..06ed0e2 --- /dev/null +++ b/backend/utils/search/documentEmbeddings/metadata.js @@ -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, +}; diff --git a/backend/utils/search/documentEmbeddings/semantic.js b/backend/utils/search/documentEmbeddings/semantic.js new file mode 100644 index 0000000..fe60fcd --- /dev/null +++ b/backend/utils/search/documentEmbeddings/semantic.js @@ -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, +}; diff --git a/backend/utils/search/documentEmbeddings/vectorId.js b/backend/utils/search/documentEmbeddings/vectorId.js new file mode 100644 index 0000000..5112611 --- /dev/null +++ b/backend/utils/search/documentEmbeddings/vectorId.js @@ -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, +}; diff --git a/frontend/src/models/document.ts b/frontend/src/models/document.ts index 700079a..2e63ef7 100644 --- a/frontend/src/models/document.ts +++ b/frontend/src/models/document.ts @@ -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; diff --git a/frontend/src/pages/DocumentView/FragmentList/SearchView/index.tsx b/frontend/src/pages/DocumentView/FragmentList/SearchView/index.tsx new file mode 100644 index 0000000..c528662 --- /dev/null +++ b/frontend/src/pages/DocumentView/FragmentList/SearchView/index.tsx @@ -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>; + FragmentItem: (props: any) => JSX.Element; + canEdit: boolean; +}) { + const formEl = useRef(null); + const [searching, setSearching] = useState(false); + const [showSearchMethods, setShowSearchMethods] = useState(false); + const [searchBy, setSearchBy] = useState('exactText'); + const [searchTerm, setSearchTerm] = useState(''); + 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) => { + 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 ( +
+
+
+
+ + +
+ + +
+ +
+
+
+ + +
+ ); +} diff --git a/frontend/src/pages/DocumentView/FragmentList/index.tsx b/frontend/src/pages/DocumentView/FragmentList/index.tsx index 596be54..5f13d33 100644 --- a/frontend/src/pages/DocumentView/FragmentList/index.tsx +++ b/frontend/src/pages/DocumentView/FragmentList/index.tsx @@ -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({ -
+ + ); } -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 ( diff --git a/frontend/src/pages/WorkspaceDashboard/DocumentsList/SearchView/index.tsx b/frontend/src/pages/WorkspaceDashboard/DocumentsList/SearchView/index.tsx index 97ca607..7cecc22 100644 --- a/frontend/src/pages/WorkspaceDashboard/DocumentsList/SearchView/index.tsx +++ b/frontend/src/pages/WorkspaceDashboard/DocumentsList/SearchView/index.tsx @@ -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(null); const [searching, setSearching] = useState(false); const [showSearchMethods, setShowSearchMethods] = useState(false); const [searchBy, setSearchBy] = useState('exactText'); const [searchTerm, setSearchTerm] = useState(''); 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) => { e.preventDefault();