Update app.py

app.py

@@ -79,34 +79,103 @@ def scroll_to_bottom():
 # Core processing functions
 @st.cache_data(show_spinner=False, ttl=3600)
 @handle_errors
-def summarize_pdf(_pdf_file_path, num_clusters=10):
-    embeddings_model = OpenAIEmbeddings(model="text-embedding-3-small", api_key=openai_api_key)
-    llm = ChatOpenAI(model="gpt-4", api_key=openai_api_key, temperature=0.3)
+def summarize_pdf(pdf_file_path, num_clusters=10):
+    # Keep track of page numbers for each chunk
+    loader = PyMuPDFLoader(pdf_file_path)
+    docs = loader.load()
+    
+    # Create chunks with page numbers
+    text_splitter = SpacyTextSplitter(chunk_size=500)
+    chunks_with_metadata = []
+    for doc in docs:
+        chunks = text_splitter.split_text(doc.page_content)
+        for chunk in chunks:
+            chunks_with_metadata.append({
+                "text": chunk,
+                "page": doc.metadata["page"] + 1  # Convert to 1-based numbering
+            })
     
+    # Modified prompt for citation formatting
     prompt = ChatPromptTemplate.from_template(
-        """Generate a comprehensive summary with these elements:
-        1. Key findings and conclusions
-        2. Main methodologies used
-        3. Important data points
-        4. Limitations mentioned
-        Context: {topic}"""
+        """Generate a summary with inline citations for each key point using [Source X] format. 
+        Structure your response as:
+        
+        ## Comprehensive Summary
+        {summary_content}
+        
+        ## Source References
+        {sources_list}
+        
+        Contexts: {topic}"""
     )
     
-    loader = PyMuPDFLoader(_pdf_file_path)
-    docs = loader.load()
-    full_text = "\n".join(doc.page_content for doc in docs)
-    cleaned_full_text = clean_text(remove_references(full_text))
+    # Create source mapping
+    sources = [f"Source {i+1}: Page {chunk['page']}" 
+              for i, chunk in enumerate(chunks_with_metadata)]
     
-    text_splitter = SpacyTextSplitter(chunk_size=500)
-    split_contents = text_splitter.split_text(cleaned_full_text)
+    # Generate summary with citations
+    chain = prompt | llm | StrOutputParser()
+    results = chain.invoke({
+        "topic": ' '.join([chunk["text"] for chunk in chunks_with_metadata]),
+        "sources_list": "\n".join(sources)
+    })
+    
+    return add_interactive_citations(results, chunks_with_metadata)
+
+
+def add_interactive_citations(summary_text, source_chunks):
+    # Create source boxes with page numbers and full text
+    sources_html = """<div style="margin-top: 20px; border-top: 2px solid #e0e0e0; padding-top: 15px;">
+                        <h4>📚 Source References</h4>"""
     
-    embeddings = embeddings_model.embed_documents(split_contents)
-    kmeans = KMeans(n_clusters=num_clusters, random_state=0).fit(embeddings)
-    closest_indices = [np.argmin(np.linalg.norm(embeddings - center, axis=1)) 
-                      for center in kmeans.cluster_centers_]
+    for idx, chunk in enumerate(source_chunks):
+        sources_html += f"""
+        <div id="source-{idx+1}" style="margin: 10px 0; padding: 10px; 
+                      border: 1px solid #e0e0e0; border-radius: 5px;
+                      transition: all 0.3s ease;">
+            <div style="display: flex; justify-content: space-between;">
+                <strong>Source {idx+1}</strong>
+                <span style="color: #666;">Page {chunk['page']}</span>
+            </div>
+            <div style="margin-top: 5px; color: #444; font-size: 0.9em;">
+                {chunk['text']}
+            </div>
+        </div>
+        """
+    sources_html += "</div>"
     
-    chain = prompt | llm | StrOutputParser()
-    return chain.invoke({"topic": ' '.join([split_contents[idx] for idx in closest_indices])})
+    # Add click interactions
+    interaction_js = """
+    <script>
+    document.querySelectorAll('[data-citation]').forEach(item => {
+        item.addEventListener('click', function(e) {
+            const sourceId = this.getAttribute('data-source');
+            const sourceDiv = document.getElementById(sourceId);
+            
+            // Highlight animation
+            sourceDiv.style.border = '2px solid #4CAF50';
+            sourceDiv.style.boxShadow = '0 2px 8px rgba(76,175,80,0.3)';
+            
+            setTimeout(() => {
+                sourceDiv.style.border = '1px solid #e0e0e0';
+                sourceDiv.style.boxShadow = 'none';
+            }, 1000);
+            
+            // Smooth scroll
+            sourceDiv.scrollIntoView({behavior: 'smooth'});
+        });
+    });
+    </script>
+    """
+    
+    # Replace citations with interactive elements
+    cited_summary = re.sub(r'\[Source (\d+)\]', 
+        lambda m: f'<a data-citation="true" data-source="source-{m.group(1)}" '
+                  f'style="cursor: pointer; color: #4CAF50; text-decoration: underline;">'
+                  f'[Source {m.group(1)}]</a>', 
+        summary_text)
+    
+    return f"{cited_summary}{sources_html}{interaction_js}"
 
 @st.cache_data(show_spinner=False, ttl=3600)
 @handle_errors
@@ -114,34 +183,50 @@ def qa_pdf(_pdf_file_path, query, num_clusters=5):
     embeddings_model = OpenAIEmbeddings(model="text-embedding-3-small", api_key=openai_api_key)
     llm = ChatOpenAI(model="gpt-4", api_key=openai_api_key, temperature=0.3)
     
-    prompt = ChatPromptTemplate.from_template(
-        """Answer this question: {question}
-        Using only this context: {context}
-        Format your answer with:
-        - Clear section headings
-        - Bullet points for lists
-        - Bold key terms
-        - Citations from the text"""
-    )
-    
+    # Load PDF with page numbers
     loader = PyMuPDFLoader(_pdf_file_path)
     docs = loader.load()
-    full_text = "\n".join(doc.page_content for doc in docs)
-    cleaned_full_text = clean_text(remove_references(full_text))
     
+    # Create chunks with page metadata
     text_splitter = SpacyTextSplitter(chunk_size=500)
-    split_contents = text_splitter.split_text(cleaned_full_text)
+    chunks_with_metadata = []
+    for doc in docs:
+        chunks = text_splitter.split_text(doc.page_content)
+        for chunk in chunks:
+            chunks_with_metadata.append({
+                "text": clean_text(chunk),
+                "page": doc.metadata["page"] + 1
+            })
     
+    # Find relevant chunks
+    embeddings = embeddings_model.embed_documents([chunk["text"] for chunk in chunks_with_metadata])
     query_embedding = embeddings_model.embed_query(query)
-    similarities = cosine_similarity([query_embedding], 
-                                   embeddings_model.embed_documents(split_contents))[0]
+    similarities = cosine_similarity([query_embedding], embeddings)[0]
     top_indices = np.argsort(similarities)[-num_clusters:]
     
+    # Prepare prompt with citation instructions
+    prompt = ChatPromptTemplate.from_template(
+        """Answer this question with inline citations using [Source X] format:
+        {question}
+        
+        Use these verified sources:
+        {context}
+        
+        Structure your answer with:
+        - Clear section headings
+        - Bullet points for lists
+        - Citations for all factual claims"""
+    )
+    
     chain = prompt | llm | StrOutputParser()
-    return chain.invoke({
+    raw_answer = chain.invoke({
         "question": query,
-        "context": ' '.join([split_contents[i] for i in top_indices])
+        "context": '\n\n'.join([f"Source {i+1} (Page {chunks_with_metadata[i]['page']}): {chunks_with_metadata[i]['text']}" 
+                              for i in top_indices])
     })
+    
+    return generate_interactive_citations(raw_answer, [chunks_with_metadata[i] for i in top_indices])
+
 
 @st.cache_data(show_spinner=False, ttl=3600)
 @handle_errors