simplify

app.py

@@ -97,7 +97,7 @@ def perform_search(query, top_k):
     """Этап 1: Поиск и возврат результатов"""
    
     if not query:
-        return None, None, [], [], "Введите вопрос для поиска"
+        return None, [], [], "Введите вопрос для поиска"
 
     # Выполняем поиск
     scores = retrieval.bm25_search(query)
@@ -111,7 +111,7 @@ def perform_search(query, top_k):
     
     status = f"Найдено {len(scores)} чанков. Top-{top_k} выбраны."
     
-    return None, scores, chunk_ids, top_k_indices, status
+    return scores, chunk_ids, top_k_indices, status
 
 def filter_chunks_by_documents(top_k_indices, all_scores, selected_docs):
     """Фильтрует чанки по выбранным документам"""
@@ -136,7 +136,75 @@ def filter_chunks_by_documents(top_k_indices, all_scores, selected_docs):
     
     return filtered_indices
 
-def format_retrieval_results(filtered_indices, all_scores, top_k_results):
+def format_selected_chunks(selected_indices):
+    """Форматирует выбранные чанки в единый текст для вывода и LLM
+    
+    Возвращает текст в формате:
+    Документ {название}:
+    Предложение1 Предложение2 ... Предложение5 Предложение6
+    """
+    if not selected_indices:
+        return ""
+    
+    # Форматируем результаты для вывода
+    chunks_with_info = []
+    for idx in selected_indices:
+        if idx >= len(retrieval.docs_metadata) or idx >= len(retrieval.chunks):
+            continue
+        
+        doc_id = retrieval.docs_metadata[idx]
+        doc_name = retrieval.docs_names[doc_id] if doc_id < len(retrieval.docs_names) else "Неизвестный документ"
+        
+        chunks_with_info.append({
+            'index': idx,
+            'doc_name': doc_name,
+            'doc_id': doc_id,
+            'chunk_text': retrieval.chunks[idx]
+        })
+    
+    if not chunks_with_info:
+        return "Нет валидных чанков"
+    
+    # Группируем чанки по документам
+    docs_chunks = {}
+    for chunk_info in chunks_with_info:
+        doc_name = chunk_info['doc_name']
+        if doc_name not in docs_chunks:
+            docs_chunks[doc_name] = []
+        docs_chunks[doc_name].append(chunk_info['index'])
+    
+    # Форматируем вывод
+    result_lines = []
+    for doc_name in sorted(docs_chunks.keys()):
+        chunk_indices = sorted(docs_chunks[doc_name])
+        
+        # Группируем подряд идущие индексы
+        groups = []
+        current_group = [chunk_indices[0]]
+        
+        for i in range(1, len(chunk_indices)):
+            if chunk_indices[i] == chunk_indices[i-1] + 1:
+                current_group.append(chunk_indices[i])
+            else:
+                groups.append(current_group)
+                current_group = [chunk_indices[i]]
+        groups.append(current_group)
+        
+        # Собираем текст для каждой группы БЕЗ запятых
+        group_texts = []
+        for group in groups:
+            sentences = [retrieval.chunks[idx] for idx in group]
+            group_texts.append(" ".join(sentences))
+        
+        # Выводим документ с многоточием между группами
+        doc_output = f"Документ {doc_name}:\n" + " ... ".join(group_texts)
+        result_lines.append(doc_output)
+        result_lines.append("")  # Пустая строка между документами
+    
+    return "\n".join(result_lines)
+
+
+def format_retrieval_results(filtered_indices, top_k_results):
     """Форматирует результаты retrieval для отображения в текстовом поле
     
     Алгоритм:
@@ -217,62 +285,7 @@ def format_retrieval_results(filtered_indices, all_scores, top_k_results):
         if added_count > top_k_results or iteration_added == 0:
             break
     
-    # Форматируем результаты для вывода
-    chunks_with_info = []
-    for idx in selected:
-        if idx >= len(retrieval.docs_metadata) or idx >= len(retrieval.chunks):
-            continue
-        
-        doc_id = retrieval.docs_metadata[idx]
-        doc_name = retrieval.docs_names[doc_id] if doc_id < len(retrieval.docs_names) else "Неизвестный документ"
-        
-        chunks_with_info.append({
-            'index': idx,
-            'doc_name': doc_name,
-            'doc_id': doc_id,
-            'chunk_text': retrieval.chunks[idx]
-        })
-    
-    if not chunks_with_info:
-        return "Нет валидных чанков"
-    
-    # Группируем чанки по документам
-    docs_chunks = {}
-    for chunk_info in chunks_with_info:
-        doc_name = chunk_info['doc_name']
-        if doc_name not in docs_chunks:
-            docs_chunks[doc_name] = []
-        docs_chunks[doc_name].append(chunk_info['index'])
-    
-    # Форматируем вывод
-    result_lines = []
-    for doc_name in sorted(docs_chunks.keys()):
-        chunk_indices = sorted(docs_chunks[doc_name])
-        
-        # Группируем подряд идущие индексы
-        groups = []
-        current_group = [chunk_indices[0]]
-        
-        for i in range(1, len(chunk_indices)):
-            if chunk_indices[i] == chunk_indices[i-1] + 1:
-                current_group.append(chunk_indices[i])
-            else:
-                groups.append(current_group)
-                current_group = [chunk_indices[i]]
-        groups.append(current_group)
-        
-        # Собираем текст для каждой группы
-        group_texts = []
-        for group in groups:
-            sentences = [retrieval.chunks[idx] for idx in group]
-            group_texts.append(", ".join(sentences))
-        
-        # Выводим документ с многоточием между группами
-        doc_output = f"Документ {doc_name}:\n" + " ... ".join(group_texts)
-        result_lines.append(doc_output)
-        result_lines.append("")  # Пустая строка между документами
-    
-    return "\n".join(result_lines)
+    return format_selected_chunks(list(selected))
 
 def ask_llm(query, filtered_indices_state):
     """Этап 2: Отправка отфильтрованных чанков в LLM с потоковой выдачей"""
@@ -287,40 +300,13 @@ def ask_llm(query, filtered_indices_state):
         yield "Нет выбранных чанков для отправки в LLM"
         return
     
-    # Сортируем чанки сначала по документу, потом по chunk_id
-    chunks_with_doc = []
-    for idx in chunks_to_use:
-        if idx >= len(retrieval.docs_metadata):
-            continue
-        doc_id = retrieval.docs_metadata[idx]
-        doc_name = retrieval.docs_names[doc_id]
-        chunks_with_doc.append((doc_name, idx, doc_id))
+    # Форматируем контекст используя ту же функцию, что и в интерфейсе
+    context = format_selected_chunks(list(chunks_to_use))
     
-    if not chunks_with_doc:
+    if not context or context == "Нет валидных чанков":
         yield "Нет валидных чанков для отправки"
         return
     
-    # Сортируем: сначала по имени документа, потом по chunk_id
-    chunks_with_doc.sort(key=lambda x: (x[0], x[1]))
-    
-    # Собираем текст выбранных чанков в правильном порядке
-    context_parts = []
-    current_doc = None
-    
-    for doc_name, idx, doc_id in chunks_with_doc:
-        # Добавляем разделитель между документами
-        if current_doc != doc_name:
-            if current_doc is not None:
-                context_parts.append("\n---\n")
-            context_parts.append(f"=== Документ: {doc_name} ===\n")
-            current_doc = doc_name
-        
-        # Добавляем чанк с его номером
-        chunk_text = retrieval.chunks[idx]
-        context_parts.append(f"[Чанк {idx}]\n{chunk_text}\n")
-    
-    context = "".join(context_parts)
-    
     # Формируем промпт и отправляем в LLM
     prompt = wrap_prompt(context, query)
     
@@ -417,14 +403,14 @@ with gr.Blocks(title="RAG Application", theme=gr.themes.Soft()) as iface:
     search_btn.click(
         fn=perform_search,
         inputs=[search_query_input, top_k_slider],
-        outputs=[None, all_scores_state, all_chunk_ids_state, top_k_indices_state, search_status]
+        outputs=[all_scores_state, all_chunk_ids_state, top_k_indices_state, search_status]
     ).then(
         fn=filter_chunks_by_documents,
         inputs=[top_k_indices_state, all_scores_state, docs_after],
         outputs=[filtered_indices_state]
     ).then(
         fn=format_retrieval_results,
-        inputs=[filtered_indices_state, all_scores_state, display_k_slider],
+        inputs=[filtered_indices_state, display_k_slider],
         outputs=[retrieval_results]
     )
     
@@ -435,14 +421,14 @@ with gr.Blocks(title="RAG Application", theme=gr.themes.Soft()) as iface:
         outputs=[filtered_indices_state]
     ).then(
         fn=format_retrieval_results,
-        inputs=[filtered_indices_state, all_scores_state, display_k_slider],
+        inputs=[filtered_indices_state, display_k_slider],
         outputs=[retrieval_results]
     )
     
     # Обработчик изменения слайдера отображения
     display_k_slider.change(
         fn=format_retrieval_results,
-        inputs=[filtered_indices_state, all_scores_state, display_k_slider],
+        inputs=[filtered_indices_state, display_k_slider],
         outputs=[retrieval_results]
     )