khoj 1.16.1.dev15__py3-none-any.whl

khoj/migrations/migrate_server_pg.py

@@ -0,0 +1,138 @@
+"""
+The application config currently looks like this:
+app:
+  should-log-telemetry: true
+content-type:
+    ...
+processor:
+  conversation:
+    conversation-logfile: ~/.khoj/processor/conversation/conversation_logs.json
+    max-prompt-size: null
+    offline-chat:
+      chat-model: mistral-7b-instruct-v0.1.Q4_0.gguf
+      enable-offline-chat: false
+    openai:
+      api-key: sk-blah
+      chat-model: gpt-3.5-turbo
+    tokenizer: null
+search-type:
+  asymmetric:
+    cross-encoder: cross-encoder/ms-marco-MiniLM-L-6-v2
+    encoder: sentence-transformers/multi-qa-MiniLM-L6-cos-v1
+    encoder-type: null
+    model-directory: /Users/si/.khoj/search/asymmetric
+  image:
+    encoder: sentence-transformers/clip-ViT-B-32
+    encoder-type: null
+    model-directory: /Users/si/.khoj/search/image
+  symmetric:
+    cross-encoder: cross-encoder/ms-marco-MiniLM-L-6-v2
+    encoder: sentence-transformers/all-MiniLM-L6-v2
+    encoder-type: null
+    model-directory: ~/.khoj/search/symmetric
+version: 0.14.0
+
+
+The new version will looks like this:
+app:
+  should-log-telemetry: true
+processor:
+  conversation:
+    offline-chat:
+      enabled: false
+    openai:
+      api-key: sk-blah
+    chat-model-options:
+      - chat-model: gpt-3.5-turbo
+        tokenizer: null
+        type: openai
+      - chat-model: mistral-7b-instruct-v0.1.Q4_0.gguf
+        tokenizer: null
+        type: offline
+search-type:
+  asymmetric:
+    cross-encoder: cross-encoder/ms-marco-MiniLM-L-6-v2
+    encoder: sentence-transformers/multi-qa-MiniLM-L6-cos-v1
+version: 0.15.0
+"""
+
+import logging
+
+from packaging import version
+
+from khoj.database.models import (
+    ChatModelOptions,
+    OpenAIProcessorConversationConfig,
+    SearchModelConfig,
+)
+from khoj.utils.yaml import load_config_from_file, save_config_to_file
+
+logger = logging.getLogger(__name__)
+
+
+def migrate_server_pg(args):
+    schema_version = "0.15.0"
+    raw_config = load_config_from_file(args.config_file)
+    previous_version = raw_config.get("version")
+
+    if previous_version is None or version.parse(previous_version) < version.parse(schema_version):
+        logger.info(
+            f"Migrating configuration used for version {previous_version} to latest version for server with postgres in {args.version_no}"
+        )
+        raw_config["version"] = schema_version
+
+        if raw_config is None:
+            return args
+
+        if "search-type" in raw_config and raw_config["search-type"]:
+            if "asymmetric" in raw_config["search-type"]:
+                # Delete all existing search models
+                SearchModelConfig.objects.filter(model_type=SearchModelConfig.ModelType.TEXT).delete()
+                # Create new search model from existing Khoj YAML config
+                asymmetric_search = raw_config["search-type"]["asymmetric"]
+                SearchModelConfig.objects.create(
+                    name="default",
+                    model_type=SearchModelConfig.ModelType.TEXT,
+                    bi_encoder=asymmetric_search.get("encoder"),
+                    cross_encoder=asymmetric_search.get("cross-encoder"),
+                )
+
+        if "processor" in raw_config and raw_config["processor"] and "conversation" in raw_config["processor"]:
+            processor_conversation = raw_config["processor"]["conversation"]
+
+            if "offline-chat" in raw_config["processor"]["conversation"]:
+                offline_chat = raw_config["processor"]["conversation"]["offline-chat"]
+                ChatModelOptions.objects.create(
+                    chat_model=offline_chat.get("chat-model"),
+                    tokenizer=processor_conversation.get("tokenizer"),
+                    max_prompt_size=processor_conversation.get("max-prompt-size"),
+                    model_type=ChatModelOptions.ModelType.OFFLINE,
+                )
+
+            if (
+                "openai" in raw_config["processor"]["conversation"]
+                and raw_config["processor"]["conversation"]["openai"]
+            ):
+                openai = raw_config["processor"]["conversation"]["openai"]
+
+                if openai.get("api-key") is None:
+                    logger.error("OpenAI API Key is not set. Will not be migrating OpenAI config.")
+                else:
+                    if openai.get("chat-model") is None:
+                        openai["chat-model"] = "gpt-3.5-turbo"
+
+                    openai_config = OpenAIProcessorConversationConfig.objects.create(
+                        api_key=openai.get("api-key"), name="default"
+                    )
+
+                    ChatModelOptions.objects.create(
+                        chat_model=openai.get("chat-model"),
+                        tokenizer=processor_conversation.get("tokenizer"),
+                        max_prompt_size=processor_conversation.get("max-prompt-size"),
+                        model_type=ChatModelOptions.ModelType.OPENAI,
+                        openai_config=openai_config,
+                    )
+
+        save_config_to_file(raw_config, args.config_file)
+
+    return args

khoj/migrations/migrate_version.py

@@ -0,0 +1,17 @@
+from khoj.utils.yaml import load_config_from_file, save_config_to_file
+
+
+def migrate_config_to_version(args):
+    schema_version = "0.9.0"
+    raw_config = load_config_from_file(args.config_file)
+
+    # Add version to khoj config schema
+    if "version" not in raw_config:
+        raw_config["version"] = schema_version
+        save_config_to_file(raw_config, args.config_file)
+
+        # regenerate khoj index on first start of this version
+        # this should refresh index and apply index corruption fixes from #325
+        args.regenerate = True
+
+    return args

khoj/processor/content/docx/docx_to_entries.py

@@ -0,0 +1,110 @@
+import logging
+import os
+from datetime import datetime
+from typing import Dict, List, Tuple
+
+from langchain_community.document_loaders import Docx2txtLoader
+
+from khoj.database.models import Entry as DbEntry
+from khoj.database.models import KhojUser
+from khoj.processor.content.text_to_entries import TextToEntries
+from khoj.utils.helpers import timer
+from khoj.utils.rawconfig import Entry
+
+logger = logging.getLogger(__name__)
+
+
+class DocxToEntries(TextToEntries):
+    def __init__(self):
+        super().__init__()
+
+    # Define Functions
+    def process(
+        self, files: dict[str, str] = None, full_corpus: bool = True, user: KhojUser = None, regenerate: bool = False
+    ) -> Tuple[int, int]:
+        # Extract required fields from config
+        if not full_corpus:
+            deletion_file_names = set([file for file in files if files[file] == b""])
+            files_to_process = set(files) - deletion_file_names
+            files = {file: files[file] for file in files_to_process}
+        else:
+            deletion_file_names = None
+
+        # Extract Entries from specified Docx files
+        with timer("Extract entries from specified DOCX files", logger):
+            file_to_text_map, current_entries = DocxToEntries.extract_docx_entries(files)
+
+        # Split entries by max tokens supported by model
+        with timer("Split entries by max token size supported by model", logger):
+            current_entries = self.split_entries_by_max_tokens(current_entries, max_tokens=256)
+
+        # Identify, mark and merge any new entries with previous entries
+        with timer("Identify new or updated entries", logger):
+            num_new_embeddings, num_deleted_embeddings = self.update_embeddings(
+                current_entries,
+                DbEntry.EntryType.DOCX,
+                DbEntry.EntrySource.COMPUTER,
+                "compiled",
+                logger,
+                deletion_file_names,
+                user,
+                regenerate=regenerate,
+                file_to_text_map=file_to_text_map,
+            )
+
+        return num_new_embeddings, num_deleted_embeddings
+
+    @staticmethod
+    def extract_docx_entries(docx_files) -> Tuple[Dict, List[Entry]]:
+        """Extract entries from specified DOCX files"""
+
+        entries: List[str] = []
+        entry_to_location_map: List[Tuple[str, str]] = []
+        file_to_text_map = dict()
+        for docx_file in docx_files:
+            try:
+                timestamp_now = datetime.utcnow().timestamp()
+                tmp_file = f"tmp_docx_file_{timestamp_now}.docx"
+                with open(tmp_file, "wb") as f:
+                    bytes_content = docx_files[docx_file]
+                    f.write(bytes_content)
+
+                # Load the content using Docx2txtLoader
+                loader = Docx2txtLoader(tmp_file)
+                docx_entries_per_file = loader.load()
+
+                # Convert the loaded entries into the desired format
+                docx_texts = [page.page_content for page in docx_entries_per_file]
+
+                entry_to_location_map += zip(docx_texts, [docx_file] * len(docx_texts))
+                entries.extend(docx_texts)
+                file_to_text_map[docx_file] = docx_texts
+            except Exception as e:
+                logger.warning(f"Unable to process file: {docx_file}. This file will not be indexed.")
+                logger.warning(e, exc_info=True)
+            finally:
+                if os.path.exists(f"{tmp_file}"):
+                    os.remove(f"{tmp_file}")
+        return file_to_text_map, DocxToEntries.convert_docx_entries_to_maps(entries, dict(entry_to_location_map))
+
+    @staticmethod
+    def convert_docx_entries_to_maps(parsed_entries: List[str], entry_to_file_map) -> List[Entry]:
+        """Convert each DOCX entry into a dictionary"""
+        entries = []
+        for parsed_entry in parsed_entries:
+            entry_filename = entry_to_file_map[parsed_entry]
+            # Append base filename to compiled entry for context to model
+            heading = f"{entry_filename}\n"
+            compiled_entry = f"{heading}{parsed_entry}"
+            entries.append(
+                Entry(
+                    compiled=compiled_entry,
+                    raw=parsed_entry,
+                    heading=heading,
+                    file=f"{entry_filename}",
+                )
+            )
+
+        logger.debug(f"Converted {len(parsed_entries)} DOCX entries to dictionaries")
+
+        return entries

khoj/processor/content/github/github_to_entries.py

@@ -0,0 +1,224 @@
+import logging
+import time
+from typing import Any, Dict, List, Tuple
+
+import requests
+from magika import Magika
+
+from khoj.database.models import Entry as DbEntry
+from khoj.database.models import GithubConfig, KhojUser
+from khoj.processor.content.markdown.markdown_to_entries import MarkdownToEntries
+from khoj.processor.content.org_mode.org_to_entries import OrgToEntries
+from khoj.processor.content.plaintext.plaintext_to_entries import PlaintextToEntries
+from khoj.processor.content.text_to_entries import TextToEntries
+from khoj.utils.helpers import timer
+from khoj.utils.rawconfig import GithubContentConfig, GithubRepoConfig
+
+logger = logging.getLogger(__name__)
+magika = Magika()
+
+
+class GithubToEntries(TextToEntries):
+    def __init__(self, config: GithubConfig):
+        super().__init__(config)
+        raw_repos = config.githubrepoconfig.all()
+        repos = []
+        for repo in raw_repos:
+            repos.append(
+                GithubRepoConfig(
+                    name=repo.name,
+                    owner=repo.owner,
+                    branch=repo.branch,
+                )
+            )
+        self.config = GithubContentConfig(
+            pat_token=config.pat_token,
+            repos=repos,
+        )
+        self.session = requests.Session()
+        self.session.headers.update({"Authorization": f"token {self.config.pat_token}"})
+
+    @staticmethod
+    def wait_for_rate_limit_reset(response, func, *args, **kwargs):
+        if response.status_code != 200 and response.headers.get("X-RateLimit-Remaining") == "0":
+            wait_time = int(response.headers.get("X-RateLimit-Reset")) - int(time.time())
+            logger.info(f"Github Rate limit reached. Waiting for {wait_time} seconds")
+            time.sleep(wait_time)
+            return func(*args, **kwargs)
+        else:
+            return
+
+    def process(
+        self, files: dict[str, str] = None, full_corpus: bool = True, user: KhojUser = None, regenerate: bool = False
+    ) -> Tuple[int, int]:
+        if self.config.pat_token is None or self.config.pat_token == "":
+            logger.error(f"Github PAT token is not set. Skipping github content")
+            raise ValueError("Github PAT token is not set. Skipping github content")
+        current_entries = []
+        for repo in self.config.repos:
+            current_entries += self.process_repo(repo)
+
+        return self.update_entries_with_ids(current_entries, user=user)
+
+    def process_repo(self, repo: GithubRepoConfig):
+        repo_url = f"https://api.github.com/repos/{repo.owner}/{repo.name}"
+        repo_shorthand = f"{repo.owner}/{repo.name}"
+        logger.info(f"Processing github repo {repo_shorthand}")
+        with timer("Download files from github repo", logger):
+            try:
+                markdown_files, org_files, plaintext_files = self.get_files(repo_url, repo)
+            except ConnectionAbortedError as e:
+                logger.error(f"Github rate limit reached. Skip indexing github repo {repo_shorthand}")
+                raise e
+            except Exception as e:
+                logger.error(f"Unable to download github repo {repo_shorthand}", exc_info=True)
+                raise e
+
+        logger.info(
+            f"Found {len(markdown_files)} md, {len(org_files)} org and {len(plaintext_files)} text files in github repo {repo_shorthand}"
+        )
+        current_entries = []
+
+        with timer(f"Extract markdown entries from github repo {repo_shorthand}", logger):
+            current_entries = MarkdownToEntries.convert_markdown_entries_to_maps(
+                *GithubToEntries.extract_markdown_entries(markdown_files)
+            )
+
+        with timer(f"Extract org entries from github repo {repo_shorthand}", logger):
+            current_entries += OrgToEntries.convert_org_nodes_to_entries(
+                *GithubToEntries.extract_org_entries(org_files)
+            )
+
+        with timer(f"Extract plaintext entries from github repo {repo_shorthand}", logger):
+            current_entries += PlaintextToEntries.convert_text_files_to_entries(
+                *GithubToEntries.extract_plaintext_entries(plaintext_files)
+            )
+
+        with timer(f"Split entries by max token size supported by model {repo_shorthand}", logger):
+            current_entries = TextToEntries.split_entries_by_max_tokens(current_entries, max_tokens=256)
+
+        return current_entries
+
+    def update_entries_with_ids(self, current_entries, user: KhojUser = None):
+        # Identify, mark and merge any new entries with previous entries
+        with timer("Identify new or updated entries", logger):
+            num_new_embeddings, num_deleted_embeddings = self.update_embeddings(
+                current_entries,
+                DbEntry.EntryType.GITHUB,
+                DbEntry.EntrySource.GITHUB,
+                key="compiled",
+                logger=logger,
+                user=user,
+            )
+
+        return num_new_embeddings, num_deleted_embeddings
+
+    def get_files(self, repo_url: str, repo: GithubRepoConfig):
+        # Get the contents of the repository
+        repo_content_url = f"{repo_url}/git/trees/{repo.branch}"
+        headers = {"Authorization": f"token {self.config.pat_token}"}
+        params = {"recursive": "true"}
+        response = requests.get(repo_content_url, headers=headers, params=params)
+        contents = response.json()
+
+        # Raise exception if hit rate limit
+        if response.status_code != 200 and response.headers.get("X-RateLimit-Remaining") == "0":
+            raise ConnectionAbortedError("Github rate limit reached")
+
+        # Extract markdown files from the repository
+        markdown_files: List[Dict[str, str]] = []
+        org_files: List[Dict[str, str]] = []
+        plaintext_files: List[Dict[str, str]] = []
+        if "tree" not in contents:
+            return markdown_files, org_files, plaintext_files
+
+        for item in contents["tree"]:
+            # Find all markdown files in the repository
+            if item["type"] == "blob" and item["path"].endswith(".md"):
+                # Create URL for each markdown file on Github
+                url_path = f'https://github.com/{repo.owner}/{repo.name}/blob/{repo.branch}/{item["path"]}'
+
+                # Add markdown file contents and URL to list
+                markdown_files += [{"content": self.get_file_contents(item["url"]), "path": url_path}]
+
+            # Find all org files in the repository
+            elif item["type"] == "blob" and item["path"].endswith(".org"):
+                # Create URL for each org file on Github
+                url_path = f'https://github.com/{repo.owner}/{repo.name}/blob/{repo.branch}/{item["path"]}'
+
+                # Add org file contents and URL to list
+                org_files += [{"content": self.get_file_contents(item["url"]), "path": url_path}]
+
+            # Find, index remaining non-binary files in the repository
+            elif item["type"] == "blob":
+                url_path = f'https://github.com/{repo.owner}/{repo.name}/blob/{repo.branch}/{item["path"]}'
+                content_bytes = self.get_file_contents(item["url"], decode=False)
+                content_type, content_str = None, None
+                try:
+                    content_type = magika.identify_bytes(content_bytes).output.group
+                except:
+                    logger.error(f"Unable to identify content type of file at {url_path}. Skip indexing it")
+                    continue
+
+                # Add non-binary file contents and URL to list
+                if content_type in ["text", "code"]:
+                    try:
+                        content_str = content_bytes.decode("utf-8")
+                    except:
+                        logger.error(f"Unable to decode content of file at {url_path}. Skip indexing it")
+                        continue
+                    plaintext_files += [{"content": content_str, "path": url_path}]
+
+        return markdown_files, org_files, plaintext_files
+
+    def get_file_contents(self, file_url, decode=True):
+        # Get text from each markdown file
+        headers = {"Accept": "application/vnd.github.v3.raw"}
+        response = self.session.get(file_url, headers=headers, stream=True)
+
+        # Stop indexing on hitting rate limit
+        if response.status_code != 200 and response.headers.get("X-RateLimit-Remaining") == "0":
+            raise ConnectionAbortedError("Github rate limit reached")
+
+        content = "" if decode else b""
+        for chunk in response.iter_content(chunk_size=2048):
+            if chunk:
+                try:
+                    content += chunk.decode("utf-8") if decode else chunk
+                except Exception as e:
+                    logger.error(f"Unable to decode chunk from {file_url}")
+                    logger.error(e)
+
+        return content
+
+    @staticmethod
+    def extract_markdown_entries(markdown_files):
+        entries = []
+        entry_to_file_map = []
+        for doc in markdown_files:
+            entries, entry_to_file_map = MarkdownToEntries.process_single_markdown_file(
+                doc["content"], doc["path"], entries, entry_to_file_map
+            )
+        return entries, dict(entry_to_file_map)
+
+    @staticmethod
+    def extract_org_entries(org_files):
+        entries = []
+        entry_to_file_map = []
+
+        for doc in org_files:
+            entries, entry_to_file_map = OrgToEntries.process_single_org_file(
+                doc["content"], doc["path"], entries, entry_to_file_map
+            )
+        return entries, dict(entry_to_file_map)
+
+    @staticmethod
+    def extract_plaintext_entries(plaintext_files):
+        entries = []
+        entry_to_file_map = []
+
+        for doc in plaintext_files:
+            entries, entry_to_file_map = PlaintextToEntries.process_single_plaintext_file(
+                doc["content"], doc["path"], entries, entry_to_file_map
+            )
+        return entries, dict(entry_to_file_map)

khoj/processor/content/images/image_to_entries.py

@@ -0,0 +1,118 @@
+import base64
+import logging
+import os
+from datetime import datetime
+from typing import Dict, List, Tuple
+
+from rapidocr_onnxruntime import RapidOCR
+
+from khoj.database.models import Entry as DbEntry
+from khoj.database.models import KhojUser
+from khoj.processor.content.text_to_entries import TextToEntries
+from khoj.utils.helpers import timer
+from khoj.utils.rawconfig import Entry
+
+logger = logging.getLogger(__name__)
+
+
+class ImageToEntries(TextToEntries):
+    def __init__(self):
+        super().__init__()
+
+    # Define Functions
+    def process(
+        self, files: dict[str, str] = None, full_corpus: bool = True, user: KhojUser = None, regenerate: bool = False
+    ) -> Tuple[int, int]:
+        # Extract required fields from config
+        if not full_corpus:
+            deletion_file_names = set([file for file in files if files[file] == b""])
+            files_to_process = set(files) - deletion_file_names
+            files = {file: files[file] for file in files_to_process}
+        else:
+            deletion_file_names = None
+
+        # Extract Entries from specified image files
+        with timer("Extract entries from specified Image files", logger):
+            file_to_text_map, current_entries = ImageToEntries.extract_image_entries(files)
+
+        # Split entries by max tokens supported by model
+        with timer("Split entries by max token size supported by model", logger):
+            current_entries = self.split_entries_by_max_tokens(current_entries, max_tokens=256)
+
+        # Identify, mark and merge any new entries with previous entries
+        with timer("Identify new or updated entries", logger):
+            num_new_embeddings, num_deleted_embeddings = self.update_embeddings(
+                current_entries,
+                DbEntry.EntryType.IMAGE,
+                DbEntry.EntrySource.COMPUTER,
+                "compiled",
+                logger,
+                deletion_file_names,
+                user,
+                regenerate=regenerate,
+                file_to_text_map=file_to_text_map,
+            )
+
+        return num_new_embeddings, num_deleted_embeddings
+
+    @staticmethod
+    def extract_image_entries(image_files) -> Tuple[Dict, List[Entry]]:  # important function
+        """Extract entries by page from specified image files"""
+        file_to_text_map = dict()
+        entries: List[str] = []
+        entry_to_location_map: List[Tuple[str, str]] = []
+        for image_file in image_files:
+            try:
+                loader = RapidOCR()
+                bytes = image_files[image_file]
+                # write the image to a temporary file
+                timestamp_now = datetime.utcnow().timestamp()
+                # use either png or jpg
+                if image_file.endswith(".png"):
+                    tmp_file = f"tmp_image_file_{timestamp_now}.png"
+                elif image_file.endswith(".jpg") or image_file.endswith(".jpeg"):
+                    tmp_file = f"tmp_image_file_{timestamp_now}.jpg"
+                with open(tmp_file, "wb") as f:
+                    bytes = image_files[image_file]
+                    f.write(bytes)
+                try:
+                    image_entries_per_file = ""
+                    result, _ = loader(tmp_file)
+                    if result:
+                        expanded_entries = [text[1] for text in result]
+                        image_entries_per_file = " ".join(expanded_entries)
+                except ImportError:
+                    logger.warning(f"Unable to process file: {image_file}. This file will not be indexed.")
+                    continue
+                entry_to_location_map.append((image_entries_per_file, image_file))
+                entries.extend([image_entries_per_file])
+                file_to_text_map[image_file] = image_entries_per_file
+            except Exception as e:
+                logger.warning(f"Unable to process file: {image_file}. This file will not be indexed.")
+                logger.warning(e, exc_info=True)
+            finally:
+                if os.path.exists(tmp_file):
+                    os.remove(tmp_file)
+        return file_to_text_map, ImageToEntries.convert_image_entries_to_maps(entries, dict(entry_to_location_map))
+
+    @staticmethod
+    def convert_image_entries_to_maps(parsed_entries: List[str], entry_to_file_map) -> List[Entry]:
+        "Convert each image entries into a dictionary"
+        entries = []
+        for parsed_entry in parsed_entries:
+            entry_filename = entry_to_file_map[parsed_entry]
+            # Append base filename to compiled entry for context to model
+            heading = f"{entry_filename}\n"
+            compiled_entry = f"{heading}{parsed_entry}"
+            entries.append(
+                Entry(
+                    compiled=compiled_entry,
+                    raw=parsed_entry,
+                    heading=heading,
+                    file=f"{entry_filename}",
+                )
+            )
+
+        logger.debug(f"Converted {len(parsed_entries)} image entries to dictionaries")
+
+        return entries