Merge branch 'main' into release/v0.4.0

backend/app/api/v1/chat.py

@@ -65,6 +65,24 @@ def chat(project_id: int, chat_request: ChatRequest, db: Session = Depends(get_d
 
         ordered_file_names = [doc_id_to_filename[doc_id] for doc_id in doc_ids]
 
+        extract_vectorstore = ChromaDB(f"panda-etl-extraction-{project_id}",
+                                       similarity_threshold=settings.chat_extraction_doc_threshold)
+
+        # Extract reference documents from the extraction results from db
+        extraction_docs = extract_vectorstore.get_relevant_docs(
+                chat_request.query,
+                k=settings.chat_extraction_max_docs
+        )
+
+        # Append text from single documents together
+        for extraction_doc in extraction_docs["metadatas"][0]:
+            index = next((i for i, item in enumerate(ordered_file_names) if item == extraction_doc["filename"]), None)
+            if index is None:
+                ordered_file_names.append(extraction_doc["filename"])
+                docs.append(extraction_doc["reference"])
+            else:
+                docs[index] = f'{extraction_doc["reference"]}\n\n{docs[index]}'
+
         docs_formatted = [
             {"filename": filename, "quote": quote}
             for filename, quote in zip(ordered_file_names, docs)

backend/app/config.py

@@ -26,6 +26,10 @@ class Settings(BaseSettings):
     openai_api_key: str = ""
     openai_embedding_model: str = "text-embedding-ada-002"
 
+    # Extraction References for chat
+    chat_extraction_doc_threshold: float = 0.5
+    chat_extraction_max_docs: int = 50
+
     class Config:
         env_file = ".env"
 

backend/app/processing/process_queue.py

@@ -46,8 +46,9 @@ def process_step_task(
         # Initial DB operations (open and fetch relevant data)
         with SessionLocal() as db:
             process = process_repository.get_process(db, process_id)
+            project_id = process.project_id
             process_step = process_repository.get_process_step(db, process_step_id)
-
+            filename =  process_step.asset.filename
             if process.status == ProcessStatus.STOPPED:
                 return False  # Stop processing if the process is stopped
 
@@ -84,6 +85,15 @@ def process_step_task(
                             output_references=data["context"],
                         )
 
+                    # vectorize extraction result
+                    try:
+                        vectorize_extraction_process_step(project_id=project_id,
+                                                        process_step_id=process_step_id,
+                                                        filename=filename,
+                                                        references=data["context"])
+                    except Exception :
+                        logger.error(f"Failed to vectorize extraction results for chat {traceback.format_exc()}")
+
                 success = True
 
             except CreditLimitExceededException:
@@ -361,3 +371,39 @@ def update_process_step_status(
     process_repository.update_process_step_status(
         db, process_step, status, output=output, output_references=output_references
     )
+
+def vectorize_extraction_process_step(project_id: int, process_step_id: int, filename: str, references: dict) -> None:
+    # Vectorize extraction result and dump in database
+    field_references = {}
+
+    # Loop to concatenate sources for each reference
+    for extraction_references in references:
+        for extraction_reference in extraction_references:
+            sources = extraction_reference.get("sources", [])
+            if sources:
+                sources_catenated = "\n".join(sources)
+                field_references.setdefault(extraction_reference["name"], "")
+                field_references[extraction_reference["name"]] += (
+                    "\n" + sources_catenated if field_references[extraction_reference["name"]] else sources_catenated
+                )
+
+    # Only proceed if there are references to add
+    if not field_references:
+        return
+
+    # Initialize Vectorstore
+    vectorstore = ChromaDB(f"panda-etl-extraction-{project_id}")
+
+    docs = [f"{filename} {key}" for key in field_references]
+    metadatas = [
+        {
+            "project_id": project_id,
+            "process_step_id": process_step_id,
+            "filename": filename,
+            "reference": reference
+        }
+        for reference in field_references.values()
+    ]
+
+    # Add documents to vectorstore
+    vectorstore.add_docs(docs=docs, metadatas=metadatas)

backend/tests/processing/test_process_queue.py

@@ -1,11 +1,12 @@
 from app.requests.schemas import ExtractFieldsResponse
 import pytest
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, Mock, patch
 from app.processing.process_queue import (
     handle_exceptions,
     extract_process,
     update_process_step_status,
     find_best_match_for_short_reference,
+    vectorize_extraction_process_step,
 )
 from app.exceptions import CreditLimitExceededException
 from app.models import ProcessStepStatus
@@ -180,3 +181,107 @@ def test_chroma_db_initialization(mock_extract_data, mock_chroma):
 
     mock_chroma.assert_called_with(f"panda-etl-{process.project_id}", similarity_threshold=3)
     assert mock_chroma.call_count >= 1
+
+@patch('app.processing.process_queue.ChromaDB')
+def test_vectorize_extraction_process_step_single_reference(mock_chroma_db):
+    # Mock ChromaDB instance
+    mock_vectorstore = MagicMock()
+    mock_chroma_db.return_value = mock_vectorstore
+
+    # Inputs
+    project_id = 123
+    process_step_id = 1
+    filename = "sample_file"
+    references = [
+        [
+            {"name": "field1", "sources": ["source1", "source2"]}
+        ]
+    ]
+
+    # Call function
+    vectorize_extraction_process_step(project_id, process_step_id, filename, references)
+
+    # Expected docs and metadata to add to ChromaDB
+    expected_docs = ["sample_file field1"]
+    expected_metadatas = [
+        {
+            "project_id": project_id,
+            "process_step_id": process_step_id,
+            "filename": filename,
+            "reference": "source1\nsource2"
+        }
+    ]
+
+    # Assertions
+    mock_vectorstore.add_docs.assert_called_once_with(
+        docs=expected_docs,
+        metadatas=expected_metadatas
+    )
+
+@patch('app.processing.process_queue.ChromaDB')
+def test_vectorize_extraction_process_step_multiple_references_concatenation(mock_chroma_db):
+    # Mock ChromaDB instance
+    mock_vectorstore = MagicMock()
+    mock_chroma_db.return_value = mock_vectorstore
+
+    # Inputs
+    project_id = 456
+    process_step_id = 2
+    filename = "test_file"
+    references = [
+        [
+            {"name": "field1", "sources": ["source1", "source2"]},
+            {"name": "field1", "sources": ["source3"]}
+        ],
+        [
+            {"name": "field2", "sources": ["source4"]}
+        ]
+    ]
+
+    # Call function
+    vectorize_extraction_process_step(project_id, process_step_id, filename, references)
+
+    # Expected docs and metadata to add to ChromaDB
+    expected_docs = ["test_file field1", "test_file field2"]
+    expected_metadatas = [
+        {
+            "project_id": project_id,
+            "process_step_id": process_step_id,
+            "filename": filename,
+            "reference": "source1\nsource2\nsource3"
+        },
+        {
+            "project_id": project_id,
+            "process_step_id": process_step_id,
+            "filename": filename,
+            "reference": "source4"
+        }
+    ]
+
+    # Assertions
+    mock_vectorstore.add_docs.assert_called_once_with(
+        docs=expected_docs,
+        metadatas=expected_metadatas
+    )
+
+@patch('app.processing.process_queue.ChromaDB')  # Replace with the correct module path
+def test_vectorize_extraction_process_step_empty_sources(mock_chroma_db):
+    # Mock ChromaDB instance
+    mock_vectorstore = MagicMock()
+    mock_chroma_db.return_value = mock_vectorstore
+
+    # Inputs
+    project_id = 789
+    process_step_id = 3
+    filename = "empty_sources_file"
+    references = [
+        [
+            {"name": "field1", "sources": []}
+        ]
+    ]
+
+    # Call function
+    vectorize_extraction_process_step(project_id, process_step_id, filename, references)
+
+    # Expected no calls to add_docs due to empty sources
+    mock_vectorstore.add_docs.assert_not_called()

frontend/.env.example

@@ -1,3 +1,4 @@
 NEXT_PUBLIC_API_URL=http://localhost:3000/api/v1
 NEXT_PUBLIC_STORAGE_URL=http://localhost:3000/api/assets
-NEXT_PUBLIC_MIXPANEL_TOKEN=f2e8a71ab2bde33ebf346c5abf6ba9fa
+NEXT_PUBLIC_MIXPANEL_TOKEN=f2e8a71ab2bde33ebf346c5abf6ba9fa
+NEXT_PUBLIC_ROLLBAR_ACCESS_TOKEN=0df0bee895044430880278e2b2a5b2d2

frontend/package.json

@@ -16,6 +16,7 @@
     "@radix-ui/react-context-menu": "^2.2.1",
     "@radix-ui/react-radio-group": "^1.2.0",
     "@react-pdf/renderer": "^4.0.0",
+    "@rollbar/react": "^0.12.0-beta",
     "@tanstack/react-query": "^5.51.1",
     "@tippyjs/react": "^4.2.6",
     "@types/mixpanel-browser": "^2.50.1",
@@ -39,6 +40,7 @@
     "react-quill": "^2.0.0",
     "rehype-sanitize": "^6.0.0",
     "remark-gfm": "^4.0.0",
+    "rollbar": "^2.26.4",
     "tailwind-merge": "^2.4.0"
   },
   "devDependencies": {

frontend/src/app/(app)/layout.tsx

@@ -4,6 +4,8 @@ import { ScrollProvider } from "@/context/ScrollContext";
 import Sidebar from "@/components/ui/Sidebar";
 import Navbar from "@/components/ui/Navbar";
 import "@/app/style/globals.css";
+import { ErrorBoundary } from "@rollbar/react";
+import { ViewErrorFallback } from "@/components/ErrorFallback";
 
 export default function RootLayout({
   children,
@@ -17,9 +19,11 @@ export default function RootLayout({
           <Sidebar />
           <div className="flex-1 flex flex-col overflow-hidden md:ml-64">
             <Navbar />
-            <main className="flex-1 overflow-x-hidden overflow-y-auto bg-gray-50 px-10 py-6">
-              {children}
-            </main>
+            <ErrorBoundary fallbackUI={ViewErrorFallback}>
+              <main className="flex-1 overflow-x-hidden overflow-y-auto bg-gray-50 px-10 py-6">
+                {children}
+              </main>
+            </ErrorBoundary>
           </div>
         </div>
       </ScrollProvider>

frontend/src/app/layout.tsx

@@ -4,9 +4,16 @@ import React from "react";
 import { ReactQueryClientProvider } from "@/components/ReactQueryClientProvider";
 import { Toaster } from "react-hot-toast";
 import "@/app/style/globals.css";
+import { Provider, ErrorBoundary } from "@rollbar/react";
+import { GlobalErrorFallback } from "@/components/ErrorFallback";
 
 const inter = Inter({ subsets: ["latin"] });
 
+const rollbarConfig = {
+  accessToken: process.env.NEXT_PUBLIC_ROLLBAR_ACCESS_TOKEN,
+  environment: process.env.NODE_ENV ?? "production",
+};
+
 export const metadata: Metadata = {
   title: "PandaETL",
   description:
@@ -22,14 +29,18 @@ export default function RootLayout({
   children: React.ReactNode;
 }>) {
   return (
-    <ReactQueryClientProvider>
-      <html lang="en">
-        <head />
-        <body className={inter.className}>
-          <Toaster position="top-right" />
-          {children}
-        </body>
-      </html>
-    </ReactQueryClientProvider>
+    <html lang="en">
+      <head />
+      <body className={inter.className}>
+        <Provider config={rollbarConfig}>
+          <ErrorBoundary fallbackUI={GlobalErrorFallback}>
+            <ReactQueryClientProvider>
+              <Toaster position="top-right" />
+              {children}
+            </ReactQueryClientProvider>
+          </ErrorBoundary>
+        </Provider>
+      </body>
+    </html>
   );
 }

frontend/src/components/ErrorFallback.tsx

@@ -0,0 +1,49 @@
+"use client";
+import React, { useEffect, useState } from "react";
+import { Button } from "./ui/Button";
+import { usePathname } from "next/navigation";
+
+interface ErrorFallbackProps {
+  error: Error | null;
+  resetError: () => void;
+}
+
+interface ErrorFallbackBaseProps extends ErrorFallbackProps {
+  textColor: string;
+}
+
+// Fallback UI component displayed on error
+const ErrorFallbackBase = ({
+  error,
+  resetError,
+  textColor,
+}: ErrorFallbackBaseProps) => {
+  const pathname = usePathname();
+  const [initialPathname, setInitialPathname] = useState<string | null>(null);
+
+  useEffect(() => {
+    if (initialPathname === null) {
+      setInitialPathname(pathname);
+    } else if (pathname !== initialPathname && error) {
+      resetError();
+    }
+  }, [pathname, initialPathname, error, resetError]);
+
+  return (
+    <div className={`p-8 text-center ${textColor}`}>
+      <h1>Oops! Something went wrong.</h1>
+      <p>{error?.message || "We're working on fixing this issue."}</p>
+      <Button onClick={resetError} className="mt-4">
+        Try Again
+      </Button>
+    </div>
+  );
+};
+
+export const GlobalErrorFallback = (props: ErrorFallbackProps) => (
+  <ErrorFallbackBase {...props} textColor="text-white" />
+);
+
+export const ViewErrorFallback = (props: ErrorFallbackProps) => (
+  <ErrorFallbackBase {...props} textColor="text-black" />
+);