khoj 1.33.3.dev32__py3-none-any.whl

khoj/routers/api_model.py

@@ -0,0 +1,123 @@
+import json
+import logging
+from typing import Dict, Optional, Union
+
+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.routers.helpers import update_telemetry_state
+
+api_model = APIRouter()
+logger = logging.getLogger(__name__)
+
+
+@api_model.get("/chat/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.name, "id": conversation_option.id})
+
+    return Response(content=json.dumps(all_conversation_options), media_type="application/json", status_code=200)
+
+
+@api_model.get("/chat")
+@requires(["authenticated"])
+def get_user_chat_model(
+    request: Request,
+    client: Optional[str] = None,
+):
+    user = request.user.object
+
+    chat_model = ConversationAdapters.get_chat_model(user)
+
+    if chat_model is None:
+        chat_model = ConversationAdapters.get_default_chat_model(user)
+
+    return Response(status_code=200, content=json.dumps({"id": chat_model.id, "chat_model": chat_model.name}))
+
+
+@api_model.post("/chat", 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_model.post("/voice", 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_model.post("/paint", 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"}

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/api_subscription.py

@@ -0,0 +1,144 @@
+import json
+import logging
+import os
+from datetime import datetime, timezone
+
+from asgiref.sync import sync_to_async
+from fastapi import APIRouter, Request, Response
+from starlette.authentication import requires
+
+from khoj.database import adapters
+from khoj.database.models import KhojUser, Subscription
+from khoj.routers.helpers import update_telemetry_state
+from khoj.utils import state
+
+# Stripe integration for Khoj Cloud Subscription
+if state.billing_enabled:
+    import stripe
+
+    stripe.api_key = os.getenv("STRIPE_API_KEY")
+endpoint_secret = os.getenv("STRIPE_SIGNING_SECRET")
+logger = logging.getLogger(__name__)
+subscription_router = APIRouter()
+
+
+@subscription_router.post("")
+async def subscribe(request: Request):
+    """Webhook for Stripe to send subscription events to Khoj Cloud"""
+    event = None
+    try:
+        payload = await request.body()
+        sig_header = request.headers["stripe-signature"]
+        event = stripe.Webhook.construct_event(payload, sig_header, endpoint_secret)
+    except ValueError as e:
+        # Invalid payload
+        raise e
+    except stripe.error.SignatureVerificationError as e:
+        # Invalid signature
+        raise e
+
+    event_type = event["type"]
+    if event_type not in {
+        "invoice.paid",
+        "customer.subscription.updated",
+        "customer.subscription.deleted",
+    }:
+        logger.warning(f"Unhandled Stripe event type: {event['type']}")
+        return {"success": False}
+
+    # Retrieve the customer's details
+    subscription = event["data"]["object"]
+    customer_id = subscription["customer"]
+    customer = stripe.Customer.retrieve(customer_id)
+    customer_email = customer["email"]
+    user = None
+    is_new = False
+
+    # Handle valid stripe webhook events
+    success = True
+    if event_type in {"invoice.paid"}:
+        # Mark the user as subscribed and update the next renewal date on payment
+        subscription = stripe.Subscription.list(customer=customer_id).data[0]
+        renewal_date = datetime.fromtimestamp(subscription["current_period_end"], tz=timezone.utc)
+        user, is_new = await adapters.set_user_subscription(
+            customer_email, is_recurring=True, renewal_date=renewal_date
+        )
+        success = user is not None
+    elif event_type in {"customer.subscription.updated"}:
+        user_subscription = await sync_to_async(adapters.get_user_subscription)(customer_email)
+
+        renewal_date = None
+        if subscription["current_period_end"]:
+            renewal_date = datetime.fromtimestamp(subscription["current_period_end"], tz=timezone.utc)
+
+        # Allow updating subscription status if paid user
+        if user_subscription and user_subscription.renewal_date:
+            # Mark user as unsubscribed or resubscribed
+            is_recurring = not subscription["cancel_at_period_end"]
+            user, is_new = await adapters.set_user_subscription(
+                customer_email, is_recurring=is_recurring, renewal_date=renewal_date
+            )
+            success = user is not None
+    elif event_type in {"customer.subscription.deleted"}:
+        # Reset the user to trial state
+        user, is_new = await adapters.set_user_subscription(
+            customer_email, is_recurring=False, renewal_date=None, type=Subscription.Type.TRIAL
+        )
+        success = user is not None
+
+    if user and is_new:
+        update_telemetry_state(
+            request=request,
+            telemetry_type="api",
+            api="create_user",
+            metadata={"server_id": str(user.user.uuid)},
+        )
+        logger.log(logging.INFO, f"🥳 New User Created: {user.user.uuid}")
+
+    logger.info(f'Stripe subscription {event["type"]} for {customer_email}')
+    return {"success": success}
+
+
+@subscription_router.patch("")
+@requires(["authenticated"])
+async def update_subscription(request: Request, operation: str):
+    # Retrieve the customer's details
+    email = request.user.object.email
+    customers = stripe.Customer.list(email=email).auto_paging_iter()
+    customer = next(customers, None)
+    if customer is None:
+        return {"success": False, "message": "Customer not found"}
+
+    if operation == "cancel":
+        customer_id = customer.id
+        for subscription in stripe.Subscription.list(customer=customer_id):
+            stripe.Subscription.modify(subscription.id, cancel_at_period_end=True)
+        return {"success": True}
+
+    elif operation == "resubscribe":
+        subscriptions = stripe.Subscription.list(customer=customer.id).auto_paging_iter()
+        # Find the subscription that is set to cancel at the end of the period
+        for subscription in subscriptions:
+            if subscription.cancel_at_period_end:
+                # Update the subscription to not cancel at the end of the period
+                stripe.Subscription.modify(subscription.id, cancel_at_period_end=False)
+                return {"success": True}
+        return {"success": False, "message": "No subscription found that is set to cancel"}
+
+    return {"success": False, "message": "Invalid operation"}
+
+
+@subscription_router.post("/trial", response_class=Response)
+@requires(["authenticated"])
+async def start_trial(request: Request) -> Response:
+    user: KhojUser = request.user.object
+
+    # Start a trial for the user
+    updated_subscription = await adapters.astart_trial_subscription(user)
+
+    # Return trial status as a JSON response
+    return Response(
+        content=json.dumps({"trial_enabled": updated_subscription is not None}),
+        media_type="application/json",
+        status_code=200,
+    )

khoj/routers/auth.py

@@ -0,0 +1,307 @@
+import asyncio
+import datetime
+import logging
+import os
+from typing import Optional
+from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
+
+import requests
+from fastapi import APIRouter, Depends
+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.app.settings import DISABLE_HTTPS
+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 (
+    EmailVerificationApiRateLimiter,
+    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)
+
+    user, is_new = await aget_or_create_user_by_email(form.email)
+
+    if user:
+        unique_id = user.email_verification_code
+        await send_magic_link_email(user.email, unique_id, request.base_url)
+        if is_new:
+            update_telemetry_state(
+                request=request,
+                telemetry_type="api",
+                api="create_user__email",
+                metadata={"server_id": str(user.uuid)},
+            )
+            logger.log(logging.INFO, f"🥳 New User Created: {user.uuid}")
+
+    return Response(status_code=200)
+
+
+@auth_router.get("/magic")
+async def sign_in_with_magic_link(
+    request: Request,
+    code: str,
+    email: str,
+    rate_limiter=Depends(
+        EmailVerificationApiRateLimiter(requests=10, window=60 * 60 * 24, slug="magic_link_verification")
+    ),
+):
+    user, code_is_expired = await aget_user_validated_by_email_verification_code(code, email)
+
+    if user:
+        if code_is_expired:
+            request.session["user"] = {}
+            return Response(status_code=403)
+
+        id_info = {
+            "email": user.email,
+        }
+
+        request.session["user"] = dict(id_info)
+        return RedirectResponse(url="/")
+    return Response(status_code=401)
+
+
+@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_post(request: Request):
+    # This is maintained for compatibility with the /login endpoint
+    form = await request.form()
+    next_url = get_next_url(request)
+    for q in request.query_params:
+        if 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__google",
+                metadata={"server_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("/redirect")
+async def auth(request: Request):
+    next_url_path = get_next_url(request)
+
+    # Add query params from request, excluding OAuth params to next URL
+    oauth_params = {"code", "state", "scope", "authuser", "prompt", "session_state", "access_type", "next"}
+    query_params = {param: value for param, value in request.query_params.items() if param not in oauth_params}
+
+    # Rebuild next URL with updated query params
+    parsed_next_url_path = urlparse(next_url_path)
+    next_url = urlunparse(
+        (
+            parsed_next_url_path.scheme,
+            parsed_next_url_path.netloc,
+            parsed_next_url_path.path,
+            parsed_next_url_path.params,
+            urlencode(query_params, doseq=True),
+            parsed_next_url_path.fragment,
+        )
+    )
+
+    # Construct the full redirect URI including domain
+    base_url = str(request.base_url).rstrip("/")
+    if not DISABLE_HTTPS:
+        base_url = base_url.replace("http://", "https://")
+    redirect_uri = f"{base_url}{request.app.url_path_for('auth')}"
+
+    # Build the payload for the token request
+    code = request.query_params.get("code")
+    payload = {
+        "code": code,
+        "client_id": os.environ["GOOGLE_CLIENT_ID"],
+        "client_secret": os.environ["GOOGLE_CLIENT_SECRET"],
+        "redirect_uri": redirect_uri,
+        "grant_type": "authorization_code",
+    }
+
+    # Request the token from Google
+    verified_data = requests.post(
+        "https://oauth2.googleapis.com/token",
+        headers={"Content-Type": "application/x-www-form-urlencoded"},
+        data=payload,
+    )
+
+    # Validate the OAuth response
+    if verified_data.status_code != 200:
+        logger.error(f"Token request failed: {verified_data.text}")
+        try:
+            error_json = verified_data.json()
+            logger.error(f"Error response JSON for Google verification: {error_json}")
+        except ValueError:
+            logger.error("Response content is not valid JSON")
+        verified_data.raise_for_status()
+
+    credential = verified_data.json().get("id_token")
+    if not credential:
+        logger.error("Missing id_token in OAuth response")
+        return RedirectResponse(url="/login?error=invalid_token", status_code=HTTP_302_FOUND)
+
+    # Validate the OAuth token
+    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>")
+
+    # Get or create the authenticated user in the database
+    khoj_user = await get_or_create_user(idinfo)
+
+    # Set the user session if the user is authenticated
+    if khoj_user:
+        request.session["user"] = dict(idinfo)
+
+        # Send a welcome email to new users
+        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__google",
+                metadata={"server_id": str(khoj_user.uuid)},
+            )
+            logger.log(logging.INFO, f"🥳 New User Created: {khoj_user.uuid}")
+
+    # Redirect the user to the next URL
+    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="/")
+
+
+@auth_router.get("/oauth/metadata")
+async def oauth_metadata(request: Request):
+    redirect_uri = str(request.app.url_path_for("auth"))
+
+    return {
+        "google": {
+            "client_id": os.environ.get("GOOGLE_CLIENT_ID"),
+            "redirect_uri": f"{redirect_uri}",
+        }
+    }