app.py

import os
import re
from dotenv import load_dotenv
load_dotenv()

from langchain.agents.openai_assistant import OpenAIAssistantRunnable
from langchain.schema import HumanMessage, AIMessage

import gradio as gr

# Load API key and assistant IDs
api_key = os.getenv('OPENAI_API_KEY')
extractor_agents = {
    "Solution Specifier A": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_A'),
    "Solution Specifier A1": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_A1'),
}

# Function to create a new extractor LLM instance
def get_extractor_llm(agent_id):
    return OpenAIAssistantRunnable(assistant_id=agent_id, api_key=api_key, as_agent=True)

# Utility function to remove citations
def remove_citation(text):
    # Define the regex pattern to match the citation format 【number†text】
    pattern = r"【\d+†\w+】"
    # Replace the pattern with an empty string
    return re.sub(pattern, "📚", text)

# Prediction function
def predict(message, history, selected_agent):
    # Get the extractor LLM for the selected agent
    agent_id = extractor_agents[selected_agent]
    extractor_llm = get_extractor_llm(agent_id)
    
    # Prepare the chat history
    history_langchain_format = []
    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))
    history_langchain_format.append(HumanMessage(content=message))
    
    # Get the response
    gpt_response = extractor_llm.invoke({"content": message})
    output = gpt_response.return_values["output"]
    non_cited_output = remove_citation(output)
    return non_cited_output

# Define the Gradio interface using Blocks for layout
def app_interface():
    with gr.Blocks() as demo:
        dropdown = gr.Dropdown(
            choices=list(extractor_agents.keys()),
            value="Solution Specifier A",
            label="Choose Extractor Agent",
            interactive=True,
        )

        chat = gr.ChatInterface(
            fn=predict,
            additional_inputs=[dropdown],
            title="Solution Specifier Chat",
            description="Test with different solution specifiers"
        )
    return demo

# Launch the app
chat_interface = app_interface()
chat_interface.launch(share=True)