Update app.py

app.py

@@ -1,5 +1,137 @@
+import os
+import gradio as gr
+import pandas as pd
+from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float, insert, text, inspect
+from smolagents import tool, CodeAgent, InferenceClientModel
+
+# --- Setup SQLite database ---
+engine = create_engine("sqlite:///data.db")
+metadata_obj = MetaData()
+
+# --- Create students table ---
+students = Table(
+    "students",
+    metadata_obj,
+    Column("student_id", Integer, primary_key=True),
+    Column("student_name", String(32), nullable=False),
+)
+metadata_obj.create_all(engine)
+
+# Insert sample student data
+student_rows = [
+    {"student_id": 1, "student_name": "Alice Johnson"},
+    {"student_id": 2, "student_name": "Bob Smith"},
+    {"student_id": 3, "student_name": "Charlie Brown"},
+    {"student_id": 4, "student_name": "Diana Prince"},
+]
+for row in student_rows:
+    stmt = insert(students).values(**row)
+    with engine.begin() as conn:
+        conn.execute(stmt)
+
+# --- Create subjects table ---
+subjects = Table(
+    "subjects",
+    metadata_obj,
+    Column("student_id", Integer, primary_key=True),
+    Column("subject_name", String(32), primary_key=True),
+    Column("marks", Float, nullable=False),
+    Column("teacher", String(32), nullable=False),
+)
+metadata_obj.create_all(engine)
+
+# Insert sample subject/marks data
+subject_rows = [
+    {"student_id": 1, "subject_name": "Math", "marks": 95.0, "teacher": "Mr. Adams"},
+    {"student_id": 1, "subject_name": "English", "marks": 88.5, "teacher": "Ms. Baker"},
+    {"student_id": 1, "subject_name": "Science", "marks": 92.0, "teacher": "Dr. Carter"},
+    {"student_id": 2, "subject_name": "Math", "marks": 72.0, "teacher": "Mr. Adams"},
+    {"student_id": 2, "subject_name": "Science", "marks": 81.0, "teacher": "Dr. Carter"},
+    {"student_id": 3, "subject_name": "Math", "marks": 85.0, "teacher": "Mr. Adams"},
+    {"student_id": 3, "subject_name": "History", "marks": 90.0, "teacher": "Mrs. Davis"},
+    {"student_id": 4, "subject_name": "English", "marks": 78.0, "teacher": "Ms. Baker"},
+    {"student_id": 4, "subject_name": "Science", "marks": 82.0, "teacher": "Dr. Carter"},
+]
+for row in subject_rows:
+    stmt = insert(subjects).values(**row)
+    with engine.begin() as conn:
+        conn.execute(stmt)
+
+# --- Convert tables to DataFrames for front-end display ---
+def fetch_table(table_name: str) -> pd.DataFrame:
+    with engine.connect() as con:
+        rows = con.execute(text(f"SELECT * FROM {table_name}"))
+        df = pd.DataFrame(rows.fetchall(), columns=rows.keys())
+    return df
+
+students_df = fetch_table("students")
+subjects_df = fetch_table("subjects")
+
+# --- Define SQL tool for the agent ---
+@tool
+def sql_engine(query: str) -> list:
+    """
+    Executes SQL queries on the available tables: students and subjects.
+
+    Args:
+        query (str): The SQL query string to execute.
+
+    Returns:
+        list: List of tuples containing the query results.
+    """
+    try:
+        print(f"🧩 Executing query: {query}")
+        with engine.connect() as con:
+            rows = con.execute(text(query))
+            results = [tuple(row) for row in rows]
+        return results or []
+    except Exception as e:
+        return [f"⚠️ SQL Error: {str(e)}"]
+
+# Dynamically describe tables for the agent
+updated_description = "You can run SQL queries on these tables:\n"
+inspector = inspect(engine)
+for table in ["students", "subjects"]:
+    columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]
+    table_description = f"\n\nTable '{table}':\nColumns:\n" + "\n".join(
+        [f" - {name}: {col_type}" for name, col_type in columns_info]
+    )
+    updated_description += table_description
+
+sql_engine.description = updated_description
+
+# --- Create the agent ---
+agent = CodeAgent(
+    tools=[sql_engine],
+    model=InferenceClientModel(
+        "meta-llama/Meta-Llama-3-8B-Instruct",
+        api_key=os.environ.get("HF_TOKEN")
+    ),
+)
+
+# --- Sample prompts to display ---
+sample_prompts = [
+    "Which student scored highest in Math?",
+    "List all subjects and marks for Alice Johnson.",
+    "Who is the teacher of Science for Diana Prince?",
+    "Show students who scored more than 80 in Science.",
+    "Average marks per subject."
+]
+
+# --- Define Gradio interface ---
+def ask_agent(question: str) -> str:
+    """Ask the AI agent a question and return its answer."""
+    try:
+        result = agent.run(question)
+        if isinstance(result, list):
+            result_str = "\n".join(str(r) for r in result)
+            return result_str or "No results."
+        return str(result)
+    except Exception as e:
+        return f"⚠️ Error: {str(e)}"
+
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏫 Student Marks SQL Agent")
+    gr.Markdown("## 🏫 Text to SQL Agent")
 
     with gr.Row():
         # --- Left column: tables (smaller) ---