summaryrefslogtreecommitdiff
path: root/rag/ui.py
blob: 37c50ddd53e7cfbf1a9dbf74e491a67c6543c30a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import streamlit as st

from langchain_community.document_loaders.blob_loaders import Blob

try:
    from rag.rag import RAG
except ModuleNotFoundError:
    from rag import RAG

rag = RAG()


def upload_pdfs():
    files = st.file_uploader(
        "Choose pdfs to add to the knowledge base",
        type="pdf",
        accept_multiple_files=True,
    )
    for file in files:
        blob = Blob.from_data(file.read())
        rag.add_pdf_from_blob(blob)


if __name__ == "__main__":
    ss = st.session_state
    st.header("RAG-UI")

    upload_pdfs()
    query = st.text_area(
        "query",
        key="query",
        height=100,
        placeholder="Enter query here",
        help="",
        label_visibility="collapsed",
        disabled=False,
    )

    (b,) = st.columns(1)
    (result_column, context_column) = st.columns(2)

    if b.button("Generate", disabled=False, type="primary", use_container_width=True):
        query = ss.get("query", "")
        with st.spinner("Generating answer..."):
            response = rag.retrieve(query)

        with result_column:
            st.markdown("### Answer")
            st.markdown(response.answer)

        with context_column:
            st.markdown("### Context")
            for c in response.context:
                st.markdown(c)
                st.markdown("---")