Update agent.py

agent.py

@@ -136,6 +136,8 @@ with open("config.yaml", "r") as f:
     config = yaml.safe_load(f)
 
 provider = config["provider"]
+model_config = config["models"][provider]
+
 #prompt_path = config["system_prompt_path"]
 enabled_tool_names = config["tools"]
 
@@ -245,8 +247,6 @@ question_retriever_tool = create_retriever_tool(
 
 
 
-
-
 tools = [
     multiply,
     add,
@@ -258,27 +258,24 @@ tools = [
     arvix_search,
 ]
 
-# Build graph function
-def build_graph(provider: str = "google"):
-    """Build the graph"""
-    # Load environment variables from .env file
+
+def get_llm(provider: str, config: dict):
     if provider == "google":
-        # Google Gemini
-        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+        return ChatGoogleGenerativeAI(model=config["model"], temperature=config["temperature"])
     elif provider == "groq":
-        # Groq https://console.groq.com/docs/models
-        llm = ChatGroq(model="gemma2-9b-it", temperature=0) # optional : qwen-qwq-32b gemma2-9b-it
+        return ChatGroq(model=config["model"], temperature=config["temperature"])
     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,
-            ),
+        return ChatHuggingFace(
+            llm=HuggingFaceEndpoint(url=config["url"], temperature=config["temperature"])
         )
     else:
-        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
-    # Bind tools to LLM
+        raise ValueError(f"Invalid provider: {provider}")
+
+
+# Build graph function
+def build_graph():
+    """Build the graph based on provider"""
+    llm = get_llm(provider, model_config)
     llm_with_tools = llm.bind_tools(tools)
 
     # Node