Update langchain_logic/agent_setup.py

langchain_logic/agent_setup.py

@@ -1,29 +1,46 @@
-import os
+import os # <-- Import the os module
 from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_google_genai import ChatGoogleGenerativeAI
-# Import the new tool
-from langchain_logic.tools import schedule_appointment, search_for_appointments, delete_appointment_records, update_appointment_record, list_all_appointments
+from langchain_logic.tools import (
+    schedule_appointment, 
+    search_for_appointments, 
+    delete_appointment_records, 
+    update_appointment_record, 
+    list_all_appointments
+)
 from datetime import datetime
 
-google_api_key=os.getenv("GOOGLE_API_KEY")
+# --- IMPORTANT ---
+# This code will now work both locally (with a .env file) and on Hugging Face Spaces.
+# On Hugging Face, os.getenv("GOOGLE_API_KEY") will read the secret you set.
+# Locally, load_dotenv() in your main file will load it from .env for os.getenv to find.
+google_api_key = os.getenv("GOOGLE_API_KEY")
+
+if not google_api_key:
+    raise ValueError("GOOGLE_API_KEY not found. Please set it as a secret in Hugging Face Spaces or in a .env file for local development.")
+# ---------------
+
 
 def create_agent_executor():
     """Creates the LangChain agent and executor."""
-    # Add the new tool to the list
     tools = [
-        schedule_appointment,
-        search_for_appointments,
-        list_all_appointments, # <-- ADDED HERE
-        delete_appointment_records,
+        schedule_appointment, 
+        search_for_appointments, 
+        list_all_appointments, 
+        delete_appointment_records, 
         update_appointment_record
     ]
-    
 
-    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-lite-preview-06-17", temperature=0, google_api_key=google_api_key)
+    # Explicitly pass the API key to the model
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-2.5-flash-lite-preview-06-17",
+        temperature=0,
+        google_api_key=google_api_key # <-- Pass the key here
+    )
 
-    # You can optionally update the prompt to mention the new capability
     prompt = ChatPromptTemplate.from_messages([
+        # ... (your prompt remains the same) ...
         ("system", """You are an advanced, helpful appointment scheduling assistant.
 
         Your tasks are to schedule, search, update, and delete appointments using the available tools.
@@ -47,7 +64,7 @@ def create_agent_executor():
     agent_executor = AgentExecutor(
         agent=agent,
         tools=tools,
-        verbose=True, # For debugging
+        verbose=True,
         handle_parsing_errors=True
     )
     return agent_executor