pages/input.py

import streamlit as st
from menu import menu_with_redirect

# Standard imports
import numpy as np
import pandas as pd

# Path manipulation
from pathlib import Path

# Custom and other imports
import project_config
from utils import load_kg

# Redirect to app.py if not logged in, otherwise show the navigation menu
menu_with_redirect()

# Header
st.image(str(project_config.MEDIA_DIR / 'input_header.svg'), use_column_width=True)

# Main content
# st.markdown(f"Hello, {st.session_state.name}!")

st.subheader("Construct Query", divider = "red")

# Checkbox to allow reverse edges
allow_reverse_edges = st.checkbox("Reverse Edges", value = False)

# Load knowledge graph
kg_nodes = load_kg()

with st.spinner('Loading knowledge graph...'):
    # kg_nodes = nodes = pd.read_csv(project_config.DATA_DIR / 'kg_nodes.csv', dtype = {'node_index': int}, low_memory = False)
    node_types = pd.read_csv(project_config.DATA_DIR / 'kg_node_types.csv')
    edge_types = pd.read_csv(project_config.DATA_DIR / 'kg_edge_types.csv')

    if not allow_reverse_edges:
       edge_types = edge_types[edge_types.direction == 'forward']

# Select source node type
source_node_type = st.selectbox("Source Node Type", node_types['node_type'],
                                format_func = lambda x: x.replace("_", " "))

# Select source node
source_node = st.selectbox("Source Node", kg_nodes[kg_nodes['node_type'] == source_node_type]['node_name'])

# Select target node type
target_node_type = st.selectbox("Target Node Type", edge_types[edge_types.x_type == source_node_type].y_type.unique(),
                                format_func = lambda x: x.replace("_", " "))

# Select relation
relation = st.selectbox("Edge Type", edge_types[(edge_types.x_type == source_node_type) & (edge_types.y_type == target_node_type)].relation.unique(),
                        format_func = lambda x: x.replace("_", "-"))

# Button to submit query
if st.button("Submit Query"):
    
    # Save query to session state
    st.session_state.query = {
        "source_node_type": source_node_type,
        "source_node": source_node,
        "target_node_type": target_node_type,
        "relation": relation
    }

    # # Write query to console
    # st.write("Current Query:")
    # st.write(st.session_state.query)
    st.write("Query submitted.")

    # Switch to the Predict page
    st.switch_page("pages/predict.py")


st.subheader("Knowledge Graph", divider = "red")    
display_data = kg_nodes[['node_id', 'node_type', 'node_name', 'node_source']].copy()
display_data = display_data.rename(columns = {'node_id': 'ID', 'node_type': 'Type', 'node_name': 'Name', 'node_source': 'Database'})
st.dataframe(display_data, use_container_width = True)