Update app.py

app.py

@@ -4,35 +4,60 @@ from dotenv import load_dotenv
 load_dotenv()
 
 from langchain.agents.openai_assistant import OpenAIAssistantRunnable
-from langchain.agents import AgentExecutor
-
 from langchain.schema import HumanMessage, AIMessage
 
-import gradio
+import gradio as gr
 
+# Load API key and assistant IDs
 api_key = os.getenv('OPENAI_API_KEY')
-extractor_agent = os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_A')
-
-extractor_llm = OpenAIAssistantRunnable(assistant_id=extractor_agent, api_key=api_key, as_agent=True)
-
-
+extractor_agents = {
+    "Solution Specifier A": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_A'),
+    "Solution Specifier B": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_B'),
+    "Solution Specifier C": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_C'),
+    "Solution Specifier D": os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_D'),
+}
+
+# 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)
 
-def predict(message, history):
+# 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
 
-#gradio.Markdown("Click [here](https://www.google.com) to visit Google.")
-chat = gradio.ChatInterface(predict, title="Solution Specifier A", description="testing for the time being")
-chat.launch(share=True)
+# Define the Gradio interface
+def app_interface():
+    dropdown = gr.Dropdown(choices=list(extractor_agents.keys()), value="Solution Specifier A", label="Choose Extractor Agent")
+    chat = gr.ChatInterface(
+        fn=lambda message, history, selected_agent: predict(message, history, selected_agent),
+        inputs=[dropdown],
+        title="Solution Specifier Chat",
+        description="Test with different solution specifiers"
+    )
+    return chat
+
+# Launch the app
+chat_interface = app_interface()
+chat_interface.launch(share=True)