Merge pull request #142 from MODSetter/dev

feat: Document Selector in Chat.

Author: MODSetter

surfsense_backend/app/agents/researcher/configuration.py

@@ -33,6 +33,7 @@ class Configuration:
     search_space_id: int
     search_mode: SearchMode
     research_mode: ResearchMode
+    document_ids_to_add_in_context: List[int]
 
     @classmethod
     def from_runnable_config(

surfsense_backend/app/agents/researcher/nodes.py

@@ -30,6 +30,7 @@ async def handle_chat_data(
     search_space_id = request.data.get('search_space_id')
     research_mode: str = request.data.get('research_mode')
     selected_connectors: List[str] = request.data.get('selected_connectors')
+    document_ids_to_add_in_context: List[int] = request.data.get('document_ids_to_add_in_context')
 
     search_mode_str = request.data.get('search_mode', "CHUNKS")
 
@@ -71,7 +72,8 @@ async def handle_chat_data(
         research_mode,
         selected_connectors,
         langchain_chat_history,
-        search_mode_str
+        search_mode_str,
+        document_ids_to_add_in_context
     ))
     response.headers['x-vercel-ai-data-stream'] = 'v1'
     return response

surfsense_backend/app/routes/chats_routes.py

@@ -17,7 +17,8 @@ async def stream_connector_search_results(
     research_mode: str, 
     selected_connectors: List[str],
     langchain_chat_history: List[Any],
-    search_mode_str: str
+    search_mode_str: str,
+    document_ids_to_add_in_context: List[int]
 ) -> AsyncGenerator[str, None]:
     """
     Stream connector search results to the client
@@ -62,7 +63,8 @@ async def stream_connector_search_results(
             "user_id": user_id_str,
             "search_space_id": search_space_id,
             "search_mode": search_mode,
-            "research_mode": research_mode
+            "research_mode": research_mode,
+            "document_ids_to_add_in_context": document_ids_to_add_in_context
         }
     }
     # Initialize state with database session and streaming service

surfsense_backend/app/tasks/stream_connector_search_results.py

@@ -4,7 +4,7 @@
 from sqlalchemy.future import select
 from app.retriver.chunks_hybrid_search import ChucksHybridSearchRetriever
 from app.retriver.documents_hybrid_search import DocumentHybridSearchRetriever
-from app.db import SearchSourceConnector, SearchSourceConnectorType, Chunk, Document
+from app.db import SearchSourceConnector, SearchSourceConnectorType, Chunk, Document, SearchSpace
 from tavily import TavilyClient
 from linkup import LinkupClient
 from sqlalchemy import func
@@ -33,7 +33,8 @@ async def initialize_counter(self):
                 result = await self.session.execute(
                     select(func.count(Chunk.id))
                     .join(Document)
-                    .filter(Document.user_id == self.user_id)
+                    .join(SearchSpace)
+                    .filter(SearchSpace.user_id == self.user_id)
                 )
                 chunk_count = result.scalar() or 0
                 self.source_id_counter = chunk_count + 1

surfsense_backend/app/utils/connector_service.py

@@ -72,7 +72,7 @@ export default function EditConnectorPage() {
   const [connector, setConnector] = useState<SearchSourceConnector | null>(null);
   const [isLoading, setIsLoading] = useState(true);
   const [isSubmitting, setIsSubmitting] = useState(false);
-  console.log("connector", connector);
+  // console.log("connector", connector);
   // Initialize the form
   const form = useForm<ApiConnectorFormValues>({
     resolver: zodResolver(apiConnectorFormSchema),

surfsense_web/app/dashboard/[search_space_id]/connectors/[connector_id]/page.tsx

@@ -283,8 +283,8 @@ export default function DocumentsTable() {
     const searchSpaceId = Number(params.search_space_id);
     const { documents, loading, error, refreshDocuments, deleteDocument } = useDocuments(searchSpaceId);
 
-    console.log("Search Space ID:", searchSpaceId);
-    console.log("Documents loaded:", documents?.length);
+    // console.log("Search Space ID:", searchSpaceId);
+    // console.log("Documents loaded:", documents?.length);
 
     useEffect(() => {
         console.log("Delete document function available:", !!deleteDocument);
@@ -315,7 +315,7 @@ export default function DocumentsTable() {
 
     const handleDeleteRows = async () => {
         const selectedRows = table.getSelectedRowModel().rows;
-        console.log("Deleting selected rows:", selectedRows.length);
+        // console.log("Deleting selected rows:", selectedRows.length);
 
         if (selectedRows.length === 0) {
             toast.error("No rows selected");
@@ -324,14 +324,14 @@ export default function DocumentsTable() {
 
         // Create an array of promises for each delete operation
         const deletePromises = selectedRows.map(row => {
-            console.log("Deleting row with ID:", row.original.id);
+            // console.log("Deleting row with ID:", row.original.id);
             return deleteDocument(row.original.id);
         });
 
         try {
             // Execute all delete operations
             const results = await Promise.all(deletePromises);
-            console.log("Delete results:", results);
+            // console.log("Delete results:", results);
 
             // Check if all deletions were successful
             const allSuccessful = results.every(result => result === true);