Xu ly SSO

search/retriever.py

@@ -1,51 +1,84 @@
 import logging
-from typing import List
+import os
+from typing import List, Optional
 from opensearchpy import OpenSearch, RequestsHttpConnection
 from core.config import settings
 from core.models import DocumentChunk
 
 logger = logging.getLogger("Retriever")
 
+
 class SearchRetriever:
     def __init__(self, index_name: str = "poc_sharepoint_docs"):
         self.index_name = index_name
         
-        # Kết nối OpenSearch
+        host = settings.opensearch_host
+        if host == "opensearch" and os.environ.get("ENV") != "docker":
+            host = "localhost"
+        
         self.client = OpenSearch(
-            hosts=[{'host': settings.opensearch_host, 'port': settings.opensearch_port}],
+            hosts=[{'host': host, 'port': settings.opensearch_port}],
             http_auth=(settings.opensearch_user, settings.opensearch_pass),
             use_ssl=False,
             verify_certs=False,
             connection_class=RequestsHttpConnection
         )
         
-        # Load Local Embedding Model (để biến câu hỏi thành vector cùng không gian với dữ liệu)
-        logger.info("Đang nạp Embedding Model cho Retriever...")
+        logger.info("Loading Embedding Model for Retriever...")
         from sentence_transformers import SentenceTransformer
         self.embedder = SentenceTransformer('keepitreal/vietnamese-sbert')
 
-    def retrieve(self, query: str, top_k: int = 5) -> List[DocumentChunk]:
+    def retrieve(self, query: str, top_k: int = 5, user_email: Optional[str] = None, is_admin: bool = False) -> List[DocumentChunk]:
         """
-        Tìm kiếm ngữ nghĩa (Semantic Search) dựa trên Vector k-NN
-        """
-        logger.info(f"Đang tìm kiếm ngữ nghĩa cho câu hỏi: '{query}'")
+        Tìm kiếm ngữ nghĩa với ACL filtering.
+        
+        Args:
+            query: Câu hỏi của user
+            top_k: Số kết quả tối đa
+            user_email: Email user để filter quyền.
+            is_admin: True = bypass ACL, thấy tất cả.
+        """
+        logger.info(f"Search: '{query[:80]}' (user={user_email or 'none'}, admin={is_admin})")
         
-        # 1. Chuyển câu hỏi thành Vector
         query_vector = self.embedder.encode(query).tolist()
         
-        # 2. Xây dựng k-NN Query cho OpenSearch
-        # Ta có thể kết hợp Hybrid Search (Vector + Text) ở đây nếu muốn
-        search_query = {
-            "size": top_k,
-            "query": {
-                "knn": {
-                    "embedding": {
-                        "vector": query_vector,
-                        "k": top_k
+        # Admin hoặc không có user_email → không filter
+        if is_admin or not user_email:
+            search_query = {
+                "size": top_k,
+                "query": {
+                    "knn": {
+                        "embedding": {
+                            "vector": query_vector,
+                            "k": top_k
+                        }
+                    }
+                }
+            }
+        else:
+            # User thường → filter theo permissions
+            search_query = {
+                "size": top_k,
+                "query": {
+                    "bool": {
+                        "must": [
+                            {
+                                "knn": {
+                                    "embedding": {
+                                        "vector": query_vector,
+                                        "k": top_k * 2
+                                    }
+                                }
+                            }
+                        ],
+                        "should": [
+                            {"term": {"permissions": "*"}},
+                            {"term": {"permissions": user_email.lower()}}
+                        ],
+                        "minimum_should_match": 1
                     }
                 }
             }
-        }
         
         try:
             response = self.client.search(
@@ -58,13 +91,12 @@ class SearchRetriever:
             
             for hit in hits:
                 source = hit["_source"]
-                # Chuyển từ JSON sang DocumentChunk model
                 chunk = DocumentChunk(**source)
                 results.append(chunk)
                 
-            logger.info(f"Tìm thấy {len(results)} đoạn văn phù hợp nhất.")
+            logger.info(f"Found {len(results)} chunks")
             return results
             
         except Exception as e:
-            logger.error(f"Lỗi khi truy vấn OpenSearch: {e}")
+            logger.error(f"OpenSearch query error: {e}")
             return []