(wip) fix memoization / memory leak

onyx-dot-app · Oct 1, 2024 · 09eaac3 · 09eaac3
1 parent d11954d
commit 09eaac3
Show file tree

Hide file tree

Showing 9 changed files with 603 additions and 459 deletions.
diff --git a/backend/danswer/configs/app_configs.py b/backend/danswer/configs/app_configs.py
@@ -135,7 +135,7 @@
     os.environ.get("POSTGRES_PASSWORD") or "password"
 )
 POSTGRES_HOST = os.environ.get("POSTGRES_HOST") or "localhost"
-POSTGRES_PORT = os.environ.get("POSTGRES_PORT") or "5432"
+POSTGRES_PORT = os.environ.get("POSTGRES_PORT") or "5433"
 POSTGRES_DB = os.environ.get("POSTGRES_DB") or "postgres"
 
 # defaults to False

diff --git a/deployment/docker_compose/docker-compose.dev.yml b/deployment/docker_compose/docker-compose.dev.yml
@@ -293,7 +293,7 @@ services:
       - POSTGRES_USER=${POSTGRES_USER:-postgres}
       - POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-password}
     ports:
-      - "5432:5432"
+      - "5433:5432"
     volumes:
       - db_volume:/var/lib/postgresql/data
 

diff --git a/deployment/docker_compose/docker-compose.gpu-dev.yml b/deployment/docker_compose/docker-compose.gpu-dev.yml
@@ -303,7 +303,7 @@ services:
       - POSTGRES_USER=${POSTGRES_USER:-postgres}
       - POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-password}
     ports:
-      - "5432:5432"
+      - "5433:5432"
     volumes:
       - db_volume:/var/lib/postgresql/data
 

diff --git a/deployment/docker_compose/docker-compose.search-testing.yml b/deployment/docker_compose/docker-compose.search-testing.yml
@@ -154,7 +154,7 @@ services:
       - POSTGRES_USER=${POSTGRES_USER:-postgres}
       - POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-password}
     ports:
-      - "5432"
+      - "5433"
     volumes:
       - db_volume:/var/lib/postgresql/data
 

diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx
@@ -162,9 +162,6 @@ export function ChatPage({
     user,
     availableAssistants
   );
-  const finalAssistants = user
-    ? orderAssistantsForUser(visibleAssistants, user)
-    : visibleAssistants;
 
   const existingChatSessionAssistantId = selectedChatSession?.persona_id;
   const [selectedAssistant, setSelectedAssistant] = useState<
@@ -219,7 +216,7 @@ export function ChatPage({
   const liveAssistant =
     alternativeAssistant ||
     selectedAssistant ||
-    finalAssistants[0] ||
+    visibleAssistants[0] ||
     availableAssistants[0];
 
   useEffect(() => {
@@ -689,7 +686,7 @@ export function ChatPage({
   useEffect(() => {
     if (messageHistory.length === 0 && chatSessionIdRef.current === null) {
       setSelectedAssistant(
-        finalAssistants.find((persona) => persona.id === defaultAssistantId)
+        visibleAssistants.find((persona) => persona.id === defaultAssistantId)
       );
     }
   }, [defaultAssistantId]);
@@ -2394,7 +2391,10 @@ export function ChatPage({
                               showDocs={() => setDocumentSelection(true)}
                               selectedDocuments={selectedDocuments}
                               // assistant stuff
-                              assistantOptions={finalAssistants}
+                              assistantOptions={orderAssistantsForUser(
+                                visibleAssistants,
+                                user
+                              )}
                               selectedAssistant={liveAssistant}
                               setSelectedAssistant={onAssistantChange}
                               setAlternativeAssistant={setAlternativeAssistant}

diff --git a/web/src/app/chat/message/CodeBlock.tsx b/web/src/app/chat/message/CodeBlock.tsx
@@ -1,7 +1,61 @@
-import React, { useState, ReactNode, useCallback, useMemo, memo } from "react";
+// CodeBlock.tsx
+import React, {
+  useState,
+  ReactNode,
+  useCallback,
+  useMemo,
+  memo,
+  useEffect,
+  CSSProperties,
+  useRef,
+} from "react";
+import { isEqual } from "lodash";
 import { FiCheck, FiCopy } from "react-icons/fi";
 
-const CODE_BLOCK_PADDING_TYPE = { padding: "1rem" };
+const CODE_BLOCK_PADDING_TYPE: CSSProperties = { padding: "1rem" };
+
+interface CodeBlockProps {
+  className?: string;
+  children?: ReactNode;
+  content: string;
+  style?: CSSProperties;
+  // Add any other specific props you need to handle explicitly
+  // Example:
+  // onClick?: (event: React.MouseEvent) => void;
+  // title?: string;
+}
+
+const arePropsEqual = (
+  prevProps: CodeBlockProps,
+  nextProps: CodeBlockProps
+) => {
+  const prevContent =
+    typeof prevProps.children === "string"
+      ? prevProps.children
+      : prevProps.content;
+  const nextContent =
+    typeof nextProps.children === "string"
+      ? nextProps.children
+      : nextProps.content;
+
+  const isContentEqual = prevContent === nextContent;
+  const isClassNameEqual = prevProps.className === nextProps.className;
+  const isStyleEqual = isEqual(prevProps.style, nextProps.style);
+  // Compare other explicit props here if added
+  // Example:
+  // const isOnClickEqual = prevProps.onClick === nextProps.onClick;
+
+  const areEqual = isContentEqual && isClassNameEqual && isStyleEqual;
+
+  // console.log("Props comparison:", {
+  //   isContentEqual,
+  //   isClassNameEqual,
+  //   isStyleEqual,
+  //   areEqual,
+  // });
+
+  return areEqual;
+};
 
 interface CodeBlockProps {
   className?: string | undefined;
@@ -10,12 +64,15 @@ interface CodeBlockProps {
   [key: string]: any;
 }
 
-export const CodeBlock = memo(function CodeBlock({
+export const CodeBlockComponent = memo(function CodeBlock({
   className = "",
   children,
   content,
   ...props
 }: CodeBlockProps) {
+  const renderCount = useRef(0);
+  renderCount.current += 1;
+  console.log("CodeBlockComponent render count:", renderCount.current);
   const [copied, setCopied] = useState(false);
 
   const language = useMemo(() => {
@@ -155,3 +212,4 @@ export const CodeBlock = memo(function CodeBlock({
     </div>
   );
 });
+export const CodeBlock = memo(CodeBlockComponent, arePropsEqual);
diff --git a/web/src/app/chat/message/MemoizedTextComponents.tsx b/web/src/app/chat/message/MemoizedTextComponents.tsx
@@ -1,5 +1,5 @@
 import { Citation } from "@/components/search/results/Citation";
-import React, { memo } from "react";
+import React, { memo, useRef } from "react";
 
 export const MemoizedLink = memo((props: any) => {
   const { node, ...rest } = props;
@@ -25,9 +25,15 @@ export const MemoizedLink = memo((props: any) => {
   }
 });
 
-export const MemoizedParagraph = memo(({ node, ...props }: any) => (
-  <p {...props} className="text-default" />
-));
+export const MemoizedParagraph = memo(({ content }: { content: string }) => {
+  const renderCount = useRef(0);
+  renderCount.current += 1;
+
+  console.log("MemoizedParagraph render count:", renderCount.current);
+  console.log("MemoizedParagraph props:", content);
+
+  return <p className="text-default">{content}</p>;
+});
 
 MemoizedLink.displayName = "MemoizedLink";
 MemoizedParagraph.displayName = "MemoizedParagraph";