Query Limits Management Feature

.gitignore

@@ -15,3 +15,5 @@ venv/
 *.egg-info/
 
 .aider*
+
+cmoocenv

src/codehelp/migrations/20250112--codehelp--add_query_limits.sql

@@ -0,0 +1,16 @@
+-- SPDX-FileCopyrightText: Mark Liffiton <[email protected]>, Rana Moeez Hassan
+--
+-- SPDX-License-Identifier: AGPL-3.0-only
+
+PRAGMA foreign_keys = OFF;
+BEGIN;
+
+-- Add columns to the classes table
+ALTER TABLE classes ADD COLUMN query_limit_enabled BOOLEAN NOT NULL DEFAULT 0;
+ALTER TABLE classes ADD COLUMN max_queries INTEGER NOT NULL DEFAULT 50;
+
+-- Ensure the users table has a column to track the number of queries used
+ALTER TABLE users ADD COLUMN queries_used INTEGER NOT NULL DEFAULT 0;
+
+COMMIT;
+PRAGMA foreign_keys = ON;

src/codehelp/templates/help_form.html

@@ -107,8 +107,22 @@
 
   </div>
 
-  <div class="column has-background-light">
-    {{ recent_queries(history) }}
+  <div class="column is-one-quarter-desktop">
+    {% if auth.cur_class and auth.cur_class.query_limit_enabled and auth.cur_class.role == 'student' %}
+      <div class="notification {% if auth.user.queries_used >= auth.cur_class.max_queries %}is-warning{% else %}is-info{% endif %} is-light mb-4">
+        <div class="level">
+          <div class="level-item has-text-centered">
+            <div>
+              <p class="heading">Query Usage</p>
+              <p class="title">
+                {{ auth.user.queries_used }}/{{ auth.cur_class.max_queries }}
+              </p>
+            </div>
+          </div>
+        </div>
+      </div>
+    {% endif %}
+    {{ recent_queries() }}
   </div>
 
 </div>

src/gened/auth.py

@@ -48,15 +48,18 @@ class UserData:
     id: int
     display_name: str
     auth_provider: AuthProvider
+    queries_used: int
     is_admin: bool = False
-    is_tester: bool = False
+    is_tester: bool = False 
 
 @dataclass(frozen=True)
 class ClassData:
     class_id: int
     class_name: str
     role_id: int
     role: RoleType
+    query_limit_enabled: int
+    max_queries: int
 
 @dataclass(frozen=True)
 class AuthData:
@@ -130,7 +133,8 @@ def _get_auth_from_session() -> AuthData:
             users.display_name,
             users.is_admin,
             users.is_tester,
-            auth_providers.name AS auth_provider
+            auth_providers.name AS auth_provider,
+            users.queries_used
         FROM users
         LEFT JOIN auth_providers ON auth_providers.id=users.auth_provider
         WHERE users.id=?
@@ -147,6 +151,7 @@ def _get_auth_from_session() -> AuthData:
         auth_provider=user_row['auth_provider'],
         is_admin=user_row['is_admin'],
         is_tester=user_row['is_tester'],
+        queries_used=user_row['queries_used']
     )
 
     # Check the database for any active roles (may be changed by another user)
@@ -158,7 +163,9 @@ def _get_auth_from_session() -> AuthData:
             roles.class_id,
             roles.role,
             classes.name,
-            classes.enabled
+            classes.enabled,
+            classes.query_limit_enabled,
+            classes.max_queries
         FROM roles
         JOIN classes ON classes.id=roles.class_id
         WHERE roles.user_id=? AND roles.active=1
@@ -177,6 +184,8 @@ def _get_auth_from_session() -> AuthData:
             class_name=row['name'],
             role_id=row['role_id'],
             role=row['role'],
+            query_limit_enabled=row['query_limit_enabled'],
+            max_queries=row['max_queries'],
         )
         if row['class_id'] == sess_class_id:
             assert cur_class is None  # sanity check: should only ever match one role/class
@@ -191,12 +200,14 @@ def _get_auth_from_session() -> AuthData:
 
     # admin gets instructor role in all classes automatically
     if user.is_admin and cur_class is None and sess_class_id is not None:
-        class_row = db.execute("SELECT name FROM classes WHERE id=?", [sess_class_id]).fetchone()
+        class_row = db.execute("SELECT name, query_limit_enabled, max_queries FROM classes WHERE id=?", [sess_class_id]).fetchone()
         cur_class = ClassData(
             class_id=sess_class_id,
             class_name=class_row['name'],
             role_id=-1,
             role='instructor',
+            query_limit_enabled=class_row['query_limit_enabled'],
+            max_queries=class_row['max_queries']
         )
 
     # return an AuthData with all collected values

src/gened/base.py

@@ -10,7 +10,7 @@
 from typing import Any
 
 import flask.app
-import pyrage
+import pyrage # type: ignore
 from dotenv import load_dotenv
 from flask import Flask, render_template, send_from_directory
 from flask.wrappers import Response

src/gened/class_config.py

@@ -50,7 +50,7 @@ def config_form() -> str:
     class_id = cur_class.class_id
 
     class_row = db.execute("""
-        SELECT classes.id, classes.enabled, classes_user.link_ident, classes_user.link_reg_expires, classes_user.link_anon_login, classes_user.llm_api_key, classes_user.model_id
+        SELECT classes.id, classes.enabled, classes.query_limit_enabled, classes.max_queries,classes_user.link_ident, classes_user.link_reg_expires, classes_user.link_anon_login, classes_user.llm_api_key, classes_user.model_id
         FROM classes
         LEFT JOIN classes_user
           ON classes.id = classes_user.class_id

src/gened/db.py

@@ -14,7 +14,7 @@
 from pathlib import Path
 
 import click
-import pyrage
+import pyrage # type: ignore
 from flask import current_app, g
 from flask.app import Flask
 from werkzeug.security import generate_password_hash
@@ -153,7 +153,7 @@ def init_db() -> None:
         db.executescript(f.read())
 
     # App-specific schema in the app's package
-    with current_app.open_resource('schema.sql', mode='r', encoding='utf-8') as f:
+    with current_app.open_resource('schema.sql', mode='r') as f:
         db.executescript(f.read())
 
     # Mark all existing migrations as applied (since this is a fresh DB)

src/gened/filters.py

@@ -82,3 +82,20 @@ def markdown_filter(value: str) -> str:
         html = markdown_processor.render(value)
         # relying on MarkdownIt's escaping (w/o HTML parsing, due to "js-default"), so mark this as safe
         return Markup(html)
+
+def fmt_button_text(button_data: dict[str, str]) -> str:
+    """Format button HTML to be displayed in a table cell. Not used right now as could not get ButtonCol to work properly. """
+    url = button_data.get('url', '')
+    icon = button_data.get('icon', '')
+    text = button_data.get('text', '')
+
+    icon_html = f'<span class="icon"><i class="fas fa-{icon}"></i></span>' if icon else ''
+
+    return Markup(f'''
+        <form method="POST" action="{url}" style="display:inline">
+            <button class="button is-small is-info" type="submit">
+                {icon_html}
+                <span>{text}</span>
+            </button>
+        </form>
+    ''')

src/gened/instructor.py

@@ -32,8 +32,9 @@
 from .csv import csv_response
 from .data_deletion import delete_class_data
 from .db import get_db
+from .auth import get_auth
 from .redir import safe_redirect
-from .tables import BoolCol, DataTable, NumCol, UserCol
+from .tables import ButtonCol, BoolCol, Col, DataTable, NumCol, UserCol
 
 bp = Blueprint('instructor', __name__, template_folder='templates')
 
@@ -85,6 +86,31 @@ def _get_class_users(*, for_export: bool = False) -> list[Row]:
 
     return users
 
+def _get_query_limits_data() -> list[Row]:
+    cur_class = get_auth_class()
+    class_id = cur_class.class_id
+    db = get_db()
+
+    users = db.execute("""
+        SELECT
+            users.id,
+            json_array(users.display_name, auth_providers.name, users.display_extra) AS user,
+            users.queries_used,
+            classes.max_queries,
+            ('<form method="POST" action="/instructor/reset_student_queries/' || users.id || 
+             '" style="display:inline"><button class="button is-small is-warning" type="submit">' ||
+             '<span class="icon"><i class="fas fa-rotate"></i></span>' ||
+             '<span>Reset</span></button></form>') AS reset_button
+        FROM users
+        LEFT JOIN auth_providers ON users.auth_provider=auth_providers.id
+        JOIN roles ON roles.user_id=users.id
+        JOIN classes ON roles.class_id=classes.id
+        WHERE roles.class_id=? AND roles.role='student'
+        ORDER BY users.display_name
+    """, [class_id]).fetchall()
+
+    return users
+
 
 @bp.route("/")
 def main() -> str | Response:
@@ -98,6 +124,12 @@ def main() -> str | Response:
         data=users,
     )
 
+    queries_limits_table = DataTable(
+        name='query_limits',
+        columns=[NumCol('id', hidden=True), UserCol('user'), NumCol('queries_used', align="center"), NumCol('max_queries', align="center"), Col('reset_button', kind='html', align="center")],
+        data=_get_query_limits_data()
+    )
+
     sel_user_name = None
     sel_user_id = request.args.get('user', type=int)
     if sel_user_id is not None:
@@ -110,7 +142,8 @@ def main() -> str | Response:
     queries_table.data = _get_class_queries(sel_user_id)
     queries_table.csv_link = url_for('instructor.get_csv', kind='queries')
 
-    return render_template("instructor_view.html", users=users_table, queries=queries_table, user=sel_user_name)
+    return render_template( "instructor_view.html", users=users_table, queries=queries_table, query_limits=queries_limits_table, user=sel_user_name
+)
 
 
 @bp.route("/csv/<string:kind>")
@@ -199,3 +232,91 @@ def delete_class() -> Response:
 
     switch_class(None)
     return redirect(url_for("profile.main"))
+
+@bp.route("/class/config/save", methods=["POST"])
+def save_class_config() -> Response:
+    db = get_db()
+    cur_class = get_auth_class()  # Use get_auth_class() instead of get_auth()
+    class_id = cur_class.class_id
+
+    query_limit_enabled = 'query_limit_enabled' in request.form
+    max_queries = int(request.form.get('max_queries', 50))  # Default to 50 if not provided
+
+    db.execute("""
+        UPDATE classes
+        SET query_limit_enabled = ?, max_queries = ?
+        WHERE id = ?
+    """, [query_limit_enabled, max_queries, class_id])
+    db.commit()
+
+    flash("Class configuration updated.", "success")
+    return redirect(url_for("class_config.config_form"))
+
+@bp.route("/class/reset_queries", methods=["POST"])
+def reset_queries() -> Response:
+    db = get_db()
+    cur_class = get_auth_class()  # Use get_auth_class() instead of get_auth()
+    class_id = cur_class.class_id
+
+    db.execute("""
+        UPDATE users 
+        SET queries_used = 0 
+        WHERE id IN (
+            SELECT user_id 
+            FROM roles 
+            WHERE class_id = ? AND role = 'student'
+        )
+    """, [class_id])
+    db.commit()
+
+    flash("Query counts reset for all students.", "success")
+    return redirect(url_for("instructor.main"))  # Redirect to instructor main page
+
+@bp.route("/reset_student_queries/<int:user_id>", methods=["POST"])
+def reset_student_queries(user_id: int) -> Response:
+    db = get_db()
+    cur_class = get_auth_class()
+    class_id = cur_class.class_id
+
+    # Verify user belongs to class and is a student
+    student = db.execute("""
+        SELECT users.id 
+        FROM users 
+        JOIN roles ON roles.user_id = users.id 
+        WHERE users.id = ? AND roles.class_id = ? AND roles.role = 'student'
+    """, [user_id, class_id]).fetchone()
+
+    if not student:
+        flash("Invalid student ID.", "error")
+        return redirect(url_for("instructor.main"))
+
+    # Reset queries for student
+    db.execute(
+        "UPDATE users SET queries_used = 0 WHERE id = ?",
+        [user_id]
+    )
+    db.commit()
+
+    flash("Query count reset for student.", "success")
+    return redirect(url_for("instructor.main"))
+
+@bp.route("/reset_all_queries", methods=["POST"])
+def reset_all_queries() -> Response:
+    db = get_db()
+    cur_class = get_auth_class()
+    class_id = cur_class.class_id
+
+    # Reset queries for all students in class
+    db.execute("""
+        UPDATE users 
+        SET queries_used = 0 
+        WHERE id IN (
+            SELECT user_id 
+            FROM roles 
+            WHERE class_id = ? AND role = 'student'
+        )
+    """, [class_id])
+    db.commit()
+
+    flash("Query counts reset for all students.", "success")
+    return redirect(url_for("instructor.main"))