app.py

import streamlit as st
from socratic import *
from tool import get_tools
from agent import Agent
from common import get_llm
import sys
import time
import random

logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)

# --- APPLICATION ---

PAGE_TITLE: str = "Socratiq"
PAGE_ICON: str = "🗨️"
N_ROUND: int = 50

st.set_page_config(page_title=PAGE_TITLE, page_icon=PAGE_ICON)

open_api_key = st.secrets["OPENAI_API_KEY"]
st.session_state.key = open_api_key


def init_session() -> None:
    tools = get_tools()
    st.session_state.agent = Agent(get_llm(model_name='gpt-4o', model_temperature=0, api_key=st.session_state.key),
                                   tools)
    st.session_state.socrates = SocraticGPT(role=SOCRATES, tools=st.session_state.agent._tools,
                                            key=st.session_state.key, n_round=N_ROUND)
    st.session_state.theaetetus = SocraticGPT(role=THEAETETUS, tools=st.session_state.agent._tools,
                                              key=st.session_state.key, n_round=N_ROUND)
    st.session_state.plato = SocraticGPT(role=PLATO, tools=st.session_state.agent._tools, key=st.session_state.key,
                                         n_round=N_ROUND)
    st.session_state.dialog_lead = None
    st.session_state.dialog_follower = None
    st.session_state.messages = []
    st.session_state.question = None
    st.session_state.user_input = None
    st.session_state.in_progress = False
    st.session_state.user_question = None


def get_random_question():
    return "What is the size of the Moon?" if random.randint(1,10) % 3 == 0 else "What is love?"


def show_intro_screen():
    st.header("Socratiq")
    st.image("./assets/intro-trio.jpg")
    description = """\
        You ask, 'What's the meaning of life?' and our trio of digital philosophers fetch real-time 
        wisdom faster than you can say 'Immanuel Kant.' Whether you’re curious about science, or even the nuances of 
        modern art, Soratiq has you covered.
    """
    st.caption(description)

    if st.session_state.question is None:
        question = st.text_input(label='Paste your question. E.g. "What is the size of the moon?"',
                                label_visibility='collapsed',
                                placeholder='Paste your question. E.g. "What is the size of the moon?"')
        col1, _, _, _, col2 = st.columns(5)
        if col1.button(label="Ask Away!", help="Push the button and dive into the dialog around your question..."):
            if not len(question) > 0:
                st.warning(
                    "Whoops! That question seems to be doing the vanishing act. Could you give it another shot? Magic words: 'Valid question, Please!' 🪄")
        if len(question) > 0:
            set_user_question(question)

        if col2.button(label="QuestionRoll", help="The button generates a random question for the user to ponder or discuss. This should be a fun and engaging experience, sparking curiosity."):
            question = get_random_question()
            set_user_question(question)

        st.divider()
        with st.expander("Who are they? What's going on?"):
            st.markdown("""\
                Imagine you've just stepped into an intellectual arena that's as engaging as a top-tier debate club and as lively as your favorite coffeehouse. Like the brainstorming session you've always dreamed of.
                - First up, **you're the Maestro**. With a simple thumbs-up or down, you steer this conversation, helping to sculpt the dialogue into something meaningful. Your feedback is the compass.

                - Next, we have **Socrates and Theaetetus**, your co-pilots in this philosophical flight. They serve up questions designed to delve deep into the core of the matter.

                - **Plato's on board too**, adding a layer of logical critique to ensure our dialogue doesn't go off the rails.

                - Finally, meet the **Agent**, our whiz at rapid information retrieval. This isn't just Googling; think of it as real-time fact-checking with a dash of AI brilliance, you'd think it had a Ph.D. in Searchology.

                Try `QuestionRoll` and get a random question to watch how it goes. [More details](https://princeton-nlp.github.io/SocraticAI/)
            """)
    else:
        if st.session_state.question is not None:
            st.subheader(f"*{st.session_state.question}*")
        st.divider()



def set_user_question(question):
    st.session_state.question = question
    st.session_state.socrates.set_question(question)
    st.session_state.theaetetus.set_question(question)
    st.session_state.plato.set_question(question)
    st.experimental_rerun()


if 'question' not in st.session_state:
    init_session()

def get_avatar(role):
    if role == SOCRATES:
        return "./assets/socrates-avatar.jpg"
    elif role == THEAETETUS:
        return "./assets/theatetus-avatar.png"
    elif role == PLATO:
        return "./assets/plato-avatar.jpg"
    elif role == 'agent':
        return "ai"
    elif role == 'system':
        return "./assets/agent-avatar.jpg"
    elif role == 'user':
        return "user"

def get_role(role):
    if role == 'agent':
        return 'ai'
    elif role == 'system':
        return 'ai'
    else:
        return 'user'

def show_chat() -> None:
    if st.session_state.messages:
        for message in st.session_state.messages:
            role = message['role']
            content = message['content']
            with st.chat_message(get_role(role), avatar=get_avatar(role)):
                if role == 'system':
                    st.markdown("*" + content + "*")
                else:
                    st.markdown(content)


def add_message(role, content):
    st.session_state.messages.append({"role": role, "content": content})
    with st.chat_message(get_role(role), avatar=get_avatar(role)):
        st.markdown(content)


def chat_input():
    if st.session_state.user_question is not None:
        prompt = st.chat_input("Share your wisdom...")
        if prompt:
            st.session_state.user_input = prompt
    elif st.session_state.question is not None:
        st.chat_input("...", disabled=True)

def main() -> None:
    show_intro_screen()
    chat_input()
    show_chat()

    if st.session_state.question is not None and st.session_state.user_question is None:
        if not st.session_state.in_progress:
            st.session_state.in_progress = True
            st.session_state.dialog_lead, st.session_state.dialog_follower = st.session_state.socrates, st.session_state.theaetetus
            # add_message(st.session_state.dialog_lead.role,
            #             f"""Hi {st.session_state.dialog_follower.role}, let's solve this problem together. Please feel free to correct me if I make any logical or mathematical mistakes.\n""")
        else:
            with st.spinner(f"{st.session_state.dialog_follower.role} is thinking..."):
                rep = st.session_state.dialog_follower.get_response()
            add_message(st.session_state.dialog_follower.role, f"{st.session_state.dialog_follower.role}: " + rep)
            st.session_state.dialog_lead.update_history(rep)
            st.session_state.plato.update_history(f"{st.session_state.dialog_follower.role}: " + rep)

            # next round the opponent answers
            st.session_state.dialog_lead, st.session_state.dialog_follower = st.session_state.dialog_follower, st.session_state.dialog_lead

            answer = SocraticGPT.get_answer(rep)
            user_question = SocraticGPT.get_user_question(rep)
            agent_question = SocraticGPT.get_agent_question(rep)

            if st.session_state.dialog_lead.role == st.session_state.theaetetus.role:
                if user_question is None and agent_question is None and answer is None:
                    with st.spinner(f"thinking critically..."):
                        pr = st.session_state.plato.get_proofread()
                    if pr:
                        add_message(st.session_state.plato.role, f"{st.session_state.plato.role}: " + pr)
                        st.session_state.socrates.add_proofread(pr)
                        st.session_state.theaetetus.add_proofread(pr)
                        user_question = SocraticGPT.get_user_question(pr) # Plato can suggest to use agent or get user feedback
                        agent_question = SocraticGPT.get_agent_question(pr)

            if agent_question:
                with st.status(f"Consulting the agent: ''' {agent_question} '''"):
                    agent_msg = st.session_state.agent.run(agent_question) # TODO: agent status callback or status polling in a loop
                st.session_state.socrates.add_agent_feedback(agent_question, agent_msg)
                st.session_state.theaetetus.add_agent_feedback(agent_question, agent_msg)
                st.session_state.plato.add_agent_feedback(agent_question, agent_msg)
                add_message('agent', f"Agent: {agent_msg}")

            if user_question:
                st.session_state.user_question = user_question
                add_message('system', f'User feedback is required. The question: **"{" ".join(user_question)}"**')
                st.experimental_rerun()

            if answer:
                st.session_state.user_question = f"Is that correct answer? - {answer}"
                add_message('system', f"""User, are you agree with the answer? - **{" ".join(answer)}**""")
                st.experimental_rerun()

    if st.session_state.user_input is not None:
        user_input = st.session_state.user_input
        if st.session_state.user_question is not None:
            user_question = st.session_state.user_question
            st.session_state.socrates.add_user_feedback(user_question, user_input)
            st.session_state.theaetetus.add_user_feedback(user_question, user_input)
            st.session_state.plato.add_user_feedback(user_question, user_input)
            st.session_state.user_question = None
            add_message("user", f"{user_input}")

        st.session_state.user_input = None

    if st.session_state.question is not None and st.session_state.user_question is None:
        time.sleep(1)
        st.experimental_rerun()


if __name__ == "__main__":
    main()


# TODO: publish/access dialog debug logs, so the user can dig into the details
# TODO: possible answers to the question - like 'double check your answer' or 'make the answer sound like a pirate' etc
# TODO: get rid of autogenerated 'Footnote' header in agent's references
# TODO: handle rate limit reject, e.g. Limit: 10000 / min. Please try again in 6ms.  - it's easy to wait for 6 ms. explicitly prompt it or ask Plato
# TODO: streaming response https://github.com/langchain-ai/streamlit-agent/blob/main/streamlit_agent/basic_streaming.py