Update utils.py

utils.py

@@ -195,73 +195,112 @@ def debug_search_tables(vector_index, search_term="С-25"):
     
     return matching
 
+GENERIC_STEEL_CONTEXT = "стандарт ГОСТ технические условия марка материал применение сварка"
+
+from config import QUERY_EXPANSION_PROMPT
 from documents_prep import normalize_text, normalize_steel_designations
 
-def expand_query_with_llm(query, llm_model):
-    """Generate 5 alternative query formulations using LLM"""
-    try:
-        from config import QUERY_EXPANSION_PROMPT
-        
-        expansion_prompt = QUERY_EXPANSION_PROMPT.format(original_query=query)
+STEEL_PRODUCT_EXPANSIONS = {
+    "08X18H10T": ["Листы", "Трубы", "Поковки", "Крепежные изделия", "Сортовой прокат", "Отливки"],
+    "12X18H10T": ["Листы", "Поковки", "Сортовой прокат"],
+    "10X17H13M2T": ["Трубы", "Арматура", "Поковки", "Фланцы"],
+    "20X23H18": ["Листы", "Сортовой прокат", "Поковки"],
+    "03X17H14M3": ["Трубы", "Листы", "Проволока"]
+}
+
+
+def enhance_query_for_steel_grades(query):
+    """Expand query with steel grade specific context"""
+    import re    
+    
+    # FIX: Use the same pattern as normalize_steel_designations
+    # Pattern for regular steel grades: 08X18H10T, 12X18H10T, etc.
+    steel_pattern = r'\b\d{1,3}(?:[A-ZА-ЯЁ]\d*)+\b'
+    # Pattern for welding wires: СВ-08X19H10, CB-08X19H10
+    wire_pattern = r'\b[СC][ВB]-\d{1,3}(?:[A-ZА-ЯЁ]\d*)+\b'
+    
+    matches = re.findall(steel_pattern, query, re.IGNORECASE)
+    wire_matches = re.findall(wire_pattern, query, re.IGNORECASE)
+    
+    all_matches = matches + wire_matches
+
+    if not all_matches:
+        return query
+    
+    # Collect context expansions
+    added_context = []
+    grades_found = []
+    
+    for match in all_matches:
+        match_upper = match.upper()
+        grades_found.append(match_upper)
         
-        log_message(f"Generating query variations for: {query}")
-        response = llm_model.complete(expansion_prompt)
+        # Check if we have specific context for this grade
+        if match_upper in STEEL_PRODUCT_EXPANSIONS:
+            context = ' '.join(STEEL_PRODUCT_EXPANSIONS[match_upper])
+            added_context.append(context)
+            log_message(f"  Found specific context for {match_upper}: {context}")
+        else:
+            # Use generic context for unknown grades
+            added_context.append(GENERIC_STEEL_CONTEXT)
+            log_message(f"  Using generic context for {match_upper}")
+    
+    # Build enhanced query
+    if added_context:
+        # Remove duplicates from context
+        unique_context = ' '.join(set(' '.join(added_context).split()))
+        enhanced = f"{query} {unique_context}"
         
-        # Parse response - split by newlines and filter empty
-        variations = [line.strip() for line in response.text.split('\n') if line.strip()]
-        variations = variations[:5]  # Take only first 5
+        log_message(f"Enhanced query for steel grades: {', '.join(grades_found)}")
+        log_message(f"Added context: {unique_context[:100]}...")
         
-        if variations:
-            log_message(f"Generated {len(variations)} query variations:")
-            for i, var in enumerate(variations, 1):
-                log_message(f"  {i}. {var}")
-            
-            # Combine original + variations
-            combined_query = query + " " + " ".join(variations)
-            return combined_query
-        else:
-            log_message("No variations generated, using original query")
-            return query
-            
-    except Exception as e:
-        log_message(f"Error generating query variations: {e}")
-        return query
+        return enhanced
+    
+    return query
 
 
 def answer_question(question, query_engine, reranker, current_model, chunks_df=None, rerank_top_k=20):
     
-    # Apply normalizations
     normalized_question = normalize_text(question)
-    normalized_question_2, query_changes, change_list = normalize_steel_designations(normalized_question)
+    normalized_question_2, query_changes, change_list = normalize_steel_designations(question)
+    
+    # Step 1: Keyword-based enhancement (existing)
+    enhanced_question = enhance_query_for_steel_grades(normalized_question_2)
+    
+    # Step 2: LLM-based query expansion (NEW)
+    try:
+        llm = get_llm_model(current_model)
+        expansion_prompt = QUERY_EXPANSION_PROMPT.format(original_query=enhanced_question)
+        expanded_queries = llm.complete(expansion_prompt).text.strip()
+        
+        # Combine original + expanded queries
+        enhanced_question = f"{enhanced_question} {expanded_queries}"
+        log_message(f"LLM expanded query: {expanded_queries[:200]}...")
+    except Exception as e:
+        log_message(f"Query expansion failed: {e}, using keyword-only enhancement")
     
     if change_list:
         log_message(f"Query changes: {', '.join(change_list)}")
-    
+    if change_list:
+        log_message(f"Query changes: {', '.join(change_list)}")    
     if query_engine is None:
         return "<div style='background-color: #e53e3e; color: white; padding: 20px; border-radius: 10px;'>Система не инициализирована</div>", "", ""
     
     try:
         start_time = time.time()
-        
-        # EXPAND QUERY USING LLM
-        from utils import get_llm_model
-        llm = get_llm_model(current_model)
-        expanded_query = expand_query_with_llm(normalized_question_2, llm)
-        
-        # Use expanded query for retrieval
-        retrieved_nodes = query_engine.retriever.retrieve(expanded_query)
-        
+        retrieved_nodes = query_engine.retriever.retrieve(enhanced_question)
         log_message(f"user query: {question}")
         log_message(f"normalized query: {normalized_question}")
         log_message(f"after steel normalization: {normalized_question_2}")
-        log_message(f"expanded query length: {len(expanded_query)} chars")
+        log_message(f"enhanced query: {enhanced_question}") 
         log_message(f"Steel grades normalized in query: {query_changes}")
         
         log_message(f"RETRIEVED: {len(retrieved_nodes)} nodes")
         
         unique_retrieved = deduplicate_nodes(retrieved_nodes)
+
+        # IMPROVED DEBUG: Log what was actually retrieved with FULL metadata
         log_message(f"RETRIEVED: unique {len(unique_retrieved)} nodes")
-        
         for i, node in enumerate(unique_retrieved):
             node_type = node.metadata.get('type', 'text')
             doc_id = node.metadata.get('document_id', 'N/A')
@@ -270,6 +309,7 @@ def answer_question(question, query_engine, reranker, current_model, chunks_df=N
                 table_num = node.metadata.get('table_number', 'N/A')
                 table_id = node.metadata.get('table_identifier', 'N/A')
                 table_title = node.metadata.get('table_title', 'N/A')
+                # Show first 200 chars of content to verify it's the right table
                 content_preview = node.text[:200].replace('\n', ' ')
                 log_message(f"  [{i+1}] {doc_id} - Table {table_num} | ID: {table_id}")
                 log_message(f"      Title: {table_title[:80]}")
@@ -280,11 +320,12 @@ def answer_question(question, query_engine, reranker, current_model, chunks_df=N
         
         log_message(f"UNIQUE NODES: {len(unique_retrieved)} nodes")
         
-        reranked_nodes = rerank_nodes(normalized_question_2, unique_retrieved, reranker, 
-                                     top_k=rerank_top_k)
+        # Simple reranking with NORMALIZED question and PARAMETERIZED top_k
+        reranked_nodes = rerank_nodes(enhanced_question, unique_retrieved, reranker, 
+                                     top_k=rerank_top_k)  # NOW PARAMETERIZED
         
-        # Use ORIGINAL normalized question for final answer generation
-        response = query_engine.query(normalized_question_2)
+        # Direct query without formatting - use normalized question
+        response = query_engine.query(enhanced_question)
         
         end_time = time.time()
         processing_time = end_time - start_time