api key added

utils.py

@@ -1,21 +1,41 @@
+import logging
+import sys
 from llama_index.llms.google_genai import GoogleGenAI
+from llama_index.llms.openai import OpenAI
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from sentence_transformers import CrossEncoder
+from config import AVAILABLE_MODELS, DEFAULT_MODEL, GOOGLE_API_KEY
+import time
+from index_retriever import rerank_nodes
 from my_logging import log_message
-import os
+from config import PROMPT_SIMPLE_POISK
 
-api_key = os.getenv('GOOGLE_API_KEY')  # or however you're loading it
-if not api_key:
-    log_message("GOOGLE_API_KEY not found in environment")
-    raise ValueError("GOOGLE_API_KEY not found in environment")
-
-
-def get_llm_model(model_name="gemini-2.5-flash"):
-    api_key = os.getenv('GOOGLE_API_KEY')
-    if not api_key:
-        log_message("GOOGLE_API_KEY not found in environment")
-        raise ValueError("GOOGLE_API_KEY not found in environment")
-    return GoogleGenAI(model=model_name, api_key=api_key)
+def get_llm_model(model_name):
+    try:
+        model_config = AVAILABLE_MODELS.get(model_name)
+        if not model_config:
+            log_message(f"Модель {model_name} не найдена, использую модель по умолчанию")
+            model_config = AVAILABLE_MODELS[DEFAULT_MODEL]
+        
+        if not model_config.get("api_key"):
+            raise Exception(f"API ключ не найден для модели {model_name}")
+        
+        if model_config["provider"] == "google":
+            return GoogleGenAI(
+                model=model_config["model_name"], 
+                api_key=model_config["api_key"]
+            )
+        elif model_config["provider"] == "openai":
+            return OpenAI(
+                model=model_config["model_name"],
+                api_key=model_config["api_key"]
+            )
+        else:
+            raise Exception(f"Неподдерживаемый провайдер: {model_config['provider']}")
+            
+    except Exception as e:
+        log_message(f"Ошибка создания модели {model_name}: {str(e)}")
+        return GoogleGenAI(model="gemini-2.0-flash", api_key=GOOGLE_API_KEY)
 
 def get_embedding_model(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"):
     return HuggingFaceEmbedding(model_name=model_name)
@@ -23,250 +43,263 @@ def get_embedding_model(model_name="sentence-transformers/paraphrase-multilingua
 def get_reranker_model(model_name='cross-encoder/ms-marco-MiniLM-L-12-v2'):
     return CrossEncoder(model_name)
 
-def format_sources(nodes):
-    sources = []
+def format_context_for_llm(nodes):
+    context_parts = []
+    
     for node in nodes:
-        meta = node.metadata
-        doc_type = meta.get('type', 'text')
-        doc_id = meta.get('document_id', 'unknown')
+        metadata = node.metadata if hasattr(node, 'metadata') else {}
+        doc_id = metadata.get('document_id', 'Неизвестный документ')
         
-        if doc_type == 'table':
-            table_num = meta.get('table_number', 'unknown')
-            title = meta.get('table_title', '')
-            sources.append(f"📊 {doc_id} - Таблица {table_num}: {title}")
-        elif doc_type == 'image':
-            img_num = meta.get('image_number', 'unknown')
-            sources.append(f"🖼️ {doc_id} - Рисунок {img_num}")
+        section_info = ""
+        
+        # Handle section information with proper hierarchy
+        if metadata.get('section_path'):
+            section_path = metadata['section_path']
+            section_text = metadata.get('section_text', '')
+            parent_section = metadata.get('parent_section', '')
+            parent_title = metadata.get('parent_title', '')
+            level = metadata.get('level', '')
+            
+            if level in ['subsection', 'sub_subsection', 'sub_sub_subsection'] and parent_section and parent_title:
+                # For subsections: раздел X (Title), пункт X.X
+                if section_text:
+                    section_info = f"раздел {parent_section} ({parent_title}), пункт {section_path} ({section_text})"
+                else:
+                    section_info = f"раздел {parent_section} ({parent_title}), пункт {section_path}"
+            elif section_text:
+                # For main sections: раздел X (Title)
+                section_info = f"раздел {section_path} ({section_text})"
+            else:
+                section_info = f"раздел {section_path}"
+                
+        elif metadata.get('section_id'):
+            section_id = metadata['section_id']
+            section_text = metadata.get('section_text', '')
+            level = metadata.get('level', '')
+            parent_section = metadata.get('parent_section', '')
+            parent_title = metadata.get('parent_title', '')
+            
+            if level in ['subsection', 'sub_subsection', 'sub_sub_subsection'] and parent_section and parent_title:
+                if section_text:
+                    section_info = f"раздел {parent_section} ({parent_title}), пункт {section_id} ({section_text})"
+                else:
+                    section_info = f"раздел {parent_section} ({parent_title}), пункт {section_id}"
+            elif section_text:
+                section_info = f"раздел {section_id} ({section_text})"
+            else:
+                section_info = f"раздел {section_id}"
+
+        # Override with table/image info if applicable
+        if metadata.get('type') == 'table' and metadata.get('table_number'):
+            table_num = metadata['table_number']
+            if not str(table_num).startswith('№'):
+                table_num = f"№{table_num}"
+            table_title = metadata.get('table_title', '')
+            # Include section context for tables
+            base_section = ""
+            if metadata.get('section_path'):
+                base_section = f", раздел {metadata['section_path']}"
+            elif metadata.get('section_id'):
+                base_section = f", раздел {metadata['section_id']}"
+            
+            if table_title:
+                section_info = f"Таблица {table_num} ({table_title}){base_section}"
+            else:
+                section_info = f"Таблица {table_num}{base_section}"
+        
+        if metadata.get('type') == 'image' and metadata.get('image_number'):
+            image_num = metadata['image_number']
+            if not str(image_num).startswith('№'):
+                image_num = f"№{image_num}"
+            image_title = metadata.get('image_title', '')
+            # Include section context for images
+            base_section = ""
+            if metadata.get('section_path'):
+                base_section = f", раздел {metadata['section_path']}"
+            elif metadata.get('section_id'):
+                base_section = f", раздел {metadata['section_id']}"
+            
+            if image_title:
+                section_info = f"Рисунок {image_num} ({image_title}){base_section}"
+            else:
+                section_info = f"Рисунок {image_num}{base_section}"
+        
+        context_text = node.text if hasattr(node, 'text') else str(node)
+        
+        if section_info:
+            formatted_context = f"[ИСТОЧНИК: {section_info}, документ {doc_id}]\n{context_text}\n"
         else:
-            section = meta.get('section_id', '')
-            sources.append(f"📄 {doc_id} - Раздел {section}")
+            formatted_context = f"[ИСТОЧНИК: документ {doc_id}]\n{context_text}\n"
+        
+        context_parts.append(formatted_context)
     
-    return "\n".join(set(sources))
-
-def create_chunks_info_for_display(nodes):
-    chunks_info = []
-    for node in nodes:
-        meta = node.metadata
-        chunk_info = {
-            'document_id': meta.get('document_id', 'unknown'),
-            'section_path': meta.get('section_path', ''),
-            'section_id': meta.get('section_id', 'unknown'),
-            'section_text': meta.get('section_text', ''),
-            'parent_section': meta.get('parent_section', ''),
-            'parent_title': meta.get('parent_title', ''),
-            'level': meta.get('level', ''),
-            'chunk_text': node.text[:500],
-            'type': meta.get('type', 'text'),
-            'table_number': meta.get('table_number', ''),
-            'image_number': meta.get('image_number', '')
-        }
-        chunks_info.append(chunk_info)
-    return chunks_info
+    return "\n".join(context_parts)
 
-def format_answer_html(answer_text, model_name):
-    html = f"""
-    <div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px;'>
-        <div style='margin-bottom: 10px;'>
-            <span style='background-color: #4a5568; padding: 5px 10px; border-radius: 5px; font-size: 12px;'>
-                Модель: {model_name}
-            </span>
-        </div>
-        <div style='line-height: 1.6;'>
-            {answer_text}
-        </div>
-    </div>
-    """
-    return html
 
-def format_sources_html(sources_text):
-    if not sources_text or sources_text == "":
-        return "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Нет источников</div>"
+def generate_sources_html(nodes, chunks_df=None):
+    html = "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; max-height: 400px; overflow-y: auto;'>"
+    html += "<h3 style='color: #63b3ed; margin-top: 0;'>Источники:</h3>"
     
-    sources_list = sources_text.strip().split('\n')
-    html = "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px;'>"
-    html += "<h4 style='color: white; margin-bottom: 15px;'>Использованные источники:</h4>"
-    html += "<div style='line-height: 2;'>"
+    sources_by_doc = {}
     
-    for source in sources_list:
-        if source.strip():
-            html += f"<div style='padding: 5px 0; border-bottom: 1px solid #4a5568;'>{source}</div>"
+    for i, node in enumerate(nodes):
+        metadata = node.metadata if hasattr(node, 'metadata') else {}
+        doc_type = metadata.get('type', 'text')
+        doc_id = metadata.get('document_id', 'unknown')
+        section_id = metadata.get('section_id', '')
+        section_text = metadata.get('section_text', '')
+        section_path = metadata.get('section_path', '')
+        
+        # Create a unique key for grouping
+        if doc_type == 'table':
+            table_num = metadata.get('table_number', 'unknown')
+            key = f"{doc_id}_table_{table_num}"
+        elif doc_type == 'image':
+            image_num = metadata.get('image_number', 'unknown')
+            key = f"{doc_id}_image_{image_num}"
+        else:
+            # For text documents, group by section path or section id
+            section_key = section_path if section_path else section_id
+            key = f"{doc_id}_text_{section_key}"
+        
+        if key not in sources_by_doc:
+            sources_by_doc[key] = {
+                'doc_id': doc_id,
+                'doc_type': doc_type,
+                'metadata': metadata,
+                'sections': set()
+            }
+        
+        # Add section information
+        if section_path:
+            sources_by_doc[key]['sections'].add(f"пункт {section_path}")
+        elif section_id and section_id != 'unknown':
+            sources_by_doc[key]['sections'].add(f"пункт {section_id}")
     
-    html += "</div></div>"
-    return html
+    # Generate HTML for each unique source
+    for source_info in sources_by_doc.values():
+        metadata = source_info['metadata']
+        doc_type = source_info['doc_type']
+        doc_id = source_info['doc_id']
+        
+        html += f"<div style='margin-bottom: 15px; padding: 15px; border: 1px solid #4a5568; border-radius: 8px; background-color: #1a202c;'>"
+        
+        if doc_type == 'text':
+            html += f"<h4 style='margin: 0 0 10px 0; color: #63b3ed;'>📄 {doc_id}</h4>"
 
-def format_chunks_html(chunks_info):
-    if not chunks_info:
-        return "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Нет данных о чанках</div>"
-    
-    html = "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; max-height: 500px; overflow-y: auto;'>"
-    html += f"<h4 style='color: white; margin-bottom: 15px;'>Найдено релевантных чанков: {len(chunks_info)}</h4>"
-    
-    for i, chunk in enumerate(chunks_info):
-        bg_color = "#4a5568" if i % 2 == 0 else "#374151"
+        elif doc_type == 'table' or doc_type == 'table_row':
+            table_num = metadata.get('table_number', 'unknown')
+            table_title = metadata.get('table_title', '')
+            if table_num and table_num != 'unknown':
+                if not str(table_num).startswith('№'):
+                    table_num = f"№{table_num}"
+                html += f"<h4 style='margin: 0 0 10px 0; color: #68d391;'>📊 Таблица {table_num} - {doc_id}</h4>"
+                if table_title and table_title != 'unknown':
+                    html += f"<p style='margin: 5px 0; color: #a0aec0; font-size: 14px;'>{table_title}</p>"
+            else:
+                html += f"<h4 style='margin: 0 0 10px 0; color: #68d391;'>📊 Таблица - {doc_id}</h4>"
+                
+        elif doc_type == 'image':
+            image_num = metadata.get('image_number', 'unknown')
+            image_title = metadata.get('image_title', '')
+            section = metadata.get('section', '')
+            if image_num and image_num != 'unknown':
+                if not str(image_num).startswith('№'):
+                    image_num = f"№{image_num}"
+                html += f"<h4 style='margin: 0 0 10px 0; color: #fbb6ce;'>🖼️ Изображение {image_num} - {doc_id}</h4>"
+                if image_title and image_title != 'unknown':
+                    html += f"<p style='margin: 5px 0; color: #a0aec0; font-size: 14px;'>{image_title}</p>"
+                if section and section != 'unknown':
+                    html += f"<p style='margin: 5px 0; color: #a0aec0; font-size: 12px;'>Раздел: {section}</p>"
+            else:
+                html += f"<h4 style='margin: 0 0 10px 0; color: #fbb6ce;'>🖼️ Изображение - {doc_id}</h4>"
         
-        from app import get_section_display, get_formatted_content
-        section_display = get_section_display(chunk)
-        formatted_content = get_formatted_content(chunk)
+        # Add file link if available
+        if chunks_df is not None and 'file_link' in chunks_df.columns and doc_type == 'text':
+            doc_rows = chunks_df[chunks_df['document_id'] == doc_id]
+            if not doc_rows.empty:
+                file_link = doc_rows.iloc[0]['file_link']
+                html += f"<a href='{file_link}' target='_blank' style='color: #68d391; text-decoration: none; font-size: 14px; display: inline-block; margin-top: 10px;'>🔗 Ссылка на документ</a><br>"
         
-        html += f"""
-        <div style='background-color: {bg_color}; padding: 10px; margin: 5px 0; border-radius: 5px; border-left: 4px solid #60a5fa;'>
-            <strong style='color: #93c5fd;'>Документ:</strong> <span style='color: white;'>{chunk['document_id']}</span><br>
-            <strong style='color: #93c5fd;'>Раздел:</strong> <span style='color: white;'>{section_display}</span><br>
-            <strong style='color: #93c5fd;'>Содержание:</strong><br>
-            <div style='background-color: #1f2937; padding: 8px; margin-top: 5px; border-radius: 3px; font-family: monospace; font-size: 12px; color: #d1d5db; max-height: 200px; overflow-y: auto;'>
-                {formatted_content}
-            </div>
-        </div>
-        """
+        html += "</div>"
     
     html += "</div>"
     return html
 
-def deduplicate_nodes(nodes):
-    """Deduplicate retrieved nodes based on unique identifiers"""
-    seen = set()
-    unique_nodes = []
+
+def answer_question(question, query_engine, reranker, current_model, chunks_df=None):
+    if query_engine is None:
+        return "<div style='background-color: #e53e3e; color: white; padding: 20px; border-radius: 10px;'>Система не инициализирована</div>", "", ""
     
-    for node in nodes:
-        # Create unique identifier from metadata
-        doc_id = node.metadata.get('document_id', '')
-        section_id = node.metadata.get('section_id', '')
-        chunk_id = node.metadata.get('chunk_id', 0)
-        node_type = node.metadata.get('type', 'text')
+    try:
+        start_time = time.time()
         
-        if node_type == 'table':
-            table_num = node.metadata.get('table_number', '')
-            identifier = f"{doc_id}|table|{table_num}|{chunk_id}"
-        elif node_type == 'image':
-            img_num = node.metadata.get('image_number', '')
-            identifier = f"{doc_id}|image|{img_num}"
-        else:
-            identifier = f"{doc_id}|{section_id}|{chunk_id}"
+        llm = get_llm_model(current_model)
         
-        if identifier not in seen:
-            seen.add(identifier)
-            unique_nodes.append(node)
-    
-    return unique_nodes
-
-
-def answer_question(question, query_engine, reranker, model_name):
-    try:
-        log_message(f"\n{'='*70}")
-        log_message(f"QUERY: {question}")
-
-        retrieved = query_engine.retrieve(question)
-        total_retrieved = len(retrieved)
-        log_message(f"RETRIEVED: {total_retrieved} nodes (before deduplication)")
+        # Direct retrieval without query expansion
+        retrieved_nodes = query_engine.retriever.retrieve(question)
         
-        # Deduplicate
-        unique_retrieved = deduplicate_nodes(retrieved)
-        duplicates_removed = total_retrieved - len(unique_retrieved)
-        log_message(f"DEDUPLICATION: {duplicates_removed} duplicates removed")
-        log_message(f"UNIQUE NODES: {len(unique_retrieved)} nodes")
-
-        reranked = rerank_nodes(question, unique_retrieved, reranker, top_k=20, min_score=-0.5)
-        log_message(f"RERANKED: {len(reranked)} nodes (after scoring)")
-
-
-        doc_groups = {}
-        for n in reranked:
-            doc_id = n.metadata.get('document_id', 'unknown')
-            if doc_id not in doc_groups:
-                doc_groups[doc_id] = {'tables': [], 'text': [], 'images': []}
-            
-            node_type = n.metadata.get('type', 'text')
-            if node_type == 'table':
-                doc_groups[doc_id]['tables'].append(n)
-            elif node_type == 'image':
-                doc_groups[doc_id]['images'].append(n)
-            else:
-                doc_groups[doc_id]['text'].append(n)
+        log_message(f"Получено {len(retrieved_nodes)} узлов")
         
-        log_message(f"Documents found: {list(doc_groups.keys())}")
-
-        context_parts = []
-        for doc_id, groups in doc_groups.items():
-            doc_section = [f"=== ДОКУМЕНТ: {doc_id} ==="]
-            
-            if groups['tables']:
-                doc_section.append("\n--- ТАБЛИЦЫ ---")
-                for n in groups['tables']:
-                    meta = n.metadata
-                    table_id = meta.get('table_identifier', meta.get('table_number', 'unknown'))
-                    title = meta.get('table_title', '')
-                    doc_section.append(f"\n[Таблица {table_id}] {title}")
-                    doc_section.append(n.text[:1500])
-                    log_message(f"  Included table {table_id} from {doc_id}")
-            
-            if groups['images']:
-                doc_section.append("\n--- ИЗОБРАЖЕНИЯ ---")
-                for n in groups['images']:
-                    meta = n.metadata
-                    img_id = meta.get('image_number', 'unknown')
-                    doc_section.append(f"\n[Рисунок {img_id}]")
-                    doc_section.append(n.text[:1000])
-                    log_message(f"  Included image {img_id} from {doc_id}")
-            
-            if groups['text']:
-                doc_section.append("\n--- ТЕКСТ ---")
-                for n in groups['text'][:3]:
-                    doc_section.append(n.text[:800])
-                    log_message(f"  Included text section from {doc_id}")
-            
-            context_parts.append("\n".join(doc_section))
-
-        context = "\n\n" + ("="*70 + "\n\n").join(context_parts)
+        reranked_nodes = rerank_nodes(
+            question,  
+            retrieved_nodes, 
+            reranker, 
+            top_k=20,
+            min_score_threshold=0.5, 
+            diversity_penalty=0.3
+        )
         
-        log_message(f"Context length: {len(context)} chars")
-
-        from config import CUSTOM_PROMPT
-        prompt = CUSTOM_PROMPT.format(context_str=context, query_str=question)
-
-        from llama_index.core import Settings
-        response = Settings.llm.complete(prompt)
-
-        sources_text = format_sources(reranked)
-        chunks_info = create_chunks_info_for_display(reranked)
+        formatted_context = format_context_for_llm(reranked_nodes)
         
-        answer_html = format_answer_html(response.text, model_name)
-        sources_html = format_sources_html(sources_text)
-        chunks_html = format_chunks_html(chunks_info)
+        enhanced_question = f"""Контекст из базы данных:
+{formatted_context}
 
-        return answer_html, sources_html, chunks_html
+Вопрос пользователя: {question}
 
-    except Exception as e:
-        log_message(f"Error: {e}")
-        import traceback
-        log_message(traceback.format_exc())
+Инструкция: Ответь на вопрос, используя ТОЛЬКО информацию из контекста выше. 
+Если информации недостаточно, четко укажи это. Цитируй конкретные источники."""
         
-        error_html = f"<div style='background-color: #2d3748; color: #ef4444; padding: 20px; border-radius: 10px;'>Ошибка: {str(e)}</div>"
-        sources_html = "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Источники недоступны из-за ошибки</div>"
-        chunks_html = "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Чанки недоступны из-за ошибки</div>"
+        response = query_engine.query(enhanced_question)
         
-        return error_html, sources_html, chunks_html
-
-def rerank_nodes(query, nodes, reranker, top_k=20, min_score=-0.5):
-    if not nodes or not reranker:
-        log_message("WARNING: No nodes or reranker available")
-        return nodes[:top_k]
+        end_time = time.time()
+        processing_time = end_time - start_time
+        
+        log_message(f"Обработка завершена за {processing_time:.2f}с")
+        
+        sources_html = generate_sources_html(reranked_nodes, chunks_df)
+        
+        answer_with_time = f"""<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; margin-bottom: 10px;'>
+        <h3 style='color: #63b3ed; margin-top: 0;'>Ответ (Модель: {current_model}):</h3>
+        <div style='line-height: 1.6; font-size: 16px;'>{response.response}</div>
+        <div style='margin-top: 15px; padding-top: 10px; border-top: 1px solid #4a5568; font-size: 14px; color: #a0aec0;'>
+        Время обработки: {processing_time:.2f} секунд
+        </div>
+        </div>"""
+        
+        chunk_info = []
+        for node in reranked_nodes:
+            metadata = node.metadata if hasattr(node, 'metadata') else {}
+            chunk_info.append({
+                'document_id': metadata.get('document_id', 'unknown'),
+                'section_id': metadata.get('section_id', metadata.get('section', 'unknown')),
+                'section_path': metadata.get('section_path', ''),
+                'section_text': metadata.get('section_text', ''),
+                'level': metadata.get('level', ''),
+                'parent_section': metadata.get('parent_section', ''),
+                'parent_title': metadata.get('parent_title', ''),
+                'type': metadata.get('type', 'text'),
+                'table_number': metadata.get('table_number', ''),
+                'image_number': metadata.get('image_number', ''),
+                'chunk_size': len(node.text),
+                'chunk_text': node.text
+            })
+        from app import create_chunks_display_html
+        chunks_html = create_chunks_display_html(chunk_info)
 
-    pairs = [[query, n.text[:500]] for n in nodes]
-    scores = reranker.predict(pairs)
-    scored = sorted(zip(nodes, scores), key=lambda x: x[1], reverse=True)
-    
-    if scored:
-        top_5_scores = [s for _, s in scored[:5]]
-        bottom_5_scores = [s for _, s in scored[-5:]]
-        log_message(f"Score range: {min(scores):.3f} to {max(scores):.3f}")
-        log_message(f"Top 5 scores: {top_5_scores}")
-        log_message(f"Bottom 5 scores: {bottom_5_scores}")
-    
-    above_threshold = sum(1 for _, s in scored if s >= min_score)
-    log_message(f"Nodes above threshold ({min_score}): {above_threshold}/{len(scored)}")
-    
-    filtered = [n for n, s in scored if s >= min_score]
-    result = filtered[:top_k] if filtered else [n for n, _ in scored[:top_k]]
-    
-    log_message(f"Returning {len(result)} nodes after reranking")
-    return result
+        return answer_with_time, sources_html, chunks_html
+        
+    except Exception as e:
+        log_message(f"Ошибка: {str(e)}")
+        error_msg = f"<div style='background-color: #e53e3e; color: white; padding: 20px; border-radius: 10px;'>Ошибка: {str(e)}</div>"
+        return error_msg, "", ""