Cognition integration provider

app.py

@@ -9,6 +9,8 @@
 )
 from submodules.model import session
 
+import traceback
+
 app = FastAPI()
 
 
@@ -17,6 +19,9 @@ async def handle_db_session(request: Request, call_next):
     session_token = general.get_ctx_token()
     try:
         response = await call_next(request)
+    except:
+        traceback.print_exc()
+        response = None
     finally:
         general.remove_and_refresh_session(session_token)
 
@@ -66,6 +71,7 @@ class MostSimilarByEmbeddingRequest(BaseModel):
     att_filter: Optional[List[Dict[str, Any]]] = None
     threshold: Optional[Union[float, int]] = None
     question: Optional[str] = None
+    user_id: Optional[str] = None
 
 
 @app.post("/most_similar_by_embedding")
@@ -99,6 +105,7 @@ def most_similar_by_embedding(
         request.att_filter,
         request.threshold,
         include_scores,
+        request.user_id,
     )
 
     if request.question:

neural_search/util.py

@@ -10,10 +10,18 @@
     embedding,
     record_label_association,
     record,
+    project,
+    user,
 )
-from submodules.model.enums import EmbeddingPlatform, LabelSource
+from submodules.model.cognition_objects import group_member
+from submodules.model.integration_objects.helper import (
+    REFINERY_ATTRIBUTE_ACCESS_GROUPS,
+    REFINERY_ATTRIBUTE_ACCESS_USERS,
+)
+from submodules.model.enums import EmbeddingPlatform, LabelSource, UserRoles
 
 from .similarity_threshold import SimilarityThreshold, NO_THRESHOLD_INDICATOR
+import traceback
 
 port = int(os.environ["QDRANT_PORT"])
 qdrant_client = QdrantClient(host="qdrant", port=port, timeout=60)
@@ -48,9 +56,25 @@ def most_similar_by_embedding(
     att_filter: Optional[List[Dict[str, Any]]] = None,
     threshold: Optional[float] = None,
     include_scores: bool = False,
+    user_id: Optional[str] = None,
 ) -> List[str]:
     if not is_filter_valid_for_embedding(project_id, embedding_id, att_filter):
         return []
+    if project.check_access_management_active(project_id):
+        if not user_id:
+            return []
+        requesting_user = user.get(user_id)
+        if not requesting_user:
+            return []
+        if requesting_user.role != UserRoles.ENGINEER.value:
+            check_access = True
+            group_members = group_member.get_by_user_id(user_id)
+            group_ids = [str(group_member.group_id) for group_member in group_members]
+        else:
+            check_access = False
+    else:
+        check_access = False
+
     tmp_limit = limit
     has_sub_key = embedding.has_sub_key(project_id, embedding_id)
     if has_sub_key:
@@ -66,14 +90,20 @@ def most_similar_by_embedding(
     elif similarity_threshold == NO_THRESHOLD_INDICATOR:
         similarity_threshold = None
     try:
+        _filter = __build_filter(att_filter)
+        if check_access:
+            _filter = __add_access_management_filter(_filter, group_ids, user_id)
+
         search_result = qdrant_client.search(
             collection_name=embedding_id,
             query_vector=query_vector,
-            query_filter=__build_filter(att_filter),
+            query_filter=_filter,
             limit=tmp_limit,
             score_threshold=similarity_threshold,
         )
-    except Exception:
+    except Exception as e:
+        print(f"Error during search in Qdrant: {e}", flush=True)
+        print(traceback.format_exc(), flush=True)
         return []
 
     if include_scores:
@@ -125,6 +155,34 @@ def __build_filter(att_filter: List[Dict[str, Any]]) -> models.Filter:
     return models.Filter(must=must)
 
 
+def __add_access_management_filter(
+    base_filter: models.Filter, group_ids, user_id
+) -> models.Filter:
+    access_management_filter = models.Filter(
+        should=[
+            models.FieldCondition(
+                key=REFINERY_ATTRIBUTE_ACCESS_GROUPS,
+                match=models.MatchAny(
+                    any=group_ids,
+                ),
+            ),
+            models.FieldCondition(
+                key=REFINERY_ATTRIBUTE_ACCESS_USERS,
+                match=models.MatchValue(
+                    value=user_id,
+                ),
+            ),
+        ]
+    )
+    if base_filter is None:
+        return access_management_filter
+    else:
+        return models.Filter(
+            must=base_filter,
+            should=access_management_filter,
+        )
+
+
 def __build_filter_item(filter_item: Dict[str, Any]) -> models.FieldCondition:
     if isinstance(filter_item["value"], list):
         if filter_item.get("type") == "between":