test_chat.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure LLM options\n",
    "#from langchain_community.chat_models import ChatOllama\n",
    "from langchain_ollama import ChatOllama\n",
    "from langchain_community.llms import DeepInfra\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "api = os.getenv(\"DEEPINFRA_API_KEY\")\n",
    "# API setup (optional)\n",
    "deepinfra_model = \"meta-llama/Llama-3.3-70B-Instruct\"\n",
    "# \"https://api.python.langchain.com/en/latest/llms/langchain_community.llms.deepinfra.DeepInfra.html\"\n",
    "\n",
    "di = DeepInfra(model_id=deepinfra_model,\n",
    "                    deepinfra_api_token=api)\n",
    "\n",
    "# Choose which model to use\n",
    "llm = di \n",
    "\n",
    "# Initialize database connection\n",
    "db = SQLDatabase.from_uri(\"sqlite:///Spring_2025_courses.db\", sample_rows_in_table_info=0)\n",
    "\n",
    "# Helper functions\n",
    "def get_schema(_):\n",
    "    \"\"\"Retrieve database schema\"\"\"\n",
    "    return db.get_table_info()\n",
    "\n",
    "def run_query(query):\n",
    "    \"\"\"Execute SQL query\"\"\"\n",
    "    return db.run(query)\n",
    "\n",
    "# print(get_schema(_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "# SQL generation prompt\n",
    "system_prompt = \"\"\" \n",
    "You are a SQLite expert. Given an input question, create one syntactically correct SQLite query to run. Generate only one query. No pre-amble.\n",
    "\n",
    "Here is the relevant table information:\n",
    "schema: {schema}\n",
    "Departments: 'Applied Math & Statistics', 'ART', 'Biomedical Science (BMS)', 'Computer Science (CS), Computer Networks and Cybersecurity (CNCS)', 'Data Science (DS)', 'Dance', 'General Education (GE)'\n",
    "New/Repeat: Have no values \n",
    "\n",
    "Tips:\n",
    "- Use LIKE instead of = in the queries\n",
    "- Don't use New/Repeat in queries\n",
    "\n",
    "Below are a number of examples of questions and their corresponding SQL queries.\n",
    "\n",
    "example 1:\n",
    "Question:List all departments.\n",
    "SQL query:SELECT Department FROM Spring_2025_courses;\n",
    "\n",
    "example 2:\n",
    "Question:Find all classes that are taught by Dr. Qu.\n",
    "SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Instructor LIKE '%Qu%';\n",
    "\n",
    "example 3:\n",
    "Question:List all courses in the Data Science department.\n",
    "SQL query:SELECT CourseCode, CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Data Science%';\n",
    "\n",
    "example 4:\n",
    "Question:Find the total number of courses offered for Spring 2025.\n",
    "SQL query:SELECT COUNT(CourseCode) FROM Spring_2025_courses;\n",
    "\n",
    "example 5:\n",
    "Question:List all professors from Biomed Department.\n",
    "SQL query:SELECT Instructor FROM Spring_2025_courses WHERE Department LIKE '%Biomed%';\n",
    "\n",
    "example 6:\n",
    "Question:Which courses are available in the Biomed department for Spring 2025?\n",
    "SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Biomed%';\n",
    "\n",
    "example 7:\n",
    "Question:How many courses are offered under the Arts Management (BFA) program?\n",
    "SQL query:SELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Department LIKE Program LIKE '%BFA%'\n",
    "\n",
    "example 8:\n",
    "Question:Which courses are available for the Master of Science in Applied Math?\n",
    "SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Program LIKE '%Master of Science%' AND Department LIKE '%Applied Math%';\n",
    "\n",
    "Write only one SQLite query that would answer the user's question. No pre-amble.\n",
    "\"\"\"\n",
    "\n",
    "human_prompt = \"\"\"Based on the table schema below, write a SQL query that would answer the user's question:\n",
    "{schema}\n",
    "\n",
    "Question: {question}\n",
    "SQL Query:\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", system_prompt),\n",
    "    (\"human\", human_prompt),\n",
    "])\n",
    "\n",
    "# Build query generation chain\n",
    "sql_generator = (\n",
    "    RunnablePassthrough.assign(schema=get_schema)\n",
    "    | prompt\n",
    "    | llm.bind(stop=[\"\\nSQLResult:\"])\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "# Natural language response generation\n",
    "nl_system_prompt = \"\"\"Given an input question and SQL response, convert it to a natural language answer. \n",
    "\n",
    "Below are a number of examples of questions and their corresponding SQL queries.\n",
    "\n",
    "example 1: \n",
    "Question: Which courses are available in the Data Science department?\n",
    "SQL Query: SELECT CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Data Science%';\n",
    "SQL Response: [('Introduction to Data Science',), ('Career Development in Data Science',), ('Data Structures',), ('Data Visualization',), ('Data Inference',), ('Data Mining',), ('Big Data Engineering',), ('Independent Study for Data Science',), ('Senior Project',), ('Mathematical Foundation for Data Science (Online)',), ('Computational Foundation for Data Science (Online)',), ('Probability for Data Science (Online)',), ('Exploratory Data Analysis and Visualization',), ('Modern Applied Statistical Learning',), ('Data Mining for Business',), ('Big Data and Data Engineering',), ('Design and Analysis of Experiments',), ('Design and Analysis of Experiments',), ('Computer Vision and Natural Language Processing (Independent Study)',), ('Capstone Project',)]\n",
    "Response: The available courses in the Data Science department are Introduction to Data Science, Career Development in Data Science, Data Structures, Data Visualization, Data Inference, Data Mining, Big Data Engineering, Independent Study for Data Science, Senior Project, Mathematical Foundation for Data Science, Computational Foundation for Data Science, Probability for Data Science, Exploratory Data Analysis and Visualization, Modern Applied Statistical Learning, Data Mining for Business, Big Data and Data Engineering, Design and Analysis of Experiments, and Capstone Project.\n",
    "\n",
    "example 2:\n",
    "Question: Find all classes that are taught by Dr. Qu.\n",
    "SQL Query: SELECT CourseTitle FROM Spring_2025_courses WHERE Instructor LIKE '%Qu%';\n",
    "SQL Response: [('Introduction to Data Science',), ('Data Structures',), ('Data Mining',), ('Computational Foundation for Data Science (Online)',), ('Data Mining for Business',), ('Computer Vision and Natural Language Processing (Independent Study)',)]\n",
    "Response:  Dr. Qu is teaching Introduction to Data Science, Data Structures, Data Mining, Computational Foundation for Data Science, Data Mining for Business, and Computer Vision and Natural Language Processing.\n",
    "\n",
    "Generate only ONE answer based on the SQL response. No pre-amble.\n",
    "\"\"\"\n",
    "\n",
    "response_prompt = \"\"\"Based on the table schema below, question, sql query, and sql response, write a natural language response:\n",
    "{schema}\n",
    "\n",
    "Question:{question}\n",
    "SQL Query:{query}\n",
    "SQL Response:{response}\n",
    "Response:\"\"\"\n",
    "\n",
    "nl_prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", nl_system_prompt),\n",
    "    (\"human\", response_prompt),\n",
    "])\n",
    "\n",
    "# Complete chain with natural language output\n",
    "complete_chain = (\n",
    "    RunnablePassthrough.assign(query=sql_generator)\n",
    "    | RunnablePassthrough.assign(\n",
    "        schema=get_schema,\n",
    "        response=lambda x: db.run(x['query']),\n",
    "    )\n",
    "    | nl_prompt\n",
    "    | llm\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def demonstrate_chain_steps(question: str):\n",
    "    \"\"\"\n",
    "    Demonstrate each step of the complete chain\n",
    "    \n",
    "    Args:\n",
    "        question (str): Natural language question\n",
    "    \"\"\"\n",
    "    print(\"\\n=== Chain Demonstration ===\\n\")\n",
    "    \n",
    "    # Step 1: Initial input\n",
    "    print(\"1. Initial Input:\")\n",
    "    print({\"question\": question})\n",
    "    \n",
    "    # Step 2: Generate SQL\n",
    "    sql = sql_generator.invoke({\"question\": question})\n",
    "    print(\"\\n2. Generated SQL:\")\n",
    "    print(sql)\n",
    "    \n",
    "    # Step 3: Execute SQL\n",
    "    db_result = db.run(sql)\n",
    "    print(\"\\n3. Database Result:\")\n",
    "    print(db_result)\n",
    "    \n",
    "    # Step 4: Get schema\n",
    "    schema = get_schema(None)\n",
    "    print(\"\\n4. Database Schema:\")\n",
    "    print(schema[:200] + \"...\")  # Truncated for readability\n",
    "    \n",
    "    # Step 5: Complete response\n",
    "    final_response = complete_chain.invoke({\"question\": question})\n",
    "    print(\"\\n5. Final Natural Language Response:\")\n",
    "    print(final_response)\n",
    "    \n",
    "    print(\"\\n=== End Demonstration ===\\n\")\n",
    "\n",
    "# Example usage\n",
    "demonstrate_chain_steps(\"Find all classes that are taught by Dr. Bess.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain_core.prompts import MessagesPlaceholder\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "# SQL generation prompt\n",
    "system_prompt = \"\"\" \n",
    "You are a SQLite expert. Given an input question, create one syntactically correct SQLite query to run. Generate only one query. No preamble.\n",
    "\n",
    "Here is the relevant table information:\n",
    "schema: {schema}\n",
    "Departments: 'Applied Math & Statistics', 'ART', 'Biomedical Science (BMS)', 'Computer Science (CS), Computer Networks and Cybersecurity (CNCS)', 'Data Science (DS)', 'Dance', 'General Education (GE)'\n",
    "New/Repeat: Have no values (Don't use New/Repeat in queries)\n",
    "\n",
    "Tips:\n",
    "- Use LIKE instead of = in the queries\n",
    "\n",
    "Below are a number of examples of questions and their corresponding SQL queries.\n",
    "\n",
    "example 1:\n",
    "Question: List all departments.\n",
    "SQL query: SELECT Department FROM Spring_2025_courses;\n",
    "\n",
    "example 2:\n",
    "Question: Find all classes that are taught by Dr. Qu.\n",
    "SQL query: SELECT CourseTitle FROM Spring_2025_courses WHERE Instructor LIKE '%Qu%';\n",
    "\n",
    "example 3:\n",
    "Question: List all courses in the Data Science department.\n",
    "SQL query: SELECT CourseCode, CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Data Science%';\n",
    "\n",
    "example 4:\n",
    "Question: Find the total number of courses offered for Spring 2025.\n",
    "SQL query: SELECT COUNT(CourseCode) FROM Spring_2025_courses;\n",
    "\n",
    "example 5:\n",
    "Question: List all professors from Biomed Department.\n",
    "SQL query: SELECT Instructor FROM Spring_2025_courses WHERE Department LIKE '%Biomed%';\n",
    "\n",
    "Write only one SQLite query that would answer the user's question. No preamble.\n",
    "\"\"\"\n",
    "\n",
    "# Initialize memory\n",
    "memory = ConversationBufferMemory(return_messages=True)\n",
    "\n",
    "# Updated prompt with memory\n",
    "memory_prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", \"Convert questions to SQL using schema: {schema}\"),\n",
    "    MessagesPlaceholder(variable_name=\"history\"),\n",
    "    (\"human\", \"{question}\"),\n",
    "])\n",
    "\n",
    "# Memory-enabled query chain\n",
    "def save_context(input_output):\n",
    "    \"\"\"Save conversation context\"\"\"\n",
    "    output = {\"output\": input_output.pop(\"output\")}\n",
    "    memory.save_context(input_output, output)\n",
    "    return output[\"output\"]\n",
    "\n",
    "sql_chain_with_memory = (\n",
    "    RunnablePassthrough.assign(\n",
    "        schema=get_schema,\n",
    "        history=lambda x: memory.load_memory_variables(x)[\"history\"],\n",
    "    )\n",
    "    | memory_prompt\n",
    "    | llm.bind(stop=[\"\\nSQLResult:\"])\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "# Final chain with memory\n",
    "sql_memory_chain = RunnablePassthrough.assign(output=sql_chain_with_memory) | save_context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "from langchain_core.messages import HumanMessage\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import START, MessagesState, StateGraph\n",
    "\n",
    "# SQL generation prompt\n",
    "system_prompt = \"\"\" \n",
    "You are a SQLite expert. Given an input question, create one syntactically correct SQLite query to run. Generate only one query. No pre-amble.\n",
    "\n",
    "Here is the relevant table information:\n",
    "schema: {schema}\n",
    "Departments: 'Applied Math & Statistics', 'ART', 'Biomedical Science (BMS)', 'Computer Science (CS), Computer Networks and Cybersecurity (CNCS)', 'Data Science (DS)', 'Dance', 'General Education (GE)'\n",
    "New/Repeat: Have no values (Don't use New/Repeat in queries)\n",
    "\n",
    "Tips:\n",
    "- Use LIKE instead of = in the queries\n",
    "\n",
    "Below are a number of examples of questions and their corresponding SQL queries.\n",
    "\n",
    "example 1:\n",
    "Question:List all departments.\n",
    "SQL query:SELECT Department FROM Spring_2025_courses;\n",
    "\n",
    "example 2:\n",
    "Question:Find all classes that are taught by Dr. Qu.\n",
    "SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Instructor LIKE '%Qu%';\n",
    "\n",
    "example 3:\n",
    "Question:List all courses in the Data Science department.\n",
    "SQL query:SELECT CourseCode, CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Data Science%';\n",
    "\n",
    "example 4:\n",
    "Question:Find the total number of courses offered for Spring 2025.\n",
    "SQL query:SELECT COUNT(CourseCode) FROM Spring_2025_courses;\n",
    "\n",
    "example 5:\n",
    "Question:List all professors from Biomed Department.\n",
    "SQL query:SELECT Instructor FROM Spring_2025_courses WHERE Department LIKE '%Biomed%';\n",
    "\n",
    "example 6:\n",
    "Question:Which courses are available in the Biomed department for Spring 2025?\n",
    "SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Biomed%';\n",
    "\n",
    "example 7:\n",
    "Question:How many courses are offered under the Arts Management (BFA) program?\n",
    "SQL query:SELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Department LIKE Program LIKE '%BFA%'\n",
    "\n",
    "example 8:\n",
    "Question:Which courses are available for the Master of Science in Applied Math?\n",
    "SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Program LIKE '%Master of Science%' AND Department LIKE '%Applied Math%';\n",
    "\n",
    "Write only one SQLite query that would answer the user's question. No pre-amble.\n",
    "\"\"\"\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(state_schema=MessagesState)\n",
    "\n",
    "# Memory initialization\n",
    "memory = ConversationBufferMemory(return_messages=True)\n",
    "\n",
    "# Define the function that calls the model\n",
    "def call_model(state: MessagesState):\n",
    "    response = llm.invoke(state[\"messages\"])\n",
    "    return {\"messages\": response}\n",
    "\n",
    "# Define the two nodes we will cycle between\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_node(\"model\", call_model)\n",
    "\n",
    "# Adding memory is straightforward in langgraph\n",
    "memory_saver = MemorySaver()\n",
    "\n",
    "app = workflow.compile(checkpointer=memory_saver)\n",
    "\n",
    "# The thread id is a unique key that identifies this particular conversation.\n",
    "thread_id = uuid.uuid4()\n",
    "config = {\"configurable\": {\"thread_id\": thread_id}}\n",
    "\n",
    "# Updated prompt with memory\n",
    "memory_prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\", system_prompt),\n",
    "    MessagesPlaceholder(variable_name=\"history\"),\n",
    "    (\"human\", \"{question}\"),\n",
    "])\n",
    "\n",
    "# Memory-enabled query chain\n",
    "def save_context(input_output):\n",
    "    \"\"\"Save conversation context\"\"\"\n",
    "    output = {\"output\": input_output.pop(\"output\")}\n",
    "    memory.save_context(input_output, output)\n",
    "    return output[\"output\"]\n",
    "\n",
    "sql_chain_with_memory = (\n",
    "    RunnablePassthrough.assign(\n",
    "        schema=get_schema,  # Placeholder function to provide schema\n",
    "        history=lambda x: memory.load_memory_variables(x)[\"history\"],\n",
    "    )\n",
    "    | memory_prompt\n",
    "    | llm.bind(stop=[\"\\nSQLResult:\"])\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "# Final chain with memory\n",
    "sql_memory_chain = RunnablePassthrough.assign(output=sql_chain_with_memory) | save_context\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Example conversation\n",
    "# input_message = HumanMessage(content=\"Which courses are available for the Bachelor of Science in Data Science?\")\n",
    "# for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
    "#     event[\"messages\"][-1].pretty_print()\n",
    "\n",
    "# # Confirming memory functionality\n",
    "# input_message = HumanMessage(content=\"Who is teaching the first course on the list?\")\n",
    "# for event in app.stream({\"messages\": [input_message]}, config, stream_mode=\"values\"):\n",
    "#     event[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_query_result(question, include_sql=False):\n",
    "    \"\"\"\n",
    "    Display query results in a formatted way\n",
    "    \n",
    "    Args:\n",
    "        question (str): Natural language question\n",
    "        include_sql (bool): Whether to show the generated SQL\n",
    "    \"\"\"\n",
    "    # Generate SQL and get response\n",
    "    sql = sql_generator.invoke({\"question\": question})\n",
    "    response = complete_chain.invoke({\"question\": question})\n",
    "    \n",
    "    # Print formatted output\n",
    "    print(\"\\n\" + \"=\"*50)\n",
    "    print(f\"Question: {question}\")\n",
    "    if include_sql:\n",
    "        print(f\"\\nGenerated SQL: {sql}\")\n",
    "    print(f\"\\nAnswer: {response}\")\n",
    "    print(\"=\"*50 + \"\\n\")\n",
    "\n",
    "# Example usage\n",
    "display_query_result(\"Which courses are available in the Data Science department for Spring 2025?\", include_sql=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the chain\n",
    "def display_conversation():\n",
    "    \"\"\"Display the current conversation history.\"\"\"\n",
    "    history = memory.load_memory_variables({})[\"history\"]\n",
    "    print(\"\\nConversation History:\")\n",
    "    print(\"=\" * 50)\n",
    "    for msg in history:\n",
    "        role = \"Human\" if msg.type == \"human\" else \"Assistant\"\n",
    "        print(f\"{role}: {msg.content}\\n\")\n",
    "\n",
    "\n",
    "# Test input and output\n",
    "response1 = sql_memory_chain.invoke({\"question\": \"Which courses are available for the Bachelor of Science in Data Science?\"})\n",
    "print(\"First Response:\", response1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response2 = sql_memory_chain.invoke({\"question\": \"Who is teaching the first course on the list?\"})\n",
    "print(\"Follow-up Response:\", response2)\n",
    "\n",
    "display_conversation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def safe_query(question):\n",
    "    \"\"\"\n",
    "    Safely execute a query with error handling\n",
    "    \n",
    "    Args:\n",
    "        question (str): Natural language question\n",
    "    Returns:\n",
    "        dict: Result and status\n",
    "    \"\"\"\n",
    "    try:\n",
    "        response = complete_chain.invoke({\"question\": question})\n",
    "        return {\n",
    "            \"success\": True,\n",
    "            \"response\": response\n",
    "        }\n",
    "    except Exception as e:\n",
    "        return {\n",
    "            \"success\": False,\n",
    "            \"error\": str(e)\n",
    "        }\n",
    "\n",
    "# Example usage\n",
    "result = safe_query(\"Which courses are available in the Data Science department for Spring 2025?\")\n",
    "if result[\"success\"]:\n",
    "    print(\"Answer:\", result[\"response\"])\n",
    "else:\n",
    "    print(\"Error:\", result[\"error\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from typing import Dict, List\n",
    "\n",
    "class InputValidator:\n",
    "    \"\"\"Input validation and sanitization for database queries\"\"\"\n",
    "    \n",
    "    # Restricted words that might indicate harmful queries\n",
    "    RESTRICTED_WORDS = {\n",
    "        'delete', 'drop', 'truncate', 'update', 'insert', \n",
    "        'alter', 'create', 'replace', 'modify', 'grant'\n",
    "    }\n",
    "    \n",
    "    # Allowed table names from our database\n",
    "    ALLOWED_TABLES = ['spring_2025_courses']  # in lowercase\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.error_messages = []\n",
    "    \n",
    "    def validate_question(self, question: str) -> bool:\n",
    "        \"\"\"\n",
    "        Validate a natural language question\n",
    "        \n",
    "        Args:\n",
    "            question (str): User's input question\n",
    "            \n",
    "        Returns:\n",
    "            bool: True if valid, False otherwise\n",
    "        \"\"\"\n",
    "        self.error_messages = []\n",
    "        \n",
    "        # Check if question is empty or too long\n",
    "        if not question or not question.strip():\n",
    "            self.error_messages.append(\"Question cannot be empty\")\n",
    "            return False\n",
    "            \n",
    "        if len(question) > 500:\n",
    "            self.error_messages.append(\"Question is too long (max 500 characters)\")\n",
    "            return False\n",
    "        \n",
    "        # Check for basic SQL injection attempts\n",
    "        question_lower = question.lower()\n",
    "        if any(word in question_lower for word in self.RESTRICTED_WORDS):\n",
    "            self.error_messages.append(\"Question contains restricted keywords\")\n",
    "            return False\n",
    "        \n",
    "        # Check for excessive special characters\n",
    "        if re.search(r'[;{}\\\\]', question):\n",
    "            self.error_messages.append(\"Question contains invalid characters\")\n",
    "            return False\n",
    "            \n",
    "        return True\n",
    "    \n",
    "    def get_error_messages(self) -> List[str]:\n",
    "        \"\"\"Get all error messages from validation\"\"\"\n",
    "        return self.error_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QueryValidator:\n",
    "    \"\"\"Validate generated SQL queries before execution\"\"\"\n",
    "\n",
    "    def __init__(self, db_connection):\n",
    "        self.db = db_connection\n",
    "        self.error_messages = []\n",
    "    \n",
    "    def get_error_messages(self) -> List[str]:\n",
    "        \"\"\"Get all error messages from validation\"\"\"\n",
    "        return self.error_messages\n",
    "    \n",
    "    def validate_sql(self, sql: str) -> bool:\n",
    "        \"\"\"Validate generated SQL query\"\"\"\n",
    "        sql_lower = sql.lower()\n",
    "        \n",
    "        # Check if query is read-only (SELECT only)\n",
    "        if not sql_lower.strip().startswith('select'):\n",
    "            self.error_messages.append(\"Only SELECT queries are allowed\")\n",
    "            return False\n",
    "        \n",
    "        # Check for multiple statements\n",
    "        if ';' in sql[:-1]:  # Allow semicolon at the end\n",
    "            self.error_messages.append(\"Multiple SQL statements are not allowed\")\n",
    "            return False\n",
    "        \n",
    "        # Validate table names\n",
    "        table = self._extract_table_names(sql_lower)\n",
    "        if table not in InputValidator.ALLOWED_TABLES:\n",
    "            self.error_messages.append(\"Query contains invalid table names\")\n",
    "            return False\n",
    "            \n",
    "        return True\n",
    "    \n",
    "    def _extract_table_names(self, sql: str) -> str:\n",
    "        \"\"\"Extract table names from SQL query\"\"\"\n",
    "        # Simple regex to extract table names\n",
    "        # Note: This is a basic implementation\n",
    "        from_matches = re.findall(r'from\\s+([a-zA-Z_][a-zA-Z0-9_]*)', sql.lower())\n",
    "        # join_matches = re.findall(r'join\\s+([a-zA-Z_][a-zA-Z0-9_]*)', sql)\n",
    "        return from_matches[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SafeQueryExecutor:\n",
    "    \"\"\"Safe execution of natural language queries\"\"\"\n",
    "    \n",
    "    def __init__(self, db_connection, llm_chain):\n",
    "        self.input_validator = InputValidator()\n",
    "        self.query_validator = QueryValidator(db_connection)\n",
    "        self.db = db_connection\n",
    "        self.chain = llm_chain\n",
    "        \n",
    "    def execute_safe_query(self, question: str) -> Dict:\n",
    "        \"\"\"\n",
    "        Safely execute a natural language query\n",
    "        \n",
    "        Args:\n",
    "            question (str): User's natural language question\n",
    "            \n",
    "        Returns:\n",
    "            dict: Query result and status\n",
    "        \"\"\"\n",
    "        # Validate input\n",
    "        if not self.input_validator.validate_question(question):\n",
    "            return {\n",
    "                'success': False,\n",
    "                'error': self.input_validator.get_error_messages(),\n",
    "                'query': None,\n",
    "                'result': None\n",
    "            }\n",
    "        \n",
    "        try:\n",
    "            # Generate SQL query\n",
    "            sql_query = sql_generator.invoke({\"question\": question})\n",
    "            \n",
    "            # Validate generated SQL\n",
    "            if not self.query_validator.validate_sql(sql_query):\n",
    "                return {\n",
    "                    'success': False,\n",
    "                    'error': self.query_validator.get_error_messages(),\n",
    "                    'query': sql_query,\n",
    "                    'result': None\n",
    "                }\n",
    "            \n",
    "            # Execute query\n",
    "            result = complete_chain.invoke({\"question\": question})\n",
    "            \n",
    "            return {\n",
    "                'success': True,\n",
    "                'error': None,\n",
    "                'query': sql_query,\n",
    "                'result': result\n",
    "            }\n",
    "            \n",
    "        except Exception as e:\n",
    "            return {\n",
    "                'success': False,\n",
    "                'error': [str(e)],\n",
    "                'query': None,\n",
    "                'result': None\n",
    "            }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the safe executor\n",
    "safe_executor = SafeQueryExecutor(db, complete_chain)\n",
    "\n",
    "result1 = safe_executor.execute_safe_query(\"Which courses are available in the Data Science department for Spring 2025?\")\n",
    "print(\"\\nValid Query Result:\")\n",
    "print(result1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Valid Query Result:\n",
      "{'success': True, 'error': None, 'query': \" \\nSELECT CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Data Science%';\", 'result': ' The available courses in the Data Science department for Spring 2025 are Introduction to Data Science, Career Development in Data Science, Data Structures, Data Visualization, Data Inference, Data Mining, Big Data Engineering, Independent Study for Data Science, Senior Project, Mathematical Foundation for Data Science, Computational Foundation for Data Science, Probability for Data Science, Exploratory Data Analysis and Visualization, Modern Applied Statistical Learning, Data Mining for Business, Big Data and Data Engineering, Design and Analysis of Experiments, and Capstone Project.'}\n",
      "\n",
      "Invalid Query Result:\n",
      "{'success': True, 'error': None, 'query': \" \\nSELECT * FROM Spring_2025_courses WHERE CourseTitle LIKE '%Introduction to Python%';\", 'result': '  The course Introduction to Python is not offered in Spring 2025.'}\n",
      "\n",
      "Restricted Query Result:\n",
      "{'success': False, 'error': ['Multiple SQL statements are not allowed'], 'query': ' \\nSELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Program LIKE \\'%Bachelor of Science%\\' AND Department LIKE \\'%Computer Science%\\'  AND \"New/Repeat\" IS NULL; \\n\\nHowever, since New/Repeat is not supposed to be used in queries according to the prompt, the query should be revised to:\\n\\nSELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Program LIKE \\'%Bachelor of Science%\\' AND Department LIKE \\'%Computer Science%\\'; \\n\\nHowever, considering the table schema, the query should be revised to:\\n\\nSELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Program LIKE \\'%Bachelor of Science%\\' AND Department LIKE \\'%Computer Science%\\'; \\n\\nHowever, this query may not be accurate because it does not account for the case where the department is \\'Computer Science (CS), Computer Networks and Cybersecurity (CNCS)\\'. Therefore, the query should be revised to:\\n\\nSELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Program LIKE \\'%Bachelor of Science%\\' AND (Department LIKE \\'%Computer Science%\\' OR Department LIKE \\'%Computer Networks and Cybersecurity%\\'); \\n\\nHowever, this query may still not be accurate because it does not account for the case where the department is \\'Computer Science (CS), Computer Networks and Cybersecurity (CNCS)\\' and the program is \\'Bachelor of Science\\'. Therefore, the query should be revised to:\\n\\nSELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Program LIKE \\'%Bachelor of Science%\\' AND Department LIKE \\'%Computer Science (CS), Computer Networks and Cybersecurity (CNCS)%\\'; \\n\\nHowever, this query may still not be accurate because it does not account for the case where the department is \\'Computer Science (CS), Computer Networks and Cybersecurity (CNCS)\\' and the program is \\'Bachelor of Science\\' and the department name may not be exact. Therefore, the query should be revised to:\\n\\nSELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Program LIKE \\'%Bachelor of Science%\\' AND Department LIKE \\'%Computer Science%\\' AND Department LIKE \\'%Computer Networks and Cybersecurity%\\'; \\n\\nHowever, the query is still not accurate because it requires the department to have both \\'Computer Science\\' and \\'Computer Networks and Cybersecurity\\' which may not be the case. Therefore, the query should be revised to:\\n\\nSELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Program LIKE \\'%Bachelor of Science%\\' AND (Department LIKE \\'%Computer Science%\\' OR Department LIKE \\'%Computer Networks and Cybersecurity%\\'); \\n\\nHowever, this query may still not be accurate because', 'result': None}\n",
      "\n",
      "Restricted Query Result:\n",
      "{'success': True, 'error': None, 'query': \" \\nSELECT CourseTitle FROM Spring_2025_courses WHERE Instructor LIKE '%Qu%';\", 'result': ' Dr. Qu is teaching Introduction to Data Science, Data Structures, Data Mining, Computational Foundation for Data Science, Data Mining for Business, and Computer Vision and Natural Language Processing in Spring 2025.'}\n"
     ]
    }
   ],
   "source": [
    "# Initialize the safe executor\n",
    "safe_executor = SafeQueryExecutor(db, complete_chain)\n",
    "\n",
    "# Example queries\n",
    "def test_queries():\n",
    "    # Valid query\n",
    "    result1 = safe_executor.execute_safe_query(\"Which courses are available in the Data Science department for Spring 2025?\")\n",
    "    print(\"\\nValid Query Result:\")\n",
    "    print(result1)\n",
    "    \n",
    "    # Invalid query with SQL injection attempt\n",
    "    result2 = safe_executor.execute_safe_query(\"\"\"What are the details of the course ‘Introduction to Python’?\"\"\")\n",
    "    print(\"\\nInvalid Query Result:\")\n",
    "    print(result2)\n",
    "    \n",
    "    # Query with restricted words\n",
    "    result3 = safe_executor.execute_safe_query(\"How many courses are offered under the Computer Science (Bachelor of Science) program?\")\n",
    "    print(\"\\nRestricted Query Result:\")\n",
    "    print(result3)\n",
    "\n",
    "    # Query with restricted words\n",
    "    result4 = safe_executor.execute_safe_query(\"Show me the courses taught by Dr. Qu in Spring 2025.\")\n",
    "    print(\"\\nRestricted Query Result:\")\n",
    "    print(result4)\n",
    "    \n",
    "# Run test queries\n",
    "test_queries()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result1 = safe_executor.execute_safe_query(\"DROP TABLE courses;\")\n",
    "result1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "COS243",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}