Skip to content

Commit 3cc8a97

Browse files
param20hparam20h
authored
Merge PR #586
fix: resolve memory leak and process crashing during concurrent PDF ingestion
2 parents 6d76216 + bed456c commit 3cc8a97

3 files changed

Lines changed: 254 additions & 216 deletions

File tree

‎backend/app/services/document_ingestion.py‎

Lines changed: 163 additions & 154 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
import traceback
33
import logging
44
from datetime import datetime, timezone
5+
import asyncio
6+
import threading
7+
import gc
58

69
from app.models import Document
710
from app.rag.agent import persist_document_keywords
@@ -12,6 +15,10 @@
1215
logger = logging.getLogger(__name__)
1316
settings = get_settings()
1417

18+
# Define semaphores to throttle parallel file ingestion (Limit concurrency to 3)
19+
ingestion_semaphore = asyncio.Semaphore(3)
20+
threading_semaphore = threading.Semaphore(3)
21+
1522

1623
def _update_progress(document_id: str, progress: int, stage: str, error: str = None):
1724
"""Update document progress fields in the database."""
@@ -39,172 +46,174 @@ def ingest_document(document_id: str, filepath: str, original_name: str, user_id
3946
"""
4047
from app.database import SessionLocal
4148

42-
db = SessionLocal()
43-
try:
44-
doc = db.query(Document).filter(
45-
Document.id == document_id,
46-
Document.is_deleted.is_(False),
47-
).first()
48-
if not doc:
49-
logger.error("Document %s not found for ingestion", document_id)
50-
return
51-
52-
doc.status = "processing"
53-
doc.processing_stage = "extracting"
54-
doc.processing_progress = 10
55-
doc.error_message = None
56-
doc.last_error_traceback = None
57-
db.commit()
58-
59-
page_count = get_page_count(filepath)
60-
doc.page_count = page_count
61-
doc.processing_progress = 20
62-
db.commit()
63-
49+
with threading_semaphore:
50+
db = SessionLocal()
6451
try:
65-
chunk_kwargs = {}
66-
if doc.chunk_size is not None:
67-
chunk_kwargs["chunk_size"] = doc.chunk_size
68-
if doc.chunk_overlap is not None:
69-
chunk_kwargs["chunk_overlap"] = doc.chunk_overlap
70-
doc.processing_stage = "chunking"
71-
doc.processing_progress = 30
72-
db.commit()
73-
chunks = chunk_document(filepath, **chunk_kwargs)
74-
except TypeError:
75-
chunks = chunk_document(filepath)
76-
77-
# ── Proximity caption pass (PDF only) ────────────────────────────────
78-
# Write bounding-box-derived captions into image chunks BEFORE store_chunks()
79-
# so generate_captions_for_chunks() in vectorstore.py only needs to handle
80-
# the OCR / placeholder fallback for any images without adjacent text.
81-
ext = filepath.rsplit(".", 1)[-1].lower()
82-
if ext == "pdf":
83-
try:
84-
from app.rag.vision import extract_captions_from_pdf
85-
86-
pdf_captions = extract_captions_from_pdf(filepath)
87-
# Build lookup: page -> [captions in figure_index order]
88-
caption_map: dict = {}
89-
for cap in pdf_captions:
90-
caption_map.setdefault(cap["page"], []).append(cap)
91-
92-
fig_counters: dict = {}
93-
for chunk in chunks:
94-
if not chunk.get("image_bytes"):
95-
continue
96-
page = chunk.get("page", 1)
97-
idx = fig_counters.get(page, 0)
98-
page_caps = caption_map.get(page, [])
99-
if idx < len(page_caps) and page_caps[idx]["caption"]:
100-
chunk["image_caption"] = page_caps[idx]["caption"]
101-
chunk["bbox"] = str(page_caps[idx]["bbox"])
102-
fig_counters[page] = idx + 1
103-
except Exception as exc:
104-
logger.warning(
105-
"Proximity caption extraction failed for %s: %s", document_id, exc
106-
)
107-
# ── End proximity caption pass ────────────────────────────────────────
108-
109-
if not chunks:
110-
doc.status = "failed"
111-
doc.processing_progress = 0
112-
doc.error_message = "No text could be extracted from the document"
52+
doc = db.query(Document).filter(
53+
Document.id == document_id,
54+
Document.is_deleted.is_(False),
55+
).first()
56+
if not doc:
57+
logger.error("Document %s not found for ingestion", document_id)
58+
return
59+
60+
doc.status = "processing"
61+
doc.processing_stage = "extracting"
62+
doc.processing_progress = 10
63+
doc.error_message = None
64+
doc.last_error_traceback = None
11365
db.commit()
114-
return
11566

116-
doc.processing_progress = 50
117-
doc.processing_stage = "indexing"
118-
db.commit()
67+
page_count = get_page_count(filepath)
68+
doc.page_count = page_count
69+
doc.processing_progress = 20
70+
db.commit()
11971

120-
try:
121-
from app.rag.graph_builder import build_graph, save_graph
72+
try:
73+
chunk_kwargs = {}
74+
if doc.chunk_size is not None:
75+
chunk_kwargs["chunk_size"] = doc.chunk_size
76+
if doc.chunk_overlap is not None:
77+
chunk_kwargs["chunk_overlap"] = doc.chunk_overlap
78+
doc.processing_stage = "chunking"
79+
doc.processing_progress = 30
80+
db.commit()
81+
chunks = chunk_document(filepath, **chunk_kwargs)
82+
except TypeError:
83+
chunks = chunk_document(filepath)
84+
85+
# ── Proximity caption pass (PDF only) ────────────────────────────────
86+
# Write bounding-box-derived captions into image chunks BEFORE store_chunks()
87+
# so generate_captions_for_chunks() in vectorstore.py only needs to handle
88+
# the OCR / placeholder fallback for any images without adjacent text.
89+
ext = filepath.rsplit(".", 1)[-1].lower()
90+
if ext == "pdf":
91+
try:
92+
from app.rag.vision import extract_captions_from_pdf
93+
94+
pdf_captions = extract_captions_from_pdf(filepath)
95+
# Build lookup: page -> [captions in figure_index order]
96+
caption_map: dict = {}
97+
for cap in pdf_captions:
98+
caption_map.setdefault(cap["page"], []).append(cap)
99+
100+
fig_counters: dict = {}
101+
for chunk in chunks:
102+
if not chunk.get("image_bytes"):
103+
continue
104+
page = chunk.get("page", 1)
105+
idx = fig_counters.get(page, 0)
106+
page_caps = caption_map.get(page, [])
107+
if idx < len(page_caps) and page_caps[idx]["caption"]:
108+
chunk["image_caption"] = page_caps[idx]["caption"]
109+
chunk["bbox"] = str(page_caps[idx]["bbox"])
110+
fig_counters[page] = idx + 1
111+
except Exception as exc:
112+
logger.warning(
113+
"Proximity caption extraction failed for %s: %s", document_id, exc
114+
)
115+
# ── End proximity caption pass ────────────────────────────────────────
116+
117+
if not chunks:
118+
doc.status = "failed"
119+
doc.processing_progress = 0
120+
doc.error_message = "No text could be extracted from the document"
121+
db.commit()
122+
return
122123

123-
graph = build_graph(chunks)
124-
save_graph(graph, user_id=user_id, document_id=document_id)
125-
except Exception as e:
126-
logger.warning("Could not build knowledge graph for document %s: %s", document_id, e)
124+
doc.processing_progress = 50
125+
doc.processing_stage = "indexing"
126+
db.commit()
127127

128-
doc.processing_progress = 70
129-
doc.processing_stage = "embedding"
130-
db.commit()
128+
try:
129+
from app.rag.graph_builder import build_graph, save_graph
131130

132-
chunk_count = store_chunks(
133-
chunks=chunks,
134-
document_id=document_id,
135-
filename=original_name,
136-
user_id=user_id,
137-
)
131+
graph = build_graph(chunks)
132+
save_graph(graph, user_id=user_id, document_id=document_id)
133+
except Exception as e:
134+
logger.warning("Could not build knowledge graph for document %s: %s", document_id, e)
138135

139-
persist_document_keywords(doc, chunks, db)
136+
doc.processing_progress = 70
137+
doc.processing_stage = "embedding"
138+
db.commit()
140139

141-
doc.processing_progress = 85
142-
db.commit()
140+
chunk_count = store_chunks(
141+
chunks=chunks,
142+
document_id=document_id,
143+
filename=original_name,
144+
user_id=user_id,
145+
)
143146

144-
try:
145-
from app.rag.summarizer import generate_document_summary
147+
persist_document_keywords(doc, chunks, db)
146148

147-
summary = generate_document_summary(filepath, max_sentences=2)
148-
if summary:
149-
doc.summary = summary
150-
db.commit()
151-
except Exception as e:
152-
logger.warning("Could not generate summary for document %s: %s", document_id, e)
153-
doc.summary = None
149+
doc.processing_progress = 85
150+
db.commit()
154151

155-
# ── URL extraction pass (PDF only) ────────────────────────────────
156-
ext = filepath.rsplit(".", 1)[-1].lower()
157-
if ext == "pdf":
158152
try:
159-
from app.rag.url_extractor import extract_urls_from_pdf
160-
import json
153+
from app.rag.summarizer import generate_document_summary
154+
155+
summary = generate_document_summary(filepath, max_sentences=2)
156+
if summary:
157+
doc.summary = summary
158+
db.commit()
159+
except Exception as e:
160+
logger.warning("Could not generate summary for document %s: %s", document_id, e)
161+
doc.summary = None
162+
163+
# ── URL extraction pass (PDF only) ────────────────────────────────
164+
ext = filepath.rsplit(".", 1)[-1].lower()
165+
if ext == "pdf":
166+
try:
167+
from app.rag.url_extractor import extract_urls_from_pdf
168+
import json
169+
170+
urls = extract_urls_from_pdf(filepath)
171+
doc.extracted_urls = json.dumps(urls) if urls else None
172+
db.commit()
173+
logger.info(
174+
"Extracted %s URLs from document %s",
175+
len(urls),
176+
document_id,
177+
)
178+
except Exception as exc:
179+
logger.warning(
180+
"URL extraction failed for document %s: %s",
181+
document_id,
182+
exc,
183+
)
184+
# ── End URL extraction pass ───────────────────────────────────────
185+
186+
doc.chunk_count = chunk_count
187+
doc.status = "ready"
188+
doc.processing_progress = 100
189+
doc.processing_stage = "completed"
190+
doc.completed_at = datetime.now(timezone.utc)
191+
doc.error_message = None
192+
db.commit()
161193

162-
urls = extract_urls_from_pdf(filepath)
163-
doc.extracted_urls = json.dumps(urls) if urls else None
164-
db.commit()
165-
logger.info(
166-
"Extracted %s URLs from document %s",
167-
len(urls),
168-
document_id,
169-
)
170-
except Exception as exc:
171-
logger.warning(
172-
"URL extraction failed for document %s: %s",
173-
document_id,
174-
exc,
175-
)
176-
# ── End URL extraction pass ───────────────────────────────────────
177-
178-
doc.chunk_count = chunk_count
179-
doc.status = "ready"
180-
doc.processing_progress = 100
181-
doc.processing_stage = "completed"
182-
doc.completed_at = datetime.now(timezone.utc)
183-
doc.error_message = None
184-
db.commit()
185-
186-
logger.info(
187-
"Document %s ingested: %s pages, %s chunks",
188-
document_id,
189-
page_count,
190-
chunk_count,
191-
)
194+
logger.info(
195+
"Document %s ingested: %s pages, %s chunks",
196+
document_id,
197+
page_count,
198+
chunk_count,
199+
)
192200

193-
except Exception as e:
194-
logger.error("Ingestion error for %s: %s", document_id, e)
195-
db.rollback()
196-
try:
197-
doc = db.query(Document).filter(
198-
Document.id == document_id,
199-
Document.is_deleted.is_(False),
200-
).first()
201-
if doc:
202-
doc.status = "failed"
203-
doc.processing_progress = 0
204-
doc.error_message = str(e)[:500]
205-
doc.last_error_traceback = traceback.format_exc()[:2000]
206-
db.commit()
207-
except Exception:
208-
logger.exception("Failed to mark document %s as failed", document_id)
209-
finally:
210-
db.close()
201+
except Exception as e:
202+
logger.error("Ingestion error for %s: %s", document_id, e)
203+
db.rollback()
204+
try:
205+
doc = db.query(Document).filter(
206+
Document.id == document_id,
207+
Document.is_deleted.is_(False),
208+
).first()
209+
if doc:
210+
doc.status = "failed"
211+
doc.processing_progress = 0
212+
doc.error_message = str(e)[:500]
213+
doc.last_error_traceback = traceback.format_exc()[:2000]
214+
db.commit()
215+
except Exception:
216+
logger.exception("Failed to mark document %s as failed", document_id)
217+
finally:
218+
db.close()
219+
gc.collect()

‎backend/app/services/layout_parser.py‎

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -81,20 +81,26 @@ def process_embedded_images(self, page_num: int, page_obj: fitz.Page) -> List[st
8181

8282
def ingest_document(self) -> List[Dict[str, Any]]:
8383
"""Executes the hybrid pipeline generating combined text and image context strings."""
84-
final_payload = []
85-
structured_chunks = self.extract_structured_text()
86-
final_payload.extend(structured_chunks)
84+
import gc
85+
try:
86+
final_payload = []
87+
structured_chunks = self.extract_structured_text()
88+
final_payload.extend(structured_chunks)
8789

88-
for page_num in range(len(self.doc)):
89-
page = self.doc.load_page(page_num)
90-
img_summaries = self.process_embedded_images(page_num, page)
91-
for summary in img_summaries:
92-
final_payload.append(
93-
{
94-
"page_number": page_num + 1,
95-
"text": f"[Visual Data Extraction Summary]: {summary}",
96-
"type": "visual_image_summary",
97-
}
98-
)
90+
for page_num in range(len(self.doc)):
91+
page = self.doc.load_page(page_num)
92+
img_summaries = self.process_embedded_images(page_num, page)
93+
for summary in img_summaries:
94+
final_payload.append(
95+
{
96+
"page_number": page_num + 1,
97+
"text": f"[Visual Data Extraction Summary]: {summary}",
98+
"type": "visual_image_summary",
99+
}
100+
)
101+
gc.collect()
99102

100-
return final_payload
103+
return final_payload
104+
finally:
105+
self.doc.close()
106+
gc.collect()

0 commit comments

Comments
 (0)