From 06213ea814a556d8a0e9500a0cfa550b111bc872 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Tue, 30 Apr 2024 17:41:17 +0530
Subject: [PATCH] Fix token retrieval when executing the job and name async job
 approriately

---
 src/khoj/database/adapters/__init__.py | 9 ++++++++-
 src/khoj/routers/auth.py               | 6 +++---
 src/khoj/routers/helpers.py            | 4 ++--
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py
index f9f2193a..ee72a76a 100644
--- a/src/khoj/database/adapters/__init__.py
+++ b/src/khoj/database/adapters/__init__.py
@@ -72,7 +72,14 @@ async def set_notion_config(token: str, user: KhojUser):
     return notion_config
 
 
-async def create_khoj_token(user: KhojUser, name=None):
+def create_khoj_token(user: KhojUser, name=None):
+    "Create Khoj API key for user"
+    token = f"kk-{secrets.token_urlsafe(32)}"
+    name = name or f"{generate_random_name().title()}"
+    return KhojApiUser.objects.create(token=token, user=user, name=name)
+
+
+async def acreate_khoj_token(user: KhojUser, name=None):
     "Create Khoj API key for user"
     token = f"kk-{secrets.token_urlsafe(32)}"
     name = name or f"{generate_random_name().title()}"
diff --git a/src/khoj/routers/auth.py b/src/khoj/routers/auth.py
index 199ccd2b..6e0e30c1 100644
--- a/src/khoj/routers/auth.py
+++ b/src/khoj/routers/auth.py
@@ -12,7 +12,7 @@ from starlette.responses import HTMLResponse, RedirectResponse, Response
 from starlette.status import HTTP_302_FOUND
 
 from khoj.database.adapters import (
-    create_khoj_token,
+    acreate_khoj_token,
     delete_khoj_token,
     get_khoj_tokens,
     get_or_create_user,
@@ -67,9 +67,9 @@ async def login(request: Request):
 async def generate_token(request: Request, token_name: Optional[str] = None):
     "Generate API token for given user"
     if token_name:
-        token = await create_khoj_token(user=request.user.object, name=token_name)
+        token = await acreate_khoj_token(user=request.user.object, name=token_name)
     else:
-        token = await create_khoj_token(user=request.user.object)
+        token = await acreate_khoj_token(user=request.user.object)
     return {
         "token": token.token,
         "name": token.name,
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index 29cf95e6..2840a5cb 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -886,9 +886,9 @@ def scheduled_chat(query_to_run: str, scheduling_request: str, subject: str, use
         # Add authorization request header in non-anonymous mode
         token = get_khoj_tokens(user)
         if is_none_or_empty(token):
-            token = create_khoj_token(user)
+            token = create_khoj_token(user).token
         else:
-            token = token[0]
+            token = token[0].token
         headers["Authorization"] = f"Bearer {token}"
 
     # Call the chat API endpoint with authenticated user token and query