fix: run ruff formatter to fix code quality

This commit is contained in:
Anish Sarkar 2025-11-30 04:15:38 +05:30
parent b98c312fb1
commit e419702ebd
4 changed files with 77 additions and 66 deletions

View file

@ -48,15 +48,19 @@ def upgrade() -> None:
from app.tasks.celery_tasks.blocknote_migration_tasks import ( from app.tasks.celery_tasks.blocknote_migration_tasks import (
populate_blocknote_for_documents_task, populate_blocknote_for_documents_task,
) )
# Queue the task to run asynchronously # Queue the task to run asynchronously
populate_blocknote_for_documents_task.apply_async() populate_blocknote_for_documents_task.apply_async()
print("✓ Queued Celery task to populate blocknote_document for existing documents") print(
"✓ Queued Celery task to populate blocknote_document for existing documents"
)
except Exception as e: except Exception as e:
# If Celery is not available or task queueing fails, log but don't fail the migration # If Celery is not available or task queueing fails, log but don't fail the migration
print(f"⚠ Warning: Could not queue blocknote population task: {e}") print(f"⚠ Warning: Could not queue blocknote population task: {e}")
print(" You can manually trigger it later with:") print(" You can manually trigger it later with:")
print(" celery -A app.celery_app call app.tasks.celery_tasks.blocknote_migration_tasks.populate_blocknote_for_documents_task") print(
" celery -A app.celery_app call app.tasks.celery_tasks.blocknote_migration_tasks.populate_blocknote_for_documents_task"
)
def downgrade() -> None: def downgrade() -> None:

View file

@ -33,7 +33,7 @@ async def get_editor_content(
attempts to generate it from chunks (lazy migration). attempts to generate it from chunks (lazy migration).
""" """
from sqlalchemy.orm import selectinload from sqlalchemy.orm import selectinload
result = await session.execute( result = await session.execute(
select(Document) select(Document)
.options(selectinload(Document.chunks)) .options(selectinload(Document.chunks))
@ -58,39 +58,39 @@ async def get_editor_content(
# Lazy migration: Try to generate blocknote_document from chunks # Lazy migration: Try to generate blocknote_document from chunks
from app.utils.blocknote_converter import convert_markdown_to_blocknote from app.utils.blocknote_converter import convert_markdown_to_blocknote
chunks = sorted(document.chunks, key=lambda c: c.id) chunks = sorted(document.chunks, key=lambda c: c.id)
if not chunks: if not chunks:
raise HTTPException( raise HTTPException(
status_code=400, status_code=400,
detail="This document has no chunks and cannot be edited. Please re-upload to enable editing.", detail="This document has no chunks and cannot be edited. Please re-upload to enable editing.",
) )
# Reconstruct markdown from chunks # Reconstruct markdown from chunks
markdown_content = "\n\n".join(chunk.content for chunk in chunks) markdown_content = "\n\n".join(chunk.content for chunk in chunks)
if not markdown_content.strip(): if not markdown_content.strip():
raise HTTPException( raise HTTPException(
status_code=400, status_code=400,
detail="This document has empty content and cannot be edited.", detail="This document has empty content and cannot be edited.",
) )
# Convert to BlockNote # Convert to BlockNote
blocknote_json = await convert_markdown_to_blocknote(markdown_content) blocknote_json = await convert_markdown_to_blocknote(markdown_content)
if not blocknote_json: if not blocknote_json:
raise HTTPException( raise HTTPException(
status_code=500, status_code=500,
detail="Failed to convert document to editable format. Please try again later.", detail="Failed to convert document to editable format. Please try again later.",
) )
# Save the generated blocknote_document (lazy migration) # Save the generated blocknote_document (lazy migration)
document.blocknote_document = blocknote_json document.blocknote_document = blocknote_json
document.content_needs_reindexing = False document.content_needs_reindexing = False
document.last_edited_at = None document.last_edited_at = None
await session.commit() await session.commit()
return { return {
"document_id": document.id, "document_id": document.id,
"title": document.title, "title": document.title,
@ -111,7 +111,7 @@ async def save_document(
Called when user clicks 'Save & Exit'. Called when user clicks 'Save & Exit'.
""" """
from app.tasks.celery_tasks.document_reindex_tasks import reindex_document_task from app.tasks.celery_tasks.document_reindex_tasks import reindex_document_task
# Verify ownership # Verify ownership
result = await session.execute( result = await session.execute(
select(Document) select(Document)
@ -119,27 +119,27 @@ async def save_document(
.filter(Document.id == document_id, SearchSpace.user_id == user.id) .filter(Document.id == document_id, SearchSpace.user_id == user.id)
) )
document = result.scalars().first() document = result.scalars().first()
if not document: if not document:
raise HTTPException(status_code=404, detail="Document not found") raise HTTPException(status_code=404, detail="Document not found")
blocknote_document = data.get("blocknote_document") blocknote_document = data.get("blocknote_document")
if not blocknote_document: if not blocknote_document:
raise HTTPException(status_code=400, detail="blocknote_document is required") raise HTTPException(status_code=400, detail="blocknote_document is required")
# Save BlockNote document # Save BlockNote document
document.blocknote_document = blocknote_document document.blocknote_document = blocknote_document
document.last_edited_at = datetime.now(UTC) document.last_edited_at = datetime.now(UTC)
document.content_needs_reindexing = True document.content_needs_reindexing = True
await session.commit() await session.commit()
# Queue reindex task # Queue reindex task
reindex_document_task.delay(document_id, str(user.id)) reindex_document_task.delay(document_id, str(user.id))
return { return {
"status": "saved", "status": "saved",
"document_id": document_id, "document_id": document_id,
"message": "Document saved and will be reindexed in the background", "message": "Document saved and will be reindexed in the background",
"last_edited_at": document.last_edited_at.isoformat() "last_edited_at": document.last_edited_at.isoformat(),
} }

View file

@ -36,7 +36,7 @@ def populate_blocknote_for_documents_task(
): ):
""" """
Celery task to populate blocknote_document for existing documents. Celery task to populate blocknote_document for existing documents.
Args: Args:
document_ids: Optional list of specific document IDs to process. document_ids: Optional list of specific document IDs to process.
If None, processes all documents with blocknote_document IS NULL. If None, processes all documents with blocknote_document IS NULL.
@ -60,7 +60,7 @@ async def _populate_blocknote_for_documents(
): ):
""" """
Async function to populate blocknote_document for documents. Async function to populate blocknote_document for documents.
Args: Args:
document_ids: Optional list of specific document IDs to process document_ids: Optional list of specific document IDs to process
batch_size: Number of documents to process per batch batch_size: Number of documents to process per batch
@ -69,75 +69,83 @@ async def _populate_blocknote_for_documents(
try: try:
# Build query for documents that need blocknote_document populated # Build query for documents that need blocknote_document populated
query = select(Document).where(Document.blocknote_document.is_(None)) query = select(Document).where(Document.blocknote_document.is_(None))
# If specific document IDs provided, filter by them # If specific document IDs provided, filter by them
if document_ids: if document_ids:
query = query.where(Document.id.in_(document_ids)) query = query.where(Document.id.in_(document_ids))
# Load chunks relationship to avoid N+1 queries # Load chunks relationship to avoid N+1 queries
query = query.options(selectinload(Document.chunks)) query = query.options(selectinload(Document.chunks))
# Execute query # Execute query
result = await session.execute(query) result = await session.execute(query)
documents = result.scalars().all() documents = result.scalars().all()
total_documents = len(documents) total_documents = len(documents)
logger.info(f"Found {total_documents} documents to process") logger.info(f"Found {total_documents} documents to process")
if total_documents == 0: if total_documents == 0:
logger.info("No documents to process") logger.info("No documents to process")
return return
# Process documents in batches # Process documents in batches
processed = 0 processed = 0
failed = 0 failed = 0
for i in range(0, total_documents, batch_size): for i in range(0, total_documents, batch_size):
batch = documents[i : i + batch_size] batch = documents[i : i + batch_size]
logger.info(f"Processing batch {i // batch_size + 1}: documents {i+1}-{min(i+batch_size, total_documents)}") logger.info(
f"Processing batch {i // batch_size + 1}: documents {i + 1}-{min(i + batch_size, total_documents)}"
)
for document in batch: for document in batch:
try: try:
# Use preloaded chunks from selectinload - no need to query again # Use preloaded chunks from selectinload - no need to query again
chunks = sorted(document.chunks, key=lambda c: c.id) chunks = sorted(document.chunks, key=lambda c: c.id)
if not chunks: if not chunks:
logger.warning( logger.warning(
f"Document {document.id} ({document.title}) has no chunks, skipping" f"Document {document.id} ({document.title}) has no chunks, skipping"
) )
failed += 1 failed += 1
continue continue
# Reconstruct markdown by concatenating chunk contents # Reconstruct markdown by concatenating chunk contents
markdown_content = "\n\n".join(chunk.content for chunk in chunks) markdown_content = "\n\n".join(
chunk.content for chunk in chunks
)
if not markdown_content or not markdown_content.strip(): if not markdown_content or not markdown_content.strip():
logger.warning( logger.warning(
f"Document {document.id} ({document.title}) has empty markdown content, skipping" f"Document {document.id} ({document.title}) has empty markdown content, skipping"
) )
failed += 1 failed += 1
continue continue
# Convert markdown to BlockNote JSON # Convert markdown to BlockNote JSON
blocknote_json = await convert_markdown_to_blocknote(markdown_content) blocknote_json = await convert_markdown_to_blocknote(
markdown_content
)
if not blocknote_json: if not blocknote_json:
logger.warning( logger.warning(
f"Failed to convert markdown to BlockNote for document {document.id} ({document.title})" f"Failed to convert markdown to BlockNote for document {document.id} ({document.title})"
) )
failed += 1 failed += 1
continue continue
# Update document with blocknote_document (other fields already have correct defaults) # Update document with blocknote_document (other fields already have correct defaults)
document.blocknote_document = blocknote_json document.blocknote_document = blocknote_json
processed += 1 processed += 1
# Commit every batch_size documents to avoid long transactions # Commit every batch_size documents to avoid long transactions
if processed % batch_size == 0: if processed % batch_size == 0:
await session.commit() await session.commit()
logger.info(f"Committed batch: {processed} documents processed so far") logger.info(
f"Committed batch: {processed} documents processed so far"
)
except Exception as e: except Exception as e:
logger.error( logger.error(
f"Error processing document {document.id} ({document.title}): {e}", f"Error processing document {document.id} ({document.title}): {e}",
@ -146,15 +154,15 @@ async def _populate_blocknote_for_documents(
failed += 1 failed += 1
# Continue with next document instead of failing entire batch # Continue with next document instead of failing entire batch
continue continue
# Commit remaining changes in the batch # Commit remaining changes in the batch
await session.commit() await session.commit()
logger.info(f"Completed batch {i // batch_size + 1}") logger.info(f"Completed batch {i // batch_size + 1}")
logger.info( logger.info(
f"Migration complete: {processed} documents processed, {failed} failed" f"Migration complete: {processed} documents processed, {failed} failed"
) )
except Exception as e: except Exception as e:
await session.rollback() await session.rollback()
logger.error(f"Error in blocknote migration task: {e}", exc_info=True) logger.error(f"Error in blocknote migration task: {e}", exc_info=True)

View file

@ -35,7 +35,7 @@ def get_celery_session_maker():
def reindex_document_task(self, document_id: int, user_id: str): def reindex_document_task(self, document_id: int, user_id: str):
""" """
Celery task to reindex a document after editing. Celery task to reindex a document after editing.
Args: Args:
document_id: ID of document to reindex document_id: ID of document to reindex
user_id: ID of user who edited the document user_id: ID of user who edited the document
@ -62,66 +62,65 @@ async def _reindex_document(document_id: int, user_id: str):
.where(Document.id == document_id) .where(Document.id == document_id)
) )
document = result.scalars().first() document = result.scalars().first()
if not document: if not document:
logger.error(f"Document {document_id} not found") logger.error(f"Document {document_id} not found")
return return
if not document.blocknote_document: if not document.blocknote_document:
logger.warning(f"Document {document_id} has no BlockNote content") logger.warning(f"Document {document_id} has no BlockNote content")
return return
logger.info(f"Reindexing document {document_id} ({document.title})") logger.info(f"Reindexing document {document_id} ({document.title})")
# 1. Convert BlockNote → Markdown # 1. Convert BlockNote → Markdown
markdown_content = await convert_blocknote_to_markdown( markdown_content = await convert_blocknote_to_markdown(
document.blocknote_document document.blocknote_document
) )
if not markdown_content: if not markdown_content:
logger.error(f"Failed to convert document {document_id} to markdown") logger.error(f"Failed to convert document {document_id} to markdown")
return return
# 2. Delete old chunks explicitly # 2. Delete old chunks explicitly
from app.db import Chunk from app.db import Chunk
await session.execute(
delete(Chunk).where(Chunk.document_id == document_id) await session.execute(delete(Chunk).where(Chunk.document_id == document_id))
)
await session.flush() # Ensure old chunks are deleted await session.flush() # Ensure old chunks are deleted
# 3. Create new chunks # 3. Create new chunks
new_chunks = await create_document_chunks(markdown_content) new_chunks = await create_document_chunks(markdown_content)
# 4. Add new chunks to session # 4. Add new chunks to session
for chunk in new_chunks: for chunk in new_chunks:
chunk.document_id = document_id chunk.document_id = document_id
session.add(chunk) session.add(chunk)
logger.info(f"Created {len(new_chunks)} chunks for document {document_id}") logger.info(f"Created {len(new_chunks)} chunks for document {document_id}")
# 5. Regenerate summary # 5. Regenerate summary
user_llm = await get_user_long_context_llm( user_llm = await get_user_long_context_llm(
session, user_id, document.search_space_id session, user_id, document.search_space_id
) )
document_metadata = { document_metadata = {
"title": document.title, "title": document.title,
"document_type": document.document_type.value, "document_type": document.document_type.value,
} }
summary_content, summary_embedding = await generate_document_summary( summary_content, summary_embedding = await generate_document_summary(
markdown_content, user_llm, document_metadata markdown_content, user_llm, document_metadata
) )
# 6. Update document # 6. Update document
document.content = summary_content document.content = summary_content
document.embedding = summary_embedding document.embedding = summary_embedding
document.content_needs_reindexing = False document.content_needs_reindexing = False
await session.commit() await session.commit()
logger.info(f"Successfully reindexed document {document_id}") logger.info(f"Successfully reindexed document {document_id}")
except Exception as e: except Exception as e:
await session.rollback() await session.rollback()
logger.error(f"Error reindexing document {document_id}: {e}", exc_info=True) logger.error(f"Error reindexing document {document_id}: {e}", exc_info=True)