app.py

import streamlit as st
from typing import List, Dict, Any, Callable
from finalthesis import rag_chains
from finalthesis import retrievers
# from langchain.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.documents import Document
# document import Document

st.title('RAG application')
# vectorstore = Chroma(
#     persist_directory = "storage_new/fact_books_fixed_size_vectorstore_chroma",
#     embedding_function = OpenAIEmbeddings()
# )

parent_retriever = retrievers.load_parent_retriever(
    "storage_final_retrievers/parent_retriever",
    4000,
    500,
    )

llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)



def flatten_list(nested_list):
    """
    Flattens a list of lists into a single list.
    If the input is already a flat list, it returns the list unchanged.

    Parameters:
    nested_list (list): The list to be flattened.

    Returns:
    list: A flattened list.
    """
    # Check if the input is a simple list (not a list of lists)
    if all(not isinstance(i, list) for i in nested_list):
        return nested_list

    # Flatten the list of lists
    flattened = []
    for sublist in nested_list:
        if isinstance(sublist, list):
            flattened.extend(sublist)
        else:
            flattened.append(sublist)
    
    return flattened

def get_page_content_as_text(doc_list:list[Document]) -> str:
    doc_list = flatten_list(doc_list)
    page_contents = [doc.page_content for doc in doc_list]
    page_contents = set(page_contents)
    return "\n\n---\n\n".join(page_contents)

def generate_response (text: str, answer_generation_func: Callable[str, dict]) -> str:
    response = answer_generation_func(text)
    return response


st.info("Enter a question that the system should answer.\n"
        "The system was trained on three books about European history, sourced from Project Gutenberg.\n"
        "The books are:\n"
        "- [A History of Germany from the Earliest Times to the Present Day by Bayard Taylor](https://www.gutenberg.org/ebooks/36484)\n"
        "- [Early European History by Hutton Webster](https://www.gutenberg.org/ebooks/7960)\n"
        "- [Medieval Europe by H. W. Carless Davis](https://www.gutenberg.org/ebooks/6369)\n"
        )
show_example_questions = st.checkbox('Show example questions', value=True)
if show_example_questions:
    st.info(
        "To give you some inspiration, here are some example questions:\n"
        "- How did the Hebrew kingdoms become part of the Persian Empire?\n"
        "- What was Otto I's decision regarding assuming the duties of an Italian King?\n"
        "- What were some of the characteristics of German cities during the later Middle Ages?\n"
        "- What was the objective of the Third Crusade in relation to the Sultanate of Egypt?\n"
        "- What are some features of Assyrian warfare?\n"
        "- What was the significance of the Council of Nicaea in the early Christian Church?\n"
        "- What impact did the assessment of heavy taxes on the Romans by the Germans have on their political, economic, and social relationships during the period of Germanic invasions in the Roman Empire?\n"
        "- What was the prevalence of city life in the Roman Empire, and what were some important cities within it?\n"
        "- What led to the decline of the Merovingian dynasty?\n"
        "- What impact did the success of the First Crusade have on Papal power and the Church's faction hostile to Henry IV?\n"
        "- How did medieval architects address the issue of outward thrust in vaulted ceilings?\n"
    )


options = ['[RAG-system](https://arxiv.org/abs/2005.11401)', 
           '[RAT-system](https://arxiv.org/abs/2403.05313)',]


st.session_state['selected_option']  = st.radio('Select an system to use:', options, )


if st.session_state['selected_option'] == '[RAT-system](https://arxiv.org/abs/2403.05313)':
    st.warning(
        "The RAT-system will create a more fully fledged answer as it iteratively " 
        "checks its answers by going over the retrieved documents. "
        "However, in the analysis that was performed, it was noted that the faithfulness "
        "of the answers was not always as high as the RAG-system's. \n"
        "Thus, the model might hallucinate more information!"
    )
    parent_retriever.search_kwargs = {'k': 8}
    chain = rag_chains.RATChain(
        retriever=parent_retriever,
        llm=llm,
        )
else:
    parent_retriever.search_kwargs = {'k': 16}
    chain = rag_chains.build_naive_rag_chain(
        retriever=parent_retriever,
        llm=llm,
    )
question = st.text_area('Enter question:')
submitted = st.button('Submit')
# submitted = st.form_submit_button('Submit')



if submitted:
    if question:
        with st.spinner(text="Response is being generated..."):
            response = generate_response(question, chain.invoke)
            st.info("QUESTION:  \n" + response["question"])
            st.info("ANSWER:  \n" + response["answer"])
            st.info("CONTEXT:  \n" + get_page_content_as_text(response["context"]))
    else:
        st.error("Please enter a question.")
        # question: