Update app.py

app.py

@@ -40,23 +40,47 @@ def wikipedia_search(query: str):
 class BasicAgent:
     def __init__(self, model_id="ollama/qwen2:7b"):
         print('BasicAgent initialized.')
-        self.tool_node = [web_search, wikipedia_search]
-        self.agent = ToolCallingAgent(
-            tools=[self.tool_node],
-            model=LiteLLMModel(
-                model_id=model_id,
-            ),
-            max_steps=15,
-            verbosity_level=2,
+        self.system_prompt = (
+            "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. "
+            "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="ollama/qwen2:7b", huggingfacehub_api_token=os.environ.get("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()
 
     def __call__(self, question: str) -> str:
-        print(f"Agent received question: {question[:50]}...")
-        try:
-            return self.agent.run(question)
-        except Exception as e:
-            print(f"Error in agent: {e}")
-            return f"AGENT ERROR: {e}"
+        messages = [
+            SystemMessage(content=self.system_prompt),
+            HumanMessage(content=question)
+        ]
+        response = self.alfred.invoke({"messages": messages})
+        model_answer = response["messages"][-1].content
+        return model_answer
+
     
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """