Update app.py

app.py

@@ -26,13 +26,20 @@ if "Mistral" not in os.environ:
     os.environ["Mistral"] = getpass.getpass("Enter your Mistral API key: ")
 
 print("Loading LLM...")
-chat = ChatMistralAI(
+chat1 = ChatMistralAI(
     model="pixtral-large-2411",
     temperature=0.01,
     max_retries=6,
     mistral_api_key = os.getenv("Mistral")
 )
-print(f"Model {chat.model} downloaded!")
+print(f"Model {chat1.model} downloaded!")
+chat2 = ChatMistralAI(
+    model="mistral-medium-2505",
+    temperature=0.01,
+    max_retries=6,
+    mistral_api_key = os.getenv("Mistral")
+)
+print(f"Model {chat2.model} downloaded!")
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -133,8 +140,7 @@ run_code_tool = Tool(
     description="Executes a python code."
 )
 
-tools = [get_file_path_tool, get_content_tool, search_web_tool, extract_text_tool, play_chess_tool, run_code_tool]
-chat_with_tools = chat.bind_tools(tools, parallel_tool_calls=False)
+
 
 
 # Generate the AgentState and Agent graph
@@ -142,59 +148,68 @@ class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
 
 
-def assistant(state: AgentState):
+def assistant(state: AgentState, chat):
     return {
         "messages": chat.invoke(state["messages"]),
     }
 
+def build_agent(chat):
+    tools = [get_file_path_tool, get_content_tool, search_web_tool, extract_text_tool, play_chess_tool, run_code_tool]
+    chat_with_tools = chat.bind_tools(tools, parallel_tool_calls=False)
+    # The graph
+    builder = StateGraph(AgentState)
 
-# The graph
-builder = StateGraph(AgentState)
-
-# Define nodes: these do the work
-builder.add_node("assistant", assistant)
-builder.add_node("tools", ToolNode([get_file_path_tool]))
-
-# Define edges: these determine how the control flow moves
-builder.add_edge(START, "assistant")
-builder.add_conditional_edges(
-    "assistant",
-    # If the latest message requires a tool, route to tools
-    # Otherwise, provide a direct response
-    tools_condition
-)
-builder.add_edge("tools", "assistant")
-alfred = builder.compile()
-system_prompt = SystemMessage(
-    content="You are a general AI assistant. \
-        I will ask you a question. Report your thoughts shortly, and finish your answer with the following template: \
-        FINAL ANSWER: YOUR FINAL ANSWER \
-        YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. \
-        If you are asked for a number, use only digits in your final answer. Don't use comma nor brackets to write your number neither use units such as $ or percent sign unless specified otherwise. \
-        If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. \
-        If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. Wrap the list in [square brackets].\
-        If there is a file attached, open the file and read it. \
-        If you don't have enough references to answer, use your tools, search the web, run your code or convert data to a data frame, whatever helps. \
-        If the question refers to an external content and there is no reference file attached, perform a web search and retrieve relevant information from the internet. \
-        If there is a code, execute it. \
-        Make sure that each final answer is preceded with 'FINAL ANSWER:' and is short: it should contain a number (without full stop at the end), a string (one or two words only, without full stop at the end) or a comma-separated list (without full stops at the end), nothing else. "
+    # Define nodes: these do the work
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode([get_file_path_tool]))
+    
+    # Define edges: these determine how the control flow moves
+    builder.add_edge(START, "assistant")
+    builder.add_conditional_edges(
+        "assistant",
+        # If the latest message requires a tool, route to tools
+        # Otherwise, provide a direct response
+        tools_condition
     )
+    builder.add_edge("tools", "assistant")
+    alfred = builder.compile()    
+    return alfred
 
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
 
-    def __call__(self, question: str) -> str:
+    def __call__(self, question: str) -> str:  
+      
         print(f"Agent received question (first 50 chars): {question[:50]}...")
+        system_prompt = SystemMessage(
+        content="You are a general AI assistant. \
+            I will ask you a question. Report your thoughts shortly, and finish your answer with the following template: \
+            FINAL ANSWER: YOUR FINAL ANSWER \
+            YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. \
+            If you are asked for a number, use only digits in your final answer. Don't use comma nor brackets to write your number neither use units such as $ or percent sign unless specified otherwise. \
+            If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. \
+            If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. Wrap the list in [square brackets].\
+            If there is a file attached, open the file and read it. \
+            If you don't have enough references to answer, use your tools, search the web, run your code or convert data to a data frame, whatever helps. \
+            If the question refers to an external content and there is no reference file attached, perform a web search and retrieve relevant information from the internet. \
+            If there is a code, execute it. \
+            Make sure that each final answer is preceded with 'FINAL ANSWER:' and is short: it should contain a number (without full stop at the end), a string (one or two words only, without full stop at the end) or a comma-separated list (without full stops at the end), nothing else. "
+        )
         message = HumanMessage(content=question)
         print(message)
-        answer = alfred.invoke(input={"messages": [system_prompt, message]},config={"recursion_limit": 6})['messages'][-1].content
+        try:
+            alfred = build_agent(chat1)
+            answer = alfred.invoke(input={"messages": [system_prompt, message]},config={"recursion_limit": 6})['messages'][-1].content
+        except Exception as e:
+            alfred = build_agent(chat2)
         answer = "".join(re.findall(r'(FINAL ANSWER.*)', answer, flags=re.M))    
         answer = answer.replace('FINAL ANSWER:', '')
         answer = answer.replace('FINAL ANSWER', '')
         answer = answer.replace('YOUR ', '')
         answer = answer.replace('*', '')
-        print(f"Agent returning fixed answer: {answer}")
+        print(f"Agent returning fixed answer: {answer}")        
+            
         return answer