Fix pdf source retrieval information

generic_rag/graphs/ret_gen.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import Any, Union
 
 from langchain import hub
@@ -53,15 +54,29 @@ class RetGenLangGraph:
         pdf_sources = {}
         for context in self.last_invoke["context"]:
             try:
-                if context.metadata["filetype"] == "application/pdf":
+                Path(context.metadata["source"]).suffix == ".pdf"
-                    source = context.metadata["source"]
+            except KeyError:
-                    page_number = context.metadata["page_number"]
+                continue
-                    if source in pdf_sources:
-                        pdf_sources[source].add(page_number)
             else:
-                        pdf_sources[source] = {page_number}
+                source = context.metadata["source"]
+
+            if source not in pdf_sources:
+                pdf_sources[source] = set()
+
+            # The page numbers are in the `page_numer` and `page` fields.
+            try:
+                page_number = context.metadata["page_number"]
             except KeyError:
                 pass
+            else:
+                pdf_sources[source].add(page_number)
+
+            try:
+                page_number = context.metadata["page"]
+            except KeyError:
+                pass
+            else:
+                pdf_sources[source].add(page_number)
 
         return pdf_sources
 
@@ -75,9 +90,10 @@ class RetGenLangGraph:
         web_sources = set()
         for context in self.last_invoke["context"]:
             try:
-                if context.metadata["filetype"] == "web":
+                context.metadata["filetype"] == "web"
-                    web_sources.add(context.metadata["source"])
             except KeyError:
-                pass
+                continue
+            else:
+                web_sources.add(context.metadata["source"])
 
         return web_sources