fix: [LAW] Fix pipeline vector search and graph search

DinhManhAVG · DinhManhAVG · commit e74e5b2b04e3 · 2025-12-05T17:06:55.000+07:00
diff --git a/ai_service/app/pipelines/answer_composer.py b/ai_service/app/pipelines/answer_composer.py
@@ -2,6 +2,7 @@
 import google.generativeai as genai
 from typing import Dict, Any, List
 from ..dependencies import get_settings
+from ..utils.stream_utils import create_metadata_chunk
 
 log = logging.getLogger(__name__)
 
@@ -12,19 +13,42 @@ def __init__(self):
              raise ValueError("GEMINI_API_KEY not found")
         genai.configure(api_key=settings.gemini_api_key)
         self.model_name = 'gemini-2.5-pro' # Or settings.GEMINI_COMPOSER_MODEL
+    
+    def _build_citation_from_node_id(self, node_id: str) -> str:
+        """Parse node_id to build citation title"""
+        import re
+        parts = []
+        
+        article_match = re.search(r'_Dieu_(\d+)', node_id)
+        clause_match = re.search(r'_Khoan_(\d+)', node_id)
+        point_match = re.search(r'_Diem_([a-z])', node_id)
+        doc_match = re.match(r'^(luat|nghi-dinh|thong-tu)-(\d+)-(\d+)-([A-Z\-]+)', node_id)
+        
+        if clause_match:
+            parts.append(f"Khoản {clause_match.group(1)}")
+        if point_match:
+            parts.append(f"Điểm {point_match.group(1)}")
+        if article_match:
+            parts.append(f"Điều {article_match.group(1)}")
+        if doc_match:
+            doc_type, number, year, issuer = doc_match.groups()
+            type_map = {"luat": "Luật", "nghi-dinh": "Nghị định", "thong-tu": "Thông tư"}
+            parts.append(f"của {type_map.get(doc_type, 'Văn bản')} {number}/{year}/{issuer}")
+        
+        return ", ".join(parts) if parts else f"Node {node_id}"
 
     def compose(
         self, 
         question: str, 
-        graph_result: Dict[str, Any], 
+        search_result: Dict[str, Any],  # Chứa kết quả từ vector + graph search
         web_context: str = ""
     ) -> str:
         """
         Tổng hợp câu trả lời cuối cùng.
         """
         log.info(f"Composing answer for question: {question}")
-        final_choice = graph_result.get("final_choice")
-        all_candidates = graph_result.get("all_candidates", [])
+        final_choice = search_result.get("final_choice")
+        all_candidates = search_result.get("all_candidates", [])
         
         # Xây dựng context từ Graph
         graph_context_str = ""
@@ -79,6 +103,31 @@ def compose(
         else:
             graph_context_str = "Không tìm thấy thông tin trong cơ sở dữ liệu luật nội bộ."
 
+        # Build citations từ all_candidates cho FE
+        citations = []
+        if all_candidates:
+            for idx, candidate in enumerate(all_candidates, 1):
+                node_id = candidate.get("source_node_id", "")
+                
+                # Format title từ metadata
+                parts = []
+                if candidate.get("clause_number"):
+                    parts.append(f"Khoản {candidate['clause_number']}")
+                if candidate.get("point_letter"):
+                    parts.append(f"Điểm {candidate['point_letter']}")
+                if candidate.get("article_number"):
+                    parts.append(f"Điều {candidate['article_number']}")
+                if candidate.get("document_code"):
+                    parts.append(f"của {candidate['document_code']}")
+                
+                title = ", ".join(parts) if parts else self._build_citation_from_node_id(node_id)
+                
+                # Add to citations for FE
+                citations.append({
+                    "node_id": node_id,
+                    "title": title
+                })
+        
         # Xây dựng context từ Web (nếu có)
         web_context_str = ""
         if web_context:
@@ -143,22 +192,22 @@ def compose(
 
         try:
             response = model.generate_content(prompt)
-            return response.text.strip()
+            return response.text.strip(), citations
         except Exception as e:
             log.error(f"Lỗi compose answer: {e}")
-            return "Xin lỗi, hệ thống gặp sự cố khi tổng hợp câu trả lời."
+            return "Xin lỗi, hệ thống gặp sự cố khi tổng hợp câu trả lời.", []
 
     def compose_stream(
         self, 
         question: str, 
-        graph_result: Dict[str, Any], 
+        search_result: Dict[str, Any],  # Chứa kết quả từ vector + graph search
         web_context: str = ""
     ):
         """
         Tổng hợp câu trả lời cuối cùng (Streaming).
         """
-        final_choice = graph_result.get("final_choice")
-        all_candidates = graph_result.get("all_candidates", [])
+        final_choice = search_result.get("final_choice")
+        all_candidates = search_result.get("all_candidates", [])
         
         # Xây dựng context từ Graph (same logic as compose)
         graph_context_str = ""
@@ -207,6 +256,33 @@ def compose_stream(
         else:
             graph_context_str = "Không tìm thấy thông tin trong cơ sở dữ liệu luật nội bộ."
 
+        # Build citations từ all_candidates (duplicated from compose for consistency)
+        citations = []
+        log.info(f"DEBUG: all_candidates count = {len(all_candidates)}")
+        if all_candidates:
+            log.info(f"DEBUG: Building citations from {len(all_candidates)} candidates")
+            for idx, candidate in enumerate(all_candidates, 1):
+                node_id = candidate.get("source_node_id", "")
+                
+                # Format title từ metadata
+                parts = []
+                if candidate.get("clause_number"):
+                    parts.append(f"Khoản {candidate['clause_number']}")
+                if candidate.get("point_letter"):
+                    parts.append(f"Điểm {candidate['point_letter']}")
+                if candidate.get("article_number"):
+                    parts.append(f"Điều {candidate['article_number']}")
+                if candidate.get("document_code"):
+                    parts.append(f"của {candidate['document_code']}")
+                
+                title = ", ".join(parts) if parts else self._build_citation_from_node_id(node_id)
+                
+                # Add to citations for FE
+                citations.append({
+                    "node_id": node_id,
+                    "title": title
+                })
+                
         # Xây dựng context từ Web (nếu có)
         web_context_str = ""
         if web_context:
@@ -270,6 +346,10 @@ def compose_stream(
         )
 
         try:
+            # Yield metadata first với citations
+            if citations:
+                yield create_metadata_chunk(citations)
+            
             response = model.generate_content(prompt, stream=True)
             for chunk in response:
                 if chunk.text:
diff --git a/ai_service/app/pipelines/graph_search.py b/ai_service/app/pipelines/graph_search.py
@@ -297,35 +297,50 @@ def run_direct_graph_query(
 GRAPH SCHEMA:
 {GraphSearchService._schema_snapshot}
 
-CẤU TRÚC GRAPH:
-- Document (văn bản luật): có properties: title, code, law_number, issuer
-- Article (điều): có properties: number, heading, content
-- Clause (khoản): có properties: number, content
-- Point (điểm): có properties: letter, content
-- Chapter (chương): có properties: number, title
+CẤU TRÚC GRAPH (Properties thực tế):
+- Document (văn bản luật):
+  • id: "luat-41-2024-QH15"
+  • code: "41/2024/QH15"
+  • doc_id: "luat-41-2024-QH15"
+  • title: "Luật Bảo hiểm xã hội 41/2024/QH15"
+  • type: "Luật"
+  • effectiveDate, issuedDate
+
+- Chapter (chương): number, title, content
+- Article (điều): number, heading, content
+- Clause (khoản): number, content
+- Point (điểm): letter, content
 
 RELATIONSHIPS:
 - (Document)-[:HAS_CHAPTER]->(Chapter)
 - (Document)-[:HAS_ARTICLE]->(Article)
 - (Article)-[:HAS_CLAUSE]->(Clause)
 - (Clause)-[:HAS_POINT]->(Point)
 
+**QUAN TRỌNG - Khi tìm kiếm Document theo tên:**
+- **TUYỆT ĐỐI KHÔNG** dùng cú pháp {{title: '...'}}
+- **BẮT BUỘC** phải dùng WHERE toLower(d.title) CONTAINS toLower('...') để tìm kiếm gần đúng
+- Ví dụ: MATCH (d:Document) WHERE toLower(d.title) CONTAINS toLower('Luật Bảo hiểm xã hội')
+
 NHIỆM VỤ:
 Tạo Cypher query để trả lời câu hỏi. Các dạng query phổ biến:
 
 1. **COUNT (Đếm số lượng)**:
-   Ví dụ: "Luật X có bao nhiêu chương/điều"
+   Ví dụ: "Luật BHXH 41/2024/QH15 có bao nhiêu chương"
    ```cypher
-   MATCH (d:Document {{code: 'X'}})-[:HAS_CHAPTER]->(ch:Chapter)
-   RETURN d.title AS document_title, d.code AS document_code, count(ch) AS chapter_count
+   MATCH (d:Document)
+   WHERE toLower(d.title) CONTAINS toLower('Luật Bảo hiểm xã hội')
+     AND toLower(d.title) CONTAINS '41/2024'
+   MATCH (d)-[:HAS_CHAPTER]->(ch:Chapter)
+   RETURN d.title, d.code, count(ch) AS chapter_count
    ```
 
 2. **LIST (Liệt kê)**:
    Ví dụ: "Danh sách văn bản luật"
    ```cypher
    MATCH (d:Document)
-   RETURN d.title AS title, d.code AS code, d.law_number AS law_number
-   ORDER BY d.law_number
+   RETURN d.title, d.code, d.type
+   ORDER BY d.issuedDate DESC
    LIMIT 20
    ```
 
@@ -335,13 +350,14 @@ def run_direct_graph_query(
    MATCH (a:Article {{number: '60'}})
    OPTIONAL MATCH (a)<-[:HAS_ARTICLE]-(d:Document)
    RETURN a.number AS article_number, a.heading AS article_heading,
-          a.content AS article_content, d.title AS document_title
+          a.content AS article_content, d.title AS document_title, d.code AS document_code
    ```
 
 YÊU CẦU:
 - Chỉ trả về Cypher query, KHÔNG THÊM text giải thích
 - Query phải trả về kết quả có ý nghĩa (sử dụng AS alias rõ ràng)
 - Nếu có extracted_facts, sử dụng làm filter (WHERE)
+- **LUÔN dùng toLower() CONTAINS cho title search**, KHÔNG dùng exact match với title
 
 HÃY TẠO CYPHER QUERY:
 """
@@ -354,7 +370,7 @@ def run_direct_graph_query(
             response = model.generate_content(prompt)
             cypher_query = self._clean(response.text)
             
-            log.info(f"Generated direct Cypher query: {cypher_query[:200]}...")
+            log.info(f"Generated direct Cypher query: {cypher_query}...")
             
             # Run query
             results = self.run_cypher(cypher_query)
diff --git a/ai_service/app/services/streaming_service.py b/ai_service/app/services/streaming_service.py
@@ -89,21 +89,36 @@ async def generate_streaming_response(question: str, is_debug: bool = False) ->
                 yield create_tool_thought(AgentTool.GRAPH_SEARCH)
                 log.info("GRAPH_LOOKUP category - using direct Cypher query (skip vector search)")
                 gs = GraphSearchService()
-                graph_result = await asyncio.to_thread(
+                search_result = await asyncio.to_thread(
                     gs.run_direct_graph_query,
                     question,
                     processed['structured']
                 )
                 
+                # Check if graph search found results
+                has_results = search_result.get('final_choice') is not None and search_result.get('results')
+                web_context = ""
+                
+                if not has_results:
+                    # Fallback to web search if graph returns no results
+                    log.info("GRAPH_LOOKUP found no results, falling back to web search")
+                    yield create_tool_thought(AgentTool.WEB_SEARCH)
+                    web_context, reference_count = await StreamingService._perform_web_search(question)
+                    yield create_tool_thought(AgentTool.WEB_SEARCH, f"Số lượng website cung cấp thông tin tham khảo: {reference_count}")
+                
                 # Go directly to answer composition
                 yield create_tool_thought(AgentTool.CONCLUSION)
                 composer = AnswerComposerService()
                 
                 if is_debug:
-                    full_response = await asyncio.to_thread(composer.compose, question, graph_result, "")
+                    full_response, citations = await asyncio.to_thread(composer.compose, question, search_result, web_context)
                     yield create_content_chunk(full_response)
+                    
+                    if citations:
+                        from ..utils.stream_utils import create_metadata_chunk
+                        yield create_metadata_chunk(citations)
                 else:
-                    iterator = composer.compose_stream(question, graph_result, "")
+                    iterator = composer.compose_stream(question, search_result, web_context)
                     for chunk_text in iterator:
                         yield create_content_chunk(chunk_text)
                 return
@@ -120,9 +135,9 @@ async def generate_streaming_response(question: str, is_debug: bool = False) ->
             # 3. Graph Search (skip rerank, Gemini filtering happens inside)
             yield create_tool_thought(AgentTool.GRAPH_SEARCH)
             gs = GraphSearchService()
-            graph_result = await asyncio.to_thread(gs.run_graph_search_for_node_ids, node_ids, question)
+            search_result = await asyncio.to_thread(gs.run_graph_search_for_node_ids, node_ids, question)
             
-            final_choice = graph_result.get('final_choice')
+            final_choice = search_result.get('final_choice')
             web_context = ""
             
             # 5. Web Search Fallback
@@ -138,10 +153,14 @@ async def generate_streaming_response(question: str, is_debug: bool = False) ->
             composer = AnswerComposerService()
 
             if is_debug:
-                content = await asyncio.to_thread(composer.compose, question, graph_result, web_context)
+                content, citations = await asyncio.to_thread(composer.compose, question, search_result, web_context)
+                # Emit metadata first
+                if citations:
+                    from ..utils.stream_utils import create_metadata_chunk
+                    yield create_metadata_chunk(citations)
                 yield create_content_chunk(content)
             else:
-                iterator = composer.compose_stream(question, graph_result, web_context)
+                iterator = composer.compose_stream(question, search_result, web_context)
                 for chunk_text in iterator:
                     yield create_content_chunk(chunk_text)
 
diff --git a/ai_service/app/utils/stream_utils.py b/ai_service/app/utils/stream_utils.py
@@ -45,4 +45,15 @@ def create_content_chunk(text: str) -> str:
         },
         timestamp=time.time()
     )
+    return f"data: {chunk.json()}\n\n"
+
+def create_metadata_chunk(citations: list) -> str:
+    """Helper để tạo chunk metadata chứa citations"""
+    chunk = StreamingChunk(
+        type="metadata",
+        data={
+            "citations": citations
+        },
+        timestamp=time.time()
+    )
     return f"data: {chunk.json()}\n\n"
