add async content processing

open_notebook/graphs/content_processing/__init__.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any, Dict
 
 import magic
 from langgraph.graph import END, START, StateGraph
@@ -21,7 +22,7 @@ from open_notebook.graphs.content_processing.video import extract_best_audio_fro
 from open_notebook.graphs.content_processing.youtube import extract_youtube_transcript
 
 
-def source_identification(state: ContentState):
+async def source_identification(state: ContentState) -> Dict[str, str]:
     """
     Identify the content source based on parameters
     """
@@ -37,7 +38,7 @@ def source_identification(state: ContentState):
     return {"source_type": doc_type}
 
 
-def file_type(state: ContentState):
+async def file_type(state: ContentState) -> Dict[str, Any]:
     """
     Identify the file using python-magic
     """
@@ -49,7 +50,7 @@ def file_type(state: ContentState):
     return return_dict
 
 
-def file_type_edge(data: ContentState):
+async def file_type_edge(data: ContentState) -> str:
     assert data.get("identified_type"), "Type not identified"
     identified_type = data["identified_type"]
 
@@ -69,7 +70,7 @@ def file_type_edge(data: ContentState):
         )
 
 
-def delete_file(data: ContentState):
+async def delete_file(data: ContentState) -> Dict[str, Any]:
     if data.get("delete_source"):
         logger.debug(f"Deleting file: {data.get('file_path')}")
         file_path = data.get("file_path")
@@ -81,9 +82,21 @@ def delete_file(data: ContentState):
                 logger.warning(f"File not found while trying to delete: {file_path}")
     else:
         logger.debug("Not deleting file")
+    return {}
 
 
+async def url_type_router(x: ContentState) -> str:
+    return x.get("identified_type", "")
+
+
+async def source_type_router(x: ContentState) -> str:
+    return x.get("source_type", "")
+
+
+# Create workflow
 workflow = StateGraph(ContentState)
+
+# Add nodes
 workflow.add_node("source", source_identification)
 workflow.add_node("url_provider", url_provider)
 workflow.add_node("file_type", file_type)
@@ -95,10 +108,12 @@ workflow.add_node("extract_best_audio_from_video", extract_best_audio_from_video
 workflow.add_node("extract_audio", extract_audio)
 workflow.add_node("extract_youtube_transcript", extract_youtube_transcript)
 workflow.add_node("delete_file", delete_file)
+
+# Add edges
 workflow.add_edge(START, "source")
 workflow.add_conditional_edges(
     "source",
-    lambda x: x.get("source_type"),
+    source_type_router,
     {
         "url": "url_provider",
         "file": "file_type",
@@ -111,7 +126,7 @@ workflow.add_conditional_edges(
 )
 workflow.add_conditional_edges(
     "url_provider",
-    lambda x: x.get("identified_type"),
+    url_type_router,
     {"article": "extract_url", "youtube": "extract_youtube_transcript"},
 )
 workflow.add_edge("url_provider", END)
@@ -125,4 +140,6 @@ workflow.add_edge("extract_office_content", "delete_file")
 workflow.add_edge("extract_best_audio_from_video", "extract_audio")
 workflow.add_edge("extract_audio", "delete_file")
 workflow.add_edge("delete_file", END)
+
+# Compile graph
 graph = workflow.compile()

open_notebook/graphs/content_processing/audio.py

@@ -1,4 +1,6 @@
+import asyncio
 import os
+from functools import partial
 from math import ceil
 
 from loguru import logger
@@ -11,45 +13,35 @@ from open_notebook.graphs.content_processing.state import ContentState
 # future: parallelize the transcription process
 
 
-def split_audio(input_file, segment_length_minutes=15, output_prefix=None):
+async def split_audio(input_file, segment_length_minutes=15, output_prefix=None):
     """
-    Split an audio file into segments of specified length.
+    Split an audio file into segments asynchronously.
-
-    Args:
-        input_file (str): Path to the input audio file
-        segment_length_minutes (int): Length of each segment in minutes
-        output_dir (str): Directory to save the segments (defaults to input file's directory)
-        output_prefix (str): Prefix for output files (defaults to input filename)
-
-    Returns:
-        list: List of paths to the created segment files
     """
+
+    def _split(input_file, segment_length_minutes, output_prefix):
         # Convert input file to absolute path
-    input_file = os.path.abspath(input_file)
+        input_file_abs = os.path.abspath(input_file)
-
+        output_dir = os.path.dirname(input_file_abs)
-    output_dir = os.path.dirname(input_file)
         os.makedirs(output_dir, exist_ok=True)
 
         # Set up output prefix
         if output_prefix is None:
-        output_prefix = os.path.splitext(os.path.basename(input_file))[0]
+            output_prefix = os.path.splitext(os.path.basename(input_file_abs))[0]
 
         # Load the audio file
-    audio = AudioSegment.from_file(input_file)
+        audio = AudioSegment.from_file(input_file_abs)
 
         # Calculate segment length in milliseconds
         segment_length_ms = segment_length_minutes * 60 * 1000
 
         # Calculate number of segments
         total_segments = ceil(len(audio) / segment_length_ms)
-    logger.debug(f"Splitting file: {input_file} into {total_segments} segments")
+        logger.debug(f"Splitting file: {input_file_abs} into {total_segments} segments")
 
-    # List to store output file paths
         output_files = []
 
         # Split the audio into segments
         for i in range(total_segments):
-        # Calculate start and end times for this segment
             start_time = i * segment_length_ms
             end_time = min((i + 1) * segment_length_ms, len(audio))
 
@@ -57,44 +49,66 @@ def split_audio(input_file, segment_length_minutes=15, output_prefix=None):
             segment = audio[start_time:end_time]
 
             # Generate output filename
-        # Format: prefix_001.mp3 (padding with zeros ensures correct ordering)
             output_filename = f"{output_prefix}_{str(i+1).zfill(3)}.mp3"
             output_path = os.path.join(output_dir, output_filename)
 
             # Export segment
             segment.export(output_path, format="mp3")
-
             output_files.append(output_path)
 
-        # Optional progress indication
             logger.debug(f"Exported segment {i+1}/{total_segments}: {output_filename}")
 
         return output_files
 
+    # Run CPU-bound audio processing in thread pool
+    return await asyncio.get_event_loop().run_in_executor(
+        None, partial(_split, input_file, segment_length_minutes, output_prefix)
+    )
 
-def extract_audio(data: ContentState):
+
+async def transcribe_audio_segment(audio_file, model):
+    """Transcribe a single audio segment asynchronously"""
+
+    def _transcribe(audio_file, model):
+        return model.transcribe(audio_file)
+
+    return await asyncio.get_event_loop().run_in_executor(
+        None, partial(_transcribe, audio_file, model)
+    )
+
+
+async def extract_audio(data: ContentState):
     SPEECH_TO_TEXT_MODEL = model_manager.speech_to_text
-
     input_audio_path = data.get("file_path")
     audio_files = []
 
     try:
-        audio_files = split_audio(input_audio_path)
+        # Split audio into segments
-        transcriptions = []
+        audio_files = await split_audio(input_audio_path)
 
-        for audio_file in audio_files:
+        # Transcribe all segments concurrently
-            transcriptions.append(SPEECH_TO_TEXT_MODEL.transcribe(audio_file))
+        transcribe_tasks = [
+            transcribe_audio_segment(audio_file, SPEECH_TO_TEXT_MODEL)
+            for audio_file in audio_files
+        ]
+        transcriptions = await asyncio.gather(*transcribe_tasks)
 
         return {"content": " ".join(transcriptions)}
 
     except Exception as e:
         logger.error(f"Error transcribing audio: {str(e)}")
         logger.exception(e)
-        raise  # Re-raise the exception after logging
+        raise
 
     finally:
-        for file in audio_files:
+        # Clean up temporary files
+        def _cleanup(files):
+            for file in files:
                 try:
                     os.remove(file)
                 except OSError as e:
                     logger.error(f"Error removing temporary file {file}: {str(e)}")
+
+        await asyncio.get_event_loop().run_in_executor(
+            None, partial(_cleanup, audio_files)
+        )

open_notebook/graphs/content_processing/office.py

@@ -1,3 +1,6 @@
+import asyncio
+from functools import partial
+
 from docx import Document
 from loguru import logger
 from openpyxl import load_workbook
@@ -12,7 +15,10 @@ SUPPORTED_OFFICE_TYPES = [
 ]
 
 
-def extract_docx_content_detailed(file_path):
+async def extract_docx_content_detailed(file_path):
+    """Extract content from DOCX file"""
+
+    def _extract():
         try:
             doc = Document(file_path)
             content = []
@@ -89,9 +95,13 @@ def extract_docx_content_detailed(file_path):
             logger.error(f"Failed to extract DOCX content: {e}")
             return None
 
+    return await asyncio.get_event_loop().run_in_executor(None, _extract)
 
-# Example of usage with metadata
+
-def get_docx_info(file_path):
+async def get_docx_info(file_path):
+    """Get DOCX metadata and content"""
+
+    def _get_info():
         try:
             doc = Document(file_path)
 
@@ -127,8 +137,13 @@ def get_docx_info(file_path):
             logger.error(f"Failed to get DOCX info: {e}")
             return None
 
+    return await asyncio.get_event_loop().run_in_executor(None, _get_info)
 
-def extract_pptx_content(file_path):
+
+async def extract_pptx_content(file_path):
+    """Extract content from PPTX file"""
+
+    def _extract():
         try:
             prs = Presentation(file_path)
             content = []
@@ -143,7 +158,9 @@ def extract_pptx_content(file_path):
                 # Extract text from all shapes
                 for shape in slide.shapes:
                     if hasattr(shape, "text") and shape.text.strip():
-                    if shape != slide.shapes.title:  # Skip title as it's already added
+                        if (
+                            shape != slide.shapes.title
+                        ):  # Skip title as it's already added
                             content.append(shape.text.strip())
 
             return "\n\n".join(content)
@@ -152,8 +169,13 @@ def extract_pptx_content(file_path):
             logger.error(f"Failed to extract PPTX content: {e}")
             return None
 
+    return await asyncio.get_event_loop().run_in_executor(None, _extract)
 
-def extract_xlsx_content(file_path, max_rows=1000, max_cols=100):
+
+async def extract_xlsx_content(file_path, max_rows=10000, max_cols=100):
+    """Extract content from XLSX file"""
+
+    def _extract():
         try:
             wb = load_workbook(file_path, data_only=True)
             content = []
@@ -180,7 +202,9 @@ def extract_xlsx_content(file_path, max_rows=1000, max_cols=100):
                     row_data = []
                     for col in range(1, max_col + 1):
                         cell_value = ws.cell(row=row, column=col).value
-                    row_data.append(str(cell_value) if cell_value is not None else "")
+                        row_data.append(
+                            str(cell_value) if cell_value is not None else ""
+                        )
                     content.append("| " + " | ".join(row_data) + " |")
 
             return "\n".join(content)
@@ -189,8 +213,13 @@ def extract_xlsx_content(file_path, max_rows=1000, max_cols=100):
             logger.error(f"Failed to extract XLSX content: {e}")
             return None
 
+    return await asyncio.get_event_loop().run_in_executor(None, partial(_extract))
 
-def get_pptx_info(file_path):
+
+async def get_pptx_info(file_path):
+    """Get PPTX metadata and content"""
+
+    def _get_info():
         try:
             prs = Presentation(file_path)
 
@@ -219,8 +248,13 @@ def get_pptx_info(file_path):
             logger.error(f"Failed to get PPTX info: {e}")
             return None
 
+    return await asyncio.get_event_loop().run_in_executor(None, _get_info)
 
-def get_xlsx_info(file_path):
+
+async def get_xlsx_info(file_path):
+    """Get XLSX metadata and content"""
+
+    def _get_info():
         try:
             wb = load_workbook(file_path, data_only=True)
 
@@ -250,14 +284,15 @@ def get_xlsx_info(file_path):
             logger.error(f"Failed to get XLSX info: {e}")
             return None
 
+    return await asyncio.get_event_loop().run_in_executor(None, _get_info)
 
-def extract_office_content(state: ContentState):
+
+async def extract_office_content(state: ContentState):
     """Universal function to extract content from Office files"""
     assert state.get("file_path"), "No file path provided"
     assert (
         state.get("identified_type") in SUPPORTED_OFFICE_TYPES
     ), "Unsupported File Type"
-
     file_path = state["file_path"]
     doc_type = state["identified_type"]
 
@@ -266,24 +301,23 @@ def extract_office_content(state: ContentState):
         == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
     ):
         logger.debug("Extracting content from DOCX file")
-        content = extract_docx_content_detailed(file_path)
+        content = await extract_docx_content_detailed(file_path)
-        info = get_docx_info(file_path)
+        info = await get_docx_info(file_path)
     elif (
         doc_type
         == "application/vnd.openxmlformats-officedocument.presentationml.presentation"
     ):
         logger.debug("Extracting content from PPTX file")
-        content = extract_pptx_content(file_path)
+        content = await extract_pptx_content(file_path)
-        info = get_pptx_info(file_path)
+        info = await get_pptx_info(file_path)
     elif (
         doc_type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
     ):
         logger.debug("Extracting content from XLSX file")
-        content = extract_xlsx_content(file_path)
+        content = await extract_xlsx_content(file_path)
-        info = get_xlsx_info(file_path)
+        info = await get_xlsx_info(file_path)
     else:
         raise Exception(f"Unsupported file format: {doc_type}")
 
     del info["content"]
-
     return {"content": content, "metadata": info}

open_notebook/graphs/content_processing/pdf.py

@@ -1,3 +1,4 @@
+import asyncio
 import re
 import unicodedata
 
@@ -114,7 +115,7 @@ def clean_pdf_text(text):
     return text.strip()
 
 
-def _extract_text_from_pdf(pdf_path):
+async def _extract_text_from_pdf(pdf_path):
     doc = fitz.open(pdf_path)
     try:
         text = ""
@@ -127,20 +128,39 @@ def _extract_text_from_pdf(pdf_path):
         doc.close()
 
 
-def extract_pdf(state: ContentState):
+async def _extract_text_from_pdf(pdf_path):
+    """Extract text from PDF asynchronously"""
+
+    def _extract():
+        doc = fitz.open(pdf_path)
+        try:
+            text = ""
+            logger.debug(f"Found {len(doc)} pages in PDF")
+            for page in doc:
+                text += page.get_text()
+            return clean_pdf_text(text)
+        finally:
+            doc.close()
+
+    # Run CPU-bound PDF processing in a thread pool
+    return await asyncio.get_event_loop().run_in_executor(None, _extract)
+
+
+async def extract_pdf(state: ContentState):
     """
-    Parse the text file and print its content.
+    Parse the PDF file and extract its content asynchronously.
     """
     return_dict = {}
     assert state.get("file_path"), "No file path provided"
     assert state.get("identified_type") in SUPPORTED_FITZ_TYPES, "Unsupported File Type"
+
     if (
         state.get("file_path") is not None
         and state.get("identified_type") in SUPPORTED_FITZ_TYPES
     ):
         file_path = state.get("file_path")
         try:
-            text = _extract_text_from_pdf(file_path)
+            text = await _extract_text_from_pdf(file_path)
             return_dict["content"] = text
         except FileNotFoundError:
             raise FileNotFoundError(f"File not found at {file_path}")

open_notebook/graphs/content_processing/text.py

@@ -1,11 +1,13 @@
+import asyncio
+
 from loguru import logger
 
 from open_notebook.graphs.content_processing.state import ContentState
 
 
-def extract_txt(state: ContentState):
+async def extract_txt(state: ContentState):
     """
-    Parse the text file and print its content.
+    Parse the text file and extract its content asynchronously.
     """
     return_dict = {}
     if (
@@ -14,12 +16,22 @@ def extract_txt(state: ContentState):
     ):
         logger.debug(f"Extracting text from {state.get('file_path')}")
         file_path = state.get("file_path")
+
         if file_path is not None:
             try:
+
+                def _read_file():
                     with open(file_path, "r", encoding="utf-8") as file:
-                    content = file.read()
+                        return file.read()
+
+                # Run file I/O in thread pool
+                content = await asyncio.get_event_loop().run_in_executor(
+                    None, _read_file
+                )
+
                 logger.debug(f"Extracted: {content[:100]}")
                 return_dict["content"] = content
+
             except FileNotFoundError:
                 raise FileNotFoundError(f"File not found at {file_path}")
             except Exception as e:

open_notebook/graphs/content_processing/url.py

@@ -1,7 +1,7 @@
 import re
 from urllib.parse import urlparse
 
-import requests  # type: ignore
+import aiohttp
 from bs4 import BeautifulSoup, Comment
 from loguru import logger
 
@@ -29,7 +29,7 @@ def url_provider(state: ContentState):
     return return_dict
 
 
-def extract_url_bs4(url: str):
+async def extract_url_bs4(url: str):
     """
     Get the title and content of a URL using bs4
     """
@@ -42,9 +42,10 @@ def extract_url_bs4(url: str):
         if url.startswith("<!DOCTYPE html>") or url.startswith("<html"):
             html_content = url
         else:
-            response = requests.get(url, headers=headers, timeout=10)
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url, headers=headers, timeout=10) as response:
                     response.raise_for_status()
-            html_content = response.text
+                    html_content = await response.text()
 
         soup = BeautifulSoup(html_content, "html.parser")
 
@@ -143,7 +144,7 @@ def extract_url_bs4(url: str):
             "url": url if not url.startswith("<!DOCTYPE html>") else None,
         }
 
-    except requests.exceptions.RequestException as e:
+    except aiohttp.ClientError as e:
         logger.error(f"Failed to fetch URL {url}: {e}")
         return None
     except Exception as e:
@@ -151,12 +152,13 @@ def extract_url_bs4(url: str):
         return None
 
 
-def extract_url_jina(url: str):
+async def extract_url_jina(url: str):
     """
     Get the content of a URL using Jina
     """
-    response = requests.get(f"https://r.jina.ai/{url}")
+    async with aiohttp.ClientSession() as session:
-    text = response.text
+        async with session.get(f"https://r.jina.ai/{url}") as response:
+            text = await response.text()
             if text.startswith("Title:") and "\n" in text:
                 title_end = text.index("\n")
                 title = text[6:title_end].strip()
@@ -166,23 +168,22 @@ def extract_url_jina(url: str):
                 )
                 return {"title": title, "content": content}
             else:
-        content = text
                 logger.debug(
-            f"Processed url: {url}, does not have Title prefix, returning full content: {content[:100]}..."
+                    f"Processed url: {url}, does not have Title prefix, returning full content: {text[:100]}..."
                 )
                 return {"content": text}
 
 
-def extract_url(state: ContentState):
+async def extract_url(state: ContentState):
     assert state.get("url"), "No URL provided"
     url = state["url"]
     try:
-        result = extract_url_bs4(url)
+        result = await extract_url_bs4(url)
         if not result or not result.get("content"):
             logger.debug(
                 f"BS4 extraction failed for url {url}, falling back to Jina extractor"
             )
-            result = extract_url_jina(url)
+            result = await extract_url_jina(url)
         return result
     except Exception as e:
         logger.error(f"URL extraction failed for URL: {url}")

open_notebook/graphs/content_processing/video.py

@@ -1,16 +1,20 @@
+import asyncio
 import json
 import os
 import subprocess
+from functools import partial
 
 from loguru import logger
 
 from open_notebook.graphs.content_processing.state import ContentState
 
 
-def extract_audio_from_video(input_file, output_file, stream_index):
+async def extract_audio_from_video(input_file, output_file, stream_index):
     """
-    Extract the specified audio stream to MP3 format
+    Extract the specified audio stream to MP3 format asynchronously
     """
+
+    def _extract(input_file, output_file, stream_index):
         try:
             cmd = [
                 "ffmpeg",
@@ -33,17 +37,22 @@ def extract_audio_from_video(input_file, output_file, stream_index):
             return True
 
         except Exception as e:
-        print(f"Error extracting audio: {str(e)}")
+            logger.error(f"Error extracting audio: {str(e)}")
             return False
 
+    return await asyncio.get_event_loop().run_in_executor(
+        None, partial(_extract, input_file, output_file, stream_index)
+    )
 
-def get_audio_streams(input_file):
+
+async def get_audio_streams(input_file):
     """
-    Analyze video file and return information about all audio streams
+    Analyze video file and return information about all audio streams asynchronously
     """
+
+    def _analyze(input_file):
         logger.debug(f"Analyzing video file {input_file} for audio streams")
         try:
-        # Get stream information in JSON format
             cmd = [
                 "ffprobe",
                 "-v",
@@ -64,14 +73,20 @@ def get_audio_streams(input_file):
             return data.get("streams", [])
 
         except Exception as e:
-        print(f"Error analyzing file: {str(e)}")
+            logger.error(f"Error analyzing file: {str(e)}")
             return []
 
+    return await asyncio.get_event_loop().run_in_executor(
+        None, partial(_analyze, input_file)
+    )
 
-def select_best_audio_stream(streams):
+
+async def select_best_audio_stream(streams):
     """
     Select the best audio stream based on various quality metrics
     """
+
+    def _select(streams):
         if not streams:
             logger.debug("No audio streams found")
             return None
@@ -101,14 +116,26 @@ def select_best_audio_stream(streams):
         # Return the stream with highest score
         return max(scored_streams, key=lambda x: x[0])[1]
 
+    return await asyncio.get_event_loop().run_in_executor(
+        None, partial(_select, streams)
+    )
 
-def extract_best_audio_from_video(data: ContentState):
+
+async def extract_best_audio_from_video(data: ContentState):
     """
-    Main function to extract the best audio stream from a video file
+    Main function to extract the best audio stream from a video file asynchronously
     """
     input_file = data.get("file_path")
     assert input_file is not None, "Input file path must be provided"
-    if not os.path.exists(input_file):
+
+    def _check_file(path):
+        return os.path.exists(path)
+
+    file_exists = await asyncio.get_event_loop().run_in_executor(
+        None, partial(_check_file, input_file)
+    )
+
+    if not file_exists:
         logger.critical(f"Input file not found: {input_file}")
         return False
 
@@ -116,20 +143,20 @@ def extract_best_audio_from_video(data: ContentState):
     output_file = f"{base_name}_audio.mp3"
 
     # Get all audio streams
-    streams = get_audio_streams(input_file)
+    streams = await get_audio_streams(input_file)
     if not streams:
         logger.debug("No audio streams found in the file")
         return False
 
     # Select best stream
-    best_stream = select_best_audio_stream(streams)
+    best_stream = await select_best_audio_stream(streams)
     if not best_stream:
         logger.error("Could not determine best audio stream")
         return False
 
     # Extract the selected stream
     stream_index = streams.index(best_stream)
-    success = extract_audio_from_video(input_file, output_file, stream_index)
+    success = await extract_audio_from_video(input_file, output_file, stream_index)
 
     if success:
         logger.debug(f"Successfully extracted audio to: {output_file}")

open_notebook/graphs/content_processing/youtube.py

@@ -1,7 +1,7 @@
 import re
 import ssl
 
-import requests
+import aiohttp
 from bs4 import BeautifulSoup
 from loguru import logger
 from youtube_transcript_api import YouTubeTranscriptApi  # type: ignore
@@ -14,11 +14,15 @@ from open_notebook.graphs.content_processing.state import ContentState
 ssl._create_default_https_context = ssl._create_unverified_context
 
 
-def get_video_title(video_id):
+async def get_video_title(video_id):
     try:
         url = f"https://www.youtube.com/watch?v={video_id}"
-        response = requests.get(url)
+        async with aiohttp.ClientSession() as session:
-        soup = BeautifulSoup(response.text, "html.parser")
+            async with session.get(url) as response:
+                html = await response.text()
+
+        # BeautifulSoup doesn't support async operations
+        soup = BeautifulSoup(html, "html.parser")
 
         # YouTube stores title in a meta tag
         title = soup.find("meta", property="og:title")["content"]
@@ -63,7 +67,7 @@ def _extract_youtube_id(url):
     return match.group(1) if match else None
 
 
-def get_best_transcript(video_id, preferred_langs=["en", "es", "pt"]):
+async def get_best_transcript(video_id, preferred_langs=["en", "es", "pt"]):
     try:
         transcript_list = YouTubeTranscriptApi.list_transcripts(video_id)
 
@@ -129,7 +133,7 @@ def get_best_transcript(video_id, preferred_langs=["en", "es", "pt"]):
         return None
 
 
-def extract_youtube_transcript(state: ContentState):
+async def extract_youtube_transcript(state: ContentState):
     """
     Parse the text file and print its content.
     """
@@ -139,12 +143,12 @@ def extract_youtube_transcript(state: ContentState):
     )
 
     video_id = _extract_youtube_id(state.get("url"))
-    transcript = get_best_transcript(video_id, languages)
+    transcript = await get_best_transcript(video_id, languages)
 
     logger.debug(f"Found transcript: {transcript}")
     formatter = TextFormatter()
     try:
-        title = get_video_title(video_id)
+        title = await get_video_title(video_id)
     except Exception as e:
         logger.critical(f"Failed to get video title for video_id: {video_id}")
         logger.exception(e)

open_notebook/graphs/source.py

@@ -16,8 +16,6 @@ from open_notebook.graphs.content_processing import graph as content_graph
 from open_notebook.graphs.multipattern import graph as transform_graph
 from open_notebook.utils import surreal_clean
 
-# todo: we can make this more efficient
-
 
 class SourceState(TypedDict):
     content_state: ContentState
@@ -32,20 +30,24 @@ class TransformationState(TypedDict):
     transformation: dict
 
 
-def content_process(state: SourceState):
+async def content_process(state: SourceState) -> dict:
     content_state = state["content_state"]
     logger.debug("Content processing started for new content")
-    return {"content_state": content_graph.invoke(content_state)}
+    processed_state = await content_graph.ainvoke(content_state)
+    return {"content_state": processed_state}
 
 
-def run_patterns(input_text, patterns):
+async def run_patterns(input_text: str, patterns: List[dict]) -> str:
-    output = transform_graph.invoke(dict(content_stack=[input_text], patterns=patterns))
+    output = await transform_graph.ainvoke(
+        dict(content_stack=[input_text], patterns=patterns)
+    )
     return output["output"]
 
 
-def save_source(state: SourceState):
+def save_source(state: SourceState) -> dict:
     logger.debug("Saving source")
     content_state = state["content_state"]
+
     source = Source(
         asset=Asset(
             url=content_state.get("url"), file_path=content_state.get("file_path")
@@ -61,9 +63,10 @@ def save_source(state: SourceState):
     return {"source": source}
 
 
-def trigger_transformations(state: SourceState, config: RunnableConfig):
+def trigger_transformations(state: SourceState, config: RunnableConfig) -> List[Send]:
     if len(state["transformations"]) == 0:
         return []
+
     transformations = Transformation.get_all()
     to_apply = [
         t
@@ -71,6 +74,7 @@ def trigger_transformations(state: SourceState, config: RunnableConfig):
         if t["name"] in state["transformations"]
     ]
     logger.debug(f"Applying transformations {to_apply}")
+
     return [
         Send(
             "transform_content",
@@ -83,24 +87,34 @@ def trigger_transformations(state: SourceState, config: RunnableConfig):
     ]
 
 
-def transform_content(state: TransformationState):
+async def transform_content(state: TransformationState) -> dict:
     source = state["source"]
     content = source.full_text
     transformation = state["transformation"]
+
     logger.debug(f"Applying transformation {transformation['name']}")
-    result = run_patterns(content, patterns=transformation["patterns"])
+    result = await run_patterns(content, patterns=transformation["patterns"])
+
     source.add_insight(transformation["name"], surreal_clean(result))
+
     return {"transformations": [{"name": transformation["name"], "content": result}]}
 
 
+# Create and compile the workflow
 workflow = StateGraph(SourceState)
+
+# Add nodes
 workflow.add_node("content_process", content_process)
 workflow.add_node("save_source", save_source)
 workflow.add_node("transform_content", transform_content)
+
+# Define the graph edges
 workflow.add_edge(START, "content_process")
 workflow.add_edge("content_process", "save_source")
 workflow.add_conditional_edges(
     "save_source", trigger_transformations, ["transform_content"]
 )
 workflow.add_edge("transform_content", END)
+
+# Compile the graph
 source_graph = workflow.compile()

pages/stream_app/source.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 from pathlib import Path
 
@@ -71,13 +72,15 @@ def add_source(notebook_id):
                         f.write(source_file.getbuffer())
 
                 st.write("Processing content..")
-                source_graph.invoke(
+                asyncio.run(
+                    source_graph.ainvoke(
                         {
                             "content_state": req,
                             "notebook_id": notebook_id,
                             "transformations": apply_transformations,
                         }
                     )
+                )
             except UnsupportedTypeException as e:
                 st.warning(
                     "This type of content is not supported yet. If you think it should be, let us know on the project Issues's page"