khoj 1.33.3.dev32__py3-none-any.whl

khoj/processor/content/notion/notion_to_entries.py

@@ -0,0 +1,259 @@
+import logging
+from enum import Enum
+from typing import Tuple
+
+import requests
+
+from khoj.database.models import Entry as DbEntry
+from khoj.database.models import KhojUser, NotionConfig
+from khoj.processor.content.text_to_entries import TextToEntries
+from khoj.utils.helpers import timer
+from khoj.utils.rawconfig import Entry, NotionContentConfig
+
+logger = logging.getLogger(__name__)
+
+
+class NotionBlockType(Enum):
+    PARAGRAPH = "paragraph"
+    HEADING_1 = "heading_1"
+    HEADING_2 = "heading_2"
+    HEADING_3 = "heading_3"
+    BULLETED_LIST_ITEM = "bulleted_list_item"
+    NUMBERED_LIST_ITEM = "numbered_list_item"
+    TO_DO = "to_do"
+    TOGGLE = "toggle"
+    CHILD_PAGE = "child_page"
+    UNSUPPORTED = "unsupported"
+    BOOKMARK = "bookmark"
+    DIVIDER = "divider"
+    PDF = "pdf"
+    IMAGE = "image"
+    EMBED = "embed"
+    VIDEO = "video"
+    FILE = "file"
+    SYNCED_BLOCK = "synced_block"
+    TABLE_OF_CONTENTS = "table_of_contents"
+    COLUMN = "column"
+    EQUATION = "equation"
+    LINK_PREVIEW = "link_preview"
+    COLUMN_LIST = "column_list"
+    QUOTE = "quote"
+    BREADCRUMB = "breadcrumb"
+    LINK_TO_PAGE = "link_to_page"
+    CHILD_DATABASE = "child_database"
+    TEMPLATE = "template"
+    CALLOUT = "callout"
+
+
+class NotionToEntries(TextToEntries):
+    def __init__(self, config: NotionConfig):
+        super().__init__(config)
+        self.config = NotionContentConfig(
+            token=config.token,
+        )
+        self.session = requests.Session()
+        if config.token:
+            self.session.headers.update({"Authorization": f"Bearer {config.token}", "Notion-Version": "2022-02-22"})
+        self.unsupported_block_types = [
+            NotionBlockType.BOOKMARK.value,
+            NotionBlockType.DIVIDER.value,
+            NotionBlockType.CHILD_DATABASE.value,
+            NotionBlockType.TEMPLATE.value,
+            NotionBlockType.CALLOUT.value,
+            NotionBlockType.UNSUPPORTED.value,
+        ]
+
+        self.display_block_block_types = [
+            NotionBlockType.PARAGRAPH.value,
+            NotionBlockType.HEADING_1.value,
+            NotionBlockType.HEADING_2.value,
+            NotionBlockType.HEADING_3.value,
+            NotionBlockType.BULLETED_LIST_ITEM.value,
+            NotionBlockType.NUMBERED_LIST_ITEM.value,
+            NotionBlockType.TO_DO.value,
+            NotionBlockType.TOGGLE.value,
+            NotionBlockType.CHILD_PAGE.value,
+            NotionBlockType.BOOKMARK.value,
+            NotionBlockType.DIVIDER.value,
+        ]
+
+        self.body_params = {"page_size": 100}
+
+    def process(self, files: dict[str, str], user: KhojUser, regenerate: bool = False) -> Tuple[int, int]:
+        current_entries = []
+
+        # Get all pages
+        with timer("Getting all pages via search endpoint", logger=logger):
+            responses = []
+
+            while True:
+                result = self.session.post(
+                    "https://api.notion.com/v1/search",
+                    json=self.body_params,
+                ).json()
+                responses.append(result)
+                if result.get("has_more", False) == False:
+                    break
+                else:
+                    self.body_params.update({"start_cursor": result["next_cursor"]})
+
+        for response in responses:
+            with timer("Processing response", logger=logger):
+                pages_or_databases = response.get("results", [])
+
+                # Get all pages content
+                for p_or_d in pages_or_databases:
+                    with timer(f"Processing {p_or_d['object']} {p_or_d['id']}", logger=logger):
+                        if p_or_d["object"] == "database":
+                            # TODO: Handle databases
+                            continue
+                        elif p_or_d["object"] == "page":
+                            page_entries = self.process_page(p_or_d)
+                            current_entries.extend(page_entries)
+
+        current_entries = TextToEntries.split_entries_by_max_tokens(current_entries, max_tokens=256)
+
+        return self.update_entries_with_ids(current_entries, user=user)
+
+    def process_page(self, page):
+        page_id = page["id"]
+        title, content = self.get_page_content(page_id)
+
+        if title == None or content == None:
+            return []
+
+        current_entries = []
+        curr_heading = ""
+        for block in content.get("results", []):
+            block_type = block.get("type")
+
+            if block_type == None:
+                continue
+            block_data = block[block_type]
+
+            if block_data.get("rich_text") == None or len(block_data["rich_text"]) == 0:
+                # There's no text to handle here.
+                continue
+
+            raw_content = ""
+            if block_type in ["heading_1", "heading_2", "heading_3"]:
+                # If the current block is a heading, we can consider the previous block processing completed.
+                # Add it as an entry and move on to processing the next chunk of the page.
+                if raw_content != "":
+                    current_entries.append(
+                        Entry(
+                            compiled=raw_content,
+                            raw=raw_content,
+                            heading=title,
+                            file=page["url"],
+                        )
+                    )
+                curr_heading = block_data["rich_text"][0]["plain_text"]
+            else:
+                if curr_heading != "":
+                    # Add the last known heading to the content for additional context
+                    raw_content = self.process_heading(curr_heading)
+            for text in block_data["rich_text"]:
+                raw_content += self.process_text(text)
+
+            if block.get("has_children", True):
+                raw_content += "\n"
+                raw_content = self.process_nested_children(
+                    self.get_block_children(block["id"]), raw_content, block_type
+                )
+
+            if raw_content != "":
+                current_entries.append(
+                    Entry(
+                        compiled=raw_content,
+                        raw=raw_content,
+                        heading=title,
+                        file=page["url"],
+                    )
+                )
+        return current_entries
+
+    def process_heading(self, heading):
+        return f"\n<b>{heading}</b>\n"
+
+    def process_nested_children(self, children, raw_content, block_type=None):
+        results = children.get("results", [])
+        for child in results:
+            child_type = child.get("type")
+            if child_type == None:
+                continue
+            child_data = child[child_type]
+            if child_data.get("rich_text") and len(child_data["rich_text"]) > 0:
+                for text in child_data["rich_text"]:
+                    raw_content += self.process_text(text, block_type)
+            if child_data.get("has_children", True):
+                return self.process_nested_children(self.get_block_children(child["id"]), raw_content, block_type)
+
+        return raw_content
+
+    def process_text(self, text, block_type=None):
+        text_type = text.get("type", None)
+        if text_type in self.unsupported_block_types:
+            return ""
+        if text.get("href", None):
+            return f"<a href='{text['href']}'>{text['plain_text']}</a>"
+        raw_text = text["plain_text"]
+        if text_type in self.display_block_block_types or block_type in self.display_block_block_types:
+            return f"\n{raw_text}\n"
+        return raw_text
+
+    def get_block_children(self, block_id):
+        try:
+            return self.session.get(f"https://api.notion.com/v1/blocks/{block_id}/children").json()
+        except Exception as e:
+            logger.error(f"Error getting children for block {block_id}: {e}")
+            return {}
+
+    def get_page(self, page_id):
+        return self.session.get(f"https://api.notion.com/v1/pages/{page_id}").json()
+
+    def get_page_children(self, page_id):
+        return self.session.get(f"https://api.notion.com/v1/blocks/{page_id}/children").json()
+
+    def get_page_content(self, page_id):
+        try:
+            page = self.get_page(page_id)
+            content = self.get_page_children(page_id)
+        except Exception as e:
+            logger.error(f"Error getting page {page_id}: {e}", exc_info=True)
+            return None, None
+        properties = page.get("properties", {})
+
+        title_field = "title"
+        if "Title" in properties:
+            title_field = "Title"
+        elif "Name" in properties:
+            title_field = "Name"
+        elif "Page" in properties:
+            title_field = "Page"
+        elif "Event" in properties:
+            title_field = "Event"
+        elif title_field not in properties:
+            logger.debug(f"Title field not found for page {page_id}. Setting title as None...")
+            title = None
+            return title, content
+        try:
+            title = page["properties"][title_field]["title"][0]["text"]["content"]
+        except Exception as e:
+            logger.warning(f"Error getting title for page {page_id}: {e}. Setting title as None...")
+            title = None
+        return title, content
+
+    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(
+                user,
+                current_entries,
+                DbEntry.EntryType.NOTION,
+                DbEntry.EntrySource.NOTION,
+                key="compiled",
+                logger=logger,
+            )
+
+        return num_new_embeddings, num_deleted_embeddings

khoj/processor/content/org_mode/org_to_entries.py

@@ -0,0 +1,226 @@
+import logging
+import re
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+from khoj.database.models import Entry as DbEntry
+from khoj.database.models import KhojUser
+from khoj.processor.content.org_mode import orgnode
+from khoj.processor.content.org_mode.orgnode import Orgnode
+from khoj.processor.content.text_to_entries import TextToEntries
+from khoj.utils import state
+from khoj.utils.helpers import timer
+from khoj.utils.rawconfig import Entry
+
+logger = logging.getLogger(__name__)
+
+
+class OrgToEntries(TextToEntries):
+    def __init__(self):
+        super().__init__()
+
+    # Define Functions
+    def process(self, files: dict[str, str], user: KhojUser, regenerate: bool = False) -> Tuple[int, int]:
+        deletion_file_names = set([file for file in files if files[file] == ""])
+        files_to_process = set(files) - deletion_file_names
+        files = {file: files[file] for file in files_to_process}
+
+        # Extract Entries from specified Org files
+        max_tokens = 256
+        with timer("Extract entries from specified Org files", logger):
+            file_to_text_map, current_entries = self.extract_org_entries(files, max_tokens=max_tokens)
+
+        with timer("Split entries by max token size supported by model", logger):
+            current_entries = self.split_entries_by_max_tokens(current_entries, max_tokens=max_tokens)
+
+        # 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(
+                user,
+                current_entries,
+                DbEntry.EntryType.ORG,
+                DbEntry.EntrySource.COMPUTER,
+                "compiled",
+                logger,
+                deletion_file_names,
+                regenerate=regenerate,
+                file_to_text_map=file_to_text_map,
+            )
+
+        return num_new_embeddings, num_deleted_embeddings
+
+    @staticmethod
+    def extract_org_entries(
+        org_files: dict[str, str], index_heading_entries: bool = False, max_tokens=256
+    ) -> Tuple[Dict, List[Entry]]:
+        "Extract entries from specified Org files"
+        file_to_text_map, entries, entry_to_file_map = OrgToEntries.extract_org_nodes(org_files, max_tokens)
+        return file_to_text_map, OrgToEntries.convert_org_nodes_to_entries(
+            entries, entry_to_file_map, index_heading_entries
+        )
+
+    @staticmethod
+    def extract_org_nodes(
+        org_files: dict[str, str], max_tokens
+    ) -> Tuple[Dict, List[List[Orgnode]], Dict[Orgnode, str]]:
+        "Extract org nodes from specified org files"
+        entries: List[List[Orgnode]] = []
+        entry_to_file_map: List[Tuple[Orgnode, str]] = []
+        file_to_text_map = {}
+        for org_file in org_files:
+            try:
+                org_content = org_files[org_file]
+                entries, entry_to_file_map = OrgToEntries.process_single_org_file(
+                    org_content, org_file, entries, entry_to_file_map, max_tokens
+                )
+                file_to_text_map[org_file] = org_content
+            except Exception as e:
+                logger.error(f"Unable to process file: {org_file}. Skipped indexing it.\nError; {e}", exc_info=True)
+
+        return file_to_text_map, entries, dict(entry_to_file_map)
+
+    @staticmethod
+    def process_single_org_file(
+        org_content: str,
+        org_file: str,
+        entries: List[List[Orgnode]],
+        entry_to_file_map: List[Tuple[Orgnode, str]],
+        max_tokens=256,
+        ancestry: Dict[int, str] = {},
+    ) -> Tuple[List[List[Orgnode]], List[Tuple[Orgnode, str]]]:
+        """Parse org_content from org_file into OrgNode entries
+
+        Recurse down org file entries, one heading level at a time,
+        until reach a leaf entry or the current entry tree fits max_tokens.
+
+        Parse recursion terminating entry (trees) into (a list of) OrgNode objects.
+        """
+        # Prepend the org section's heading ancestry
+        ancestry_string = "\n".join([f"{'*' * key} {ancestry[key]}" for key in sorted(ancestry.keys())])
+        org_content_with_ancestry = f"{ancestry_string}{org_content}"
+
+        # If content is small or content has no children headings, save it as a single entry
+        # Note: This is the terminating condition for this recursive function
+        if len(TextToEntries.tokenizer(org_content_with_ancestry)) <= max_tokens or not re.search(
+            rf"^\*{{{len(ancestry)+1},}}\s", org_content, re.MULTILINE
+        ):
+            orgnode_content_with_ancestry = orgnode.makelist(org_content_with_ancestry, org_file)
+            entry_to_file_map += zip(orgnode_content_with_ancestry, [org_file] * len(orgnode_content_with_ancestry))
+            entries.extend([orgnode_content_with_ancestry])
+            return entries, entry_to_file_map
+
+        # Split this entry tree into sections by the next heading level in it
+        # Increment heading level until able to split entry into sections or reach max heading level
+        # A successful split will result in at least 2 sections
+        max_heading_level = 100
+        next_heading_level = len(ancestry)
+        sections: List[str] = []
+        while len(sections) < 2 and next_heading_level < max_heading_level:
+            next_heading_level += 1
+            sections = re.split(rf"(\n|^)(?=[*]{{{next_heading_level}}} .+\n?)", org_content, flags=re.MULTILINE)
+
+        # If unable to split entry into sections, log error and skip indexing it
+        if next_heading_level == max_heading_level:
+            logger.error(f"Unable to split current entry chunk: {org_content_with_ancestry[:20]}. Skip indexing it.")
+            return entries, entry_to_file_map
+
+        # Recurse down each non-empty section after parsing its body, heading and ancestry
+        for section in sections:
+            # Skip empty sections
+            if section.strip() == "":
+                continue
+
+            # Extract the section body and (when present) the heading
+            current_ancestry = ancestry.copy()
+            first_non_empty_line = [line for line in section.split("\n") if line.strip() != ""][0]
+            # If first non-empty line is a heading with expected heading level
+            if re.search(rf"^\*{{{next_heading_level}}}\s", first_non_empty_line):
+                # Extract the section body without the heading
+                current_section_body = "\n".join(section.split(first_non_empty_line, 1)[1:])
+                # Parse the section heading into current section ancestry
+                current_section_title = first_non_empty_line[next_heading_level:].strip()
+                current_ancestry[next_heading_level] = current_section_title
+            # Else process the section as just body text
+            else:
+                current_section_body = section
+
+            # Recurse down children of the current entry
+            OrgToEntries.process_single_org_file(
+                current_section_body,
+                org_file,
+                entries,
+                entry_to_file_map,
+                max_tokens,
+                current_ancestry,
+            )
+
+        return entries, entry_to_file_map
+
+    @staticmethod
+    def convert_org_nodes_to_entries(
+        parsed_entries: List[List[Orgnode]],
+        entry_to_file_map: Dict[Orgnode, str],
+        index_heading_entries: bool = False,
+    ) -> List[Entry]:
+        """
+        Convert OrgNode lists into list of Entry objects
+
+        Each list of OrgNodes is a parsed parent org tree or leaf node.
+        Convert each list of these OrgNodes into a single Entry.
+        """
+        entries: List[Entry] = []
+        for entry_group in parsed_entries:
+            entry_heading, entry_compiled, entry_raw = "", "", ""
+            for parsed_entry in entry_group:
+                if not parsed_entry.hasBody and not index_heading_entries:
+                    # Ignore title notes i.e notes with just headings and empty body
+                    continue
+
+                todo_str = f"{parsed_entry.todo} " if parsed_entry.todo else ""
+
+                # Set base level to current org-node tree's root heading level
+                if not entry_heading and parsed_entry.level > 0:
+                    base_level = parsed_entry.level
+                # Indent entry by 1 heading level as ancestry is prepended as top level heading
+                heading = f"{'*' * (parsed_entry.level-base_level+2)} {todo_str}" if parsed_entry.level > 0 else ""
+                if parsed_entry.heading:
+                    heading += f"{parsed_entry.heading}."
+
+                # Prepend ancestor headings, filename as top heading to root parent entry for context
+                # Children nodes do not need ancestors trail as root parent node will have it
+                if not entry_heading:
+                    ancestors_trail = " / ".join(parsed_entry.ancestors) or Path(entry_to_file_map[parsed_entry])
+                    heading = f"* {ancestors_trail}\n{heading}" if heading else f"* {ancestors_trail}."
+
+                compiled = heading
+
+                if parsed_entry.tags:
+                    tags_str = " ".join(parsed_entry.tags)
+                    compiled += f"\t {tags_str}."
+
+                if parsed_entry.closed:
+                    compiled += f'\n Closed on {parsed_entry.closed.strftime("%Y-%m-%d")}.'
+
+                if parsed_entry.scheduled:
+                    compiled += f'\n Scheduled for {parsed_entry.scheduled.strftime("%Y-%m-%d")}.'
+
+                if parsed_entry.hasBody:
+                    compiled += f"\n {parsed_entry.body}"
+
+                # Add the sub-entry contents to the entry
+                entry_compiled += compiled
+                entry_raw += f"{parsed_entry}"
+                if not entry_heading:
+                    entry_heading = heading
+
+            if entry_compiled:
+                entries.append(
+                    Entry(
+                        compiled=entry_compiled,
+                        raw=entry_raw,
+                        heading=entry_heading,
+                        file=entry_to_file_map[parsed_entry],
+                    )
+                )
+
+        return entries