Add OpenAI API status check and memory usage by module to status page

OWASP-BLT · Jan 1, 2025 · 11b58bf · 11b58bf
1 parent 5af7222
commit 11b58bf
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 3 deletions.
diff --git a/website/templates/status_page.html b/website/templates/status_page.html
@@ -52,14 +52,14 @@
         .status-dot.red {
             background-color: #dc3545;
         }
-        .memory-info, .memory-consumers {
+        .memory-info, .memory-consumers, .memory-by-module {
             margin-top: 10px;
         }
-        .memory-consumers ul {
+        .memory-consumers ul, .memory-by-module ul {
             list-style-type: none;
             padding: 0;
         }
-        .memory-consumers li {
+        .memory-consumers li, .memory-by-module li {
             margin-bottom: 5px;
         }
     </style>
@@ -147,6 +147,19 @@ <h3>Top Memory Consumers</h3>
                 </ul>
             </div>
         </div>
+        <!-- Memory Usage by Module -->
+        <div class="status-card">
+            <h3>Memory Usage by Module</h3>
+            <div class="memory-by-module">
+                <ul>
+                    {% for module, memory in status.memory_by_module.items %}
+                        <li>
+                            <strong>{{ module }}</strong>: {{ memory|filesizeformat }}
+                        </li>
+                    {% endfor %}
+                </ul>
+            </div>
+        </div>
         <div class="status-card">
             <h3>Database Connection Count</h3>
             <p>{{ status.db_connection_count }}</p>
@@ -161,5 +174,16 @@ <h3>Redis Stats</h3>
                 {% endfor %}
             </ul>
         </div>
+        <div class="status-card">
+            <h3>OpenAI API</h3>
+            <div class="status-item">
+                <div class="status-dot {% if status.openai %}green{% else %}red{% endif %}"></div>
+                {% if status.openai %}
+                    <span>Operational</span>
+                {% else %}
+                    <span>Not Operational</span>
+                {% endif %}
+            </div>
+        </div>
     </div>
 {% endblock content %}
diff --git a/website/views/core.py b/website/views/core.py
@@ -1,6 +1,7 @@
 import json
 import os
 import subprocess
+import sys
 import tracemalloc
 import urllib
 from datetime import datetime, timezone
@@ -73,6 +74,7 @@ def check_status(request):
             "bitcoin_block": None,
             "sendgrid": False,
             "github": False,
+            "openai": False,
             "memory_info": psutil.virtual_memory()._asdict(),
             "top_memory_consumers": [],
             "memory_profiling": {},
@@ -138,6 +140,26 @@ def check_status(request):
                 status["github"] = False
                 print(f"GitHub API Error: {e}")
 
+        openai_api_key = os.getenv("OPENAI_API_KEY")
+        if openai_api_key:
+            try:
+                headers = {"Authorization": f"Bearer {openai_api_key}"}
+                response = requests.get(
+                    "https://api.openai.com/v1/models", headers=headers, timeout=5
+                )
+
+                if response.status_code == 200:
+                    status["openai"] = True
+                else:
+                    status["openai"] = False
+                    print(
+                        f"OpenAI API token check failed with status code {response.status_code}: {response.json().get('message', 'No message provided')}"
+                    )
+
+            except requests.exceptions.RequestException as e:
+                status["openai"] = False
+                print(f"OpenAI API Error: {e}")
+
         # Get the top 5 processes by memory usage
         for proc in psutil.process_iter(["pid", "name", "memory_info"]):
             try:
@@ -147,6 +169,23 @@ def check_status(request):
             except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                 pass
 
+        # Get memory usage by module for the current Python process
+        current_process = psutil.Process(os.getpid())
+        memory_info_by_module = {}
+        for module in sys.modules.values():
+            try:
+                module_name = getattr(module, "__name__", "unknown")
+                module_memory = sum(
+                    obj.nbytes for obj in gc.get_objects() if isinstance(obj, type(module))
+                )
+                memory_info_by_module[module_name] = module_memory
+            except Exception:
+                pass
+
+        status["memory_by_module"] = sorted(
+            memory_info_by_module.items(), key=lambda x: x[1], reverse=True
+        )[:10]  # Top 10 modules by memory usage
+
         status["top_memory_consumers"] = sorted(
             status["top_memory_consumers"],
             key=lambda x: x["memory_info"]["rss"],