Extracts weld type from query

index_retriever.py

@@ -65,9 +65,96 @@ def rerank_nodes(query, nodes, reranker, top_k=25, min_score_threshold=0.5):
         log_message(f"Ошибка переранжировки: {str(e)}")
         return nodes[:top_k]
 
+def extract_weld_type_from_query(query):
+    """Extract welded joint type (С-XX, У-XX, etc.) from query"""
+    import re
+    
+    # Pattern for Russian weld types: С-25, У-12, Т-5, etc.
+    patterns = [
+        r'[СУТ]-\d+',  # Matches С-25, У-12, Т-5
+        r'(?:тип|тип[а-я]*)\s+([СУТ]-\d+)',  # "тип С-25" or "тип: С-25"
+    ]
+    
+    for pattern in patterns:
+        match = re.search(pattern, query, re.IGNORECASE)
+        if match:
+            if '-' in match.group(0):
+                return match.group(0).upper()
+            elif len(match.groups()) > 0:
+                return match.group(1).upper()
+    
+    return None
+
+
+def retrieve_nodes_with_weld_type_priority(query, vector_index, hybrid_retriever, reranker, top_k=20):
+    """
+    Enhanced retrieval that prioritizes welded joint type matches
+    """
+    from utils import deduplicate_nodes
+    
+    log_message(f"Enhanced retrieval for query: {query}")
+    
+    # Step 1: Try to extract weld type from query
+    weld_type = extract_weld_type_from_query(query)
+    
+    if weld_type:
+        log_message(f"Detected weld type in query: {weld_type}")
+        
+        # Step 2: Direct lookup in docstore for this weld type
+        direct_matches = []
+        all_nodes = list(vector_index.docstore.docs.values())
+        
+        for node in all_nodes:
+            metadata = node.metadata if hasattr(node, 'metadata') else {}
+            
+            # Check if this is a table node with matching weld type
+            if metadata.get('type') == 'table':
+                table_num = metadata.get('table_number', '')
+                table_title = metadata.get('table_title', '')
+                
+                # Check multiple fields for the weld type
+                if (weld_type in str(table_num) or 
+                    weld_type in str(table_title) or 
+                    weld_type in str(metadata.get('section', ''))):
+                    direct_matches.append(node)
+                    log_message(f"  Direct match found: {metadata.get('document_id')} - {table_title}")
+        
+        if direct_matches:
+            # Remove duplicates
+            direct_matches = deduplicate_nodes(direct_matches)
+            log_message(f"Found {len(direct_matches)} direct matches for {weld_type}")
+            
+            # Add some context from hybrid retriever
+            hybrid_results = hybrid_retriever.retrieve(query)
+            
+            # Combine: prioritize direct matches, supplement with hybrid results
+            combined = direct_matches + hybrid_results
+            combined = deduplicate_nodes(combined)
+            
+            # Rerank combined results
+            reranked = rerank_nodes(query, combined, reranker, top_k=top_k)
+            log_message(f"Combined retrieval: {len(direct_matches)} direct + hybrid, returning {len(reranked)} reranked")
+            
+            return reranked
+    
+    # Step 3: Fall back to normal hybrid retrieval if no weld type found
+    log_message("No weld type detected, using standard hybrid retrieval")
+    retrieved_nodes = hybrid_retriever.retrieve(query)
+    retrieved_nodes = deduplicate_nodes(retrieved_nodes)
+    reranked_nodes = rerank_nodes(query, retrieved_nodes, reranker, top_k=top_k)
+    
+    return reranked_nodes
+
+
+# Update create_query_engine to use the enhanced retrieval
 def create_query_engine(vector_index):
     try:
         from config import CUSTOM_PROMPT
+        from llama_index.core.prompts import PromptTemplate
+        from llama_index.core.response_synthesizers import get_response_synthesizer, ResponseMode
+        from llama_index.core.query_engine import RetrieverQueryEngine
+        from llama_index.retrievers.bm25 import BM25Retriever
+        from llama_index.core.retrievers import QueryFusionRetriever, VectorIndexRetriever
         
         bm25_retriever = BM25Retriever.from_defaults(
             docstore=vector_index.docstore, 
@@ -92,14 +179,29 @@ def create_query_engine(vector_index):
             text_qa_template=custom_prompt_template
         )
         
-        query_engine = RetrieverQueryEngine(
+        # Create custom query engine with enhanced retrieval
+        class EnhancedRetrieverQueryEngine(RetrieverQueryEngine):
+            def __init__(self, retriever, response_synthesizer, vector_index, reranker):
+                super().__init__(retriever=retriever, response_synthesizer=response_synthesizer)
+                self.vector_index = vector_index
+                self.reranker = reranker
+            
+            def retrieve(self, query):
+                """Override retrieve to use enhanced weld-type-aware retrieval"""
+                return retrieve_nodes_with_weld_type_priority(
+                    query, self.vector_index, self.retriever, self.reranker, top_k=20
+                )
+        
+        query_engine = EnhancedRetrieverQueryEngine(
             retriever=hybrid_retriever,
-            response_synthesizer=response_synthesizer
+            response_synthesizer=response_synthesizer,
+            vector_index=vector_index,
+            reranker=None  # Will be passed in later
         )
         
-        log_message("Query engine успешно создан")
+        log_message("Enhanced query engine created with weld-type prioritization")
         return query_engine
         
     except Exception as e:
-        log_message(f"Ошибка создания query engine: {str(e)}")
+        log_message(f"Error creating enhanced query engine: {str(e)}")
         raise