Update app.py

app.py

@@ -15,7 +15,7 @@ from langchain_community.utilities.arxiv import ArxivAPIWrapper
 
 from langgraph.graph import StateGraph, END
 
-from langchain_core.messages import BaseMessage, FunctionMessage, HumanMessage, AIMessage
+from langchain_core.messages import BaseMessage, FunctionMessage, HumanMessage, AIMessage, SystemMessage
 from langchain_openai import ChatOpenAI
 
 # --- Constants ---
@@ -49,16 +49,31 @@ class AgentState(TypedDict):
     next_action: Optional[str] # To decide if we need to call tools or respond
 
 class LangGraphAgent:
-    def __init__(self):
-        print("LangGraphAgent initializing...")
+    def __init__(self, llm_choice: str = "gemini"):
+        print(f"LangGraphAgent initializing with {llm_choice}...")
         if not OPENROUTER_API_KEY:
             raise ValueError("OPENROUTER_API_KEY is not set. Cannot initialize LLM.")
 
-        self.llm = ChatOpenAI(
-            model="google/gemini-2.0-flash-001",
-            api_key=OPENROUTER_API_KEY,
-            base_url="https://openrouter.ai/api/v1"
-        )
+        if llm_choice == "llama":
+            self.llm = ChatOpenAI(
+                model="meta-llama/llama-3.1-8b-instruct:free",
+                api_key=OPENROUTER_API_KEY,
+                base_url="https://openrouter.ai/api/v1",
+                temperature=0.1, # Llama models can be sensitive to temperature
+                # max_tokens=150 # Llama 8B might benefit from a smaller max_token for concise answers
+            )
+            print("Initialized Llama 3.1 8B Instruct.")
+        elif llm_choice == "gemini":
+            self.llm = ChatOpenAI(
+                model="google/gemini-2.0-flash-001",
+                api_key=OPENROUTER_API_KEY,
+                base_url="https://openrouter.ai/api/v1",
+                temperature=0.1 # Adding temperature for consistency
+            )
+            print("Initialized Gemini 2.0 Flash.")
+        else:
+            raise ValueError(f"Unsupported LLM choice: {llm_choice}. Choose 'gemini' or 'llama'.")
+
         self.tools_map = {tool.name: tool for tool in tools}
         self.graph = self._build_graph()
         print("LangGraphAgent initialized.")
@@ -127,41 +142,48 @@ class LangGraphAgent:
         return {"messages": tool_messages}
 
     def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:100]}...")
-        initial_state = {"messages": [HumanMessage(content=question)]}
+        print(f"Agent received question (first 100 chars): {question[:100]}...")
+        
+        system_prompt = (
+            "You are an AI assistant designed to answer questions concisely. "
+            "Your goal is to provide only the direct answer to the question, without any additional explanations, conversation, or prefixes like 'FINAL ANSWER:'. "
+            "For example, if the question is 'What is the capital of France?', you should respond with 'Paris'. "
+            "If the question asks for a list, provide it comma-separated, e.g., 'apple, banana, cherry'. "
+            "If the question asks for a number, provide only the number, e.g., '42'."
+        )
+        initial_state = {"messages": [SystemMessage(content=system_prompt), HumanMessage(content=question)]}
         
-        # The GAIA prompt example suggests not including "FINAL ANSWER" and just replying with the answer.
-        # We need to ensure the LLM is prompted to provide a direct answer after tool use.
-        # For simplicity in this template, we will take the last AI message content as the answer.
-        # A more robust solution might involve a specific "final answer" node or prompt engineering.
-
         final_graph_state = None
         try:
             for event in self.graph.stream(initial_state, {"recursion_limit": 100}): # Added recursion limit
-                # print(f"Graph event: {event}") # For debugging stream
                 if END in event:
                     final_graph_state = event[END]
                     break
-                # Update final_graph_state with the latest state from any node
-                # This ensures we have the latest messages even if END is not directly reached by llm
-                # (e.g. if recursion limit is hit)
                 for key in event:
                     if key != END:
                         final_graph_state = event[key]
 
             if final_graph_state and final_graph_state["messages"]:
-                # Get the last AI message as the answer
                 for msg in reversed(final_graph_state["messages"]):
                     if isinstance(msg, AIMessage) and not msg.tool_calls:
                         answer = msg.content.strip()
-                        # Ensure no "FINAL ANSWER:" prefix as per GAIA instructions
-                        if answer.upper().startswith("FINAL ANSWER:"):
-                            answer = answer[len("FINAL ANSWER:"):].strip()
+                        # Remove common prefixes that LLMs might add despite instructions
+                        prefixes_to_remove = [
+                            "FINAL ANSWER:", "The answer is", "Here is the answer:", 
+                            "The final answer is", "Answer:", "Solution:"
+                        ]
+                        for prefix in prefixes_to_remove:
+                            if answer.upper().startswith(prefix.upper()):
+                                answer = answer[len(prefix):].strip()
+                        
+                        # Remove potential quotation marks if the answer is a single word/phrase
+                        if len(answer.split()) < 5: # Heuristic for short answers
+                            if answer.startswith(('"', "'")) and answer.endswith(('"', "'")):
+                                answer = answer[1:-1]
+
                         print(f"Agent returning answer: {answer}")
                         return answer
-                # Fallback if no suitable AI message is found
                 print("No suitable AI message found for final answer. Returning last message content.")
-                # This might be a tool call or an intermediate thought, not ideal.
                 return str(final_graph_state["messages"][-1].content) if final_graph_state["messages"] else "Error: No messages in final state."
             else:
                 print("Error: Agent did not reach a final state or no messages found.")
@@ -199,7 +221,8 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     submit_url = f"{api_url}/submit"
 
     try:
-        agent = LangGraphAgent()
+        # Default to Llama for now, can be made configurable later (e.g., via Gradio input)
+        agent = LangGraphAgent(llm_choice="llama") 
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None