interface/components.py

import streamlit as st
import core.pipelines as pipelines_functions
from inspect import getmembers, isfunction
from networkx.drawing.nx_agraph import to_agraph


def component_select_pipeline(container):
    pipeline_names, pipeline_funcs = list(
        zip(*getmembers(pipelines_functions, isfunction))
    )
    pipeline_names = [
        " ".join([n.capitalize() for n in name.split("_")]) for name in pipeline_names
    ]
    with container:
        selected_pipeline = st.selectbox(
            "Select pipeline",
            pipeline_names,
            index=pipeline_names.index("Keyword Search")
            if "Keyword Search" in pipeline_names
            else 0,
        )
        (
            st.session_state["search_pipeline"],
            st.session_state["index_pipeline"],
        ) = pipeline_funcs[pipeline_names.index(selected_pipeline)]()


def component_show_pipeline(container, pipeline):
    """Draw the pipeline"""
    with st.expander("Show pipeline"):
        graphviz = to_agraph(pipeline.graph)
        graphviz.layout("dot")
        st.graphviz_chart(graphviz.string())


def component_show_search_result(container, results):
    with container:
        for idx, document in enumerate(results):
            st.markdown(f"### Match {idx+1}")
            st.markdown(f"**Text**: {document['text']}")
            st.markdown(f"**Document**: {document['id']}")
            if document["score"] is not None:
                st.markdown(f"**Score**: {document['score']:.3f}")
            st.markdown("---")


def component_text_input(container):
    """Draw the Text Input widget"""
    with container:
        texts = []
        doc_id = 1
        with st.expander("Enter documents"):
            while True:
                text = st.text_input(f"Document {doc_id}", key=doc_id)
                if text != "":
                    texts.append({"text": text})
                    doc_id += 1
                    st.markdown("---")
                else:
                    break
        corpus = [
            {"text": doc["text"], "id": doc_id} for doc_id, doc in enumerate(texts)
        ]
        return corpus