khoj 1.33.3.dev32__py3-none-any.whl

khoj/utils/initialization.py

@@ -0,0 +1,301 @@
+import logging
+import os
+from typing import Tuple
+
+import openai
+
+from khoj.database.adapters import ConversationAdapters
+from khoj.database.models import (
+    AiModelApi,
+    ChatModel,
+    KhojUser,
+    SpeechToTextModelOptions,
+    TextToImageModelConfig,
+)
+from khoj.processor.conversation.utils import model_to_prompt_size, model_to_tokenizer
+from khoj.utils.constants import (
+    default_anthropic_chat_models,
+    default_gemini_chat_models,
+    default_offline_chat_models,
+    default_openai_chat_models,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def initialization(interactive: bool = True):
+    def _create_admin_user():
+        logger.info(
+            "👩‍✈️ Setting up admin user. These credentials will allow you to configure your server at /server/admin."
+        )
+        if not interactive and (not os.getenv("KHOJ_ADMIN_EMAIL") or not os.getenv("KHOJ_ADMIN_PASSWORD")):
+            logger.error(
+                "🚨 Admin user cannot be created. Please set the KHOJ_ADMIN_EMAIL, KHOJ_ADMIN_PASSWORD environment variables or start server in interactive mode."
+            )
+            exit(1)
+        email_addr = os.getenv("KHOJ_ADMIN_EMAIL") or input("Email: ")
+        password = os.getenv("KHOJ_ADMIN_PASSWORD") or input("Password: ")
+        admin_user = KhojUser.objects.create_superuser(email=email_addr, username=email_addr, password=password)
+        logger.info(f"👩‍✈️ Created admin user: {admin_user.email}")
+
+    def _create_chat_configuration():
+        logger.info(
+            "🗣️ Configure chat models available to your server. You can always update these at /server/admin using your admin account"
+        )
+
+        openai_api_base = os.getenv("OPENAI_API_BASE")
+        provider = "Ollama" if openai_api_base and openai_api_base.endswith(":11434/v1/") else "OpenAI"
+        openai_api_key = os.getenv("OPENAI_API_KEY", "placeholder" if openai_api_base else None)
+        default_chat_models = default_openai_chat_models
+        if openai_api_base:
+            # Get available chat models from OpenAI compatible API
+            try:
+                openai_client = openai.OpenAI(api_key=openai_api_key, base_url=openai_api_base)
+                default_chat_models = [model.id for model in openai_client.models.list()]
+                # Put the available default OpenAI models at the top
+                valid_default_models = [model for model in default_openai_chat_models if model in default_chat_models]
+                other_available_models = [model for model in default_chat_models if model not in valid_default_models]
+                default_chat_models = valid_default_models + other_available_models
+            except Exception as e:
+                logger.warning(
+                    f"⚠️ Failed to fetch {provider} chat models. Fallback to default models. Error: {str(e)}"
+                )
+
+        # Set up OpenAI's online chat models
+        openai_configured, openai_provider = _setup_chat_model_provider(
+            ChatModel.ModelType.OPENAI,
+            default_chat_models,
+            default_api_key=openai_api_key,
+            api_base_url=openai_api_base,
+            vision_enabled=True,
+            is_offline=False,
+            interactive=interactive,
+            provider_name=provider,
+        )
+
+        # Setup OpenAI speech to text model
+        if openai_configured:
+            default_speech2text_model = "whisper-1"
+            if interactive:
+                openai_speech2text_model = input(
+                    f"Enter the OpenAI speech to text model you want to use (default: {default_speech2text_model}): "
+                )
+                openai_speech2text_model = openai_speech2text_model or default_speech2text_model
+            else:
+                openai_speech2text_model = default_speech2text_model
+            SpeechToTextModelOptions.objects.create(
+                model_name=openai_speech2text_model, model_type=SpeechToTextModelOptions.ModelType.OPENAI
+            )
+
+        # Setup OpenAI text to image model
+        if openai_configured:
+            default_text_to_image_model = "dall-e-3"
+            if interactive:
+                openai_text_to_image_model = input(
+                    f"Enter the OpenAI text to image model you want to use (default: {default_text_to_image_model}): "
+                )
+                openai_text_to_image_model = openai_text_to_image_model or default_text_to_image_model
+            else:
+                openai_text_to_image_model = default_text_to_image_model
+            TextToImageModelConfig.objects.create(
+                model_name=openai_text_to_image_model,
+                model_type=TextToImageModelConfig.ModelType.OPENAI,
+                ai_model_api=openai_provider,
+            )
+
+        # Set up Google's Gemini online chat models
+        _setup_chat_model_provider(
+            ChatModel.ModelType.GOOGLE,
+            default_gemini_chat_models,
+            default_api_key=os.getenv("GEMINI_API_KEY"),
+            vision_enabled=True,
+            is_offline=False,
+            interactive=interactive,
+            provider_name="Google Gemini",
+        )
+
+        # Set up Anthropic's online chat models
+        _setup_chat_model_provider(
+            ChatModel.ModelType.ANTHROPIC,
+            default_anthropic_chat_models,
+            default_api_key=os.getenv("ANTHROPIC_API_KEY"),
+            vision_enabled=True,
+            is_offline=False,
+            interactive=interactive,
+        )
+
+        # Set up offline chat models
+        _setup_chat_model_provider(
+            ChatModel.ModelType.OFFLINE,
+            default_offline_chat_models,
+            default_api_key=None,
+            vision_enabled=False,
+            is_offline=True,
+            interactive=interactive,
+        )
+
+        # Explicitly set default chat model
+        chat_models_configured = ChatModel.objects.count()
+        if chat_models_configured > 0:
+            default_chat_model_name = ChatModel.objects.first().name
+            # If there are multiple chat models, ask the user to choose the default chat model
+            if chat_models_configured > 1 and interactive:
+                user_chat_model_name = input(
+                    f"Enter the default chat model to use (default: {default_chat_model_name}): "
+                )
+            else:
+                user_chat_model_name = None
+
+            # If the user's choice is valid, set it as the default chat model
+            if user_chat_model_name and ChatModel.objects.filter(name=user_chat_model_name).exists():
+                default_chat_model_name = user_chat_model_name
+
+            logger.info("🗣️ Chat model configuration complete")
+
+        # Set up offline speech to text model
+        use_offline_speech2text_model = "n" if not interactive else input("Use offline speech to text model? (y/n): ")
+        if use_offline_speech2text_model == "y":
+            logger.info("🗣️ Setting up offline speech to text model")
+            # Delete any existing speech to text model options. There can only be one.
+            SpeechToTextModelOptions.objects.all().delete()
+
+            default_offline_speech2text_model = "base"
+            offline_speech2text_model = input(
+                f"Enter the Whisper model to use Offline (default: {default_offline_speech2text_model}): "
+            )
+            offline_speech2text_model = offline_speech2text_model or default_offline_speech2text_model
+            SpeechToTextModelOptions.objects.create(
+                model_name=offline_speech2text_model, model_type=SpeechToTextModelOptions.ModelType.OFFLINE
+            )
+
+            logger.info(f"🗣️  Offline speech to text model configured to {offline_speech2text_model}")
+
+    def _setup_chat_model_provider(
+        model_type: ChatModel.ModelType,
+        default_chat_models: list,
+        default_api_key: str,
+        interactive: bool,
+        api_base_url: str = None,
+        vision_enabled: bool = False,
+        is_offline: bool = False,
+        provider_name: str = None,
+    ) -> Tuple[bool, AiModelApi]:
+        supported_vision_models = (
+            default_openai_chat_models + default_anthropic_chat_models + default_gemini_chat_models
+        )
+        provider_name = provider_name or model_type.name.capitalize()
+        default_use_model = {True: "y", False: "n"}[default_api_key is not None or is_offline]
+        use_model_provider = (
+            default_use_model if not interactive else input(f"Add {provider_name} chat models? (y/n): ")
+        )
+
+        if use_model_provider != "y":
+            return False, None
+
+        logger.info(f"️💬 Setting up your {provider_name} chat configuration")
+
+        ai_model_api = None
+        if not is_offline:
+            if interactive:
+                user_api_key = input(f"Enter your {provider_name} API key (default: {default_api_key}): ")
+                api_key = user_api_key if user_api_key != "" else default_api_key
+            else:
+                api_key = default_api_key
+            ai_model_api = AiModelApi.objects.create(api_key=api_key, name=provider_name, api_base_url=api_base_url)
+
+        if interactive:
+            user_chat_models = input(
+                f"Enter the {provider_name} chat models you want to use (default: {','.join(default_chat_models)}): "
+            )
+            chat_models = user_chat_models.split(",") if user_chat_models != "" else default_chat_models
+            chat_models = [model.strip() for model in chat_models]
+        else:
+            chat_models = default_chat_models
+
+        for chat_model in chat_models:
+            default_max_tokens = model_to_prompt_size.get(chat_model)
+            default_tokenizer = model_to_tokenizer.get(chat_model)
+            vision_enabled = vision_enabled and chat_model in supported_vision_models
+
+            chat_model_options = {
+                "name": chat_model,
+                "model_type": model_type,
+                "max_prompt_size": default_max_tokens,
+                "vision_enabled": vision_enabled,
+                "tokenizer": default_tokenizer,
+                "ai_model_api": ai_model_api,
+            }
+
+            ChatModel.objects.create(**chat_model_options)
+
+        logger.info(f"🗣️ {provider_name} chat model configuration complete")
+        return True, ai_model_api
+
+    def _update_chat_model_options():
+        """Update available chat models for OpenAI-compatible APIs"""
+        try:
+            # Get OpenAI configs with custom base URLs
+            custom_configs = AiModelApi.objects.exclude(api_base_url__isnull=True)
+
+            # Only enable for whitelisted provider names (i.e Ollama) for now
+            # TODO: This is hacky. Will be replaced with more robust solution based on provider type enum
+            custom_configs = custom_configs.filter(name__in=["Ollama"])
+
+            for config in custom_configs:
+                try:
+                    # Create OpenAI client with custom base URL
+                    openai_client = openai.OpenAI(api_key=config.api_key, base_url=config.api_base_url)
+
+                    # Get available models
+                    available_models = [model.id for model in openai_client.models.list()]
+
+                    # Get existing chat model options for this config
+                    existing_models = ChatModel.objects.filter(
+                        ai_model_api=config, model_type=ChatModel.ModelType.OPENAI
+                    )
+
+                    # Add new models
+                    for model_name in available_models:
+                        if not existing_models.filter(name=model_name).exists():
+                            ChatModel.objects.create(
+                                name=model_name,
+                                model_type=ChatModel.ModelType.OPENAI,
+                                max_prompt_size=model_to_prompt_size.get(model_name),
+                                vision_enabled=model_name in default_openai_chat_models,
+                                tokenizer=model_to_tokenizer.get(model_name),
+                                ai_model_api=config,
+                            )
+
+                    # Remove models that are no longer available
+                    existing_models.exclude(name__in=available_models).delete()
+
+                except Exception as e:
+                    logger.warning(f"Failed to update models for {config.name}: {str(e)}")
+
+        except Exception as e:
+            logger.error(f"Failed to update chat model options: {str(e)}")
+
+    admin_user = KhojUser.objects.filter(is_staff=True).first()
+    if admin_user is None:
+        while True:
+            try:
+                _create_admin_user()
+                break
+            except Exception as e:
+                logger.error(f"🚨 Failed to create admin user: {e}", exc_info=True)
+
+    chat_config = ConversationAdapters.get_default_chat_model()
+    if admin_user is None and chat_config is None:
+        while True:
+            try:
+                _create_chat_configuration()
+                break
+            # Some environments don't support interactive input. We catch the exception and return if that's the case.
+            # The admin can still configure their settings from the admin page.
+            except EOFError:
+                return
+            except Exception as e:
+                logger.error(f"🚨 Failed to create chat configuration: {e}", exc_info=True)
+    else:
+        _update_chat_model_options()
+        logger.info("🗣️ Chat model configuration updated")

khoj/utils/jsonl.py

@@ -0,0 +1,43 @@
+import gzip
+import json
+import logging
+
+from khoj.utils.constants import empty_escape_sequences
+from khoj.utils.helpers import get_absolute_path
+
+logger = logging.getLogger(__name__)
+
+
+def load_jsonl(input_path):
+    "Read List of JSON objects from JSON line file"
+    # Initialize Variables
+    data = []
+    jsonl_file = None
+
+    # Open JSONL file
+    if input_path.suffix == ".gz":
+        jsonl_file = gzip.open(get_absolute_path(input_path), "rt", encoding="utf-8")
+    else:
+        jsonl_file = open(get_absolute_path(input_path), "r", encoding="utf-8")
+
+    # Read JSONL file
+    for line in jsonl_file:
+        data.append(json.loads(line.strip(empty_escape_sequences)))
+
+    # Close JSONL file
+    jsonl_file.close()
+
+    # Log JSONL entries loaded
+    logger.debug(f"Loaded {len(data)} records from {input_path}")
+
+    return data
+
+
+def compress_jsonl_data(jsonl_data, output_path):
+    # Create output directory, if it doesn't exist
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with gzip.open(output_path, "wt", encoding="utf-8") as gzip_file:
+        gzip_file.write(jsonl_data)
+
+    logger.debug(f"Wrote jsonl data to gzip compressed jsonl at {output_path}")

khoj/utils/models.py

@@ -0,0 +1,47 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+import openai
+import torch
+from tqdm import trange
+
+
+class BaseEncoder(ABC):
+    @abstractmethod
+    def __init__(self, model_name: str, device: torch.device = None, **kwargs):
+        ...
+
+    @abstractmethod
+    def encode(self, entries: List[str], device: torch.device = None, **kwargs) -> torch.Tensor:
+        ...
+
+
+class OpenAI(BaseEncoder):
+    def __init__(self, model_name, client: openai.OpenAI, device=None):
+        self.model_name = model_name
+        self.openai_client = client
+        self.embedding_dimensions = None
+
+    def encode(self, entries, device=None, **kwargs):
+        embedding_tensors = []
+
+        for index in trange(0, len(entries)):
+            # OpenAI models create better embeddings for entries without newlines
+            processed_entry = entries[index].replace("\n", " ")
+
+            try:
+                response = self.openai_client.embeddings.create(input=processed_entry, model=self.model_name)
+                embedding_tensors += [torch.tensor(response.data[0].embedding, device=device)]
+                # Use current models embedding dimension, once available
+                # Else default to embedding dimensions of the text-embedding-ada-002 model
+                self.embedding_dimensions = len(response.data[0].embedding) if not self.embedding_dimensions else 1536
+            except Exception as e:
+                print(
+                    f"Failed to encode entry {index} of length: {len(entries[index])}\n\n{entries[index][:1000]}...\n\n{e}"
+                )
+                # Use zero embedding vector for entries with failed embeddings
+                # This ensures entry embeddings match the order of the source entries
+                # And they have minimal similarity to other entries (as zero vectors are always orthogonal to other vector)
+                embedding_tensors += [torch.zeros(self.embedding_dimensions, device=device)]
+
+        return torch.stack(embedding_tensors)

khoj/utils/rawconfig.py

@@ -0,0 +1,208 @@
+# System Packages
+import json
+import uuid
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+from khoj.utils.helpers import to_snake_case_from_dash
+
+
+class ConfigBase(BaseModel):
+    class Config:
+        alias_generator = to_snake_case_from_dash
+        populate_by_name = True
+
+    def __getitem__(self, item):
+        return getattr(self, item)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+
+class LocationData(BaseModel):
+    city: Optional[str]
+    region: Optional[str]
+    country: Optional[str]
+    country_code: Optional[str]
+
+    def __str__(self):
+        parts = []
+        if self.city:
+            parts.append(self.city)
+        if self.region:
+            parts.append(self.region)
+        if self.country:
+            parts.append(self.country)
+        return ", ".join(parts)
+
+
+class FileFilterRequest(BaseModel):
+    filename: str
+    conversation_id: str
+
+
+class FilesFilterRequest(BaseModel):
+    filenames: List[str]
+    conversation_id: str
+
+
+class TextConfigBase(ConfigBase):
+    compressed_jsonl: Path
+    embeddings_file: Path
+
+
+class TextContentConfig(ConfigBase):
+    input_files: Optional[List[Path]] = None
+    input_filter: Optional[List[str]] = None
+    index_heading_entries: Optional[bool] = False
+
+
+class GithubRepoConfig(ConfigBase):
+    name: str
+    owner: str
+    branch: Optional[str] = "master"
+
+
+class GithubContentConfig(ConfigBase):
+    pat_token: Optional[str] = None
+    repos: List[GithubRepoConfig]
+
+
+class NotionContentConfig(ConfigBase):
+    token: str
+
+
+class ContentConfig(ConfigBase):
+    org: Optional[TextContentConfig] = None
+    markdown: Optional[TextContentConfig] = None
+    pdf: Optional[TextContentConfig] = None
+    plaintext: Optional[TextContentConfig] = None
+    github: Optional[GithubContentConfig] = None
+    notion: Optional[NotionContentConfig] = None
+    image: Optional[TextContentConfig] = None
+    docx: Optional[TextContentConfig] = None
+
+
+class ImageSearchConfig(ConfigBase):
+    encoder: str
+    encoder_type: Optional[str] = None
+    model_directory: Optional[Path] = None
+
+    class Config:
+        protected_namespaces = ()
+
+
+class SearchConfig(ConfigBase):
+    image: Optional[ImageSearchConfig] = None
+
+
+class OpenAIProcessorConfig(ConfigBase):
+    api_key: str
+    chat_model: Optional[str] = "gpt-4o-mini"
+
+
+class OfflineChatProcessorConfig(ConfigBase):
+    chat_model: Optional[str] = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF"
+
+
+class ConversationProcessorConfig(ConfigBase):
+    openai: Optional[OpenAIProcessorConfig] = None
+    offline_chat: Optional[OfflineChatProcessorConfig] = None
+    max_prompt_size: Optional[int] = None
+    tokenizer: Optional[str] = None
+
+
+class ProcessorConfig(ConfigBase):
+    conversation: Optional[ConversationProcessorConfig] = None
+
+
+class AppConfig(ConfigBase):
+    should_log_telemetry: bool = True
+
+
+class FullConfig(ConfigBase):
+    content_type: Optional[ContentConfig] = None
+    search_type: Optional[SearchConfig] = None
+    processor: Optional[ProcessorConfig] = None
+    app: Optional[AppConfig] = AppConfig()
+    version: Optional[str] = None
+
+
+class SearchResponse(ConfigBase):
+    entry: str
+    score: float
+    cross_score: Optional[float] = None
+    additional: Optional[dict] = None
+    corpus_id: str
+
+
+class FileData(BaseModel):
+    name: str
+    content: bytes
+    file_type: str
+    encoding: str | None = None
+
+
+class FileAttachment(BaseModel):
+    name: str
+    content: str
+    file_type: str
+    size: int
+
+
+class ChatRequestBody(BaseModel):
+    q: str
+    n: Optional[int] = 7
+    d: Optional[float] = None
+    stream: Optional[bool] = False
+    title: Optional[str] = None
+    conversation_id: Optional[str] = None
+    turn_id: Optional[str] = None
+    city: Optional[str] = None
+    region: Optional[str] = None
+    country: Optional[str] = None
+    country_code: Optional[str] = None
+    timezone: Optional[str] = None
+    images: Optional[list[str]] = None
+    files: Optional[list[FileAttachment]] = []
+    create_new: Optional[bool] = False
+
+
+class Entry:
+    raw: str
+    compiled: str
+    heading: Optional[str]
+    file: Optional[str]
+    corpus_id: str
+
+    def __init__(
+        self,
+        raw: str = None,
+        compiled: str = None,
+        heading: Optional[str] = None,
+        file: Optional[str] = None,
+        corpus_id: uuid.UUID = None,
+    ):
+        self.raw = raw
+        self.compiled = compiled
+        self.heading = heading
+        self.file = file
+        self.corpus_id = str(corpus_id)
+
+    def to_json(self) -> str:
+        return json.dumps(self.__dict__, ensure_ascii=False)
+
+    def __repr__(self) -> str:
+        return self.__dict__.__repr__()
+
+    @classmethod
+    def from_dict(cls, dictionary: dict):
+        return cls(
+            raw=dictionary["raw"],
+            compiled=dictionary["compiled"],
+            file=dictionary.get("file", None),
+            heading=dictionary.get("heading", None),
+            corpus_id=dictionary.get("corpus_id", None),
+        )

khoj/utils/state.py

@@ -0,0 +1,48 @@
+import os
+import threading
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, Dict, List
+
+from apscheduler.schedulers.background import BackgroundScheduler
+from openai import OpenAI
+from whisper import Whisper
+
+from khoj.database.models import ProcessLock
+from khoj.processor.embeddings import CrossEncoderModel, EmbeddingsModel
+from khoj.utils import config as utils_config
+from khoj.utils.config import OfflineChatProcessorModel, SearchModels
+from khoj.utils.helpers import LRU, get_device, is_env_var_true
+from khoj.utils.rawconfig import FullConfig
+
+# Application Global State
+config = FullConfig()
+search_models = SearchModels()
+embeddings_model: Dict[str, EmbeddingsModel] = None
+cross_encoder_model: Dict[str, CrossEncoderModel] = None
+openai_client: OpenAI = None
+offline_chat_processor_config: OfflineChatProcessorModel = None
+whisper_model: Whisper = None
+config_file: Path = None
+verbose: int = 0
+host: str = None
+port: int = None
+ssl_config: Dict[str, str] = None
+cli_args: List[str] = None
+query_cache: Dict[str, LRU] = defaultdict(LRU)
+chat_lock = threading.Lock()
+SearchType = utils_config.SearchType
+scheduler: BackgroundScheduler = None
+schedule_leader_process_lock: ProcessLock = None
+telemetry: List[Dict[str, str]] = []
+telemetry_disabled: bool = is_env_var_true("KHOJ_TELEMETRY_DISABLE")
+khoj_version: str = None
+device = get_device()
+chat_on_gpu: bool = True
+anonymous_mode: bool = False
+pretrained_tokenizers: Dict[str, Any] = dict()
+billing_enabled: bool = (
+    os.getenv("STRIPE_API_KEY") is not None
+    and os.getenv("STRIPE_SIGNING_SECRET") is not None
+    and os.getenv("KHOJ_CLOUD_SUBSCRIPTION_URL") is not None
+)

khoj/utils/yaml.py

@@ -0,0 +1,47 @@
+from pathlib import Path
+
+import yaml
+
+from khoj.utils import state
+from khoj.utils.rawconfig import FullConfig
+
+# Do not emit tags when dumping to YAML
+yaml.emitter.Emitter.process_tag = lambda self, *args, **kwargs: None  # type: ignore[assignment]
+
+
+def save_config_to_file_updated_state():
+    with open(state.config_file, "w") as outfile:
+        yaml.dump(yaml.safe_load(state.config.json(by_alias=True)), outfile)
+        outfile.close()
+    return state.config
+
+
+def save_config_to_file(yaml_config: dict, yaml_config_file: Path):
+    "Write config to YML file"
+    # Create output directory, if it doesn't exist
+    yaml_config_file.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(yaml_config_file, "w", encoding="utf-8") as config_file:
+        yaml.safe_dump(yaml_config, config_file, allow_unicode=True)
+
+
+def load_config_from_file(yaml_config_file: Path) -> dict:
+    "Read config from YML file"
+    config_from_file = None
+    with open(yaml_config_file, "r", encoding="utf-8") as config_file:
+        config_from_file = yaml.safe_load(config_file)
+    return config_from_file
+
+
+def parse_config_from_string(yaml_config: dict) -> FullConfig:
+    "Parse and validate config in YML string"
+    return FullConfig.model_validate(yaml_config)
+
+
+def parse_config_from_file(yaml_config_file):
+    "Parse and validate config in YML file"
+    return parse_config_from_string(load_config_from_file(yaml_config_file))
+
+
+def yaml_dump(data):
+    return yaml.dump(data, allow_unicode=True, sort_keys=False, default_flow_style=False)