import streamlit as st
import latex2markdown
from langchain.docstore.document import Document

def format_document(document: Document):
    """TODO: Implement a nice style"""
    return document.dict()

def format_result_set(result):
    st.write(latex2markdown.LaTeX2Markdown(result["result"]).to_markdown())

    agree = st.checkbox('Show source documents')
    source_documents = result["source_documents"]
    if agree:
        st.write('Source Documents:')
        for document in source_documents:
            st.write(format_document(document))