simplest version

utils.py

@@ -15,6 +15,9 @@ def get_reranker_model(model_name='cross-encoder/ms-marco-MiniLM-L-12-v2'):
     """Get reranker model"""
     return CrossEncoder(model_name)
 
+
+
+
 def format_sources(nodes):
     """Format retrieved sources for display"""
     sources = []
@@ -40,74 +43,96 @@ def preprocess_query(question):
     import re
     
     question_lower = question.lower()
-    
-    # Extract document ID and normalize
-    doc_match = re.search(r'(гост|нп|му)\s*р?\s*[№-]*\s*([0-9\.-]+)', question_lower)
-    
     enhanced_query = question
     
+    # Detect "list all tables" queries - handle differently
+    if any(phrase in question_lower for phrase in ['какие таблиц', 'список таблиц', 'перечисл', 'все таблиц']):
+        # For listing queries, just extract document ID
+        doc_match = re.search(r'(гост|нп|му)[^\s]*\s*р?\s*[№-]*\s*([0-9\.-]+)', question_lower)
+        if doc_match:
+            doc_id = f"{doc_match.group(1).upper()} Р {doc_match.group(2)}"
+            enhanced_query = f"документ {doc_id} таблица"
+        return enhanced_query
+    
+    # For specific table queries
+    table_match = re.search(r'табли[цу]\w*\s+(?:№|номер)?\s*([а-яa-z0-9\.]+)', question_lower)
+    if table_match:
+        table_num = table_match.group(1).upper()
+        enhanced_query += f" таблица {table_num}"
+    
+    # Document detection
+    doc_match = re.search(r'(гост|нп|му)[^\s]*\s*р?\s*[№-]*\s*([0-9\.-]+)', question_lower)
     if doc_match:
-        doc_type = doc_match.group(1).upper()
-        doc_num = doc_match.group(2)
-        
-        # Add normalized versions
-        enhanced_query += f" {doc_type} Р {doc_num}"
+        doc_id = f"{doc_match.group(1).upper()} Р {doc_match.group(2)}"
+        enhanced_query += f" документ {doc_id}"
     
     return enhanced_query
 
 def answer_question(question, query_engine, reranker):
     try:
-        log_message(f"Query: {question}")
+        log_message(f"\n{'='*70}")
+        log_message(f"QUERY: {question}")
         
         enhanced_query = preprocess_query(question)
-        if enhanced_query != question:
-            log_message(f"Enhanced query: {enhanced_query}")
+        log_message(f"Enhanced: {enhanced_query}")
+        
+        # Detect listing queries - need MORE chunks
+        is_listing_query = any(phrase in question.lower() 
+                              for phrase in ['какие таблиц', 'список', 'перечисл', 'все таблиц'])
         
         retrieved = query_engine.retriever.retrieve(enhanced_query)
-        log_message(f"Retrieved {len(retrieved)} nodes")
+        log_message(f"\nRETRIEVED: {len(retrieved)} nodes")
         
+        # Log retrieved docs
         doc_stats = {}
         for n in retrieved:
             doc_id = n.metadata.get('document_id', 'unknown')
-            doc_type = n.metadata.get('type', 'text')
+            doc_group = n.metadata.get('document_group', doc_id)
             
-            if doc_id not in doc_stats:
-                doc_stats[doc_id] = {'tables': set(), 'text': 0, 'images': 0}
+            if doc_group not in doc_stats:
+                doc_stats[doc_group] = {'tables': set(), 'text': 0}
             
-            if doc_type == 'table':
+            if n.metadata.get('type') == 'table':
                 table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
-                doc_stats[doc_id]['tables'].add(table_id)
-            elif doc_type == 'image':
-                doc_stats[doc_id]['images'] += 1
+                doc_stats[doc_group]['tables'].add(table_id)
             else:
-                doc_stats[doc_id]['text'] += 1
+                doc_stats[doc_group]['text'] += 1
         
         for doc_id in sorted(doc_stats.keys()):
             stats = doc_stats[doc_id]
-            parts = []
+            log_message(f"  {doc_id}: {len(stats['tables'])} tables, {stats['text']} text")
             if stats['tables']:
-                parts.append(f"tables={list(stats['tables'])[:5]}")
-            if stats['text']:
-                parts.append(f"text={stats['text']}")
-            if stats['images']:
-                parts.append(f"images={stats['images']}")
-            log_message(f"  {doc_id}: {', '.join(parts)}")
-
-        doc_ids = [n.metadata.get('document_id', 'unknown') for n in retrieved]
-        table_nums = [n.metadata.get('table_number', '') for n in retrieved if n.metadata.get('type') == 'table']
-        log_message(f"Retrieved from documents: {set(doc_ids)}")
-        if table_nums:
-            log_message(f"Retrieved tables: {set(table_nums)}")
+                log_message(f"    Tables: {sorted(stats['tables'])}")
         
-        reranked = rerank_nodes(question, retrieved, reranker, top_k=25)
-        log_message(f"Reranked to {len(reranked)} nodes")
+        # Adjust reranking based on query type
+        if is_listing_query:
+            reranked = rerank_nodes(question, retrieved, reranker, top_k=50, min_score=0.2)
+        else:
+            reranked = rerank_nodes(question, retrieved, reranker, top_k=25, min_score=0.3)
         
-        doc_ids_reranked = [n.metadata.get('document_id', 'unknown') for n in reranked]
-        table_nums_reranked = [n.metadata.get('table_number', '') for n in reranked if n.metadata.get('type') == 'table']
-        log_message(f"After reranking - documents: {set(doc_ids_reranked)}")
-        if table_nums_reranked:
-            log_message(f"After reranking - tables: {set(table_nums_reranked)}")
+        log_message(f"\nRERANKED: {len(reranked)} nodes")
         
+        # Log reranked
+        doc_stats_reranked = {}
+        for n in reranked:
+            doc_group = n.metadata.get('document_group', n.metadata.get('document_id', 'unknown'))
+            
+            if doc_group not in doc_stats_reranked:
+                doc_stats_reranked[doc_group] = {'tables': set(), 'text': 0}
+            
+            if n.metadata.get('type') == 'table':
+                table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
+                doc_stats_reranked[doc_group]['tables'].add(table_id)
+            else:
+                doc_stats_reranked[doc_group]['text'] += 1
+        
+        for doc_id in sorted(doc_stats_reranked.keys()):
+            stats = doc_stats_reranked[doc_id]
+            log_message(f"  {doc_id}: {len(stats['tables'])} tables, {stats['text']} text")
+            if stats['tables']:
+                log_message(f"    Tables: {sorted(stats['tables'])}")
+        
+        # Build context
         context_parts = []
         for n in reranked:
             meta = n.metadata
@@ -115,103 +140,48 @@ def answer_question(question, query_engine, reranker):
             doc_type = meta.get('type', 'text')
             
             if doc_type == 'table':
-                table_num = meta.get('table_number', 'unknown')
+                table_id = meta.get('table_identifier', meta.get('table_number', 'unknown'))
                 title = meta.get('table_title', '')
-                source_label = f"[ТАБЛИЦА {table_num} - {doc_id}]"
+                source_label = f"[{doc_id} - Таблица {table_id}]"
                 if title:
                     source_label += f" {title}"
-            elif doc_type == 'image':
-                img_num = meta.get('image_number', 'unknown')
-                source_label = f"[РИСУНОК {img_num} - {doc_id}]"
             else:
-                section = meta.get('section_id', '')
-                source_label = f"[{doc_id} - {section}]"
+                source_label = f"[{doc_id}]"
             
-            context_parts.append(f"{source_label}\n{n.text}")
+            context_parts.append(f"{source_label}\n{n.text[:500]}")  # Limit context per chunk
         
-        context = "\n\n" + ("="*70 + "\n\n").join(context_parts)
-        from config import CUSTOM_PROMPT
-        prompt = f"""
-Вы являетесь высокоспециализированным Ассистентом для анализа нормативных документов (AIEXP). Ваша цель - предоставлять точные, корректные и контекстно релевантные ответы исключительно на основе предоставленного контекста из нормативной документации.
-
-ПРАВИЛА АНАЛИЗА ЗАПРОСА:
-
-1. ПРЯМЫЕ ВОПРОСЫ БЕЗ ДОКУМЕНТАЛЬНОГО КОНТЕКСТА:
-   Если пользователь задает вопрос типа "В каких случаях могут быть признаны протоколы испытаний?" без предоставления дополнительных документов, найдите соответствующую информацию в доступном контексте и предоставьте полный ответ с указанием источников.
-
-2. ОПРЕДЕЛЕНИЕ ТИПА ЗАДАЧИ:
-
-   а) ПОИСК И ОТВЕТ НА ВОПРОС (ключевые слова: "в каких случаях", "когда", "кто", "что", "как", "почему"):
-      - Найдите релевантную информацию в контексте
-      - Предоставьте развернутый ответ
-      - Обязательно укажите конкретные документы и разделы
-      - Процитируйте ключевые положения
-
-   б) КРАТКОЕ САММАРИ (ключевые слова: "кратко", "суммировать", "резюме", "основные моменты"):
-      - Предоставьте структурированное резюме
-      - Выделите ключевые требования
-      - Используйте нумерованный список
-
-   в) ПОИСК ДОКУМЕНТА И ПУНКТА (ключевые слова: "найти", "где", "какой документ", "в каком разделе"):
-      - Укажите конкретный документ и структурное расположение
-      - Предоставьте точные номера разделов/пунктов
-
-   г) ПРОВЕРКА КОРРЕКТНОСТИ (ключевые слова: "правильно ли", "соответствует ли", "проверить"):
-      - Четко укажите: "СООТВЕТСТВУЕТ" или "НЕ СООТВЕТСТВУЕТ"
-      - Перечислите конкретные требования
-
-   д) ПЛАН ДЕЙСТВИЙ (ключевые слова: "план", "алгоритм", "пошагово"):
-      - Создайте пронумерованный план
-      - Укажите ссылки на соответствующие пункты НД
-
-ПРАВИЛА ФОРМИРОВАНИЯ ОТВЕТОВ:
-
-Работай исключительно с информацией из предоставленного контекста. Запрещено использовать:
-- Общие знания
-- Информацию из интернета
-- Данные из предыдущих диалогов
-- Собственные предположения
-
-1. СТРУКТУРА ОТВЕТА:
-   - Начинайте с прямого ответа на вопрос
-   - Затем указывайте нормативные основания
-   - Завершайте ссылками на конкретные документы и разделы
-
-2. РАБОТА С КОНТЕКСТОМ:
-   - Если информация найдена в контексте - предоставьте полный ответ
-   - Если информация не найдена: "Информация по вашему запросу не найдена в доступной нормативной документации"
-   - Не делайте предположений за пределами контекста
-   - Не используйте общие знания
-
-3. ТЕРМИНОЛОГИЯ И ЦИТИРОВАНИЕ:
-   - Сохраняйте официальную терминологию НД
-   - Цитируйте точные формулировки ключевых требований
-   - При множественных источниках - укажите все релевантные
-
-4. ФОРМАТИРОВАНИЕ:
-   - Для перечислений: используйте нумерованные списки
-   - Выделяйте критически важные требования
-   - Структурируйте ответ логически
-
-# КАК РАБОТАТЬ С ЗАПРОСОМ
+        context = "\n\n" + ("="*50 + "\n\n").join(context_parts)
+        
+        # Adjust prompt for listing queries
+        if is_listing_query:
+            prompt = f"""Контекст содержит информацию о таблицах из документов.
 
-**Шаг 1:** Определи, что именно ищет пользователь (термин, требование, процедура, условие)
+КОНТЕКСТ:
+{context}
 
-**Шаг 2:** Найди релевантную информацию в контексте
+ВОПРОС: {question}
 
-**Шаг 3:** Сформируй ответ:
-   - Если нашел: укажи документ и пункт, процитируй нужную часть
-   - Если не нашел: четко сообщи об отсутствии информации
+ИНСТРУКЦИИ:
+1. Перечисли ВСЕ таблицы, найденные в контексте для запрошенного документа
+2. Укажи номер таблицы и название (если есть)
+3. Если таблиц нет - скажи прямо
 
-**Шаг 4:** При наличии нескольких источников:
-   - Представь их последовательно с указанием источника каждого
-   - Если источников много (>4) — сначала дай их список, потом цитаты   
+ОТВЕТ (список таблиц):"""
+        else:
+            prompt = f"""Ты эксперт по технической документации.
 
 КОНТЕКСТ:
 {context}
 
 ВОПРОС: {question}
-"""
+
+ИНСТРУКЦИИ:
+1. Отвечай ТОЛЬКО на основе контекста
+2. Укажи источник (документ, таблицу)
+3. Если нужно показать содержимое таблицы - покажи ВСЕ данные
+4. Если информации нет - скажи прямо
+
+ОТВЕТ:"""
         
         response = query_engine.query(prompt)
         sources = format_sources(reranked)
@@ -223,54 +193,47 @@ def answer_question(question, query_engine, reranker):
         import traceback
         log_message(traceback.format_exc())
         return f"Ошибка: {e}", ""
+    
 
-def rerank_nodes(query, nodes, reranker, top_k=30, min_score=0.3):
-    """Rerank nodes with diversity - MORE LENIENT"""
+def rerank_nodes(query, nodes, reranker, top_k=25, min_score=0.3):
+    """Rerank with document grouping awareness"""
     if not nodes:
         return []
     
-    # Score all nodes
     pairs = [[query, n.text] for n in nodes]
     scores = reranker.predict(pairs)
     
-    # Sort by score
     scored = sorted(zip(nodes, scores), key=lambda x: x[1], reverse=True)
     
-    # More lenient threshold
+    log_message(f"Top 10 reranking scores: {[f'{s:.3f}' for _, s in scored[:10]]}")
+    
+    # More lenient filtering
     filtered = [(n, s) for n, s in scored if s >= min_score]
     
     if not filtered:
-        # Fallback: take top 50% if nothing passes threshold
-        cutoff = max(scores) * 0.5
+        cutoff = max(scores) * 0.4
         filtered = [(n, s) for n, s in scored if s >= cutoff][:top_k]
     
-    # Log top scores for debugging
-    log_message(f"Top 5 reranking scores: {[f'{s:.3f}' for _, s in scored[:5]]}")
+    # Group by document for diversity
+    doc_groups = {}
+    for node, score in filtered:
+        doc_group = node.metadata.get('document_group', node.metadata.get('document_id', 'unknown'))
+        if doc_group not in doc_groups:
+            doc_groups[doc_group] = []
+        doc_groups[doc_group].append((node, score))
     
-    # Diversity selection - but prioritize tables if query mentions them
+    # Take top chunks from each document group
     selected = []
-    seen_docs = set()
-    table_nodes = []
-    other_nodes = []
-    
-    for node, score in filtered:
-        if node.metadata.get('type') == 'table':
-            table_nodes.append((node, score))
-        else:
-            other_nodes.append((node, score))
+    group_limits = max(3, top_k // max(1, len(doc_groups)))
     
-    # If query mentions "таблица", prioritize table nodes
-    if 'таблиц' in query.lower():
-        combined = table_nodes + other_nodes
-    else:
-        combined = filtered
+    for doc_group in doc_groups:
+        selected.extend([n for n, s in doc_groups[doc_group][:group_limits]])
     
-    for node, score in combined[:top_k]:
-        if len(selected) >= top_k:
-            break
-        selected.append(node)
-        seen_docs.add(node.metadata.get('document_id', 'unknown'))
+    # Fill remaining slots with highest scores
+    if len(selected) < top_k:
+        remaining = [n for n, s in filtered if n not in selected]
+        selected.extend(remaining[:top_k - len(selected)])
     
-    log_message(f"Reranked: {len(filtered)} → {len(selected)} (from {len(seen_docs)} docs)")
+    log_message(f"Reranked: {len(filtered)} → {len(selected)} (from {len(doc_groups)} doc groups)")
     
-    return selected
+    return selected[:top_k]