Pebblo: Policy enforcement in Safe Retriever

Author: Raj725

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

@@ -102,18 +102,27 @@ def _call(
         _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
         question = inputs[self.input_key]
         auth_context = inputs.get(self.auth_context_key)
-        semantic_context = inputs.get(self.semantic_context_key)
+        is_privileged_user = self.pb_client.is_privileged_user(auth_context)
+        semantic_context = self.determine_semantic_context(
+            is_privileged_user, auth_context, inputs
+        )
         _, prompt_entities = self.pb_client.check_prompt_validity(question)
 
         accepts_run_manager = (
             "run_manager" in inspect.signature(self._get_docs).parameters
         )
         if accepts_run_manager:
             docs = self._get_docs(
-                question, auth_context, semantic_context, run_manager=_run_manager
+                question,
+                auth_context,
+                semantic_context,
+                is_privileged_user,
+                run_manager=_run_manager,
             )
         else:
-            docs = self._get_docs(question, auth_context, semantic_context)  # type: ignore[call-arg]
+            docs = self._get_docs(
+                question, auth_context, semantic_context, is_privileged_user
+            )  # type: ignore[call-arg]
         answer = self.combine_documents_chain.run(
             input_documents=docs, question=question, callbacks=_run_manager.get_child()
         )
@@ -155,7 +164,10 @@ async def _acall(
         _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
         question = inputs[self.input_key]
         auth_context = inputs.get(self.auth_context_key)
-        semantic_context = inputs.get(self.semantic_context_key)
+        is_privileged_user = self.pb_client.is_privileged_user(auth_context)
+        semantic_context = self.determine_semantic_context(
+            is_privileged_user, auth_context, inputs
+        )
         accepts_run_manager = (
             "run_manager" in inspect.signature(self._aget_docs).parameters
         )
@@ -164,10 +176,16 @@ async def _acall(
 
         if accepts_run_manager:
             docs = await self._aget_docs(
-                question, auth_context, semantic_context, run_manager=_run_manager
+                question,
+                auth_context,
+                semantic_context,
+                is_privileged_user,
+                run_manager=_run_manager,
             )
         else:
-            docs = await self._aget_docs(question, auth_context, semantic_context)  # type: ignore[call-arg]
+            docs = await self._aget_docs(
+                question, auth_context, semantic_context, is_privileged_user
+            )  # type: ignore[call-arg]
         answer = await self.combine_documents_chain.arun(
             input_documents=docs, question=question, callbacks=_run_manager.get_child()
         )
@@ -254,6 +272,7 @@ def from_chain_type(
             api_key=api_key,
             classifier_location=classifier_location,
             classifier_url=classifier_url,
+            app_name=app_name,
         )
         # send app discovery request
         pb_client.send_app_discover(app)
@@ -289,11 +308,13 @@ def _get_docs(
         question: str,
         auth_context: Optional[AuthContext],
         semantic_context: Optional[SemanticContext],
+        is_privileged_user: bool = False,
         *,
         run_manager: CallbackManagerForChainRun,
     ) -> List[Document]:
         """Get docs."""
-        set_enforcement_filters(self.retriever, auth_context, semantic_context)
+        if not is_privileged_user:
+            set_enforcement_filters(self.retriever, auth_context, semantic_context)
         return self.retriever.get_relevant_documents(
             question, callbacks=run_manager.get_child()
         )
@@ -303,15 +324,42 @@ async def _aget_docs(
         question: str,
         auth_context: Optional[AuthContext],
         semantic_context: Optional[SemanticContext],
+        is_privileged_user: bool = False,
         *,
         run_manager: AsyncCallbackManagerForChainRun,
     ) -> List[Document]:
         """Get docs."""
-        set_enforcement_filters(self.retriever, auth_context, semantic_context)
+        if not is_privileged_user:
+            set_enforcement_filters(self.retriever, auth_context, semantic_context)
         return await self.retriever.aget_relevant_documents(
             question, callbacks=run_manager.get_child()
         )
 
+    def determine_semantic_context(
+        self,
+        is_privileged_user: bool,
+        auth_context: Optional[AuthContext],
+        inputs: Dict[str, Any],
+    ) -> Optional[SemanticContext]:
+        """
+        Determine semantic context based on the auth_context or inputs.
+
+        Args:
+            is_privileged_user (bool): If the user is a privileged user.
+            auth_context (Optional[AuthContext]): Authentication context.
+            inputs (Dict[str, Any]): Input dictionary containing various parameters.
+
+        Returns:
+            Optional[SemanticContext]: Resolved semantic context.
+        """
+        semantic_context = None
+        if not is_privileged_user:
+            # Get semantic context from policy if present otherwise from inputs
+            semantic_context = self.pb_client.get_semantic_context(
+                auth_context
+            ) or inputs.get(self.semantic_context_key)
+        return semantic_context
+
     @staticmethod
     def _get_app_details(  # type: ignore
         app_name: str, owner: str, description: str, llm: BaseLanguageModel, **kwargs

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

@@ -1,6 +1,7 @@
 """Models for the PebbloRetrievalQA chain."""
 
-from typing import Any, List, Optional, Union
+from enum import Enum
+from typing import Any, List, Optional, Set, Union
 
 from pydantic import BaseModel
 
@@ -10,7 +11,7 @@ class AuthContext(BaseModel):
 
     name: Optional[str] = None
     user_id: str
-    user_auth: List[str]
+    user_auth: List[str] = []
     """List of user authorizations, which may include their User ID and 
     the groups they are part of"""
 
@@ -149,3 +150,54 @@ class Qa(BaseModel):
     user: str
     user_identities: Optional[List[str]]
     classifier_location: str
+
+
+class PolicyType(Enum):
+    """Enums for policy types"""
+
+    IDENTITY = "identity"
+    APPLICATION = "application"
+    COST = "cost"
+
+
+class SemanticGuardrail(BaseModel):
+    """
+    Semantic Guardrail for Entities and Topics (Restrictions).
+
+    Attributes:
+        entities (Optional[Set[str]]): A set of entity restrictions.
+        topics (Optional[Set[str]]): A set of topic restrictions.
+    """
+
+    entities: Set[str] = set()
+    topics: Set[str] = set()
+
+
+class Policy(BaseModel):
+    """
+    Policy base class.
+
+    Attributes:
+        schema_version (int): The schema version of the policy.
+        type (PolicyType): The type of policy.
+    """
+
+    schema_version: int = 1
+    type: PolicyType
+
+    class Config:
+        extra = "ignore"
+
+
+class IdentityPolicy(Policy):
+    """
+    Policy for access control.
+
+    Attributes:
+        privileged_identities (Set[str]): List of identities with privileged access.
+        user_semantic_guardrail (dict[str, SemanticGuardrail]): Mapping of identities to
+            semantic guardrail restrictions.
+    """
+
+    privileged_identities: Set[str] = set()
+    user_semantic_guardrail: dict[str, SemanticGuardrail] = {}