simplest version

documents_prep.py

@@ -7,57 +7,71 @@ from llama_index.core.text_splitter import SentenceSplitter
 from my_logging import log_message
 
 # Configuration
-CHUNK_SIZE = 1500
+CHUNK_SIZE = 1024
 CHUNK_OVERLAP = 256
+
 def chunk_text_documents(documents):
-    """Chunk with deduplication"""
     text_splitter = SentenceSplitter(
         chunk_size=CHUNK_SIZE,
-        chunk_overlap=300  # Increased overlap
+        chunk_overlap=CHUNK_OVERLAP
     )
     
-    seen_texts = set()
     chunked = []
-    
     for doc in documents:
-        text_normalized = doc.text.strip()
-        if len(text_normalized) < 50 or text_normalized in seen_texts:
-            continue
-        
-        seen_texts.add(text_normalized)
-        
         chunks = text_splitter.get_nodes_from_documents([doc])
         for i, chunk in enumerate(chunks):
             chunk.metadata.update({
                 'chunk_id': i,
                 'total_chunks': len(chunks),
-                'chunk_size': len(chunk.text),
-                'document_group': normalize_doc_id(doc.metadata.get('document_id', 'unknown'))
+                'chunk_size': len(chunk.text)  # Add chunk size
             })
             chunked.append(chunk)
     
+    # Log statistics
     if chunked:
         avg_size = sum(len(c.text) for c in chunked) / len(chunked)
-        log_message(f"✓ Text: {len(documents)} docs → {len(chunked)} chunks (avg: {avg_size:.0f} chars)")
+        min_size = min(len(c.text) for c in chunked)
+        max_size = max(len(c.text) for c in chunked)
+        log_message(f"✓ Text: {len(documents)} docs → {len(chunked)} chunks")
+        log_message(f"  Size stats: avg={avg_size:.0f}, min={min_size}, max={max_size} chars")
     
     return chunked
 
-def chunk_table_by_rows(table_data, doc_id, max_chars=2000):
-    """Chunk tables by content size, not fixed rows"""
+
+def normalize_doc_id(doc_id):
+    """Normalize document ID for consistent matching"""
+    if not doc_id or doc_id == 'unknown':
+        return doc_id
+    
+    doc_id = str(doc_id).strip()
+    
+    # Normalize spacing: "ГОСТ Р" variations
+    import re
+    doc_id = re.sub(r'ГОСТ\s*Р', 'ГОСТ Р', doc_id, flags=re.IGNORECASE)
+    doc_id = re.sub(r'НП\s*-', 'НП-', doc_id, flags=re.IGNORECASE)
+    
+    return doc_id
+
+
+def chunk_table_by_rows(table_data, doc_id, max_rows=2):
     headers = table_data.get('headers', [])
     rows = table_data.get('data', [])
     table_num = table_data.get('table_number', 'unknown')
     table_title = table_data.get('table_title', '')
     section = table_data.get('section', '')
     
+    # NORMALIZE document ID
+    doc_id = normalize_doc_id(doc_id)
+    
     table_num_clean = str(table_num).strip()
     
-    # Create unique identifier
+    # Create section-aware identifier
     import re
     if 'приложени' in section.lower():
         appendix_match = re.search(r'приложени[еия]\s*(\d+|[а-яА-Я])', section.lower())
         if appendix_match:
-            table_identifier = f"{table_num_clean} (Приложение {appendix_match.group(1).upper()})"
+            appendix_num = appendix_match.group(1).upper()
+            table_identifier = f"{table_num_clean} Приложение {appendix_num}"
         else:
             table_identifier = table_num_clean
     else:
@@ -66,161 +80,128 @@ def chunk_table_by_rows(table_data, doc_id, max_chars=2000):
     if not rows:
         return []
     
-    # Estimate base metadata size
-    base_content = f"Документ: {doc_id}\nТаблица: {table_identifier}\n"
-    if table_title:
-        base_content += f"Название: {table_title}\n"
-    if section:
-        base_content += f"Раздел: {section}\n"
-    
-    header_content = ""
-    if headers:
-        header_content = "Столбцы: " + " | ".join(str(h) for h in headers) + "\n\n"
-    
-    base_size = len(base_content) + len(header_content)
-    
-    # Group rows by size
-    chunks = []
-    current_rows = []
-    current_size = base_size
+    log_message(f"  📊 Processing: {doc_id} - {table_identifier} ({len(rows)} rows)")
     
-    for row in rows:
-        # Estimate row size
-        if isinstance(row, dict):
-            row_str = " | ".join(f"{k}: {v}" for k, v in row.items() 
-                                if v and str(v).strip() and str(v).lower() not in ['nan', 'none', ''])
-        elif isinstance(row, list):
-            row_str = " | ".join(str(v) for v in row 
-                                if v and str(v).strip() and str(v).lower() not in ['nan', 'none', ''])
-        else:
-            row_str = str(row)
+    if len(rows) <= max_rows:
+        content = format_table_content(table_data, headers, rows, doc_id, table_identifier)
+        chunk_size = len(content)
         
-        row_size = len(row_str) + 2  # +2 for newline
+        metadata = {
+            'type': 'table',
+            'document_id': doc_id,
+            'table_number': table_num_clean,
+            'table_identifier': table_identifier,
+            'table_title': table_title,
+            'section': section,
+            'total_rows': len(rows),
+            'chunk_size': chunk_size,
+            'is_complete_table': True
+        }
         
-        # If adding this row exceeds limit and we have rows, create chunk
-        if current_size + row_size > max_chars and current_rows:
-            chunks.append(current_rows[:])
-            current_rows = []
-            current_size = base_size
+        log_message(f"    Chunk: 1/1, {chunk_size} chars, doc={doc_id}, table={table_identifier}")
         
-        current_rows.append(row)
-        current_size += row_size
+        return [Document(text=content, metadata=metadata)]
     
-    # Add remaining rows
-    if current_rows:
-        chunks.append(current_rows)
+    chunks = []
+    overlap = 1
     
-    # Create documents
-    documents = []
-    for chunk_idx, chunk_rows in enumerate(chunks):
-        content = base_content
-        content += f"Таблица {table_identifier} документа {doc_id}\n"
-        if len(chunks) > 1:
-            content += f"Часть {chunk_idx+1} из {len(chunks)}\n"
-        content += "\n" + header_content
+    for i in range(0, len(rows), max_rows - overlap):
+        chunk_rows = rows[i:min(i+max_rows, len(rows))]
+        chunk_num = i // (max_rows - overlap)
+        
+        content = format_table_content(
+            table_data, 
+            headers, 
+            chunk_rows,
+            doc_id,
+            table_identifier,
+            chunk_info=f"Строки {i+1}-{i+len(chunk_rows)} из {len(rows)}"
+        )
         
-        for idx, row in enumerate(chunk_rows, 1):
-            if isinstance(row, dict):
-                parts = [f"{k}: {v}" for k, v in row.items() 
-                        if v and str(v).strip() and str(v).lower() not in ['nan', 'none', '']]
-                if parts:
-                    content += f"{idx}. {' | '.join(parts)}\n"
-            elif isinstance(row, list):
-                parts = [str(v) for v in row if v and str(v).strip() and str(v).lower() not in ['nan', 'none', '']]
-                if parts:
-                    content += f"{idx}. {' | '.join(parts)}\n"
+        chunk_size = len(content)
         
         metadata = {
             'type': 'table',
             'document_id': doc_id,
-            'document_group': normalize_doc_id(doc_id),
             'table_number': table_num_clean,
             'table_identifier': table_identifier,
             'table_title': table_title,
             'section': section,
-            'chunk_id': chunk_idx,
-            'total_chunks': len(chunks),
-            'chunk_size': len(content),
-            'is_complete_table': len(chunks) == 1
+            'chunk_id': chunk_num,
+            'row_start': i,
+            'row_end': i + len(chunk_rows),
+            'total_rows': len(rows),
+            'chunk_size': chunk_size,
+            'total_chunks': (len(rows) + max_rows - overlap - 1) // (max_rows - overlap),
+            'is_complete_table': False
         }
         
-        documents.append(Document(text=content, metadata=metadata))
+        log_message(f"    Chunk: {chunk_num+1}, rows {i}-{i+len(chunk_rows)}, {chunk_size} chars")
         
-        log_message(f"    Chunk {chunk_idx+1}: {len(chunk_rows)} rows, {len(content)} chars")
-        log_message(f"      Meta: doc={doc_id}, table={table_identifier}, group={metadata['document_group']}")
+        chunks.append(Document(text=content, metadata=metadata))
     
-    log_message(f"  Table {table_identifier} ({doc_id}): {len(rows)} rows → {len(chunks)} chunks")
-    
-    return documents
+    return chunks
 
 
-def normalize_doc_id(doc_id):
-    import re
-    normalized = re.sub(r'\s+', ' ', str(doc_id).strip().upper())
-    normalized = normalized.replace('ГОСТ Р', 'ГОСТР').replace('ГОСТР', 'ГОСТ Р')
-    return normalized
-
-
-def format_table_content(table_data, headers, rows, doc_id, table_identifier, chunk_info=""):
+def format_table_content(table_data, headers, rows, table_identifier, chunk_info=""):
+    doc_id = table_data.get('document_id', table_data.get('document', 'unknown'))
     table_num = table_data.get('table_number', 'unknown')
     table_title = table_data.get('table_title', '')
     section = table_data.get('section', '')
     
-    # Build content with multiple search variations
+    # Use enhanced identifier
     content = f"ДОКУМЕНТ: {doc_id}\n"
     content += f"ТАБЛИЦА: {table_identifier}\n"
-    
-    # Add search variations for document ID
-    doc_variations = [doc_id]
-    if 'Р' in doc_id:
-        doc_variations.append(doc_id.replace(' Р ', ' Р'))
-        doc_variations.append(doc_id.replace(' Р ', 'Р'))
-    
-    for var in set(doc_variations):
-        content += f"ДОКУМЕНТ_ВАРИАНТ: {var}\n"
-    
+    content += f"ПОЛНОЕ НАЗВАНИЕ: {table_identifier}\n"
+    content += f"НОМЕР ТАБЛИЦЫ: {table_num}\n"
     if table_title:
         content += f"НАЗВАНИЕ: {table_title}\n"
     if section:
         content += f"РАЗДЕЛ: {section}\n"
-    
     content += f"{'='*70}\n\n"
     
-    # Enhanced search text
-    content += f"Документ {doc_id}. "
-    content += f"Таблица {table_identifier}. "
-    content += f"Номер таблицы {table_num}. "
-    
-    if table_title:
-        content += f"Название: {table_title}. "
+    # Enhanced search keywords
+    content += f"Это таблица {table_identifier} из документа {doc_id}. "
+    content += f"Идентификатор таблицы: {table_identifier}. "
+    content += f"Номер: {table_num}. "
+    content += f"Документ: {doc_id}. "
     
     if section:
-        content += f"Раздел: {section}. "
+        content += f"Находится в разделе: {section}. "
+        if 'приложени' in section.lower():
+            content += f"Таблица из приложения. "
+    
+    if table_title:
+        content += f"Название таблицы: {table_title}. "
+        content += f"Таблица о: {table_title}. "
     
-    # Add more search patterns
-    content += f"Таблицы документа {doc_id}. "
-    content += f"Содержание {doc_id}. "
+    content += f"Поиск: таблица {table_identifier} {doc_id}. "
     
     if chunk_info:
-        content += f"{chunk_info}. "
+        content += f"\n{chunk_info}\n"
     
-    content += f"\n\nДАННЫЕ ТАБЛИЦЫ {table_identifier}:\n{'='*70}\n\n"
+    content += f"\n\nСОДЕРЖИМОЕ ТАБЛИЦЫ {table_identifier}:\n"
+    content += f"="*70 + "\n\n"
     
     if headers:
-        content += f"СТОЛБЦЫ: {' | '.join(str(h) for h in headers)}\n\n"
+        header_str = ' | '.join(str(h) for h in headers)
+        content += f"ЗАГОЛОВКИ СТОЛБЦОВ:\n{header_str}\n\n"
     
+    content += f"ДАННЫЕ ТАБЛИЦЫ:\n"
     for idx, row in enumerate(rows, 1):
         if isinstance(row, dict):
             parts = [f"{k}: {v}" for k, v in row.items() 
-                    if v and str(v).strip().lower() not in ['nan', 'none', '', 'null']]
+                    if v and str(v).strip() and str(v).lower() not in ['nan', 'none', '']]
             if parts:
                 content += f"{idx}. {' | '.join(parts)}\n"
         elif isinstance(row, list):
-            parts = [str(v) for v in row 
-                    if v and str(v).strip().lower() not in ['nan', 'none', '', 'null']]
+            parts = [str(v) for v in row if v and str(v).strip() and str(v).lower() not in ['nan', 'none', '']]
             if parts:
                 content += f"{idx}. {' | '.join(parts)}\n"
     
+    content += f"\n{'='*70}\n"
+    content += f"КОНЕЦ ТАБЛИЦЫ {table_identifier} ИЗ {doc_id}\n"
+    
     return content
 
 def load_json_documents(repo_id, hf_token, json_dir):
@@ -352,6 +333,7 @@ def load_json_documents(repo_id, hf_token, json_dir):
     return documents
 
 def extract_sections_from_json(json_path):
+    """Extract sections from a single JSON file"""
     documents = []
     
     try:
@@ -359,8 +341,8 @@ def extract_sections_from_json(json_path):
             data = json.load(f)
         
         doc_id = data.get('document_metadata', {}).get('document_id', 'unknown')
-        doc_id = normalize_doc_id(doc_id)  # NORMALIZE
         
+        # Extract all section levels
         for section in data.get('sections', []):
             if section.get('section_text', '').strip():
                 documents.append(Document(
@@ -368,11 +350,11 @@ def extract_sections_from_json(json_path):
                     metadata={
                         'type': 'text',
                         'document_id': doc_id,
-                        'section_id': section.get('section_id', ''),
-                        'chunk_size': len(section['section_text'])
+                        'section_id': section.get('section_id', '')
                     }
                 ))
             
+            # Subsections
             for subsection in section.get('subsections', []):
                 if subsection.get('subsection_text', '').strip():
                     documents.append(Document(
@@ -380,11 +362,11 @@ def extract_sections_from_json(json_path):
                         metadata={
                             'type': 'text',
                             'document_id': doc_id,
-                            'section_id': subsection.get('subsection_id', ''),
-                            'chunk_size': len(subsection['subsection_text'])
+                            'section_id': subsection.get('subsection_id', '')
                         }
                     ))
                 
+                # Sub-subsections
                 for sub_sub in subsection.get('sub_subsections', []):
                     if sub_sub.get('sub_subsection_text', '').strip():
                         documents.append(Document(
@@ -392,8 +374,7 @@ def extract_sections_from_json(json_path):
                             metadata={
                                 'type': 'text',
                                 'document_id': doc_id,
-                                'section_id': sub_sub.get('sub_subsection_id', ''),
-                                'chunk_size': len(sub_sub['sub_subsection_text'])
+                                'section_id': sub_sub.get('sub_subsection_id', '')
                             }
                         ))
     
@@ -404,17 +385,13 @@ def extract_sections_from_json(json_path):
 
 
 def load_table_documents(repo_id, hf_token, table_dir):
-    """Load ALL tables including from multi-document files"""
+    """Load and chunk tables"""
     log_message("Loading tables...")
     
     files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token)
     table_files = [f for f in files if f.startswith(table_dir) and f.endswith('.json')]
     
-    log_message(f"Found {len(table_files)} table files")
-    
     all_chunks = []
-    doc_id_stats = {}
-    
     for file_path in table_files:
         try:
             local_path = hf_hub_download(
@@ -427,40 +404,32 @@ def load_table_documents(repo_id, hf_token, table_dir):
             with open(local_path, 'r', encoding='utf-8') as f:
                 data = json.load(f)
             
+            # Extract file-level document_id
             file_doc_id = data.get('document_id', data.get('document', 'unknown'))
             
             for sheet in data.get('sheets', []):
+                # Use sheet-level document_id if available, otherwise use file-level
                 sheet_doc_id = sheet.get('document_id', sheet.get('document', file_doc_id))
                 
-                # Track which documents we're loading
-                if sheet_doc_id not in doc_id_stats:
-                    doc_id_stats[sheet_doc_id] = 0
-                
+                # CRITICAL: Pass document_id to chunk function
                 chunks = chunk_table_by_rows(sheet, sheet_doc_id)
                 all_chunks.extend(chunks)
-                doc_id_stats[sheet_doc_id] += len(chunks)
                 
         except Exception as e:
             log_message(f"Error loading {file_path}: {e}")
     
-    # Log what we loaded
-    log_message(f"\nTable loading summary:")
-    for doc_id, count in sorted(doc_id_stats.items()):
-        log_message(f"  {doc_id}: {count} chunks")
-    
-    log_message(f"\n✓ Total table chunks: {len(all_chunks)}")
+    log_message(f"✓ Loaded {len(all_chunks)} table chunks")
     return all_chunks
 
+
 def load_image_documents(repo_id, hf_token, image_dir):
-    """Load with proper linking"""
+    """Load image descriptions"""
     log_message("Loading images...")
     
     files = list_repo_files(repo_id=repo_id, repo_type="dataset", token=hf_token)
     csv_files = [f for f in files if f.startswith(image_dir) and f.endswith('.csv')]
     
     documents = []
-    seen = set()
-    
     for file_path in csv_files:
         try:
             local_path = hf_hub_download(
@@ -473,28 +442,22 @@ def load_image_documents(repo_id, hf_token, image_dir):
             df = pd.read_csv(local_path)
             
             for _, row in df.iterrows():
-                doc_id = str(row.get('Обозначение документа', 'unknown'))
-                img_num = str(row.get('№ Изображения', 'unknown'))
-                
-                key = f"{doc_id}_{img_num}"
-                if key in seen:
-                    continue
-                seen.add(key)
-                
-                content = f"Документ: {doc_id}\n"
-                content += f"Рисунок: {img_num}\n"
+                content = f"Документ: {row.get('Обозначение документа', 'unknown')}\n"
+                content += f"Рисунок: {row.get('№ Изображения', 'unknown')}\n"
                 content += f"Название: {row.get('Название изображения', '')}\n"
                 content += f"Описание: {row.get('Описание изображение', '')}\n"
+                content += f"Раздел: {row.get('Раздел документа', '')}\n"
+                
+                chunk_size = len(content)
                 
                 documents.append(Document(
                     text=content,
                     metadata={
                         'type': 'image',
-                        'document_id': doc_id,
-                        'document_group': normalize_doc_id(doc_id),
-                        'image_number': img_num,
+                        'document_id': str(row.get('Обозначение документа', 'unknown')),
+                        'image_number': str(row.get('№ Изображения', 'unknown')),
                         'section': str(row.get('Раздел документа', '')),
-                        'chunk_size': len(content)
+                        'chunk_size': chunk_size
                     }
                 ))
         except Exception as e:
@@ -502,12 +465,11 @@ def load_image_documents(repo_id, hf_token, image_dir):
     
     if documents:
         avg_size = sum(d.metadata['chunk_size'] for d in documents) / len(documents)
-        log_message(f"✓ Images: {len(documents)} loaded (avg: {avg_size:.0f} chars)")
+        log_message(f"✓ Loaded {len(documents)} images (avg size: {avg_size:.0f} chars)")
     
     return documents
 
 
-
 def load_all_documents(repo_id, hf_token, json_dir, table_dir, image_dir):
     """Main loader - combines all document types"""
     log_message("="*60)

index_retriever.py

@@ -35,19 +35,19 @@ def create_query_engine(vector_index):
     # Vector retriever
     vector_retriever = VectorIndexRetriever(
         index=vector_index,
-        similarity_top_k=40
+        similarity_top_k=50
     )
     
     # BM25 retriever
     bm25_retriever = BM25Retriever.from_defaults(
         docstore=vector_index.docstore,
-        similarity_top_k=40
+        similarity_top_k=50
     )
     
     # Hybrid fusion
     hybrid_retriever = QueryFusionRetriever(
         [vector_retriever, bm25_retriever],
-        similarity_top_k=50,
+        similarity_top_k=60,
         num_queries=1
     )
     

utils.py

@@ -40,96 +40,74 @@ 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_id = f"{doc_match.group(1).upper()} Р {doc_match.group(2)}"
-        enhanced_query += f" документ {doc_id}"
+        doc_type = doc_match.group(1).upper()
+        doc_num = doc_match.group(2)
+        
+        # Add normalized versions
+        enhanced_query += f" {doc_type} Р {doc_num}"
     
     return enhanced_query
 
 def answer_question(question, query_engine, reranker):
     try:
-        log_message(f"\n{'='*70}")
-        log_message(f"QUERY: {question}")
+        log_message(f"Query: {question}")
         
         enhanced_query = preprocess_query(question)
-        log_message(f"Enhanced: {enhanced_query}")
-        
-        # Detect listing queries - need MORE chunks
-        is_listing_query = any(phrase in question.lower() 
-                              for phrase in ['какие таблиц', 'список', 'перечисл', 'все таблиц'])
+        if enhanced_query != question:
+            log_message(f"Enhanced query: {enhanced_query}")
         
         retrieved = query_engine.retriever.retrieve(enhanced_query)
-        log_message(f"\nRETRIEVED: {len(retrieved)} nodes")
+        log_message(f"Retrieved {len(retrieved)} nodes")
         
-        # Log retrieved docs
         doc_stats = {}
         for n in retrieved:
             doc_id = n.metadata.get('document_id', 'unknown')
-            doc_group = n.metadata.get('document_group', doc_id)
+            doc_type = n.metadata.get('type', 'text')
             
-            if doc_group not in doc_stats:
-                doc_stats[doc_group] = {'tables': set(), 'text': 0}
+            if doc_id not in doc_stats:
+                doc_stats[doc_id] = {'tables': set(), 'text': 0, 'images': 0}
             
-            if n.metadata.get('type') == 'table':
+            if doc_type == 'table':
                 table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
-                doc_stats[doc_group]['tables'].add(table_id)
+                doc_stats[doc_id]['tables'].add(table_id)
+            elif doc_type == 'image':
+                doc_stats[doc_id]['images'] += 1
             else:
-                doc_stats[doc_group]['text'] += 1
+                doc_stats[doc_id]['text'] += 1
         
         for doc_id in sorted(doc_stats.keys()):
             stats = doc_stats[doc_id]
-            log_message(f"  {doc_id}: {len(stats['tables'])} tables, {stats['text']} text")
+            parts = []
             if stats['tables']:
-                log_message(f"    Tables: {sorted(stats['tables'])}")
-        
-        # 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)
-        
-        log_message(f"\nRERANKED: {len(reranked)} nodes")
+                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 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
+        reranked = rerank_nodes(question, retrieved, reranker, top_k=25)
+        log_message(f"Reranked to {len(reranked)} nodes")
         
-        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'])}")
+        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)}")
         
-        # Build context
         context_parts = []
         for n in reranked:
             meta = n.metadata
@@ -137,48 +115,103 @@ def answer_question(question, query_engine, reranker):
             doc_type = meta.get('type', 'text')
             
             if doc_type == 'table':
-                table_id = meta.get('table_identifier', meta.get('table_number', 'unknown'))
+                table_num = meta.get('table_number', 'unknown')
                 title = meta.get('table_title', '')
-                source_label = f"[{doc_id} - Таблица {table_id}]"
+                source_label = f"[ТАБЛИЦА {table_num} - {doc_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:
-                source_label = f"[{doc_id}]"
+                section = meta.get('section_id', '')
+                source_label = f"[{doc_id} - {section}]"
             
-            context_parts.append(f"{source_label}\n{n.text[:500]}")  # Limit context per chunk
-        
-        context = "\n\n" + ("="*50 + "\n\n").join(context_parts)
+            context_parts.append(f"{source_label}\n{n.text}")
         
-        # Adjust prompt for listing queries
-        if is_listing_query:
-            prompt = f"""Контекст содержит информацию о таблицах из документов.
+        context = "\n\n" + ("="*70 + "\n\n").join(context_parts)
+        from config import CUSTOM_PROMPT
+        prompt = f"""
+Вы являетесь высокоспециализированным Ассистентом для анализа нормативных документов (AIEXP). Ваша цель - предоставлять точные, корректные и контекстно релевантные ответы исключительно на основе предоставленного контекста из нормативной документации.
 
-КОНТЕКСТ:
-{context}
+ПРАВИЛА АНАЛИЗА ЗАПРОСА:
 
-ВОПРОС: {question}
+1. ПРЯМЫЕ ВОПРОСЫ БЕЗ ДОКУМЕНТАЛЬНОГО КОНТЕКСТА:
+   Если пользователь задает вопрос типа "В каких случаях могут быть признаны протоколы испытаний?" без предоставления дополнительных документов, найдите соответствующую информацию в доступном контексте и предоставьте полный ответ с указанием источников.
 
-ИНСТРУКЦИИ:
-1. Перечисли ВСЕ таблицы, найденные в контексте для запрошенного документа
-2. Укажи номер таблицы и название (если есть)
-3. Если таблиц нет - скажи прямо
+2. ОПРЕДЕЛЕНИЕ ТИПА ЗАДАЧИ:
 
-ОТВЕТ (список таблиц):"""
-        else:
-            prompt = f"""Ты эксперт по технической документации.
+   а) ПОИСК И ОТВЕТ НА ВОПРОС (ключевые слова: "в каких случаях", "когда", "кто", "что", "как", "почему"):
+      - Найдите релевантную информацию в контексте
+      - Предоставьте развернутый ответ
+      - Обязательно укажите конкретные документы и разделы
+      - Процитируйте ключевые положения
+
+   б) КРАТКОЕ САММАРИ (ключевые слова: "кратко", "суммировать", "резюме", "основные моменты"):
+      - Предоставьте структурированное резюме
+      - Выделите ключевые требования
+      - Используйте нумерованный список
+
+   в) ПОИСК ДОКУМЕНТА И ПУНКТА (ключевые слова: "найти", "где", "какой документ", "в каком разделе"):
+      - Укажите конкретный документ и структурное расположение
+      - Предоставьте точные номера разделов/пунктов
+
+   г) ПРОВЕРКА КОРРЕКТНОСТИ (ключевые слова: "правильно ли", "соответствует ли", "проверить"):
+      - Четко укажите: "СООТВЕТСТВУЕТ" или "НЕ СООТВЕТСТВУЕТ"
+      - Перечислите конкретные требования
+
+   д) ПЛАН ДЕЙСТВИЙ (ключевые слова: "план", "алгоритм", "пошагово"):
+      - Создайте пронумерованный план
+      - Укажите ссылки на соответствующие пункты НД
+
+ПРАВИЛА ФОРМИРОВАНИЯ ОТВЕТОВ:
+
+Работай исключительно с информацией из предоставленного контекста. Запрещено использовать:
+- Общие знания
+- Информацию из интернета
+- Данные из предыдущих диалогов
+- Собственные предположения
+
+1. СТРУКТУРА ОТВЕТА:
+   - Начинайте с прямого ответа на вопрос
+   - Затем указывайте нормативные основания
+   - Завершайте ссылками на конкретные документы и разделы
+
+2. РАБОТА С КОНТЕКСТОМ:
+   - Если информация найдена в контексте - предоставьте полный ответ
+   - Если информация не найдена: "Информация по вашему запросу не найдена в доступной нормативной документации"
+   - Не делайте предположений за пределами контекста
+   - Не используйте общие знания
+
+3. ТЕРМИНОЛОГИЯ И ЦИТИРОВАНИЕ:
+   - Сохраняйте официальную терминологию НД
+   - Цитируйте точные формулировки ключевых требований
+   - При множественных источниках - укажите все релевантные
+
+4. ФОРМАТИРОВАНИЕ:
+   - Для перечислений: используйте нумерованные списки
+   - Выделяйте критически важные требования
+   - Структурируйте ответ логически
+
+# КАК РАБОТАТЬ С ЗАПРОСОМ
+
+**Шаг 1:** Определи, что именно ищет пользователь (термин, требование, процедура, условие)
+
+**Шаг 2:** Найди релевантную информацию в контексте
+
+**Шаг 3:** Сформируй ответ:
+   - Если нашел: укажи документ и пункт, процитируй нужную часть
+   - Если не нашел: четко сообщи об отсутствии информации
+
+**Шаг 4:** При наличии нескольких источников:
+   - Представь их последовательно с указанием источника каждого
+   - Если источников много (>4) — сначала дай их список, потом цитаты   
 
 КОНТЕКСТ:
 {context}
 
 ВОПРОС: {question}
-
-ИНСТРУКЦИИ:
-1. Отвечай ТОЛЬКО на основе контекста
-2. Укажи источник (документ, таблицу)
-3. Если нужно показать содержимое таблицы - покажи ВСЕ данные
-4. Если информации нет - скажи прямо
-
-ОТВЕТ:"""
+"""
         
         response = query_engine.query(prompt)
         sources = format_sources(reranked)
@@ -190,45 +223,54 @@ 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=25, min_score=0.3):
-    """Rerank with document grouping awareness"""
+
+def rerank_nodes(query, nodes, reranker, top_k=30, min_score=0.3):
+    """Rerank nodes with diversity - MORE LENIENT"""
     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)
     
-    log_message(f"Top 10 reranking scores: {[f'{s:.3f}' for _, s in scored[:10]]}")
-    
-    # More lenient filtering
+    # More lenient threshold
     filtered = [(n, s) for n, s in scored if s >= min_score]
     
     if not filtered:
-        cutoff = max(scores) * 0.4
+        # Fallback: take top 50% if nothing passes threshold
+        cutoff = max(scores) * 0.5
         filtered = [(n, s) for n, s in scored if s >= cutoff][:top_k]
     
-    # 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))
+    # Log top scores for debugging
+    log_message(f"Top 5 reranking scores: {[f'{s:.3f}' for _, s in scored[:5]]}")
     
-    # Take top chunks from each document group
+    # Diversity selection - but prioritize tables if query mentions them
     selected = []
-    group_limits = max(3, top_k // max(1, len(doc_groups)))
+    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))
     
-    for doc_group in doc_groups:
-        selected.extend([n for n, s in doc_groups[doc_group][:group_limits]])
+    # If query mentions "таблица", prioritize table nodes
+    if 'таблиц' in query.lower():
+        combined = table_nodes + other_nodes
+    else:
+        combined = filtered
     
-    # 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)])
+    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'))
     
-    log_message(f"Reranked: {len(filtered)} → {len(selected)} (from {len(doc_groups)} doc groups)")
+    log_message(f"Reranked: {len(filtered)} → {len(selected)} (from {len(seen_docs)} docs)")
     
-    return selected[:top_k]
+    return selected