Merge dev into main

cognee-frontend/src/app/dashboard/AddDataToCognee.tsx

@@ -1,6 +1,8 @@
 import { FormEvent, useCallback, useState } from "react";
-import { CloseIcon, PlusIcon } from "@/ui/Icons";
+
+import { LoadingIndicator } from "@/ui/App";
 import { useModal } from "@/ui/elements/Modal";
+import { CloseIcon, PlusIcon } from "@/ui/Icons";
 import { CTAButton, GhostButton, IconButton, Modal, NeutralButton, Select } from "@/ui/elements";
 
 import addData from "@/modules/ingestion/addData";
@@ -44,14 +46,19 @@ export default function AddDataToCognee({ datasets, refreshDatasets, useCloud =
     )
       .then(({ dataset_id, dataset_name }) => {
         refreshDatasets();
-        setFilesForUpload(null);
-
-        return cognifyDataset({
-          id: dataset_id,
-          name: dataset_name,
-          data: [],  // not important, just to mimick Dataset
-          status: "",  // not important, just to mimick Dataset
-        }, useCloud);
+
+        return cognifyDataset(
+          {
+            id: dataset_id,
+            name: dataset_name,
+            data: [],  // not important, just to mimick Dataset
+            status: "",  // not important, just to mimick Dataset
+          },
+          useCloud,
+        )
+          .then(() => {
+            setFilesForUpload(null);
+          });
       });
   }, [filesForUpload, refreshDatasets, useCloud]);
 
@@ -76,7 +83,7 @@ export default function AddDataToCognee({ datasets, refreshDatasets, useCloud =
             <span className="text-2xl">Add new data to a dataset?</span>
             <IconButton disabled={isProcessingDataWithCognee} onClick={closeAddDataModal}><CloseIcon /></IconButton>
           </div>
-          <div className="mt-8 mb-6">Please select a dataset to add data in.<br/> If you don&apos;t have any, don&apos;t worry, we will create one for you.</div>
+          <div className="mt-8 mb-6">Please select a {useCloud ? "cloud" : "local"} dataset to add data in.<br/> If you don&apos;t have any, don&apos;t worry, we will create one for you.</div>
           <form onSubmit={submitDataToCognee}>
             <div className="max-w-md flex flex-col gap-4">
               <Select name="datasetName">
@@ -105,7 +112,8 @@ export default function AddDataToCognee({ datasets, refreshDatasets, useCloud =
             <div className="flex flex-row gap-4 mt-4 justify-end">
               <GhostButton disabled={isProcessingDataWithCognee} type="button" onClick={() => closeAddDataModal()}>cancel</GhostButton>
               <CTAButton disabled={isProcessingDataWithCognee} type="submit">
-                {isProcessingDataWithCognee ? "processing..." : "add"}
+                {isProcessingDataWithCognee && <LoadingIndicator color="white" />}
+                add
               </CTAButton>
             </div>
           </form>

cognee-frontend/src/app/dashboard/Dashboard.tsx

@@ -54,10 +54,12 @@ export default function Dashboard({ accessToken }: DashboardProps) {
   const [selectedNotebookId, setSelectedNotebookId] = useState<string | null>(null);
 
   const handleNotebookRemove = useCallback((notebookId: string) => {
-    setSelectedNotebookId((currentSelectedNotebookId) => (
-      currentSelectedNotebookId === notebookId ? null : currentSelectedNotebookId
-    ));
-    return removeNotebook(notebookId);
+    return removeNotebook(notebookId)
+      .then(() => {
+        setSelectedNotebookId((currentSelectedNotebookId) => (
+          currentSelectedNotebookId === notebookId ? null : currentSelectedNotebookId
+        ));
+      });
   }, [removeNotebook]);
 
   const saveNotebookTimeoutRef = useRef<number | null>(null);
@@ -158,7 +160,6 @@ export default function Dashboard({ accessToken }: DashboardProps) {
                 key={selectedNotebook.id}
                 notebook={selectedNotebook}
                 updateNotebook={handleNotebookUpdate}
-                saveNotebook={saveNotebook}
                 runCell={runCell}
               />
             )}

cognee-frontend/src/app/dashboard/InstanceDatasetsAccordion.tsx

@@ -36,7 +36,7 @@ export default function InstanceDatasetsAccordion({ onDatasetsChange }: Instance
     };
 
     checkConnectionToLocalCognee();
-  }, [checkConnectionToCloudCognee, setCloudCogneeConnected, setLocalCogneeConnected]);
+  }, [setCloudCogneeConnected, setLocalCogneeConnected]);
 
   const {
     value: isCloudConnectedModalOpen,

cognee-frontend/src/app/dashboard/NotebooksAccordion.tsx

@@ -56,8 +56,8 @@ export default function NotebooksAccordion({
     notebookLoading();
     removeNotebook(notebookToRemove!.id)
       .finally(notebookLoaded)
-      .finally(closeRemoveNotebookModal);
-      setNotebookToRemove(null);
+      .finally(closeRemoveNotebookModal)
+      .finally(() => setNotebookToRemove(null));
   };
 
   const handleNotebookAdd = useCallback((_: object, formEvent?: FormEvent<HTMLFormElement>) => {

cognee-frontend/src/ui/App/Loading/DefaultLoadingIndicator/LoadingIndicator.module.css

@@ -3,9 +3,9 @@
   width: 1rem;
   height: 1rem;
   border-radius: 50%;
-  border: 0.18rem solid var(--color-indigo-600);;
-  border-top-color: transparent;
-  border-bottom-color: transparent;
+  border: 0.18rem solid var(--color-indigo-600);
+  border-top-color: transparent !important;
+  border-bottom-color: transparent !important;
   animation: spin 2s linear infinite;
 }
 

cognee-frontend/src/ui/App/Loading/DefaultLoadingIndicator/LoadingIndicator.tsx

@@ -1,5 +1,6 @@
-import styles from './LoadingIndicator.module.css';
+import classNames from "classnames";
+import styles from "./LoadingIndicator.module.css";
 
-export default function LoadingIndicator() {
-  return <div className={styles.loadingIndicator} />
+export default function LoadingIndicator({ color = "" }) {
+  return <div className={classNames(styles.loadingIndicator, `!border-${color}`)} />
 }

cognee-frontend/src/ui/elements/Notebook/Notebook.tsx

@@ -16,22 +16,9 @@ interface NotebookProps {
   notebook: NotebookType;
   runCell: (notebook: NotebookType, cell: Cell, cogneeInstance: string) => Promise<void>;
   updateNotebook: (updatedNotebook: NotebookType) => void;
-  saveNotebook: (notebook: NotebookType) => void;
 }
 
-export default function Notebook({ notebook, updateNotebook, saveNotebook, runCell }: NotebookProps) {
-  const saveCells = useCallback(() => {
-    saveNotebook(notebook);
-  }, [notebook, saveNotebook]);
-
-  useEffect(() => {
-    window.addEventListener("beforeunload", saveCells);
-
-    return () => {
-      window.removeEventListener("beforeunload", saveCells);
-    };
-  }, [saveCells]);
-
+export default function Notebook({ notebook, updateNotebook, runCell }: NotebookProps) {
   useEffect(() => {
     if (notebook.cells.length === 0) {
       const newCell: Cell = {
@@ -44,8 +31,9 @@ export default function Notebook({ notebook, updateNotebook, saveNotebook, runCe
        ...notebook,
         cells: [newCell],
       });
+      toggleCellOpen(newCell.id)
     }
-  }, [notebook, saveNotebook, updateNotebook]);
+  }, [notebook, updateNotebook]);
 
   const handleCellRun = useCallback((cell: Cell, cogneeInstance: string) => {
     return runCell(notebook, cell, cogneeInstance);
@@ -289,9 +277,14 @@ function CellResult({ content }: { content: [] }) {
         }
       }
       if (typeof(line) === "object" && line["result"]) {
+        const datasets = Array.from(
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          new Set(Object.values(line["datasets"]).map((dataset: any) => dataset.name))
+        ).join(", ");
+
         parsedContent.push(
           <div className="w-full h-full bg-white">
-            <span className="text-sm pl-2 mb-4">query response (dataset: {line["dataset_name"]})</span>
+            <span className="text-sm pl-2 mb-4">query response (datasets: {datasets})</span>
             <span className="block px-2 py-2">{line["result"]}</span>
           </div>
         );
@@ -303,7 +296,7 @@ function CellResult({ content }: { content: [] }) {
                 data={transformToVisualizationData(line["graphs"]["*"])}
                 ref={graphRef as MutableRefObject<GraphVisualizationAPI>}
                 graphControls={graphControls}
-                className="min-h-48"
+                className="min-h-80"
               />
             </div>
           );

cognee-frontend/src/ui/elements/Notebook/NotebookCellHeader.tsx

@@ -1,9 +1,9 @@
 import { useState } from "react";
 import classNames from "classnames";
 
-import { useBoolean } from "@/utils";
+import { isCloudEnvironment, useBoolean } from "@/utils";
 import { PlayIcon } from "@/ui/Icons";
-import { PopupMenu, IconButton, Select } from "@/ui/elements";
+import { PopupMenu, IconButton } from "@/ui/elements";
 import { LoadingIndicator } from "@/ui/App";
 
 import { Cell } from "./types";
@@ -33,7 +33,7 @@ export default function NotebookCellHeader({
     setFalse: setIsNotRunningCell,
   } = useBoolean(false);
 
-  const [runInstance, setRunInstance] = useState<string>("local");
+  const [runInstance, setRunInstance] = useState<string>(isCloudEnvironment() ? "cloud" : "local");
 
   const handleCellRun = () => {
     setIsRunningCell();
@@ -50,14 +50,23 @@ export default function NotebookCellHeader({
         <span className="ml-4">{cell.name}</span>
       </div>
       <div className="pr-4 flex flex-row items-center gap-8">
-        <Select name="cogneeInstance" onChange={(event) => setRunInstance(event.currentTarget.value)} className="!bg-transparent outline-none cursor-pointer !hover:bg-gray-50">
+          {isCloudEnvironment() ? (
+            <div>
+              cloud cognee
+            </div>
+          ) : (
+            <div>
+              local cognee
+            </div>
+          )}
+        {/* <Select name="cogneeInstance" onChange={(event) => setRunInstance(event.currentTarget.value)} className="!bg-transparent outline-none cursor-pointer !hover:bg-gray-50">
           <option value="local" className="flex flex-row items-center gap-2">
             local cognee
           </option>
           <option value="cloud" className="flex flex-row items-center gap-2">
             cloud cognee
           </option>
-        </Select>
+        </Select> */}
         <PopupMenu>
           <div className="flex flex-col gap-0.5">
             <button onClick={() => moveCellUp(cell)} className="hover:bg-gray-100 w-full text-left px-2 cursor-pointer">move cell up</button>

cognee/api/v1/notebooks/routers/get_notebooks_router.py

@@ -31,7 +31,8 @@ def get_notebooks_router():
 
     @router.get("")
     async def get_notebooks_endpoint(user: User = Depends(get_authenticated_user)):
-        return await get_notebooks(user.id)
+        async with get_async_session() as session:
+            return await get_notebooks(user.id, session)
 
     @router.post("")
     async def create_notebook_endpoint(

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -58,8 +58,8 @@ def __init__(self, connection_string: str):
         else:
             self.engine = create_async_engine(
                 connection_string,
-                pool_size=5,
-                max_overflow=10,
+                pool_size=20,
+                max_overflow=20,
                 pool_recycle=280,
                 pool_pre_ping=True,
                 pool_timeout=280,

cognee/infrastructure/databases/vector/config.py

@@ -39,7 +39,7 @@ def validate_paths(cls, values):
             values.vector_db_url = ensure_absolute_path(
                 values.vector_db_url,
             )
-        else:
+        elif not values.vector_db_url:
             # Default path
             databases_directory_path = os.path.join(base_config.system_root_directory, "databases")
             values.vector_db_url = os.path.join(databases_directory_path, "cognee.lancedb")

cognee/infrastructure/utils/run_async.py

@@ -1,13 +1,18 @@
 import asyncio
 from functools import partial
+import inspect
 
 
 async def run_async(func, *args, loop=None, executor=None, **kwargs):
     if loop is None:
         try:
-            running_loop = asyncio.get_running_loop()
+            loop = asyncio.get_running_loop()
         except RuntimeError:
-            running_loop = asyncio.get_event_loop()
+            loop = asyncio.get_event_loop()
 
-    pfunc = partial(func, *args, **kwargs)
-    return await running_loop.run_in_executor(executor, pfunc)
+    if "loop" in inspect.signature(func).parameters:
+        pfunc = partial(func, *args, loop=loop, **kwargs)
+    else:
+        pfunc = partial(func, *args, **kwargs)
+
+    return await loop.run_in_executor(executor, pfunc)

cognee/infrastructure/utils/run_sync.py

@@ -2,16 +2,17 @@
 import threading
 
 
-def run_sync(coro, timeout=None):
+def run_sync(coro, running_loop=None, timeout=None):
     result = None
     exception = None
 
     def runner():
-        nonlocal result, exception
+        nonlocal result, exception, running_loop
 
         try:
             try:
-                running_loop = asyncio.get_running_loop()
+                if not running_loop:
+                    running_loop = asyncio.get_running_loop()
 
                 result = asyncio.run_coroutine_threadsafe(coro, running_loop).result(timeout)
             except RuntimeError:

cognee/modules/notebooks/methods/get_notebook.py

@@ -1,6 +1,6 @@
 from uuid import UUID
 from typing import Optional
-from sqlalchemy import select
+from sqlalchemy import and_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from cognee.infrastructure.databases.relational import with_async_session
@@ -15,7 +15,7 @@ async def get_notebook(
     session: AsyncSession,
 ) -> Optional[Notebook]:
     result = await session.execute(
-        select(Notebook).where(Notebook.owner_id == user_id and Notebook.id == notebook_id)
+        select(Notebook).where(and_(Notebook.owner_id == user_id, Notebook.id == notebook_id))
     )
 
     return result.scalar()

cognee/modules/notebooks/methods/update_notebook.py

@@ -1,4 +1,3 @@
-from typing import Callable, AsyncContextManager
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from cognee.infrastructure.databases.relational import with_async_session

cognee/modules/notebooks/operations/run_in_local_sandbox.py

@@ -5,16 +5,18 @@
 
 def wrap_in_async_handler(user_code: str) -> str:
     return (
-        "from cognee.infrastructure.utils.run_sync import run_sync\n\n"
-        "async def __user_main__():\n"
+        "import asyncio\n"
+        + "asyncio.set_event_loop(running_loop)\n\n"
+        + "from cognee.infrastructure.utils.run_sync import run_sync\n\n"
+        + "async def __user_main__():\n"
         + "\n".join("    " + line for line in user_code.strip().split("\n"))
         + "\n"
-        "    globals().update(locals())\n\n"
-        "run_sync(__user_main__())\n"
+        + "    globals().update(locals())\n\n"
+        + "run_sync(__user_main__(), running_loop)\n"
     )
 
 
-def run_in_local_sandbox(code, environment=None):
+def run_in_local_sandbox(code, environment=None, loop=None):
     environment = environment or {}
     code = wrap_in_async_handler(code.replace("\xa0", "\n"))
 
@@ -31,6 +33,7 @@ def customPrintFunction(output):
         printOutput.append(output)
 
     environment["print"] = customPrintFunction
+    environment["running_loop"] = loop
 
     try:
         exec(code, environment)

cognee/modules/retrieval/graph_completion_context_extension_retriever.py

@@ -48,7 +48,7 @@ async def get_completion(
         query: str,
         context: Optional[List[Edge]] = None,
         context_extension_rounds=4,
-    ) -> str:
+    ) -> List[str]:
         """
         Extends the context for a given query by retrieving related triplets and generating new
         completions based on them.

cognee/modules/retrieval/graph_completion_cot_retriever.py

@@ -58,7 +58,7 @@ async def get_completion(
         query: str,
         context: Optional[List[Edge]] = None,
         max_iter=4,
-    ) -> str:
+    ) -> List[str]:
         """
         Generate completion responses based on a user query and contextual information.
 

cognee/modules/retrieval/graph_completion_retriever.py

@@ -135,7 +135,7 @@ async def get_completion(
         self,
         query: str,
         context: Optional[List[Edge]] = None,
-    ) -> Any:
+    ) -> List[str]:
         """
         Generates a completion using graph connections context based on a query.
 

cognee/modules/retrieval/temporal_retriever.py

@@ -136,7 +136,7 @@ async def get_context(self, query: str) -> Any:
 
         return self.descriptions_to_string(top_k_events)
 
-    async def get_completion(self, query: str, context: Optional[str] = None) -> str:
+    async def get_completion(self, query: str, context: Optional[str] = None) -> List[str]:
         """Generates a response using the query and optional context."""
         if not context:
             context = await self.get_context(query=query)