app.py

import gradio as gr
from ontochat.functions import set_openai_api_key, user_story_generator, cq_generator, load_example_user_story, clustering_generator, ontology_testing, load_example

user_story_template = """**Persona:**\n\n- Name: -\n- Age: -\n- Occupation: -\n- Skills: -\n- Interests: -\n\n**Goal:**\n\n- Description: -\n- Keywords: -\n\n**Scenario:**\n\n- Before: -\n- During: -\n- After: -\n\n**Example Data:**\n\n- Category: -\n- Data: -\n\n**Resources:**\n\n- Resource Name: -\n- Link: -"""

with gr.Blocks() as set_api_key:
    gr.Markdown(
        """
        # Welcome to OntoChat! 👋

        **Hello! I'm OntoChat, your conversational ontology engineering assistant.** 🎉

        I'm here to help you streamline the complex process of building and refining ontologies. Whether you're collecting requirements, generating user stories, extracting competency questions, or testing early versions of your ontology, I've got you covered! You can use me to:

        - **Collaboratively create user stories** tailored to your domain.
        - **Automatically extract and refine competency questions** from your stories.
        - **Cluster and analyze competency questions** to identify patterns and gaps in your requirements.
        - **Test and verify your ontology's design** without needing to write queries manually.

        Let's work together to simplify your ontology engineering process!
        Visit [OntoChat on GitHub](https://github.com/King-s-Knowledge-Graph-Lab/OntoChat) for more information.
        """
    )

    with gr.Group():
        api_key = gr.Textbox(
            label="OpenAI API Key",
            info="Please input your OpenAI API Key if you don't have it set up on your own machine. Please note that "
                 "the key will only be used for this demo and will not be uploaded or used anywhere else."
        )
        api_key_btn = gr.Button(value="Set API Key")
        api_key_btn.click(fn=set_openai_api_key, inputs=api_key, outputs=api_key)

with gr.Blocks() as user_story_interface:
    with gr.Row():
        with gr.Column(scale=1):
            user_story_chatbot = gr.Chatbot(
                value=[
                    {"role": "assistant", "content": ( 
                        "Hello! I'm OntoChat, your trusted assistant for creating an ontology user story. Here's how I can assist you: \n\n1. I'll guide you through a sequence of questions, and you can use the provided prompt templates to craft your answers. \n\n2. You're welcome to share comments on my responses, and I'll refine them based on your feedback to ensure continuous improvement. 😊\n\nLet's get started! Could you tell me a little about the purpose of this ontology?"
                    )}
                ],
                height="472px",
                type="messages"
            )
            user_story_input = gr.Textbox(
                label="Message OntoChat",
                placeholder="Please type your message here and press Shift + Enter to interact with the chatbot:",
                max_lines=4,
                lines = 4
            )
            elicitation_questions_dataset = gr.Dataset(
                components=[user_story_input],
                label="Prompt Templates",
                type="index",
                samples=[
                    ["Create Persona"],
                    ["Create User Goal"],
                    ["Create Actions"],
                    ["Create Keywords"],
                    ["Create Current Methods"],
                    ["Create Challenges"],
                    ["Create New Methods"],
                    ["Create Outcomes"]
                ],
                samples_per_page = 10
            )

    user_story_input.submit(
        fn=user_story_generator,
        inputs=[user_story_input, user_story_chatbot],
        outputs=[user_story_chatbot, user_story_input]
    )
    elicitation_questions_dataset.click(
        fn=load_example, 
        inputs=[elicitation_questions_dataset], 
        outputs=[user_story_input]
    ) 

with gr.Blocks() as cq_interface:
    with gr.Row():
        with gr.Column():
            cq_chatbot = gr.Chatbot(
                value=[
                    {
                        "role": "assistant",
                        "content": (
                            "I am OntoChat, your conversational ontology engineering assistant. Here is the second step of "
                            "the system. Please give me your user story and tell me how many competency questions you want "
                            "me to generate from the user story."
                        )
                    }
                ],
                type="messages" 
            )
            cq_input = gr.Textbox(
                label="Chatbot input",
                placeholder="Please type your message here and press Enter to interact with the chatbot:"
            )
            gr.Markdown(
                """
                ### User story examples
                Click the button below to use an example user story from 
                [Linka](https://github.com/polifonia-project/stories/tree/main/Linka_Computer_Scientist) in Polifonia.
                """
            )
            example_btn = gr.Button(value="Use example user story")
            example_btn.click(
                fn=load_example_user_story,
                inputs=[],
                outputs=[cq_input]
            )
        cq_output = gr.TextArea(
            label="Competency questions",
            interactive=True
        )
    cq_input.submit(
        fn=cq_generator,
        inputs=[
            cq_input, cq_chatbot
        ],
        outputs=[
            cq_output, cq_chatbot, cq_input
        ]
    )

clustering_interface = gr.Interface(
    fn=clustering_generator,
    inputs=[
        gr.TextArea(
            label="Competency questions",
            info="Please copy the previously generated competency questions and paste it here. You can also modify "
                 "the questions before submitting them."
        ),
        gr.Dropdown(
            value="LLM clustering",
            choices=["LLM clustering", "Agglomerative clustering"],
            label="Clustering method",
            info="Please select the clustering method."
        ),
        gr.Textbox(
            label="Number of clusters (optional for LLM clustering)",
            info="Please input the number of clusters you want to generate. And please do not input a number that "
                 "exceeds the total number of competency questions."
        )
    ],
    outputs=[
        gr.Image(label="Visualization"),
        gr.Code(
            language='json',
            label="Competency Question clusters"
        )
    ],
    title="OntoChat",
    description="This is the third step of OntoChat. Please copy the generated competency questions from the previous "
                "step and run the clustering algorithm to group the competency questions based on their topics. From "
                "our experience, LLM clustering has the best performance.",
    flagging_mode="never"
)

with gr.Blocks() as testing_interface:
    gr.Markdown(
        """
        # OntoChat
        This is the final part of OntoChat which performs ontology testing based on the input ontology file and CQs. 
        """
    )

    with gr.Group():
        api_key = gr.Textbox(
            label="OpenAI API Key",
            placeholder="If you have set the key in other tabs, you don't have to set it again.",
            info="Please input your OpenAI API Key if you don't have it set up on your own machine. Please note that "
                 "the key will only be used for this demo and will not be uploaded or used anywhere else."
        )
        api_key_btn = gr.Button(value="Set API Key")
        api_key_btn.click(fn=set_openai_api_key, inputs=api_key, outputs=api_key)

    ontology_file = gr.File(label="Ontology file")
    ontology_desc = gr.Textbox(
        label="Ontology description",
        placeholder="Please provide a description of the ontology uploaded to provide basic information and "
                    "additional context."
    )
    cq_testing_input = gr.Textbox(
        label="Competency questions",
        placeholder="Please provide the competency questions that you want to test with."
    )
    testing_btn = gr.Button(value="Test")
    testing_output = gr.TextArea(label="Ontology testing output")
    testing_btn.click(
        fn=ontology_testing,
        inputs=[
            ontology_file, ontology_desc, cq_testing_input
        ],
        outputs=[
            testing_output
        ]
    )

demo = gr.TabbedInterface(
    [set_api_key, user_story_interface, cq_interface, clustering_interface, testing_interface],
    ["Set API Key", "User Story Generation", "Competency Question Extraction", "Competency Question Analysis", "Ontology Testing"]
)

if __name__ == "__main__":
    demo.launch(share=True)