agent_tracker/database.py
nickviljoen 32b08f8b0c Add Prompt Audit & Auto-Classification feature with Gemini integration
Adds Gemini-powered agent classification system that analyzes agent instructions
to determine category, risk level, discipline, and client detection. Includes
admin Prompt Audit tab, audit review workflow, and auto-classification on sync.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 18:29:35 +02:00

40 lines
1.5 KiB
Python

from motor.motor_asyncio import AsyncIOMotorClient
from dotenv import load_dotenv
import os
load_dotenv()
MONGODB_URI = os.getenv("MONGODB_URI", "mongodb://localhost:27017")
MONGODB_DBNAME = os.getenv("MONGODB_DBNAME", "agenthub_db")
client = AsyncIOMotorClient(MONGODB_URI)
db = client[MONGODB_DBNAME]
users_collection = db.get_collection("users")
agents_collection = db.get_collection("agents")
agent_usage_collection = db.get_collection("agent_usage")
notifications_collection = db.get_collection("token_notifications")
agent_ratings_collection = db.get_collection("agent_ratings")
audit_history_collection = db.get_collection("audit_history")
async def ensure_indexes():
"""Create database indexes for performance"""
try:
await agent_ratings_collection.create_index(
[("agent_id", 1), ("user_id", 1)],
unique=True
)
await agents_collection.create_index([("verification_status", 1)])
await agents_collection.create_index([("audit_status", 1)])
await audit_history_collection.create_index([("agent_id", 1), ("audit_date", -1)])
print("Database indexes ensured successfully")
except Exception as e:
print(f"Warning: Failed to create indexes: {e}")
async def check_database_health():
"""Check MongoDB connection health"""
try:
# Attempt to ping the database
await client.admin.command('ping')
return {"status": "connected", "healthy": True}
except Exception as e:
return {"status": "disconnected", "healthy": False, "error": str(e)}