"""
FastAPI application for ClientSphere RAG Backend.
Provides endpoints for knowledge base management and chat.
"""
from fastapi import FastAPI, File, UploadFile, HTTPException, Form, BackgroundTasks, Request, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from pathlib import Path
import shutil
import uuid
from datetime import datetime
from typing import Optional
import logging
from app.config import settings
from app.middleware.auth import get_auth_context, require_auth
from app.middleware.rate_limit import (
limiter,
get_tenant_rate_limit_key,
RateLimitExceeded,
_rate_limit_exceeded_handler
)
from app.models.schemas import (
UploadResponse,
ChatRequest,
ChatResponse,
KnowledgeBaseStats,
HealthResponse,
DocumentStatus,
Citation,
)
from app.models.billing_schemas import (
UsageResponse,
PlanLimitsResponse,
CostReportResponse,
SetPlanRequest
)
from app.rag.ingest import parser
from app.rag.chunking import chunker
from app.rag.embeddings import get_embedding_service
from app.rag.vectorstore import get_vector_store
from app.rag.retrieval import get_retrieval_service
from app.rag.answer import get_answer_service
from app.db.database import get_db, init_db
from app.billing.quota import check_quota, ensure_tenant_exists
from app.billing.usage_tracker import track_usage
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Initialize FastAPI app
app = FastAPI(
title=settings.APP_NAME,
description="RAG-based customer support chatbot API",
version="1.0.0",
)
# Initialize database on startup
@app.on_event("startup")
async def startup_event():
"""Initialize database on application startup."""
init_db()
logger.info("Database initialized")
# Configure CORS - SECURITY: Restrict in production
if settings.ALLOWED_ORIGINS == "*":
allowed_origins = ["*"]
else:
# Split by comma and strip whitespace
allowed_origins = [origin.strip() for origin in settings.ALLOWED_ORIGINS.split(",") if origin.strip()]
# Default to allowing localhost if no origins specified
if not allowed_origins or allowed_origins == ["*"]:
allowed_origins = ["*"] # Allow all in dev mode
logger.info(f"CORS configured with origins: {allowed_origins}")
app.add_middleware(
CORSMiddleware,
allow_origins=allowed_origins,
allow_credentials=True,
allow_methods=["GET", "POST", "DELETE", "OPTIONS"], # Include OPTIONS for preflight
allow_headers=["Content-Type", "Authorization", "X-Tenant-Id", "X-User-Id"], # Include auth headers
)
# Configure rate limiting
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
# Add exception handler for validation errors
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
"""Handle request validation errors with detailed logging."""
body = await request.body()
logger.error(f"Request validation error: {exc.errors()}")
logger.error(f"Request body (raw): {body}")
logger.error(f"Request headers: {dict(request.headers)}")
return JSONResponse(
status_code=422,
content={"detail": exc.errors(), "body": body.decode('utf-8', errors='ignore')}
)
# Add exception handler for validation errors
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
"""Handle request validation errors with detailed logging."""
logger.error(f"Request validation error: {exc.errors()}")
logger.error(f"Request body: {await request.body()}")
return JSONResponse(
status_code=422,
content={"detail": exc.errors(), "body": str(await request.body())}
)
# ============== Health & Status Endpoints ==============
@app.get("/", response_model=HealthResponse)
async def root():
"""Root endpoint with basic info."""
return HealthResponse(
status="ok",
version="1.0.0",
vector_db_connected=True,
llm_configured=bool(settings.GEMINI_API_KEY or settings.OPENAI_API_KEY)
)
@app.get("/health", response_model=HealthResponse)
async def health_check():
"""Health check endpoint."""
try:
vector_store = get_vector_store()
stats = vector_store.get_stats()
return HealthResponse(
status="healthy",
version="1.0.0",
vector_db_connected=True,
llm_configured=bool(settings.GEMINI_API_KEY or settings.OPENAI_API_KEY)
)
except Exception as e:
logger.error(f"Health check failed: {e}")
return HealthResponse(
status="unhealthy",
version="1.0.0",
vector_db_connected=False,
llm_configured=False
)
@app.get("/health/live")
async def liveness():
"""Kubernetes liveness probe - always returns alive."""
return {"status": "alive"}
@app.get("/health/ready")
async def readiness():
"""Kubernetes readiness probe - checks dependencies."""
checks = {
"vector_db": False,
"llm_configured": bool(settings.GEMINI_API_KEY or settings.OPENAI_API_KEY)
}
# Check vector DB connection
try:
vector_store = get_vector_store()
vector_store.get_stats()
checks["vector_db"] = True
except Exception as e:
logger.warning(f"Vector DB check failed: {e}")
checks["vector_db"] = False
# All checks must pass
if all(checks.values()):
return {"status": "ready", "checks": checks}
else:
from fastapi import HTTPException
raise HTTPException(status_code=503, detail={"status": "not_ready", "checks": checks})
# ============== Knowledge Base Endpoints ==============
@app.post("/kb/upload", response_model=UploadResponse)
@limiter.limit("20/hour", key_func=get_tenant_rate_limit_key)
async def upload_document(
background_tasks: BackgroundTasks,
request: Request,
file: UploadFile = File(...),
tenant_id: Optional[str] = Form(None), # Optional in dev, ignored in prod
user_id: Optional[str] = Form(None), # Optional in dev, ignored in prod
kb_id: str = Form(...)
):
"""
Upload a document to the knowledge base.
- Saves file to disk
- Parses and chunks the document
- Generates embeddings
- Stores in vector database
"""
# SECURITY: Extract tenant_id from auth token in production
if settings.ENV == "prod":
auth_context = await require_auth(request)
tenant_id = auth_context.get("tenant_id")
if not tenant_id:
raise HTTPException(
status_code=403,
detail="tenant_id must come from authentication token in production mode"
)
elif not tenant_id:
raise HTTPException(
status_code=400,
detail="tenant_id is required"
)
# Validate file type
file_ext = Path(file.filename).suffix.lower()
if file_ext not in parser.SUPPORTED_EXTENSIONS:
raise HTTPException(
status_code=400,
detail=f"Unsupported file type: {file_ext}. Supported: {parser.SUPPORTED_EXTENSIONS}"
)
# Validate file size (SECURITY)
file.file.seek(0, 2) # Seek to end
file_size = file.file.tell()
file.file.seek(0) # Reset to start
max_size_bytes = settings.MAX_FILE_SIZE_MB * 1024 * 1024
if file_size > max_size_bytes:
raise HTTPException(
status_code=400,
detail=f"File too large. Maximum size: {settings.MAX_FILE_SIZE_MB}MB"
)
# Generate document ID
doc_id = f"{tenant_id}_{kb_id}_{uuid.uuid4().hex[:8]}"
# Save file to uploads directory
upload_path = settings.UPLOADS_DIR / f"{doc_id}_{file.filename}"
try:
with open(upload_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
logger.info(f"Saved file: {upload_path}")
except Exception as e:
logger.error(f"Error saving file: {e}")
raise HTTPException(status_code=500, detail="Failed to save file")
# Process document in background
background_tasks.add_task(
process_document,
upload_path,
tenant_id, # CRITICAL: Multi-tenant isolation
user_id,
kb_id,
file.filename,
doc_id
)
return UploadResponse(
success=True,
message="Document upload started. Processing in background.",
document_id=doc_id,
file_name=file.filename,
chunks_created=0,
status=DocumentStatus.PROCESSING
)
async def process_document(
file_path: Path,
tenant_id: str, # CRITICAL: Multi-tenant isolation
user_id: str,
kb_id: str,
original_filename: str,
document_id: str
):
"""
Background task to process an uploaded document.
"""
try:
logger.info(f"Processing document: {original_filename}")
# Parse document
parsed_doc = parser.parse(file_path)
logger.info(f"Parsed document: {len(parsed_doc.text)} characters")
# Chunk document
chunks = chunker.chunk_text(
parsed_doc.text,
page_numbers=parsed_doc.page_map
)
logger.info(f"Created {len(chunks)} chunks")
if not chunks:
logger.warning(f"No chunks created from {original_filename}")
return
# Create metadata for each chunk
metadatas = []
chunk_ids = []
chunk_texts = []
for chunk in chunks:
metadata = chunker.create_chunk_metadata(
chunk=chunk,
tenant_id=tenant_id, # CRITICAL: Multi-tenant isolation
kb_id=kb_id,
user_id=user_id,
file_name=original_filename,
file_type=parsed_doc.file_type,
total_chunks=len(chunks),
document_id=document_id
)
metadatas.append(metadata)
chunk_ids.append(metadata["chunk_id"])
chunk_texts.append(chunk.content)
# Generate embeddings
embedding_service = get_embedding_service()
embeddings = embedding_service.embed_texts(chunk_texts)
logger.info(f"Generated {len(embeddings)} embeddings")
# Store in vector database
vector_store = get_vector_store()
vector_store.add_documents(
documents=chunk_texts,
embeddings=embeddings,
metadatas=metadatas,
ids=chunk_ids
)
logger.info(f"Successfully processed {original_filename}: {len(chunks)} chunks stored")
except Exception as e:
logger.error(f"Error processing document {original_filename}: {e}")
raise
@app.get("/kb/stats", response_model=KnowledgeBaseStats)
async def get_kb_stats(
request: Request,
tenant_id: Optional[str] = None, # Optional in dev, ignored in prod
kb_id: Optional[str] = None,
user_id: Optional[str] = None # Optional in dev, ignored in prod
):
"""Get statistics for a knowledge base."""
# SECURITY: Get tenant_id and user_id from auth context
auth_context = await get_auth_context(request)
tenant_id_from_auth = auth_context.get("tenant_id")
user_id_from_auth = auth_context.get("user_id")
if settings.ENV == "prod":
if not tenant_id_from_auth or not user_id_from_auth:
raise HTTPException(
status_code=403,
detail="tenant_id and user_id must come from authentication token in production mode"
)
tenant_id = tenant_id_from_auth
user_id = user_id_from_auth
else:
tenant_id = tenant_id or tenant_id_from_auth
user_id = user_id or user_id_from_auth
if not tenant_id or not kb_id or not user_id:
raise HTTPException(
status_code=400,
detail="tenant_id, kb_id, and user_id are required"
)
try:
vector_store = get_vector_store()
stats = vector_store.get_stats(tenant_id=tenant_id, kb_id=kb_id, user_id=user_id)
return KnowledgeBaseStats(
tenant_id=tenant_id, # CRITICAL: Multi-tenant isolation
kb_id=kb_id,
user_id=user_id,
total_documents=len(stats.get("file_names", [])),
total_chunks=stats.get("total_chunks", 0),
file_names=stats.get("file_names", []),
last_updated=datetime.utcnow()
)
except Exception as e:
logger.error(f"Error getting KB stats: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.delete("/kb/document")
async def delete_document(
request: Request,
tenant_id: Optional[str] = None, # Optional in dev, ignored in prod
kb_id: Optional[str] = None,
user_id: Optional[str] = None, # Optional in dev, ignored in prod
file_name: Optional[str] = None
):
"""Delete a document from the knowledge base."""
# SECURITY: Get tenant_id and user_id from auth context
auth_context = await get_auth_context(request)
tenant_id_from_auth = auth_context.get("tenant_id")
user_id_from_auth = auth_context.get("user_id")
if settings.ENV == "prod":
if not tenant_id_from_auth or not user_id_from_auth:
raise HTTPException(
status_code=403,
detail="tenant_id and user_id must come from authentication token in production mode"
)
tenant_id = tenant_id_from_auth
user_id = user_id_from_auth
else:
tenant_id = tenant_id or tenant_id_from_auth
user_id = user_id or user_id_from_auth
if not tenant_id or not kb_id or not user_id or not file_name:
raise HTTPException(
status_code=400,
detail="tenant_id, kb_id, user_id, and file_name are required (provide via headers or query params)"
)
try:
vector_store = get_vector_store()
deleted = vector_store.delete_by_filter({
"tenant_id": tenant_id, # CRITICAL: Multi-tenant isolation
"kb_id": kb_id,
"user_id": user_id,
"file_name": file_name
})
return {
"success": True,
"message": f"Deleted {deleted} chunks",
"file_name": file_name
}
except Exception as e:
logger.error(f"Error deleting document: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.delete("/kb/clear")
async def clear_kb(
request: Request,
tenant_id: Optional[str] = None, # Optional in dev, ignored in prod
kb_id: Optional[str] = None,
user_id: Optional[str] = None # Optional in dev, ignored in prod
):
"""Clear all documents from a knowledge base."""
# SECURITY: Get tenant_id and user_id from auth context
auth_context = await get_auth_context(request)
tenant_id_from_auth = auth_context.get("tenant_id")
user_id_from_auth = auth_context.get("user_id")
if settings.ENV == "prod":
if not tenant_id_from_auth or not user_id_from_auth:
raise HTTPException(
status_code=403,
detail="tenant_id and user_id must come from authentication token in production mode"
)
tenant_id = tenant_id_from_auth
user_id = user_id_from_auth
else:
tenant_id = tenant_id or tenant_id_from_auth
user_id = user_id or user_id_from_auth
if not tenant_id or not kb_id or not user_id:
raise HTTPException(
status_code=400,
detail="tenant_id, kb_id, and user_id are required"
)
try:
vector_store = get_vector_store()
deleted = vector_store.delete_by_filter({
"tenant_id": tenant_id, # CRITICAL: Multi-tenant isolation
"kb_id": kb_id,
"user_id": user_id
})
return {
"success": True,
"message": f"Cleared knowledge base. Deleted {deleted} chunks.",
"kb_id": kb_id
}
except Exception as e:
logger.error(f"Error clearing KB: {e}")
raise HTTPException(status_code=500, detail=str(e))
# ============== Chat Endpoints ==============
@app.post("/chat", response_model=ChatResponse)
@limiter.limit("10/minute", key_func=get_tenant_rate_limit_key)
async def chat(chat_request: ChatRequest, request: Request):
"""
Process a chat message using RAG.
- Retrieves relevant context from knowledge base
- Generates answer using LLM
- Returns answer with citations
"""
conversation_id = "unknown"
try:
logger.info(f"=== CHAT REQUEST RECEIVED ===")
logger.info(f"Request body: tenant_id={chat_request.tenant_id}, user_id={chat_request.user_id}, kb_id={chat_request.kb_id}, question_length={len(chat_request.question)}")
logger.info(f"Request headers: {dict(request.headers)}")
# SECURITY: Get tenant_id and user_id from auth context
# In PROD: MUST come from JWT token (never from request body)
try:
auth_context = await get_auth_context(request)
except Exception as e:
logger.error(f"Error getting auth context: {e}", exc_info=True)
raise HTTPException(status_code=401, detail=f"Authentication error: {str(e)}")
tenant_id_from_auth = auth_context.get("tenant_id")
user_id_from_auth = auth_context.get("user_id")
if settings.ENV == "prod":
if not tenant_id_from_auth or not user_id_from_auth:
raise HTTPException(
status_code=403,
detail="tenant_id and user_id must come from authentication token in production mode"
)
# Override request values with auth context (security enforcement)
chat_request.tenant_id = tenant_id_from_auth
chat_request.user_id = user_id_from_auth
else:
# DEV mode: use from request if provided, otherwise from auth context
if not chat_request.tenant_id:
chat_request.tenant_id = tenant_id_from_auth
if not chat_request.user_id:
chat_request.user_id = user_id_from_auth
if not chat_request.tenant_id or not chat_request.user_id:
raise HTTPException(
status_code=400,
detail="tenant_id and user_id are required (provide via X-Tenant-Id/X-User-Id headers or request body)"
)
# Log without PII in production
if settings.ENV == "prod":
logger.info(f"Chat request: tenant={chat_request.tenant_id}, user={chat_request.user_id}, kb={chat_request.kb_id}, q_length={len(chat_request.question)}")
else:
logger.info(f"Chat request: tenant={chat_request.tenant_id}, user={chat_request.user_id}, kb={chat_request.kb_id}, q={chat_request.question[:50]}...")
# Generate conversation ID if not provided
conversation_id = chat_request.conversation_id or f"conv_{uuid.uuid4().hex[:12]}"
# Get database session
try:
db = next(get_db())
except Exception as e:
logger.error(f"Database connection error: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=f"Database error: {str(e)}")
try:
# Ensure tenant exists in billing DB
ensure_tenant_exists(db, chat_request.tenant_id)
# Check quota BEFORE making LLM call
has_quota, quota_error = check_quota(db, chat_request.tenant_id)
if not has_quota:
logger.warning(f"Quota exceeded for tenant {chat_request.tenant_id}")
raise HTTPException(
status_code=402,
detail=quota_error or "AI quota exceeded. Upgrade your plan."
)
# Retrieve relevant context
retrieval_service = get_retrieval_service()
results, confidence, has_relevant = retrieval_service.retrieve(
query=chat_request.question,
tenant_id=chat_request.tenant_id, # CRITICAL: Multi-tenant isolation
kb_id=chat_request.kb_id,
user_id=chat_request.user_id
)
logger.info(f"Retrieval results: {len(results)} results, confidence={confidence:.3f}, has_relevant={has_relevant}")
# Format context for LLM
context, citations_info = retrieval_service.get_context_for_llm(results)
logger.info(f"Formatted context length: {len(context)} chars, citations: {len(citations_info)}")
# Generate answer
answer_service = get_answer_service()
answer_result = answer_service.generate_answer(
question=chat_request.question,
context=context,
citations_info=citations_info,
confidence=confidence,
has_relevant_results=has_relevant
)
# Track usage if LLM was called (usage info present)
usage_info = answer_result.get("usage")
if usage_info:
try:
track_usage(
db=db,
tenant_id=chat_request.tenant_id,
user_id=chat_request.user_id,
kb_id=chat_request.kb_id,
provider=settings.LLM_PROVIDER,
model=usage_info.get("model_used", settings.GEMINI_MODEL if settings.LLM_PROVIDER == "gemini" else settings.OPENAI_MODEL),
prompt_tokens=usage_info.get("prompt_tokens", 0),
completion_tokens=usage_info.get("completion_tokens", 0)
)
except Exception as e:
logger.error(f"Failed to track usage: {e}", exc_info=True)
# Don't fail the request if usage tracking fails
# Build metadata with refusal info
metadata = {
"chunks_retrieved": len(results),
"kb_id": chat_request.kb_id
}
if "refused" in answer_result:
metadata["refused"] = answer_result["refused"]
if "refusal_reason" in answer_result:
metadata["refusal_reason"] = answer_result["refusal_reason"]
if "verifier_passed" in answer_result:
metadata["verifier_passed"] = answer_result["verifier_passed"]
return ChatResponse(
success=True,
answer=answer_result["answer"],
citations=answer_result["citations"],
confidence=answer_result["confidence"],
from_knowledge_base=answer_result["from_knowledge_base"],
escalation_suggested=answer_result["escalation_suggested"],
conversation_id=conversation_id,
refused=answer_result.get("refused", False),
metadata=metadata
)
except ValueError as e:
# API key or configuration error
error_msg = str(e)
logger.error(f"Configuration error: {error_msg}")
if "API key" in error_msg.lower():
return ChatResponse(
success=False,
answer="⚠️ LLM API key not configured. Please set GEMINI_API_KEY in your .env file. Retrieval is working, but answer generation requires an API key.",
citations=[],
confidence=0.0,
from_knowledge_base=False,
escalation_suggested=True,
conversation_id=conversation_id,
metadata={"error": error_msg, "error_type": "configuration"}
)
else:
return ChatResponse(
success=False,
answer=f"Configuration error: {error_msg}",
citations=[],
confidence=0.0,
from_knowledge_base=False,
escalation_suggested=True,
conversation_id=conversation_id,
metadata={"error": error_msg}
)
except HTTPException:
# Re-raise HTTP exceptions (they have proper status codes)
raise
except Exception as e:
logger.error(f"Chat error: {e}", exc_info=True)
logger.error(f"Error type: {type(e).__name__}", exc_info=True)
return ChatResponse(
success=False,
answer=f"I encountered an error processing your request: {str(e)}. Please check the server logs for details.",
citations=[],
confidence=0.0,
from_knowledge_base=False,
escalation_suggested=True,
conversation_id=conversation_id,
metadata={"error": str(e), "error_type": type(e).__name__}
)
except HTTPException:
# Re-raise HTTP exceptions from outer try block
raise
except Exception as e:
logger.error(f"Outer chat error: {e}", exc_info=True)
return ChatResponse(
success=False,
answer=f"I encountered an error processing your request: {str(e)}. Please check the server logs for details.",
citations=[],
confidence=0.0,
from_knowledge_base=False,
escalation_suggested=True,
conversation_id=conversation_id,
metadata={"error": str(e), "error_type": type(e).__name__}
)
# ============== Utility Endpoints ==============
@app.get("/kb/search")
@limiter.limit("30/minute", key_func=get_tenant_rate_limit_key)
async def search_kb(
request: Request,
query: str,
tenant_id: Optional[str] = None, # Optional in dev, ignored in prod
kb_id: Optional[str] = None,
user_id: Optional[str] = None, # Optional in dev, ignored in prod
top_k: int = 5
):
"""
Search the knowledge base without generating an answer.
Useful for debugging and testing retrieval.
"""
# SECURITY: Extract tenant_id from auth token in production
if settings.ENV == "prod":
auth_context = await require_auth(request)
tenant_id = auth_context.get("tenant_id")
user_id = auth_context.get("user_id")
if not tenant_id or not user_id:
raise HTTPException(
status_code=403,
detail="tenant_id and user_id must come from authentication token in production mode"
)
elif not tenant_id or not kb_id or not user_id:
raise HTTPException(
status_code=400,
detail="tenant_id, kb_id, and user_id are required"
)
try:
retrieval_service = get_retrieval_service()
results, confidence, has_relevant = retrieval_service.retrieve(
query=query,
tenant_id=tenant_id, # CRITICAL: Multi-tenant isolation
kb_id=kb_id,
user_id=user_id,
top_k=top_k
)
return {
"success": True,
"results": [
{
"chunk_id": r.chunk_id,
"content": r.content[:500] + "..." if len(r.content) > 500 else r.content,
"metadata": r.metadata,
"similarity_score": r.similarity_score
}
for r in results
],
"confidence": confidence,
"has_relevant_results": has_relevant
}
except Exception as e:
logger.error(f"Search error: {e}")
raise HTTPException(status_code=500, detail=str(e))
# ============== Billing & Usage Endpoints ==============
@app.get("/billing/usage", response_model=UsageResponse)
async def get_usage(
request: Request,
range: str = "month", # "day" or "month"
year: Optional[int] = None,
month: Optional[int] = None,
day: Optional[int] = None
):
"""
Get usage statistics for the current tenant.
Args:
range: "day" or "month"
year: Year (optional, defaults to current)
month: Month 1-12 (optional, defaults to current)
day: Day 1-31 (optional, defaults to current, only for range="day")
"""
# Get tenant from auth
auth_context = await get_auth_context(request)
tenant_id = auth_context.get("tenant_id")
if not tenant_id:
raise HTTPException(status_code=403, detail="tenant_id required")
db = next(get_db())
try:
from app.db.models import UsageDaily, UsageMonthly
from datetime import datetime
from calendar import monthrange
now = datetime.utcnow()
target_year = year or now.year
target_month = month or now.month
if range == "day":
target_day = day or now.day
date_start = datetime(target_year, target_month, target_day)
daily = db.query(UsageDaily).filter(
UsageDaily.tenant_id == tenant_id,
UsageDaily.date == date_start
).first()
if not daily:
return UsageResponse(
tenant_id=tenant_id,
period="day",
total_requests=0,
total_tokens=0,
total_cost_usd=0.0,
start_date=date_start,
end_date=date_start
)
return UsageResponse(
tenant_id=tenant_id,
period="day",
total_requests=daily.total_requests,
total_tokens=daily.total_tokens,
total_cost_usd=daily.total_cost_usd,
gemini_requests=daily.gemini_requests,
openai_requests=daily.openai_requests,
start_date=daily.date,
end_date=daily.date
)
else: # month
monthly = db.query(UsageMonthly).filter(
UsageMonthly.tenant_id == tenant_id,
UsageMonthly.year == target_year,
UsageMonthly.month == target_month
).first()
if not monthly:
# Calculate date range for the month
_, last_day = monthrange(target_year, target_month)
start_date = datetime(target_year, target_month, 1)
end_date = datetime(target_year, target_month, last_day)
return UsageResponse(
tenant_id=tenant_id,
period="month",
total_requests=0,
total_tokens=0,
total_cost_usd=0.0,
start_date=start_date,
end_date=end_date
)
_, last_day = monthrange(monthly.year, monthly.month)
start_date = datetime(monthly.year, monthly.month, 1)
end_date = datetime(monthly.year, monthly.month, last_day)
return UsageResponse(
tenant_id=tenant_id,
period="month",
total_requests=monthly.total_requests,
total_tokens=monthly.total_tokens,
total_cost_usd=monthly.total_cost_usd,
gemini_requests=monthly.gemini_requests,
openai_requests=monthly.openai_requests,
start_date=start_date,
end_date=end_date
)
except Exception as e:
logger.error(f"Error getting usage: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
@app.get("/billing/limits", response_model=PlanLimitsResponse)
async def get_limits(request: Request):
"""Get current plan limits and usage for the tenant."""
# Get tenant from auth
auth_context = await get_auth_context(request)
tenant_id = auth_context.get("tenant_id")
if not tenant_id:
raise HTTPException(status_code=403, detail="tenant_id required")
db = next(get_db())
try:
from app.billing.quota import get_tenant_plan, get_monthly_usage
from datetime import datetime
plan = get_tenant_plan(db, tenant_id)
if not plan:
# Default to starter
plan_name = "starter"
monthly_limit = 500
else:
plan_name = plan.plan_name
monthly_limit = plan.monthly_chat_limit
# Get current month usage
now = datetime.utcnow()
monthly_usage = get_monthly_usage(db, tenant_id, now.year, now.month)
current_usage = monthly_usage.total_requests if monthly_usage else 0
remaining = None if monthly_limit == -1 else max(0, monthly_limit - current_usage)
return PlanLimitsResponse(
tenant_id=tenant_id,
plan_name=plan_name,
monthly_chat_limit=monthly_limit,
current_month_usage=current_usage,
remaining_chats=remaining
)
except Exception as e:
logger.error(f"Error getting limits: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
@app.post("/billing/plan")
async def set_plan(request_body: SetPlanRequest, http_request: Request):
"""
Set tenant's subscription plan (admin only in production).
In dev mode, allows any tenant to set their plan.
In prod mode, should be restricted to admin users.
"""
# Get tenant from auth
auth_context = await get_auth_context(http_request)
auth_tenant_id = auth_context.get("tenant_id")
# In prod, verify admin role (placeholder - implement actual admin check)
if settings.ENV == "prod":
# TODO: Add admin role check
if auth_tenant_id != request_body.tenant_id:
raise HTTPException(status_code=403, detail="Cannot set plan for other tenants")
db = next(get_db())
try:
from app.billing.quota import set_tenant_plan
plan = set_tenant_plan(db, request_body.tenant_id, request_body.plan_name)
return {
"success": True,
"tenant_id": request_body.tenant_id,
"plan_name": plan.plan_name,
"monthly_chat_limit": plan.monthly_chat_limit
}
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
except Exception as e:
logger.error(f"Error setting plan: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
@app.get("/billing/cost-report", response_model=CostReportResponse)
async def get_cost_report(
request: Request,
range: str = "month",
year: Optional[int] = None,
month: Optional[int] = None
):
"""Get cost report with breakdown by provider and model."""
# Get tenant from auth
auth_context = await get_auth_context(request)
tenant_id = auth_context.get("tenant_id")
if not tenant_id:
raise HTTPException(status_code=403, detail="tenant_id required")
db = next(get_db())
try:
from app.db.models import UsageEvent
from datetime import datetime
from sqlalchemy import func, and_
now = datetime.utcnow()
target_year = year or now.year
target_month = month or now.month
# Query usage events for the period
if range == "month":
query = db.query(UsageEvent).filter(
and_(
UsageEvent.tenant_id == tenant_id,
func.extract('year', UsageEvent.request_timestamp) == target_year,
func.extract('month', UsageEvent.request_timestamp) == target_month
)
)
else: # all time
query = db.query(UsageEvent).filter(UsageEvent.tenant_id == tenant_id)
events = query.all()
# Calculate totals
total_cost = sum(e.estimated_cost_usd for e in events)
total_requests = len(events)
total_tokens = sum(e.total_tokens for e in events)
# Breakdown by provider
breakdown_by_provider = {}
for event in events:
provider = event.provider
if provider not in breakdown_by_provider:
breakdown_by_provider[provider] = {
"requests": 0,
"tokens": 0,
"cost_usd": 0.0
}
breakdown_by_provider[provider]["requests"] += 1
breakdown_by_provider[provider]["tokens"] += event.total_tokens
breakdown_by_provider[provider]["cost_usd"] += event.estimated_cost_usd
# Breakdown by model
breakdown_by_model = {}
for event in events:
model = event.model
if model not in breakdown_by_model:
breakdown_by_model[model] = {
"requests": 0,
"tokens": 0,
"cost_usd": 0.0
}
breakdown_by_model[model]["requests"] += 1
breakdown_by_model[model]["tokens"] += event.total_tokens
breakdown_by_model[model]["cost_usd"] += event.estimated_cost_usd
return CostReportResponse(
tenant_id=tenant_id,
period=range,
total_cost_usd=total_cost,
total_requests=total_requests,
total_tokens=total_tokens,
breakdown_by_provider=breakdown_by_provider,
breakdown_by_model=breakdown_by_model
)
except Exception as e:
logger.error(f"Error getting cost report: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
