File size: 13,022 Bytes
b30f068 6163d4f b30f068 22d5317 b30f068 22d5317 b30f068 6163d4f 22d5317 6163d4f 22d5317 b30f068 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 | """
Document Loader
Scans PDF folder, processes PDFs, and stores in database + Qdrant
"""
import sys
from pathlib import Path
# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))
from typing import Dict
from datetime import datetime
from src.cdms.pdf_processor import PDFProcessor
from src.cdms.schema import DatabaseManager, Document, DocumentChunk
from src.rag.embeddings import OpenAIEmbeddingService
from src.rag.vector_store import QdrantVectorStore, get_shared_vector_store
from src.config.credentials import CredentialsManager
from src.config.paths import PDF_DIR
from src.cdms.product_catalog import normalize_filename
class DocumentLoader:
"""
Loads and indexes PDF documents for RAG
Usage:
loader = DocumentLoader()
loader.load_all_pdfs() # Process all PDFs in data/pdfs/
"""
def __init__(self, pdf_folder: str = None, vector_store=None):
"""
Initialize document loader
Args:
pdf_folder: Folder containing PDF files (defaults to the CDMS PDF dir
PDF_DIR = <root>/data/pdfs/cdms, where the labels live)
vector_store: Optional shared QdrantVectorStore. The embedded (on-disk)
Qdrant allows only ONE client per path per process, so in
online mode the loader MUST reuse the RAG searcher's store
instead of opening a second one (which would lock-fail and
silently drop the freshly fetched label's chunks).
"""
self.pdf_folder = Path(pdf_folder) if pdf_folder is not None else PDF_DIR
self.pdf_folder.mkdir(parents=True, exist_ok=True)
# Initialize components
self.pdf_processor = PDFProcessor()
self.db_manager = DatabaseManager()
# Initialize embedding service (will load OpenAI key)
try:
creds = CredentialsManager()
openai_key = creds.get_api_key("openai")
self.embedding_service = OpenAIEmbeddingService(api_key=openai_key)
except Exception as e:
print(f"⚠️ Warning: Could not initialize OpenAI embeddings: {e}")
self.embedding_service = None
# Reuse the provided store, else the process-wide singleton (never open a
# second client to the embedded store — see get_shared_vector_store).
if vector_store is not None:
self.vector_store = vector_store
else:
try:
self.vector_store = get_shared_vector_store()
except Exception as e:
print(f"⚠️ Warning: Could not initialize Qdrant: {e}")
self.vector_store = None
def load_pdf(self, pdf_path: str, force_reprocess: bool = False, pdf_url: str = None) -> Dict:
"""
Load a single PDF file
Args:
pdf_path: Path to PDF file
force_reprocess: If True, reprocess even if already indexed
pdf_url: Optional original PDF URL (for CDMS labels from Tavily)
Returns:
Dict with processing result
"""
pdf_path = Path(pdf_path)
if not pdf_path.exists():
return {
"success": False,
"error": f"File not found: {pdf_path}"
}
# Check if already processed
session = self.db_manager.get_session()
try:
doc_id = Document.generate_id(str(pdf_path))
existing_doc = session.query(Document).filter_by(id=doc_id).first()
if existing_doc and existing_doc.processed == 1 and not force_reprocess:
return {
"success": True,
"message": f"PDF already processed: {pdf_path.name}",
"document_id": doc_id,
"skipped": True
}
# If force reprocess, delete old chunks and document
if force_reprocess and existing_doc:
# Delete old chunks
session.query(DocumentChunk).filter_by(document_id=doc_id).delete()
# Delete old document
session.delete(existing_doc)
session.commit()
print(f"🔄 Re-processing: {pdf_path.name}")
# Process PDF
print(f"📄 Processing: {pdf_path.name}")
result = self.pdf_processor.process_pdf(str(pdf_path))
if not result["success"]:
return result
# Store in database. Persist the source URL in doc_metadata so
# citations survive even if Qdrant is rebuilt from SQLite alone.
doc = Document(
id=doc_id,
filename=pdf_path.name,
filepath=str(pdf_path),
file_size=pdf_path.stat().st_size,
num_pages=result["num_pages"],
num_chunks=result["num_chunks"],
processed=1,
last_processed=datetime.utcnow(),
doc_metadata={"pdf_url": pdf_url} if pdf_url else None
)
session.merge(doc)
# Store chunks with accurate page numbers
chunks = result.get("chunks", [])
page_numbers = result.get("page_numbers", [])
# PHASE 2 FIX: Validate page numbers with warnings
if not page_numbers:
print(f"⚠️ Warning: No page numbers provided for {pdf_path.name}, estimating...")
page_numbers = [(idx // 3) + 1 for idx in range(len(chunks))]
elif len(page_numbers) != len(chunks):
print(f"⚠️ Warning: Page numbers count ({len(page_numbers)}) doesn't match chunks count ({len(chunks)}) for {pdf_path.name}")
# Fix by padding or truncating to match
if len(page_numbers) < len(chunks):
last_page = page_numbers[-1] if page_numbers else 1
page_numbers.extend([last_page] * (len(chunks) - len(page_numbers)))
else:
page_numbers = page_numbers[:len(chunks)]
# URL hash for reliable matching (computed once per PDF).
url_hash = ""
if pdf_url:
import hashlib
url_hash = hashlib.md5(pdf_url.encode()).hexdigest()[:12]
# Normalized product name (e.g. "roundup_bdc94bbee383.pdf" -> "roundup")
# stored in the payload so retrieval can filter by product AT the vector
# level, instead of post-filtering a global (Roundup-dominated) top-k.
product = normalize_filename(pdf_path.name)
# Pass 1: write SQLite chunk rows and collect texts for BATCH embedding.
# (Previously embeddings were generated one OpenAI call per chunk; batching
# collapses ~N calls per PDF into a single request -- ~2100 calls -> ~53.)
chunk_ids = []
chunk_texts = []
chunk_pages = []
chunk_indices = []
for idx, (chunk_text, page_num) in enumerate(zip(chunks, page_numbers)):
# PHASE 2 FIX: Validate page number is positive
if page_num <= 0:
page_num = (idx // 3) + 1
page_numbers[idx] = page_num
chunk_id = DocumentChunk.generate_id(doc_id, idx)
# Calculate token count (rough estimate: 1 token ≈ 4 chars)
token_count = len(chunk_text) // 4
chunk = DocumentChunk(
id=chunk_id,
document_id=doc_id,
chunk_index=idx,
content=chunk_text,
page_number=page_num, # ENHANCED: Accurate page number from PDF processor
char_count=len(chunk_text),
token_count=token_count
)
session.merge(chunk)
chunk_ids.append(chunk_id)
chunk_texts.append(chunk_text)
chunk_pages.append(page_num)
chunk_indices.append(idx)
chunks_stored = len(chunk_ids)
embeddings_generated = 0
# Pass 2: one batched embedding request per PDF, then upsert to Qdrant.
if self.embedding_service and self.vector_store and chunk_texts:
try:
embeddings = self.embedding_service.generate_embeddings_batch(chunk_texts)
for chunk_id, embedding, chunk_text, page_num, idx in zip(
chunk_ids, embeddings, chunk_texts, chunk_pages, chunk_indices
):
metadata = {
"document_id": doc_id,
"document_name": pdf_path.name,
"chunk_index": idx,
"content": chunk_text, # Store full content for search
"page_number": page_num, # ENHANCED: Accurate page number
"source_file": pdf_path.name,
"product": product, # normalized product for Qdrant-level filtering
"pdf_url": pdf_url if pdf_url else "", # PHASE 1 FIX: Store PDF URL in metadata
"url_hash": url_hash # PHASE 1 FIX: Store URL hash for reliable matching
}
self.vector_store.add_document_chunk(
chunk_id=chunk_id,
embedding=embedding,
metadata=metadata
)
embeddings_generated += 1
except Exception as e:
print(f"⚠️ Warning: Could not generate embeddings for {pdf_path.name}: {e}")
session.commit()
return {
"success": True,
"document_id": doc_id,
"filename": pdf_path.name,
"chunks_stored": chunks_stored,
"embeddings_generated": embeddings_generated,
"num_pages": result["num_pages"]
}
finally:
session.close()
def load_all_pdfs(self, force_reprocess: bool = False) -> Dict:
"""
Load all PDFs from the pdf_folder
Returns:
Dict with summary of processing
"""
if not self.pdf_folder.exists():
return {
"success": False,
"error": f"PDF folder not found: {self.pdf_folder}"
}
pdf_files = list(self.pdf_folder.glob("*.pdf"))
if not pdf_files:
return {
"success": False,
"error": f"No PDF files found in {self.pdf_folder}"
}
print(f"📚 Found {len(pdf_files)} PDF file(s)")
print("-" * 70)
results = []
for pdf_file in pdf_files:
result = self.load_pdf(str(pdf_file), force_reprocess=force_reprocess)
results.append(result)
if result.get("success") and not result.get("skipped"):
print(f"✅ {pdf_file.name}: {result.get('chunks_stored', 0)} chunks, {result.get('embeddings_generated', 0)} embeddings")
elif result.get("skipped"):
print(f"⏭️ {pdf_file.name}: Already processed (skipped)")
print()
# Summary
successful = sum(1 for r in results if r.get("success"))
total_chunks = sum(r.get("chunks_stored", 0) for r in results)
total_embeddings = sum(r.get("embeddings_generated", 0) for r in results)
return {
"success": True,
"total_files": len(pdf_files),
"successful": successful,
"total_chunks": total_chunks,
"total_embeddings": total_embeddings,
"results": results
}
# Test function
if __name__ == "__main__":
print("Testing Document Loader...")
print("=" * 70)
loader = DocumentLoader()
# Load all PDFs
result = loader.load_all_pdfs()
if result.get("success"):
print("\n" + "=" * 70)
print("✅ Processing Complete!")
print(f" Files processed: {result['successful']}/{result['total_files']}")
print(f" Total chunks: {result['total_chunks']}")
print(f" Embeddings generated: {result['total_embeddings']}")
else:
print(f"\n❌ Error: {result.get('error')}")
print("\n💡 To use:")
print(" 1. Create data/pdfs/ folder")
print(" 2. Add PDF files")
print(" 3. Make sure OPENAI_API_KEY is in .env")
print(" 4. Make sure Qdrant is running (Docker)")
|