Update agent.py

agent.py

@@ -152,25 +152,41 @@ tools = [
 ]
 
 # Build graph function
-def build_graph(provider: str = "openai"):
+def build_graph(provider: str = "huggingface", huggingface_model: str = "mistral"):
     """Build the graph"""
-    # Load environment variables from .env file
+
     if provider == "google":
         # Google Gemini
-        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    elif provider == "openai":
-        # OpenAI (e.g., GPT-4 or GPT-3.5 free)
-        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+        llm = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash",
+            temperature=0,
+            google_api_key=os.getenv("GOOGLE_API_KEY")
+        )
+
     elif provider == "huggingface":
-        # Hugging Face endpoint
+        # Choose between supported Hugging Face models
+        if huggingface_model == "mistral":
+            model_url = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
+        elif huggingface_model == "llama":
+            model_url = "https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf"
+        else:
+            raise ValueError("Unsupported Hugging Face model")
+
+        hf_token = os.getenv("HUGGINGFACE_API_TOKEN")
+        headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
+
         llm = ChatHuggingFace(
             llm=HuggingFaceEndpoint(
-                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
+                url=model_url,
                 temperature=0,
-            ),
+                headers=headers
+            )
         )
+
     else:
-        raise ValueError("Invalid provider. Choose 'google', 'openai' or 'huggingface'.")
+        raise ValueError("Invalid provider. Choose 'google' or 'huggingface'.")
+
+    return llm
         
     # Bind tools to LLM
     llm_with_tools = llm.bind_tools(tools)