3e9ccafad2
- Model renames: gpt-5.2 → gpt-5.4-2026-03-05, gemini-3-pro-preview → gemini-3.1-pro-preview; retire gpt-4.1 via alias fallback - New: llm_usage_context.py (ContextVar-based attribution), model_pricing.py (tiered pricing + 60s cache), usage_event.py (append-only telemetry), quota.py (user/FG quota enforcement with 80% warning) - Wire _record_usage into all 3 LLM methods; set_llm_context at every service entry point - Fix admin_required decorator (was sync, never awaited User.find_by_id); add active_required and with_user_context decorators - Inject user_id into ContextVar from JWT on every authenticated request - Add DB indexes for usage_events, model_pricing, users collections - Seed script for model pricing (gpt-5.4 single-tier, gemini-3.1 two-tier 200k threshold) - Fix parse_json_response NameError (logger undefined at module level) - 70 passing tests: conftest.py with sys.modules stubs, test_usage_infrastructure.py (52 tests), rewrite stale test_llm_service.py (18 tests) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
89 lines
No EOL
3.9 KiB
Python
Executable file
89 lines
No EOL
3.9 KiB
Python
Executable file
from motor.motor_asyncio import AsyncIOMotorClient
|
|
from pymongo import MongoClient
|
|
import os
|
|
import logging
|
|
|
|
# Global Motor client singleton - per event loop
|
|
_motor_clients = {} # event_loop_id -> (client, database)
|
|
|
|
async def get_db():
|
|
"""Get database connection using singleton Motor client per event loop."""
|
|
import asyncio
|
|
|
|
# Get current event loop to ensure Motor client affinity
|
|
try:
|
|
current_loop = asyncio.get_running_loop()
|
|
loop_id = id(current_loop)
|
|
except RuntimeError:
|
|
raise RuntimeError("get_db() must be called from within an async context")
|
|
|
|
# Return cached database for this event loop if available
|
|
if loop_id in _motor_clients:
|
|
client, database = _motor_clients[loop_id]
|
|
return database
|
|
|
|
# Read MongoDB connection from environment
|
|
mongo_uri = os.environ.get('MONGO_URI')
|
|
mongo_user = os.environ.get('MONGO_USER')
|
|
mongo_pass = os.environ.get('MONGO_PASS')
|
|
mongo_host = os.environ.get('MONGO_HOST', 'localhost')
|
|
mongo_port = os.environ.get('MONGO_PORT', '27017')
|
|
|
|
# Build URI: prefer MONGO_URI, fall back to host+port with optional credentials
|
|
if not mongo_uri:
|
|
if mongo_user and mongo_pass:
|
|
mongo_uri = f"mongodb://{mongo_user}:{mongo_pass}@{mongo_host}:{mongo_port}/semblance_db?authSource=admin"
|
|
else:
|
|
mongo_uri = f"mongodb://{mongo_host}:{mongo_port}"
|
|
|
|
try:
|
|
motor_client = AsyncIOMotorClient(mongo_uri, serverSelectionTimeoutMS=5000)
|
|
database = motor_client.semblance_db
|
|
await database.command('ping')
|
|
logging.info("Successfully connected to MongoDB")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to connect to MongoDB: {e}. Check MONGO_URI in backend/.env.") from e
|
|
|
|
_motor_clients[loop_id] = (motor_client, database)
|
|
|
|
# Ensure indexes exist (idempotent)
|
|
try:
|
|
await database.users.create_index("username", unique=True, background=True)
|
|
await database.users.create_index("email", unique=True, background=True)
|
|
await database.users.create_index("role", background=True)
|
|
await database.users.create_index([("is_active", 1), ("username", 1)], background=True)
|
|
await database.personas.create_index("created_by", background=True)
|
|
await database.focus_groups.create_index("created_by", background=True)
|
|
await database.folders.create_index("created_by", background=True)
|
|
await database.folders.create_index("parent_folder_id", background=True)
|
|
# usage_events indexes
|
|
await database.usage_events.create_index([("user_id", 1), ("ts", -1)], background=True)
|
|
await database.usage_events.create_index([("focus_group_id", 1), ("ts", -1)], background=True)
|
|
await database.usage_events.create_index([("ts", -1)], background=True)
|
|
await database.usage_events.create_index([("feature", 1), ("ts", -1)], background=True)
|
|
await database.usage_events.create_index([("model", 1), ("ts", -1)], background=True)
|
|
await database.usage_events.create_index([("status", 1), ("ts", -1)], background=True)
|
|
# model_pricing indexes
|
|
await database.model_pricing.create_index([("model", 1), ("effective_from", -1)], background=True)
|
|
except Exception as e:
|
|
logging.warning(f"Index creation warning (non-fatal): {e}")
|
|
|
|
return database
|
|
|
|
|
|
def close_db_connections():
|
|
"""Close all Motor clients and their PyMongo background threads."""
|
|
global _motor_clients
|
|
|
|
closed_count = 0
|
|
for loop_id, (client, database) in _motor_clients.items():
|
|
try:
|
|
client.close()
|
|
closed_count += 1
|
|
except Exception as e:
|
|
logging.warning(f"Error closing Motor client for loop {loop_id}: {e}")
|
|
|
|
if closed_count > 0:
|
|
logging.info(f"🗄️ Closed {closed_count} Motor clients - PyMongo threads should stop")
|
|
|
|
_motor_clients.clear() |