Project Completion Commit

.dockerignore

@@ -0,0 +1,6 @@
+.git/
+.gitignore
+.env
+.venv/
+__pycache__/
+*.pyc

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.py[cod]
+*$py.class
+.venv/
+.env/
+venv/.env
+.env

.python-version

@@ -0,0 +1 @@
+3.12

Dockerfile

@@ -0,0 +1,41 @@
+# 1. Base Image
+FROM python:3.11-slim
+
+# 2. Environment Variables for Hugging Face compatibility
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    HOME=/app \
+    PATH=/app/.local/bin:$PATH
+
+WORKDIR /app
+
+# 3. Install System Dependencies
+# libpq-dev is for PostgreSQL, curl is for Streamlit health checks
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    libpq-dev \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# 4. Install uv (The blazing fast package manager)
+RUN pip install uv
+
+# 5. Cache & Install Python Dependencies
+COPY pyproject.toml uv.lock ./
+RUN uv pip install --system -r pyproject.toml
+
+# 6. Copy your application code
+COPY . .
+
+# 7. Permissions: Make the script executable
+RUN chmod +x start.sh
+
+# 8. Permissions: Hugging Face runs as user 1000, not root!
+RUN chown -R 1000:1000 /app
+USER 1000
+
+# 9. Expose Ports (7860 for UI, 8000 for internal API)
+EXPOSE 7860 8000
+
+# 10. Start the application
+CMD ["./start.sh"]

app.py

@@ -0,0 +1,114 @@
+import streamlit as st
+import requests
+import uuid
+import time
+
+# --- CONFIGURATION ---
+API_URL = "http://localhost:8000" # Your FastAPI server URL
+
+st.set_page_config(page_title="Text@SQL Agent", page_icon="🤖", layout="centered")
+
+# --- SESSION STATE INITIALIZATION ---
+# This ensures variables survive when Streamlit re-renders the page
+if "thread_id" not in st.session_state:
+    st.session_state.thread_id = str(uuid.uuid4()) # Unique session ID for LangGraph memory
+if "user_id" not in st.session_state:
+    st.session_state.user_id = "tenant_" + str(uuid.uuid4())[:8]
+if "is_db_connected" not in st.session_state:
+    st.session_state.is_db_connected = False
+if "connection_url" not in st.session_state:
+    st.session_state.connection_url = ""
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+
+# --- SIDEBAR: DATABASE CONNECTION ---
+with st.sidebar:
+    st.header("⚙️ Database Setup")
+    
+    # If already connected, disable the input to enforce ONE database connection
+    db_input = st.text_input(
+        "Enter Database URL:",
+        disabled=st.session_state.is_db_connected
+    )
+    
+    if not st.session_state.is_db_connected:
+        if st.button("Connect & Initialize", type="primary", use_container_width=True):
+            if not db_input:
+                st.error("Please enter a valid URL.")
+            else:
+                with st.spinner("Building embeddings and initializing agent..."):
+                    try:
+                        # 1. Hit your FastAPI upload endpoint
+                        payload = {"connection_url": db_input, "user_id": st.session_state.user_id}
+                        response = requests.post(f"{API_URL}/upload_url", json=payload)
+                        
+                        if response.status_code == 200:
+                            # 2. Lock the connection and unlock the chat
+                            st.session_state.is_db_connected = True
+                            st.session_state.connection_url = db_input
+                            
+                            # Because your FastAPI upload uses BackgroundTasks, it returns instantly. 
+                            # We add a 2-second UI buffer here so the Qdrant embeddings have time to finish 
+                            # before the user fires off their first chat question.
+                            time.sleep(15) 
+                            
+                            st.success("Database connected securely!")
+                            st.rerun() # Refresh UI to unlock the chat window
+                        else:
+                            st.error(f"Failed to connect: {response.text}")
+                    except requests.exceptions.ConnectionError:
+                        st.error("🚨 Cannot connect to backend. Is FastAPI running?")
+    else:
+        st.success("✅ Connected to Database")
+        st.caption(f"URL: {st.session_state.connection_url}")
+        
+        # Add a reset button just in case they want to start completely over
+        if st.button("Disconnect & Reset", use_container_width=True):
+            st.session_state.clear()
+            st.rerun()
+
+# --- MAIN CHAT INTERFACE ---
+st.title("🗣️ Text2SQL Agent")
+
+# The Lock: Do not render the chat if DB is not connected
+if not st.session_state.is_db_connected:
+    st.info("👈 Please connect your database in the sidebar to begin analyzing data.")
+else:
+    # 1. Display previous chat messages from session state
+    for msg in st.session_state.chat_history:
+        with st.chat_message(msg["role"]):
+            st.markdown(msg["content"])
+
+    # 2. The Chat Input box
+    if user_query := st.chat_input("Ask a question about your data..."):
+        
+        # Immediately display the user's question in the UI
+        st.session_state.chat_history.append({"role": "user", "content": user_query})
+        with st.chat_message("user"):
+            st.markdown(user_query)
+
+        # 3. Call the LangGraph Backend
+        with st.chat_message("assistant"):
+            with st.spinner("Analyzing schema and generating SQL..."):
+                try:
+                    payload = {
+                        "message": user_query,
+                        "thread_id": st.session_state.thread_id,
+                        "user_id": st.session_state.user_id,
+                        "connection_url": st.session_state.connection_url
+                    }
+                    
+                    response = requests.post(f"{API_URL}/chat", json=payload)
+                    
+                    if response.status_code == 200:
+                        # Extract the final_result from your FastAPI JSON response
+                        answer = response.json().get("response", "No response found.")
+                        st.markdown(answer)
+                        
+                        # Save the assistant's answer to the UI history
+                        st.session_state.chat_history.append({"role": "assistant", "content": answer})
+                    else:
+                        st.error(f"Agent Error: {response.text}")
+                        
+                except requests.exceptions.ConnectionError:
+                    st.error("🚨 Connection dropped. Ensure FastAPI is running.")

pyproject.toml

@@ -0,0 +1,24 @@
+[project]
+name = "text-to-sql-agent"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "faker>=40.21.0",
+    "fastapi>=0.136.3",
+    "fastembed>=0.8.0",
+    "langchain>=1.3.4",
+    "langchain-community>=0.4.2",
+    "langchain-core>=1.4.0",
+    "langchain-openai>=1.2.2",
+    "langgraph>=1.2.4",
+    "langgraph-checkpoint-postgres>=3.1.0",
+    "langsmith>=0.8.8",
+    "psycopg-binary>=3.3.4",
+    "python-dotenv>=1.2.2",
+    "qdrant-client>=1.18.0",
+    "sqlalchemy>=2.0.50",
+    "streamlit>=1.58.0",
+    "pymysql",
+]

requirements.txt

@@ -0,0 +1,20 @@
+sqlalchemy
+langchain-core
+qdrant-client
+fastembed
+python-dotenv
+langchain
+langchain-classic
+langchain-community
+langgraph
+langchain-openai
+pydantic
+fastapi
+langgraph-checkpoint-postgres
+uvicorn
+python-multipart
+streamlit
+requests
+psycopg-pool
+langsmith
+pymysql

src/embedding.py

@@ -0,0 +1,71 @@
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams, SparseVectorParams, PointStruct
+from fastembed import TextEmbedding, SparseTextEmbedding
+import uuid
+from dotenv import load_dotenv
+import os
+from src.scheme import create_scheme
+
+COLLECTION_NAME = "Text2SQL"
+
+load_dotenv()
+
+qdrant_api = os.getenv("QDRANT_API_KEY")
+qdrant_url = os.getenv("QDRANT_URL")
+
+def create_embeddings(connection_url : str , user_id : str) :
+    client = QdrantClient(api_key=qdrant_api , url=qdrant_url)
+
+    dense_model = TextEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    sparse_model = SparseTextEmbedding(model_name="Qdrant/bm25")
+
+    if not client.collection_exists(COLLECTION_NAME) :
+        client.create_collection(collection_name=COLLECTION_NAME,
+                                 vectors_config={"dense": VectorParams(size=384, distance=Distance.COSINE)},
+                                 sparse_vectors_config={"sparse": SparseVectorParams()})
+        
+    try:
+        client.create_payload_index(
+            collection_name=COLLECTION_NAME,
+            field_name="user_id",
+            field_schema="keyword",
+        )
+    except Exception:
+        pass
+    
+    docs = create_scheme(connection_url)
+    text = [doc.page_content for doc in docs]
+
+    dense_vectors = list(dense_model.embed(text))
+    sparse_vectors = list(sparse_model.embed(text))
+
+    points = []
+
+    for i , doc in enumerate(docs) :
+        dense_vector = dense_vectors[i].tolist()
+
+        sparse_embeddings = sparse_vectors[i]
+
+        sparse_vector = {
+            'indices' : sparse_embeddings.indices.tolist(),
+            'values' : sparse_embeddings.values.tolist()
+        }
+
+        table_id = str(uuid.uuid4())
+
+        point = PointStruct(
+            id = table_id ,
+            vector = {
+                "dense" : dense_vector ,
+                "sparse" : sparse_vector
+            },
+            payload = {
+                'user_id' : user_id,
+                'text' : doc.page_content,
+                'table_name' : doc.metadata.get("table_name")
+            }
+        )
+
+        points.append(point)
+
+    client.upsert(collection_name=COLLECTION_NAME, points=points)

src/graph.py

@@ -0,0 +1,211 @@
+from typing import TypedDict , Annotated , List , Optional
+from langgraph.graph.message import add_messages 
+from langchain_core.messages import SystemMessage , HumanMessage
+from langchain_openai import ChatOpenAI
+from src.retrieval import retrieve
+import os
+from dotenv import load_dotenv
+from langgraph.graph import StateGraph, START ,END
+from pydantic import BaseModel , Field
+import datetime
+from langchain_community.utilities import SQLDatabase
+
+load_dotenv()
+
+class State(TypedDict) :
+    connection_url : str 
+    user_id : str
+    messages : Annotated[List , add_messages]
+    scheme : str
+    sql_query : str
+    query_result : str
+    error : Optional[str]
+    retry : int
+    final_result : str
+
+
+llm = ChatOpenAI(
+    model="openai/gpt-4o-mini",
+    openai_api_key=os.getenv("OPENROUTER_API_KEY"),
+    openai_api_base="https://openrouter.ai/api/v1",
+    temperature=0
+)
+
+class sql_query(BaseModel) :
+    generated_sql_query : str = Field(...,description="The raw, valid executable SQL query text. Contain absolutely NO markdown wrapping, code blocks, or conversational formatting.")
+
+def retrieve_node(state : State) :
+    messages = state.get("messages")
+    db_url = state.get("connection_url")
+    user_id = state.get("user_id")
+
+    query = messages[-1].content
+
+    scheme = retrieve(user_id , query , db_url)
+
+    return {'scheme' : scheme}
+
+def generate_node(state : State) :
+    messages = state.get("messages")
+    scheme = state.get("scheme")
+    error = state.get("error")
+    wrong_query = state.get('sql_query')
+
+    llm_with_structured_output = llm.with_structured_output(sql_query)
+
+    history_messages = messages[:-1]
+    current_query_string = messages[-1].content
+
+    current_date = datetime.datetime.now().strftime("%Y-%m-%d")
+
+    if history_messages:
+        history_text = "\n".join([
+            f"{msg.type.capitalize()}: {msg.content}" 
+            for msg in history_messages
+        ])
+    else:
+        history_text = "This is the first user request. No history exists."
+
+    if error and wrong_query :
+        error_context = f"""
+=== 🚨 ERROR CORRECTION MODE 🚨 ===
+Your previous attempt to answer this request failed.
+[PREVIOUS BROKEN QUERY]: 
+{wrong_query}
+
+[DATABASE ERROR MESSAGE]: 
+{error}
+
+INSTRUCTION: Analyze the error message and the schema carefully. Fix the syntax, column names, or logic, and generate a CORRECTED query.
+"""
+    else :
+        error_context = ""
+
+        system_prompt = SystemMessage(content=f"""You are an expert Data Analyst and Database Engineer. 
+Your job is to write highly optimized, perfectly accurate database queries based on user requests.
+
+=== DATABASE SCHEMA & DIALECT ===
+Look at the metadata below to identify the targeted database engine dialect and table layout:
+{scheme}
+
+=== CONVERSATION HISTORY ===
+Use this previous context to resolve ambiguous terms (e.g., if the user says "filter those by...", look here to see what "those" refers to):
+{history_text}
+{error_context}
+
+=== CRITICAL RULES ===
+1. ALIGNMENT: Only use the tables and columns provided in the schema above. Do not hallucinate column names.
+2. DIALECT MATCHING: Look at the 'Dialect:' specified above and write strict queries matching that exact syntax. 
+3. JOINS: Pay close attention to the FOREIGN KEY constraints provided in the schema to perform accurate JOINs.
+4. CURRENT DATE: Today's date is {current_date}. Use this exact date for any relative time filters (e.g., "last month", "this year").
+5. CASE SENSITIVITY: When filtering by strings, use case-insensitive comparisons (e.g., LOWER(column) = LOWER('value')) unless instructed otherwise.
+6. SECURITY: NEVER generate DML queries (INSERT, UPDATE, DELETE, DROP). Only generate SELECT statements.
+
+=== OUTPUT SELECTION RULES ===
+1. If the user asks WHO / WHICH / WHAT IS THE NAME / identify a person, customer, user, product, company, or entity, return the human-readable name field, not just the ID.
+2. If the schema has both an ID column and a name column, prefer selecting the name column in the final output.
+3. If the name is in another table, use the required JOIN to fetch it.
+4. Only return an ID alone when the user explicitly asks for the ID, or when no name-like field exists in the schema.
+5. For count/number questions, return an aggregate numeric result, not a list of rows.
+6. For "who/which" questions, do not answer with only identifiers if a readable label exists in the schema.
+
+=== INSTRUCTIONS ===
+First, think through the necessary tables, filters, joins, and the exact type of answer expected.
+Then, provide the final executable SQL query specifically for the LATEST USER REQUEST.""")
+    
+    final_msg = [
+        system_prompt,
+        HumanMessage(content=f"LATEST USER REQUEST:\n{current_query_string}")
+    ]
+
+    response = llm_with_structured_output.invoke(final_msg)
+
+    return {'sql_query' : response.generated_sql_query , "error" : None}
+
+def execute_node(state : State) :
+    url = state.get("connection_url")
+    sql_query = state.get("sql_query")
+    retry = state.get("retry" , 0)
+
+    try :
+        db = SQLDatabase.from_uri(url)
+
+        result = db.run(sql_query)
+
+        return {"query_result" : result , "error" : None , "retry" : 0}
+    
+    except Exception as e :
+        return {'error' : str(e) , "retry" : retry+1}
+    
+def routing(state : State) :
+    error = state.get("error")
+    retry = state.get('retry' , 0)
+
+    if error and retry<3 :
+        return "generate_node"
+    else :
+        return "answer_node"
+    
+def answer_node(state : State) :
+    messages = state.get("messages")
+    query_result = state.get("query_result" , "No records found.")
+    error = state.get("error")
+
+    history_messages = messages[:-1]
+    user_query = messages[-1].content
+
+    if history_messages:
+        history_text = "\n".join([
+            f"{msg.type.capitalize()}: {msg.content}" 
+            for msg in history_messages
+        ])
+    else:
+        history_text = "This is the first user request. No history exists."
+
+    system_prompt = f"""You are a helpful Data Analyst communicating directly with a user.
+
+=== CONVERSATION HISTORY ===
+Use this to maintain the context and tone of the conversation:
+{history_text}
+
+=== EXECUTION CONTEXT ===\n"""
+
+    if error:
+        system_prompt += f"""Unfortunately, the database returned an error and the data could not be retrieved. 
+Error details: {error}
+INSTRUCTION: Politely apologize to the user and briefly explain that you encountered a technical issue retrieving their specific request."""
+    else:
+        system_prompt += f"""The database returned this raw data: {query_result}
+
+INSTRUCTIONS:
+1. Answer using ONLY the returned data.
+2. Never invent a name, value, or entity that is not present in the result.
+3. If the result contains both an ID and a name, use the name in the final answer and mention the ID only if helpful.
+4. If the result contains only an ID and the user asked for a name/person/entity, say that the returned data only contains an identifier and no readable name.
+5. Do not substitute or guess a name from a customer_id or any other identifier.
+6. Do not mention SQL, the database, schemas, or how you got the data.
+7. Give a clean, professional, and conversational response."""
+        
+    final_msg = [
+        SystemMessage(content=system_prompt),
+        HumanMessage(content=f"LATEST USER REQUEST:\n{user_query}")
+    ]
+    
+    response = llm.invoke(final_msg)
+
+    return {"messages": [response], "final_result": response.content}
+
+workflow = StateGraph(State)
+
+workflow.add_node("retrieve_node" , retrieve_node)
+workflow.add_node("generate_node" , generate_node)
+workflow.add_node("execute_node" , execute_node)
+workflow.add_node("answer_node" , answer_node)
+
+workflow.add_edge(START , "retrieve_node")
+workflow.add_edge("retrieve_node" , "generate_node")
+workflow.add_edge("generate_node" , "execute_node")
+workflow.add_conditional_edges("execute_node" , routing , {
+    "answer_node" : "answer_node" , "generate_node" : "generate_node"
+})
+workflow.add_edge("answer_node" , END)

src/main.py

@@ -0,0 +1,73 @@
+from fastapi import FastAPI , HTTPException , BackgroundTasks
+from src.embedding import create_embeddings
+from src.graph import workflow
+from pydantic import BaseModel , Field
+from langgraph.checkpoint.postgres import PostgresSaver
+from langchain_core.messages import HumanMessage
+import os 
+
+app = FastAPI(
+    title="Text2SQL Agent API",
+    description="A production-grade backend powering LangGraph agent.",
+    version="1.0.0"
+)
+
+class UploadRequest(BaseModel):
+    connection_url: str = Field(..., description="Database URL")
+    user_id: str = Field(..., description="The unique identifier for the tenant context.")
+
+class ChatRequest(BaseModel) :
+    message : str = Field(...,description="Input message by the user.")
+    thread_id : str = Field(...,description="Unique session ID to maintain short term memory.")
+    user_id : str = Field(...,description="The unique identifier for the tenant context.")
+    connection_url : str = Field(...,description="Database URL")
+
+@app.post("/upload_url" , summary="Recieve database URL and invoke ingestion pipeline.")
+def upload(request : UploadRequest , background_tasks : BackgroundTasks) :
+    background_tasks.add_task(create_embeddings , request.connection_url , request.user_id)
+
+    return {
+        "status" : "success",
+        "message" : "Ingestion Pipeline started !"
+    }
+
+@app.post("/chat" , summary="Return the response generated by the agent for the given user query.")
+def chat_endpoint(request : ChatRequest) :
+    db_uri = os.getenv("DATABASE_URI")
+
+    with PostgresSaver.from_conn_string(db_uri) as checkpointer:
+        checkpointer.setup()
+
+        agent = workflow.compile(
+        checkpointer=checkpointer
+        )
+        config = {
+            "configurable" : {
+                'thread_id' : request.thread_id
+            }
+        }
+
+        initial_state = {
+            'connection_url' : request.connection_url ,
+            'user_id' : request.user_id ,
+            'messages' : [HumanMessage(content=request.message)],
+            'retry' : 0
+        }
+        try :
+            result = agent.invoke(initial_state , config=config)
+
+            final_result = result.get("final_result")
+
+            print("*"*50 , flush=True)
+            print(f"\n\n Scheme : {result['scheme']}\n\n" , flush=True)
+            print(f"\n\nSql Query : {result['sql_query']}\n\n" , flush=True)
+            print(f"\n\nQuery Result : {result['query_result']}\n\n" , flush=True)
+
+            return {
+                    "status": "success",
+                    "thread_id": request.thread_id,
+                    "response": final_result
+                }
+
+        except Exception as e :
+            raise HTTPException(status_code=500 , detail=f"Error : {str(e)}")

src/retrieval.py

@@ -0,0 +1,70 @@
+import os
+from dotenv import load_dotenv
+from qdrant_client import QdrantClient
+from qdrant_client import models
+from fastembed import TextEmbedding, SparseTextEmbedding
+from langchain_community.utilities import SQLDatabase
+
+load_dotenv()
+
+qdrant_api = os.getenv("QDRANT_API_KEY")
+qdrant_url = os.getenv("QDRANT_URL")
+
+COLLECTION_NAME = "Text2SQL"
+
+def retrieve(user_id : str , query : str , connection_url: str) :
+
+    client = QdrantClient(api_key=qdrant_api , url=qdrant_url)
+
+    dense_model = TextEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    sparse_model = SparseTextEmbedding(model_name="Qdrant/bm25")
+
+    dense_query_vector = list(dense_model.embed([query]))[0]
+
+    sparse_query = list(sparse_model.embed([query]))[0]
+
+    sparse_query_vector = models.SparseVector(indices=sparse_query.indices,
+                                              values=sparse_query.values)
+    
+    user_filter = models.Filter(
+            must=[
+                models.FieldCondition(
+                    key="user_id",
+                    match=models.MatchValue(value=user_id)
+                )
+            ]
+        )
+    
+    results = client.query_points(
+            collection_name=COLLECTION_NAME,
+            prefetch=[
+                models.Prefetch(
+                    query=dense_query_vector,
+                    limit=10,
+                    using="dense",
+                    filter=user_filter
+                ),
+                models.Prefetch(
+                    query=sparse_query_vector,
+                    using="sparse",
+                    limit=10,
+                    filter=user_filter
+                )
+            ],
+            query=models.FusionQuery(fusion=models.Fusion.RRF),
+            limit=10
+        )
+    
+    tables = []
+    for point in results.points :
+        table = point.payload['table_name']
+        if table not in tables :
+            tables.append(table)
+
+    db = SQLDatabase.from_uri(connection_url , sample_rows_in_table_info=0)
+
+    dialect = db.dialect
+    
+    final_schemes = f"Dialect : {dialect}\n {db.get_table_info(table_names=tables)}"
+    
+    return final_schemes

src/scheme.py

@@ -0,0 +1,35 @@
+from sqlalchemy import create_engine , inspect
+from langchain_core.documents import Document
+
+def create_scheme(database_url : str) -> list[dict] :
+
+    scheme_docs = []
+
+    engine = create_engine(database_url)
+    inspector = inspect(engine)
+
+    tables = inspector.get_table_names()
+
+    for table in tables :
+        clean_table = table.replace("_" , " ")
+
+        columns = inspector.get_columns(table)
+
+        clean_columns = []
+
+        for col in columns :
+            clean_col = col['name'].replace("_" , " ")
+            clean_columns.append(clean_col)
+        
+        doc = f"Table: {clean_table}.\nColumns: {', '.join(clean_columns)}"
+
+        scheme_docs.append(
+            Document(
+                page_content=doc,
+                metadata={
+                    "table_name" : table
+                }
+            )
+        )
+    
+    return scheme_docs

start.sh

@@ -0,0 +1,7 @@
+echo "Starting FastAPI Backend..."
+uv run uvicorn src.main:app --host 0.0.0.0 --port 8000 &
+
+sleep 3
+
+echo "Starting Streamlit Frontend..."
+uv run streamlit run app.py --server.port=7860 --server.address=0.0.0.0