sqlAgent.py

from langchain.sql_database import SQLDatabase
from langchain.agents import create_sql_agent, AgentExecutor
from langgraph.graph.message import add_messages
from langchain_groq import ChatGroq
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
import os
import logging
from dotenv import load_dotenv
from pathlib import Path

# ── Path resolution — works for HF Spaces (/app/), gunicorn, and local dev ───
BASE_DIR = Path(__file__).resolve().parent

# ── Load .env ONLY in local development ───────────────────────────────────────
# In HF Spaces, secrets are injected as env vars automatically.
# load_dotenv() will simply do nothing if .env doesn't exist — that's fine.
dotenv_path = BASE_DIR / ".env"
load_dotenv(dotenv_path=dotenv_path, override=False)
# override=False means HF-injected env vars take priority over .env

# ── Validate key ──────────────────────────────────────────────────────────────
groq_key = os.getenv("GROQ_API_KEY")
if not groq_key:
    raise EnvironmentError(
        "GROQ_API_KEY not set.\n"
        "  → On HF Spaces: go to Settings → Variables and secrets → add GROQ_API_KEY\n"
        "  → Locally: add GROQ_API_KEY=your_key to your .env file"
    )

# ── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")
logger = logging.getLogger(__name__)
logger.info(f"BASE_DIR : {BASE_DIR}")

# ── Database — absolute path so it works wherever the process runs ────────────
db_path = BASE_DIR / "customer_orders.db"
if not db_path.exists():
    raise FileNotFoundError(
        f"Database not found at: {db_path}\n"
        f"Files in BASE_DIR: {list(BASE_DIR.iterdir())}"  # helps debug on HF
    )

db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
logger.info(f"Database loaded: {db_path}")


# LLM Setup
llm = ChatGroq(
    model="meta-llama/llama-4-scout-17b-16e-instruct",
    temperature=0,
    max_retries=2,
    groq_api_key = os.getenv("GROQ_API_KEY")
)

_AGENT_SYSTEM_PROMPT = """
You are OrderBot, a helpful and polite AI customer-support assistant for a food-delivery company.

Rules you MUST follow:
1. You can ONLY answer questions about orders. Do NOT answer anything unrelated to orders.
2. You ONLY query the database for the specific order IDs the customer mentions.
   Never return data for all orders or for other customers.
3. If the customer asks about cancellation:
   - "delivered"        → politely say it cannot be cancelled; suggest a return.
   - "out for delivery" → escalate to a human agent.
   - "placed/preparing" → confirm cancellation and mention a 3-5 day refund.
4. If the customer asks where their order is, provide the order_status and ETA fields.
5. Always be concise, warm, and professional.
6. Never expose raw SQL, database errors, or internal system details.
7. If you cannot find an order, politely ask the customer to double-check the order ID.
8. If you cannot find a customer, politely ask the customer to double-check the customer ID.
"""


def build_sql_agent() -> AgentExecutor:
    """Construct and return a LangChain SQL agent."""
    toolkit = SQLDatabaseToolkit(db=db, llm=llm)
    tools = toolkit.get_tools()


#Create the SQL agent with the system message

    agent = create_sql_agent(
        llm=llm,
        toolkit=toolkit,
        verbose=True,
        return_intermediate_steps=True,
        system_message=_AGENT_SYSTEM_PROMPT,
        max_iterations=6,
        handle_parsing_errors=True,
    )

    logger.info("SQL agent built successfully")
    return agent