mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-04 22:02:16 +02:00
fix: run ruff formatter to fix code quality
This commit is contained in:
parent
b98c312fb1
commit
e419702ebd
4 changed files with 77 additions and 66 deletions
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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(),
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue