admin-valentin/open-notebook
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add async content processing

Browse source
This commit is contained in:
LUIS NOVO 2024-11-11 17:32:35 -03:00
parent ac2ea9e554
commit 00f070a644
10 changed files with 541 additions and 395 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
open_notebook/graphs/content_processing/__init__.py
View file

@ -1,4 +1,5 @@
import os import os
from typing import Any, Dict
import magic import magic
from langgraph.graph import END, START, StateGraph 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 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 Identify the content source based on parameters
""" """
@ -37,7 +38,7 @@ def source_identification(state: ContentState):
return {"source_type": doc_type} 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 Identify the file using python-magic
""" """
@ -49,7 +50,7 @@ def file_type(state: ContentState):
return return_dict 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" assert data.get("identified_type"), "Type not identified"
identified_type = data["identified_type"] 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"): if data.get("delete_source"):
logger.debug(f"Deleting file: {data.get('file_path')}") logger.debug(f"Deleting file: {data.get('file_path')}")
file_path = 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}") logger.warning(f"File not found while trying to delete: {file_path}")
else: else:
logger.debug("Not deleting file") 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) workflow = StateGraph(ContentState)
# Add nodes
workflow.add_node("source", source_identification) workflow.add_node("source", source_identification)
workflow.add_node("url_provider", url_provider) workflow.add_node("url_provider", url_provider)
workflow.add_node("file_type", file_type) 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_audio", extract_audio)
workflow.add_node("extract_youtube_transcript", extract_youtube_transcript) workflow.add_node("extract_youtube_transcript", extract_youtube_transcript)
workflow.add_node("delete_file", delete_file) workflow.add_node("delete_file", delete_file)
# Add edges
workflow.add_edge(START, "source") workflow.add_edge(START, "source")
workflow.add_conditional_edges( workflow.add_conditional_edges(
"source", "source",
lambda x: x.get("source_type"), source_type_router,
{ {
"url": "url_provider", "url": "url_provider",
"file": "file_type", "file": "file_type",
@ -111,7 +126,7 @@ workflow.add_conditional_edges(
) )
workflow.add_conditional_edges( workflow.add_conditional_edges(
"url_provider", "url_provider",
lambda x: x.get("identified_type"), url_type_router,
{"article": "extract_url", "youtube": "extract_youtube_transcript"}, {"article": "extract_url", "youtube": "extract_youtube_transcript"},
) )
workflow.add_edge("url_provider", END) 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_best_audio_from_video", "extract_audio")
workflow.add_edge("extract_audio", "delete_file") workflow.add_edge("extract_audio", "delete_file")
workflow.add_edge("delete_file", END) workflow.add_edge("delete_file", END)
# Compile graph
graph = workflow.compile() graph = workflow.compile()

126
open_notebook/graphs/content_processing/audio.py
View file

@ -1,4 +1,6 @@
import asyncio
import os import os
from functools import partial
from math import ceil from math import ceil
from loguru import logger from loguru import logger
@ -11,90 +13,102 @@ from open_notebook.graphs.content_processing.state import ContentState
# future: parallelize the transcription process # 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
""" """
# Convert input file to absolute path
input_file = os.path.abspath(input_file)
output_dir = os.path.dirname(input_file) def _split(input_file, segment_length_minutes, output_prefix):
os.makedirs(output_dir, exist_ok=True) # Convert input file to absolute path
input_file_abs = os.path.abspath(input_file)
output_dir = os.path.dirname(input_file_abs)
os.makedirs(output_dir, exist_ok=True)
# Set up output prefix # Set up output prefix
if output_prefix is None: 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 # Load the audio file
audio = AudioSegment.from_file(input_file) audio = AudioSegment.from_file(input_file_abs)
# Calculate segment length in milliseconds # Calculate segment length in milliseconds
segment_length_ms = segment_length_minutes * 60 * 1000 segment_length_ms = segment_length_minutes * 60 * 1000
# Calculate number of segments # Calculate number of segments
total_segments = ceil(len(audio) / segment_length_ms) 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 = []
output_files = []
# Split the audio into segments # Split the audio into segments
for i in range(total_segments): for i in range(total_segments):
# Calculate start and end times for this segment start_time = i * segment_length_ms
start_time = i * segment_length_ms end_time = min((i + 1) * segment_length_ms, len(audio))
end_time = min((i + 1) * segment_length_ms, len(audio))
# Extract segment # Extract segment
segment = audio[start_time:end_time] segment = audio[start_time:end_time]
# Generate output filename # 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_filename = f"{output_prefix}_{str(i+1).zfill(3)}.mp3" output_path = os.path.join(output_dir, output_filename)
output_path = os.path.join(output_dir, output_filename)
# Export segment # Export segment
segment.export(output_path, format="mp3") segment.export(output_path, format="mp3")
output_files.append(output_path)
output_files.append(output_path) logger.debug(f"Exported segment {i+1}/{total_segments}: {output_filename}")
# Optional progress indication return output_files
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 SPEECH_TO_TEXT_MODEL = model_manager.speech_to_text
input_audio_path = data.get("file_path") input_audio_path = data.get("file_path")
audio_files = [] audio_files = []
try: 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)} return {"content": " ".join(transcriptions)}
except Exception as e: except Exception as e:
logger.error(f"Error transcribing audio: {str(e)}") logger.error(f"Error transcribing audio: {str(e)}")
logger.exception(e) logger.exception(e)
raise # Re-raise the exception after logging raise
finally: finally:
for file in audio_files: # Clean up temporary files
try: def _cleanup(files):
os.remove(file) for file in files:
except OSError as e: try:
logger.error(f"Error removing temporary file {file}: {str(e)}") 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)
)

428
open_notebook/graphs/content_processing/office.py
View file

@ -1,3 +1,6 @@
import asyncio
from functools import partial
from docx import Document from docx import Document
from loguru import logger from loguru import logger
from openpyxl import load_workbook from openpyxl import load_workbook
@ -12,252 +15,284 @@ SUPPORTED_OFFICE_TYPES = [
] ]
def extract_docx_content_detailed(file_path): async def extract_docx_content_detailed(file_path):
try: """Extract content from DOCX file"""
doc = Document(file_path)
content = []
for paragraph in doc.paragraphs: def _extract():
if not paragraph.text.strip(): try:
continue doc = Document(file_path)
content = []
style = paragraph.style.name if paragraph.style else "Normal" for paragraph in doc.paragraphs:
text = paragraph.text.strip() if not paragraph.text.strip():
continue
# Get paragraph formatting style = paragraph.style.name if paragraph.style else "Normal"
p_format = paragraph.paragraph_format text = paragraph.text.strip()
indent = p_format.left_indent or 0
# Convert indent to spaces (1 level = 4 spaces) # Get paragraph formatting
indent_level = 0 p_format = paragraph.paragraph_format
if hasattr(indent, "pt"): indent = p_format.left_indent or 0
indent_level = int(indent.pt / 72) # 72 points = 1 inch
indent_spaces = " " * (indent_level * 4)
# Handle different types of formatting # Convert indent to spaces (1 level = 4 spaces)
if "Heading" in style: indent_level = 0
level = style[-1] if style[-1].isdigit() else "1" if hasattr(indent, "pt"):
heading_marks = "#" * int(level) indent_level = int(indent.pt / 72) # 72 points = 1 inch
content.append(f"\n{heading_marks} {text}\n") indent_spaces = " " * (indent_level * 4)
# Handle bullet points # Handle different types of formatting
elif ( if "Heading" in style:
paragraph.style level = style[-1] if style[-1].isdigit() else "1"
and hasattr(paragraph.style, "name") heading_marks = "#" * int(level)
and paragraph.style.name.startswith("List") content.append(f"\n{heading_marks} {text}\n")
):
# Numbered list # Handle bullet points
if ( elif (
hasattr(paragraph._p, "pPr") paragraph.style
and paragraph._p.pPr is not None and hasattr(paragraph.style, "name")
and hasattr(paragraph._p.pPr, "numPr") and paragraph.style.name.startswith("List")
and paragraph._p.pPr.numPr is not None
): ):
# Try to get the actual number # Numbered list
try: if (
if ( hasattr(paragraph._p, "pPr")
hasattr(paragraph._p.pPr.numPr, "numId") and paragraph._p.pPr is not None
and paragraph._p.pPr.numPr.numId is not None and hasattr(paragraph._p.pPr, "numPr")
and hasattr(paragraph._p.pPr.numPr.numId, "val") and paragraph._p.pPr.numPr is not None
): ):
number = paragraph._p.pPr.numPr.numId.val # Try to get the actual number
content.append(f"{indent_spaces}{number}. {text}") try:
else: if (
hasattr(paragraph._p.pPr.numPr, "numId")
and paragraph._p.pPr.numPr.numId is not None
and hasattr(paragraph._p.pPr.numPr.numId, "val")
):
number = paragraph._p.pPr.numPr.numId.val
content.append(f"{indent_spaces}{number}. {text}")
else:
content.append(f"{indent_spaces}1. {text}")
except Exception:
content.append(f"{indent_spaces}1. {text}") content.append(f"{indent_spaces}1. {text}")
except Exception: # Bullet list
content.append(f"{indent_spaces}1. {text}")
# Bullet list
else:
content.append(f"{indent_spaces}* {text}")
else:
# Handle text formatting
formatted_text = []
for run in paragraph.runs:
if run.bold:
formatted_text.append(f"**{run.text}**")
elif run.italic:
formatted_text.append(f"*{run.text}*")
else: else:
formatted_text.append(run.text) content.append(f"{indent_spaces}* {text}")
content.append(f"{indent_spaces}{''.join(formatted_text)}") else:
# Handle text formatting
formatted_text = []
for run in paragraph.runs:
if run.bold:
formatted_text.append(f"**{run.text}**")
elif run.italic:
formatted_text.append(f"*{run.text}*")
else:
formatted_text.append(run.text)
return "\n\n".join(content) content.append(f"{indent_spaces}{''.join(formatted_text)}")
except Exception as e: return "\n\n".join(content)
logger.error(f"Failed to extract DOCX content: {e}")
return None except Exception as e:
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 async def get_docx_info(file_path):
def get_docx_info(file_path): """Get DOCX metadata and content"""
try:
doc = Document(file_path)
# Extract core properties if available def _get_info():
core_props = { try:
"author": doc.core_properties.author, doc = Document(file_path)
"created": doc.core_properties.created,
"modified": doc.core_properties.modified,
"title": doc.core_properties.title,
"subject": doc.core_properties.subject,
"keywords": doc.core_properties.keywords,
"category": doc.core_properties.category,
"comments": doc.core_properties.comments,
}
# Get document content # Extract core properties if available
content = extract_docx_content_detailed(file_path) core_props = {
"author": doc.core_properties.author,
"created": doc.core_properties.created,
"modified": doc.core_properties.modified,
"title": doc.core_properties.title,
"subject": doc.core_properties.subject,
"keywords": doc.core_properties.keywords,
"category": doc.core_properties.category,
"comments": doc.core_properties.comments,
}
# Get document statistics # Get document content
stats = { content = extract_docx_content_detailed(file_path)
"paragraph_count": len(doc.paragraphs),
"word_count": sum(
len(p.text.split()) for p in doc.paragraphs if p.text.strip()
),
"character_count": sum(
len(p.text) for p in doc.paragraphs if p.text.strip()
),
}
return {"metadata": core_props, "content": content, "statistics": stats} # Get document statistics
stats = {
"paragraph_count": len(doc.paragraphs),
"word_count": sum(
len(p.text.split()) for p in doc.paragraphs if p.text.strip()
),
"character_count": sum(
len(p.text) for p in doc.paragraphs if p.text.strip()
),
}
except Exception as e: return {"metadata": core_props, "content": content, "statistics": stats}
logger.error(f"Failed to get DOCX info: {e}")
return None except Exception as e:
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):
try: """Extract content from PPTX file"""
prs = Presentation(file_path)
content = []
for slide_number, slide in enumerate(prs.slides, 1): def _extract():
content.append(f"\n# Slide {slide_number}\n") try:
prs = Presentation(file_path)
content = []
# Extract title for slide_number, slide in enumerate(prs.slides, 1):
if slide.shapes.title: content.append(f"\n# Slide {slide_number}\n")
content.append(f"## {slide.shapes.title.text}\n")
# Extract text from all shapes # Extract title
for shape in slide.shapes: if slide.shapes.title:
if hasattr(shape, "text") and shape.text.strip(): content.append(f"## {slide.shapes.title.text}\n")
if shape != slide.shapes.title: # Skip title as it's already added
content.append(shape.text.strip())
return "\n\n".join(content) # 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
content.append(shape.text.strip())
except Exception as e: return "\n\n".join(content)
logger.error(f"Failed to extract PPTX content: {e}")
return None except Exception as e:
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):
try: """Extract content from XLSX file"""
wb = load_workbook(file_path, data_only=True)
content = []
for sheet in wb.sheetnames: def _extract():
ws = wb[sheet] try:
content.append(f"\n# Sheet: {sheet}\n") wb = load_workbook(file_path, data_only=True)
content = []
# Get the maximum row and column with data for sheet in wb.sheetnames:
max_row = min(ws.max_row, max_rows) ws = wb[sheet]
max_col = min(ws.max_column, max_cols) content.append(f"\n# Sheet: {sheet}\n")
# Create markdown table header # Get the maximum row and column with data
headers = [] max_row = min(ws.max_row, max_rows)
for col in range(1, max_col + 1): max_col = min(ws.max_column, max_cols)
cell_value = ws.cell(row=1, column=col).value
headers.append(str(cell_value) if cell_value is not None else "")
content.append("| " + " | ".join(headers) + " |") # Create markdown table header
content.append("| " + " | ".join(["---"] * len(headers)) + " |") headers = []
# Add table content
for row in range(2, max_row + 1):
row_data = []
for col in range(1, max_col + 1): for col in range(1, max_col + 1):
cell_value = ws.cell(row=row, column=col).value cell_value = ws.cell(row=1, column=col).value
row_data.append(str(cell_value) if cell_value is not None else "") headers.append(str(cell_value) if cell_value is not None else "")
content.append("| " + " | ".join(row_data) + " |")
return "\n".join(content) content.append("| " + " | ".join(headers) + " |")
content.append("| " + " | ".join(["---"] * len(headers)) + " |")
except Exception as e: # Add table content
logger.error(f"Failed to extract XLSX content: {e}") for row in range(2, max_row + 1):
return None 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 ""
)
content.append("| " + " | ".join(row_data) + " |")
return "\n".join(content)
except Exception as e:
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):
try: """Get PPTX metadata and content"""
prs = Presentation(file_path)
# Extract basic properties def _get_info():
props = { try:
"slide_count": len(prs.slides), prs = Presentation(file_path)
"title": "", # PowerPoint doesn't have built-in metadata like Word
}
# Get document content # Extract basic properties
content = extract_pptx_content(file_path) props = {
"slide_count": len(prs.slides),
"title": "", # PowerPoint doesn't have built-in metadata like Word
}
# Get presentation statistics # Get document content
stats = { content = extract_pptx_content(file_path)
"slide_count": len(prs.slides),
"shape_count": sum(len(slide.shapes) for slide in prs.slides),
"text_frame_count": sum(
sum(1 for shape in slide.shapes if hasattr(shape, "text"))
for slide in prs.slides
),
}
return {"metadata": props, "content": content, "statistics": stats} # Get presentation statistics
stats = {
"slide_count": len(prs.slides),
"shape_count": sum(len(slide.shapes) for slide in prs.slides),
"text_frame_count": sum(
sum(1 for shape in slide.shapes if hasattr(shape, "text"))
for slide in prs.slides
),
}
except Exception as e: return {"metadata": props, "content": content, "statistics": stats}
logger.error(f"Failed to get PPTX info: {e}")
return None except Exception as e:
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):
try: """Get XLSX metadata and content"""
wb = load_workbook(file_path, data_only=True)
# Extract basic properties def _get_info():
props = { try:
"sheet_count": len(wb.sheetnames), wb = load_workbook(file_path, data_only=True)
"sheets": wb.sheetnames,
"title": wb.properties.title,
"creator": wb.properties.creator,
"created": wb.properties.created,
"modified": wb.properties.modified,
}
# Get document content # Extract basic properties
content = extract_xlsx_content(file_path) props = {
"sheet_count": len(wb.sheetnames),
"sheets": wb.sheetnames,
"title": wb.properties.title,
"creator": wb.properties.creator,
"created": wb.properties.created,
"modified": wb.properties.modified,
}
# Get workbook statistics # Get document content
stats = { content = extract_xlsx_content(file_path)
"sheet_count": len(wb.sheetnames),
"total_rows": sum(sheet.max_row for sheet in wb.worksheets),
"total_columns": sum(sheet.max_column for sheet in wb.worksheets),
}
return {"metadata": props, "content": content, "statistics": stats} # Get workbook statistics
stats = {
"sheet_count": len(wb.sheetnames),
"total_rows": sum(sheet.max_row for sheet in wb.worksheets),
"total_columns": sum(sheet.max_column for sheet in wb.worksheets),
}
except Exception as e: return {"metadata": props, "content": content, "statistics": stats}
logger.error(f"Failed to get XLSX info: {e}")
return None except Exception as e:
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""" """Universal function to extract content from Office files"""
assert state.get("file_path"), "No file path provided" assert state.get("file_path"), "No file path provided"
assert ( assert (
state.get("identified_type") in SUPPORTED_OFFICE_TYPES state.get("identified_type") in SUPPORTED_OFFICE_TYPES
), "Unsupported File Type" ), "Unsupported File Type"
file_path = state["file_path"] file_path = state["file_path"]
doc_type = state["identified_type"] doc_type = state["identified_type"]
@ -266,24 +301,23 @@ def extract_office_content(state: ContentState):
== "application/vnd.openxmlformats-officedocument.wordprocessingml.document" == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
): ):
logger.debug("Extracting content from DOCX file") 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 ( elif (
doc_type doc_type
== "application/vnd.openxmlformats-officedocument.presentationml.presentation" == "application/vnd.openxmlformats-officedocument.presentationml.presentation"
): ):
logger.debug("Extracting content from PPTX file") 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 ( elif (
doc_type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" doc_type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
): ):
logger.debug("Extracting content from XLSX file") 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: else:
raise Exception(f"Unsupported file format: {doc_type}") raise Exception(f"Unsupported file format: {doc_type}")
del info["content"] del info["content"]
return {"content": content, "metadata": info} return {"content": content, "metadata": info}

28
open_notebook/graphs/content_processing/pdf.py
View file

@ -1,3 +1,4 @@
import asyncio
import re import re
import unicodedata import unicodedata
@ -114,7 +115,7 @@ def clean_pdf_text(text):
return text.strip() return text.strip()
def _extract_text_from_pdf(pdf_path): async def _extract_text_from_pdf(pdf_path):
doc = fitz.open(pdf_path) doc = fitz.open(pdf_path)
try: try:
text = "" text = ""
@ -127,20 +128,39 @@ def _extract_text_from_pdf(pdf_path):
doc.close() 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 = {} return_dict = {}
assert state.get("file_path"), "No file path provided" assert state.get("file_path"), "No file path provided"
assert state.get("identified_type") in SUPPORTED_FITZ_TYPES, "Unsupported File Type" assert state.get("identified_type") in SUPPORTED_FITZ_TYPES, "Unsupported File Type"
if ( if (
state.get("file_path") is not None state.get("file_path") is not None
and state.get("identified_type") in SUPPORTED_FITZ_TYPES and state.get("identified_type") in SUPPORTED_FITZ_TYPES
): ):
file_path = state.get("file_path") file_path = state.get("file_path")
try: try:
text = _extract_text_from_pdf(file_path) text = await _extract_text_from_pdf(file_path)
return_dict["content"] = text return_dict["content"] = text
except FileNotFoundError: except FileNotFoundError:
raise FileNotFoundError(f"File not found at {file_path}") raise FileNotFoundError(f"File not found at {file_path}")

24
open_notebook/graphs/content_processing/text.py
View file

@ -1,11 +1,13 @@
import asyncio
from loguru import logger from loguru import logger
from open_notebook.graphs.content_processing.state import ContentState 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 = {} return_dict = {}
if ( if (
@ -14,12 +16,22 @@ def extract_txt(state: ContentState):
): ):
logger.debug(f"Extracting text from {state.get('file_path')}") logger.debug(f"Extracting text from {state.get('file_path')}")
file_path = state.get("file_path") file_path = state.get("file_path")
if file_path is not None: if file_path is not None:
try: try:
with open(file_path, "r", encoding="utf-8") as file:
content = file.read() def _read_file():
logger.debug(f"Extracted: {content[:100]}") with open(file_path, "r", encoding="utf-8") as file:
return_dict["content"] = content 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: except FileNotFoundError:
raise FileNotFoundError(f"File not found at {file_path}") raise FileNotFoundError(f"File not found at {file_path}")
except Exception as e: except Exception as e:

53
open_notebook/graphs/content_processing/url.py
View file

@ -1,7 +1,7 @@
import re import re
from urllib.parse import urlparse from urllib.parse import urlparse
import requests # type: ignore import aiohttp
from bs4 import BeautifulSoup, Comment from bs4 import BeautifulSoup, Comment
from loguru import logger from loguru import logger
@ -29,7 +29,7 @@ def url_provider(state: ContentState):
return return_dict 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 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"): if url.startswith("<!DOCTYPE html>") or url.startswith("<html"):
html_content = url html_content = url
else: else:
response = requests.get(url, headers=headers, timeout=10) async with aiohttp.ClientSession() as session:
response.raise_for_status() async with session.get(url, headers=headers, timeout=10) as response:
html_content = response.text response.raise_for_status()
html_content = await response.text()
soup = BeautifulSoup(html_content, "html.parser") 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, "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}") logger.error(f"Failed to fetch URL {url}: {e}")
return None return None
except Exception as e: except Exception as e:
@ -151,38 +152,38 @@ def extract_url_bs4(url: str):
return None return None
def extract_url_jina(url: str): async def extract_url_jina(url: str):
""" """
Get the content of a URL using Jina 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:
if text.startswith("Title:") and "\n" in text: text = await response.text()
title_end = text.index("\n") if text.startswith("Title:") and "\n" in text:
title = text[6:title_end].strip() title_end = text.index("\n")
content = text[title_end + 1 :].strip() title = text[6:title_end].strip()
logger.debug( content = text[title_end + 1 :].strip()
f"Processed url: {url}, found title: {title}, content: {content[:100]}..." logger.debug(
) f"Processed url: {url}, found title: {title}, content: {content[:100]}..."
return {"title": title, "content": content} )
else: return {"title": title, "content": content}
content = text else:
logger.debug( 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} return {"content": text}
def extract_url(state: ContentState): async def extract_url(state: ContentState):
assert state.get("url"), "No URL provided" assert state.get("url"), "No URL provided"
url = state["url"] url = state["url"]
try: try:
result = extract_url_bs4(url) result = await extract_url_bs4(url)
if not result or not result.get("content"): if not result or not result.get("content"):
logger.debug( logger.debug(
f"BS4 extraction failed for url {url}, falling back to Jina extractor" 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 return result
except Exception as e: except Exception as e:
logger.error(f"URL extraction failed for URL: {url}") logger.error(f"URL extraction failed for URL: {url}")

179
open_notebook/graphs/content_processing/video.py
View file

@ -1,114 +1,141 @@
import asyncio
import json import json
import os import os
import subprocess import subprocess
from functools import partial
from loguru import logger from loguru import logger
from open_notebook.graphs.content_processing.state import ContentState 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
""" """
try:
cmd = [
"ffmpeg",
"-i",
input_file,
"-map",
f"0:a:{stream_index}", # Select specific audio stream
"-codec:a",
"libmp3lame", # Use MP3 codec
"-q:a",
"2", # High quality setting
"-y", # Overwrite output file if exists
output_file,
]
result = subprocess.run(cmd, capture_output=True, text=True) def _extract(input_file, output_file, stream_index):
if result.returncode != 0: try:
raise Exception(f"FFmpeg failed: {result.stderr}") cmd = [
"ffmpeg",
"-i",
input_file,
"-map",
f"0:a:{stream_index}", # Select specific audio stream
"-codec:a",
"libmp3lame", # Use MP3 codec
"-q:a",
"2", # High quality setting
"-y", # Overwrite output file if exists
output_file,
]
return True result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
raise Exception(f"FFmpeg failed: {result.stderr}")
except Exception as e: return True
print(f"Error extracting audio: {str(e)}")
return False except Exception as 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
""" """
logger.debug(f"Analyzing video file {input_file} for audio streams")
try:
# Get stream information in JSON format
cmd = [
"ffprobe",
"-v",
"quiet",
"-print_format",
"json",
"-show_streams",
"-select_streams",
"a",
input_file,
]
result = subprocess.run(cmd, capture_output=True, text=True) def _analyze(input_file):
if result.returncode != 0: logger.debug(f"Analyzing video file {input_file} for audio streams")
raise Exception(f"FFprobe failed: {result.stderr}") try:
cmd = [
"ffprobe",
"-v",
"quiet",
"-print_format",
"json",
"-show_streams",
"-select_streams",
"a",
input_file,
]
data = json.loads(result.stdout) result = subprocess.run(cmd, capture_output=True, text=True)
return data.get("streams", []) if result.returncode != 0:
raise Exception(f"FFprobe failed: {result.stderr}")
except Exception as e: data = json.loads(result.stdout)
print(f"Error analyzing file: {str(e)}") return data.get("streams", [])
return []
except Exception as 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 Select the best audio stream based on various quality metrics
""" """
if not streams:
logger.debug("No audio streams found")
return None
else:
logger.debug(f"Found {len(streams)} audio streams")
# Score each stream based on various factors def _select(streams):
scored_streams = [] if not streams:
for stream in streams: logger.debug("No audio streams found")
score = 0 return None
else:
logger.debug(f"Found {len(streams)} audio streams")
# Prefer higher bit rates # Score each stream based on various factors
bit_rate = stream.get("bit_rate") scored_streams = []
if bit_rate: for stream in streams:
score += int(int(bit_rate) / 1000000) # Convert to Mbps and ensure int score = 0
# Prefer more channels (stereo over mono) # Prefer higher bit rates
channels = stream.get("channels", 0) bit_rate = stream.get("bit_rate")
score += channels * 10 if bit_rate:
score += int(int(bit_rate) / 1000000) # Convert to Mbps and ensure int
# Prefer higher sample rates # Prefer more channels (stereo over mono)
sample_rate = stream.get("sample_rate", "0") channels = stream.get("channels", 0)
score += int(int(sample_rate) / 48000) score += channels * 10
scored_streams.append((score, stream)) # Prefer higher sample rates
sample_rate = stream.get("sample_rate", "0")
score += int(int(sample_rate) / 48000)
# Return the stream with highest score scored_streams.append((score, stream))
return max(scored_streams, key=lambda x: x[0])[1]
# 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") input_file = data.get("file_path")
assert input_file is not None, "Input file path must be provided" 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}") logger.critical(f"Input file not found: {input_file}")
return False return False
@ -116,20 +143,20 @@ def extract_best_audio_from_video(data: ContentState):
output_file = f"{base_name}_audio.mp3" output_file = f"{base_name}_audio.mp3"
# Get all audio streams # Get all audio streams
streams = get_audio_streams(input_file) streams = await get_audio_streams(input_file)
if not streams: if not streams:
logger.debug("No audio streams found in the file") logger.debug("No audio streams found in the file")
return False return False
# Select best stream # Select best stream
best_stream = select_best_audio_stream(streams) best_stream = await select_best_audio_stream(streams)
if not best_stream: if not best_stream:
logger.error("Could not determine best audio stream") logger.error("Could not determine best audio stream")
return False return False
# Extract the selected stream # Extract the selected stream
stream_index = streams.index(best_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: if success:
logger.debug(f"Successfully extracted audio to: {output_file}") logger.debug(f"Successfully extracted audio to: {output_file}")

20
open_notebook/graphs/content_processing/youtube.py
View file

@ -1,7 +1,7 @@
import re import re
import ssl import ssl
import requests import aiohttp
from bs4 import BeautifulSoup from bs4 import BeautifulSoup
from loguru import logger from loguru import logger
from youtube_transcript_api import YouTubeTranscriptApi # type: ignore 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 ssl._create_default_https_context = ssl._create_unverified_context
def get_video_title(video_id): async def get_video_title(video_id):
try: try:
url = f"https://www.youtube.com/watch?v={video_id}" 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 # YouTube stores title in a meta tag
title = soup.find("meta", property="og:title")["content"] 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 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: try:
transcript_list = YouTubeTranscriptApi.list_transcripts(video_id) transcript_list = YouTubeTranscriptApi.list_transcripts(video_id)
@ -129,7 +133,7 @@ def get_best_transcript(video_id, preferred_langs=["en", "es", "pt"]):
return None return None
def extract_youtube_transcript(state: ContentState): async def extract_youtube_transcript(state: ContentState):
""" """
Parse the text file and print its content. 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")) 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}") logger.debug(f"Found transcript: {transcript}")
formatter = TextFormatter() formatter = TextFormatter()
try: try:
title = get_video_title(video_id) title = await get_video_title(video_id)
except Exception as e: except Exception as e:
logger.critical(f"Failed to get video title for video_id: {video_id}") logger.critical(f"Failed to get video title for video_id: {video_id}")
logger.exception(e) logger.exception(e)

34
open_notebook/graphs/source.py
View file

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

15
pages/stream_app/source.py
View file

@ -1,3 +1,4 @@
import asyncio
import os import os
from pathlib import Path from pathlib import Path
@ -71,12 +72,14 @@ def add_source(notebook_id):
f.write(source_file.getbuffer()) f.write(source_file.getbuffer())
st.write("Processing content..") st.write("Processing content..")
source_graph.invoke( asyncio.run(
{ source_graph.ainvoke(
"content_state": req, {
"notebook_id": notebook_id, "content_state": req,
"transformations": apply_transformations, "notebook_id": notebook_id,
} "transformations": apply_transformations,
}
)
) )
except UnsupportedTypeException as e: except UnsupportedTypeException as e:
st.warning( st.warning(