app/main.py
Rama main · réplica de lectura del repositorio original.
from __future__ import annotations
import asyncio
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any
from fastapi import FastAPI, Request
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.responses import FileResponse, JSONResponse, Response
from fastapi.staticfiles import StaticFiles
from app.chat import (
ChatOrchestrator,
build_chat_router,
build_playground_chat_router,
error_payload,
)
from app.config import Settings, provider_key_is_configured
from app.debug.routes import build_debug_router
from app.identity import DefaultPrincipalResolver, PrincipalContext
from app.logging_config import configure_logging
from app.media.routes import build_media_router
from app.media.sanitizer import MediaSanitizer
from app.observability import PrivacyMetrics
from app.policies.policy_engine import PolicyEngine
from app.privacy.detection import DetectorEnsemble, LegacyEntityDetectorAdapter
from app.privacy.models import DetectorProfile
from app.privacy.output_guard import OutputPrivacyGuard
from app.privacy.pipeline import PrivacyPipeline
from app.privacy.policy import LegacyPolicyAdapter, PolicyEngineV2, load_policy_v2
from app.privacy.risk import PrivacyRiskEngine
from app.privacy.runtime import PrivacyRuntime
from app.privacy.wire import FinalWirePrivacyGuard
from app.providers import (
AdapterPolicy,
GeminiGenerateContentAdapter,
OpenAIChatCompletionsAdapter,
)
from app.proxy.gemini_client import GeminiClient
from app.proxy.llm_client import LLMClient
from app.security import (
FixedWindowRateLimiter,
RequestBodyLimitMiddleware,
bearer_token_is_allowed,
requires_proxy_auth,
)
from app.storage.conversation_store import InMemoryConversationStore
from app.storage.mapping_store import InMemoryMappingStore
PLAYGROUND_DIR = Path(__file__).resolve().parents[1] / "playground-react" / "dist"
def create_app(
settings: Settings | None = None,
*,
llm_client: Any | None = None,
ocr_adapter: Any | None = None,
face_detector: Any | None = None,
playground_dir: Path | None = None,
) -> FastAPI:
settings = settings or Settings.from_env()
playground_dir = playground_dir or PLAYGROUND_DIR
configure_logging(settings.log_level)
privacy_detector = DetectorEnsemble(profile=DetectorProfile(settings.detector_profile))
detector = LegacyEntityDetectorAdapter(privacy_detector)
policy = PolicyEngine(
settings.policy_file,
block_api_keys=settings.block_api_keys,
block_secrets=settings.block_secrets,
block_credit_cards=settings.block_credit_cards,
public_email_allowlist=settings.public_email_allowlist,
public_email_domain_allowlist=settings.public_email_domain_allowlist,
public_person_allowlist=settings.public_person_allowlist,
)
policy_document = (
load_policy_v2(settings.policy_v2_file)
if settings.policy_v2_file is not None
else LegacyPolicyAdapter().adapt(policy.rules)
)
policy_v2 = PolicyEngineV2(policy_document)
risk_engine = PrivacyRiskEngine()
privacy_metrics = PrivacyMetrics()
privacy_runtime = PrivacyRuntime(
privacy_detector,
risk_engine,
policy_v2,
privacy_metrics,
)
mapping_store = InMemoryMappingStore(settings.mapping_ttl_seconds)
conversation_store = InMemoryConversationStore(
settings.conversation_ttl_seconds,
settings.conversation_max_messages,
settings.conversation_max_chars,
settings.conversation_max_count,
)
media_sanitizer = MediaSanitizer(
detector,
policy,
settings.max_media_file_bytes,
ocr_adapter=ocr_adapter,
face_detector=face_detector,
)
if llm_client is not None:
upstream = llm_client
elif settings.llm_provider.lower() == "gemini":
upstream = GeminiClient(
settings.gemini_base_url,
settings.gemini_api_key or settings.llm_api_key,
settings.upstream_timeout_seconds,
max_response_bytes=settings.max_upstream_response_bytes,
)
else:
upstream = LLMClient(
settings.llm_base_url,
settings.llm_api_key,
settings.upstream_timeout_seconds,
max_response_bytes=settings.max_upstream_response_bytes,
)
provider_name = settings.llm_provider.lower()
provider_secret = settings.gemini_api_key if provider_name == "gemini" else settings.llm_api_key
if provider_name == "gemini" and not provider_secret:
provider_secret = settings.llm_api_key
provider_ready = llm_client is not None or provider_key_is_configured(provider_secret)
wire_guard = FinalWirePrivacyGuard(detector)
output_guard = OutputPrivacyGuard(detector)
openai_adapter_policy = AdapterPolicy(
safe_extension_fields=frozenset({"metadata", "user"})
)
ingress_adapter = OpenAIChatCompletionsAdapter(openai_adapter_policy)
egress_adapter = (
GeminiGenerateContentAdapter()
if provider_name == "gemini"
else OpenAIChatCompletionsAdapter(openai_adapter_policy)
)
privacy_pipeline = PrivacyPipeline(
upstream,
provider_name,
wire_guard,
output_guard,
ingress_adapter=ingress_adapter,
egress_adapter=egress_adapter,
)
rate_limiter = FixedWindowRateLimiter(
settings.rate_limit_requests,
settings.rate_limit_window_seconds,
)
principal_resolver = DefaultPrincipalResolver(settings.application_id)
chat_orchestrator = ChatOrchestrator(
settings=settings,
detector=detector,
policy=policy,
privacy_runtime=privacy_runtime,
mapping_store=mapping_store,
conversation_store=conversation_store,
upstream=upstream,
provider_name=provider_name,
provider_ready=provider_ready,
privacy_pipeline=privacy_pipeline,
output_guard=output_guard,
principal_resolver=principal_resolver,
)
@asynccontextmanager
async def lifespan(_: FastAPI):
async def sweep_expired_state() -> None:
while True:
await asyncio.sleep(
min(settings.mapping_ttl_seconds, settings.conversation_ttl_seconds, 60)
)
mapping_store.cleanup()
conversation_store.cleanup()
sweeper = asyncio.create_task(sweep_expired_state())
try:
yield
finally:
sweeper.cancel()
try:
await sweeper
except asyncio.CancelledError:
pass
app = FastAPI(
title="MaskGate",
version="0.1.0",
description="Self-hosted OpenAI-compatible privacy proxy for LLM applications.",
lifespan=lifespan,
docs_url=None if settings.is_production else "/docs",
redoc_url=None if settings.is_production else "/redoc",
openapi_url=None if settings.is_production else "/openapi.json",
)
app.add_middleware(
RequestBodyLimitMiddleware,
max_bytes=settings.max_request_body_bytes,
media_file_max_bytes=settings.max_media_file_bytes,
)
app.add_middleware(
TrustedHostMiddleware,
allowed_hosts=list(settings.trusted_hosts),
)
app.state.settings = settings
app.state.detector = detector
app.state.privacy_detector = privacy_detector
app.state.policy = policy
app.state.policy_v2 = policy_v2
app.state.risk_engine = risk_engine
app.state.privacy_metrics = privacy_metrics
app.state.privacy_runtime = privacy_runtime
app.state.mapping_store = mapping_store
app.state.conversation_store = conversation_store
app.state.llm_client = upstream
app.state.provider_ready = provider_ready
app.state.media_sanitizer = media_sanitizer
app.state.principal_resolver = principal_resolver
app.state.chat_orchestrator = chat_orchestrator
@app.middleware("http")
async def proxy_authentication(request: Request, call_next: Any) -> Response:
if requires_proxy_auth(request.url.path):
authorization = request.headers.get("authorization")
if settings.api_keys and not bearer_token_is_allowed(authorization, settings.api_keys):
return JSONResponse(
status_code=401,
content=error_payload("Invalid or missing proxy API key", "unauthorized"),
headers={
"WWW-Authenticate": "Bearer",
"Cache-Control": "no-store",
},
)
principal = principal_resolver.resolve(
authorization=authorization if settings.api_keys else None,
client_host=request.client.host if request.client else None,
)
request.state.principal = principal
request.state.security_identity = principal.vault_namespace
retry_after = rate_limiter.retry_after(principal.rate_limit_key)
if retry_after is not None:
return JSONResponse(
status_code=429,
content=error_payload("Rate limit exceeded", "rate_limit_exceeded"),
headers={
"Retry-After": str(retry_after),
"Cache-Control": "no-store",
},
)
return await call_next(request)
# Registered after authentication so it is the outermost application
# middleware and also decorates early 401/413/429/host-rejection responses.
@app.middleware("http")
async def privacy_response_headers(request: Request, call_next: Any) -> Response:
response = await call_next(request)
response.headers.setdefault("X-Content-Type-Options", "nosniff")
response.headers.setdefault("X-Frame-Options", "DENY")
response.headers.setdefault("Referrer-Policy", "no-referrer")
response.headers.setdefault(
"Permissions-Policy",
"camera=(), microphone=(), geolocation=()",
)
response.headers.setdefault("Cross-Origin-Opener-Policy", "same-origin")
response.headers.setdefault(
"Content-Security-Policy",
"default-src 'self'; "
"script-src 'self'; "
"style-src 'self'; "
"img-src 'self' data:; "
"font-src 'self'; "
"connect-src 'self'; "
"object-src 'none'; "
"base-uri 'none'; "
"frame-ancestors 'none'; "
"form-action 'self'",
)
if request.url.path.startswith(("/v1", "/playground/api", "/debug")):
response.headers["Cache-Control"] = "no-store"
return response
def request_principal(request: Request) -> PrincipalContext:
principal = getattr(request.state, "principal", None)
if isinstance(principal, PrincipalContext):
return principal
return principal_resolver.resolve(
authorization=None,
client_host=request.client.host if request.client else None,
)
@app.get("/health")
@app.get("/health/live")
async def health() -> dict[str, str]:
return {"status": "ok"}
@app.get("/health/ready")
async def readiness() -> JSONResponse:
if not provider_ready:
return JSONResponse(
status_code=503,
content={"status": "not_ready", "provider_ready": False},
headers={"Cache-Control": "no-store"},
)
return JSONResponse(content={"status": "ready", "provider_ready": True})
app.include_router(
build_chat_router(chat_orchestrator, request_principal)
)
app.include_router(
build_media_router(
media_sanitizer,
max_concurrency=settings.media_max_concurrency,
)
)
if settings.enable_debug_endpoints:
app.include_router(build_debug_router(detector, policy, settings.masking_mode))
if settings.enable_playground and playground_dir.is_dir():
app.mount(
"/playground/assets",
StaticFiles(directory=playground_dir / "assets"),
name="playground-assets",
)
@app.get("/playground", include_in_schema=False)
@app.get("/playground/", include_in_schema=False)
async def playground() -> FileResponse:
return FileResponse(playground_dir / "index.html")
@app.get("/playground/config", include_in_schema=False)
async def playground_config() -> dict[str, object]:
return {
"provider": settings.llm_provider,
"model": settings.gemini_model
if provider_name == "gemini"
else settings.llm_default_model,
"masking_mode": settings.masking_mode,
"conversation_store": "ram",
"conversation_ttl_seconds": settings.conversation_ttl_seconds,
"provider_ready": provider_ready,
}
app.include_router(
build_playground_chat_router(chat_orchestrator, request_principal)
)
return app
app = create_app()