app.py

import json
from typing import List, Tuple
import os

import gradio as gr
from dotenv import load_dotenv
from slugify import slugify

from config import STUDY_FILES
from rag.rag_pipeline import RAGPipeline
from utils.helpers import generate_follow_up_questions, append_to_study_files
from utils.prompts import (
    highlight_prompt,
    evidence_based_prompt,
    sample_questions,
)
import openai

from config import STUDY_FILES, OPENAI_API_KEY
from utils.zotero_manager import ZoteroManager

load_dotenv()

openai.api_key = OPENAI_API_KEY

# Cache for RAG pipelines
rag_cache = {}

def process_zotero_library_items(zotero_library_id: str, zotero_api_access_key: str) -> str:
    if not zotero_library_id or not zotero_api_access_key:
        return "Please enter your zotero library Id and API Access Key"

    zotero_library_id = zotero_library_id
    zotero_library_type = "user"  # or "group"
    zotero_api_access_key = zotero_api_access_key

    message = ""

    try:
        zotero_manager = ZoteroManager(
            zotero_library_id, zotero_library_type, zotero_api_access_key
        )

        zotero_collections = zotero_manager.get_collections()
        zotero_collection_lists = zotero_manager.list_zotero_collections(zotero_collections)
        filtered_zotero_collection_lists = (
            zotero_manager.filter_and_return_collections_with_items(zotero_collection_lists)
        )

        for collection in filtered_zotero_collection_lists:
            collection_name = collection.get("name")
            if collection_name not in STUDY_FILES:
                collection_key = collection.get("key")
                collection_items = zotero_manager.get_collection_items(collection_key)
                zotero_collection_items = (
                    zotero_manager.get_collection_zotero_items_by_key(collection_key)
                )
                #### Export zotero collection items to json ####
                zotero_items_json = zotero_manager.zotero_items_to_json(zotero_collection_items)
                export_file = f"{slugify(collection_name)}_zotero_items.json"
                zotero_manager.write_zotero_items_to_json_file(
                    zotero_items_json, f"data/{export_file}"
                )
                append_to_study_files("study_files.json", collection_name, f"data/{export_file}")
        message = "Successfully processed items in your zotero library"
    except Exception as e:
        message = f"Error process your zotero library: {str(e)}"
    
    return message


def get_rag_pipeline(study_name: str) -> RAGPipeline:
    """Get or create a RAGPipeline instance for the given study."""
    if study_name not in rag_cache:
        study_file = STUDY_FILES.get(study_name)
        if not study_file:
            raise ValueError(f"Invalid study name: {study_name}")
        rag_cache[study_name] = RAGPipeline(study_file)
    return rag_cache[study_name]


def chat_function(
    message: str, study_name: str, prompt_type: str
) -> str:
    """Process a chat message and generate a response using the RAG pipeline."""

    if not message.strip():
        return "Please enter a valid query."

    rag = get_rag_pipeline(study_name)
    prompt = {
        "Highlight": highlight_prompt,
        "Evidence-based": evidence_based_prompt,
    }.get(prompt_type)

    response = rag.query(message, prompt_template=prompt)
    return response.response


def get_study_info(study_name: str) -> str:
    """Retrieve information about the specified study."""

    study_file = STUDY_FILES.get(study_name)
    if not study_file:
        return "Invalid study name"

    with open(study_file, "r") as f:
        data = json.load(f)
    return f"### Number of documents: {len(data)}"


def update_interface(study_name: str) -> Tuple[str, gr.update, gr.update, gr.update]:
    """Update the interface based on the selected study."""

    study_info = get_study_info(study_name)
    questions = sample_questions.get(study_name, [])[:3]
    if not questions:
        questions = sample_questions.get("General", [])[:3]
    visible_questions = [gr.update(visible=True, value=q) for q in questions]
    hidden_questions = [gr.update(visible=False) for _ in range(3 - len(questions))]
    return (study_info, *visible_questions, *hidden_questions)


def set_question(question: str) -> str:
    return question.lstrip("✨ ")

def process_multi_input(text, study_name, prompt_type):
    # Split input based on commas and strip any extra spaces
    variable_list = [word.strip().upper() for word in text.split(',')]
    user_message =f"Extract and present in a tabular format the following variables for each {study_name} study: {', '.join(variable_list)}"
    response = chat_function(user_message, study_name, prompt_type)
    return response


def create_gr_interface() -> gr.Blocks:
    """
    Create and configure the Gradio interface for the RAG platform.

    This function sets up the entire user interface, including:
    - Chat interface with message input and display
    - Study selection dropdown
    - Sample and follow-up question buttons
    - Prompt type selection
    - Event handlers for user interactions

    Returns:
        gr.Blocks: The configured Gradio interface ready for launching.
    """

    with gr.Blocks() as demo:
        gr.Markdown("# ACRES RAG Platform")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Zotero Credentials")
                zotero_library_id = gr.Textbox(label="Zotero Library ID", type="password", placeholder="Enter Your Zotero Library ID here...")
                zotero_api_access_key = gr.Textbox(label="Zotero API Access Key", type="password", placeholder="Enter Your Zotero API Access Key...")
                process_zotero_btn = gr.Button("Process your Zotero Library")
                zotero_output = gr.Markdown(label="Zotero")

                gr.Markdown("### Study Information")
                study_dropdown = gr.Dropdown(
                    choices=list(STUDY_FILES.keys()),
                    label="Select Study",
                    value=list(STUDY_FILES.keys())[0],
                )
                study_info = gr.Markdown(label="Study Details")

                gr.Markdown("### Settings")
                prompt_type = gr.Radio(
                    ["Default", "Highlight", "Evidence-based"],
                    label="Prompt Type",
                    value="Default",
                )
                # clear = gr.Button("Clear Chat")
            
            with gr.Column(scale=3):
                gr.Markdown("### Study Variables")
                with gr.Row():
                    study_variables = gr.Textbox(
                        show_label=False,
                        placeholder="Type your variables separated by commas e.g (Study ID, Study Title, Authors etc)",
                        scale=4,
                        lines=1,
                        autofocus=True,
                    )
                    submit_btn = gr.Button("Submit", scale=1)
                answer_output = gr.Markdown(label="Answer")

        def user(
            user_message: str, history: List[List[str]]
        ) -> Tuple[str, List[List[str]]]:
            return "", (
                history + [[user_message, None]] if user_message.strip() else history
            )

        def bot(
            history: List[List[str]], study_name: str, prompt_type: str
        ) -> List[List[str]]:
            """
            Generate bot response and update the interface.

            This function:
            1. Processes the latest user message
            2. Generates a response using the RAG pipeline
            3. Updates the chat history
            4. Generates follow-up questions
            5. Prepares interface updates for follow-up buttons

            Args:
                history (List[List[str]]): The current chat history.
                study_name (str): The name of the current study.
                prompt_type (str): The type of prompt being used.

            Returns:
                Tuple[List[List[str]], gr.update, gr.update, gr.update]:
                Updated chat history and interface components for follow-up questions.
            """
            if not history:
                return history, [], [], []

            user_message = history[-1][0]
            bot_message = chat_function(user_message, history, study_name, prompt_type)
            history[-1][1] = bot_message

            return history

        # msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        #     bot,
        #     [chatbot, study_dropdown, prompt_type],
        #     [chatbot, *follow_up_btns],
        # )
        # send_btn.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        #     bot,
        #     [chatbot, study_dropdown, prompt_type],
        #     [chatbot, *follow_up_btns],
        # )
        # for btn in follow_up_btns + sample_btns:
        #     btn.click(set_question, inputs=[btn], outputs=[msg])

        # clear.click(lambda: None, None, chatbot, queue=False)

        study_dropdown.change(
            fn=get_study_info,
            inputs=study_dropdown,
            outputs=[study_info],
        )

        process_zotero_btn.click(process_zotero_library_items, inputs=[zotero_library_id, zotero_api_access_key], outputs=[zotero_output], queue=False)
        submit_btn.click(process_multi_input, inputs=[study_variables, study_dropdown, prompt_type], outputs=[answer_output], queue=False)

    return demo


demo = create_gr_interface()

if __name__ == "__main__":
    # demo = create_gr_interface()
    demo.launch(share=True, debug=True)