Merge pull request #30 from pangeacyber/kenany/response-side-panel

Add right-hand side panel for API responses
pangeacyber · Dec 18, 2024 · dcbd420 · dcbd420
2 parents 9a8b39d + e24cd70
commit dcbd420
Show file tree

Hide file tree

Showing 19 changed files with 2,554 additions and 130 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -10,6 +10,7 @@
     "create-authz-tuples": "tsx scripts/create-authz-tuples.ts"
   },
   "dependencies": {
+    "@emotion/css": "11.13.5",
     "@emotion/react": "^11.13.5",
     "@emotion/styled": "^11.13.5",
     "@googleapis/drive": "8.14.0",
@@ -20,12 +21,15 @@
     "@mui/material": "6.1.10",
     "@pangeacyber/react-auth": "^0.0.14",
     "@pangeacyber/react-mui-audit-log-viewer": "^1.0.5",
+    "hast-util-to-jsx-runtime": "2.3.2",
     "langchain": "0.3.7",
     "markdown-table": "3.0.4",
     "next": "14.2.20",
-    "pangea-node-sdk": "4.1.0",
+    "pangea-node-sdk": "4.2.0-beta.2",
     "react": "^18.3.1",
-    "react-dom": "^18.3.1"
+    "react-dom": "^18.3.1",
+    "react-json-view-lite": "2.0.1",
+    "shiki": "1.24.2"
   },
   "devDependencies": {
     "@dotenvx/dotenvx": "1.29.0",

diff --git a/src/app/api/ai/route.ts b/src/app/api/ai/route.ts
@@ -3,13 +3,15 @@ import { BedrockEmbeddings, ChatBedrockConverse } from "@langchain/aws";
 import { HumanMessage, SystemMessage } from "@langchain/core/messages";
 import { StringOutputParser } from "@langchain/core/output_parsers";
 import { MemoryVectorStore } from "langchain/vectorstores/memory";
+import type { AuthZ } from "pangea-node-sdk";
 
 import {
   auditLogRequest,
   auditSearchRequest,
   authzCheckRequest,
   validateToken,
 } from "../requests";
+import type { PangeaResponse } from "@src/types";
 import { rateLimitQuery } from "@src/utils";
 import { DAILY_MAX_MESSAGES, PROMPT_MAX_CHARS } from "@src/const";
 import { GoogleDriveRetriever } from "@src/google";
@@ -88,6 +90,7 @@ export async function POST(request: NextRequest) {
     let docs = await retriever.invoke(body.userPrompt);
 
     // Filter documents based on user's permissions in AuthZ.
+    const authzResponses: PangeaResponse<AuthZ.CheckResult>[] = [];
     if (body.authz) {
       docs = await Promise.all(
         docs.map(async (doc) => {
@@ -97,6 +100,16 @@ export async function POST(request: NextRequest) {
             resource: { type: "file", id: doc.id },
             debug: true,
           });
+          if ("request_id" in response) {
+            authzResponses.push({
+              request_id: response.request_id,
+              request_time: response.request_time,
+              response_time: response.response_time,
+              result: response.result,
+              status: response.status,
+              summary: response.summary,
+            });
+          }
           return response.result.allowed ? doc : null;
         }),
       ).then((results) => results.filter((doc) => doc !== null));
@@ -130,7 +143,15 @@ Context: ${context}`),
 
     await auditLogRequest(auditLogData);
 
-    return Response.json({ content: text });
+    return Response.json({
+      content: text,
+      authzResponses,
+      documents: docs.map(({ id, metadata, pageContent }) => ({
+        id,
+        metadata,
+        pageContent,
+      })),
+    });
   } catch (err) {
     console.log("Error:", err);
     return new Response(`{"error": "ConverseCommand failed"}`, {

diff --git a/src/app/api/data/route.ts b/src/app/api/data/route.ts
@@ -43,8 +43,8 @@ export async function POST(request: NextRequest) {
       await auditLogRequest(auditLogData);
     } catch (_) {}
 
-    return new Response(JSON.stringify(response.result));
+    return Response.json(response);
   } else {
-    return new Response(JSON.stringify(response.result), { status: 400 });
+    return Response.json(response, { status: 400 });
   }
 }
diff --git a/src/app/api/prompt/route.ts b/src/app/api/prompt/route.ts
@@ -38,8 +38,8 @@ export async function POST(request: NextRequest) {
       await auditLogRequest(auditLogData);
     } catch (_) {}
 
-    return new Response(JSON.stringify(response.result));
+    return Response.json(response);
   } else {
-    return new Response(JSON.stringify(response.result), { status: 400 });
+    return Response.json(response, { status: 400 });
   }
 }
diff --git a/src/app/api/requests.ts b/src/app/api/requests.ts
@@ -2,6 +2,7 @@ import type { NextRequest } from "next/server";
 import type { AuthZ } from "pangea-node-sdk";
 
 import { delay } from "@src/utils";
+import { PangeaResponse } from "@src/types";
 
 type ValidatedToken =
   | { success: true; username: string; profile: Record<string, string> }
@@ -29,15 +30,6 @@ export async function validateToken(
   };
 }
 
-export interface ResponseObject<T> {
-  request_id: string;
-  request_time: string;
-  response_time: string;
-  status: string;
-  result: T;
-  summary: string;
-}
-
 export async function auditLogRequest(data: { event: Record<string, string> }) {
   const url = getUrl("audit", "v1/log");
   const now = new Date();
@@ -74,7 +66,7 @@ export async function auditSearchRequest(data: object) {
 export async function authzCheckRequest(
   data: AuthZ.CheckRequest,
 ): Promise<
-  ResponseObject<AuthZ.CheckResult> | { result: { allowed: boolean } }
+  PangeaResponse<AuthZ.CheckResult> | { result: { allowed: boolean } }
 > {
   const url = getUrl("authz", "v1/check");
   const { success, response } = await postRequest(url, data);
@@ -127,7 +119,7 @@ export async function getRequest(url: string) {
 
 async function handleAsync(response: Response): Promise<Response> {
   const data = await response.json();
-  const url = `https://${process.env.NEXT_PUBLIC_PANGEA_BASE_DOMAIN}/request/${data?.request_id}`;
+  const url = data.result.location;
   const maxRetries = 3;
   let retryCount = 1;
 

diff --git a/src/app/components/PanelHeader/index.ts b/src/app/components/PanelHeader/index.ts
@@ -0,0 +1,13 @@
+"use client";
+
+import { styled } from "@mui/material/styles";
+
+const PanelHeader = styled("div")(({ theme }) => ({
+  display: "flex",
+  alignItems: "center",
+  padding: theme.spacing(0, 1),
+  ...theme.mixins.toolbar,
+  justifyContent: "space-between",
+}));
+
+export default PanelHeader;
diff --git a/src/app/context.tsx b/src/app/context.tsx
@@ -1,3 +1,5 @@
+import type { DocumentInterface } from "@langchain/core/documents";
+import { PangeaResponse } from "@src/types";
 import {
   FC,
   ReactNode,
@@ -19,6 +21,13 @@ export interface ChatContextProps {
   sidePanelOpen: boolean;
   auditPanelOpen: boolean;
   loginOpen: boolean;
+  promptGuardResponse: Partial<PangeaResponse<unknown>>;
+  aiGuardResponses: readonly [
+    Partial<PangeaResponse<unknown>>,
+    Partial<PangeaResponse<unknown>>,
+  ];
+  authzResponses: readonly PangeaResponse<unknown>[];
+  documents: readonly DocumentInterface[];
   setLoading: (value: boolean) => void;
   setSystemPrompt: (value: string) => void;
   setUserPrompt: (value: string) => void;
@@ -28,6 +37,12 @@ export interface ChatContextProps {
   setSidePanelOpen: (value: boolean) => void;
   setAuditPanelOpen: (value: boolean) => void;
   setLoginOpen: (value: boolean) => void;
+  setPromptGuardResponse: (value: PangeaResponse<unknown>) => void;
+  setAiGuardResponses: (
+    value: readonly [PangeaResponse<unknown>, PangeaResponse<unknown>],
+  ) => void;
+  setAuthzResponses: (value: readonly PangeaResponse<unknown>[]) => void;
+  setDocuments: (value: readonly DocumentInterface[]) => void;
 }
 
 const ChatContext = createContext<ChatContextProps>({
@@ -40,6 +55,10 @@ const ChatContext = createContext<ChatContextProps>({
   sidePanelOpen: true,
   auditPanelOpen: false,
   loginOpen: false,
+  promptGuardResponse: {},
+  aiGuardResponses: [{}, {}],
+  authzResponses: [],
+  documents: [],
   setLoading: () => {},
   setSystemPrompt: () => {},
   setUserPrompt: () => {},
@@ -49,6 +68,10 @@ const ChatContext = createContext<ChatContextProps>({
   setSidePanelOpen: () => {},
   setAuditPanelOpen: () => {},
   setLoginOpen: () => {},
+  setPromptGuardResponse: () => {},
+  setAiGuardResponses: () => {},
+  setAuthzResponses: () => {},
+  setDocuments: () => {},
 });
 
 export interface ChatProviderProps {
@@ -82,6 +105,17 @@ export const ChatProvider: FC<ChatProviderProps> = ({ children }) => {
   const [sidePanelOpen, setSidePanelOpen] = useState(true);
   const [auditPanelOpen, setAuditPanelOpen] = useState(false);
   const [loginOpen, setLoginOpen] = useState(false);
+  const [promptGuardResponse, setPromptGuardResponse] = useState({});
+  const [aiGuardResponses, setAiGuardResponses] = useState<
+    readonly [
+      Partial<PangeaResponse<unknown>>,
+      Partial<PangeaResponse<unknown>>,
+    ]
+  >([{}, {}]);
+  const [authzResponses, setAuthzResponses] = useState<
+    readonly PangeaResponse<unknown>[]
+  >([]);
+  const [documents, setDocuments] = useState<readonly DocumentInterface[]>([]);
 
   useEffect(() => {
     if (!mounted.current) {
@@ -118,6 +152,10 @@ export const ChatProvider: FC<ChatProviderProps> = ({ children }) => {
       sidePanelOpen,
       auditPanelOpen,
       loginOpen,
+      promptGuardResponse,
+      aiGuardResponses,
+      authzResponses,
+      documents,
       setLoading,
       setSystemPrompt,
       setUserPrompt,
@@ -127,6 +165,10 @@ export const ChatProvider: FC<ChatProviderProps> = ({ children }) => {
       setSidePanelOpen,
       setAuditPanelOpen,
       setLoginOpen,
+      setPromptGuardResponse,
+      setAiGuardResponses,
+      setAuthzResponses,
+      setDocuments,
     }),
     [
       loading,
@@ -138,6 +180,10 @@ export const ChatProvider: FC<ChatProviderProps> = ({ children }) => {
       sidePanelOpen,
       auditPanelOpen,
       loginOpen,
+      promptGuardResponse,
+      aiGuardResponses,
+      authzResponses,
+      documents,
       setLoading,
       setSystemPrompt,
       setUserPrompt,

diff --git a/src/app/features/AuditViewer/index.tsx b/src/app/features/AuditViewer/index.tsx
@@ -121,8 +121,8 @@ const AuditViewer = () => {
           gap={2}
           sx={{
             width: sidePanelOpen
-              ? "calc(100vw - 375px)"
-              : "calc(100vw - 130px)",
+              ? "calc(100vw - 860px)"
+              : "calc(100vw - 580px)",
             marginBottom: "20px",
             padding: "20px",
             borderRadius: "10px",

diff --git a/src/app/features/ChatWindow/index.tsx b/src/app/features/ChatWindow/index.tsx
@@ -14,6 +14,7 @@ import {
 import { useTheme } from "@mui/material/styles";
 import SendIcon from "@mui/icons-material/Send";
 import { useAuth } from "@pangeacyber/react-auth";
+import type { AIGuard } from "pangea-node-sdk";
 
 import { DAILY_MAX_MESSAGES, PROMPT_MAX_CHARS } from "@src/const";
 import { useChatContext, ChatMessage } from "@src/app/context";
@@ -27,6 +28,7 @@ import {
 } from "./utils";
 import ChatScroller from "./components/ChatScroller";
 import { Colors } from "@src/app/theme";
+import type { PangeaResponse } from "@src/types";
 import { rateLimitQuery } from "@src/utils";
 
 function hashCode(str: string) {
@@ -50,6 +52,10 @@ const ChatWindow = () => {
     setUserPrompt,
     setLoading,
     setLoginOpen,
+    setPromptGuardResponse,
+    setAiGuardResponses,
+    setAuthzResponses,
+    setDocuments,
   } = useChatContext();
   const { authenticated, user, logout } = useAuth();
   const [messages, setMessages] = useState<ChatMessage[]>([]);
@@ -119,6 +125,7 @@ const ChatWindow = () => {
 
       try {
         const promptResp = await callPromptGuard(token, userPrompt, "");
+        setPromptGuardResponse(promptResp);
         const pgMsg: ChatMessage = {
           hash: hashCode(JSON.stringify(promptResp)),
           type: "prompt_guard",
@@ -127,7 +134,7 @@ const ChatWindow = () => {
         setMessages((prevMessages) => [...prevMessages, pgMsg]);
 
         // don't send to the llm if prompt is malicious
-        if (promptResp?.detected) {
+        if (promptResp.result.detected) {
           processingError("Processing halted: suspicious prompt");
           return;
         }
@@ -139,20 +146,25 @@ const ChatWindow = () => {
     }
 
     let llmUserPrompt = userPrompt;
+    let guardedInput: PangeaResponse<AIGuard.TextGuardResult>;
 
     if (dataGuardEnabled) {
       setProcessing("Checking user prompt with AI Guard");
 
       try {
-        const dataResp = await callInputDataGuard(token, userPrompt);
+        guardedInput = await callInputDataGuard(token, userPrompt);
+        setAiGuardResponses([
+          guardedInput,
+          {} as PangeaResponse<AIGuard.TextGuardResult>,
+        ]);
         const dgiMsg: ChatMessage = {
-          hash: hashCode(JSON.stringify(dataResp)),
+          hash: hashCode(JSON.stringify(guardedInput)),
           type: "ai_guard",
-          findings: JSON.stringify(dataResp.findings),
+          findings: JSON.stringify(guardedInput.result.findings),
         };
         setMessages((prevMessages) => [...prevMessages, dgiMsg]);
 
-        llmUserPrompt = dataResp.redacted_prompt;
+        llmUserPrompt = guardedInput.result.redacted_prompt;
       } catch (err) {
         const status = err instanceof Response ? err.status : 0;
         processingError("AI Guard call failed, please try again", status);
@@ -172,12 +184,15 @@ const ChatWindow = () => {
     let llmResponse = "";
 
     try {
-      llmResponse = await sendUserMessage(
+      const llmResponseObj = await sendUserMessage(
         token,
         llmUserPrompt,
         systemPrompt,
         authzEnabled,
       );
+      llmResponse = llmResponseObj.content;
+      setAuthzResponses(llmResponseObj.authzResponses);
+      setDocuments(llmResponseObj.documents);
 
       // decrement daily remaining count
       setRemaining((curVal) => curVal - 1);
@@ -192,14 +207,15 @@ const ChatWindow = () => {
 
       try {
         const dataResp = await callResponseDataGuard(token, llmResponse);
+        setAiGuardResponses([guardedInput!, dataResp]);
         const dgrMsg: ChatMessage = {
           hash: hashCode(JSON.stringify(dataResp)),
           type: "ai_guard",
-          findings: JSON.stringify(dataResp.findings),
+          findings: JSON.stringify(dataResp.result.findings),
         };
         dataGuardMessages.push(dgrMsg);
 
-        llmResponse = dataResp.redacted_prompt;
+        llmResponse = dataResp.result.redacted_prompt;
       } catch (err) {
         const status = err instanceof Response ? err.status : 0;
         processingError("AI Guard call failed, please try again", status);