From 74a3220e9fa18a13fd1b28f55b2079059d2ec266 Mon Sep 17 00:00:00 2001
From: Boris <boris@topoteretes.com>
Date: Thu, 11 Sep 2025 16:16:03 +0200
Subject: [PATCH] fix: graph view with search (#1368)

<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
---
 .../45957f0a9849_add_notebook_table.py        |  46 ++++++++
 .../src/app/dashboard/AddDataToCognee.tsx     |   3 +-
 .../src/app/dashboard/Dashboard.tsx           |  24 +++-
 .../dashboard/InstanceDatasetsAccordion.tsx   |  29 +++--
 cognee-frontend/src/app/dashboard/page.tsx    |  12 +-
 cognee-frontend/src/app/page.tsx              |   4 +-
 cognee-frontend/src/modules/auth/getUser.ts   |   6 +
 cognee-frontend/src/modules/auth/index.ts     |   2 +-
 cognee-frontend/src/modules/auth/types.ts     |   2 +-
 .../src/modules/auth/useAuthenticatedUser.ts  |   2 +-
 .../src/modules/cloud/checkCloudConnection.ts |   5 +-
 .../src/modules/notebooks/useNotebooks.ts     |  67 +++++------
 cognee-frontend/src/ui/Layout/Header.tsx      |  15 ++-
 .../src/ui/elements/Notebook/Notebook.tsx     | 109 ++++++++++--------
 .../elements/Notebook/NotebookCellHeader.tsx  |  24 ++--
 cognee-frontend/src/utils/fetch.ts            |  13 ++-
 cognee-frontend/src/utils/index.ts            |   1 +
 .../src/utils/isCloudEnvironment.ts           |   4 +
 .../notebooks/routers/get_notebooks_router.py |   3 +-
 cognee/infrastructure/utils/run_sync.py       |   9 +-
 .../operations/run_in_local_sandbox.py        |  14 ---
 .../utils/brute_force_triplet_search.py       |   3 +-
 cognee/modules/search/methods/search.py       |   8 +-
 .../search/utils/prepare_search_result.py     |   4 +-
 24 files changed, 268 insertions(+), 141 deletions(-)
 create mode 100644 alembic/versions/45957f0a9849_add_notebook_table.py
 create mode 100644 cognee-frontend/src/modules/auth/getUser.ts
 create mode 100644 cognee-frontend/src/utils/isCloudEnvironment.ts

diff --git a/alembic/versions/45957f0a9849_add_notebook_table.py b/alembic/versions/45957f0a9849_add_notebook_table.py
new file mode 100644
index 000000000..ded7ccd9c
--- /dev/null
+++ b/alembic/versions/45957f0a9849_add_notebook_table.py
@@ -0,0 +1,46 @@
+"""Add notebook table
+
+Revision ID: 45957f0a9849
+Revises: 9e7a3cb85175
+Create Date: 2025-09-10 17:47:58.201319
+
+"""
+
+from datetime import datetime, timezone
+from uuid import uuid4
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "45957f0a9849"
+down_revision: Union[str, None] = "9e7a3cb85175"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+
+    if "notebooks" not in inspector.get_table_names():
+        # Define table with all necessary columns including primary key
+        op.create_table(
+            "notebooks",
+            sa.Column("id", sa.UUID, primary_key=True, default=uuid4),  # Critical for SQLite
+            sa.Column("owner_id", sa.UUID, index=True),
+            sa.Column("name", sa.String(), nullable=False),
+            sa.Column("cells", sa.JSON(), nullable=False),
+            sa.Column("deletable", sa.Boolean(), default=True),
+            sa.Column("created_at", sa.DateTime(), default=lambda: datetime.now(timezone.utc)),
+        )
+
+
+def downgrade() -> None:
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+
+    if "notebooks" in inspector.get_table_names():
+        op.drop_table("notebooks")
diff --git a/cognee-frontend/src/app/dashboard/AddDataToCognee.tsx b/cognee-frontend/src/app/dashboard/AddDataToCognee.tsx
index a74547b5e..4f9a3179e 100644
--- a/cognee-frontend/src/app/dashboard/AddDataToCognee.tsx
+++ b/cognee-frontend/src/app/dashboard/AddDataToCognee.tsx
@@ -39,7 +39,8 @@ export default function AddDataToCognee({ datasets, refreshDatasets, useCloud =
       } : {
         name: "main_dataset",
       },
-      Array.from(filesForUpload)
+      Array.from(filesForUpload),
+      useCloud
     )
       .then(({ dataset_id, dataset_name }) => {
         refreshDatasets();
diff --git a/cognee-frontend/src/app/dashboard/Dashboard.tsx b/cognee-frontend/src/app/dashboard/Dashboard.tsx
index a785f8e00..97576f3b4 100644
--- a/cognee-frontend/src/app/dashboard/Dashboard.tsx
+++ b/cognee-frontend/src/app/dashboard/Dashboard.tsx
@@ -5,16 +5,31 @@ import { useCallback, useEffect, useRef, useState } from "react";
 import { Header } from "@/ui/Layout";
 import { SearchIcon } from "@/ui/Icons";
 import { Notebook } from "@/ui/elements";
+import { fetch, isCloudEnvironment } from "@/utils";
 import { Notebook as NotebookType } from "@/ui/elements/Notebook/types";
+import { useAuthenticatedUser } from "@/modules/auth";
 import { Dataset } from "@/modules/ingestion/useDatasets";
 import useNotebooks from "@/modules/notebooks/useNotebooks";
 
+import AddDataToCognee from "./AddDataToCognee";
 import NotebooksAccordion from "./NotebooksAccordion";
 import CogneeInstancesAccordion from "./CogneeInstancesAccordion";
-import AddDataToCognee from "./AddDataToCognee";
 import InstanceDatasetsAccordion from "./InstanceDatasetsAccordion";
 
-export default function Dashboard() {
+interface DashboardProps {
+  user?: {
+    id: string;
+    name: string;
+    email: string;
+    picture: string;
+  };
+  accessToken: string;
+}
+
+export default function Dashboard({ accessToken }: DashboardProps) {
+  fetch.setAccessToken(accessToken);
+  const { user } = useAuthenticatedUser();
+
   const {
     notebooks,
     refreshNotebooks,
@@ -91,6 +106,8 @@ export default function Dashboard() {
     setDatasets(datasets);
   }, []);
 
+  const isCloudEnv = isCloudEnvironment();
+
   return (
     <div className="h-full flex flex-col bg-gray-200">
       <video
@@ -104,7 +121,7 @@ export default function Dashboard() {
         Your browser does not support the video tag.
       </video>
 
-      <Header />
+      <Header user={user} />
 
       <div className="relative flex-1 flex flex-row gap-2.5 items-start w-full max-w-[1920px] max-h-[calc(100% - 3.5rem)] overflow-hidden mx-auto px-2.5 py-2.5">
         <div className="px-5 py-4 lg:w-96 bg-white rounded-xl min-h-full">
@@ -116,6 +133,7 @@ export default function Dashboard() {
           <AddDataToCognee
             datasets={datasets}
             refreshDatasets={refreshDatasetsRef.current}
+            useCloud={isCloudEnv}
           />
 
           <NotebooksAccordion
diff --git a/cognee-frontend/src/app/dashboard/InstanceDatasetsAccordion.tsx b/cognee-frontend/src/app/dashboard/InstanceDatasetsAccordion.tsx
index dad9f48e4..cfb0db2e8 100644
--- a/cognee-frontend/src/app/dashboard/InstanceDatasetsAccordion.tsx
+++ b/cognee-frontend/src/app/dashboard/InstanceDatasetsAccordion.tsx
@@ -1,6 +1,7 @@
+import classNames from "classnames";
 import { useCallback, useEffect } from "react";
 
-import { fetch, useBoolean } from "@/utils";
+import { fetch, isCloudEnvironment, useBoolean } from "@/utils";
 import { checkCloudConnection } from "@/modules/cloud";
 import { CaretIcon, CloseIcon, CloudIcon, LocalCogneeIcon } from "@/ui/Icons";
 import { CTAButton, GhostButton, IconButton, Input, Modal } from "@/ui/elements";
@@ -18,11 +19,13 @@ export default function InstanceDatasetsAccordion({ onDatasetsChange }: Instance
   const {
     value: isCloudCogneeConnected,
     setTrue: setCloudCogneeConnected,
-  } = useBoolean(false);
+  } = useBoolean(isCloudEnvironment());
 
-  const checkConnectionToCloudCognee = useCallback((apiKey: string) => {
-      fetch.setApiKey(apiKey);
-      return checkCloudConnection(apiKey)
+  const checkConnectionToCloudCognee = useCallback((apiKey?: string) => {
+      if (apiKey) {
+        fetch.setApiKey(apiKey);
+      }
+      return checkCloudConnection()
         .then(setCloudCogneeConnected)
     }, [setCloudCogneeConnected]);
 
@@ -33,8 +36,6 @@ export default function InstanceDatasetsAccordion({ onDatasetsChange }: Instance
     };
 
     checkConnectionToLocalCognee();
-
-    checkConnectionToCloudCognee("");
   }, [checkConnectionToCloudCognee, setCloudCogneeConnected, setLocalCogneeConnected]);
 
   const {
@@ -54,8 +55,12 @@ export default function InstanceDatasetsAccordion({ onDatasetsChange }: Instance
       });
   };
 
+  const isCloudEnv = isCloudEnvironment();
+
   return (
-    <>
+    <div className={classNames("flex flex-col", {
+      "flex-col-reverse": isCloudEnv,
+    })}>
       <DatasetsAccordion
         title={(
           <div className="flex flex-row items-center justify-between">
@@ -69,7 +74,7 @@ export default function InstanceDatasetsAccordion({ onDatasetsChange }: Instance
         switchCaretPosition={true}
         className="pt-3 pb-1.5"
         contentClassName="pl-4"
-        onDatasetsChange={onDatasetsChange}
+        onDatasetsChange={!isCloudEnv ? onDatasetsChange : () => {}}
       />
 
       {isCloudCogneeConnected ? (
@@ -86,11 +91,11 @@ export default function InstanceDatasetsAccordion({ onDatasetsChange }: Instance
           switchCaretPosition={true}
           className="pt-3 pb-1.5"
           contentClassName="pl-4"
-          onDatasetsChange={onDatasetsChange}
+          onDatasetsChange={isCloudEnv ? onDatasetsChange : () => {}}
           useCloud={true}
         />
       ) : (
-        <button className="w-full flex flex-row items-center justify-between py-1.5 cursor-pointer" onClick={!isCloudCogneeConnected ? openCloudConnectionModal : () => {}}>
+        <button className="w-full flex flex-row items-center justify-between py-1.5 cursor-pointer pt-3" onClick={!isCloudCogneeConnected ? openCloudConnectionModal : () => {}}>
           <div className="flex flex-row items-center gap-1.5">
             <CaretIcon className="rotate-[-90deg]" />
             <div className="flex flex-row items-center gap-2">
@@ -120,6 +125,6 @@ export default function InstanceDatasetsAccordion({ onDatasetsChange }: Instance
           </form>
         </div>
       </Modal>
-    </>
+    </div>
   );
 }
diff --git a/cognee-frontend/src/app/dashboard/page.tsx b/cognee-frontend/src/app/dashboard/page.tsx
index 2ab67cdd6..7ff3abaf4 100644
--- a/cognee-frontend/src/app/dashboard/page.tsx
+++ b/cognee-frontend/src/app/dashboard/page.tsx
@@ -1 +1,11 @@
-export { default } from "./Dashboard";
+"use server";
+
+import Dashboard from "./Dashboard";
+
+export default async function Page() {
+  const accessToken = "";
+
+  return (
+    <Dashboard accessToken={accessToken} />
+  );
+}
diff --git a/cognee-frontend/src/app/page.tsx b/cognee-frontend/src/app/page.tsx
index 523a90f6b..80db75b20 100644
--- a/cognee-frontend/src/app/page.tsx
+++ b/cognee-frontend/src/app/page.tsx
@@ -1,3 +1,3 @@
-export { default } from "./(graph)/GraphView";
+export { default } from "./dashboard/page";
 
-export const dynamic = "force-dynamic";
+// export const dynamic = "force-dynamic";
diff --git a/cognee-frontend/src/modules/auth/getUser.ts b/cognee-frontend/src/modules/auth/getUser.ts
new file mode 100644
index 000000000..c77821bd0
--- /dev/null
+++ b/cognee-frontend/src/modules/auth/getUser.ts
@@ -0,0 +1,6 @@
+import fetch from "@/utils/fetch";
+
+export default function getUser() {
+  return fetch("/v1/auth/me")
+      .then((response) => response.json());
+}
diff --git a/cognee-frontend/src/modules/auth/index.ts b/cognee-frontend/src/modules/auth/index.ts
index ea21cce2b..32ea0d8ad 100644
--- a/cognee-frontend/src/modules/auth/index.ts
+++ b/cognee-frontend/src/modules/auth/index.ts
@@ -1,2 +1,2 @@
-export { default as useAuthenticatedUser } from "./useAuthenticatedUser";
 export { type User } from "./types";
+export { default as useAuthenticatedUser } from "./useAuthenticatedUser";
diff --git a/cognee-frontend/src/modules/auth/types.ts b/cognee-frontend/src/modules/auth/types.ts
index 3441b4149..76799c36f 100644
--- a/cognee-frontend/src/modules/auth/types.ts
+++ b/cognee-frontend/src/modules/auth/types.ts
@@ -2,5 +2,5 @@ export interface User {
   id: string;
   name: string;
   email: string;
-  avatarImagePath: string;
+  picture: string;
 }
diff --git a/cognee-frontend/src/modules/auth/useAuthenticatedUser.ts b/cognee-frontend/src/modules/auth/useAuthenticatedUser.ts
index f789f3de4..c1fdd9810 100644
--- a/cognee-frontend/src/modules/auth/useAuthenticatedUser.ts
+++ b/cognee-frontend/src/modules/auth/useAuthenticatedUser.ts
@@ -3,7 +3,7 @@ import { fetch } from "@/utils";
 import { User } from "./types";
 
 export default function useAuthenticatedUser() {
-  const [user, setUser] = useState<User | null>(null);
+  const [user, setUser] = useState<User>();
 
   useEffect(() => {
     if (!user) {
diff --git a/cognee-frontend/src/modules/cloud/checkCloudConnection.ts b/cognee-frontend/src/modules/cloud/checkCloudConnection.ts
index dfc40767d..a76ec09a2 100644
--- a/cognee-frontend/src/modules/cloud/checkCloudConnection.ts
+++ b/cognee-frontend/src/modules/cloud/checkCloudConnection.ts
@@ -1,10 +1,7 @@
 import { fetch } from "@/utils";
 
-export default function checkCloudConnection(apiKey: string) {
+export default function checkCloudConnection() {
   return fetch("/v1/checks/connection", {
     method: "POST",
-    headers: {
-      "X-Api-Key": apiKey,
-    },
   });
 }
diff --git a/cognee-frontend/src/modules/notebooks/useNotebooks.ts b/cognee-frontend/src/modules/notebooks/useNotebooks.ts
index f36f97448..e427f85ee 100644
--- a/cognee-frontend/src/modules/notebooks/useNotebooks.ts
+++ b/cognee-frontend/src/modules/notebooks/useNotebooks.ts
@@ -1,5 +1,5 @@
 import { useCallback, useState } from "react";
-import { fetch } from "@/utils";
+import { fetch, isCloudEnvironment } from "@/utils";
 import { Cell, Notebook } from "@/ui/elements/Notebook/types";
 
 function useNotebooks() {
@@ -12,7 +12,7 @@ function useNotebooks() {
         headers: {
           "Content-Type": "application/json",
         },
-      })
+      }, isCloudEnvironment())
       .then((response) => response.json())
       .then((notebook) => {
         setNotebooks((notebooks) => [
@@ -26,30 +26,31 @@ function useNotebooks() {
 
   const removeNotebook = useCallback((notebookId: string) => {
     return fetch(`/v1/notebooks/${notebookId}`, {
-        method: "DELETE",
-      })
-      .then(() => {
-        setNotebooks((notebooks) =>
-          notebooks.filter((notebook) => notebook.id !== notebookId)
-        );
-      });
+      method: "DELETE",
+    }, isCloudEnvironment())
+    .then(() => {
+      setNotebooks((notebooks) =>
+        notebooks.filter((notebook) => notebook.id !== notebookId)
+      );
+    });
   }, []);
 
   const fetchNotebooks = useCallback(() => {
     return fetch("/v1/notebooks", {
-        headers: {
-          "Content-Type": "application/json",
-        },
-      })
-      .then((response) => response.json())
-      .then((notebooks) => {
-        setNotebooks(notebooks);
+      headers: {
+        "Content-Type": "application/json",
+      },
+    }, isCloudEnvironment())
+    .then((response) => response.json())
+    .then((notebooks) => {
+      setNotebooks(notebooks);
 
-        return notebooks;
-      })
-      .catch((error) => {
-        console.error("Error fetching notebooks:", error);
-      });
+      return notebooks;
+    })
+    .catch((error) => {
+      console.error("Error fetching notebooks:", error);
+      throw error
+    });
   }, []);
 
   const updateNotebook = useCallback((updatedNotebook: Notebook) => {
@@ -64,19 +65,19 @@ function useNotebooks() {
 
   const saveNotebook = useCallback((notebook: Notebook) => {
     return fetch(`/v1/notebooks/${notebook.id}`, {
-        body: JSON.stringify({
-          name: notebook.name,
-          cells: notebook.cells,
-        }),
-        method: "PUT",
-        headers: {
-          "Content-Type": "application/json",
-        },
-      })
-      .then((response) => response.json())
+      body: JSON.stringify({
+        name: notebook.name,
+        cells: notebook.cells,
+      }),
+      method: "PUT",
+      headers: {
+        "Content-Type": "application/json",
+      },
+    }, isCloudEnvironment())
+    .then((response) => response.json())
   }, []);
 
-  const runCell = useCallback((notebook: Notebook, cell: Cell) => {
+  const runCell = useCallback((notebook: Notebook, cell: Cell, cogneeInstance: string) => {
     setNotebooks((existingNotebooks) =>
       existingNotebooks.map((existingNotebook) =>
         existingNotebook.id === notebook.id ? {
@@ -100,7 +101,7 @@ function useNotebooks() {
         headers: {
           "Content-Type": "application/json",
         },
-      })
+      }, cogneeInstance === "cloud")
       .then((response) => response.json())
       .then((response) => {
         setNotebooks((existingNotebooks) =>
diff --git a/cognee-frontend/src/ui/Layout/Header.tsx b/cognee-frontend/src/ui/Layout/Header.tsx
index a615eeb4f..2f26433e9 100644
--- a/cognee-frontend/src/ui/Layout/Header.tsx
+++ b/cognee-frontend/src/ui/Layout/Header.tsx
@@ -6,12 +6,17 @@ import { useBoolean } from "@/utils";
 
 import { CloseIcon, CloudIcon, CogneeIcon } from "../Icons";
 import { CTAButton, GhostButton, IconButton, Modal } from "../elements";
-import { useAuthenticatedUser } from "@/modules/auth";
 import syncData from "@/modules/cloud/syncData";
 
-export default function Header() {
-  const { user } = useAuthenticatedUser();
+interface HeaderProps {
+  user?: {
+    name: string;
+    email: string;
+    picture: string;
+  };
+}
 
+export default function Header({ user }: HeaderProps) {
   const {
     value: isSyncModalOpen,
     setTrue: openSyncModal,
@@ -45,8 +50,8 @@ export default function Header() {
             <SettingsIcon />
           </div> */}
           <Link href="/account" className="bg-indigo-600 w-8 h-8 rounded-full overflow-hidden">
-            {user?.avatarImagePath ? (
-              <Image width="32" height="32" alt="Name of the user" src={user.avatarImagePath} />
+            {user?.picture ? (
+              <Image width="32" height="32" alt="Name of the user" src={user.picture} />
             ) : (
               <div className="w-8 h-8 rounded-full text-white flex items-center justify-center">
                 {user?.email?.charAt(0) || "C"}
diff --git a/cognee-frontend/src/ui/elements/Notebook/Notebook.tsx b/cognee-frontend/src/ui/elements/Notebook/Notebook.tsx
index 0e037890a..0789a8695 100644
--- a/cognee-frontend/src/ui/elements/Notebook/Notebook.tsx
+++ b/cognee-frontend/src/ui/elements/Notebook/Notebook.tsx
@@ -14,7 +14,7 @@ import { Cell, Notebook as NotebookType } from "./types";
 
 interface NotebookProps {
   notebook: NotebookType;
-  runCell: (notebook: NotebookType, cell: Cell) => Promise<void>;
+  runCell: (notebook: NotebookType, cell: Cell, cogneeInstance: string) => Promise<void>;
   updateNotebook: (updatedNotebook: NotebookType) => void;
   saveNotebook: (notebook: NotebookType) => void;
 }
@@ -47,8 +47,8 @@ export default function Notebook({ notebook, updateNotebook, saveNotebook, runCe
     }
   }, [notebook, saveNotebook, updateNotebook]);
 
-  const handleCellRun = useCallback((cell: Cell) => {
-    return runCell(notebook, cell);
+  const handleCellRun = useCallback((cell: Cell, cogneeInstance: string) => {
+    return runCell(notebook, cell, cogneeInstance);
   }, [notebook, runCell]);
 
   const handleCellAdd = useCallback((afterCellIndex: number, cellType: "markdown" | "code") => {
@@ -244,7 +244,7 @@ export default function Notebook({ notebook, updateNotebook, saveNotebook, runCe
 }
 
 
-function CellResult({ content = [] }) {
+function CellResult({ content }: { content: [] }) {
   const parsedContent = [];
 
   const graphRef = useRef<GraphVisualizationAPI>();
@@ -256,6 +256,7 @@ function CellResult({ content = [] }) {
   for (const line of content) {
     try {
       if (Array.isArray(line)) {
+        // @ts-expect-error line can be Array or string
         for (const item of line) {
           if (typeof item === "string") {
             parsedContent.push(
@@ -264,15 +265,13 @@ function CellResult({ content = [] }) {
               </pre>
             );
           }
-          if (typeof item === "object" && item["search_result"] && Array.isArray(item["search_result"])) {
-            for (const result of item["search_result"]) {
-              parsedContent.push(
-                <div className="w-full h-full bg-white">
-                  <span className="text-sm pl-2 mb-4">query response (dataset: {item["dataset_name"]})</span>
-                  <span className="block px-2 py-2">{result}</span>
-                </div>
-              );
-            }
+          if (typeof item === "object" && item["search_result"]) {
+            parsedContent.push(
+              <div className="w-full h-full bg-white">
+                <span className="text-sm pl-2 mb-4">query response (dataset: {item["dataset_name"]})</span>
+                <span className="block px-2 py-2">{item["search_result"]}</span>
+              </div>
+            );
           }
           if (typeof item === "object" && item["graph"] && typeof item["graph"] === "object") {
             parsedContent.push(
@@ -289,6 +288,27 @@ function CellResult({ content = [] }) {
           }
         }
       }
+      if (typeof(line) === "object" && line["result"]) {
+        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="block px-2 py-2">{line["result"]}</span>
+          </div>
+        );
+        if (line["graphs"]) {
+          parsedContent.push(
+            <div className="w-full h-full bg-white">
+              <span className="text-sm pl-2 mb-4">reasoning graph</span>
+              <GraphVisualization
+                data={transformToVisualizationData(line["graphs"]["*"])}
+                ref={graphRef as MutableRefObject<GraphVisualizationAPI>}
+                graphControls={graphControls}
+                className="min-h-48"
+              />
+            </div>
+          );
+        }
+      }
     } catch (error) {
       console.error(error);
       parsedContent.push(line);
@@ -298,45 +318,44 @@ function CellResult({ content = [] }) {
   return parsedContent.map((item, index) => (
     <div key={index} className="px-2 py-1">
       {item}
-      {/* {typeof item === "object" && item["search_result"] && Array.isArray(item["search_result"]) && (
-        (item["search_result"] as []).map((result: string) => (<pre key={result.slice(0, -10)}>{result}</pre>))
-      )}
-      {typeof item === "object" && item["graph"] && typeof item["graph"] === "object" && (
-        (item["graph"])
-      )} */}
     </div>
   ));
 
 };
 
-function transformToVisualizationData(triplets) {
+function transformToVisualizationData(graph: { nodes: [], edges: [] }) {
   // Implementation to transform triplet to visualization data
 
-  const nodes = {};
-  const links = {};
-
-  for (const triplet of triplets) {
-    nodes[triplet.source.id] = {
-      id: triplet.source.id,
-      label: triplet.source.attributes.name,
-      type: triplet.source.attributes.type,
-      attributes: triplet.source.attributes,
-    };
-    nodes[triplet.destination.id] = {
-      id: triplet.destination.id,
-      label: triplet.destination.attributes.name,
-      type: triplet.destination.attributes.type,
-      attributes: triplet.destination.attributes,
-    };
-    links[`${triplet.source.id}_${triplet.attributes.relationship_name}_${triplet.destination.id}`] = {
-      source: triplet.source.id,
-      target: triplet.destination.id,
-      label: triplet.attributes.relationship_name,
-    }
-  }
-
   return {
-    nodes: Object.values(nodes),
-    links: Object.values(links),
+    nodes: graph.nodes,
+    links: graph.edges,
   };
+
+  // const nodes = {};
+  // const links = {};
+
+  // for (const triplet of triplets) {
+  //   nodes[triplet.source.id] = {
+  //     id: triplet.source.id,
+  //     label: triplet.source.attributes.name,
+  //     type: triplet.source.attributes.type,
+  //     attributes: triplet.source.attributes,
+  //   };
+  //   nodes[triplet.destination.id] = {
+  //     id: triplet.destination.id,
+  //     label: triplet.destination.attributes.name,
+  //     type: triplet.destination.attributes.type,
+  //     attributes: triplet.destination.attributes,
+  //   };
+  //   links[`${triplet.source.id}_${triplet.attributes.relationship_name}_${triplet.destination.id}`] = {
+  //     source: triplet.source.id,
+  //     target: triplet.destination.id,
+  //     label: triplet.attributes.relationship_name,
+  //   }
+  // }
+
+  // return {
+  //   nodes: Object.values(nodes),
+  //   links: Object.values(links),
+  // };
 }
diff --git a/cognee-frontend/src/ui/elements/Notebook/NotebookCellHeader.tsx b/cognee-frontend/src/ui/elements/Notebook/NotebookCellHeader.tsx
index e5a41d9c2..b1fbee5b0 100644
--- a/cognee-frontend/src/ui/elements/Notebook/NotebookCellHeader.tsx
+++ b/cognee-frontend/src/ui/elements/Notebook/NotebookCellHeader.tsx
@@ -1,16 +1,16 @@
+import { useState } from "react";
 import classNames from "classnames";
 
 import { useBoolean } from "@/utils";
-import { LocalCogneeIcon, PlayIcon } from "@/ui/Icons";
-import { PopupMenu } from "@/ui/elements";
+import { PlayIcon } from "@/ui/Icons";
+import { PopupMenu, IconButton, Select } from "@/ui/elements";
 import { LoadingIndicator } from "@/ui/App";
 
 import { Cell } from "./types";
-import IconButton from "../IconButton";
 
 interface NotebookCellHeaderProps {
   cell: Cell;
-  runCell: (cell: Cell) => Promise<void>;
+  runCell: (cell: Cell, cogneeInstance: string) => Promise<void>;
   renameCell: (cell: Cell) => void;
   removeCell: (cell: Cell) => void;
   moveCellUp: (cell: Cell) => void;
@@ -33,9 +33,11 @@ export default function NotebookCellHeader({
     setFalse: setIsNotRunningCell,
   } = useBoolean(false);
 
+  const [runInstance, setRunInstance] = useState<string>("local");
+
   const handleCellRun = () => {
     setIsRunningCell();
-    runCell(cell)
+    runCell(cell, runInstance)
       .then(() => {
         setIsNotRunningCell();
       });
@@ -48,10 +50,14 @@ export default function NotebookCellHeader({
         <span className="ml-4">{cell.name}</span>
       </div>
       <div className="pr-4 flex flex-row items-center gap-8">
-        <div className="flex flex-row items-center gap-2">
-          <LocalCogneeIcon className="text-indigo-700" />
-          <span className="text-xs">local cognee</span>
-        </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>
         <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>
diff --git a/cognee-frontend/src/utils/fetch.ts b/cognee-frontend/src/utils/fetch.ts
index 029c92d2e..9d3d24ebd 100644
--- a/cognee-frontend/src/utils/fetch.ts
+++ b/cognee-frontend/src/utils/fetch.ts
@@ -1,4 +1,5 @@
 import handleServerErrors from "./handleServerErrors";
+import isCloudEnvironment from "./isCloudEnvironment";
 
 let numberOfRetries = 0;
 
@@ -6,9 +7,10 @@ const isAuth0Enabled = process.env.USE_AUTH0_AUTHORIZATION?.toLowerCase() === "t
 
 const backendApiUrl = process.env.NEXT_PUBLIC_BACKEND_API_URL || "http://localhost:8000/api";
 
-const cloudApiUrl = process.env.NEXT_PUBLIC_CLOUD_API_URL || "https://api.cognee.ai/api";
+const cloudApiUrl = process.env.NEXT_PUBLIC_CLOUD_API_URL || "http://localhost:8001/api";
 
 let apiKey: string | null = null;
+let accessToken: string | null = null;
 
 export default async function fetch(url: string, options: RequestInit = {}, useCloud = false): Promise<Response> {
   function retry(lastError: Response) {
@@ -34,7 +36,10 @@ export default async function fetch(url: string, options: RequestInit = {}, useC
       ...options,
       headers: {
         ...options.headers,
-        ...(useCloud ? {"X-Api-Key": apiKey!} : {}),
+        ...(useCloud && !isCloudEnvironment()
+          ? {"X-Api-Key": apiKey!}
+          : {"Authorization": `Bearer ${accessToken}`}
+        ),
       },
       credentials: "include",
     },
@@ -66,3 +71,7 @@ fetch.checkHealth = () => {
 fetch.setApiKey = (newApiKey: string) => {
   apiKey = newApiKey;
 };
+
+fetch.setAccessToken = (newAccessToken: string) => {
+  accessToken = newAccessToken;
+};
diff --git a/cognee-frontend/src/utils/index.ts b/cognee-frontend/src/utils/index.ts
index 902331cfe..702a34e76 100644
--- a/cognee-frontend/src/utils/index.ts
+++ b/cognee-frontend/src/utils/index.ts
@@ -2,3 +2,4 @@ export { default as fetch } from "./fetch";
 export { default as handleServerErrors } from "./handleServerErrors";
 export { default as useBoolean } from "./useBoolean";
 export { default as useOutsideClick } from "./useOutsideClick";
+export { default as isCloudEnvironment } from "./isCloudEnvironment";
diff --git a/cognee-frontend/src/utils/isCloudEnvironment.ts b/cognee-frontend/src/utils/isCloudEnvironment.ts
new file mode 100644
index 000000000..539551da6
--- /dev/null
+++ b/cognee-frontend/src/utils/isCloudEnvironment.ts
@@ -0,0 +1,4 @@
+
+export default function isCloudEnvironment() {
+  return process.env.NEXT_PUBLIC_IS_CLOUD_ENVIRONMENT?.toLowerCase() === "true";
+}
diff --git a/cognee/api/v1/notebooks/routers/get_notebooks_router.py b/cognee/api/v1/notebooks/routers/get_notebooks_router.py
index 5979f68ff..1d2861bf7 100644
--- a/cognee/api/v1/notebooks/routers/get_notebooks_router.py
+++ b/cognee/api/v1/notebooks/routers/get_notebooks_router.py
@@ -7,6 +7,7 @@ from fastapi import APIRouter, Depends
 
 from cognee.api.DTO import InDTO
 from cognee.infrastructure.databases.relational import get_async_session
+from cognee.infrastructure.utils.run_async import run_async
 from cognee.modules.notebooks.models import Notebook, NotebookCell
 from cognee.modules.notebooks.operations import run_in_local_sandbox
 from cognee.modules.users.models import User
@@ -74,7 +75,7 @@ def get_notebooks_router():
             if notebook is None:
                 return JSONResponse(status_code=404, content={"error": "Notebook not found"})
 
-            result, error = run_in_local_sandbox(run_code.content)
+            result, error = await run_async(run_in_local_sandbox, run_code.content)
 
             return JSONResponse(
                 status_code=200, content={"result": jsonable_encoder(result), "error": error}
diff --git a/cognee/infrastructure/utils/run_sync.py b/cognee/infrastructure/utils/run_sync.py
index 4945b0b98..913df4666 100644
--- a/cognee/infrastructure/utils/run_sync.py
+++ b/cognee/infrastructure/utils/run_sync.py
@@ -8,8 +8,15 @@ def run_sync(coro, timeout=None):
 
     def runner():
         nonlocal result, exception
+
         try:
-            result = asyncio.run(coro)
+            try:
+                running_loop = asyncio.get_running_loop()
+
+                result = asyncio.run_coroutine_threadsafe(coro, running_loop).result(timeout)
+            except RuntimeError:
+                result = asyncio.run(coro)
+
         except Exception as e:
             exception = e
 
diff --git a/cognee/modules/notebooks/operations/run_in_local_sandbox.py b/cognee/modules/notebooks/operations/run_in_local_sandbox.py
index f3955c256..b0c4da4b6 100644
--- a/cognee/modules/notebooks/operations/run_in_local_sandbox.py
+++ b/cognee/modules/notebooks/operations/run_in_local_sandbox.py
@@ -5,7 +5,6 @@ import traceback
 
 def wrap_in_async_handler(user_code: str) -> str:
     return (
-        "import asyncio\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"))
@@ -28,19 +27,6 @@ def run_in_local_sandbox(code, environment=None):
 
     printOutput = []
 
-    # def process_output(output):
-    #     try:
-    #         result = json.loads(
-    #             re.sub(
-    #                 r"'([^']*)'", r'"\1"',
-    #                 re.sub(r"\bNone\b", "null", output)
-    #             )
-    #         )
-    #         result = json.loads(output)
-    #         return result
-    #     except json.JSONDecodeError:
-    #         return output
-
     def customPrintFunction(output):
         printOutput.append(output)
 
diff --git a/cognee/modules/retrieval/utils/brute_force_triplet_search.py b/cognee/modules/retrieval/utils/brute_force_triplet_search.py
index 9e57ddc67..aa3233776 100644
--- a/cognee/modules/retrieval/utils/brute_force_triplet_search.py
+++ b/cognee/modules/retrieval/utils/brute_force_triplet_search.py
@@ -63,9 +63,10 @@ async def get_memory_fragment(
     if properties_to_project is None:
         properties_to_project = ["id", "description", "name", "type", "text"]
 
+    memory_fragment = CogneeGraph()
+
     try:
         graph_engine = await get_graph_engine()
-        memory_fragment = CogneeGraph()
 
         await memory_fragment.project_graph_from_db(
             graph_engine,
diff --git a/cognee/modules/search/methods/search.py b/cognee/modules/search/methods/search.py
index 8e4d41509..405207114 100644
--- a/cognee/modules/search/methods/search.py
+++ b/cognee/modules/search/methods/search.py
@@ -98,7 +98,9 @@ async def search(
         query.id,
         json.dumps(
             jsonable_encoder(
-                await prepare_search_result(search_results)
+                await prepare_search_result(
+                    search_results[0] if isinstance(search_results, list) else search_results
+                )
                 if use_combined_context
                 else [
                     await prepare_search_result(search_result) for search_result in search_results
@@ -109,7 +111,9 @@ async def search(
     )
 
     if use_combined_context:
-        prepared_search_results = await prepare_search_result(search_results)
+        prepared_search_results = await prepare_search_result(
+            search_results[0] if isinstance(search_results, list) else search_results
+        )
         result = prepared_search_results["result"]
         graphs = prepared_search_results["graphs"]
         context = prepared_search_results["context"]
diff --git a/cognee/modules/search/utils/prepare_search_result.py b/cognee/modules/search/utils/prepare_search_result.py
index bdcfa9928..6165e3166 100644
--- a/cognee/modules/search/utils/prepare_search_result.py
+++ b/cognee/modules/search/utils/prepare_search_result.py
@@ -13,10 +13,10 @@ async def prepare_search_result(search_result):
     context_texts = {}
 
     if isinstance(context, List) and len(context) > 0 and isinstance(context[0], Edge):
-        result_graph = transform_context_to_graph(context)
+        context_graph = transform_context_to_graph(context)
 
         graphs = {
-            "*": result_graph,
+            "*": context_graph,
         }
         context_texts = {
             "*": await resolve_edges_to_text(context),