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

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

# --- 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-4', 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, model="gpt-4")
    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?"


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. Your feedback shapes the conversation, making it an educational journey 
        tailored just for you. So why settle for small talk? Dive into Soratiq today and elevate your discourse!
    """
    st.caption(description)
    st.divider()

    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)
    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: return sources used by the agent - explicitly publish them in the dialog