app.py

import streamlit as st
from PyPDF2 import PdfReader
from anthropic import Anthropic
from prompts import DIFFERENTIATE_PROMPT
from concurrent.futures import ProcessPoolExecutor



def extract_differences(input_text):
        input_text = input_text.strip()
        qa_pairs = input_text.split('</difference>')
        
        # Initialize an empty list to hold the parsed dictionary objects
        parsed_data = []
        
        # Iterate over each QA pair
        for pair in qa_pairs:
            # Check if the pair has both question and answer (ignoring the last one)
            if '<text1_section>' in pair and '</text1_section>' in pair and '<text2_section>' in pair and '</text2_section>' in pair and '<explanation>' in pair and '</explanation>' in pair:
                # Extract the question and answer text
                text1 = pair.split('<text1_section>')[1].split('</text1_section>')[0]
                text2 = pair.split('<text2_section>')[1].split('</text2_section>')[0]
                explanation = pair.split('<explanation>')[1].split('</explanation>')[0]
                
                # Create a dictionary for the current pair and append it to the list
                parsed_data.append({'text1': text1.strip(), 'text2': text2.strip(), 'explanation': explanation.strip()})
        
        return parsed_data

st.cache_data()
def make_llm_api_call(prompt):
    client = Anthropic()
    message = client.messages.create(
        model="claude-3-haiku-20240307",
        max_tokens=4096,
        temperature=0,
        messages=[{"role": "user", "content": [{"type": "text", "text": prompt}]}],
    )
    return message


def get_llm_response(extractedtext1, extractedtext2):
    prompt = DIFFERENTIATE_PROMPT.format(text1=extractedtext1, text2=extractedtext2)
    
    message = make_llm_api_call(prompt)
    
    message_text = message.content[0].text

    try:
        try:
            before_differences = message_text.split("<differences>")[0]
            after_differences = message_text.split("</differences>")[1]
            differences_list = extract_differences(message_text.split("<differences>")[1].split("</differences>")[0].strip())
        except Exception as e:
            differences_list = extract_differences(message_text)
    except Exception as e:
        print("Error:", e)
        return message_text, []
    
    difference_content = "\n\n\n".join([f"**Text1:**\n\n{d['text1']}\n\n**Text2:**\n\n{d['text2']}\n\n**Explanation:**\n\n{d['explanation']}\n\n----------------------\n" for d in differences_list])
    # display_text = f"{before_differences}\n\n{difference_content}\n\n{after_differences}"
    display_text = difference_content
    return display_text, differences_list

def process_text_pair(pair):
    etext1, etext2 = pair
    dtext, pdata = get_llm_response(etext1, etext2)
    return dtext, pdata

def process_concurrently(extracted_text1, extracted_text2):
    # Create a pool of worker processes
    with ProcessPoolExecutor(max_workers=5) as executor:
        # Submit tasks to the pool
        futures = [executor.submit(process_text_pair, (etext1, etext2)) 
                   for etext1, etext2 in zip(extracted_text1, extracted_text2)]
        
        # Collect results
        display_text = ""
        parsed_data = []
        for future in futures:
            result = future.result()
            display_text += result[0]
            parsed_data += result[1]
    
    # Combine results
    # display_text = ''.join(display_texts)
    # parsed_data = ''.join(parsed_datas)
    
    return display_text, parsed_data


def extract_text_with_pypdf(pdf_path):
    reader = PdfReader(pdf_path)
    text = []
    for page in reader.pages:
        text.append(page.extract_text())
    return text


def main():
    st.set_page_config(layout="wide")  # Enable wide layout
    st.markdown('<div style="text-align: center;">' + '<h1>PDF Upload and Compare App</h1>' + '</div>', unsafe_allow_html=True)
    
    # Create columns for side-by-side buttons
    col1, col2 = st.columns([2, 2])
        
    # Add upload button to left column
    uploaded_file1 = col1.file_uploader("**Text 1**", type="pdf")
    
    # Add upload button to right column
    uploaded_file2 = col2.file_uploader("**Text 2**", type="pdf")
    
    # Check if both files are uploaded
    if uploaded_file1 and uploaded_file2:
        # Get filenames from uploaded files
        filename1 = uploaded_file1.name
        filename2 = uploaded_file2.name
        
        try:
            extracted_text1 = extract_text_with_pypdf(uploaded_file1)
            extracted_text2 = extract_text_with_pypdf(uploaded_file2)
            with col1.expander(filename1):
                st.write("\n\n".join("\n\n------------------------\n\n".join(extracted_text1).splitlines()))
            
            with col2.expander(filename2):
                st.write("\n\n".join("\n\n------------------------\n\n".join(extracted_text2).splitlines()))
                
            st.success(f"Content of files **{filename1}** and **{filename2}** have been extracted successfully.")
        except Exception as e:
            st.error(f"Error saving files: {str(e)}")

        # Add button at the bottom to run Find Differences function
        if st.button("Find Differences"):
            try:
                display_text = ""
                parsed_data = []
                # for etext1, etext2 in zip(extracted_text1, extracted_text2):
                #     dtext, pdata = get_llm_response(etext1, etext2)
                #     display_text += dtext
                #     parsed_data += pdata
                display_text, parsed_data = process_concurrently(extracted_text1, extracted_text2)
                # display_text, parsed_data = get_llm_response(extracted_text1, extracted_text2)
                display_text1 = "\n\n------------------------\n\n".join(extracted_text1)
                display_text2 = "\n\n------------------------\n\n".join(extracted_text2)

                for diff in parsed_data:
                    diff_text1 = diff['text1'].strip()
                    diff_text2 = diff['text2'].strip()          
                    if diff_text1 == "" or diff_text2 == "": continue              
                    diff_text1_phrase = "\n".join([f"<span style='background-color: grey;'>{t}</span>" for t in diff_text1.splitlines()])
                    diff_text2_phrase = "\n".join([f"<span style='background-color: grey;'>{t}</span>" for t in diff_text2.splitlines()])
                    display_text1 = diff_text1_phrase.join(display_text1.split(diff_text1)) if diff_text1 in display_text1 else display_text1
                    display_text2 = diff_text2_phrase.join(display_text2.split(diff_text2)) if diff_text2 in display_text2 else display_text2
                with col1.expander(f"{filename1} Difference Highlighted"):
                    st.write("\n\n".join(display_text1.splitlines()), unsafe_allow_html=True)

                with col2.expander(f"{filename2} Difference Highlighted"):
                    st.write("\n\n".join(display_text2.splitlines()), unsafe_allow_html=True)
                    
                st.markdown(display_text)
            except Exception as e:
                st.error(f"Error finding differences: {str(e)}")

if __name__ == "__main__":
    main()