Update app.py

app.py

@@ -50,16 +50,15 @@ if "vector_store" not in st.session_state:
 
 # ----------------- Metadata Extraction -----------------
 def extract_metadata_llm(pdf_path):
-    """Extracts metadata using LLM"""
+    """Extracts metadata using LLM instead of regex."""
     with pdfplumber.open(pdf_path) as pdf:
         first_page_text = pdf.pages[0].extract_text() if pdf.pages else "No text found."
 
-    # Define metadata_prompt
+    # Define metadata prompt
     metadata_prompt = PromptTemplate(
         input_variables=["text"],
         template="""
         Given the following first page of a research paper, extract metadata **strictly in JSON format**.
-
         - If no data is found for a field, return `"Unknown"` instead.
         - The response must be valid JSON.
 
@@ -81,28 +80,34 @@ def extract_metadata_llm(pdf_path):
         """
     )
 
-
+    # Run LLM Metadata Extraction
     metadata_chain = LLMChain(llm=llm_judge, prompt=metadata_prompt, output_key="metadata")
     metadata_response = metadata_chain.invoke({"text": first_page_text})
 
-    # Ensure LLM Output is a Proper JSON String
-    json_match = re.search(r"```json\n(.*?)\n```", metadata_response["metadata"], re.DOTALL)
-    json_text = json_match.group(1) if json_match else metadata_response["metadata"]
-
+    # Handle JSON extraction from LLM response
     try:
-        metadata_dict = json.loads(json_text)
-        # Ensure all expected fields are present
-        required_fields = ["Title", "Author", "Emails", "Affiliations"]
-        for field in required_fields:
-            if field not in metadata_dict:
-                metadata_dict[field] = "Unknown"
+        # Try parsing directly as JSON
+        metadata_dict = json.loads(metadata_response["metadata"])
     except json.JSONDecodeError:
-        metadata_dict = {
-            "Title": "Unknown",
-            "Author": "Unknown",
-            "Emails": "No emails found",
-            "Affiliations": "No affiliations found"
-        }
+        # Fallback: Extract JSON using regex
+        json_match = re.search(r"```json\n(.*?)\n```", metadata_response["metadata"], re.DOTALL)
+        json_text = json_match.group(1) if json_match else metadata_response["metadata"]
+
+        try:
+            metadata_dict = json.loads(json_text)
+        except json.JSONDecodeError:
+            metadata_dict = {
+                "Title": "Unknown",
+                "Author": "Unknown",
+                "Emails": "No emails found",
+                "Affiliations": "No affiliations found"
+            }
+
+    # Ensure all required fields exist
+    required_fields = ["Title", "Author", "Emails", "Affiliations"]
+    for field in required_fields:
+        if field not in metadata_dict or not metadata_dict[field].strip():
+            metadata_dict[field] = "Unknown"
 
     return metadata_dict