Update app.py

app.py

@@ -21,68 +21,62 @@ import uuid
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+SERPER_API_KEY=os.environ.get("SERPER_API_KEY")
+OPENAI_KEY=os.environ.get("OPENAI_KEY")
 
 
-
-wikipedia = WikipediaAPIWrapper()
-def wikipedia_search(query: str) -> str:
-    return wikipedia.run(query)
-
 # --- Basic Agent Definition ---
 # ----- THIS IS WHERE YOU CAN BUILD WHAT YOU WANT ------
-def web_search(query: str):
-    """Search the web for information."""
-    return DuckDuckGoSearchRun().run(query)
+serper_tool = GoogleSerperAPIWrapper()
 
-def wikipedia_search(query: str):
-    """Search Wikipedia for information."""
-    return wikipedia.run(query)
-
-HUGGINGFACEHUB_API_TOKEN=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
-
-class BasicAgent:
-    def __init__(self, model_id="ollama/qwen2:7b"):
-        print('BasicAgent initialized.')
-        self.system_prompt = (
-            "You are a general AI assistant. I will ask you a question. "
+def wikipedia_search(query: str) -> str:
+    return WikipediaAPIWrapper().run(query)
+
+serper_tool = GoogleSerperAPIWrapper()
+
+tools = [
+    Tool(
+        name = 'wikipedia_search',
+        func = wikipedia_search,
+        description = 'Search Wikipedia for factual information.'
+    ),
+    Tool(
+        name = 'google_search',
+        func = serper_tool.run,
+        description = 'Search the web using Google Serper.'
+    )
+]
+
+llm = ChatOpenAI(
+    model = 'gpt-4o-mini',
+    temperature = 0,
+    openai_api_key = OPENAI_KEY,
+    max_tokens = 4096
+)
+
+agent = initialize_agent(
+    tools=tools,
+    llm=llm,
+    agent=AgentType.OPENAI_FUNCTIONS,
+    verbose=True,
+    agent_kwargs={
+        "system_message": SystemMessage(content="You are a general AI assistant. I will ask you a question. "
             "Report your thoughts, 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. "
+            "[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, don't use comma 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."
-        )
-        web_search_tool = DuckDuckGoSearchRun()
-        #wikipedia_search_tool = wikipedia_search()
-        tools = [web_search_tool]
-        llm = HuggingFaceEndpoint(repo_id="Qwen/Qwen2.5-Coder-32B-Instruct", huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN)
-
-        chat = ChatHuggingFace(llm=llm, verbose=True)
-        chat_with_tools = chat.bind_tools(tools)
-
-        class AgentState(TypedDict):
-            messages: Annotated[list[AnyMessage], add_messages]
-
-        def assistant(state: AgentState):
-            return {
-                "messages": [chat_with_tools.invoke(state["messages"])],
-            }
-
-        builder = StateGraph(AgentState)
-        builder.add_node("assistant", assistant)
-        builder.add_node("tools", ToolNode(tools))
-        builder.add_edge(START, "assistant")
-        builder.add_conditional_edges("assistant", tools_condition)
-        builder.add_edge("tools", "assistant")
-        self.alfred = builder.compile()
+            "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.")
+    }
+)
+
+
+class BasicAgent:
+    def __init__(self):
+        self.agent = agent
 
     def __call__(self, question: str) -> str:
-        messages = [
-            SystemMessage(content=self.system_prompt),
-            HumanMessage(content=question)
-        ]
-        response = self.alfred.invoke({"messages": messages})
-        model_answer = response["messages"][-1].content
-        return model_answer
+        print(f"\nQuestion: {question}")
+        return self.agent.run(question)
 
     
 def run_and_submit_all( profile: gr.OAuthProfile | None):