add system message, ruff format, change pro to flash

agent.py

@@ -1,17 +1,19 @@
 import os
 from dotenv import load_dotenv
 
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.tools import tool
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
 
-from langgraph.graph import StateGraph, START, END, MessagesState
+from langgraph.graph import StateGraph, START, MessagesState
 from langgraph.prebuilt import ToolNode, tools_condition
 
 load_dotenv()
 
+SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question. Report your thoughts, and output only your final answer, no prefixes, suffixes, or extra text. 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."""
+
 
 @tool
 def add(a: float, b: float) -> float:
@@ -80,7 +82,7 @@ def power(a: float, b: float) -> float:
         a: Base number
         b: Exponent
     """
-    return a ** b
+    return a**b
 
 
 @tool
@@ -92,7 +94,7 @@ def square_root(a: float) -> float:
     """
     if a < 0:
         return "Error: Cannot calculate square root of negative number"
-    return a ** 0.5
+    return a**0.5
 
 
 @tool
@@ -111,9 +113,9 @@ def web_search(query: str) -> str:
 
         formatted_results = []
         for i, result in enumerate(results, 1):
-            title = result.get('title', 'No title')
-            content = result.get('content', 'No content')
-            url = result.get('url', 'No URL')
+            title = result.get("title", "No title")
+            content = result.get("content", "No content")
+            url = result.get("url", "No URL")
             formatted_results.append(f"{i}. {title}\n{content}\nSource: {url}")
 
         return "\n\n ==== \n\n".join(formatted_results)
@@ -137,7 +139,7 @@ def wikipedia_search(query: str) -> str:
 
         formatted_docs = []
         for i, doc in enumerate(docs, 1):
-            title = doc.metadata.get('title', 'No title')
+            title = doc.metadata.get("title", "No title")
             content = doc.page_content
             formatted_docs.append(f"{i}. {title}\n{content}")
 
@@ -147,32 +149,41 @@ def wikipedia_search(query: str) -> str:
 
 
 tools = [
-    add, subtract, multiply, divide, modulo, power, square_root,
-    web_search, wikipedia_search
+    add,
+    subtract,
+    multiply,
+    divide,
+    modulo,
+    power,
+    square_root,
+    web_search,
+    wikipedia_search,
 ]
 
 
 def get_llm():
     """Initialize the llm"""
     return ChatGoogleGenerativeAI(
-        model="gemini-2.5-pro",
-        temperature=0,
-        api_key=os.getenv("GEMINI_API_KEY")
+        model="gemini-2.5-flash", temperature=0, api_key=os.getenv("GEMINI_API_KEY")
     )
 
 
 def call_model(state: MessagesState):
     """Call the LLM with the current state.
-
+    
     Args:
         state: Current state containing messages
     """
     llm = get_llm()
     llm_with_tools = llm.bind_tools(tools)
-
-    messages = state['messages']
+    
+    messages = state["messages"]
+    
+    if not messages or not isinstance(messages[0], SystemMessage):
+        messages = [SystemMessage(content=SYSTEM_PROMPT)] + messages
+    
     response = llm_with_tools.invoke(messages)
-
+    
     return {"messages": [response]}