app.py

import gradio as gr
from dotenv import load_dotenv

# Import from other modules
from ibfs import start_ibfs, handle_choice
from zero_shot import start_zero_shot

# Load environment variables
load_dotenv()



# Define dependent variables for user evaluation
DV_QUESTIONS = [
    "How satisfied are you with the answer? (1-5)",
    "How clear was the explanation? (1-5)",
    "How relevant was the answer to your query? (1-5)",
    "How confident are you in the accuracy of the answer? (1-5)",
    "Would you use this method again for similar questions? (Yes/No)"
]


def save_dv_responses(user_id, method, responses):
    """Save user's responses to dependent variable questions."""
    from utils import save_results
    import json
    from datetime import datetime

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"ibfs_results/{user_id}_{method}_dv_{timestamp}.json"

    with open(filename, "w") as f:
        json.dump(responses, f, indent=2)

    return filename


def process_dv_responses(state, *responses):
    """Process and save the user's responses to dependent variables."""
    if not state:
        return "No active session"

    user_id = state.get("user_id", "unknown")
    method = state.get("method", "unknown")

    # Create a dictionary of responses
    response_dict = {
        "user_id": user_id,
        "method": method,
        "responses": {}
    }

    for i, response in enumerate(responses):
        response_dict["responses"][f"question_{i + 1}"] = {
            "question": DV_QUESTIONS[i],
            "response": response
        }

    # Save responses
    save_path = save_dv_responses(user_id, method, response_dict)

    return f"Thank you for your feedback! Responses saved to {save_path}"


def create_ibfs_interface():
    """Create the IBFS tab interface."""
    with gr.Column():
        gr.Markdown("# IBFS")
        gr.Markdown("Enter your query and set parameters to explore different strategies.")

        with gr.Row():
            with gr.Column(scale=3):
                ibfs_query_input = gr.Textbox(
                    label="Query",
                    placeholder="Enter your question here...",
                    lines=3
                )

            with gr.Column(scale=1):
                k_slider = gr.Slider(
                    minimum=2,
                    maximum=5,
                    step=1,
                    value=3,
                    label="k (Branching Factor - options per step)"
                )
                m_slider = gr.Slider(
                    minimum=1,
                    maximum=3,
                    step=1,
                    value=2,
                    label="m (Depth - number of iterations)"
                )

        ibfs_start_btn = gr.Button("Start IBFS Process", variant="primary")

        # State for maintaining context between steps
        ibfs_state = gr.State(None)

        # Chat display
        ibfs_chatbot = gr.Chatbot(
            label="Interactive Search Process",
            height=600,
            type="messages"
        )

        # User choice input for selecting strategies
        ibfs_choice_input = gr.Textbox(
            label="Enter the number of your choice (e.g., 1, 2, 3...)",
            placeholder="Type a number and press Enter",
            lines=1
        )

        # Event handlers for IBFS
        ibfs_start_btn.click(
            fn=start_ibfs,
            inputs=[ibfs_query_input, k_slider, m_slider],
            outputs=[ibfs_state, ibfs_chatbot],
            show_progress="minimal"
        )

        ibfs_choice_input.submit(
            fn=handle_choice,
            inputs=[ibfs_state, ibfs_choice_input],
            outputs=[ibfs_state, ibfs_chatbot],
            show_progress="minimal"
        )

        ibfs_choice_input.submit(
            fn=lambda: "",
            inputs=[],
            outputs=[ibfs_choice_input]
        )

        # Dependent Variables Section (initially hidden)
        with gr.Accordion("Evaluation Questions", open=False, visible=False) as ibfs_dv_accordion:
            ibfs_dv_inputs = []
            for i, question in enumerate(DV_QUESTIONS):
                if "1-5" in question:
                    dv_input = gr.Slider(
                        minimum=1,
                        maximum=5,
                        step=1,
                        value=3,
                        label=question
                    )
                elif "Yes/No" in question:
                    dv_input = gr.Radio(
                        choices=["Yes", "No"],
                        label=question
                    )
                else:
                    dv_input = gr.Textbox(
                        label=question
                    )
                ibfs_dv_inputs.append(dv_input)

            ibfs_submit_dv_btn = gr.Button("Submit Evaluation", variant="primary")
            ibfs_dv_result = gr.Markdown("")

            # Function to show DV accordion after final answer is generated
            def show_dv_section(state):
                if state and state.get("current_step", 0) == 0 and state.get("strategy_path", []):
                    # This means we've completed a cycle and have a final answer
                    state["method"] = "ibfs"  # Add method to state for DV processing
                    return gr.update(visible=True, open=True)
                return gr.update(visible=False, open=False)

            # Update visibility of DV section when state changes
            ibfs_state.change(
                fn=show_dv_section,
                inputs=[ibfs_state],
                outputs=[ibfs_dv_accordion]
            )

            # Handle DV submission
            ibfs_submit_dv_btn.click(
                fn=process_dv_responses,
                inputs=[ibfs_state] + ibfs_dv_inputs,
                outputs=ibfs_dv_result
            )

    return (ibfs_query_input, k_slider, m_slider, ibfs_start_btn,
            ibfs_state, ibfs_chatbot, ibfs_choice_input)


def create_zero_shot_interface():
    """Create the Zero-Shot tab interface."""
    with gr.Column():
        gr.Markdown("# Zero-Shot Direct Answer")
        gr.Markdown("Enter your query to get a direct answer without interactive exploration.")

        zero_query_input = gr.Textbox(
            label="Query",
            placeholder="Enter your question here...",
            lines=3
        )

        zero_shot_btn = gr.Button("Get Direct Answer", variant="primary")

        # State for maintaining context for zero-shot
        zero_state = gr.State(None)

        # Chat display
        zero_chatbot = gr.Chatbot(
            label="Direct Answer",
            height=600,
            type="messages"
        )

        # Modified zero-shot function to return state
        def zero_shot_with_state(query):
            """Wrapper for zero_shot to also return state"""
            from utils import generate_user_id
            chat_history = start_zero_shot(query)
            # Create state with user_id and method for DV processing
            state = {
                "user_id": generate_user_id(),
                "method": "zero_shot",
                "has_answer": True  # Flag to indicate answer is ready
            }
            return state, chat_history

        # Event handler for zero-shot
        zero_shot_btn.click(
            fn=zero_shot_with_state,
            inputs=zero_query_input,
            outputs=[zero_state, zero_chatbot],
            show_progress="minimal"
        )

        # Dependent Variables Section (initially hidden)
        with gr.Accordion("Evaluation Questions", open=False, visible=False) as zero_dv_accordion:
            zero_dv_inputs = []
            for i, question in enumerate(DV_QUESTIONS):
                if "1-5" in question:
                    dv_input = gr.Slider(
                        minimum=1,
                        maximum=5,
                        step=1,
                        value=3,
                        label=question
                    )
                elif "Yes/No" in question:
                    dv_input = gr.Radio(
                        choices=["Yes", "No"],
                        label=question
                    )
                else:
                    dv_input = gr.Textbox(
                        label=question
                    )
                zero_dv_inputs.append(dv_input)

            zero_submit_dv_btn = gr.Button("Submit Evaluation", variant="primary")
            zero_dv_result = gr.Markdown("")

            # Function to show DV accordion after answer is generated
            def show_zero_dv_section(state):
                if state and state.get("has_answer", False):
                    return gr.update(visible=True, open=True)
                return gr.update(visible=False, open=False)

            # Update visibility of DV section when state changes
            zero_state.change(
                fn=show_zero_dv_section,
                inputs=[zero_state],
                outputs=[zero_dv_accordion]
            )

            # Handle DV submission
            zero_submit_dv_btn.click(
                fn=process_dv_responses,
                inputs=[zero_state] + zero_dv_inputs,
                outputs=zero_dv_result
            )

    return zero_query_input, zero_shot_btn, zero_chatbot


def create_gradio_app():
    """Create the main Gradio application with tabs."""
    with gr.Blocks() as app:
        with gr.Tabs() as tabs:
            # IBFS Tab
            with gr.Tab("IBFS"):
                create_ibfs_interface()

            # Zero-Shot Tab
            with gr.Tab("Zero-Shot"):
                create_zero_shot_interface()

    return app


# Create and launch the app
if __name__ == "__main__":
    ibfs_app = create_gradio_app()
    ibfs_app.launch(share=True)