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 {
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),