diff --git a/backend/endpoints/v1/documents/index.js b/backend/endpoints/v1/documents/index.js index c32434c..3a55afb 100644 --- a/backend/endpoints/v1/documents/index.js +++ b/backend/endpoints/v1/documents/index.js @@ -22,6 +22,7 @@ const { readJSON } = require("../../../utils/storage"); const { validEmbedding } = require("../../../utils/tokenizer"); const { documentDeletedJob } = require("../../../utils/jobs/documentDeleteJob"); const { cloneDocumentJob } = require("../../../utils/jobs/cloneDocumentJob"); +const { selectConnector } = require("../../../utils/vectordatabases/providers"); process.env.NODE_ENV === "development" ? require("dotenv").config({ path: `.env.${process.env.NODE_ENV}` }) @@ -224,12 +225,13 @@ function documentEndpoints(app) { } ); - app.get( - "/v1/document/:id/source", + app.post( + "/v1/document/:id/metadatas", [validSessionForUser], async function (request, response) { try { const { id } = request.params; + const { vectorIds } = reqBody(request); const user = await userFromSession(request); if (!user) { response.sendStatus(403).end(); @@ -237,19 +239,26 @@ function documentEndpoints(app) { } const document = await WorkspaceDocument.get(`id = ${id}`); - const filepath = WorkspaceDocument.vectorFilepath(document); - console.log(filepath); - const source = await readJSON(filepath).then((res) => { - const data = {}; - Object.values(res).map((d) => { - data[d.vectorDbId] = { ...d }; - }); - return data; - }); + if (!document) { + response.sendStatus(404).end(); + return; + } - // Set cache-control to retain this file for 1 hour. - response.set("Cache-control", `public, max-age=${60 * 60}`); - response.status(200).json({ ...source }); + const workspace = await OrganizationWorkspace.get( + `id = ${document.workspace_id}` + ); + const connector = await OrganizationConnection.get( + `organization_id = ${document.organization_id}` + ); + + const VectorDb = selectConnector(connector); + const results = await VectorDb.getMetadata(workspace.slug, vectorIds); + const items = {}; + + results?.forEach((metadata) => { + items[metadata.vectorId] = { metadata }; + }); + response.status(200).json(items); } catch (e) { console.log(e); response.sendStatus(500).end(); diff --git a/backend/utils/vectordatabases/providers/chroma/index.js b/backend/utils/vectordatabases/providers/chroma/index.js index 10433e9..b4bf5cb 100644 --- a/backend/utils/vectordatabases/providers/chroma/index.js +++ b/backend/utils/vectordatabases/providers/chroma/index.js @@ -211,6 +211,22 @@ class Chroma { return { success: false, message: e.message }; } } + + async getMetadata(namespace = "", vectorIds = []) { + const { client } = await this.connect(); + const collection = await client.getCollection({ name: namespace }); + const result = await collection.get({ + ids: vectorIds, + include: ["metadatas", "documents"], + }); + + result?.metadatas?.forEach((metadata, i) => { + metadata.vectorId = vectorIds[i]; + metadata.text = result.documents[i]; + }); + + return result.metadatas; + } } module.exports.Chroma = Chroma; diff --git a/backend/utils/vectordatabases/providers/pinecone/index.js b/backend/utils/vectordatabases/providers/pinecone/index.js index ed300a1..e8d42fa 100644 --- a/backend/utils/vectordatabases/providers/pinecone/index.js +++ b/backend/utils/vectordatabases/providers/pinecone/index.js @@ -363,6 +363,25 @@ class Pinecone { return { success: false, message: e.message }; } } + + async getMetadata(namespace = "", vectorIds = []) { + const { pineconeIndex } = await this.connect(); + + const { vectors } = await pineconeIndex.fetch({ + ids: vectorIds, + namespace, + }); + const metadatas = []; + + Object.values(vectors)?.forEach((vector, i) => { + metadatas.push({ + vectorId: vector.id, + ...(vector?.metadata || {}), + }); + }); + + return metadatas; + } } module.exports.Pinecone = Pinecone; diff --git a/frontend/src/models/document.ts b/frontend/src/models/document.ts index b57f005..700079a 100644 --- a/frontend/src/models/document.ts +++ b/frontend/src/models/document.ts @@ -15,10 +15,11 @@ const Document = { return null; }); }, - source: async (id: string | number) => { - return fetch(`${API_BASE}/v1/document/${id}/source`, { - method: 'GET', - cache: 'default', + metadatas: async (id: string, vectorIds: string[]) => { + return fetch(`${API_BASE}/v1/document/${id}/metadatas`, { + method: 'POST', + cache: 'no-cache', + body: JSON.stringify({ vectorIds }), headers: baseHeaders(), }) .then((res) => res.json()) diff --git a/frontend/src/pages/DocumentView/FragmentList/EditEmbeddingConfirmation/index.tsx b/frontend/src/pages/DocumentView/FragmentList/EditEmbeddingConfirmation/index.tsx index f73f8c2..e5bbbff 100644 --- a/frontend/src/pages/DocumentView/FragmentList/EditEmbeddingConfirmation/index.tsx +++ b/frontend/src/pages/DocumentView/FragmentList/EditEmbeddingConfirmation/index.tsx @@ -134,13 +134,7 @@ const EditEmbeddingConfirmation = memo( const debouncedTokenLengthCheck = debounce(checkTokenSize, 500); return ( - - event.target == event.currentTarget && event.currentTarget?.close() - } - > +

Edit embedding

diff --git a/frontend/src/pages/DocumentView/FragmentList/index.tsx b/frontend/src/pages/DocumentView/FragmentList/index.tsx index e4ac70c..596be54 100644 --- a/frontend/src/pages/DocumentView/FragmentList/index.tsx +++ b/frontend/src/pages/DocumentView/FragmentList/index.tsx @@ -54,21 +54,15 @@ export default function FragmentList({ page, PAGE_SIZE ); + const vectorIds = _fragments.map((fragment) => fragment.vectorId); + const metadataForIds = await Document.metadatas(document.id, vectorIds); + setFragments(_fragments); setTotalFragments(totalFragments); + setSourceDoc(metadataForIds); setLoading(false); }; - // Only load source document on page load. - useEffect(() => { - async function downloadSource() { - if (!document?.id) return; - const _src = await Document.source(document.id); - setSourceDoc(_src); - } - downloadSource(); - }, [document]); - useEffect(() => { getFragments(currentPage); }, [document, currentPage]);