Update agent.py

agent.py

@@ -100,11 +100,18 @@ tools = [
 ]
 
 
-def build_graph(provider: str = "google"):
-  llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-preview-05-20", temperature=0, api_key=google_api_key)
-
-  llm_with_tools = llm.bind_tools(tools)
-
+def build_graph(provider: str):
+    if provider == "google":
+        # Google Gemini
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0,api_key=google_api_key)
+    elif provider == "huggingface":
+        # TODO: Add huggingface endpoint
+        llm = ChatHuggingFace(
+            llm=HuggingFaceEndpoint(
+                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
+                temperature=0,
+            ),
+        )
   def assistant(state: MessagesState):
       """ Use the tools to answer the query. you have add,subtract,multiply,divide,web_search,wikipedia_search,arxiv_search tools."""
       response = llm_with_tools.invoke([system_message]+state["messages"])
@@ -128,7 +135,7 @@ def build_graph(provider: str = "google"):
 if __name__ == "__main__":
     question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
     # Build the graph
-    graph = build_graph(provider="google")
+    graph = build_graph(provider="huggingface")
     # Run the graph
     messages = [HumanMessage(content=question)]
     messages = graph.invoke({"messages": messages})