k

backend/onyx/db/chat.py

@@ -942,6 +942,7 @@ def translate_db_message_to_chat_message_detail(
         sub_questions=translate_db_sub_questions_to_server_objects(
             chat_message.sub_questions
         ),
+        refined_answer_improvement=chat_message.refined_answer_improvement,
     )
 
     return chat_msg_detail

backend/onyx/server/query_and_chat/models.py

@@ -239,6 +239,7 @@ class ChatMessageDetail(BaseModel):
     sub_questions: list[SubQuestionDetail] | None = None
     files: list[FileDescriptor]
     tool_call: ToolCallFinalResult | None
+    refined_answer_improvement: bool | None = None
 
     def model_dump(self, *args: list, **kwargs: dict[str, Any]) -> dict[str, Any]:  # type: ignore
         initial_dict = super().model_dump(mode="json", *args, **kwargs)  # type: ignore

web/src/app/chat/ChatPage.tsx

@@ -1404,6 +1404,14 @@ export function ChatPage({
               isImprovement = (packet as RefinedAnswerImprovement)
                 .refined_answer_improvement;
             }
+
+            if (Object.hasOwn(packet, "stream_type")) {
+              if ((packet as any).stream_type == "main_answer") {
+                is_generating = false;
+                second_level_generating = true;
+              }
+            }
+
             // Continuously refine the sub_questions based on the packets that we receive
             if (
               Object.hasOwn(packet, "stop_reason") &&

web/src/app/chat/interfaces.ts

@@ -146,6 +146,7 @@ export interface BackendMessage {
   // Keeping existing properties
   comments: any;
   parentMessageId: number | null;
+  refined_answer_improvement: boolean | null;
 }
 
 export interface MessageResponseIDInfo {

web/src/app/chat/lib.tsx

@@ -497,6 +497,8 @@ export function processRawChatHistory(
       latestChildMessageId: messageInfo.latest_child_message,
       overridden_model: messageInfo.overridden_model,
       sub_questions: subQuestions,
+      isImprovement:
+        (messageInfo.refined_answer_improvement as unknown as boolean) || false,
     };
 
     messages.set(messageInfo.message_id, message);

web/src/app/chat/message/AgenticMessage.tsx

@@ -372,33 +372,35 @@ export const AgenticMessage = ({
     onMessageSelection &&
     otherMessagesCanSwitchTo &&
     otherMessagesCanSwitchTo.length > 1;
+  const currentIndexRef = useRef(0);
 
   useEffect(() => {
     if (!allowStreaming) {
-      //   if (typeof content === "string") {
-      //     setStreamedContent(finalContent);
-      //     setLastKnownContentLength(finalContent.length);
-      //   }
       return;
     }
 
     if (typeof finalContent !== "string") return;
 
-    let currentIndex = streamedContent.length;
     let intervalId: NodeJS.Timeout | null = null;
 
-    // if (finalContent.length > currentIndex) {
     intervalId = setInterval(() => {
       setStreamedContent((prev) => {
-        if (prev.length < finalContent.length) {
-          const nextLength = Math.min(prev.length + 5, finalContent.length);
+        if (streamedContent.length == finalContent.length) {
+          return finalContent;
+        }
+        if (currentIndexRef.current < finalContent.length) {
+          const nextLength = Math.min(
+            currentIndexRef.current + 5,
+            finalContent.length
+          );
+          currentIndexRef.current = nextLength;
           return finalContent.slice(0, nextLength);
         } else {
           if (intervalId) clearInterval(intervalId);
           return finalContent;
         }
       });
-    }, 10);
+    }, 5);
     // } else {
     //   setStreamedContent(finalContent);
     // }
@@ -472,47 +474,81 @@ export const AgenticMessage = ({
                             Answer
                           </div>
 
-                          {isImprovement == null &&
-                          subQuestions &&
-                          subQuestions.length > 0 ? (
-                            <RefinemenetBadge
-                              finished={!secondLevelGenerating}
-                              overallAnswer={secondLevelAssistantMessage || ""}
-                              secondLevelSubquestions={secondLevelSubquestions}
-                              toggleInitialAnswerVieinwg={() => {
-                                setIsViewingInitialAnswer(
-                                  !isViewingInitialAnswer
-                                );
-                              }}
-                              isViewingInitialAnswer={isViewingInitialAnswer}
-                            />
-                          ) : secondLevelAssistantMessage ? (
-                            isImprovement ? (
-                              <Badge
-                                // NOTE: This is a hack to make the badge slightly higher
-                                className="cursor-pointer mt-[1px]"
-                                variant="agent"
-                                onClick={() => {
-                                  const viewInitialAnswer =
-                                    !isViewingInitialAnswer;
-                                  setIsViewingInitialAnswer(viewInitialAnswer);
-                                  toggleDocDisplay &&
-                                    toggleDocDisplay(isViewingInitialAnswer);
-                                  if (viewInitialAnswer) {
-                                    setIsViewingInitialAnswer(true);
-                                  }
+                          {/* {
+                            <div className="bg-gray-100 p-4 rounded-md shadow-sm">
+                              <p>
+                                isGenerating: {isGenerating ? "true" : "false"}
+                              </p>
+                              <p>
+                                secondLevelGenerating:{" "}
+                                {secondLevelGenerating ? "true" : "false"}
+                              </p>
+                              <p>
+                                isImprovement:{" "}
+                                {isImprovement === null ? "null" : "not null"}
+                              </p>
+                              <p>
+                                subQuestions:{" "}
+                                {subQuestions ? "exists" : "does not exist"}
+                              </p>
+                              <p>
+                                subQuestions length:{" "}
+                                {subQuestions ? subQuestions.length : "N/A"}
+                              </p>
+                            </div>
+                          } */}
+                          {/* {isImprovement ? "true" : "false"} */}
+                          {
+                            // !isGenerating &&
+                            secondLevelGenerating &&
+                            isImprovement == null &&
+                            subQuestions &&
+                            subQuestions.length > 0 ? (
+                              <RefinemenetBadge
+                                finished={!secondLevelGenerating}
+                                overallAnswer={
+                                  secondLevelAssistantMessage || ""
+                                }
+                                secondLevelSubquestions={
+                                  secondLevelSubquestions
+                                }
+                                toggleInitialAnswerVieinwg={() => {
+                                  setIsViewingInitialAnswer(
+                                    !isViewingInitialAnswer
+                                  );
                                 }}
-                              >
-                                {isViewingInitialAnswer
-                                  ? "See Refined Answer"
-                                  : "See Original Answer"}
-                              </Badge>
+                                isViewingInitialAnswer={isViewingInitialAnswer}
+                              />
+                            ) : secondLevelAssistantMessage ? (
+                              isImprovement ? (
+                                <Badge
+                                  // NOTE: This is a hack to make the badge slightly higher
+                                  className="cursor-pointer mt-[1px]"
+                                  variant="agent"
+                                  onClick={() => {
+                                    const viewInitialAnswer =
+                                      !isViewingInitialAnswer;
+                                    setIsViewingInitialAnswer(
+                                      viewInitialAnswer
+                                    );
+                                    toggleDocDisplay &&
+                                      toggleDocDisplay(isViewingInitialAnswer);
+                                    if (viewInitialAnswer) {
+                                      setIsViewingInitialAnswer(true);
+                                    }
+                                  }}
+                                >
+                                  {isViewingInitialAnswer
+                                    ? "See Refined Answer"
+                                    : "See Original Answer"}
+                                </Badge>
+                              ) : (
+                                <NoNewAnswerMessage />
+                              )
                             ) : (
-                              <NoNewAnswerMessage />
+                              <></>
                             )
-                          ) : (
-                            <></>
-                          )}
+                          }
                         </div>
 
                         <div className="px-4">

web/src/app/chat/message/SubQuestionsDisplay.tsx

@@ -35,7 +35,7 @@ export const StatusIndicator = ({ status }: { status: ToggleState }) => {
     bg-background border-3 border-background-900  "
     ${
       status === ToggleState.Todo
-        ? "!border-4 border border-background-900 bg-background"
+        ? "!border-[3px] border border-background-900 bg-background"
         : false
           ? "bg-background border-3 border border-background-900 rotating-border"
           : "bg-background-900 flex items-center  justify-center"

web/src/app/chat/refinmentBadge.tsx

@@ -231,7 +231,7 @@ export default function RefinemenetBadge({
     <TooltipProvider delayDuration={0}>
       <Tooltip>
         <div
-          className="relative w-full max-w-sm"
+          className="relative w-fit max-w-sm"
           onMouseEnter={() => setIsHovered(true)}
           onMouseLeave={() => setIsHovered(false)}
         >
@@ -254,7 +254,7 @@ export default function RefinemenetBadge({
           <TooltipContent
             side="bottom"
             align="start"
-            className="w-80 p-4 bg-white shadow-lg rounded-md"
+            className="w-fit p-4 bg-white border-2 border-border shadow-lg rounded-md"
           >
             {/* If not done, show the "Refining" box + a chevron */}
 
@@ -496,7 +496,7 @@ export const NoNewAnswerMessage = () => {
 
   return (
     <div className="text-gray-600 text-sm" style={{ opacity: opacity }}>
-      No new answer found...
+      {/* No new answer found... */}
     </div>
   );
 };

web/src/components/search/results/Citation.tsx

@@ -76,7 +76,10 @@ export function Citation({
             </span>
           </div>
         </TooltipTrigger>
-        <TooltipContent width="mb-2 max-w-lg" className="bg-background">
+        <TooltipContent
+          width="mb-2 max-w-lg"
+          className="border-2 border-border shadow-lg bg-white"
+        >
           {document_info?.document ? (
             <CompactDocumentCard
               updatePresentingDocument={document_info.updatePresentingDocument}