Add page limitation for PDF url extraction

app.py

@@ -58,8 +58,9 @@ class Query(BaseModel):
 class DocumentID(BaseModel):
     doc_id: str
 
-class WebPDF(BaseModel):
+class PDF(BaseModel):
     url: str
+    page_num: str = -1
 
 @app.post("/search")
 async def get_articles(query: Query):
@@ -113,11 +114,11 @@ async def extract_arxiv_pdf(document: DocumentID):
         return {"error": True, "message": "Error while downloading PDF: HTTP/" + str(pdf_req.status_code)}
 
 @app.post("/extract_pdf/url")
-async def extract_pdf(pdf: WebPDF):
+async def extract_pdf(pdf: PDF):
     pdf_req = requests.get(pdf.url)
     if pdf_req.status_code == 200:
         pdf_data = BytesIO(pdf_req.content)
-        doc = fitz.open(stream=pdf_data, filetype="pdf")
+        doc = fitz.open(stream=pdf_data, filetype="pdf")[:pdf.page_num]
         pdf_text = " ".join([page.get_text("text") for page in doc])
         pdf_metadata = doc.metadata
         print(pdf_metadata)