khoj 1.16.1.dev15__py3-none-any.whl

khoj/routers/indexer.py

@@ -0,0 +1,349 @@
+import asyncio
+import logging
+from typing import Dict, Optional, Union
+
+from fastapi import APIRouter, Depends, Header, Request, Response, UploadFile
+from pydantic import BaseModel
+from starlette.authentication import requires
+
+from khoj.database.models import GithubConfig, KhojUser, NotionConfig
+from khoj.processor.content.docx.docx_to_entries import DocxToEntries
+from khoj.processor.content.github.github_to_entries import GithubToEntries
+from khoj.processor.content.images.image_to_entries import ImageToEntries
+from khoj.processor.content.markdown.markdown_to_entries import MarkdownToEntries
+from khoj.processor.content.notion.notion_to_entries import NotionToEntries
+from khoj.processor.content.org_mode.org_to_entries import OrgToEntries
+from khoj.processor.content.pdf.pdf_to_entries import PdfToEntries
+from khoj.processor.content.plaintext.plaintext_to_entries import PlaintextToEntries
+from khoj.routers.helpers import ApiIndexedDataLimiter, update_telemetry_state
+from khoj.search_type import text_search
+from khoj.utils import constants, state
+from khoj.utils.config import SearchModels
+from khoj.utils.helpers import LRU, get_file_type
+from khoj.utils.rawconfig import ContentConfig, FullConfig, SearchConfig
+from khoj.utils.yaml import save_config_to_file_updated_state
+
+logger = logging.getLogger(__name__)
+
+indexer = APIRouter()
+
+
+class File(BaseModel):
+    path: str
+    content: Union[str, bytes]
+
+
+class IndexBatchRequest(BaseModel):
+    files: list[File]
+
+
+class IndexerInput(BaseModel):
+    org: Optional[dict[str, str]] = None
+    markdown: Optional[dict[str, str]] = None
+    pdf: Optional[dict[str, bytes]] = None
+    plaintext: Optional[dict[str, str]] = None
+    image: Optional[dict[str, bytes]] = None
+    docx: Optional[dict[str, bytes]] = None
+
+
+@indexer.post("/update")
+@requires(["authenticated"])
+async def update(
+    request: Request,
+    files: list[UploadFile],
+    force: bool = False,
+    t: Optional[Union[state.SearchType, str]] = state.SearchType.All,
+    client: Optional[str] = None,
+    user_agent: Optional[str] = Header(None),
+    referer: Optional[str] = Header(None),
+    host: Optional[str] = Header(None),
+    indexed_data_limiter: ApiIndexedDataLimiter = Depends(
+        ApiIndexedDataLimiter(
+            incoming_entries_size_limit=10,
+            subscribed_incoming_entries_size_limit=25,
+            total_entries_size_limit=10,
+            subscribed_total_entries_size_limit=100,
+        )
+    ),
+):
+    user = request.user.object
+    index_files: Dict[str, Dict[str, str]] = {
+        "org": {},
+        "markdown": {},
+        "pdf": {},
+        "plaintext": {},
+        "image": {},
+        "docx": {},
+    }
+    try:
+        logger.info(f"📬 Updating content index via API call by {client} client")
+        for file in files:
+            file_content = file.file.read()
+            file_type, encoding = get_file_type(file.content_type, file_content)
+            if file_type in index_files:
+                index_files[file_type][file.filename] = file_content.decode(encoding) if encoding else file_content
+            else:
+                logger.warning(f"Skipped indexing unsupported file type sent by {client} client: {file.filename}")
+
+        indexer_input = IndexerInput(
+            org=index_files["org"],
+            markdown=index_files["markdown"],
+            pdf=index_files["pdf"],
+            plaintext=index_files["plaintext"],
+            image=index_files["image"],
+            docx=index_files["docx"],
+        )
+
+        if state.config == None:
+            logger.info("📬 Initializing content index on first run.")
+            default_full_config = FullConfig(
+                content_type=None,
+                search_type=SearchConfig.model_validate(constants.default_config["search-type"]),
+                processor=None,
+            )
+            state.config = default_full_config
+            default_content_config = ContentConfig(
+                org=None,
+                markdown=None,
+                pdf=None,
+                docx=None,
+                image=None,
+                github=None,
+                notion=None,
+                plaintext=None,
+            )
+            state.config.content_type = default_content_config
+            save_config_to_file_updated_state()
+            configure_search(state.search_models, state.config.search_type)
+
+        # Extract required fields from config
+        loop = asyncio.get_event_loop()
+        success = await loop.run_in_executor(
+            None,
+            configure_content,
+            indexer_input.model_dump(),
+            force,
+            t,
+            False,
+            user,
+        )
+        if not success:
+            raise RuntimeError("Failed to update content index")
+        logger.info(f"Finished processing batch indexing request")
+    except Exception as e:
+        logger.error(f"Failed to process batch indexing request: {e}", exc_info=True)
+        logger.error(
+            f'🚨 Failed to {"force " if force else ""}update {t} content index triggered via API call by {client} client: {e}',
+            exc_info=True,
+        )
+        return Response(content="Failed", status_code=500)
+
+    indexing_metadata = {
+        "num_org": len(index_files["org"]),
+        "num_markdown": len(index_files["markdown"]),
+        "num_pdf": len(index_files["pdf"]),
+        "num_plaintext": len(index_files["plaintext"]),
+        "num_image": len(index_files["image"]),
+        "num_docx": len(index_files["docx"]),
+    }
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="index/update",
+        client=client,
+        user_agent=user_agent,
+        referer=referer,
+        host=host,
+        metadata=indexing_metadata,
+    )
+
+    logger.info(f"📪 Content index updated via API call by {client} client")
+
+    indexed_filenames = ",".join(file for ctype in index_files for file in index_files[ctype]) or ""
+    return Response(content=indexed_filenames, status_code=200)
+
+
+def configure_search(search_models: SearchModels, search_config: Optional[SearchConfig]) -> Optional[SearchModels]:
+    # Run Validation Checks
+    if search_models is None:
+        search_models = SearchModels()
+
+    return search_models
+
+
+def configure_content(
+    files: Optional[dict[str, dict[str, str]]],
+    regenerate: bool = False,
+    t: Optional[state.SearchType] = state.SearchType.All,
+    full_corpus: bool = True,
+    user: KhojUser = None,
+) -> bool:
+    success = True
+    if t == None:
+        t = state.SearchType.All
+
+    if t is not None and t in [type.value for type in state.SearchType]:
+        t = state.SearchType(t)
+
+    if t is not None and not t.value in [type.value for type in state.SearchType]:
+        logger.warning(f"🚨 Invalid search type: {t}")
+        return False
+
+    search_type = t.value if t else None
+
+    no_documents = all([not files.get(file_type) for file_type in files])
+
+    if files is None:
+        logger.warning(f"🚨 No files to process for {search_type} search.")
+        return True
+
+    try:
+        # Initialize Org Notes Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Org.value) and files["org"]:
+            logger.info("🦄 Setting up search for orgmode notes")
+            # Extract Entries, Generate Notes Embeddings
+            text_search.setup(
+                OrgToEntries,
+                files.get("org"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup org: {e}", exc_info=True)
+        success = False
+
+    try:
+        # Initialize Markdown Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Markdown.value) and files[
+            "markdown"
+        ]:
+            logger.info("💎 Setting up search for markdown notes")
+            # Extract Entries, Generate Markdown Embeddings
+            text_search.setup(
+                MarkdownToEntries,
+                files.get("markdown"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup markdown: {e}", exc_info=True)
+        success = False
+
+    try:
+        # Initialize PDF Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Pdf.value) and files["pdf"]:
+            logger.info("🖨️ Setting up search for pdf")
+            # Extract Entries, Generate PDF Embeddings
+            text_search.setup(
+                PdfToEntries,
+                files.get("pdf"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup PDF: {e}", exc_info=True)
+        success = False
+
+    try:
+        # Initialize Plaintext Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Plaintext.value) and files[
+            "plaintext"
+        ]:
+            logger.info("📄 Setting up search for plaintext")
+            # Extract Entries, Generate Plaintext Embeddings
+            text_search.setup(
+                PlaintextToEntries,
+                files.get("plaintext"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup plaintext: {e}", exc_info=True)
+        success = False
+
+    try:
+        if no_documents:
+            github_config = GithubConfig.objects.filter(user=user).prefetch_related("githubrepoconfig").first()
+            if (
+                search_type == state.SearchType.All.value or search_type == state.SearchType.Github.value
+            ) and github_config is not None:
+                logger.info("🐙 Setting up search for github")
+                # Extract Entries, Generate Github Embeddings
+                text_search.setup(
+                    GithubToEntries,
+                    None,
+                    regenerate=regenerate,
+                    full_corpus=full_corpus,
+                    user=user,
+                    config=github_config,
+                )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup GitHub: {e}", exc_info=True)
+        success = False
+
+    try:
+        if no_documents:
+            # Initialize Notion Search
+            notion_config = NotionConfig.objects.filter(user=user).first()
+            if (
+                search_type == state.SearchType.All.value or search_type == state.SearchType.Notion.value
+            ) and notion_config:
+                logger.info("🔌 Setting up search for notion")
+                text_search.setup(
+                    NotionToEntries,
+                    None,
+                    regenerate=regenerate,
+                    full_corpus=full_corpus,
+                    user=user,
+                    config=notion_config,
+                )
+
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup Notion: {e}", exc_info=True)
+        success = False
+
+    try:
+        # Initialize Image Search
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Image.value) and files[
+            "image"
+        ]:
+            logger.info("🖼️ Setting up search for images")
+            # Extract Entries, Generate Image Embeddings
+            text_search.setup(
+                ImageToEntries,
+                files.get("image"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup images: {e}", exc_info=True)
+        success = False
+    try:
+        if (search_type == state.SearchType.All.value or search_type == state.SearchType.Docx.value) and files["docx"]:
+            logger.info("📄 Setting up search for docx")
+            text_search.setup(
+                DocxToEntries,
+                files.get("docx"),
+                regenerate=regenerate,
+                full_corpus=full_corpus,
+                user=user,
+            )
+    except Exception as e:
+        logger.error(f"🚨 Failed to setup docx: {e}", exc_info=True)
+        success = False
+
+    # Invalidate Query Cache
+    if user:
+        state.query_cache[user.uuid] = LRU()
+
+    return success

khoj/routers/notion.py

@@ -0,0 +1,91 @@
+import asyncio
+import base64
+import json
+import logging
+import os
+from concurrent.futures import ThreadPoolExecutor
+
+import requests
+from fastapi import APIRouter, BackgroundTasks, Request, Response
+from starlette.responses import RedirectResponse
+
+from khoj.database.adapters import aget_user_by_uuid
+from khoj.database.models import KhojUser, NotionConfig
+from khoj.routers.indexer import configure_content
+from khoj.utils.state import SearchType
+
+NOTION_OAUTH_CLIENT_ID = os.getenv("NOTION_OAUTH_CLIENT_ID")
+NOTION_OAUTH_CLIENT_SECRET = os.getenv("NOTION_OAUTH_CLIENT_SECRET")
+NOTION_REDIRECT_URI = os.getenv("NOTION_REDIRECT_URI")
+
+notion_router = APIRouter()
+
+executor = ThreadPoolExecutor()
+
+logger = logging.getLogger(__name__)
+
+
+def get_notion_auth_url(user: KhojUser):
+    if not NOTION_OAUTH_CLIENT_ID or not NOTION_OAUTH_CLIENT_SECRET or not NOTION_REDIRECT_URI:
+        return None
+    return f"https://api.notion.com/v1/oauth/authorize?client_id={NOTION_OAUTH_CLIENT_ID}&redirect_uri={NOTION_REDIRECT_URI}&response_type=code&state={user.uuid}"
+
+
+async def run_in_executor(func, *args):
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(executor, func, *args)
+
+
+@notion_router.get("/auth/callback")
+async def notion_auth_callback(request: Request, background_tasks: BackgroundTasks):
+    code = request.query_params.get("code")
+    state = request.query_params.get("state")
+    if not code or not state:
+        return Response("Missing code or state", status_code=400)
+
+    user: KhojUser = await aget_user_by_uuid(state)
+
+    await NotionConfig.objects.filter(user=user).adelete()
+
+    if not user:
+        raise Exception("User not found")
+
+    bearer_token = f"{NOTION_OAUTH_CLIENT_ID}:{NOTION_OAUTH_CLIENT_SECRET}"
+    base64_encoded_token = base64.b64encode(bearer_token.encode()).decode()
+
+    headers = {
+        "Accept": "application/json",
+        "Content-Type": "application/json",
+        "Authorization": f"Basic {base64_encoded_token}",
+    }
+
+    data = {
+        "grant_type": "authorization_code",
+        "code": code,
+        "redirect_uri": NOTION_REDIRECT_URI,
+    }
+
+    response = requests.post("https://api.notion.com/v1/oauth/token", data=json.dumps(data), headers=headers)
+
+    final_response = response.json()
+
+    logger.info(f"Notion auth callback response: {final_response}")
+
+    access_token = final_response.get("access_token")
+    await NotionConfig.objects.acreate(token=access_token, user=user)
+
+    owner = final_response.get("owner")
+    workspace_id = final_response.get("workspace_id")
+    workspace_name = final_response.get("workspace_name")
+    bot_id = final_response.get("bot_id")
+
+    logger.info(
+        f"Notion integration. Owner: {owner}, Workspace ID: {workspace_id}, Workspace Name: {workspace_name}, Bot ID: {bot_id}"
+    )
+
+    notion_redirect = str(request.app.url_path_for("notion_config_page"))
+
+    # Trigger an async job to configure_content. Let it run without blocking the response.
+    background_tasks.add_task(run_in_executor, configure_content, {}, False, SearchType.Notion, True, user)
+
+    return RedirectResponse(notion_redirect)

khoj/routers/storage.py

@@ -0,0 +1,35 @@
+import logging
+import os
+import uuid
+
+logger = logging.getLogger(__name__)
+
+AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY")
+AWS_SECRET_KEY = os.getenv("AWS_SECRET_KEY")
+# S3 supports serving assets via your domain. Khoj expects this to be used in production. To enable it:
+# 1. Your bucket name for images should be of the form sub.domain.tld. For example, generated.khoj.dev
+# 2. Add CNAME entry to your domain's DNS records pointing to the S3 bucket. For example, CNAME generated.khoj.dev generated-khoj-dev.s3.amazonaws.com
+AWS_UPLOAD_IMAGE_BUCKET_NAME = os.getenv("AWS_IMAGE_UPLOAD_BUCKET")
+
+aws_enabled = AWS_ACCESS_KEY is not None and AWS_SECRET_KEY is not None and AWS_UPLOAD_IMAGE_BUCKET_NAME is not None
+
+if aws_enabled:
+    from boto3 import client
+
+    s3_client = client("s3", aws_access_key_id=AWS_ACCESS_KEY, aws_secret_access_key=AWS_SECRET_KEY)
+
+
+def upload_image(image: bytes, user_id: uuid.UUID):
+    """Upload the image to the S3 bucket"""
+    if not aws_enabled:
+        logger.info("AWS is not enabled. Skipping image upload")
+        return None
+
+    image_key = f"{user_id}/{uuid.uuid4()}.webp"
+    try:
+        s3_client.put_object(Bucket=AWS_UPLOAD_IMAGE_BUCKET_NAME, Key=image_key, Body=image, ACL="public-read")
+        url = f"https://{AWS_UPLOAD_IMAGE_BUCKET_NAME}/{image_key}"
+        return url
+    except Exception as e:
+        logger.error(f"Failed to upload image to S3: {e}")
+        return None

khoj/routers/subscription.py

@@ -0,0 +1,104 @@
+import logging
+import os
+from datetime import datetime, timezone
+
+from asgiref.sync import sync_to_async
+from fastapi import APIRouter, Request
+from starlette.authentication import requires
+
+from khoj.database import adapters
+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"]
+
+    # 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 = 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)
+        # 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"]
+            updated_user = await adapters.set_user_subscription(customer_email, is_recurring=is_recurring)
+            success = updated_user is not None
+    elif event_type in {"customer.subscription.deleted"}:
+        # Reset the user to trial state
+        user = await adapters.set_user_subscription(
+            customer_email, is_recurring=False, renewal_date=False, type="trial"
+        )
+        success = user is not None
+
+    logger.info(f'Stripe subscription {event["type"]} for {customer_email}')
+    return {"success": success}
+
+
+@subscription_router.patch("")
+@requires(["authenticated"])
+async def update_subscription(request: Request, email: str, operation: str):
+    # Retrieve the customer's details
+    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"}

khoj/routers/twilio.py

@@ -0,0 +1,36 @@
+import logging
+import os
+
+from khoj.database.models import KhojUser
+
+logger = logging.getLogger(__name__)
+
+account_sid = os.getenv("TWILIO_ACCOUNT_SID")
+auth_token = os.getenv("TWILIO_AUTH_TOKEN")
+verification_service_sid = os.getenv("TWILIO_VERIFICATION_SID")
+
+twilio_enabled = account_sid is not None and auth_token is not None and verification_service_sid is not None
+if twilio_enabled:
+    from twilio.rest import Client
+
+    client = Client(account_sid, auth_token)
+
+
+def is_twilio_enabled():
+    return twilio_enabled
+
+
+def create_otp(user: KhojUser):
+    """Create a new OTP for the user"""
+    verification = client.verify.v2.services(verification_service_sid).verifications.create(
+        to=str(user.phone_number), channel="whatsapp"
+    )
+    return verification.sid is not None
+
+
+def verify_otp(user: KhojUser, code: str):
+    """Verify the OTP for the user"""
+    verification_check = client.verify.v2.services(verification_service_sid).verification_checks.create(
+        to=str(user.phone_number), code=code
+    )
+    return verification_check.status == "approved"