better_responses.py

# ref: https://github.com/twy80/LangChain_llm_Agent/tree/main
import streamlit as st
import os, base64, re, requests, datetime, time, json
import matplotlib.pyplot as plt
from io import BytesIO
from functools import partial
from tempfile import NamedTemporaryFile
from audio_recorder_streamlit import audio_recorder
from PIL import Image, UnidentifiedImageError
from openai import OpenAI
from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from langchain_anthropic import ChatAnthropic
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_google_community import GoogleSearchAPIWrapper
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.schema import HumanMessage, AIMessage
from langchain_community.utilities import BingSearchAPIWrapper
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import Docx2txtLoader
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.tools import Tool, tool
from langchain.tools.retriever import create_retriever_tool
# from langchain.agents import create_openai_tools_agent
from langchain.agents import create_tool_calling_agent
from langchain.agents import create_react_agent
from langchain.agents import AgentExecutor
from langchain_community.agent_toolkits.load_tools import load_tools
# from langchain_experimental.tools import PythonREPLTool
from langchain_experimental.utilities import PythonREPL
from langchain.callbacks.base import BaseCallbackHandler
from pydantic import BaseModel, Field
# The following are for type annotations
from typing import Union, List, Literal, Optional, Dict, Any, Annotated
from matplotlib.figure import Figure
from streamlit.runtime.uploaded_file_manager import UploadedFile
from openai._legacy_response import HttpxBinaryResponseContent
from tempfile import NamedTemporaryFile, TemporaryDirectory

# Load API keys from Hugging Face secrets
try:
    os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
    os.environ["BING_SUBSCRIPTION_KEY"] = st.secrets.get("BING_SUBSCRIPTION_KEY", "")
    os.environ["GOOGLE_API_KEY"] = st.secrets.get("GOOGLE_API_KEY", "")
    os.environ["GOOGLE_CSE_ID"] = st.secrets.get("GOOGLE_CSE_ID", "")
except KeyError as e:
    st.error(f"Missing required secret: {e}. Please set it in Hugging Face Space secrets.")
    st.stop()

def initialize_session_state_variables() -> None:
    """
    Initialize all the session state variables.
    """
    default_values = {
        "ready": False,
        "openai": None,
        "history": [],
        "model_type": "GPT Models from OpenAI",
        "agent_type": 2 * ["Tool Calling"],
        "ai_role": 2 * ["You are a helpful AI assistant."],
        "prompt_exists": False,
        "temperature": [0.7, 0.7],
        "audio_bytes": None,
        "mic_used": False,
        "audio_response": None,
        "image_url": None,
        "image_description": None,
        "uploader_key": 0,
        "tool_names": [[], []],
        "bing_subscription_validity": False,
        "google_cse_id_validity": False,
        "vector_store_message": None,
        "retriever_tool": None,
        "show_uploader": False
    }

    for key, value in default_values.items():
        if key not in st.session_state:
            st.session_state[key] = value



class StreamHandler(BaseCallbackHandler):
    def __init__(self, container, initial_text=""):
        self.container = container
        self.text = initial_text

    def on_llm_new_token(self, token: Any, **kwargs) -> None:
        new_text = self._extract_text(token)
        if new_text:
            self.text += new_text
            self.container.markdown(self.text)

    def _extract_text(self, token: Any) -> str:
        if isinstance(token, str):
            return token
        elif isinstance(token, list):
            return ''.join(self._extract_text(t) for t in token)
        elif isinstance(token, dict):
            return token.get('text', '')
        else:
            return str(token)


def check_api_keys() -> None:
    # Unset this flag to check the validity of the OpenAI API key
    st.session_state.ready = False


def message_history_to_string(extra_space: bool=True) -> str:
    """
    Return a string of the chat history contained in
    st.session_state.history.
    """

    history_list = []
    for msg in st.session_state.history:
        if isinstance(msg, HumanMessage):
            history_list.append(f"Human: {msg.content}")
        else:
            history_list.append(f"AI: {msg.content}")
    new_lines = "\n\n" if extra_space else "\n"

    return new_lines.join(history_list)


def get_chat_model(
    model: str,
    temperature: float,
    callbacks: List[BaseCallbackHandler]
) -> Union[ChatOpenAI, ChatAnthropic, ChatGoogleGenerativeAI, None]:

    """
    Get the appropriate chat model based on the given model name.
    """

    model_map = {
        "gpt-": ChatOpenAI,
    }
    for prefix, ModelClass in model_map.items():
        if model.startswith(prefix):
            return ModelClass(
                model=model,
                temperature=temperature,
                streaming=True,
                callbacks=callbacks
            )
    return None


def process_with_images(
    llm: Union[ChatOpenAI, ChatAnthropic, ChatGoogleGenerativeAI],
    message_content: str,
    image_urls: List[str]
) -> str:

    """
    Process the given history query with associated images using a language model.
    """

    content_with_images = (
        [{"type": "text", "text": message_content}] +
        [{"type": "image_url", "image_url": {"url": url}} for url in image_urls]
    )
    message_with_images = [HumanMessage(content=content_with_images)]

    return llm.invoke(message_with_images).content


def process_with_tools(
    llm: Union[ChatOpenAI, ChatAnthropic, ChatGoogleGenerativeAI],
    tools: List[Tool],
    agent_type: str,
    agent_prompt: str,
    history_query: dict
) -> str:

    """
    Create an AI agent based on the specified agent type and tools,
    then use this agent to process the given history query.
    """

    if agent_type == "Tool Calling":
        agent = create_tool_calling_agent(llm, tools, agent_prompt)
    else:
        agent = create_react_agent(llm, tools, agent_prompt)

    agent_executor = AgentExecutor(
        agent=agent, tools=tools, max_iterations=5, verbose=False,
        handle_parsing_errors=True,
    )

    return agent_executor.invoke(history_query)["output"]


def run_agent(
    query: str,
    model: str,
    tools: List[Tool],
    image_urls: List[str],
    temperature: float=0.7,
    agent_type: Literal["Tool Calling", "ReAct"]="Tool Calling",
) -> Union[str, None]:
    """
    Generate text based on user queries.
    Args:
        query: User's query
        model: LLM like "gpt-4o"
        tools: list of tools such as Search and Retrieval
        image_urls: List of URLs for images
        temperature: Value between 0 and 1. Defaults to 0.7
        agent_type: 'Tool Calling' or 'ReAct'
    Return:
        generated text
    """

    try:
        # Ensure retriever tool is included when "Retrieval" is selected
        if "Retrieval" in st.session_state.tool_names[0]:
            if st.session_state.retriever_tool:
                retriever_tool_name = "retriever"  # Ensure naming consistency
                if retriever_tool_name not in [tool.name for tool in tools]:
                    tools.append(st.session_state.retriever_tool)
                    st.write(f"✅ **{retriever_tool_name} tool has been added successfully.**")
            else:
                st.error("❌ Retriever tool is not initialized. Please create a vector store first.")
                return None  # Exit early to avoid broken tool usage

        # Debugging: Print final tools list
        st.write("**Final Tools Being Used:**", [tool.name for tool in tools])

        if "retriever" in [tool.name for tool in tools]:
            st.success("✅ Retriever tool is confirmed and ready for use.")
        elif "Retrieval" in st.session_state.tool_names[0]:
            st.warning("⚠️ 'Retrieval' was selected but the retriever tool is missing!")

        # Initialize the LLM model
        llm = get_chat_model(model, temperature, [StreamHandler(st.empty())])
        if llm is None:
            st.error(f"❌ Unsupported model: {model}", icon="🚨")
            return None
        
        # Prepare chat history
        if agent_type == "Tool Calling":
            chat_history = st.session_state.history
        else:
            chat_history = message_history_to_string()

        history_query = {"chat_history": chat_history, "input": query}

        # Generate message content
        message_with_no_image = st.session_state.chat_prompt.invoke(history_query)
        message_content = message_with_no_image.messages[0].content

        if image_urls:
            # Handle images if provided
            generated_text = process_with_images(llm, message_content, image_urls)
            human_message = HumanMessage(
                content=query, additional_kwargs={"image_urls": image_urls}
            )
        elif tools:
            # Use tools for query execution
            generated_text = process_with_tools(
                llm, tools, agent_type, st.session_state.agent_prompt, history_query
            )
            human_message = HumanMessage(content=query)
        else:
            # Fall back to basic query execution without tools
            generated_text = llm.invoke(message_with_no_image).content
            human_message = HumanMessage(content=query)

        # Convert response into plain text
        if isinstance(generated_text, list):
            generated_text = generated_text[0]["text"]

        # Update conversation history
        st.session_state.history.append(human_message)
        st.session_state.history.append(AIMessage(content=generated_text))

        return generated_text

    except Exception as e:
        st.error(f"An error occurred: {e}", icon="🚨")
        return None


def openai_create_image(
    description: str, model: str="dall-e-3", size: str="1024x1024"
) -> Optional[str]:

    """
    Generate image based on user description.
    Args:
        description: User description
        model: Default set to "dall-e-3"
        size: Pixel size of the generated image
    Return:
        URL of the generated image
    """

    try:
        with st.spinner("AI is generating..."):
            response = st.session_state.openai.images.generate(
                model=model,
                prompt=description,
                size=size,
                quality="standard",
                n=1,
            )
        image_url = response.data[0].url
    except Exception as e:
        image_url = None
        st.error(f"An error occurred: {e}", icon="🚨")

    return image_url


def get_vector_store(uploaded_files: List[UploadedFile]) -> Optional[FAISS]:
    """
    Take a list of UploadedFile objects as input, and return a FAISS vector store.
    """
    if not uploaded_files:
        return None

    documents = []
    loader_map = {
        ".pdf": PyPDFLoader,
        ".txt": TextLoader,
        ".docx": Docx2txtLoader
    }

    try:
        # Use a temporary directory instead of a fixed 'files/' directory
        with TemporaryDirectory() as temp_dir:
            for uploaded_file in uploaded_files:
                # Create a temporary file in the system's temporary directory
                with NamedTemporaryFile(dir=temp_dir, delete=False) as temp_file:
                    temp_file.write(uploaded_file.getbuffer())
                    filepath = temp_file.name
                
                file_ext = os.path.splitext(uploaded_file.name.lower())[1]
                loader_class = loader_map.get(file_ext)
                if not loader_class:
                    st.error(f"Unsupported file type: {file_ext}", icon="🚨")
                    return None

                # Load the document using the selected loader
                loader = loader_class(filepath)
                documents.extend(loader.load())

        with st.spinner("Vector store in preparation..."):
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=1000, chunk_overlap=200
            )
            doc = text_splitter.split_documents(documents)

            # Choose embeddings
            if st.session_state.model_type == "GPT Models from OpenAI":
                embeddings = OpenAIEmbeddings(model="text-embedding-3-large", dimensions=1536)
            else:
                embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")

            # Create FAISS vector database
            vector_store = FAISS.from_documents(doc, embeddings)

    except Exception as e:
        vector_store = None
        st.error(f"An error occurred: {e}", icon="🚨")

    return vector_store



def get_retriever() -> None:
    """
    Upload document(s), create a vector store, prepare a retriever tool,
    save the tool to the variable st.session_state.retriever_tool.
    """

    # Section Title
    st.write("")
    st.write("**Query Document(s)**")

    # File Upload Input
    uploaded_files = st.file_uploader(
        label="Upload an article",
        type=["txt", "pdf", "docx"],
        accept_multiple_files=True,
        label_visibility="collapsed",
        key="document_upload_" + str(st.session_state.uploader_key),
    )

    # Check if files are uploaded
    if uploaded_files:
        # Use a unique button key to avoid duplicate presses
        if st.button(label="Create the vector store", key=f"create_vector_{st.session_state.uploader_key}"):
            st.info("Creating the vector store and initializing the retriever tool...")

            # Attempt to create the vector store
            vector_store = get_vector_store(uploaded_files)

            if vector_store:
                uploaded_file_names = [file.name for file in uploaded_files]
                st.session_state.vector_store_message = (
                    f"Vector store for :blue[[{', '.join(uploaded_file_names)}]] is ready!"
                )

                # Initialize retriever and create tool
                retriever = vector_store.as_retriever()
                st.session_state.retriever_tool = create_retriever_tool(
                    retriever,
                    name="retriever",
                    description="Search uploaded documents for information when queried.",
                )

                # Add "Retrieval" to the tools list if not already present
                if "Retrieval" not in st.session_state.tool_names[0]:
                    st.session_state.tool_names[0].append("Retrieval")

                st.success("✅ Retriever tool has been successfully initialized and is ready to use.")

                # Debugging output
                st.write("**Current Tools:**", st.session_state.tool_names[0])
            else:
                st.error("❌ Failed to create vector store. Please check the uploaded files (supported formats: txt, pdf, docx).")
    else:
        st.info("Please upload document(s) to create the vector store.")




def display_text_with_equations(text: str):
    # Replace inline LaTeX equation delimiters \\( ... \\) with $
    modified_text = text.replace("\\(", "$").replace("\\)", "$")

    # Replace block LaTeX equation delimiters \\[ ... \\] with $$
    modified_text = modified_text.replace("\\[", "$$").replace("\\]", "$$")

    # Use st.markdown to display the formatted text with equations
    st.markdown(modified_text)


def read_audio(audio_bytes: bytes) -> Optional[str]:
    """
    Read audio bytes and return the corresponding text.
    """
    try:
        audio_data = BytesIO(audio_bytes)
        audio_data.name = "recorded_audio.wav"  # dummy name

        transcript = st.session_state.openai.audio.transcriptions.create(
            model="whisper-1", file=audio_data
        )
        text = transcript.text
    except Exception as e:
        text = None
        st.error(f"An error occurred: {e}", icon="🚨")

    return text


def input_from_mic() -> Optional[str]:
    """
    Convert audio input from mic to text and return it.
    If there is no audio input, None is returned.
    """

    time.sleep(0.5)
    audio_bytes = audio_recorder(
        pause_threshold=3.0, text="Speak", icon_size="2x",
        recording_color="#e87070", neutral_color="#6aa36f"        
    )

    if audio_bytes == st.session_state.audio_bytes or audio_bytes is None:
        return None
    else:
        st.session_state.audio_bytes = audio_bytes
        return read_audio(audio_bytes)


def perform_tts(text: str) -> Optional[HttpxBinaryResponseContent]:
    """
    Take text as input, perform text-to-speech (TTS),
    and return an audio_response.
    """

    try:
        with st.spinner("TTS in progress..."):
            audio_response = st.session_state.openai.audio.speech.create(
                model="tts-1",
                voice="fable",
                input=text,
            )
    except Exception as e:
        audio_response = None
        st.error(f"An error occurred: {e}", icon="🚨")

    return audio_response


def play_audio(audio_response: HttpxBinaryResponseContent) -> None:
    """
    Take an audio response (a bytes-like object)
    from TTS as input, and play the audio.
    """

    audio_data = audio_response.read()

    # Encode audio data to base64
    b64 = base64.b64encode(audio_data).decode("utf-8")

    # Create a markdown string to embed the audio player with the base64 source
    md = f"""
        <audio controls autoplay style="width: 100%;">
        <source src="data:audio/mp3;base64,{b64}" type="audio/mp3">
        Your browser does not support the audio element.
        </audio>
        """

    # Use Streamlit to render the audio player
    st.markdown(md, unsafe_allow_html=True)


def image_to_base64(image: Image) -> str:
    """
    Convert an image object from PIL to a base64-encoded image,
    and return the resulting encoded image as a string to be used
    in place of a URL.
    """

    # Convert the image to RGB mode if necessary
    if image.mode != "RGB":
        image = image.convert("RGB")

    # Save the image to a BytesIO object
    buffered_image = BytesIO()
    image.save(buffered_image, format="JPEG")

    # Convert BytesIO to bytes and encode to base64
    img_str = base64.b64encode(buffered_image.getvalue())

    # Convert bytes to string
    base64_image = img_str.decode("utf-8")

    return f"data:image/jpeg;base64,{base64_image}"


def shorten_image(image: Image, max_pixels: int=1024) -> Image:
    """
    Take an Image object as input, and shorten the image size
    if the image is greater than max_pixels x max_pixels.
    """

    if max(image.width, image.height) > max_pixels:
        if image.width > image.height:
            new_width, new_height = 1024, image.height * 1024 // image.width
        else:
            new_width, new_height = image.width * 1024 // image.height, 1024

        image = image.resize((new_width, new_height))

    return image


def upload_image_files_return_urls(
    type: List[str]=["jpg", "jpeg", "png", "bmp"]
) -> List[str]:

    """
    Upload image files, convert them to base64-encoded images, and
    return the list of the resulting encoded images to be used
    in place of URLs.
    """

    st.write("")
    st.write("**Query Image(s)**")
    source = st.radio(
        label="Image selection",
        options=("Uploaded", "From URL"),
        horizontal=True,
        label_visibility="collapsed",
    )
    image_urls = []

    if source == "Uploaded":
        uploaded_files = st.file_uploader(
            label="Upload images",
            type=type,
            accept_multiple_files=True,
            label_visibility="collapsed",
            key="image_upload_" + str(st.session_state.uploader_key),
        )
        if uploaded_files:
            try:
                for image_file in uploaded_files:
                    image = Image.open(image_file)
                    thumbnail = shorten_image(image, 300)
                    st.image(thumbnail)
                    image = shorten_image(image, 1024)
                    image_urls.append(image_to_base64(image))
            except UnidentifiedImageError as e:
                st.error(f"An error occurred: {e}", icon="🚨")
    else:
        image_url = st.text_input(
            label="URL of the image",
            label_visibility="collapsed",
            key="image_url_" + str(st.session_state.uploader_key),
        )
        if image_url:
            if is_url(image_url):
                st.image(image_url)
                image_urls = [image_url]
            else:
                st.error("Enter a proper URL", icon="🚨")

    return image_urls


def fig_to_base64(fig: Figure) -> str:
    """
    Convert a Figure object to a base64-encoded image, and return
    the resulting encoded image to be used in place of a URL.
    """

    with BytesIO() as buffer:
        fig.savefig(buffer, format="JPEG")
        buffer.seek(0)
        image = Image.open(buffer)

        return image_to_base64(image)


def is_url(text: str) -> bool:
    """
    Determine whether text is a URL or not.
    """

    regex = r"(http|https)://([\w_-]+(?:\.[\w_-]+)+)(:\S*)?"
    p = re.compile(regex)
    match = p.match(text)
    if match:
        return True
    else:
        return False


def reset_conversation() -> None:
    """
    Reset the session_state variables for resetting the conversation.
    """

    st.session_state.history = []
    st.session_state.ai_role[1] = st.session_state.ai_role[0]
    st.session_state.prompt_exists = False
    st.session_state.temperature[1] = st.session_state.temperature[0]
    st.session_state.audio_response = None
    st.session_state.vector_store_message = None
    st.session_state.tool_names[1] = st.session_state.tool_names[0]
    st.session_state.agent_type[1] = st.session_state.agent_type[0]
    st.session_state.retriever_tool = None
    st.session_state.uploader_key = 0


def switch_between_apps() -> None:
    """
    Keep the chat settings when switching the mode.
    """

    st.session_state.temperature[1] = st.session_state.temperature[0]
    st.session_state.ai_role[1] = st.session_state.ai_role[0]
    st.session_state.tool_names[1] = st.session_state.tool_names[0]
    st.session_state.agent_type[1] = st.session_state.agent_type[0]


@tool
def python_repl(
    code: Annotated[str, "The python code to execute to generate your chart."],
):
    """Use this to execute python code. If you want to see the output of a value,
    you should print it out with `print(...)`. This is visible to the user."""
    try:
        result = PythonREPL().run(code)
    except BaseException as e:
        return f"Failed to execute. Error: {repr(e)}"
    result_str = f"Successfully executed:\n```python\n{code}\n```\nStdout: {result}"
    return (
        result_str + "\n\nIf you have completed all tasks, respond with FINAL ANSWER."
    )


def set_tools() -> List[Tool]:
    """
    Set and return the tools for the agent. Tools that can be selected
    are internet_search, arxiv, wikipedia, python_repl, and retrieval.
    A Bing Subscription Key or Google CSE ID is required for internet_search.
    """

    class MySearchToolInput(BaseModel):
        query: str = Field(description="search query to look up")

    # Load tools
    arxiv = load_tools(["arxiv"])[0]
    wikipedia = load_tools(["wikipedia"])[0]
    # Python REPL is directly used here
    tool_dictionary = {
        "ArXiv": arxiv,
        "Wikipedia": wikipedia,
        "Python_REPL": python_repl,
        "Retrieval": st.session_state.retriever_tool if st.session_state.retriever_tool else None
    }
    tool_options = ["ArXiv", "Wikipedia", "Python_REPL", "Retrieval"]

    # Add Search tool dynamically if credentials are valid
    if st.session_state.bing_subscription_validity:
        search = BingSearchAPIWrapper()
    elif st.session_state.google_cse_id_validity:
        search = GoogleSearchAPIWrapper()
    else:
        search = None

    if search is not None:
        internet_search = Tool(
            name="internet_search",
            description=(
                "A search engine for comprehensive, accurate, and trusted results. "
                "Useful for when you need to answer questions about current events. "
                "Input should be a search query."
            ),
            func=partial(search.results, num_results=5),
            args_schema=MySearchToolInput,
        )
        tool_options.insert(0, "Search")
        tool_dictionary["Search"] = internet_search

    # UI for selecting tools
    st.write("")
    st.write("**Tools**")
    tool_names = st.multiselect(
        label="assistant tools",
        options=tool_options,
        default=st.session_state.tool_names[1],
        label_visibility="collapsed",
    )

    # Instructions if Search tool is unavailable
    if "Search" not in tool_options:
        st.write(
            "<small>Tools are disabled when images are uploaded and queried. "
            "To search the internet, obtain your Bing Subscription Key "
            "[here](https://portal.azure.com/) or Google CSE ID "
            "[here](https://programmablesearchengine.google.com/about/), "
            "and enter it in the sidebar. Once entered, 'Search' will be displayed "
            "in the list of tools. Note also that PythonREPL from LangChain is still "
            "in the experimental phase, so caution is advised.</small>",
            unsafe_allow_html=True,
        )
    else:
        st.write(
            "<small>Tools are disabled when images are uploaded and queried. "
            "Note also that PythonREPL from LangChain is still in the experimental phase, "
            "so caution is advised.</small>",
            unsafe_allow_html=True,
        )

    # Handle Retrieval tool initialization
    if "Retrieval" in tool_names:
        if not st.session_state.retriever_tool:
            st.info("Creating the vector store and initializing the retriever tool...")
            get_retriever()
        if st.session_state.retriever_tool:
            st.success("Retriever tool is ready for querying.")
            tool_dictionary["Retrieval"] = st.session_state.retriever_tool
        else:
            st.error("Failed to initialize the retriever tool. Please upload the document again.")
            tool_names.remove("Retrieval")  # Prevent broken Retrieval tool

    # Final tool selection
    tools = [
        tool_dictionary[key]
        for key in tool_names if tool_dictionary[key] is not None
    ]

    st.write("**Tools selected in set_tools:**", [tool.name for tool in tools])
    st.session_state.tool_names[0] = tool_names

    return tools



def set_prompts(agent_type: Literal["Tool Calling", "ReAct"]) -> None:
    """
    Set chat and agent prompts for two different types of agents:
    Tool Calling and ReAct.
    """

    if agent_type == "Tool Calling":
        st.session_state.chat_prompt = ChatPromptTemplate.from_messages([
            (
                "system",
                f"{st.session_state.ai_role[0]} Your goal is to provide "
                "answers to human inquiries. Should the information not "
                "be available, inform the human explicitly that "
                "the answer could not be found."
            ),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
        ])
        st.session_state.agent_prompt = ChatPromptTemplate.from_messages([
            (
                "system",
                f"{st.session_state.ai_role[0]} Your goal is to provide answers to human inquiries. "
                "You should specify the source of your answers, whether they are based on internet search "
                "results ('internet_search'), scientific articles from arxiv.org ('arxiv'), Wikipedia documents ('wikipedia'), "
                "uploaded documents ('retriever'), or your general knowledge. "
                "Use the 'retriever' tool to answer questions specifically related to uploaded documents. "
                "If you cannot find relevant information in the documents using the 'retriever' tool, explicitly inform the user. "
                "Use Markdown syntax and include relevant sources, such as links (URLs)."
            ),
            MessagesPlaceholder(variable_name="chat_history", optional=True),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])
    else:
        st.session_state.chat_prompt = ChatPromptTemplate.from_template(
            f"{st.session_state.ai_role[0]} "
            "Your goal is to provide answers to human inquiries. "
            "Should the information not be available, inform the human "
            "explicitly that the answer could not be found.\n\n"
            "{chat_history}\n\nHuman: {input}\n\n"
            "AI: "
        )
        st.session_state.agent_prompt = ChatPromptTemplate.from_template(
            f"{st.session_state.ai_role[0]} "
            "Your goal is to provide answers to human inquiries. "
            "When giving your answers, tell the human what your response "
            "is based on and which tools you use. Use Markdown syntax "
            "and include relevant sources, such as links (URLs), following "
            "MLA format. Should the information not be available, inform "
            "the human explicitly that the answer could not be found.\n\n"
            "TOOLS:\n"
            "------\n\n"
            "You have access to the following tools:\n\n"
            "{tools}\n\n"
            "To use a tool, please use the following format:\n\n"
            "Thought: Do I need to use a tool? Yes\n"
            "Action: the action to take, should be one of [{tool_names}]\n"
            "Action Input: the input to the action\n"
            "Observation: the result of the action\n\n"
            "When you have a response to say to the Human, "
            "or if you do not need to use a tool, you MUST use "
            "the format:\n\n"
            "Thought: Do I need to use a tool? No\n"
            "Final Answer: [your response here]\n\n"
            "Begin!\n\n"
            "Previous conversation history:\n\n"
            "{chat_history}\n\n"
            "New input: {input}\n"
            "{agent_scratchpad}"
        )


def print_conversation(no_of_msgs: Union[Literal["All"], int]) -> None:
    """
    Print the conversation stored in st.session_state.history.
    """

    if no_of_msgs == "All":
        no_of_msgs = len(st.session_state.history)

    for msg in st.session_state.history[-no_of_msgs:]:
        if isinstance(msg, HumanMessage):
            with st.chat_message("human"):
                st.write(msg.content)
        else:
            with st.chat_message("ai"):
                display_text_with_equations(msg.content)

        if urls := msg.additional_kwargs.get("image_urls"):
            for url in urls:
                st.image(url)

    # Play TTS
    if (
        st.session_state.model_type == "GPT Models from OpenAI"
        and st.session_state.audio_response is not None
    ):
        play_audio(st.session_state.audio_response)
        st.session_state.audio_response = None


def serialize_messages(
    messages: List[Union[HumanMessage, AIMessage]]
) -> List[Dict]:

    """
    Serialize the list of messages into a list of dicts
    """

    return [msg.dict() for msg in messages]


def deserialize_messages(
    serialized_messages: List[Dict]
) -> List[Union[HumanMessage, AIMessage]]:

    """
    Deserialize the list of messages from a list of dicts
    """

    deserialized_messages = []
    for msg in serialized_messages:
        if msg['type'] == 'human':
            deserialized_messages.append(HumanMessage(**msg))
        elif msg['type'] == 'ai':
            deserialized_messages.append(AIMessage(**msg))
    return deserialized_messages


def show_uploader() -> None:
    """
    Set the flag to show the uploader.
    """

    st.session_state.show_uploader = True


def check_conversation_keys(lst: List[Dict[str, Any]]) -> bool:
    """
    Check if all items in the given list are valid conversation entries.
    """

    return all(
        isinstance(item, dict) and
        isinstance(item.get("content"), str) and
        isinstance(item.get("type"), str) and
        isinstance(item.get("additional_kwargs"), dict)
        for item in lst
    )


def load_conversation() -> bool:
    """
    Load the conversation from a JSON file
    """

    st.write("")
    st.write("**Choose a (JSON) conversation file**")
    uploaded_file = st.file_uploader(
        label="Load conversation", type="json", label_visibility="collapsed"
    )
    if uploaded_file:
        try:
            data = json.load(uploaded_file)
            if isinstance(data, list) and check_conversation_keys(data):
                st.session_state.history = deserialize_messages(data)
                return True
            st.error(
                f"The uploaded data does not conform to the expected format.", icon="🚨"
            )
        except Exception as e:
            st.error(f"An error occurred: {e}", icon="🚨")

    return False


def create_text(model: str) -> None:
    """
    Take an LLM as input and generate text based on user input
    by calling run_agent().
    """

    # initial system prompts
    general_role = "You are a helpful AI assistant."
    english_teacher = (
        "You are an AI English teacher who analyzes texts and corrects "
        "any grammatical issues if necessary."
    )
    translator = (
        "You are an AI translator who translates English into Korean "
        "and Korean into English."
    )
    coding_adviser = (
        "You are an AI expert in coding who provides advice on "
        "good coding styles."
    )
    science_assistant = "You are an AI science assistant."
    roles = (
        general_role, english_teacher, translator,
        coding_adviser, science_assistant
    )

    with st.sidebar:
        st.write("")
        type_options = ("Tool Calling", "ReAct")
        st.write("**Agent Type**")
        st.session_state.agent_type[0] = st.sidebar.radio(
            label="Agent Type",
            options=type_options,
            index=type_options.index(st.session_state.agent_type[1]),
            label_visibility="collapsed",
        )
        agent_type = st.session_state.agent_type[0]
        if st.session_state.model_type == "GPT Models from OpenAI":
            st.write("")
            st.write("**Text to Speech**")
            st.session_state.tts = st.radio(
                label="TTS",
                options=("Enabled", "Disabled", "Auto"),
                # horizontal=True,
                index=1,
                label_visibility="collapsed",
            )
        st.write("")
        st.write("**Temperature**")
        st.session_state.temperature[0] = st.slider(
            label="Temperature (higher $\Rightarrow$ more random)",
            min_value=0.0,
            max_value=1.0,
            value=st.session_state.temperature[1],
            step=0.1,
            format="%.1f",
            label_visibility="collapsed",
        )
        st.write("")
        st.write("**Messages to Show**")
        no_of_msgs = st.radio(
            label="$\\textsf{Messages to show}$",
            options=("All", 20, 10),
            label_visibility="collapsed",
            horizontal=True,
            index=2,
        )

    st.write("")
    st.write("##### Message to AI")
    st.session_state.ai_role[0] = st.selectbox(
        label="AI's role",
        options=roles,
        index=roles.index(st.session_state.ai_role[1]),
        label_visibility="collapsed",
    )

    if st.session_state.ai_role[0] != st.session_state.ai_role[1]:
        reset_conversation()
        st.rerun()

    st.write("")
    st.write("##### Conversation with AI")

    # Print conversation
    print_conversation(no_of_msgs)

    # Reset, download, or load the conversation
    c1, c2, c3 = st.columns(3)
    c1.button(
        label="$~\:\,\,$Reset$~\:\,\,$",
        on_click=reset_conversation
    )
    c2.download_button(
        label="Download",
        data=json.dumps(serialize_messages(st.session_state.history), indent=4),
        file_name="conversation_with_agent.json",
        mime="application/json",
    )
    c3.button(
        label="$~~\:\,$Load$~~\:\,$",
        on_click=show_uploader,
    )

    if st.session_state.show_uploader and load_conversation():
        st.session_state.show_uploader = False
        st.rerun()

    # Set the agent prompts and tools
    set_prompts(agent_type)
    tools = set_tools()
    st.write("**Tools passed to run_agent:**", [tool.name for tool in tools])


    image_urls = []
    with st.sidebar:
        image_urls = upload_image_files_return_urls()

    if st.session_state.model_type == "GPT Models from OpenAI":
        audio_input = input_from_mic()
        if audio_input is not None:
            query = audio_input
            st.session_state.prompt_exists = True
            st.session_state.mic_used = True

    # Use your keyboard
    text_input = st.chat_input(placeholder="Enter your query")

    if text_input:
        query = text_input.strip()
        st.session_state.prompt_exists = True

    if st.session_state.prompt_exists:
        with st.chat_message("human"):
            st.write(query)

        with st.chat_message("ai"):
            generated_text = run_agent(
                query=query,
                model=model,
                tools=tools,
                image_urls=image_urls,
                temperature=st.session_state.temperature[0],
                agent_type=agent_type,
            )
            fig = plt.gcf()
            if fig and fig.get_axes():
                generated_image_url = fig_to_base64(fig)
                st.session_state.history[-1].additional_kwargs["image_urls"] = [
                    generated_image_url
                ]
        if (
            st.session_state.model_type == "GPT Models from OpenAI"
            and generated_text is not None
        ):
            # TTS under two conditions
            cond1 = st.session_state.tts == "Enabled"
            cond2 = st.session_state.tts == "Auto" and st.session_state.mic_used
            if cond1 or cond2:
                st.session_state.audio_response = perform_tts(generated_text)
            st.session_state.mic_used = False

        st.session_state.prompt_exists = False

        if generated_text is not None:
            st.session_state.uploader_key += 1
            st.rerun()


def create_image(model: str) -> None:
    """
    Generate image based on user description by calling openai_create_image().
    """

    # Set the image size
    with st.sidebar:
        st.write("")
        st.write("**Pixel size**")
        image_size = st.radio(
            label="$\\hspace{0.1em}\\texttt{Pixel size}$",
            options=("1024x1024", "1792x1024", "1024x1792"),
            # horizontal=True,
            index=0,
            label_visibility="collapsed",
        )

    st.write("")
    st.write("##### Description for your image")

    if st.session_state.image_url is not None:
        st.info(st.session_state.image_description)
        st.image(image=st.session_state.image_url, use_column_width=True)
    
    # Get an image description using the microphone
    if st.session_state.model_type == "GPT Models from OpenAI":
        audio_input = input_from_mic()
        if audio_input is not None:
            st.session_state.image_description = audio_input
            st.session_state.prompt_exists = True

    # Get an image description using the keyboard
    text_input = st.chat_input(
        placeholder="Enter a description for your image",
    )
    if text_input:
        st.session_state.image_description = text_input.strip()
        st.session_state.prompt_exists = True

    if st.session_state.prompt_exists:
        st.session_state.image_url = openai_create_image(
            st.session_state.image_description, model, image_size
        )
        st.session_state.prompt_exists = False
        if st.session_state.image_url is not None:
            st.rerun()


def create_text_image() -> None:
    """
    Generate text or image by using LLM models like 'gpt-4o'.
    """

    page_title = "LangChain LLM Agent"
    page_icon = "📚"

    st.set_page_config(
        page_title=page_title,
        page_icon=page_icon,
        layout="centered"
    )

    st.write(f"## {page_icon} $\,${page_title}")

    # Initialize all the session state variables
    initialize_session_state_variables()

    # Define model options directly here
    model_options = ["gpt-4o-mini", "gpt-4o", "dall-e-3"]

    # Sidebar content
    with st.sidebar:
        st.write("**Select a Model**")
        model = st.radio(
            label="Models",
            options=model_options,
            index=1,  # Default to the second option
            label_visibility="collapsed",
            on_change=switch_between_apps,
        )

        st.write("---")
        st.write("xyz", unsafe_allow_html=True)

    # Main logic for generating text or image
    if model == "dall-e-3":
        create_image(model)
    else:
        create_text(model)

if __name__ == "__main__":
    create_text_image()