Merge branch 'main' into feature/cog-2923-create-ci-test-for-fastembed

2025-09-12 11:24:14 -07:00 · 2025-09-12 11:24:14 -07:00 · 0b2c093144
commit 0b2c093144
parent 6c523e12d6 288218b209
23 changed files with 97 additions and 76 deletions
--- a/cognee-frontend/src/app/dashboard/AddDataToCognee.tsx
+++ b/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>
--- a/cognee-frontend/src/app/dashboard/Dashboard.tsx
+++ b/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}
              />
            )}
--- a/cognee-frontend/src/app/dashboard/InstanceDatasetsAccordion.tsx
+++ b/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,
--- a/cognee-frontend/src/app/dashboard/NotebooksAccordion.tsx
+++ b/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>) => {
--- a/cognee-frontend/src/ui/App/Loading/DefaultLoadingIndicator/LoadingIndicator.module.css
+++ b/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;
 }

--- a/cognee-frontend/src/ui/App/Loading/DefaultLoadingIndicator/LoadingIndicator.tsx
+++ b/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}`)} />
 }
--- a/cognee-frontend/src/ui/elements/Notebook/Notebook.tsx
+++ b/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>
          );
--- a/cognee-frontend/src/ui/elements/Notebook/NotebookCellHeader.tsx
+++ b/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>
--- a/cognee/api/v1/notebooks/routers/get_notebooks_router.py
+++ b/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(
--- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
+++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
@ -58,8 +58,8 @@ class SQLAlchemyAdapter:
        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,
--- a/cognee/infrastructure/databases/vector/config.py
+++ b/cognee/infrastructure/databases/vector/config.py
@ -39,7 +39,7 @@ class VectorConfig(BaseSettings):
            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")
--- a/cognee/infrastructure/utils/run_async.py
+++ b/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)
--- a/cognee/infrastructure/utils/run_sync.py
+++ b/cognee/infrastructure/utils/run_sync.py
@ -2,16 +2,17 @@ import asyncio
 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:
--- a/cognee/modules/notebooks/methods/get_notebook.py
+++ b/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()
--- a/cognee/modules/notebooks/methods/update_notebook.py
+++ b/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
--- a/cognee/modules/notebooks/operations/run_in_local_sandbox.py
+++ b/cognee/modules/notebooks/operations/run_in_local_sandbox.py
@ -5,16 +5,18 @@ import traceback

 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 run_in_local_sandbox(code, environment=None):
        printOutput.append(output)

    environment["print"] = customPrintFunction
+    environment["running_loop"] = loop

    try:
        exec(code, environment)
--- a/cognee/modules/retrieval/graph_completion_context_extension_retriever.py
+++ b/cognee/modules/retrieval/graph_completion_context_extension_retriever.py
@ -48,7 +48,7 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
        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.
--- a/cognee/modules/retrieval/graph_completion_cot_retriever.py
+++ b/cognee/modules/retrieval/graph_completion_cot_retriever.py
@ -58,7 +58,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
        query: str,
        context: Optional[List[Edge]] = None,
        max_iter=4,
-    ) -> str:
+    ) -> List[str]:
        """
        Generate completion responses based on a user query and contextual information.

--- a/cognee/modules/retrieval/graph_completion_retriever.py
+++ b/cognee/modules/retrieval/graph_completion_retriever.py
@ -135,7 +135,7 @@ class GraphCompletionRetriever(BaseGraphRetriever):
        self,
        query: str,
        context: Optional[List[Edge]] = None,
-    ) -> Any:
+    ) -> List[str]:
        """
        Generates a completion using graph connections context based on a query.

--- a/cognee/modules/retrieval/temporal_retriever.py
+++ b/cognee/modules/retrieval/temporal_retriever.py
@ -136,7 +136,7 @@ class TemporalRetriever(GraphCompletionRetriever):

        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)
--- a/cognee/modules/retrieval/user_qa_feedback.py
+++ b/cognee/modules/retrieval/user_qa_feedback.py
@ -1,4 +1,4 @@
-from typing import Any, Optional, List
+from typing import Optional, List

 from uuid import NAMESPACE_OID, uuid5, UUID
 from cognee.infrastructure.databases.graph import get_graph_engine
--- a/cognee/modules/search/utils/prepare_search_result.py
+++ b/cognee/modules/search/utils/prepare_search_result.py
@ -6,7 +6,7 @@ from cognee.modules.search.utils.transform_context_to_graph import transform_con


 async def prepare_search_result(search_result):
-    result, context, datasets = search_result
+    results, context, datasets = search_result

    graphs = None
    result_graph = None
@ -30,11 +30,11 @@ async def prepare_search_result(search_result):
            "*": "\n".join(cast(List[str], context)),
        }

-    if isinstance(result, List) and len(result) > 0 and isinstance(result[0], Edge):
-        result_graph = transform_context_to_graph(result)
+    if isinstance(results, List) and len(results) > 0 and isinstance(results[0], Edge):
+        result_graph = transform_context_to_graph(results)

    return {
-        "result": result_graph or result,
+        "result": result_graph or results[0] if len(results) == 1 else results,
        "graphs": graphs,
        "context": context_texts,
        "datasets": datasets,
--- a/pyproject.toml
+++ b/pyproject.toml
@ -1,7 +1,7 @@
 [project]
 name = "cognee"

-version = "0.3.2"
+version = "0.3.3"
 description = "Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning."
 authors = [
    { name = "Vasilije Markovic" },
@ -147,7 +147,7 @@ Homepage = "https://www.cognee.ai"
 Repository = "https://github.com/topoteretes/cognee"

 [project.scripts]
-cognee = "cognee.cli._cognee:main"
+cognee-cli = "cognee.cli._cognee:main"

 [build-system]
 requires = ["hatchling"]