File size: 5,700 Bytes
8f78c8f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
import pandas as pd

from src.llm_calls import llm_extract_table
from src.parse_response import extract_and_return_data_table
import gradio as gr


def __update_df_state(df_before, df_state, updated_df):
    new_df = pd.DataFrame(updated_df)

    if df_before is not None:
        new_df_before = df_before + [df_state]
    else:
        new_df_before = [df_state]

    new_df_after = []  # Clear redo history
    new_df_state = new_df.copy()

    return new_df_before, new_df_state, new_df_after


def extract_table_from_chat(
    chat_output, df_before, df_state, df_after, llm_type, api_key, key="Medications"
):
    try:
        updated_df = extract_and_return_data_table(chat_output=chat_output, key=key)
    except ValueError:
        try:
            json_str = llm_extract_table(chat_output, llm_type, api_key)
            updated_df = extract_and_return_data_table(chat_output=json_str, key=key)
        except KeyError:
            gr.Error(
                "Cannot extract table information from chat. "
                "Please ask the LLM to provide the dataset in JSON format.",
                duration=None,
            )
            updated_df = df_before
        except ValueError:
            gr.Error(
                "Cannot extract table information from chat. "
                "Please ask the LLM to provide the dataset in JSON format.",
                duration=None,
            )
            updated_df = df_before

    new_df_before, new_df_state, new_df_after = __update_df_state(
        df_before, df_state, updated_df
    )
    return (
        new_df_state,
        new_df_before,
        new_df_state,
        new_df_after,
        gr.update(interactive=True),
        gr.update(interactive=False),
    )


def update_llm_selection(selected_llm):
    if selected_llm == "OpenAI":
        return gr.update(label="OpenAI API Key", placeholder="Enter OpenAI API Key")
    elif selected_llm == "Perplexity":
        return gr.update(
            label="Perplexity API Key", placeholder="Enter Perplexity API Key"
        )
    else:
        raise ValueError("Invalid LLM type selected.")


def edit_or_save_changes(updated_df, df_before, df_state, df_after, current_edit_mode):
    """Save user changes, update undo history."""

    new_df = pd.DataFrame(updated_df)
    new_df_before = df_before + [df_state.copy()]
    new_df_after = []  # Clear redo history
    new_df_state = new_df.copy()

    if current_edit_mode == "Save":
        # User wants to move from save to edit
        return (
            new_df,
            new_df_before,
            new_df_state,
            new_df_after,
            gr.update(
                interactive=True
            ),  # prev button is now enabled as there was a change
            gr.update(interactive=False),  # next button
            gr.update(interactive=True),  # df display
            gr.update(value="Edit"),  # edit button
            "Edit",
        )
    elif current_edit_mode == "Edit":
        return (
            new_df_state,
            new_df_before,
            new_df_state,
            new_df_after,
            gr.update(interactive=False),  # prev button
            gr.update(interactive=False),  # next button
            gr.update(interactive=True),  # df display
            gr.update(value="Save"),  # edit button
            "Save",
        )

    else:
        raise ValueError(f"Wrong edit mode selected: {current_edit_mode}. ")


def undo(df_before, df_state, df_after):
    """Undo user change without enabling Save button."""
    if not df_before:
        return (
            df_state,
            df_before,
            df_state,
            df_after,
            gr.update(interactive=False),
            gr.update(interactive=(len(df_after) > 0)),
        )

    new_df_after = df_after + [df_state.copy()]
    new_df_state = df_before[-1]
    new_df_before = df_before[:-1]

    return (
        new_df_state,
        new_df_before,
        new_df_state,
        new_df_after,
        gr.update(interactive=(len(new_df_before) > 0)),  # prev button
        gr.update(interactive=(len(new_df_after) > 0)),  # next button
    )


def redo(df_before, df_state, df_after):
    """Redo user change without enabling Save button."""
    if not df_after:
        return (
            df_state,
            df_before,
            df_state,
            df_after,
            gr.update(interactive=(len(df_before) > 0)),
            gr.update(interactive=False),
        )

    if df_state is None:
        df_state = df_after

    new_df_before = df_before + [df_state.copy()]
    new_df_state = df_after[-1]
    new_df_after = df_after[:-1]

    return (
        new_df_state,
        new_df_before,
        new_df_state,
        new_df_after,
        gr.update(interactive=(len(new_df_before) > 0)),
        gr.update(interactive=(len(new_df_after) > 0)),
    )


# def toggle_save_edit(button_state, dataframe_display):
#     if button_state == "Edit":
#         return "Save", gr.update(interactive=True), gr.update(interactive=True)  # Enable DataFrame editing
#     else:
#         return "Edit", gr.update(interactive=True), gr.update(interactive=False)  # Disable DataFrame editing
#


def upload_file(file, df_before, df_state, df_after):
    if file is None:
        return gr.update()

    df = pd.read_excel(file.name, engine="openpyxl")

    new_df_before, new_df_state, new_df_after = __update_df_state(
        df_before, df_state, df
    )

    # print("Uploaded DataFrame:\n", df)  # Print DataFrame to console
    return (
        df,
        new_df_before,
        df,
        new_df_after,
        gr.update(interactive=False),
        gr.update(interactive=False),
    )