Merge pull request microsoft#323 from thivy/fix_docId_issue

fix doc id being exposed to client

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 {

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);

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) => {

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<Array<ServerActionResponse<ChatCitationModel>>> => {
   const items: Array<Promise<ServerActionResponse<ChatCitationModel>>> = [];
 
@@ -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<ServerActionResponse<ChatCitationModel>> => {
   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;
+};

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;
 };