community: Handle classification failure in PebbloSafeLoader and PebbloRetrievalQA

libs/community/langchain_community/chains/pebblo_retrieval/enforcement_filters.py

@@ -21,6 +21,7 @@
     AuthContext,
     SemanticContext,
 )
+from langchain_community.utilities.pebblo import CLASSIFICATION_UNAVAILABLE
 
 logger = logging.getLogger(__name__)
 
@@ -522,10 +523,28 @@ def _set_semantic_enforcement_filter(
     This method sets the semantic enforcement filter in the search_kwargs
     of the retriever based on the type of the vectorstore.
     """
+    # Add CLASSIFICATION_UNAVAILABLE to deny list if it's not empty
+    add_unavailable_to_deny_list(semantic_context)
+    # Apply semantic filter
     search_kwargs = retriever.search_kwargs
     if retriever.vectorstore.__class__.__name__ == PINECONE:
         _apply_pinecone_semantic_filter(search_kwargs, semantic_context)
     elif retriever.vectorstore.__class__.__name__ == QDRANT:
         _apply_qdrant_semantic_filter(search_kwargs, semantic_context)
     elif retriever.vectorstore.__class__.__name__ == PGVECTOR:
         _apply_pgvector_semantic_filter(search_kwargs, semantic_context)
+
+
+def add_unavailable_to_deny_list(sem_ctx: Optional[SemanticContext]) -> None:
+    """
+    Add CLASSIFICATION_UNAVAILABLE to deny list if it's not empty.
+    This function handles documents with missing semantic metadata.
+    """
+    if sem_ctx is None:
+        return
+    if sem_ctx.pebblo_semantic_entities and sem_ctx.pebblo_semantic_entities.deny:
+        if CLASSIFICATION_UNAVAILABLE not in sem_ctx.pebblo_semantic_entities.deny:
+            sem_ctx.pebblo_semantic_entities.deny.append(CLASSIFICATION_UNAVAILABLE)
+    if sem_ctx.pebblo_semantic_topics and sem_ctx.pebblo_semantic_topics.deny:
+        if CLASSIFICATION_UNAVAILABLE not in sem_ctx.pebblo_semantic_topics.deny:
+            sem_ctx.pebblo_semantic_topics.deny.append(CLASSIFICATION_UNAVAILABLE)

libs/community/langchain_community/document_loaders/pebblo.py

@@ -13,6 +13,7 @@
 from langchain_community.document_loaders.base import BaseLoader
 from langchain_community.utilities.pebblo import (
     APP_DISCOVER_URL,
+    CLASSIFICATION_UNAVAILABLE,
     CLASSIFIER_URL,
     LOADER_DOC_URL,
     PEBBLO_CLOUD_URL,
@@ -481,10 +482,15 @@ def _add_semantic_to_docs(self, classified_docs: Dict) -> List[Document]:
             for doc in self.docs_with_id
         }
 
-        for classified_doc in classified_docs.values():
-            doc_id = classified_doc.get("pb_id")
-            if doc_id in indexed_docs:
-                self._add_semantic_to_doc(indexed_docs[doc_id], classified_doc)
+        if not classified_docs:
+            # Add CLASSIFICATION_UNAVAILABLE to semantic metadata if no classification
+            for _, doc in indexed_docs.items():
+                self._add_semantic_to_doc(doc, {})
+        else:
+            for classified_doc in classified_docs.values():
+                doc_id = classified_doc.get("pb_id")
+                if doc_id in indexed_docs:
+                    self._add_semantic_to_doc(indexed_docs[doc_id], classified_doc)
 
         semantic_metadata_docs = [doc for doc in indexed_docs.values()]
 
@@ -506,6 +512,7 @@ def _unindex_docs(self) -> List[Document]:
     def _add_semantic_to_doc(self, doc: Document, classified_doc: dict) -> Document:
         """
         Adds semantic metadata to the given document in-place.
+        If classified_doc is empty, adds "unavailable" to semantic data.
 
         Args:
             doc (Document): A Document object.
@@ -514,12 +521,16 @@ def _add_semantic_to_doc(self, doc: Document, classified_doc: dict) -> Document:
         Returns:
             Document: The Document object with added semantic metadata.
         """
-        doc.metadata["pebblo_semantic_entities"] = list(
-            classified_doc.get("entities", {}).keys()
-        )
-        doc.metadata["pebblo_semantic_topics"] = list(
-            classified_doc.get("topics", {}).keys()
-        )
+        if classified_doc:
+            doc.metadata["pebblo_semantic_entities"] = list(
+                classified_doc.get("entities", {}).keys()
+            )
+            doc.metadata["pebblo_semantic_topics"] = list(
+                classified_doc.get("topics", {}).keys()
+            )
+        else:
+            doc.metadata["pebblo_semantic_entities"] = [CLASSIFICATION_UNAVAILABLE]
+            doc.metadata["pebblo_semantic_topics"] = [CLASSIFICATION_UNAVAILABLE]
         return doc
 
     def _add_pebblo_specific_metadata(self, classified_docs: dict) -> None:

libs/community/langchain_community/utilities/pebblo.py

@@ -59,6 +59,7 @@
 }
 
 SUPPORTED_LOADERS = (*file_loader, *dir_loader, *in_memory)
+CLASSIFICATION_UNAVAILABLE = "unavailable"
 
 logger = logging.getLogger(__name__)