khoj 1.16.1.dev15__py3-none-any.whl

khoj/routers/api_config.py

@@ -0,0 +1,434 @@
+import json
+import logging
+import math
+from typing import Dict, List, Optional, Union
+
+from asgiref.sync import sync_to_async
+from fastapi import APIRouter, HTTPException, Request
+from fastapi.requests import Request
+from fastapi.responses import Response
+from starlette.authentication import has_required_scope, requires
+
+from khoj.database import adapters
+from khoj.database.adapters import ConversationAdapters, EntryAdapters
+from khoj.database.models import Entry as DbEntry
+from khoj.database.models import (
+    GithubConfig,
+    KhojUser,
+    LocalMarkdownConfig,
+    LocalOrgConfig,
+    LocalPdfConfig,
+    LocalPlaintextConfig,
+    NotionConfig,
+    Subscription,
+)
+from khoj.routers.helpers import CommonQueryParams, update_telemetry_state
+from khoj.utils import constants, state
+from khoj.utils.rawconfig import (
+    FullConfig,
+    GithubContentConfig,
+    NotionContentConfig,
+    SearchConfig,
+)
+from khoj.utils.state import SearchType
+
+api_config = APIRouter()
+logger = logging.getLogger(__name__)
+
+
+def map_config_to_object(content_source: str):
+    if content_source == DbEntry.EntrySource.GITHUB:
+        return GithubConfig
+    if content_source == DbEntry.EntrySource.NOTION:
+        return NotionConfig
+    if content_source == DbEntry.EntrySource.COMPUTER:
+        return "Computer"
+
+
+async def map_config_to_db(config: FullConfig, user: KhojUser):
+    if config.content_type:
+        if config.content_type.org:
+            await LocalOrgConfig.objects.filter(user=user).adelete()
+            await LocalOrgConfig.objects.acreate(
+                input_files=config.content_type.org.input_files,
+                input_filter=config.content_type.org.input_filter,
+                index_heading_entries=config.content_type.org.index_heading_entries,
+                user=user,
+            )
+        if config.content_type.markdown:
+            await LocalMarkdownConfig.objects.filter(user=user).adelete()
+            await LocalMarkdownConfig.objects.acreate(
+                input_files=config.content_type.markdown.input_files,
+                input_filter=config.content_type.markdown.input_filter,
+                index_heading_entries=config.content_type.markdown.index_heading_entries,
+                user=user,
+            )
+        if config.content_type.pdf:
+            await LocalPdfConfig.objects.filter(user=user).adelete()
+            await LocalPdfConfig.objects.acreate(
+                input_files=config.content_type.pdf.input_files,
+                input_filter=config.content_type.pdf.input_filter,
+                index_heading_entries=config.content_type.pdf.index_heading_entries,
+                user=user,
+            )
+        if config.content_type.plaintext:
+            await LocalPlaintextConfig.objects.filter(user=user).adelete()
+            await LocalPlaintextConfig.objects.acreate(
+                input_files=config.content_type.plaintext.input_files,
+                input_filter=config.content_type.plaintext.input_filter,
+                index_heading_entries=config.content_type.plaintext.index_heading_entries,
+                user=user,
+            )
+        if config.content_type.github:
+            await adapters.set_user_github_config(
+                user=user,
+                pat_token=config.content_type.github.pat_token,
+                repos=config.content_type.github.repos,
+            )
+        if config.content_type.notion:
+            await adapters.set_notion_config(
+                user=user,
+                token=config.content_type.notion.token,
+            )
+
+
+def _initialize_config():
+    if state.config is None:
+        state.config = FullConfig()
+        state.config.search_type = SearchConfig.model_validate(constants.default_config["search-type"])
+
+
+@api_config.post("/data/content-source/github", status_code=200)
+@requires(["authenticated"])
+async def set_content_config_github_data(
+    request: Request,
+    updated_config: Union[GithubContentConfig, None],
+    client: Optional[str] = None,
+):
+    _initialize_config()
+
+    user = request.user.object
+
+    try:
+        await adapters.set_user_github_config(
+            user=user,
+            pat_token=updated_config.pat_token,
+            repos=updated_config.repos,
+        )
+    except Exception as e:
+        logger.error(e, exc_info=True)
+        raise HTTPException(status_code=500, detail="Failed to set Github config")
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="set_content_config",
+        client=client,
+        metadata={"content_type": "github"},
+    )
+
+    return {"status": "ok"}
+
+
+@api_config.post("/data/content-source/notion", status_code=200)
+@requires(["authenticated"])
+async def set_content_config_notion_data(
+    request: Request,
+    updated_config: Union[NotionContentConfig, None],
+    client: Optional[str] = None,
+):
+    _initialize_config()
+
+    user = request.user.object
+
+    try:
+        await adapters.set_notion_config(
+            user=user,
+            token=updated_config.token,
+        )
+    except Exception as e:
+        logger.error(e, exc_info=True)
+        raise HTTPException(status_code=500, detail="Failed to set Github config")
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="set_content_config",
+        client=client,
+        metadata={"content_type": "notion"},
+    )
+
+    return {"status": "ok"}
+
+
+@api_config.delete("/data/content-source/{content_source}", status_code=200)
+@requires(["authenticated"])
+async def remove_content_source_data(
+    request: Request,
+    content_source: str,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="delete_content_config",
+        client=client,
+        metadata={"content_source": content_source},
+    )
+
+    content_object = map_config_to_object(content_source)
+    if content_object is None:
+        raise ValueError(f"Invalid content source: {content_source}")
+    elif content_object != "Computer":
+        await content_object.objects.filter(user=user).adelete()
+    await sync_to_async(EntryAdapters.delete_all_entries)(user, file_source=content_source)
+
+    enabled_content = await sync_to_async(EntryAdapters.get_unique_file_types)(user)
+    return {"status": "ok"}
+
+
+@api_config.delete("/data/file", status_code=200)
+@requires(["authenticated"])
+async def remove_file_data(
+    request: Request,
+    filename: str,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="delete_file",
+        client=client,
+    )
+
+    await EntryAdapters.adelete_entry_by_file(user, filename)
+
+    return {"status": "ok"}
+
+
+@api_config.get("/data/{content_source}", response_model=List[str])
+@requires(["authenticated"])
+async def get_all_filenames(
+    request: Request,
+    content_source: str,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="get_all_filenames",
+        client=client,
+    )
+
+    return await sync_to_async(list)(EntryAdapters.get_all_filenames_by_source(user, content_source))  # type: ignore[call-arg]
+
+
+@api_config.get("/data/conversation/model/options", response_model=Dict[str, Union[str, int]])
+def get_chat_model_options(
+    request: Request,
+    client: Optional[str] = None,
+):
+    conversation_options = ConversationAdapters.get_conversation_processor_options().all()
+
+    all_conversation_options = list()
+    for conversation_option in conversation_options:
+        all_conversation_options.append({"chat_model": conversation_option.chat_model, "id": conversation_option.id})
+
+    return Response(content=json.dumps(all_conversation_options), media_type="application/json", status_code=200)
+
+
+@api_config.get("/data/conversation/model")
+@requires(["authenticated"])
+def get_user_chat_model(
+    request: Request,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    chat_model = ConversationAdapters.get_conversation_config(user)
+
+    if chat_model is None:
+        chat_model = ConversationAdapters.get_default_conversation_config()
+
+    return Response(status_code=200, content=json.dumps({"id": chat_model.id, "chat_model": chat_model.chat_model}))
+
+
+@api_config.post("/data/conversation/model", status_code=200)
+@requires(["authenticated", "premium"])
+async def update_chat_model(
+    request: Request,
+    id: str,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    new_config = await ConversationAdapters.aset_user_conversation_processor(user, int(id))
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="set_conversation_chat_model",
+        client=client,
+        metadata={"processor_conversation_type": "conversation"},
+    )
+
+    if new_config is None:
+        return {"status": "error", "message": "Model not found"}
+
+    return {"status": "ok"}
+
+
+@api_config.post("/data/voice/model", status_code=200)
+@requires(["authenticated", "premium"])
+async def update_voice_model(
+    request: Request,
+    id: str,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    new_config = await ConversationAdapters.aset_user_voice_model(user, id)
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="set_voice_model",
+        client=client,
+    )
+
+    if new_config is None:
+        return Response(status_code=404, content=json.dumps({"status": "error", "message": "Model not found"}))
+
+    return Response(status_code=202, content=json.dumps({"status": "ok"}))
+
+
+@api_config.post("/data/search/model", status_code=200)
+@requires(["authenticated"])
+async def update_search_model(
+    request: Request,
+    id: str,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    prev_config = await adapters.aget_user_search_model(user)
+    new_config = await adapters.aset_user_search_model(user, int(id))
+
+    if prev_config and int(id) != prev_config.id and new_config:
+        await EntryAdapters.adelete_all_entries(user)
+
+    if not prev_config:
+        # If the use was just using the default config, delete all the entries and set the new config.
+        await EntryAdapters.adelete_all_entries(user)
+
+    if new_config is None:
+        return {"status": "error", "message": "Model not found"}
+    else:
+        update_telemetry_state(
+            request=request,
+            telemetry_type="api",
+            api="set_search_model",
+            client=client,
+            metadata={"search_model": new_config.setting.name},
+        )
+
+    return {"status": "ok"}
+
+
+@api_config.post("/data/paint/model", status_code=200)
+@requires(["authenticated"])
+async def update_paint_model(
+    request: Request,
+    id: str,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+    subscribed = has_required_scope(request, ["premium"])
+
+    if not subscribed:
+        raise HTTPException(status_code=403, detail="User is not subscribed to premium")
+
+    new_config = await ConversationAdapters.aset_user_text_to_image_model(user, int(id))
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="set_paint_model",
+        client=client,
+        metadata={"paint_model": new_config.setting.model_name},
+    )
+
+    if new_config is None:
+        return {"status": "error", "message": "Model not found"}
+
+    return {"status": "ok"}
+
+
+@api_config.get("/index/size", response_model=Dict[str, int])
+@requires(["authenticated"])
+async def get_indexed_data_size(request: Request, common: CommonQueryParams):
+    user = request.user.object
+    indexed_data_size_in_mb = await sync_to_async(EntryAdapters.get_size_of_indexed_data_in_mb)(user)
+    return Response(
+        content=json.dumps({"indexed_data_size_in_mb": math.ceil(indexed_data_size_in_mb)}),
+        media_type="application/json",
+        status_code=200,
+    )
+
+
+@api_config.post("/user/name", status_code=200)
+@requires(["authenticated"])
+def set_user_name(
+    request: Request,
+    name: str,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    split_name = name.split(" ")
+
+    if len(split_name) > 2:
+        raise HTTPException(status_code=400, detail="Name must be in the format: Firstname Lastname")
+
+    if len(split_name) == 1:
+        first_name = split_name[0]
+        last_name = ""
+    else:
+        first_name, last_name = split_name[0], split_name[-1]
+
+    adapters.set_user_name(user, first_name, last_name)
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="set_user_name",
+        client=client,
+    )
+
+    return {"status": "ok"}
+
+
+@api_config.get("/types", response_model=List[str])
+@requires(["authenticated"])
+def get_config_types(
+    request: Request,
+):
+    user = request.user.object
+    enabled_file_types = EntryAdapters.get_unique_file_types(user)
+    configured_content_types = list(enabled_file_types)
+
+    if state.config and state.config.content_type:
+        for ctype in state.config.content_type.model_dump(exclude_none=True):
+            configured_content_types.append(ctype)
+
+    return [
+        search_type.value
+        for search_type in SearchType
+        if (search_type.value in configured_content_types) or search_type == SearchType.All
+    ]

khoj/routers/api_phone.py

@@ -0,0 +1,86 @@
+import logging
+from typing import Optional
+
+from fastapi import APIRouter, Depends, HTTPException, Request
+from starlette.authentication import requires
+
+from khoj.database import adapters
+from khoj.database.models import KhojUser
+from khoj.routers.helpers import ApiUserRateLimiter, update_telemetry_state
+from khoj.routers.twilio import create_otp, verify_otp
+
+api_phone = APIRouter()
+logger = logging.getLogger(__name__)
+
+
+@api_phone.post("", status_code=200)
+@requires(["authenticated"])
+async def update_phone_number(
+    request: Request,
+    phone_number: str,
+    client: Optional[str] = None,
+    rate_limiter_per_day=Depends(
+        ApiUserRateLimiter(requests=5, subscribed_requests=5, window=60 * 60 * 24, slug="update_phone")
+    ),
+):
+    user = request.user.object
+
+    await adapters.aset_user_phone_number(user, phone_number)
+    create_otp(user)
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="set_phone_number",
+        client=client,
+        metadata={"phone_number": phone_number},
+    )
+
+    return {"status": "ok"}
+
+
+@api_phone.delete("", status_code=200)
+@requires(["authenticated"])
+async def delete_phone_number(
+    request: Request,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    await adapters.aremove_phone_number(user)
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="delete_phone_number",
+        client=client,
+    )
+
+    return {"status": "ok"}
+
+
+@api_phone.post("/verify", status_code=200)
+@requires(["authenticated"])
+async def verify_mobile_otp(
+    request: Request,
+    code: str,
+    client: Optional[str] = None,
+    rate_limiter_per_day=Depends(
+        ApiUserRateLimiter(requests=5, subscribed_requests=5, window=60 * 60 * 24, slug="verify_phone")
+    ),
+):
+    user: KhojUser = request.user.object
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="verify_phone_number",
+        client=client,
+    )
+
+    if not verify_otp(user, code):
+        raise HTTPException(status_code=400, detail="Invalid OTP")
+
+    user.verified_phone_number = True
+    await user.asave()
+    return {"status": "ok"}

khoj/routers/auth.py

@@ -0,0 +1,181 @@
+import asyncio
+import datetime
+import logging
+import os
+from typing import Optional
+
+from fastapi import APIRouter
+from pydantic import BaseModel, EmailStr
+from starlette.authentication import requires
+from starlette.config import Config
+from starlette.requests import Request
+from starlette.responses import HTMLResponse, RedirectResponse, Response
+from starlette.status import HTTP_302_FOUND
+
+from khoj.database.adapters import (
+    acreate_khoj_token,
+    aget_or_create_user_by_email,
+    aget_user_validated_by_email_verification_code,
+    delete_khoj_token,
+    get_khoj_tokens,
+    get_or_create_user,
+)
+from khoj.routers.email import send_magic_link_email, send_welcome_email
+from khoj.routers.helpers import get_next_url, update_telemetry_state
+from khoj.utils import state
+
+logger = logging.getLogger(__name__)
+
+auth_router = APIRouter()
+
+
+class MagicLinkForm(BaseModel):
+    email: EmailStr
+
+
+if not state.anonymous_mode:
+    missing_requirements = []
+    from authlib.integrations.starlette_client import OAuth, OAuthError
+
+    try:
+        from google.auth.transport import requests as google_requests
+        from google.oauth2 import id_token
+    except ImportError:
+        missing_requirements += ["Install the Khoj production package with `pip install khoj[prod]`"]
+    if not os.environ.get("RESEND_API_KEY") and (
+        not os.environ.get("GOOGLE_CLIENT_ID") or not os.environ.get("GOOGLE_CLIENT_SECRET")
+    ):
+        missing_requirements += [
+            "Set your RESEND_API_KEY or GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET as environment variables"
+        ]
+    if missing_requirements:
+        requirements_string = "\n   - " + "\n   - ".join(missing_requirements)
+        error_msg = f"🚨 Start Khoj with --anonymous-mode flag or to enable authentication:{requirements_string}"
+        logger.error(error_msg)
+
+    config = Config(environ=os.environ)
+
+    oauth = OAuth(config)
+
+    CONF_URL = "https://accounts.google.com/.well-known/openid-configuration"
+    oauth.register(name="google", server_metadata_url=CONF_URL, client_kwargs={"scope": "openid email profile"})
+
+
+@auth_router.get("/login")
+async def login_get(request: Request):
+    redirect_uri = str(request.app.url_path_for("auth"))
+    return await oauth.google.authorize_redirect(request, redirect_uri)
+
+
+@auth_router.post("/login")
+async def login(request: Request):
+    redirect_uri = str(request.app.url_path_for("auth"))
+    return await oauth.google.authorize_redirect(request, redirect_uri)
+
+
+@auth_router.post("/magic")
+async def login_magic_link(request: Request, form: MagicLinkForm):
+    if request.user.is_authenticated:
+        # Clear the session if user is already authenticated
+        request.session.pop("user", None)
+
+    email = form.email
+    user = await aget_or_create_user_by_email(email)
+    unique_id = user.email_verification_code
+
+    if user:
+        await send_magic_link_email(email, unique_id, request.base_url)
+
+    return Response(status_code=200)
+
+
+@auth_router.get("/magic")
+async def sign_in_with_magic_link(request: Request, code: str):
+    user = await aget_user_validated_by_email_verification_code(code)
+    if user:
+        id_info = {
+            "email": user.email,
+        }
+
+        request.session["user"] = dict(id_info)
+        return RedirectResponse(url="/")
+    return RedirectResponse(request.app.url_path_for("login_page"))
+
+
+@auth_router.post("/token")
+@requires(["authenticated"], redirect="login_page")
+async def generate_token(request: Request, token_name: Optional[str] = None):
+    "Generate API token for given user"
+    if token_name:
+        token = await acreate_khoj_token(user=request.user.object, name=token_name)
+    else:
+        token = await acreate_khoj_token(user=request.user.object)
+    return {
+        "token": token.token,
+        "name": token.name,
+    }
+
+
+@auth_router.get("/token")
+@requires(["authenticated"], redirect="login_page")
+def get_tokens(request: Request):
+    "Get API tokens enabled for given user"
+    tokens = get_khoj_tokens(user=request.user.object)
+    return tokens
+
+
+@auth_router.delete("/token")
+@requires(["authenticated"], redirect="login_page")
+async def delete_token(request: Request, token: str):
+    "Delete API token for given user"
+    return await delete_khoj_token(user=request.user.object, token=token)
+
+
+@auth_router.post("/redirect")
+async def auth(request: Request):
+    form = await request.form()
+    next_url = get_next_url(request)
+    for q in request.query_params:
+        if not q == "next":
+            next_url += f"&{q}={request.query_params[q]}"
+
+    credential = form.get("credential")
+
+    csrf_token_cookie = request.cookies.get("g_csrf_token")
+    if not csrf_token_cookie:
+        logger.info("Missing CSRF token. Redirecting user to login page")
+        return RedirectResponse(url=next_url)
+    csrf_token_body = form.get("g_csrf_token")
+    if not csrf_token_body:
+        logger.info("Missing CSRF token body. Redirecting user to login page")
+        return RedirectResponse(url=next_url)
+    if csrf_token_cookie != csrf_token_body:
+        return Response("Invalid CSRF token", status_code=400)
+
+    try:
+        idinfo = id_token.verify_oauth2_token(credential, google_requests.Request(), os.environ["GOOGLE_CLIENT_ID"])
+    except OAuthError as error:
+        return HTMLResponse(f"<h1>{error.error}</h1>")
+    khoj_user = await get_or_create_user(idinfo)
+
+    if khoj_user:
+        request.session["user"] = dict(idinfo)
+
+        if datetime.timedelta(minutes=3) > (datetime.datetime.now(datetime.timezone.utc) - khoj_user.date_joined):
+            asyncio.create_task(send_welcome_email(idinfo["name"], idinfo["email"]))
+            update_telemetry_state(
+                request=request,
+                telemetry_type="api",
+                api="create_user",
+                metadata={"user_id": str(khoj_user.uuid)},
+            )
+            logger.log(logging.INFO, f"🥳 New User Created: {khoj_user.uuid}")
+            return RedirectResponse(url=next_url, status_code=HTTP_302_FOUND)
+
+    return RedirectResponse(url=next_url, status_code=HTTP_302_FOUND)
+
+
+@auth_router.get("/logout")
+async def logout(request: Request):
+    request.session.pop("user", None)
+    return RedirectResponse(url="/")