Query Limits Management Feature

.gitignore

@@ -14,4 +14,4 @@ venv/
 
 *.egg-info/
 
-.aider*
+.aider*

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

@@ -0,0 +1,14 @@
+-- SPDX-FileCopyrightText: Mark Liffiton <[email protected]>, Rana Moeez Hassan
+--
+-- SPDX-License-Identifier: AGPL-3.0-only
+
+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;

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/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/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,32 @@ 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"><img src="data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSIyNCIgaGVpZ2h0PSIyNCIgdmlld0JveD0iMCAwIDI0IDI0IiBmaWxsPSJub25lIiBzdHJva2U9ImN1cnJlbnRDb2xvciIgc3Ryb2tlLXdpZHRoPSIyIiBzdHJva2UtbGluZWNhcD0icm91bmQiIHN0cm9rZS1saW5lam9pbj0icm91bmQiIGNsYXNzPSJsdWNpZGUgbHVjaWRlLWEtYXJyb3ctZG93biI+PHBhdGggZD0iTTMuNSAxM2g2Ii8+PHBhdGggZD0ibTIgMTYgNC41LTkgNC41IDkiLz48cGF0aCBkPSJNMTggN3Y5Ii8+PHBhdGggZD0ibTE0IDEyIDQgNCA0LTQiLz48L3N2Zz4=" alt="reset icon" style="width: 1em; height: 1em;"></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 +125,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 +143,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 +233,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"))

src/gened/llm.py

@@ -67,6 +67,9 @@ class NoKeyFoundError(Exception):
 class NoTokensError(Exception):
     pass
 
+class MaxQueriesUsed(Exception):
+    pass
+
 def _get_llm(*, use_system_key: bool, spend_token: bool) -> LLM:
     ''' Get an LLM object configured based on the arguments and the current
     context (user and class).
@@ -113,6 +116,11 @@ def make_system_client(tokens_remaining: int | None = None) -> LLM:
         class_row = db.execute("""
             SELECT
                 classes.enabled,
+                classes.query_limit_enabled,
+                classes.max_queries,
+                classes.name as class_name,
+                roles.role,
+                users.queries_used,
                 COALESCE(consumers.llm_api_key, classes_user.llm_api_key) AS llm_api_key,
                 COALESCE(consumers.model_id, classes_user.model_id) AS _model_id,
                 models.model
@@ -125,14 +133,29 @@ def make_system_client(tokens_remaining: int | None = None) -> LLM:
               ON classes.id = classes_user.class_id
             LEFT JOIN models
               ON models.id = _model_id
+            LEFT JOIN roles 
+              ON roles.class_id = classes.id AND roles.user_id = ?
+            LEFT JOIN users
+              ON users.id = ?
             WHERE classes.id = ?
-        """, [auth.cur_class.class_id]).fetchone()
+        """, [auth.user_id, auth.user_id, auth.cur_class.class_id]).fetchone()
 
         if not class_row['enabled']:
             raise ClassDisabledError
 
         if not class_row['llm_api_key']:
             raise NoKeyFoundError
+
+        if class_row['query_limit_enabled'] and class_row['role'] == 'student':
+            if class_row['queries_used'] >= class_row['max_queries']:
+                raise MaxQueriesUsed
+
+            if spend_token:
+                db.execute(
+                    "UPDATE users SET queries_used = queries_used + 1 WHERE id = ?",
+                    [auth.user_id]
+                )
+                db.commit()
 
         return LLM(
             provider='openai',
@@ -203,6 +226,9 @@ def decorated_function(*args: P.args, **kwargs: P.kwargs) -> str | R:
             except NoTokensError:
                 flash("You have used all of your free queries.  If you are using this application in a class, please connect using the link from your class for continued access.  Otherwise, you can create a class and add an API key or contact us if you want to continue using this application.", "warning")
                 return render_template("error.html")
+            except MaxQueriesUsed:
+                flash("You have used the maximum of queries alloted to you by the professor. Please see the professor to reset the number of queries and get access")
+                return render_template("error.html")
 
             kwargs['llm'] = llm
             return f(*args, **kwargs)