Spaces:
Sleeping
Sleeping
File size: 6,272 Bytes
ad94b49 | 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 | import os
import tempfile
import httpx
import logging
from collections.abc import AsyncGenerator
logger = logging.getLogger(__name__)
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from services.qdrant_service import store_chunks
def _split_text(documents, chunk_size: int = 500, chunk_overlap: int = 80) -> list:
logger.info(f"Splitting {len(documents)} page(s) with chunk_size={chunk_size}, overlap={chunk_overlap}...")
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
result = splitter.split_documents(documents)
logger.info(f"Splitting complete → {len(result)} chunks created.")
return result
async def process_pdf_file_progress(
paper_id: str,
pdf_bytes: bytes,
) -> AsyncGenerator[dict, None]:
"""
Process PDF bytes with progress events: downloading -> splitting -> embedding -> storing.
Yields progress dicts for SSE streaming.
"""
# Write to temp file for loaders
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
tmp.write(pdf_bytes)
tmp_path = tmp.name
try:
import asyncio
# Stage: Parsing PDF
logger.info(f"[paper={paper_id}] Stage: Parsing PDF pages with PyMuPDF...")
yield {
"type": "paper_progress",
"stage": "parsing_pdf",
"label": "Parsing PDF pages",
"paper_id": paper_id,
}
def _parse_pdf_sync():
loader = PyMuPDFLoader(tmp_path)
documents = loader.load()
logger.info(f"[paper={paper_id}] Loaded {len(documents)} page(s) from PDF.")
chunks = _split_text(documents)
return [
{
"id": i,
"text": c.page_content,
"metadata": {
"page_number": c.metadata.get("page", 0),
"source": c.metadata.get("source", ""),
"chunk_id": f"{paper_id}_{i}",
},
}
for i, c in enumerate(chunks)
]
# Run CPU-bound parsing in a background thread
chunk_dicts = await asyncio.to_thread(_parse_pdf_sync)
logger.info(f"[paper={paper_id}] Stage complete: Parsing → {len(chunk_dicts)} chunks ready.")
# Stage: Generating embeddings
logger.info(f"[paper={paper_id}] Stage: Generating embeddings for {len(chunk_dicts)} chunks...")
yield {
"type": "paper_progress",
"stage": "generating_embeddings",
"label": f"Generating embeddings ({len(chunk_dicts)} chunks)",
"paper_id": paper_id,
}
# Stage: Storing vectors in Qdrant
logger.info(f"[paper={paper_id}] Stage: Storing embedded chunks in Qdrant...")
yield {
"type": "paper_progress",
"stage": "storing_vectors",
"label": "Storing vectors in database",
"paper_id": paper_id,
}
await store_chunks(paper_id, chunk_dicts)
logger.info(f"[paper={paper_id}] Paper processing complete!")
yield {
"type": "paper_progress",
"stage": "complete",
"label": "Paper ready for use",
"paper_id": paper_id,
}
finally:
os.unlink(tmp_path)
async def process_pdf_file(
paper_id: str,
pdf_bytes: bytes,
) -> list[dict]:
"""
Legacy method — same logic but no progress yield.
Process PDF bytes: chunk text -> store in Qdrant
Returns an empty list for backward compatibility with router signatures.
"""
# Write to temp file for loaders
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
tmp.write(pdf_bytes)
tmp_path = tmp.name
try:
import asyncio
logger.info(f"[paper={paper_id}] Processing PDF (legacy mode)...")
def _parse_pdf_sync():
loader = PyMuPDFLoader(tmp_path)
documents = loader.load()
logger.info(f"[paper={paper_id}] Loaded {len(documents)} page(s).")
chunks = _split_text(documents)
return [
{
"id": i,
"text": c.page_content,
"metadata": {
"page_number": c.metadata.get("page", 0),
"source": c.metadata.get("source", ""),
"chunk_id": f"{paper_id}_{i}",
},
}
for i, c in enumerate(chunks)
]
chunk_dicts = await asyncio.to_thread(_parse_pdf_sync)
logger.info(f"[paper={paper_id}] Generated {len(chunk_dicts)} chunks.")
await store_chunks(paper_id, chunk_dicts)
logger.info(f"[paper={paper_id}] Processing complete (legacy).")
return []
finally:
os.unlink(tmp_path)
async def download_pdf_from_url(url: str) -> bytes:
"""Download PDF from URL, handling arXiv abs->pdf redirect."""
logger.info(f"Downloading PDF from: {url}")
# Handle arXiv abs -> pdf link conversion
if "arxiv.org/abs/" in url:
url = url.replace("/abs/", "/pdf/")
if not url.endswith(".pdf") and "arxiv.org" in url and "/pdf/" not in url:
pass
async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client:
try:
headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"}
r = await client.get(url, headers=headers)
r.raise_for_status()
content_type = r.headers.get("content-type", "").lower()
if "application/pdf" not in content_type and len(r.content) < 1000:
logger.warning(f"Downloaded content from {url} might not be a PDF (Content-Type: {content_type})")
logger.info(f"Successfully downloaded PDF ({len(r.content)} bytes)")
return r.content
except Exception as e:
logger.error(f"Failed to download PDF from {url}: {str(e)}")
raise Exception(f"Could not download PDF: {str(e)}") |