a lot of changes

frontend/package.json

@@ -3,7 +3,7 @@
   "version": "0.1.0",
   "private": true,
   "dependencies": {
-    "@emotion/react": "^11.14.0",
+ "@emotion/react": "^11.14.0",
     "@emotion/styled": "^11.14.0",
     "@fortawesome/fontawesome-free": "^6.7.2",
     "@google/generative-ai": "^0.21.0",
@@ -11,12 +11,20 @@
     "@mui/material": "^6.4.3",
     "@mui/styled-engine-sc": "^6.4.2",
     "cra-template": "1.2.0",
+    "katex": "^0.16.4",
     "react": "^19.0.0",
     "react-dom": "^19.0.0",
     "react-icons": "^5.4.0",
-    "react-markdown": "^9.0.3",
+    "react-markdown": "^9.1.0",
     "react-router-dom": "^7.1.3",
     "react-scripts": "5.0.1",
+    "react-syntax-highlighter": "^15.5.0",
+    "rehype-highlight": "^7.0.2",
+    "rehype-katex": "^6.0.2",
+    "rehype-raw": "^6.1.1",
+    "rehype-sanitize": "^5.0.1",
+    "remark-gfm": "^3.0.1",
+    "remark-math": "^5.1.1",
     "styled-components": "^6.1.14",
     "web-vitals": "^4.2.4"
   },

frontend/src/App.js

@@ -1,6 +1,8 @@
-import React, { useState, useEffect } from 'react';
+import React, { useState } from 'react';
 import { BrowserRouter, Routes, Route } from 'react-router-dom';
 import CircularProgress from '@mui/material/CircularProgress';
+import Snackbar from '@mui/material/Snackbar';
+import Alert from '@mui/material/Alert';
 import logo from './Icons/settings-2.svg';
 import './App.css';
 import IntialSetting from './Components/IntialSetting.js';
@@ -20,13 +22,28 @@ function App() {
 function Home() {
   const [showSettings, setShowSettings] = useState(false);
   const [initializing, setInitializing] = useState(false);
+  // Snackbar state
+  const [snackbar, setSnackbar] = useState({
+    open: false,
+    message: "",
+    severity: "success",
+  });
 
-  // This callback is passed to IntialSetting and called when the user clicks Save.
-  // It changes the underlying header content to the initializing state.
   const handleInitializationStart = () => {
     setInitializing(true);
   };
 
+  // Function to open the snackbar
+  const openSnackbar = (message, severity = "success") => {
+    setSnackbar({ open: true, message, severity });
+  };
+
+  // Function to close the snackbar
+  const closeSnackbar = (event, reason) => {
+    if (reason === 'clickaway') return;
+    setSnackbar(prev => ({ ...prev, open: false }));
+  };
+
   return (
     <div className="App">
       <header className="App-header">
@@ -48,17 +65,31 @@ function Home() {
           </>
         )}
 
-        {/* Always render the settings modal if showSettings is true */}
+        {/* InitialSetting */}
         {showSettings && (
           <IntialSetting 
             trigger={showSettings} 
             setTrigger={setShowSettings}
             onInitializationStart={handleInitializationStart}
+            openSnackbar={openSnackbar}
+            closeSnackbar={closeSnackbar}
           />
         )}
       </header>
+      
+      {/* Render the Snackbar*/}
+      <Snackbar
+        open={snackbar.open}
+        autoHideDuration={snackbar.severity === 'success' ? 3000 : null}
+        onClose={closeSnackbar}
+        anchorOrigin={{ vertical: 'top', horizontal: 'center' }}
+      >
+        <Alert onClose={closeSnackbar} severity={snackbar.severity} variant="filled" sx={{ width: '100%' }}>
+          {snackbar.message}
+        </Alert>
+      </Snackbar>
     </div>
   );
 }
 
-export default App;
+export default App;

frontend/src/Components/AiComponents/{Evaluate.css → ChatComponents/Evaluate.css}

@@ -86,6 +86,12 @@
   color: #ddd !important;
 }
 
+/* No metrics message */
+.no-metrics-message {
+  text-align: center;
+  color: red;
+}
+
 /* Spinner styling */
 .custom-spinner {
   width: 1.35rem;

frontend/src/Components/AiComponents/{Evaluate.js → ChatComponents/Evaluate.js}

@@ -16,20 +16,9 @@ const MenuProps = {
   PaperProps: {
     className: 'evaluate-menu',
   },
+  disableScrollLock: true
 };
 
-const names = [
-  "Bias",
-  "Toxicity",
-  "Answer Correctness",
-  "Summarization",
-  "Faithfulness",
-  "Hallucination",
-  "Answer Relevancy",
-  "Contextual Relevancy",
-  "Contextual Recall"
-];
-
 function getStyles(name, selectedNames, theme) {
   return {
     fontWeight: selectedNames.includes(name.toLowerCase())
@@ -45,6 +34,32 @@ export default function MultipleSelectChip({ evaluation }) {
   const [evaluationResult, setEvaluationResult] = React.useState("");
   const [isEvaluating, setIsEvaluating] = React.useState(false);
   const [localLoading, setLocalLoading] = React.useState(false);
+  const [noMetricsError, setNoMetricsError] = React.useState("");
+  const [metricOptions, setMetricOptions] = React.useState([]);
+
+  React.useEffect(() => {
+    // If 'contents' is undefined in the payload
+    if (evaluation && evaluation.contents === undefined) {
+      setMetricOptions([
+        "Bias",
+        "Toxicity",
+        "Summarization",
+        "Answer Correctness",
+      ]);
+    } else {
+      // Else, all except "Answer Correctness"
+      setMetricOptions([
+        "Bias",
+        "Toxicity",
+        "Summarization",
+        "Faithfulness",
+        "Hallucination",
+        "Answer Relevancy",
+        "Contextual Relevancy",
+        "Contextual Recall",
+      ]);
+    }
+  }, [evaluation]);
 
   // Reset the form fields
   React.useEffect(() => {
@@ -53,6 +68,7 @@ export default function MultipleSelectChip({ evaluation }) {
     setSelectedMetrics([]);
     setEvaluationResult("");
     setLocalLoading(true);
+    setNoMetricsError("");
     
     // Simulate a loading delay
     const timer = setTimeout(() => {
@@ -66,6 +82,7 @@ export default function MultipleSelectChip({ evaluation }) {
     const metrics = typeof value === 'string' ? value.split(',') : value;
     setPersonName(metrics);
     setSelectedMetrics(metrics);
+    setNoMetricsError("");
   };
 
   const handleDelete = (chipToDelete) => {
@@ -83,6 +100,14 @@ export default function MultipleSelectChip({ evaluation }) {
   const handleEvaluateClick = async () => {
     // Clear previous evaluation result immediately.
     setEvaluationResult("");
+
+    // Check if no metrics selected
+    if (selectedMetrics.length === 0) {
+      setNoMetricsError("No metrics selected");
+      return;
+    }
+
+    setNoMetricsError("");
     setIsEvaluating(true);
 
     const payload = { ...evaluation, metrics: selectedMetrics };
@@ -130,7 +155,7 @@ export default function MultipleSelectChip({ evaluation }) {
 
   // Finds the matching display name for a metric.
   const getDisplayName = (lowerValue) => {
-    const found = names.find(n => n.toLowerCase() === lowerValue);
+    const found = metricOptions.find(n => n.toLowerCase() === lowerValue);
     return found ? found : lowerValue;
   };
 
@@ -172,7 +197,7 @@ export default function MultipleSelectChip({ evaluation }) {
               )}
               MenuProps={MenuProps}
             >
-              {names.map((name) => (
+              {metricOptions.map((name) => (
                 <MenuItem
                   key={name}
                   value={name.toLowerCase()}  // underlying value is lowercase
@@ -194,6 +219,12 @@ export default function MultipleSelectChip({ evaluation }) {
             </Button>
           </Box>
 
+          {noMetricsError && (
+            <Box className="no-metrics-message">
+              {noMetricsError}
+            </Box>
+          )}
+
           {isEvaluating && (
             <Box mt={1} display="flex" alignItems="center">
               <Box className="custom-spinner" />

frontend/src/Components/AiComponents/{Graph.js → ChatComponents/Graph.js}

@@ -2,19 +2,21 @@ import React, { useState, useEffect } from 'react';
 import { FaTimes } from 'react-icons/fa';
 import './Graph.css';
 
-export default function Graph({ open, onClose, onError }) {
+export default function Graph({ open, onClose, payload, onError }) {
   const [graphHtml, setGraphHtml] = useState("");
   const [loading, setLoading] = useState(true);
   const [error, setError] = useState("");
 
   useEffect(() => {
+    // if (open && payload) {
     if (open) {
       setLoading(true);
       setError("");
       fetch("/action/graph", {
         method: "POST",
         headers: { "Content-Type": "application/json" },
-        body: JSON.stringify({})
+        body: JSON.stringify()
+        // body: JSON.stringify(payload)
       })
         .then(res => res.json())
         .then(data => {
@@ -37,6 +39,7 @@ export default function Graph({ open, onClose, onError }) {
         });
     }
   }, [open, onError]);
+  // }, [open, payload, onError]);
 
   if (!open) return null;
 

frontend/src/Components/AiComponents/{RightSidebar.css → ChatComponents/RightSidebar.css}

@@ -7,6 +7,8 @@
 
 /* Main sidebar container */
 .right-side-bar {
+  display: flex;
+  flex-direction: column;
   position: fixed;
   top: 0;
   right: 0;
@@ -14,11 +16,16 @@
   background-color: var(--sidebar-background);  /* Keep background uniform */
   color: var(--text-light);
   box-shadow: -2px 0 8px rgba(0, 0, 0, 0.5);
-  transition: transform var(--transition-speed);
+  transition: width 0.4s ease;
   overflow-y: auto;
   z-index: 1000;
 }
 
+/* Sidebar resizing */
+.right-side-bar.resizing {
+  transition: none;
+}
+
 /* When the sidebar is closed */
 .right-side-bar.closed {
   width: 0;
@@ -59,6 +66,8 @@
 .sidebar-content {
   padding: 16px;
   background: transparent;
+  overflow-x: hidden;
+  overflow-y: auto;
 }
 
 /* Also clear any default marker via the pseudo-element */
@@ -126,4 +135,4 @@
 @keyframes spin {
   from { transform: rotate(0deg); }
   to { transform: rotate(360deg); }
-}
+}

frontend/src/Components/AiComponents/{RightSidebar.js → ChatComponents/RightSidebar.js}

@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { useRef } from 'react';
 import { FaTimes, FaCheck, FaSpinner } from 'react-icons/fa';
 import { BsChevronLeft } from 'react-icons/bs';
 import CircularProgress from '@mui/material/CircularProgress';
@@ -22,10 +22,12 @@ function RightSidebar({
 }) {
   const minWidth = 200;
   const maxWidth = 450;
+  const sidebarRef = useRef(null);
 
   // Called when the user starts resizing the sidebar.
   const startResize = (e) => {
     e.preventDefault();
+    sidebarRef.current.classList.add("resizing"); // Add the "resizing" class to the sidebar when resizing
     document.addEventListener("mousemove", resizeSidebar);
     document.addEventListener("mouseup", stopResize);
   };
@@ -38,6 +40,7 @@ function RightSidebar({
   };
 
   const stopResize = () => {
+    sidebarRef.current.classList.remove("resizing"); // Remove the "resizing" class from the sidebar when resizing stops
     document.removeEventListener("mousemove", resizeSidebar);
     document.removeEventListener("mouseup", stopResize);
   };
@@ -72,8 +75,9 @@ function RightSidebar({
   return (
     <>
       <nav
+        ref={sidebarRef}
         className={`right-side-bar ${isOpen ? "open" : "closed"}`}
-        style={{ width: isOpen ? rightSidebarWidth : undefined }}
+        style={{ width: isOpen ? rightSidebarWidth : 0 }}
       >
         <div className="sidebar-header">
           <h3>

frontend/src/Components/AiComponents/ChatComponents/Streaming.css

@@ -0,0 +1,120 @@
+.streaming-content {
+    font-family: inherit;
+    line-height: 2rem;
+    white-space: pre-wrap;
+    word-wrap: break-word;
+    margin: 0;          
+    padding: 0;
+  }
+  
+  /* Reset margin/padding for all descendants */
+  .streaming-content * {
+    margin: 0;
+    padding: 0;
+  }
+  
+  /* Top-level elements */
+  .streaming-content > * {
+    margin-top: 0.5rem;
+    margin-bottom: 0.5rem;
+  }
+  
+  /* VERY FIRST element in an AI answer */
+  .streaming-content > *:first-child {
+    margin-top: 0 !important;
+  }
+  
+  /* Headings */
+  .streaming-content h1,
+  .streaming-content h2,
+  .streaming-content h3,
+  .streaming-content h4,
+  .streaming-content h5,
+  .streaming-content h6 {
+    margin-top: 1rem;
+    margin-bottom: 0.75rem; 
+  }
+  
+  /* If heading is the very first element */
+  .streaming-content > h1:first-child,
+  .streaming-content > h2:first-child,
+  .streaming-content > h3:first-child,
+  .streaming-content > h4:first-child,
+  .streaming-content > h5:first-child,
+  .streaming-content > h6:first-child {
+    margin-top: 0 !important;
+  }
+  
+  /* Paragraphs */
+  .streaming-content p {
+    margin-top: 0.25rem;
+    margin-bottom: 0.25rem;
+  }
+  
+  /* Lists */
+  .streaming-content ul,
+  .streaming-content ol {
+    margin-top: 0.25rem;
+    margin-bottom: 0.25rem;
+    padding-left: 1.25rem;   
+    white-space: normal;     
+  }
+  
+  .streaming-content li {
+    margin-bottom: 0.25rem;
+  }
+  
+  .streaming-content li ul,
+  .streaming-content li ol {
+    margin-top: 0.15rem;
+    margin-bottom: 0.15rem;
+  }
+  
+  /* Code Blocks */
+  .code-block-container {
+    margin: 0.5rem 0;
+    border-radius: 4px;
+    background-color: #2b2b2b;
+    overflow: hidden;
+  }
+  
+  .code-block-header {
+    background-color: #1e1e1e;
+    color: #ffffff;
+    padding: 0.5rem;
+    font-size: 0.85rem;
+    font-weight: bold;
+  }
+  
+  /* Table Container */
+  .table-container {
+    margin: 0.5rem 0;
+    width: 100%;
+    overflow-x: auto;
+    border: 1px solid #ddd;
+    border-radius: 4px;
+  }
+  
+  .table-container th,
+  .table-container td {
+    border: 1px solid #ddd;
+    padding: 0.5rem;
+  }
+  
+  /* Markdown Links */
+  .markdown-link {
+    color: #1a73e8;
+    text-decoration: none;
+  }
+  .markdown-link:hover {
+    text-decoration: underline;
+  }
+  
+  /* Blockquotes */
+  .markdown-blockquote {
+    border-left: 4px solid #ccc;
+    padding-left: 0.75rem;
+    margin: 0.5rem 0; 
+    color: #555;
+    font-style: italic;
+  }

frontend/src/Components/AiComponents/ChatComponents/Streaming.js

@@ -0,0 +1,82 @@
+import React, { useEffect, useRef } from 'react';
+import ReactMarkdown from 'react-markdown';
+import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
+import { atomDark } from 'react-syntax-highlighter/dist/esm/styles/prism';
+import remarkGfm from 'remark-gfm';
+import rehypeRaw from 'rehype-raw';
+import './Streaming.css';
+
+// Streaming component for rendering markdown content
+const Streaming = ({ content, isStreaming, onContentRef }) => {
+  const contentRef = useRef(null);
+  
+  useEffect(() => {
+    if (contentRef.current && onContentRef) {
+      onContentRef(contentRef.current);
+    }
+  }, [content, onContentRef]);
+
+  const displayContent = isStreaming ? `${content}▌` : (content || '');
+
+  return (
+    <div className="streaming-content" ref={contentRef}>
+      <ReactMarkdown
+        remarkPlugins={[remarkGfm]}
+        rehypePlugins={[rehypeRaw]}
+        components={{
+          code({node, inline, className, children, ...props}) {
+            const match = /language-(\w+)/.exec(className || '');
+            return !inline ? (
+              <div className="code-block-container">
+                <div className="code-block-header">
+                  <span>{match ? match[1] : 'code'}</span>
+                </div>
+                <SyntaxHighlighter
+                  style={atomDark}
+                  language={match ? match[1] : 'text'}
+                  PreTag="div"
+                  {...props}
+                >
+                  {String(children).replace(/\n$/, '')}
+                </SyntaxHighlighter>
+              </div>
+            ) : (
+              <code className={className} {...props}>
+                {children}
+              </code>
+            );
+          },
+          table({node, ...props}) {
+            return (
+              <div className="table-container">
+                <table {...props} />
+              </div>
+            );
+          },
+          a({node, children, href, ...props}) {
+            return (
+              <a 
+                href={href}
+                target="_blank"
+                rel="noopener noreferrer"
+                className="markdown-link"
+                {...props}
+              >
+                {children}
+              </a>
+            );
+          },
+          blockquote({node, ...props}) {
+            return (
+              <blockquote className="markdown-blockquote" {...props} />
+            );
+          }
+        }}
+      >
+        {displayContent}
+      </ReactMarkdown>
+    </div>
+  );
+};
+
+export default Streaming;

frontend/src/Components/AiComponents/ChatWindow.css

@@ -122,11 +122,6 @@
 .answer {
   margin: 0;
   line-height: 1.85;
-}
-
-.markdown {
-  margin: -1rem 0 -0.8rem 0;
-  line-height: 2rem;
   white-space: pre-wrap;
 }
 
@@ -161,13 +156,6 @@
   filter: brightness(0.65);
 }
 
-/* .post-icons .copy-icon:active img,
-.post-icons .evaluate-icon:active img,
-.post-icons .sources-icon:active img,
-.post-icons .graph-icon:active img {
-  filter: brightness(0.35);
-} */
-
 /* Tooltip styling */
 .tooltip {
   position: absolute;

frontend/src/Components/AiComponents/ChatWindow.js

@@ -1,11 +1,11 @@
-import React, { useRef, useEffect, useState } from 'react';
+import React, { useRef, useState, useEffect } from 'react';
 import Box from '@mui/material/Box';
 import Snackbar from '@mui/material/Snackbar';
 import Slide from '@mui/material/Slide';
 import IconButton from '@mui/material/IconButton';
 import { FaTimes } from 'react-icons/fa';
-import ReactMarkdown from 'react-markdown';
-import GraphDialog from './Graph';
+import GraphDialog from './ChatComponents/Graph';
+import Streaming from './ChatComponents/Streaming';
 import './ChatWindow.css';
 
 import bot from '../../Icons/bot.png';
@@ -23,6 +23,7 @@ function SlideTransition(props) {
 function ChatWindow({
   blockId,
   userMessage,
+  tokenChunks,
   aiAnswer,
   thinkingTime,
   thoughtLabel,
@@ -60,12 +61,13 @@ function ChatWindow({
     setSnackbarOpen(false);
   };
 
-  // Determine if any tokens (partial or full answer) have been received.
-  const hasTokens = aiAnswer && aiAnswer.length > 0;
+  // Combine partial chunks (tokenChunks) if present; else fall back to the aiAnswer string.
+  const combinedAnswer = (tokenChunks && tokenChunks.length > 0)
+  ? tokenChunks.join("")
+  : aiAnswer;
+  const hasTokens = combinedAnswer && combinedAnswer.length > 0;
   // Assume streaming is in progress if thinkingTime is not set.
   const isStreaming = thinkingTime === null || thinkingTime === undefined;
-  // Append a trailing cursor if streaming.
-  const displayAnswer = hasTokens ? (isStreaming ? aiAnswer + "▌" : aiAnswer) : "";
 
   // Helper to render the thought label.
   const renderThoughtLabel = () => {
@@ -94,6 +96,11 @@ function ChatWindow({
     prevTasksRef.current = tasks;
   }, [tasks, blockId, openRightSidebar]);
 
+  // Handle getting the reference to the content for copy functionality
+  const handleContentRef = (ref) => {
+    answerRef.current = ref;
+  };
+
   return (
     <>
       { !hasTokens ? (
@@ -178,10 +185,12 @@ function ChatWindow({
                   <img src={bot} alt="bot icon" />
                 </div>
                 <div className="message-bubble bot-bubble">
-                  <div className="answer" ref={answerRef}>
-                    <ReactMarkdown className="markdown">
-                      {displayAnswer}
-                    </ReactMarkdown>
+                  <div className="answer">
+                    <Streaming 
+                      content={combinedAnswer}
+                      isStreaming={isStreaming}
+                      onContentRef={handleContentRef}
+                    />
                   </div>
                 </div>
                 <div className="post-icons">

frontend/src/Components/AiPage.css

@@ -24,7 +24,6 @@ html, body {
   display: flex;
   min-height: 100vh;
   position: relative;
-  overflow-x: hidden;
   overflow-y: auto;
 }
 
@@ -38,7 +37,7 @@ html, body {
   padding: 2rem;
   transition: 0.1s;
   width: 99%;
-  max-width: 800px;
+  max-width: 900px;
   margin: 0 auto;
 }
 
@@ -52,8 +51,7 @@ html, body {
 }
 
 .search-area {
-  width: 99%;
-  max-width: 800px;
+  width: 83%;
 }
 
 .search-bar {
@@ -146,8 +144,7 @@ button.send-btn.stop-btn:hover {
   bottom: 1.5rem;
   left: 50%;
   transform: translateX(-50%);
-  width: 100%;
-  max-width: 800px;
+  width: 48%;
   background-color: #21212f;
   border-radius: 0.35rem;
 }
@@ -247,4 +244,4 @@ button.send-btn.stop-btn:hover {
     margin: 0;
     padding: 1rem;
   }
-}
+}

frontend/src/Components/AiPage.js

@@ -1,9 +1,12 @@
-import React, { useState, useEffect, useRef } from 'react';
-import './AiPage.css';
+import React, { useState, useEffect, useRef, useCallback, useMemo } from 'react';
+import { flushSync } from 'react-dom';
+import Snackbar from '@mui/material/Snackbar';
+import Alert from '@mui/material/Alert';
 import { FaCog, FaPaperPlane, FaStop } from 'react-icons/fa';
 import IntialSetting from './IntialSetting';
 import ChatWindow from './AiComponents/ChatWindow';
-import RightSidebar from './AiComponents/RightSidebar';
+import RightSidebar from './AiComponents/ChatComponents/RightSidebar';
+import './AiPage.css';
 
 function AiPage() {
   // Sidebar and other states
@@ -29,6 +32,24 @@ function AiPage() {
   const [activeBlockId, setActiveBlockId] = useState(null);
   const activeEventSourceRef = useRef(null);
 
+  // Snackbar state
+  const [snackbar, setSnackbar] = useState({
+    open: false,
+    message: "",
+    severity: "success",
+  });
+
+  // Function to open the snackbar
+  const openSnackbar = (message, severity = "success") => {
+    setSnackbar({ open: true, message, severity });
+  };
+
+  // Function to close the snackbar
+  const closeSnackbar = (event, reason) => {
+    if (reason === 'clickaway') return;
+    setSnackbar(prev => ({ ...prev, open: false }));
+  };
+
   useEffect(() => {
     localStorage.setItem("rightSidebarState", isRightSidebarOpen);
   }, [isRightSidebarOpen]);
@@ -37,6 +58,7 @@ function AiPage() {
     document.documentElement.style.setProperty('--right-sidebar-width', rightSidebarWidth + 'px');
   }, [rightSidebarWidth]);
 
+  // Dynamically increase height of chat input field based on newlines entered
   useEffect(() => {
     if (textAreaRef.current) {
       if (!defaultChatHeight) {
@@ -68,14 +90,16 @@ function AiPage() {
   }, [searchText, defaultChatHeight]);
 
   const handleOpenRightSidebar = (content, chatBlockId = null) => {
-    if (chatBlockId) {
-      setSelectedChatBlockId(chatBlockId);
-    }
-    setSidebarContent(content ? content : "default");
-    setRightSidebarOpen(true);
+    flushSync(() => {
+      if (chatBlockId) {
+        setSelectedChatBlockId(chatBlockId);
+      }
+      setSidebarContent(content ? content : "default");
+      setRightSidebarOpen(true);
+    });
   };
 
-  const handleEvaluationError = (blockId, errorMsg) => {
+  const handleEvaluationError = useCallback((blockId, errorMsg) => {
     setChatBlocks(prev =>
       prev.map(block =>
         block.id === blockId
@@ -83,7 +107,7 @@ function AiPage() {
           : block
       )
     );
-  };  
+  }, []);
 
   // Initiate the SSE
   const initiateSSE = (query, blockId) => {
@@ -93,12 +117,25 @@ function AiPage() {
     activeEventSourceRef.current = eventSource;
 
     eventSource.addEventListener("token", (e) => {
-      setChatBlocks(prev => prev.map(block =>
-        block.id === blockId
-          ? { ...block, aiAnswer: block.aiAnswer + e.data }
-          : block
-      ));
-    });
+      const { chunk, index } = JSON.parse(e.data);
+      console.log("[SSE token chunk]", JSON.stringify(chunk));
+      console.log("[SSE token index]", JSON.stringify(index));
+    
+      setChatBlocks(prevBlocks => {
+        return prevBlocks.map(block => {
+          if (block.id === blockId) {
+            const newTokenArray = block.tokenChunks ? [...block.tokenChunks] : [];
+            newTokenArray[index] = chunk;
+    
+            return {
+              ...block,
+              tokenChunks: newTokenArray
+            };
+          }
+          return block;
+        });
+      });
+    });    
 
     eventSource.addEventListener("final_message", (e) => {
       const endTime = Date.now();
@@ -148,7 +185,13 @@ function AiPage() {
       console.error("Error from SSE:", e.data);
       setChatBlocks(prev => prev.map(block =>
         block.id === blockId
-          ? { ...block, isError: true, errorMessage: e.data, aiAnswer: "" }
+          ? { 
+              ...block, 
+              isError: true, 
+              errorMessage: e.data, 
+              aiAnswer: "",
+              tasks: [] 
+            }
           : block
       ));
       eventSource.close();
@@ -228,6 +271,7 @@ function AiPage() {
       { 
         id: blockId, 
         userMessage: searchText, 
+        tokenChunks: [],
         aiAnswer: "", 
         thinkingTime: null, 
         thoughtLabel: "", 
@@ -299,6 +343,16 @@ function AiPage() {
     ? selectedBlock.actions.find(a => a.name === "evaluate")
     : null;
 
+  // Memoized evaluation object
+  const evaluation = useMemo(() => {
+    if (!evaluateAction) return null;
+    return {
+      ...evaluateAction.payload,
+      blockId: selectedBlock?.id,
+      onError: handleEvaluationError,
+    };
+  }, [evaluateAction, selectedBlock?.id, handleEvaluationError]);
+
   return (
     <div 
       className="app-container"
@@ -324,11 +378,7 @@ function AiPage() {
               if (!source || !source.link) return;
               window.open(source.link, '_blank');
             }}
-            evaluation={
-              evaluateAction
-                ? { ...evaluateAction.payload, blockId: selectedBlock?.id, onError: handleEvaluationError }
-                : null
-            }            
+            evaluation={evaluation}            
           />
         </div>
       )}
@@ -342,6 +392,7 @@ function AiPage() {
                   key={block.id}
                   blockId={block.id}
                   userMessage={block.userMessage}
+                  tokenChunks={block.tokenChunks}
                   aiAnswer={block.aiAnswer}
                   thinkingTime={block.thinkingTime}
                   thoughtLabel={block.thoughtLabel}
@@ -430,8 +481,20 @@ function AiPage() {
           trigger={true}
           setTrigger={() => setShowSettingsModal(false)}
           fromAiPage={true}
+          openSnackbar={openSnackbar}
+          closeSnackbar={closeSnackbar}
         />
       )}
+      <Snackbar
+        open={snackbar.open}
+        autoHideDuration={snackbar.severity === 'success' ? 3000 : null}
+        onClose={closeSnackbar}
+        anchorOrigin={{ vertical: 'top', horizontal: 'center' }}
+      >
+        <Alert onClose={closeSnackbar} severity={snackbar.severity} variant="filled" sx={{ width: '100%' }}>
+          {snackbar.message}
+        </Alert>
+      </Snackbar>
     </div>
   );
 }

frontend/src/Components/IntialSetting.css

@@ -171,4 +171,4 @@ input, select, textarea {
       width: 90%;
       max-height: 75vh; /* Adjust height for smaller screens */
   }
-}
+}

frontend/src/Components/IntialSetting.js

@@ -5,32 +5,21 @@ import Slider from '@mui/material/Slider';
 import Stack from '@mui/material/Stack';
 import Button from '@mui/material/Button';
 import IconButton from '@mui/material/IconButton';
-import Snackbar from '@mui/material/Snackbar';
-import Alert from '@mui/material/Alert';
 import './IntialSetting.css';
 import { FaTimes, FaEye, FaEyeSlash } from 'react-icons/fa';
 
 function IntialSetting(props) {
-  // Controlled states for Model Provider and sliders
+  // State variables for form controls
   const [selectedProvider, setSelectedProvider] = useState("OpenAI");
   const [modelTemperature, setModelTemperature] = useState(0.0);
   const [modelTopP, setModelTopP] = useState(1.0);
   const [showPassword, setShowPassword] = useState(false);
 
-  // Snackbar state for notifications
-  const [snackbar, setSnackbar] = useState({
-    open: false,
-    message: "",
-    severity: "success", // "success", "error", "info", "warning"
-  });
-
-  // Ref for the form element
+  // Ref for the form and navigation hook
   const formRef = useRef(null);
-
-  // React Router hook to navigate programmatically
   const navigate = useNavigate();
 
-  // Define model options
+  // Model options for different providers
   const modelOptions = {
     OpenAI: {
       "GPT-4 Turbo": "gpt-4-turbo",
@@ -60,7 +49,7 @@ function IntialSetting(props) {
     },
   };
 
-  // Reset handler: resets both the form (uncontrolled fields) and controlled states
+  // Reset form and state variables
   const handleReset = (e) => {
     e.preventDefault();
     if (formRef.current) {
@@ -71,98 +60,66 @@ function IntialSetting(props) {
     setModelTopP(1.0);
   };
 
-  // Snackbar close handler
-  const handleSnackbarClose = (event, reason) => {
-    if (reason === 'clickaway') return;
-    setSnackbar((prev) => ({ ...prev, open: false }));
-  };
-
-  // Save handler: validates the form, shows notifications, calls the parent's callback
-  // to update the underlying page (spinner/initializing text), sends the API request, and
-  // navigates to /AiPage when the backend returns success.
+  // Handle form submission and save settings
   const handleSave = async (e) => {
     e.preventDefault();
     const form = formRef.current;
 
-    // Retrieve values from form fields using their name attribute
+    // Retrieve form values
     const modelProvider = form.elements["model-provider"].value;
     const modelName = form.elements["model-name"].value;
     const modelAPIKeys = form.elements["model-api"].value;
     const braveAPIKey = form.elements["brave-api"].value;
     const proxyList = form.elements["proxy-list"].value;
-    // const neo4jURL = form.elements["neo4j-url"].value;
-    // const neo4jUsername = form.elements["neo4j-username"].value;
-    // const neo4jPassword = form.elements["neo4j-password"].value;
 
-    // Validate required fields and collect missing field names
+    // Check for missing required fields
     const missingFields = [];
     if (!modelProvider || modelProvider.trim() === "") missingFields.push("Model Provider");
     if (!modelName || modelName.trim() === "") missingFields.push("Model Name");
     if (!modelAPIKeys || modelAPIKeys.trim() === "") missingFields.push("Model API Key");
     if (!braveAPIKey || braveAPIKey.trim() === "") missingFields.push("Brave Search API Key");
-    // if (!neo4jURL || neo4jURL.trim() === "") missingFields.push("Neo4j URL");
-    // if (!neo4jUsername || neo4jUsername.trim() === "") missingFields.push("Neo4j Username");
-    // if (!neo4jPassword || neo4jPassword.trim() === "") missingFields.push("Neo4j Password");
 
-    // If any required fields are missing, show an error notification
     if (missingFields.length > 0) {
-      setSnackbar({
-        open: true,
-        message: "Please fill in the following required fields: " + missingFields.join(", "),
-        severity: "error",
-      });
+      props.openSnackbar(
+        "Please fill in the following required fields: " + missingFields.join(", "),
+        "error"
+      );
       return;
     }
 
-    // Build the JSON payload
+    // Build payload for backend
     const payload = {
       "Model_Provider": modelProvider.toLowerCase(),
       "Model_Name": modelName,
       "Model_API_Keys": modelAPIKeys,
       "Brave_Search_API_Key": braveAPIKey,
-      // "Neo4j_URL": neo4jURL,
-      // "Neo4j_Username": neo4jUsername,
-      // "Neo4j_Password": neo4jPassword,
       "Model_Temperature": modelTemperature,
       "Model_Top_P": modelTopP,
     };
 
-    // Include Proxy List if provided
     if (proxyList && proxyList.trim() !== "") {
       payload["Proxy_List"] = proxyList;
     }
 
-    // If opened from AiPage, show "Re-applying settings..." info notification with spinner
+    // Show appropriate notification based on context
     if (props.fromAiPage) {
-      setSnackbar({
-        open: true,
-        message: (
-          <Box mt={1} display="flex" alignItems="center">
-            <Box className="re-applying-settings-custom-spinner" />
-              <Box ml={1} className="re-applying-settings-text">
-                <span>
-                  Re-applying settings. This may take a few minutes...
-                </span>
-              </Box>
+      props.openSnackbar(
+        <Box mt={1} display="flex" alignItems="center">
+          <Box className="re-applying-settings-custom-spinner" />
+          <Box ml={1} className="re-applying-settings-text">
+            <span>Re-applying settings. This may take a few minutes...</span>
           </Box>
-        ),
-        severity: "info",
-      });
+        </Box>,
+        "info"
+      );
     } else {
-      // Original immediate success notification if opened from Home/App
-      setSnackbar({
-        open: true,
-        message: "Settings saved successfully!",
-        severity: "success",
-      });
-
-      // Call the parent's callback to change the underlying page's content (spinner/text)
+      props.openSnackbar("Settings saved successfully!", "success");
       if (props.onInitializationStart) {
         props.onInitializationStart();
       }
     }
 
-    // Send the payload to the backend
+    // Send settings to backend
     try {
       const response = await fetch("/settings", {
         method: "POST",
@@ -174,44 +131,24 @@ function IntialSetting(props) {
 
       if (response.ok) {
         const data = await response.json();
-        // When the backend returns {"success": true}, navigate to /AiPage
         if (data.success === true) {
-          // If from AiPage, show the final "Settings saved successfully!" success notification
           if (props.fromAiPage) {
-            setSnackbar({
-              open: true,
-              message: "Settings saved successfully!",
-              severity: "success",
-            });
+            props.openSnackbar("Settings saved successfully!", "success");
           }
           navigate("/AiPage");
         } else {
-          // If the response is OK but success is not true
-          setSnackbar({
-            open: true,
-            message: "Error saving settings. Please try again.",
-            severity: "error",
-          });
+          props.openSnackbar("Error saving settings. Please try again.", "error");
         }
       } else {
-        // If response is not OK, display error notification
-        setSnackbar({
-          open: true,
-          message: "Error saving settings. Please try again.",
-          severity: "error",
-        });
+        props.openSnackbar("Error saving settings. Please try again.", "error");
       }
     } catch (error) {
       console.error("Error saving settings:", error);
-      // Show error notification
-      setSnackbar({
-        open: true,
-        message: "Error saving settings. Please try again.",
-        severity: "error",
-      });
+      props.openSnackbar("Error saving settings. Please try again.", "error");
     }
   };
 
+  // Render the settings modal
   return props.trigger ? (
     <div className="showSetting" onClick={() => props.setTrigger(false)}>
       <div className="showSetting-inner" onClick={(e) => e.stopPropagation()}>
@@ -220,8 +157,6 @@ function IntialSetting(props) {
           <FaTimes />
         </button>
         <form ref={formRef}>
-
-          {/* Model Provider Selection */}
           <div className="form-group">
             <label htmlFor="model-provider">Model Provider</label>
             <select
@@ -237,8 +172,6 @@ function IntialSetting(props) {
               ))}
             </select>
           </div>
-
-          {/* Model Name Selection */}
           <div className="form-group">
             <label htmlFor="model-name">Model Name</label>
             <select id="model-name" name="model-name">
@@ -251,8 +184,6 @@ function IntialSetting(props) {
               )}
             </select>
           </div>
-
-          {/* API Key Inputs */}
           <div className="form-group">
             <label htmlFor="model-api">Model API Key</label>
             <textarea
@@ -270,8 +201,6 @@ function IntialSetting(props) {
               placeholder="Enter API Key"
             />
           </div>
-
-          {/* Proxy List */}
           <div className="form-group">
             <label htmlFor="proxy-list">Proxy List</label>
             <textarea
@@ -280,8 +209,7 @@ function IntialSetting(props) {
               placeholder="Enter proxies, one per line"
             ></textarea>
           </div>
-
-          {/* Neo4j Configuration */}
+          {/* Commented Neo4j configuration fields */}
           {/* <div className="form-group">
             <label htmlFor="neo4j-url">Neo4j URL</label>
             <input
@@ -313,17 +241,15 @@ function IntialSetting(props) {
                 onClick={() => setShowPassword(prev => !prev)}
                 className="password-toggle"
                 sx={{ 
-                  color: "white",  // Change the color of the icon
-                  p: 0,           // Remove internal padding
-                  m: 0            // Remove any margin
+                  color: "white", 
+                  p: 0, 
+                  m: 0 
                 }}
               >
                 {showPassword ? <FaEyeSlash /> : <FaEye />}
               </IconButton>
             </div>
           </div> */}
-
-          {/* Model Temperature and Top-P */}
           <div className="form-group">
             <div className="sliders-container">
               <div className="slider-item">
@@ -354,8 +280,6 @@ function IntialSetting(props) {
               </div>
             </div>
           </div>
-
-          {/* Buttons */}
           <Stack direction="row" spacing={2} sx={{ justifyContent: 'flex-end' }}>
             <Button 
               type="button" 
@@ -376,18 +300,6 @@ function IntialSetting(props) {
             </Button>
           </Stack>
         </form>
-        
-        {/* Notifications */}
-        <Snackbar
-          open={snackbar.open}
-          autoHideDuration={snackbar.severity === 'success' ? 3000 : null}
-          onClose={handleSnackbarClose}
-          anchorOrigin={{ vertical: 'top', horizontal: 'center' }}
-        >
-          <Alert onClose={handleSnackbarClose} severity={snackbar.severity} variant="filled" sx={{ width: '100%' }}>
-            {snackbar.message}
-          </Alert>
-        </Snackbar>
         {props.children}
       </div>
     </div>

main.py

@@ -18,8 +18,6 @@ from google.api_core.exceptions import ResourceExhausted
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 
-ENV_FILE_PATH = os.getenv("WRITABLE_DIR", "/tmp") + "/.env"
-
 CONTEXT_LENGTH = 128000
 BUFFER = 10000
 MAX_TOKENS_ALLOWED = CONTEXT_LENGTH - BUFFER
@@ -36,12 +34,21 @@ def format_error_sse(event_type: str, data: str) -> str:
     sse_message += "\n"
     return sse_message
 
+# Stop the task on error (non-fastapi)
+def stop_on_error():
+    state = SESSION_STORE
+
+    if "process_task" in state:
+        state["process_task"].cancel()
+        del state["process_task"]
+
 # Initialize the components
 def initialize_components():
-    load_dotenv(ENV_FILE_PATH, override=True)
+    load_dotenv(override=True)
 
     from src.search.search_engine import SearchEngine
     from src.query_processing.query_processor import QueryProcessor
+    # from src.rag.neo4j_graphrag import Neo4jGraphRAG
     from src.rag.graph_rag import GraphRAG
     from src.evaluation.evaluator import Evaluator
     from src.reasoning.reasoner import Reasoner
@@ -54,6 +61,7 @@ def initialize_components():
     SESSION_STORE['search_engine'] = SearchEngine()
     SESSION_STORE['query_processor'] = QueryProcessor()
     SESSION_STORE['crawler'] = CustomCrawler(max_concurrent_requests=1000)
+    # SESSION_STORE['graph_rag'] = Neo4jGraphRAG(num_workers=os.cpu_count() * 2)
     SESSION_STORE['graph_rag'] = GraphRAG(num_workers=os.cpu_count() * 2)
     SESSION_STORE['evaluator'] = Evaluator()
     SESSION_STORE['reasoner'] = Reasoner()
@@ -74,11 +82,13 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
 
         user_query = re.sub(r'category:.*', '', user_query, flags=re.IGNORECASE).strip()
 
-        if cat_lower == "internal knowledge base":
+        if cat_lower == "basic":
             response = ""
+            chunk_counter = 1
             async for chunk in state["reasoner"].reason(user_query):
+                await sse_queue.put(("token", json.dumps({"chunk": chunk, "index": chunk_counter})))
                 response += chunk
-                await sse_queue.put(("token", chunk))
+                chunk_counter += 1
 
             await sse_queue.put(("final_message", response))
             SESSION_STORE["chat_history"].append({"query": user_query, "response": response})
@@ -90,7 +100,7 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
 
             await sse_queue.put(("complete", "done"))
 
-        elif cat_lower == "simple external lookup":
+        elif cat_lower == "advanced":
             await sse_queue.put(("step", "Searching..."))
 
             optimized_query = await state['search_engine'].generate_optimized_query(user_query)
@@ -126,9 +136,11 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
                 await sse_queue.put(("sources_read", len(search_contents)))
 
                 response = ""
+                chunk_counter = 1
                 async for chunk in state["reasoner"].reason(user_query, contents):
+                    await sse_queue.put(("token", json.dumps({"chunk": chunk, "index": chunk_counter})))
                     response += chunk
-                    await sse_queue.put(("token", chunk))
+                    chunk_counter += 1
 
                 await sse_queue.put(("final_message", response))
                 SESSION_STORE["chat_history"].append({"query": user_query, "response": response})
@@ -146,7 +158,7 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
             else:
                 await sse_queue.put(("error", "No results found."))
 
-        elif cat_lower == "complex moderate decomposition":
+        elif cat_lower == "pro":
             current_search_results = []
             current_search_contents = []
 
@@ -202,7 +214,11 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
                     await sse_queue.put(("task", (sub_query, "FAILED")))
                     return ""
 
-            tasks = [sub_query_task(sub_query) for sub_query in sub_queries]
+            tasks = []
+            if len(sub_queries) > 1 and sub_queries[0] != user_query:
+                for sub_query in sub_queries:
+                    tasks.append(sub_query_task(sub_query))
+                    
             results = await asyncio.gather(*tasks)
             end = time.time()
 
@@ -234,14 +250,16 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
                 await sse_queue.put(("sources_read", len(current_search_contents)))
                 
                 response = ""
+                chunk_counter = 1
                 is_first_chunk = True
                 async for chunk in state['reasoner'].reason(user_query, contents):
                     if is_first_chunk:
                         await sse_queue.put(("step", f"Thought and searched for {int(end - start)} seconds"))
                         is_first_chunk = False
 
+                    await sse_queue.put(("token", json.dumps({"chunk": chunk, "index": chunk_counter})))
                     response += chunk
-                    await sse_queue.put(("token", chunk))
+                    chunk_counter += 1
 
                 await sse_queue.put(("final_message", response))
                 SESSION_STORE["chat_history"].append({"query": user_query, "response": response})
@@ -262,7 +280,7 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
             else:
                 await sse_queue.put(("error", "No results found."))
 
-        elif cat_lower == "complex advanced decomposition":
+        elif cat_lower == "super":
             current_search_results = []
             current_search_contents = []
 
@@ -314,7 +332,11 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
 
                     sub_sub_queries, _ = await state['query_processor'].decompose_query(sub_query)
 
-                    tasks = [sub_sub_query_task(sub_sub_query) for sub_sub_query in sub_sub_queries]
+                    tasks = []
+                    if len(sub_sub_queries) > 1 and sub_sub_queries[0] != user_query:
+                        for sub_sub_query in sub_sub_queries:
+                            tasks.append(sub_sub_query_task(sub_sub_query))
+
                     results = await asyncio.gather(*tasks)
 
                     if any(result.strip() for result in results):
@@ -328,7 +350,11 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
                     await sse_queue.put(("task", (sub_query, "FAILED")))
                     return []
 
-            tasks = [sub_query_task(sub_query) for sub_query in sub_queries]
+            tasks = []
+            if len(sub_queries) > 1 and sub_queries[0] != user_query:
+                for sub_query in sub_queries:
+                    tasks.append(sub_query_task(sub_query))
+
             results = await asyncio.gather(*tasks)
             end = time.time()
 
@@ -366,14 +392,16 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
                 await sse_queue.put(("sources_read", len(current_search_contents)))
                 
                 response = ""
+                chunk_counter = 1
                 is_first_chunk = True
                 async for chunk in state['reasoner'].reason(user_query, contents):
                     if is_first_chunk:
                         await sse_queue.put(("step", f"Thought and searched for {int(end - start)} seconds"))
                         is_first_chunk = False
 
+                    await sse_queue.put(("token", json.dumps({"chunk": chunk, "index": chunk_counter})))
                     response += chunk
-                    await sse_queue.put(("token", chunk))
+                    chunk_counter += 1
 
                 await sse_queue.put(("final_message", response))
                 SESSION_STORE["chat_history"].append({"query": user_query, "response": response})
@@ -394,7 +422,7 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
             else:
                 await sse_queue.put(("error", "No results found."))
 
-        elif cat_lower == "extensive research dynamic structuring":
+        elif cat_lower == "ultra":
             current_search_results = []
             current_search_contents = []
 
@@ -454,6 +482,7 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
             state['graph_rag'].set_on_event_callback(on_event_callback)
 
             start = time.time()
+            # state['graph_rag'].initialize_schema()
             await state['graph_rag'].process_graph(
                 user_query,
                 similarity_threshold=0.8,
@@ -489,14 +518,16 @@ async def process_query(user_query: str, sse_queue: asyncio.Queue):
                 await sse_queue.put(("sources_read", len(current_search_contents)))
 
                 response = ""
+                chunk_counter = 1
                 is_first_chunk = True
                 async for chunk in state['reasoner'].reason(user_query, answer):
                     if is_first_chunk:
                         await sse_queue.put(("step", f"Thought and searched for {int(end - start)} seconds"))
                         is_first_chunk = False
 
+                    await sse_queue.put(("token", json.dumps({"chunk": chunk, "index": chunk_counter})))
                     response += chunk
-                    await sse_queue.put(("token", chunk))
+                    chunk_counter += 1
 
                 await sse_queue.put(("final_message", response))
                 SESSION_STORE["chat_history"].append({"query": user_query, "response": response})
@@ -677,7 +708,7 @@ async def sse_message(request: Request, user_message: str):
 
     async def event_generator():
         # Build the prompt
-        context = state["chat_history"][-5:]
+        context = state["chat_history"][-3:]
         if context:
             prompt = \
 f"""This is the previous context of the conversation:
@@ -705,6 +736,7 @@ Current Query:
                     yield f"event: final_message\ndata: {data}\n\n"
 
                 elif event_type == "error":
+                    stop_on_error()
                     yield format_error_sse("error", data)
 
                 elif event_type == "step":

src/crawl/crawler.py

@@ -778,7 +778,7 @@ Query:"""
         self, 
         url: str, 
         query: Optional[str] = None,
-        timeout: float = 10.0,
+        timeout: float = 5.0,
         return_type: str = "markdown",
         rotate_headers: bool = True,
         proxies: Optional[List[str]] = None,
@@ -813,12 +813,8 @@ Query:"""
         headers = self.get_headers() if rotate_headers else {}
         proxy = self.get_proxy(proxies) if proxies else None
 
-        # Split timeout into connection and processing phases
-        timeout_config = aiohttp.ClientTimeout(
-            total=None,  # No total timeout
-            connect=timeout,  # Connection timeout
-            sock_read=timeout  # Socket read timeout
-        )
+        # Total connection timeout
+        timeout_config = aiohttp.ClientTimeout(total=timeout)
 
         try:
             # Use provided session if available

src/helpers/helper.py

@@ -5,7 +5,7 @@ import torch
 import transformers
 from langchain.text_splitter import RecursiveCharacterTextSplitter, TokenTextSplitter
 
-ENV_FILE_PATH = os.path.join(os.getenv("WRITABLE_DIR", "/tmp"), ".env")
+ENV_FILE_PATH = os.path.join(os.path.dirname(__file__), "../../.env").replace("\\", "/")
 
 def remove_markdown(text: str) -> str:
         # Remove code block format type and the code block itself
@@ -48,6 +48,26 @@ def remove_markdown(text: str) -> str:
         
         return text.strip()
 
+def remove_outer_markdown_block(chunk, _acc={"b":""}):
+    _acc["b"] += chunk
+    p = re.compile(r'```markdown\s*\n(.*?)\n?```', re.DOTALL|re.IGNORECASE)
+    o = []
+
+    while True:
+        m = p.search(_acc["b"])
+        if not m:
+             break
+        
+        s,e = m.span()
+        o.append(_acc["b"][:s]+m.group(1))
+        _acc["b"] = _acc["b"][e:]
+
+    if '```markdown' not in _acc["b"].lower():
+        o.append(_acc["b"])
+        _acc["b"] = ""
+
+    return "".join(o)
+
 def clear_gpu_memory():
         # Clear GPU memory and cache if available
         if torch.cuda.is_available():

src/query_processing/query_processor.py

@@ -16,11 +16,11 @@ class QueryProcessor:
     async def classify_query(self, query, *, llm):
         template = \
 """You are an expert at classifying queries. Your task is to classify the given user query (descriptions are provided below) into one of the following categories:        
-1. Internal Knowledge Base -> The query is a simple question that can be answered without external internet resources. It is the simplest category.
-2. Simple External Lookup -> The query requires up to 5 external internet resource references for a quick fact-check before answering. It is a moderate category.
-3. Complex Moderate Decomposition -> The query needs complex reasoning and up to 20 external internet resource references to answer accurately. It is a complex category.
-4. Complex Advanced Decomposition -> The query requires very complex reasoning, dynamic sub-query generation, and up to 50 external internet resource references to answer accurately. It is a very complex category.
-5. Extensive Research Dynamic Structuring -> The query requires extensive research, extremely complex reasoning, and up to 500 external internet resource references to answer accurately. It is the most complex category.
+1. Basic -> The query is a simple question that can be answered without external internet resources. It is the simplest category.
+2. Advanced -> The query requires up to 5 external internet resource references for a quick fact-check before answering. It is a moderate category.
+3. Pro -> The query needs complex reasoning and up to 20 external internet resource references to answer accurately. It is a complex category.
+4. Super -> The query requires very complex reasoning, dynamic sub-query generation, and up to 50 external internet resource references to answer accurately. It is a very complex category.
+5. Ultra -> The query requires extensive research, extremely complex reasoning, and up to 500 external internet resource references to answer accurately. It is the most complex category.
 
 Rules:
 1. Only classify the query into one of the categories provided above.
@@ -28,7 +28,7 @@ Rules:
 3. If the query contains "category: <category_name>" in it, then the query should be classified into the category that is mentioned in the query.
 
 Examples:
-1. Internal Knowledge Base
+1. Basic
     - "What is the capital of France?"
     - "How many continents are there in the world?"
     - "What is the chemical formula for water?"
@@ -37,7 +37,7 @@ Examples:
     - "What is Newton's first law of motion?"
     - "What is the sum of the angles in a triangle?"
 
-2. Simple External Lookup
+2. Advanced
     - "What is the current market price of gold?"
     - "What is the current population of Japan?"
     - "Who won the Nobel Peace Prize in 2023?"
@@ -45,7 +45,7 @@ Examples:
     - "How many countries are in the European Union?"
     - "What are the top five best-selling novels of 2024?"
 
-3. Complex Moderate Decomposition
+3. Pro
     - "Compare the political structures of ancient Greece and ancient Rome"
     - "How do climate change and deforestation contribute to the decline in biodiversity?"
     - "What are the major differences in healthcare systems between the U.S., Canada, and the U.K.?"
@@ -53,7 +53,7 @@ Examples:
     - "What are the economic and social impacts of artificial intelligence on the job market?"
     - "What are the pros and cons of genetically modified organisms in agriculture?"
 
-4. Complex Advanced Decomposition
+4. Super
     - "What are the long-term psychological and societal effects of internet addiction on teenagers?"
     - "How do various world religions view the concept of life after death, and what cultural practices reflect these beliefs?"
     - "What are the implications of quantum computing on encryption technologies, and how might they evolve in the next decade?"
@@ -61,7 +61,7 @@ Examples:
     - "What are the top five electric SUVs available in the market as of October 2024 for under $60,000, and how do they compare in terms of range, features, and performance?"
     - "What are the major obstacles to achieving carbon neutrality in heavy industries like steel and cement? What are the potential solutions?"
 
-5. Extensive Research Dynamic Structuring
+5. Ultra
     - "Create a comprehensive study on the economic, cultural, and environmental impacts of megacities around the world, including case studies from at least five different continents."
     - "Conduct a detailed comparative analysis of various renewable energy policies across G20 countries, evaluating their effectiveness, challenges, and future projections."
     - "Research the evolution of space exploration programs in the last 60 years, highlighting the key milestones, technological advancements, and geopolitical factors that shaped them."

src/reasoning/reasoner.py

@@ -1,7 +1,6 @@
 from langchain.prompts import ChatPromptTemplate
 from langchain_core.prompts import ChatPromptTemplate
-from src.utils.api_key_manager import APIKeyManager
-from src.utils.api_key_manager import with_api_manager
+from src.utils.api_key_manager import APIKeyManager, with_api_manager
 from src.query_processing.late_chunking.late_chunker import LateChunker
 
 class Reasoner:
@@ -23,7 +22,7 @@ class Reasoner:
 Your task is to reason about the given user query and provide an answer.
 
 Rules:
-1. Your response should only be the answer in plain text without any formatting.
+1. Your response should only be the answer in valid markdown format.
 2. You must use proper reasoning and logic to answer the query for your internal use but do not show your reasoning process in the response.
 
 Query: 
@@ -36,7 +35,7 @@ Query:
 Given the user query and the relevant context, your task is to reason and provide an answer.
 
 Rules:
-1. Your response should only be the answer in plain text without any formatting.
+1. Your response should only be the answer in valid markdown format.
 2. You must use proper reasoning and logic to answer the query for your internal use but do not show your reasoning process in the response.
 3. You must not mention the context/documents provided to you in the response. Make it sound like you are the one who is answering the query.
 
@@ -116,5 +115,4 @@ if __name__ == "__main__":
         rotate_proxy=False,
         return_html=True
     ))
-    print(contents)
-
+    print(contents)