sequential citation numbering fix

.gitignore

@@ -1,2 +1,4 @@
 .DS_Store
-.env
+.envapp.log
+__pycache__/
+.env

utils/generator.py

@@ -90,7 +90,7 @@ def _extract_sources(processed_results: List[Dict[str, Any]], cited_numbers: Lis
         logger.info(f"Processing citation [{citation_num}] -> source_index: {source_index}")
         
         if 0 <= source_index < len(processed_results):
-            source = processed_results[source_index]
+            source = processed_results[source_index].copy()  # Make a copy to avoid modifying original
             cited_sources.append(source)
             logger.info(f"✓ Added source {citation_num}: filename='{source.get('filename', 'Unknown')}', page='{source.get('page', 'Unknown')}'")
         else:
@@ -249,6 +249,32 @@ async def _call_llm_streaming(messages: list) -> AsyncGenerator[str, None]:
 # ---------------------------------------------------------------------
 # Main Generation Functions
 # ---------------------------------------------------------------------
+def _renumber_citations(response: str, cited_numbers: List[int]) -> str:
+    """
+    Renumber citations in response to be sequential to match ChatUI display
+    We do this because otherwise when not all sources are cited, we would get not sequential source listing which is weird for the user
+    """
+    if not cited_numbers:
+        return response
+    
+    # Create mapping from original citation numbers to sequential numbers
+    citation_mapping = {str(original): str(i+1) for i, original in enumerate(cited_numbers)}
+    
+    logger.info(f"=== CITATION RENUMBERING DEBUG ===")
+    logger.info(f"Original citation numbers: {cited_numbers}")
+    logger.info(f"Citation mapping: {citation_mapping}")
+    
+    # Replace citations in response text
+    updated_response = response
+    for original, sequential in citation_mapping.items():
+        # Replace [original] with [sequential]
+        pattern = rf'\[{re.escape(original)}\]'
+        replacement = f'[{sequential}]'
+        updated_response = re.sub(pattern, replacement, updated_response)
+        logger.info(f"Replacing [{original}] with [{sequential}]")
+    
+    return updated_response
+
 async def generate(query: str, context: Union[str, List[Dict[str, Any]]], chatui_format: bool = False) -> Union[str, Dict[str, Any]]:
     """Generate an answer to a query using provided context through RAG"""
     if not query.strip():
@@ -270,6 +296,11 @@ async def generate(query: str, context: Union[str, List[Dict[str, Any]]], chatui
             if processed_results:
                 cited_numbers = _parse_citations(answer)
                 cited_sources = _extract_sources(processed_results, cited_numbers)
+                
+                # Renumber citations to match sequential source numbering
+                answer = _renumber_citations(answer, cited_numbers)
+                result["answer"] = answer
+                
                 result["sources"] = _create_sources_list(cited_sources)
             return result
         else:
@@ -312,6 +343,13 @@ async def generate_streaming(query: str, context: Union[str, List[Dict[str, Any]
         if chatui_format and processed_results:
             cited_numbers = _parse_citations(cleaned_response)
             cited_sources = _extract_sources(processed_results, cited_numbers)
+            
+            # Renumber citations to match sequential source numbering
+            final_response = _renumber_citations(cleaned_response, cited_numbers)
+            
+            # Send the renumbered response as a correction
+            yield {"event": "correction", "data": final_response}
+            
             sources = _create_sources_list(cited_sources)
             yield {"event": "sources", "data": {"sources": sources}}