main.py

from fastapi import FastAPI
from pydantic import BaseModel
import sqlite3
import os
from langchain_core.messages import HumanMessage, AIMessage
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langgraph.prebuilt import create_react_agent
from langchain_google_genai import ChatGoogleGenerativeAI
import shutil

# Copy auth.db to a writable path (only once)
AUTH_DB_PATH = "/tmp/auth.db"
if not os.path.exists(AUTH_DB_PATH):
    shutil.copy("Databases/auth.db", AUTH_DB_PATH)

# ========= LLM + DB Setup =========
os.environ["GOOGLE_API_KEY"] = "AIzaSyA8ue-NHZ_Fbak6UxoQWYizv-6JUcg7QbA"

llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
db = SQLDatabase.from_uri("sqlite:///Databases/categories_with_details.db")
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()

system_message = """
You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct {dialect} query to run,
then look at the results of the query and return the answer. Unless the user
specifies a specific number of examples they wish to obtain, always limit your
query to at most {top_k} results.

You can order the results by a relevant column to return the most interesting
examples in the database. Never query for all the columns from a specific table,
only ask for the relevant columns given the question.

You MUST double check your query before executing it. If you get an error while
executing a query, rewrite the query and try again.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the
database.

To start you should ALWAYS look at the tables in the database to see what you
can query. Do NOT skip this step.

Then you should query the schema of the most relevant tables.
""".format(dialect="SQLite", top_k=5)

agent_executor = create_react_agent(llm, tools, prompt=system_message)

# ========= FastAPI App =========
app = FastAPI()

# ========= Pydantic Models =========
class User(BaseModel):
    username: str
    password: str

class Query(BaseModel):
    question: str

# ========= Endpoints =========

@app.post("/signup")
def signup(user: User):
    conn = sqlite3.connect(AUTH_DB_PATH)
    cursor = conn.cursor()
    cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (user.username, user.password))
    conn.commit()
    conn.close()
    return {"message": "Signup successful"}

@app.post("/signin")
def signin(user: User):
    conn = sqlite3.connect(AUTH_DB_PATH)
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM users WHERE username = ? AND password = ?", (user.username, user.password))
    result = cursor.fetchone()
    conn.close()
    return {"message": "Login successful" if result else "Login failed"}

@app.post("/ask")
def ask_question(query: Query):
    response = agent_executor.invoke({"messages": [HumanMessage(content=query.question)]})
    for msg in reversed(response["messages"]):
        if isinstance(msg, AIMessage):
            return str(msg.content)
    return "No valid AI response"


# uvicorn main:app --reload