Leopoldadmin/leocrm
1
0
Fork 0
Code Issues 363 Pull Requests Actions 1 Packages Projects Releases 1 Wiki Activity

feat(B.13-B.17+B.16): Error-Handling-Infra, Observability, Graceful Shutdown, Cost-Cap, API Versioning

Browse Source 
B.13 Error-Handling-Infrastruktur:
- ErrorCategory Enum (TRANSIENT/PERMANENT/PARTIAL), ApiError erweitert
- Einheitliches Error-Response-Format: {code, detail, field, trace_id, retryable, category}
- 3 FastAPI Exception-Handler (ApiError, HTTPException, unhandled)
- classify_exception() Helper, 6 neue Error-Codes
- 28 Tests in test_error_handling.py

B.14 Observability & trace_id-Korrelation:
- trace_id pro Request (UUID4 short) in structlog contextvars
- X-Trace-Id Response-Header
- Sensitive Fields structlog processor
- llm_complete()/llm_embed() akzeptieren trace_id kwarg
- 12 Tests in test_observability.py

B.15 Graceful Shutdown & Connection Draining:
- _shutdown_event + _inflight_requests Tracking in main.py
- drain_all_connections() in ws_helpers.py
- Worker on_shutdown pausiert WorkflowInstances (status=paused)
- 8 Tests in test_graceful_shutdown.py

B.16 API Versioning Strategie:
- Plugin-Dev-Guide Kapitel 30: URL-basiertes Versioning, Breaking Change Prozess

B.17 Cost Overrun Protection:
- llm_monthly_budget_usd + llm_hard_cutoff Settings
- _check_tenant_budget() vor jedem LLM-Call
- _track_tenant_cost() in Redis (INCRBYFLOAT)
- _check_cost_alerts() bei 50%/80%/100% -> post_system_message()
- 20 Tests in test_cost_protection.py

Total: 68 neue Tests, alle grün. Keine Regressionen.
This commit is contained in:
Agent Zero
2026-08-13 21:33:14 +02:00
parent 1baa9481a2
commit b5546ea7bd
13 changed files with 1461 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/ai/llm_client.py
+209 -4
View File
@@ -21,6 +21,7 @@ import json
import logging
import os
import uuid
from datetime import datetime, timezone
from typing import Any, TYPE_CHECKING
import litellm
@@ -94,6 +95,175 @@ _PERMANENT_KEYWORDS = frozenset(
# ──────────────────────────────────────────────────────────────────────────
# ── Cost Overrun Protection (B.17) ───────────────────────────────────────────
def _get_cost_key(tenant_id: uuid.UUID | str) -> str:
"""Build the Redis cost-tracking key for the current month."""
now = datetime.now(timezone.utc)
month_str = now.strftime("%Y-%m")
return f"cost:tenant:{tenant_id}:month:{month_str}"
async def get_tenant_monthly_cost(tenant_id: uuid.UUID | str) -> float:
"""Get the accumulated LLM cost for a tenant in the current month.
Reads from Redis key ``cost:tenant:{tenant_id}:month:{YYYY-MM}``.
Returns 0.0 if Redis is unavailable or the key doesn't exist.
"""
try:
from app.core.auth import get_redis
r = get_redis()
key = _get_cost_key(tenant_id)
val = await r.get(key)
return float(val) if val else 0.0
except Exception:
logger.debug("Failed to get tenant monthly cost from Redis")
return 0.0
async def _check_tenant_budget(
tenant_id: uuid.UUID | str | None,
db: AsyncSession | None = None,
) -> None:
"""Check if tenant has exceeded their LLM budget.
Raises ``ValueError("Tenant LLM budget exceeded")`` if the tenant's
accumulated monthly cost exceeds the configured budget and
``llm_hard_cutoff`` is enabled.
"""
if tenant_id is None:
return
from app.config import get_settings
settings = get_settings()
if settings.llm_monthly_budget_usd <= 0:
return # No budget limit configured
current_cost = await get_tenant_monthly_cost(tenant_id)
if current_cost >= settings.llm_monthly_budget_usd:
if settings.llm_hard_cutoff:
logger.warning(
"Tenant %s LLM budget exceeded: $%.2f >= $%.2f (hard cutoff)",
tenant_id, current_cost, settings.llm_monthly_budget_usd,
)
raise ValueError("Tenant LLM budget exceeded")
else:
logger.warning(
"Tenant %s LLM budget exceeded: $%.2f >= $%.2f (soft limit, no cutoff)",
tenant_id, current_cost, settings.llm_monthly_budget_usd,
)
async def _track_tenant_cost(
tenant_id: uuid.UUID | str | None,
cost_usd: float,
db: AsyncSession | None = None,
) -> float:
"""Track LLM cost in Redis and check for alert thresholds.
Increments ``cost:tenant:{tenant_id}:month:{YYYY-MM}`` by ``cost_usd``.
Returns the new total cost for the month.
"""
if tenant_id is None or cost_usd <= 0:
return 0.0
try:
from app.core.auth import get_redis
r = get_redis()
key = _get_cost_key(tenant_id)
new_total = await r.incrbyfloat(key, cost_usd)
# Set TTL to 35 days so old months auto-expire
await r.expire(key, 35 * 24 * 3600)
# Check cost alert thresholds
await _check_cost_alerts(tenant_id, new_total, db)
return float(new_total)
except Exception:
logger.debug("Failed to track tenant cost in Redis")
return 0.0
async def _check_cost_alerts(
tenant_id: uuid.UUID | str,
current_cost: float,
db: AsyncSession | None = None,
) -> None:
"""Send cost alerts at 50%, 80%, and 100% of the tenant budget.
Each threshold alert is sent only once per month (tracked via Redis flag
``cost:alerted:{tenant_id}:{threshold}``).
"""
from app.config import get_settings
settings = get_settings()
budget = settings.llm_monthly_budget_usd
if budget <= 0:
return
thresholds = [0.50, 0.80, 1.00]
now = datetime.now(timezone.utc)
month_str = now.strftime("%Y-%m")
try:
from app.core.auth import get_redis
r = get_redis()
for threshold in thresholds:
threshold_cost = budget * threshold
if current_cost >= threshold_cost:
alert_key = f"cost:alerted:{tenant_id}:{threshold}:{month_str}"
already_alerted = await r.set(alert_key, "1", nx=True, ex=35 * 24 * 3600)
if already_alerted:
# This is a new alert — send notification
pct = int(threshold * 100)
logger.info(
"Cost alert: tenant %s reached %d%% of budget ($%.2f / $%.2f)",
tenant_id, pct, current_cost, budget,
)
# Best-effort system notification
if db is not None:
try:
from app.core.notifications import post_system_message
# Need a user_id — try to find an admin for this tenant
from sqlalchemy import select as sa_select
from app.models.user import User, UserTenant
from app.models.role import Role
async with db.begin_nested() if db.in_transaction() else _NoopCtx():
result = await db.execute(
sa_select(User.id).join(UserTenant, UserTenant.user_id == User.id)
.where(UserTenant.tenant_id == tenant_id)
.where(User.is_active == True) # noqa: E712
.limit(1)
)
admin_row = result.first()
if admin_row:
await post_system_message(
db=db,
tenant_id=tenant_id if isinstance(tenant_id, uuid.UUID) else uuid.UUID(str(tenant_id)),
user_id=admin_row[0],
message_type="cost_alert",
title=f"LLM Cost Alert: {pct}% of budget reached",
body=f"Current monthly LLM cost: ${current_cost:.2f} / ${budget:.2f} ({pct}%).",
severity="warning" if pct < 100 else "error",
)
except Exception:
logger.debug("Failed to send cost alert notification")
except Exception:
logger.debug("Failed to check cost alerts")
class _NoopCtx:
"""No-op async context manager for optional transaction nesting."""
async def __aenter__(self):
return self
async def __aexit__(self, *args):
return False
async def get_api_credentials(
db: AsyncSession | None,
tenant_id: uuid.UUID | None,
@@ -291,6 +461,9 @@ async def llm_complete(
response_format: dict[str, Any] | None = None,
timeout: int = DEFAULT_TIMEOUT,
max_retries: int = DEFAULT_MAX_RETRIES,
trace_id: str | None = None,
tenant_id: uuid.UUID | str | None = None,
db: AsyncSession | None = None,
) -> dict[str, Any]:
"""Generic LLM chat completion via LiteLLM with retry and cost tracking.
@@ -310,17 +483,28 @@ async def llm_complete(
response_format: Optional response format spec (e.g. JSON mode).
timeout: Request timeout in seconds (default 30).
max_retries: Max retry attempts for transient errors (default 2).
trace_id: Optional trace ID for request correlation/logging.
tenant_id: Optional tenant ID for cost tracking and budget enforcement.
db: Optional DB session for cost alert notifications.
Returns:
Dict with keys: ``content``, ``usage``, ``cost_usd``, ``model``,
``raw_response`` (the LiteLLM response object for advanced use).
Raises:
ValueError: If tenant LLM budget is exceeded (hard cutoff).
Exception: Permanent errors or after exhausting retries.
"""
# Check tenant budget before making the call
await _check_tenant_budget(tenant_id, db)
# Build LiteLLM model string
litellm_model = build_model(model, provider)
# Log trace_id correlation if provided
if trace_id:
logger.debug("llm_complete trace_id=%s model=%s", trace_id, litellm_model)
# Build kwargs
kwargs: dict[str, Any] = {
"model": litellm_model,
@@ -348,12 +532,16 @@ async def llm_complete(
cost_usd = _extract_cost_usd(response, litellm_model)
logger.debug(
"llm_complete success: model=%s tokens=%d cost=$%.6f attempt=%d",
"llm_complete success: model=%s tokens=%d cost=$%.6f attempt=%d trace_id=%s",
litellm_model,
usage["total_tokens"],
cost_usd,
attempt + 1,
trace_id or "-",
)
# Track cost in Redis for tenant budget enforcement
if tenant_id is not None and cost_usd > 0:
await _track_tenant_cost(tenant_id, cost_usd, db)
return {
"content": content,
"usage": usage,
@@ -400,6 +588,7 @@ async def llm_embed(
provider: str | None = None,
dimensions: int | None = None,
timeout: int = DEFAULT_TIMEOUT,
trace_id: str | None = None,
) -> list[list[float]]:
"""Generic text embedding via LiteLLM.
@@ -409,18 +598,25 @@ async def llm_embed(
Args:
texts: Single text string or list of texts to embed.
model: Embedding model name (default: ``openai/text-embedding-3-small``).
db: Optional DB session for API key lookup.
tenant_id: Optional tenant ID for API key lookup.
db: Optional DB session for API key lookup and cost alert notifications.
tenant_id: Optional tenant ID for cost tracking and budget enforcement.
api_key: Override API key. If ``None``, uses env/DB lookup.
api_base: Override API base URL.
provider: Provider prefix override.
dimensions: Override embedding dimensions.
timeout: Request timeout in seconds (default 30).
trace_id: Optional trace ID for request correlation/logging.
Returns:
List of embedding vectors (each a list of floats). For a single
text input, returns a one-element list.
"""
# Check tenant budget before making the call
await _check_tenant_budget(tenant_id, db)
if trace_id:
logger.debug("llm_embed trace_id=%s", trace_id)
# Normalise to list input
single_input = isinstance(texts, str)
text_list = [texts] if single_input else texts
@@ -467,11 +663,20 @@ async def llm_embed(
try:
response = await litellm.aembedding(**litellm_kwargs)
embeddings = [d["embedding"] for d in response.data]
# Track embedding cost for tenant budget enforcement
if tenant_id is not None:
try:
emb_cost = _extract_cost_usd(response, litellm_model)
if emb_cost > 0:
await _track_tenant_cost(tenant_id, emb_cost, db)
except Exception:
pass
logger.debug(
"llm_embed success: model=%s count=%d dims=%d",
"llm_embed success: model=%s count=%d dims=%d trace_id=%s",
litellm_model,
len(embeddings),
len(embeddings[0]) if embeddings else 0,
trace_id or "-",
)
return embeddings
except Exception: