diff --git a/cognee-frontend/src/app/dashboard/AddDataToCognee.tsx b/cognee-frontend/src/app/dashboard/AddDataToCognee.tsx index 4f9a3179e..60e2a4204 100644 --- 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 = Add new data to a dataset? -
Please select a dataset to add data in.
If you don't have any, don't worry, we will create one for you.
+
Please select a {useCloud ? "cloud" : "local"} dataset to add data in.
If you don't have any, don't worry, we will create one for you.
setRunInstance(event.currentTarget.value)} className="!bg-transparent outline-none cursor-pointer !hover:bg-gray-50"> + {isCloudEnvironment() ? ( +
+ cloud cognee +
+ ) : ( +
+ local cognee +
+ )} + {/* + */}
diff --git a/cognee/api/v1/notebooks/routers/get_notebooks_router.py b/cognee/api/v1/notebooks/routers/get_notebooks_router.py index 1d2861bf7..832b72a94 100644 --- 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( diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py index 9369c0653..88d2abc7e 100644 --- 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, diff --git a/cognee/infrastructure/databases/vector/config.py b/cognee/infrastructure/databases/vector/config.py index f8fad473e..23678cca4 100644 --- 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") diff --git a/cognee/infrastructure/utils/run_async.py b/cognee/infrastructure/utils/run_async.py index 5480ea110..b6b428fc8 100644 --- 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) diff --git a/cognee/infrastructure/utils/run_sync.py b/cognee/infrastructure/utils/run_sync.py index 913df4666..aa5e28b2e 100644 --- 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: diff --git a/cognee/modules/notebooks/methods/get_notebook.py b/cognee/modules/notebooks/methods/get_notebook.py index ca666edaa..a3d62e02f 100644 --- 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() diff --git a/cognee/modules/notebooks/methods/update_notebook.py b/cognee/modules/notebooks/methods/update_notebook.py index 4d194202f..8d01a06f3 100644 --- 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 diff --git a/cognee/modules/notebooks/operations/run_in_local_sandbox.py b/cognee/modules/notebooks/operations/run_in_local_sandbox.py index b0c4da4b6..071deafb7 100644 --- 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) diff --git a/cognee/modules/retrieval/graph_completion_context_extension_retriever.py b/cognee/modules/retrieval/graph_completion_context_extension_retriever.py index f4c37fc93..6aca21808 100644 --- 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. diff --git a/cognee/modules/retrieval/graph_completion_cot_retriever.py b/cognee/modules/retrieval/graph_completion_cot_retriever.py index f51433751..16d573351 100644 --- 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. diff --git a/cognee/modules/retrieval/graph_completion_retriever.py b/cognee/modules/retrieval/graph_completion_retriever.py index 29b1e9d19..cd4944ca9 100644 --- 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. diff --git a/cognee/modules/retrieval/temporal_retriever.py b/cognee/modules/retrieval/temporal_retriever.py index 36cdbd33f..c5fe2094c 100644 --- 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) diff --git a/cognee/modules/retrieval/user_qa_feedback.py b/cognee/modules/retrieval/user_qa_feedback.py index f667f785f..21e3d8e59 100644 --- 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 diff --git a/cognee/modules/search/utils/prepare_search_result.py b/cognee/modules/search/utils/prepare_search_result.py index 6165e3166..19cbe07ac 100644 --- 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, diff --git a/pyproject.toml b/pyproject.toml index 6e6d9e1a0..6e3bd19f7 100644 --- 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"]