From 7e38864e6bab9a35b074318e37e2b83c8a536822 Mon Sep 17 00:00:00 2001
From: Christopher Odoom <your.email@example.com>
Date: Fri, 24 Jul 2026 19:24:18 -0400
Subject: [PATCH] Fix AI integration readiness

---
 routes/admin_routes.py           | 94 ++++++++++++++++++++++++++------
 routes/chatbot_routes.py         |  4 +-
 templates/admin/ai_settings.html | 63 +++++++++++++++++++--
 tests/test_integrations.py       | 65 +++++++++++++++++++++-
 utils/ai_usage.py                | 21 +++++++
 5 files changed, 222 insertions(+), 25 deletions(-)

diff --git a/routes/admin_routes.py b/routes/admin_routes.py
index 194dc5a..e3f57f8 100644
--- a/routes/admin_routes.py
+++ b/routes/admin_routes.py
@@ -1,14 +1,27 @@
 from flask import Blueprint, render_template, request, redirect, url_for, flash, jsonify
 from flask_login import login_required, current_user
-from models import db, User, ExpertApplication, Analysis, Consultation
+from models import (
+    db,
+    User,
+    ExpertApplication,
+    Analysis,
+    Consultation,
+)
 from functools import wraps
 from datetime import datetime, timedelta, timezone
 from utils.email_service import send_expert_approved_email, send_expert_rejected_email
 from utils.ai_service import (
     call_openai_api, is_ai_enabled, OpenAIServiceError, get_openai_config
 )
+from utils.ai_usage import (
+    ai_usage_storage_ready,
+    initialize_ai_usage_storage,
+)
+from sqlalchemy.exc import SQLAlchemyError
+import logging
 import re
 admin = Blueprint('admin', __name__, url_prefix='/admin')
+logger = logging.getLogger(__name__)
 # Custom decorator for admin access
 def admin_required(f):
     @wraps(f)
@@ -239,28 +252,39 @@ def ai_settings():
     """Display non-secret AI integration status."""
     current_api_key, current_model = get_openai_config()
     current_enabled = is_ai_enabled()
+    storage_ready = ai_usage_storage_ready()
     current_settings = {
         'api_key_configured': bool(current_api_key),
         'model': current_model,
         'enabled': current_enabled,
+        'storage_ready': storage_ready,
     }
-    api_status = {
-        'status': (
-            'configured'
-            if current_enabled and current_api_key
-            else 'disabled'
-            if not current_enabled
-            else 'no_key'
-        ),
-        'last_error': (
-            None
-            if current_enabled and current_api_key
-            else 'AI enhancement is disabled.'
-            if not current_enabled
-            else 'OPENAI_API_KEY is missing.'
-        ),
-        'credit_warning': False,
-    }
+    if not current_enabled:
+        api_status = {
+            'status': 'disabled',
+            'label': 'Disabled',
+            'last_error': 'AI enhancement is disabled.',
+        }
+    elif not current_api_key:
+        api_status = {
+            'status': 'no_key',
+            'label': 'Missing API token',
+            'last_error': 'OPENAI_API_KEY is missing.',
+        }
+    elif not storage_ready:
+        api_status = {
+            'status': 'needs_setup',
+            'label': 'Needs initialization',
+            'last_error': (
+                'The AI usage database table has not been initialized.'
+            ),
+        }
+    else:
+        api_status = {
+            'status': 'ready',
+            'label': 'Ready',
+            'last_error': None,
+        }
     return render_template(
         'admin/ai_settings.html',
         current_settings=current_settings,
@@ -268,6 +292,40 @@ def ai_settings():
     )
 
 
+@admin.route('/initialize-ai-storage', methods=['POST'])
+@login_required
+@admin_required
+def initialize_ai_storage():
+    """Initialize durable AI usage storage for serverless deployments."""
+    data = request.get_json(silent=True)
+    if not isinstance(data, dict) or data.get('confirm') is not True:
+        return jsonify({
+            'success': False,
+            'error': 'Explicit initialization confirmation is required.',
+        }), 400
+
+    try:
+        initialize_ai_usage_storage()
+    except SQLAlchemyError:
+        logger.exception(
+            "AI storage initialization failed for admin user %s.",
+            current_user.id,
+        )
+        return jsonify({
+            'success': False,
+            'error': (
+                'The AI usage table could not be created. Verify that '
+                'DATABASE_URL is correct and permits schema changes.'
+            ),
+        }), 503
+
+    return jsonify({
+        'success': True,
+        'message': 'AI usage storage is initialized.',
+        'storage_ready': ai_usage_storage_ready(),
+    })
+
+
 @admin.route('/test-ai-integration', methods=['POST'])
 @login_required
 @admin_required
diff --git a/routes/chatbot_routes.py b/routes/chatbot_routes.py
index d25b309..99c6a66 100644
--- a/routes/chatbot_routes.py
+++ b/routes/chatbot_routes.py
@@ -64,8 +64,8 @@ def ask_question():
             success=False,
             message="AI usage tracking is unavailable.",
             response=(
-                "The AI assistant is not initialized yet. "
-                "Please contact the administrator."
+                "The AI assistant's usage storage is not ready. "
+                "An administrator can initialize it from AI Integration."
             ),
         ), 503
 
diff --git a/templates/admin/ai_settings.html b/templates/admin/ai_settings.html
index 253df36..ef2a887 100644
--- a/templates/admin/ai_settings.html
+++ b/templates/admin/ai_settings.html
@@ -7,7 +7,7 @@
     <div class="row mb-4">
         <div class="col-12">
             <h2>AI Integration</h2>
-            <p class="lead">OpenAI Responses API configuration and connection test.</p>
+            <p class="lead">OpenAI configuration, usage-storage readiness, and live provider test.</p>
         </div>
     </div>
 
@@ -43,10 +43,19 @@ <h3 class="mb-0">Server-side configuration</h3>
                 <dt class="col-sm-4">Model</dt>
                 <dd class="col-sm-8"><code>{{ current_settings.model }}</code></dd>
 
-                <dt class="col-sm-4">Status</dt>
+                <dt class="col-sm-4">Usage storage</dt>
                 <dd class="col-sm-8">
-                    <span class="badge bg-{{ 'success' if api_status.status == 'configured' else 'secondary' }}">
-                        {{ api_status.status }}
+                    {% if current_settings.storage_ready %}
+                        <span class="badge bg-success">Ready</span>
+                    {% else %}
+                        <span class="badge bg-warning text-dark">Not initialized</span>
+                    {% endif %}
+                </dd>
+
+                <dt class="col-sm-4">Overall readiness</dt>
+                <dd class="col-sm-8">
+                    <span class="badge bg-{{ 'success' if api_status.status == 'ready' else 'warning' if api_status.status == 'needs_setup' else 'secondary' }}">
+                        {{ api_status.label }}
                     </span>
                     {% if api_status.last_error %}
                         <span class="text-muted ms-2">{{ api_status.last_error }}</span>
@@ -62,6 +71,25 @@ <h3 class="mb-0">Server-side configuration</h3>
         </div>
     </div>
 
+    {% if not current_settings.storage_ready %}
+    <div class="card shadow mb-4 border-warning">
+        <div class="card-header bg-light">
+            <h4 class="mb-0">Initialize AI usage storage</h4>
+        </div>
+        <div class="card-body">
+            <p>
+                The assistant requires a durable PostgreSQL table to enforce
+                per-user usage limits. This action creates only the missing
+                <code>ai_usage_events</code> table and its indexes.
+            </p>
+            <button type="button" id="initialize-storage-button" class="btn btn-primary">
+                Initialize AI storage
+            </button>
+            <div id="storage-result" class="alert mt-3 d-none" role="status"></div>
+        </div>
+    </div>
+    {% endif %}
+
     <div class="card shadow">
         <div class="card-header bg-light">
             <h4 class="mb-0">Connection test</h4>
@@ -80,10 +108,37 @@ <h4 class="mb-0">Connection test</h4>
 
 <script>
 document.addEventListener('DOMContentLoaded', () => {
+    const storageButton = document.getElementById('initialize-storage-button');
+    const storageResult = document.getElementById('storage-result');
     const button = document.getElementById('test-button');
     const prompt = document.getElementById('test-prompt');
     const result = document.getElementById('test-result');
 
+    if (storageButton) {
+        storageButton.addEventListener('click', async () => {
+            storageButton.disabled = true;
+            storageResult.className = 'alert alert-info mt-3';
+            storageResult.textContent = 'Initializing durable AI usage storage…';
+            try {
+                const response = await fetch('{{ url_for("admin.initialize_ai_storage") }}', {
+                    method: 'POST',
+                    headers: {'Content-Type': 'application/json'},
+                    body: JSON.stringify({confirm: true})
+                });
+                const data = await response.json();
+                storageResult.className = `alert ${response.ok && data.success ? 'alert-success' : 'alert-danger'} mt-3`;
+                storageResult.textContent = data.message || data.error || 'Storage initialization failed.';
+                if (response.ok && data.success) {
+                    window.setTimeout(() => window.location.reload(), 800);
+                }
+            } catch (error) {
+                storageResult.className = 'alert alert-danger mt-3';
+                storageResult.textContent = 'Storage initialization could not be completed.';
+                storageButton.disabled = false;
+            }
+        });
+    }
+
     button.addEventListener('click', async () => {
         const value = prompt.value.trim();
         if (!value) return;
diff --git a/tests/test_integrations.py b/tests/test_integrations.py
index c0f97cc..ed29182 100644
--- a/tests/test_integrations.py
+++ b/tests/test_integrations.py
@@ -5,8 +5,9 @@
 import httpx
 from openai import NotFoundError, RateLimitError
 import pytest
+from sqlalchemy import inspect
 
-from models import AIUsageEvent
+from models import AIUsageEvent, db
 from utils.ai_service import (
     OpenAIServiceError,
     call_openai_api,
@@ -198,6 +199,68 @@ def test_admin_ai_page_never_renders_api_key(client, admin_user, monkeypatch):
     assert b"Configured" in response.data
 
 
+def test_admin_can_initialize_missing_ai_usage_storage(
+    app,
+    admin_client,
+    monkeypatch,
+):
+    monkeypatch.setenv("AI_ENHANCEMENT_ENABLED", "true")
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
+    with app.app_context():
+        AIUsageEvent.__table__.drop(bind=db.engine)
+
+    status_response = admin_client.get("/admin/ai_settings")
+
+    assert status_response.status_code == 200
+    assert b"Needs initialization" in status_response.data
+    assert b"Not initialized" in status_response.data
+
+    initialize_response = admin_client.post(
+        "/admin/initialize-ai-storage",
+        json={"confirm": True},
+    )
+
+    assert initialize_response.status_code == 200
+    assert initialize_response.get_json()["storage_ready"] is True
+    with app.app_context():
+        assert inspect(db.engine).has_table("ai_usage_events")
+
+    ready_response = admin_client.get("/admin/ai_settings")
+    assert b"Overall readiness" in ready_response.data
+    assert b"Ready" in ready_response.data
+
+
+def test_ai_storage_initialization_requires_explicit_confirmation(admin_client):
+    response = admin_client.post(
+        "/admin/initialize-ai-storage",
+        json={"confirm": False},
+    )
+
+    assert response.status_code == 400
+    assert response.get_json()["success"] is False
+
+
+def test_chatbot_explains_missing_ai_usage_storage(
+    app,
+    client,
+    test_user,
+    monkeypatch,
+):
+    monkeypatch.setenv("AI_ENHANCEMENT_ENABLED", "true")
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
+    _login(client, test_user)
+    with app.app_context():
+        AIUsageEvent.__table__.drop(bind=db.engine)
+
+    response = client.post(
+        "/chatbot/ask",
+        json={"question": "Which model?", "context": "Model selection"},
+    )
+
+    assert response.status_code == 503
+    assert "usage storage is not ready" in response.get_json()["response"]
+
+
 def test_questionnaire_ai_enhancement_requires_login(client, monkeypatch):
     monkeypatch.setenv("AI_ENHANCEMENT_ENABLED", "true")
 
diff --git a/utils/ai_usage.py b/utils/ai_usage.py
index c6669c7..b1bf069 100644
--- a/utils/ai_usage.py
+++ b/utils/ai_usage.py
@@ -1,11 +1,32 @@
 """Durable per-user usage controls for paid AI features."""
 
 import os
+import logging
 from datetime import datetime, timedelta
 
+from sqlalchemy import inspect
+from sqlalchemy.exc import SQLAlchemyError
+
 from models import AIUsageEvent, db
 
 
+logger = logging.getLogger(__name__)
+
+
+def ai_usage_storage_ready() -> bool:
+    """Return whether the durable AI usage table is available."""
+    try:
+        return inspect(db.engine).has_table(AIUsageEvent.__tablename__)
+    except SQLAlchemyError:
+        logger.exception("Could not inspect AI usage storage.")
+        return False
+
+
+def initialize_ai_usage_storage() -> None:
+    """Create only the durable AI usage table when it is missing."""
+    AIUsageEvent.__table__.create(bind=db.engine, checkfirst=True)
+
+
 def hourly_ai_limit() -> int:
     """Return the configured per-user hourly unit budget."""
     try: