PaperSearch arXiv Retrieval Pipeline

CREATE TABLE IF NOT EXISTS query (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    query TEXT UNIQUE,
    datetime TEXT
);

CREATE TABLE IF NOT EXISTS paper_search (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    query TEXT,
    pdf_url TEXT,
    filename TEXT,
    pdf_data BLOB
);

CREATE TABLE IF NOT EXISTS document_page (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    paper_search_id INTEGER,
    page_number INTEGER,
    text TEXT,
    FOREIGN KEY (paper_search_id) REFERENCES paper_search(id)
);

CREATE TABLE IF NOT EXISTS document_page_split (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    document_page_id INTEGER,
    text TEXT,
    FOREIGN KEY (document_page_id) REFERENCES document_page(id)
);

import os
import sqlite3
import xml.etree.ElementTree as ET

import PyPDF2
import requests
from nltk.tokenize import sent_tokenize


class PaperSearch:
    def __init__(self, db_name="papers.db", data_dir="data", max_results=10, chunk_size=1200):
        self.base_url = "http://export.arxiv.org/api/query"
        self.max_results = max_results
        self.data_dir = data_dir
        self.chunk_size = chunk_size
        self.conn = sqlite3.connect(db_name)
        self.cursor = self.conn.cursor()

    def search(self, query):
        params = {"search_query": query, "start": 0, "max_results": self.max_results}
        response = requests.get(self.base_url, params=params, timeout=30)
        response.raise_for_status()
        os.makedirs(self.data_dir, exist_ok=True)

        downloaded = []
        for pdf_url in self.get_pdf_links(response.text):
            filename = self.download_pdf(pdf_url, query)
            if filename:
                downloaded.append(filename)
        return downloaded

    def get_pdf_links(self, response_text):
        root = ET.fromstring(response_text)
        links = []
        for entry in root.findall("{http://www.w3.org/2005/Atom}entry"):
            for link in entry.findall("{http://www.w3.org/2005/Atom}link"):
                if link.attrib.get("title") == "pdf":
                    links.append(link.attrib["href"])
        return links

    def download_pdf(self, pdf_url, query):
        self.cursor.execute("SELECT id FROM paper_search WHERE pdf_url = ?", (pdf_url,))
        if self.cursor.fetchone():
            return None

        filename = os.path.join(self.data_dir, pdf_url.rsplit("/", 1)[-1] + ".pdf")
        response = requests.get(pdf_url, timeout=60)
        response.raise_for_status()

        with open(filename, "wb") as f:
            f.write(response.content)

        self.cursor.execute(
            "INSERT INTO paper_search (query, pdf_url, filename, pdf_data) VALUES (?, ?, ?, ?)",
            (query, pdf_url, filename, response.content),
        )
        self.conn.commit()
        self.extract_text_from_pdf(filename)
        return filename

    def extract_text_from_pdf(self, pdf_file):
        self.cursor.execute("SELECT id FROM paper_search WHERE filename = ?", (pdf_file,))
        paper_search_id = self.cursor.fetchone()[0]

        with open(pdf_file, "rb") as file:
            reader = PyPDF2.PdfReader(file)
            for page_number, page in enumerate(reader.pages, start=1):
                text = page.extract_text() or ""
                self.cursor.execute(
                    "INSERT INTO document_page (paper_search_id, page_number, text) VALUES (?, ?, ?)",
                    (paper_search_id, page_number, text),
                )
                page_id = self.cursor.lastrowid
                self.split_text_on_sentences(page_id, text)
        self.conn.commit()

    def split_text_on_sentences(self, page_id, text):
        chunks, current = [], ""
        for sentence in sent_tokenize(text):
            if len(current) + len(sentence) <= self.chunk_size:
                current += " " + sentence
            else:
                chunks.append(current.strip())
                current = sentence
        if current:
            chunks.append(current.strip())

        self.cursor.executemany(
            "INSERT INTO document_page_split(document_page_id, text) VALUES (?, ?)",
            [(page_id, chunk) for chunk in chunks if chunk],
        )
        return chunks

search = PaperSearch(db_name="papers.db", data_dir="papers", max_results=20)
downloaded_files = search.search("retrieval augmented generation")