From 33f533c4f9ce4b1b11ad8558116b81d926a65f77 Mon Sep 17 00:00:00 2001 From: thivy Date: Tue, 13 Feb 2024 17:33:01 +1100 Subject: [PATCH] fix doc id being exposed to client --- .../chat-api/chat-api-dynamic-extensions.ts | 8 ++++- .../chat-api/chat-api-rag-extension.ts | 18 ++++++++--- .../chat-services/chat-api/chat-api-rag.ts | 22 ++----------- .../chat-services/citation-service.ts | 32 +++++++++++++++++-- .../chat-page/chat-services/models.ts | 1 + 5 files changed, 55 insertions(+), 26 deletions(-) diff --git a/src/features/chat-page/chat-services/chat-api/chat-api-dynamic-extensions.ts b/src/features/chat-page/chat-services/chat-api/chat-api-dynamic-extensions.ts index dc3560658..032eae6ff 100644 --- a/src/features/chat-page/chat-services/chat-api/chat-api-dynamic-extensions.ts +++ b/src/features/chat-page/chat-services/chat-api/chat-api-dynamic-extensions.ts @@ -3,6 +3,7 @@ import "server-only"; import { ServerActionResponse } from "@/features/common/server-action-response"; +import { userHashedId } from "@/features/auth-page/helpers"; import { FindAllExtensionForCurrentUser, FindSecureHeaderValue, @@ -84,6 +85,12 @@ async function executeFunction(props: { const headerItems = await Promise.all(headerPromise); + // we need to add the user id to the headers as this is expected by the function and does not have context of the user + headerItems.push({ + id: "authorization", + key: "authorization", + value: await userHashedId(), + }); // map the headers to a dictionary const headers: { [key: string]: string } = headerItems.reduce( (acc: { [key: string]: string }, header) => { @@ -119,7 +126,6 @@ async function executeFunction(props: { if (!response.ok) { return `There was an error calling the api: ${response.statusText}`; } - const result = await response.json(); return { diff --git a/src/features/chat-page/chat-services/chat-api/chat-api-rag-extension.ts b/src/features/chat-page/chat-services/chat-api/chat-api-rag-extension.ts index 64f5080fd..35c61e7f4 100644 --- a/src/features/chat-page/chat-services/chat-api/chat-api-rag-extension.ts +++ b/src/features/chat-page/chat-services/chat-api/chat-api-rag-extension.ts @@ -1,16 +1,16 @@ import { ExtensionSimilaritySearch } from "../azure-ai-search/azure-ai-search"; -import { CreateCitations } from "../citation-service"; +import { CreateCitations, FormatCitations } from "../citation-service"; export const SearchAzureAISimilarDocuments = async (req: Request) => { try { const body = await req.json(); - const search = body.search as string; const vectors = req.headers.get("vectors") as string; const apiKey = req.headers.get("apiKey") as string; const searchName = req.headers.get("searchName") as string; const indexName = req.headers.get("indexName") as string; + const userId = req.headers.get("authorization") as string; const result = await ExtensionSimilaritySearch({ apiKey, @@ -25,10 +25,20 @@ export const SearchAzureAISimilarDocuments = async (req: Request) => { return new Response(JSON.stringify(result)); } - const citationResponse = await CreateCitations(result.response); + const withoutEmbedding = FormatCitations(result.response); + const citationResponse = await CreateCitations(withoutEmbedding, userId); // only get the citations that are ok - const allCitations = citationResponse.filter((c) => c.status === "OK"); + const allCitations = []; + for (const citation of citationResponse) { + if (citation.status === "OK") { + allCitations.push({ + id: citation.response.id, + content: citation.response.content, + }); + } + } + return new Response(JSON.stringify(allCitations)); } catch (e) { console.error("🔴 Retrieving documents", e); diff --git a/src/features/chat-page/chat-services/chat-api/chat-api-rag.ts b/src/features/chat-page/chat-services/chat-api/chat-api-rag.ts index 00d007c62..9b4b37e4f 100644 --- a/src/features/chat-page/chat-services/chat-api/chat-api-rag.ts +++ b/src/features/chat-page/chat-services/chat-api/chat-api-rag.ts @@ -8,11 +8,8 @@ import { ChatCompletionStreamParams, } from "openai/resources/beta/chat/completions"; import { ChatCompletionMessageParam } from "openai/resources/chat/completions"; -import { - DocumentSearchResponse, - SimilaritySearch, -} from "../azure-ai-search/azure-ai-search"; -import { CreateCitations } from "../citation-service"; +import { SimilaritySearch } from "../azure-ai-search/azure-ai-search"; +import { CreateCitations, FormatCitations } from "../citation-service"; import { ChatCitationModel, ChatThreadModel } from "../models"; export const ChatApiRAG = async (props: { @@ -34,20 +31,7 @@ export const ChatApiRAG = async (props: { const documents: ChatCitationModel[] = []; if (documentResponse.status === "OK") { - const withoutEmbedding: DocumentSearchResponse[] = []; - documentResponse.response.forEach((d) => { - withoutEmbedding.push({ - score: d.score, - document: { - metadata: d.document.metadata, - pageContent: d.document.pageContent, - chatThreadId: d.document.chatThreadId, - id: d.document.id, - user: d.document.user, - }, - }); - }); - + const withoutEmbedding = FormatCitations(documentResponse.response); const citationResponse = await CreateCitations(withoutEmbedding); citationResponse.forEach((c) => { diff --git a/src/features/chat-page/chat-services/citation-service.ts b/src/features/chat-page/chat-services/citation-service.ts index 4c5f0f71f..d6163d24b 100644 --- a/src/features/chat-page/chat-services/citation-service.ts +++ b/src/features/chat-page/chat-services/citation-service.ts @@ -1,3 +1,4 @@ +import { userHashedId } from "@/features/auth-page/helpers"; import { ServerActionResponse } from "@/features/common/server-action-response"; import { HistoryContainer } from "@/features/common/services/cosmos"; import { uniqueId } from "@/features/common/util"; @@ -31,8 +32,11 @@ export const CreateCitation = async ( } }; +// Create citations for the documents with a user as optional parameter +// when calling this method from the extension, you must provide the user as the REST API does not have access to the user export const CreateCitations = async ( - models: DocumentSearchResponse[] + models: DocumentSearchResponse[], + userId?: string ): Promise>> => { const items: Array>> = []; @@ -41,6 +45,7 @@ export const CreateCitations = async ( content: model, id: uniqueId(), type: CHAT_CITATION_ATTRIBUTE, + userId: userId || (await userHashedId()), }); items.push(res); @@ -54,7 +59,8 @@ export const FindCitationByID = async ( ): Promise> => { try { const querySpec: SqlQuerySpec = { - query: "SELECT * FROM root r WHERE r.type=@type AND r.id=@id", + query: + "SELECT * FROM root r WHERE r.type=@type AND r.id=@id AND r.userId=@userId ", parameters: [ { name: "@type", @@ -64,6 +70,10 @@ export const FindCitationByID = async ( name: "@id", value: id, }, + { + name: "@userId", + value: await userHashedId(), + }, ], }; @@ -89,3 +99,21 @@ export const FindCitationByID = async ( }; } }; + +export const FormatCitations = (citation: DocumentSearchResponse[]) => { + const withoutEmbedding: DocumentSearchResponse[] = []; + citation.forEach((d) => { + withoutEmbedding.push({ + score: d.score, + document: { + metadata: d.document.metadata, + pageContent: d.document.pageContent, + chatThreadId: d.document.chatThreadId, + id: "", + user: "", + }, + }); + }); + + return withoutEmbedding; +}; diff --git a/src/features/chat-page/chat-services/models.ts b/src/features/chat-page/chat-services/models.ts index 31d76926d..c8e02d208 100644 --- a/src/features/chat-page/chat-services/models.ts +++ b/src/features/chat-page/chat-services/models.ts @@ -67,6 +67,7 @@ export type MenuItemsGroup = { export type ChatCitationModel = { id: string; content: any; + userId: string; type: typeof CHAT_CITATION_ATTRIBUTE; };