src/utils/data.py

import json
from typing import Any

import pandas as pd
import streamlit as st
from functools import reduce


def get_filter_values(df: pd.DataFrame, column_name: str) -> list:
    return df[column_name].unique().tolist()


def build_filter(
    meta_data: pd.DataFrame,
    authors_filter: list[str],
    draft_cats_filter: list[str],
    round_filter: list[int],
) -> dict[str, int | str] | dict:
    authors = authors_filter
    round_number = round_filter
    draft_cats = draft_cats_filter

    # set authors_flag to True if not empty list
    authors_flag = True if len(authors) > 0 else False
    draft_cats_flag = True if len(draft_cats) > 0 else False
    round_number_flag = True if len(round_number) > 0 else False

    conditions = []

    if authors_flag:
        authors_condition = (meta_data[col] == 1 for col in authors)
        authors_conditions_list = reduce(lambda a, b: a | b, authors_condition)
        conditions.append(authors_conditions_list)

    if draft_cats_flag:
        draft_cat_condition = (meta_data[col] for col in draft_cats)
        draft_cat_conditions_list = reduce(lambda a, b: a | b, draft_cat_condition)
        conditions.append(draft_cat_conditions_list)

    if round_number_flag:
        round_condition = meta_data["round"].isin(round_number)
        conditions.append(round_condition)

    if len(conditions) == 0:
        filtered_retriever_ids = []
    else:
        final_condition = reduce(lambda a, b: a & b, conditions)
        filtered_retriever_ids = meta_data[final_condition]["retriever_id"].tolist()
    if len(filtered_retriever_ids) == 0:
        return {}
    else:
        return {"retriever_id": filtered_retriever_ids}


def load_json(file_path: str) -> dict:
    with open(file_path, "r") as f:
        return json.load(f)


def save_json(file_path: str, data: dict) -> None:
    with open(file_path, "w") as f:
        json.dump(data, f, indent=4)


def get_meta(result: dict[str, Any]) -> list[dict[str, Any]]:
    meta_data = []
    for doc in result["documents"]:
        current_meta = doc.meta
        current_meta["content"] = doc.content
        meta_data.append(current_meta)
    return meta_data


def load_css(file_name) -> None:
    with open(file_name) as f:
        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)