Update agent.py

agent.py

@@ -112,6 +112,22 @@ def arvix_search(query: str) -> str:
 
 
 
+system_prompt=  '''
+You are a helpful assistant with access to tools for answering questions. For each question, think step by step and clearly explain your reasoning. Conclude your response using the following format:
+
+FINAL ANSWER: [YOUR FINAL ANSWER]
+
+Formatting rules for YOUR FINAL ANSWER:
+
+If the answer is a number, write it as digits only. Do not include commas or units (e.g. $, %, etc.) unless explicitly required.
+
+If the answer is a string, avoid articles ("a", "an", "the"), do not use abbreviations (e.g., for cities), and spell out all digits (e.g., "two thousand").
+
+If the answer is a comma-separated list, apply the above rules to each item in the list, depending on whether they are numbers or strings.
+
+Your answer should only begin with "FINAL ANSWER: " followed directly by the formatted result. Nothing else should follow.
+'''
+sys_msg = SystemMessage(content=system_prompt)
 
 tools = [
     multiply,
@@ -155,20 +171,25 @@ def build_graph(provider: str = "huggingface"):
     def assistant(state: MessagesState):
         """Assistant node"""
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
+
+    def retriever(state: MessagesState):
+        """Retriever node"""
+       
+        return {"messages": [sys_msg] }
     
 
 
     builder = StateGraph(MessagesState)
-  
+    builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
-
-    builder.add_edge(START, "assistant")
+    builder.add_edge(START, "retriever")
+    builder.add_edge("retriever", "assistant")
     builder.add_conditional_edges(
         "assistant",
         tools_condition,
     )
-    builder.add_edge("tools", "assistant")
+    builder.add_edge("tools", "assistant"))
 
     # Compile graph
     return builder.compile()