File size: 3,717 Bytes
6993919 | 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 | #!/usr/bin/env python3
"""Ingest research papers from /root/backend/data/papers/ into RAG forensic_reports."""
import hashlib
import json
import os
import httpx
PAPERS_DIR = "/root/backend/data/papers"
METADATA_FILE = f"{PAPERS_DIR}/papers_metadata.json"
RAG_API = "http://localhost:8000/api/v1/rag/ingest"
CHUNK_SIZE = 2500
CHUNK_OVERLAP = 300
def main():
with open(METADATA_FILE) as f:
metadata = json.load(f)
pdfs = sorted([f for f in os.listdir(PAPERS_DIR) if f.endswith(".pdf")])
print(f"Found {len(pdfs)} PDFs, {len(metadata)} metadata entries")
ingested = 0
chunks_total = 0
errors = 0
for pdf_file in pdfs:
arxiv_id = pdf_file.replace(".pdf", "")
meta = metadata.get(arxiv_id, {})
# Extract text with pymupdf
try:
import fitz
doc = fitz.open(os.path.join(PAPERS_DIR, pdf_file))
text = ""
for page in doc:
text += page.get_text() + "\n"
doc.close()
except Exception as e:
print(f" SKIP {arxiv_id}: extract failed: {e}")
errors += 1
continue
if len(text.strip()) < 200:
print(f" SKIP {arxiv_id}: too short ({len(text)} chars)")
errors += 1
continue
# Chunk it
chunks = []
start = 0
while start < len(text):
end = min(start + CHUNK_SIZE, len(text))
chunk_text = text[start:end].strip()
if len(chunk_text) > 100:
chunks.append(chunk_text)
start += CHUNK_SIZE - CHUNK_OVERLAP
# Ingest each chunk
title = meta.get("title", arxiv_id)
attack_types = meta.get("attack_types", [])
authors_list = meta.get("authors", [])
collection = meta.get("collection", "forensic_reports")
success = 0
for i, chunk in enumerate(chunks):
doc_id = hashlib.sha256(f"{arxiv_id}:chunk:{i}".encode()).hexdigest()[:16]
payload = {
"collection": collection,
"content": chunk,
"metadata": {
"arxiv_id": arxiv_id,
"title": title,
"authors": ", ".join(authors_list[:3]) if isinstance(authors_list, list) else str(authors_list),
"attack_types": ", ".join(attack_types) if isinstance(attack_types, list) else str(attack_types),
"chunk_index": i,
"total_chunks": len(chunks),
"source": "research_paper",
"doc_id": doc_id,
},
}
try:
resp = httpx.post(RAG_API, json=payload, timeout=30)
if resp.status_code == 200:
success += 1
else:
if errors < 5:
print(f" ERR {arxiv_id} chunk {i}: {resp.status_code} {resp.text[:100]}")
errors += 1
except Exception as e:
if errors < 5:
print(f" ERR {arxiv_id} chunk {i}: {e}")
errors += 1
chunks_total += success
ingested += 1
if arxiv_id not in metadata:
metadata[arxiv_id] = {}
metadata[arxiv_id]["ingested"] = True
metadata[arxiv_id]["chunk_count"] = success
print(f" {arxiv_id}: {len(chunks)} chunks, {success} ingested — {title[:50]}")
# Save updated metadata
with open(METADATA_FILE, "w") as f:
json.dump(metadata, f, indent=2)
print(f"\nDONE: {ingested} papers, {chunks_total} total chunks, {errors} errors")
if __name__ == "__main__":
main()
|