feat: Add presence heartbeat for Matrix online status
- matrix-gateway: POST /internal/matrix/presence/online endpoint - usePresenceHeartbeat hook with activity tracking - Auto away after 5 min inactivity - Offline on page close/visibility change - Integrated in MatrixChatRoom component
This commit is contained in:
Apple
221
services/usage-engine/main.py
Normal file
221
services/usage-engine/main.py
Normal file
@@ -0,0 +1,221 @@
|
||||
"""
|
||||
DAARION Usage Engine
|
||||
Port: 7013
|
||||
Collects and reports usage metrics (LLM, Tools, Agents, Messages)
|
||||
"""
|
||||
import os
|
||||
import asyncio
|
||||
import asyncpg
|
||||
import nats
|
||||
from contextlib import asynccontextmanager
|
||||
from fastapi import FastAPI, HTTPException, Query
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from typing import Optional
|
||||
|
||||
from models import UsageQueryRequest, UsageQueryResponse
|
||||
from collectors import UsageCollector
|
||||
from aggregators import UsageAggregator
|
||||
|
||||
# ============================================================================
|
||||
# Configuration
|
||||
# ============================================================================
|
||||
|
||||
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://postgres:postgres@localhost:5432/daarion")
|
||||
NATS_URL = os.getenv("NATS_URL", "nats://nats:4222")
|
||||
|
||||
# ============================================================================
|
||||
# Global State
|
||||
# ============================================================================
|
||||
|
||||
db_pool: Optional[asyncpg.Pool] = None
|
||||
nc: Optional[nats.NATS] = None
|
||||
collector: Optional[UsageCollector] = None
|
||||
aggregator: Optional[UsageAggregator] = None
|
||||
|
||||
# ============================================================================
|
||||
# App Setup
|
||||
# ============================================================================
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Startup and shutdown"""
|
||||
global db_pool, nc, collector, aggregator
|
||||
|
||||
print("🚀 Starting Usage Engine...")
|
||||
|
||||
# Database
|
||||
db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
|
||||
print("✅ Database pool created")
|
||||
|
||||
# NATS
|
||||
try:
|
||||
nc = await nats.connect(NATS_URL)
|
||||
print(f"✅ Connected to NATS at {NATS_URL}")
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to connect to NATS: {e}")
|
||||
nc = None
|
||||
|
||||
# Collector
|
||||
if nc:
|
||||
collector = UsageCollector(nc, db_pool)
|
||||
await collector.start()
|
||||
else:
|
||||
print("⚠️ NATS not available, collector disabled")
|
||||
|
||||
# Aggregator
|
||||
aggregator = UsageAggregator(db_pool)
|
||||
print("✅ Aggregator ready")
|
||||
|
||||
print("✅ Usage Engine ready")
|
||||
|
||||
yield
|
||||
|
||||
# Shutdown
|
||||
print("🛑 Shutting down Usage Engine...")
|
||||
if collector:
|
||||
await collector.stop()
|
||||
if nc:
|
||||
await nc.close()
|
||||
if db_pool:
|
||||
await db_pool.close()
|
||||
|
||||
app = FastAPI(
|
||||
title="DAARION Usage Engine",
|
||||
version="1.0.0",
|
||||
description="Usage tracking and reporting for LLM, Tools, Agents",
|
||||
lifespan=lifespan
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# ============================================================================
|
||||
# API Endpoints
|
||||
# ============================================================================
|
||||
|
||||
@app.get("/internal/usage/summary", response_model=UsageQueryResponse)
|
||||
async def get_usage_summary(
|
||||
microdao_id: Optional[str] = Query(None),
|
||||
agent_id: Optional[str] = Query(None),
|
||||
period_hours: int = Query(24, ge=1, le=720)
|
||||
):
|
||||
"""
|
||||
Get aggregated usage summary
|
||||
|
||||
Query parameters:
|
||||
- microdao_id: Filter by microDAO (optional)
|
||||
- agent_id: Filter by agent (optional)
|
||||
- period_hours: Time period (1-720 hours, default 24)
|
||||
"""
|
||||
|
||||
if not aggregator:
|
||||
raise HTTPException(500, "Aggregator not initialized")
|
||||
|
||||
# Get summary
|
||||
summary = await aggregator.get_summary(
|
||||
microdao_id=microdao_id,
|
||||
agent_id=agent_id,
|
||||
period_hours=period_hours
|
||||
)
|
||||
|
||||
# Get breakdowns
|
||||
models = await aggregator.get_model_breakdown(
|
||||
microdao_id=microdao_id,
|
||||
period_hours=period_hours
|
||||
)
|
||||
|
||||
agents = await aggregator.get_agent_breakdown(
|
||||
microdao_id=microdao_id,
|
||||
period_hours=period_hours
|
||||
)
|
||||
|
||||
tools = await aggregator.get_tool_breakdown(
|
||||
microdao_id=microdao_id,
|
||||
period_hours=period_hours
|
||||
)
|
||||
|
||||
return UsageQueryResponse(
|
||||
summary=summary,
|
||||
models=models,
|
||||
agents=agents,
|
||||
tools=tools
|
||||
)
|
||||
|
||||
@app.get("/internal/usage/models")
|
||||
async def get_model_usage(
|
||||
microdao_id: Optional[str] = Query(None),
|
||||
period_hours: int = Query(24, ge=1, le=720)
|
||||
):
|
||||
"""Get usage breakdown by model"""
|
||||
|
||||
if not aggregator:
|
||||
raise HTTPException(500, "Aggregator not initialized")
|
||||
|
||||
models = await aggregator.get_model_breakdown(
|
||||
microdao_id=microdao_id,
|
||||
period_hours=period_hours
|
||||
)
|
||||
|
||||
return {"models": models}
|
||||
|
||||
@app.get("/internal/usage/agents")
|
||||
async def get_agent_usage(
|
||||
microdao_id: Optional[str] = Query(None),
|
||||
period_hours: int = Query(24, ge=1, le=720)
|
||||
):
|
||||
"""Get usage breakdown by agent"""
|
||||
|
||||
if not aggregator:
|
||||
raise HTTPException(500, "Aggregator not initialized")
|
||||
|
||||
agents = await aggregator.get_agent_breakdown(
|
||||
microdao_id=microdao_id,
|
||||
period_hours=period_hours
|
||||
)
|
||||
|
||||
return {"agents": agents}
|
||||
|
||||
@app.get("/internal/usage/tools")
|
||||
async def get_tool_usage(
|
||||
microdao_id: Optional[str] = Query(None),
|
||||
period_hours: int = Query(24, ge=1, le=720)
|
||||
):
|
||||
"""Get usage breakdown by tool"""
|
||||
|
||||
if not aggregator:
|
||||
raise HTTPException(500, "Aggregator not initialized")
|
||||
|
||||
tools = await aggregator.get_tool_breakdown(
|
||||
microdao_id=microdao_id,
|
||||
period_hours=period_hours
|
||||
)
|
||||
|
||||
return {"tools": tools}
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
"""Health check"""
|
||||
return {
|
||||
"status": "ok",
|
||||
"service": "usage-engine",
|
||||
"nats_connected": nc is not None,
|
||||
"collector_active": collector is not None,
|
||||
"aggregator_ready": aggregator is not None
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Run
|
||||
# ============================================================================
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=7013)
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user