privaite 0.2.3__py3-none-any.whl

privaite/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.3"

privaite/__main__.py

@@ -0,0 +1,3 @@
+from privaite.cli import main
+
+main()

privaite/api/chat.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from fastapi import APIRouter, Depends, Request
+from fastapi.responses import StreamingResponse
+
+from privaite.api.dependencies import get_config, get_pii_engine, get_provider_router
+from privaite.config.schema import PrivAiTeConfig
+from privaite.pii.engine import UnsupportedContentError
+from privaite.providers.router import ProviderRouter
+from privaite.utils.errors import openai_error, provider_error_response
+
+logger = logging.getLogger("privaite.api.chat")
+
+router = APIRouter(prefix="/v1")
+
+
+@router.post("/chat/completions", response_model=None)
+async def chat_completions(
+    request: Request,
+    config: PrivAiTeConfig = Depends(get_config),
+    pii_engine: Any = Depends(get_pii_engine),
+    provider_router: ProviderRouter = Depends(get_provider_router),
+):
+    body = await request.json()
+    model = body.get("model")
+    messages = body.get("messages", [])
+    stream = body.get("stream", False)
+
+    if not model:
+        return openai_error("model is required", "invalid_request_error", 400)
+
+    if not provider_router.has_model(model):
+        return openai_error(f"Model '{model}' not found", "not_found_error", 404)
+
+    mapping = None
+
+    if config.pii.enabled and pii_engine is not None:
+        try:
+            messages, mapping = await pii_engine.process_request(messages)
+            tracker = getattr(request.app.state, "pii_tracker", None)
+            if tracker and mapping and not mapping.is_empty:
+                session_id = request.headers.get(
+                    "x-session-id",
+                    request.headers.get("authorization", "anonymous"),
+                )
+                counts: dict[str, int] = {}
+                for orig in mapping._original_to_fake:
+                    t = mapping.get_entity_type(orig)
+                    if t:
+                        counts[t] = counts.get(t, 0) + 1
+                tracker.record(session_id, counts)
+        except UnsupportedContentError as exc:
+            return openai_error(str(exc), "invalid_request_error", 400)
+        except Exception:
+            logger.exception("PII processing failed")
+            if config.pii.on_error == "block":
+                return openai_error(
+                    "PII anonymization failed. Request blocked for privacy.",
+                    "server_error",
+                    500,
+                    "pii_error",
+                )
+
+    kwargs = {k: v for k, v in body.items() if k not in ("model", "messages", "stream")}
+
+    try:
+        if stream:
+            litellm_stream = await provider_router.streaming_completion(
+                model_alias=model, messages=messages, **kwargs
+            )
+
+            from privaite.streaming.handler import StreamingHandler
+
+            deanon_config = config.pii.deanonymization if config.pii.enabled else None
+            generator = StreamingHandler.stream_response(
+                litellm_stream=litellm_stream,
+                mapping=mapping,
+                deanonymizer_config=deanon_config,
+            )
+
+            return StreamingResponse(
+                generator,
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "X-Accel-Buffering": "no",
+                },
+            )
+
+        response = await provider_router.completion(
+            model_alias=model, messages=messages, **kwargs
+        )
+    except Exception as exc:
+        logger.exception("Provider error for model %s", model)
+        return provider_error_response(exc)
+
+    if (
+        mapping
+        and config.pii.deanonymization.enabled
+        and pii_engine is not None
+    ):
+        response_dict = response.model_dump() if hasattr(response, "model_dump") else dict(response)
+        for choice in response_dict.get("choices", []):
+            msg = choice.get("message", {})
+            content = msg.get("content")
+            if content:
+                msg["content"] = await pii_engine.process_response(content, mapping)
+            tool_calls = msg.get("tool_calls")
+            if tool_calls:
+                msg["tool_calls"] = await pii_engine.process_response_tool_calls(
+                    tool_calls, mapping
+                )
+            function_call = msg.get("function_call")
+            if function_call:
+                msg["function_call"] = await pii_engine.process_response_function_call(
+                    function_call, mapping
+                )
+        return response_dict
+
+    if hasattr(response, "model_dump"):
+        return response.model_dump()
+    return dict(response)

privaite/api/completions.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from fastapi import APIRouter, Depends, Request
+from fastapi.responses import StreamingResponse
+
+from privaite.api.dependencies import get_config, get_pii_engine, get_provider_router
+from privaite.config.schema import PrivAiTeConfig
+from privaite.pii.engine import UnsupportedContentError
+from privaite.providers.router import ProviderRouter
+from privaite.utils.errors import openai_error, provider_error_response
+
+logger = logging.getLogger("privaite.api.completions")
+
+router = APIRouter(prefix="/v1")
+
+
+@router.post("/completions", response_model=None)
+async def completions(
+    request: Request,
+    config: PrivAiTeConfig = Depends(get_config),
+    pii_engine: Any = Depends(get_pii_engine),
+    provider_router: ProviderRouter = Depends(get_provider_router),
+):
+    body = await request.json()
+    model = body.get("model")
+    prompt = body.get("prompt", "")
+    stream = body.get("stream", False)
+
+    if not model:
+        return openai_error("model is required", "invalid_request_error", 400)
+
+    if not provider_router.has_model(model):
+        return openai_error(f"Model '{model}' not found", "not_found_error", 404)
+
+    mapping = None
+
+    if config.pii.enabled and pii_engine is not None:
+        try:
+            if isinstance(prompt, list) and all(isinstance(p, str) for p in prompt):
+                msgs = [{"role": "user", "content": p} for p in prompt]
+                msgs, mapping = await pii_engine.process_request(msgs)
+                prompt = [m["content"] for m in msgs]
+            else:
+                msgs = [{"role": "user", "content": prompt}]
+                msgs, mapping = await pii_engine.process_request(msgs)
+                prompt = msgs[0]["content"]
+        except UnsupportedContentError as exc:
+            return openai_error(str(exc), "invalid_request_error", 400)
+        except Exception:
+            logger.exception("PII processing failed")
+            if config.pii.on_error == "block":
+                return openai_error(
+                    "PII anonymization failed. Request blocked for privacy.",
+                    "server_error", 500, "pii_error",
+                )
+
+    kwargs = {k: v for k, v in body.items() if k not in ("model", "prompt", "stream")}
+
+    try:
+        if stream:
+            litellm_stream = await provider_router.streaming_completion(
+                model_alias=model,
+                messages=[{"role": "user", "content": prompt}],
+                **kwargs,
+            )
+
+            from privaite.streaming.handler import StreamingHandler
+
+            deanon_config = config.pii.deanonymization if config.pii.enabled else None
+            generator = StreamingHandler.stream_response(
+                litellm_stream=litellm_stream,
+                mapping=mapping,
+                deanonymizer_config=deanon_config,
+            )
+
+            return StreamingResponse(
+                generator,
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "X-Accel-Buffering": "no",
+                },
+            )
+
+        response = await provider_router.completion(
+            model_alias=model,
+            messages=[{"role": "user", "content": prompt}],
+            **kwargs,
+        )
+    except Exception as exc:
+        logger.exception("Provider error for model %s", model)
+        return provider_error_response(exc)
+
+    if mapping and config.pii.deanonymization.enabled and pii_engine is not None:
+        response_dict = response.model_dump() if hasattr(response, "model_dump") else dict(response)
+        for choice in response_dict.get("choices", []):
+            text = choice.get("text")
+            if text:
+                choice["text"] = await pii_engine.process_response(text, mapping)
+        return response_dict
+
+    if hasattr(response, "model_dump"):
+        return response.model_dump()
+    return dict(response)

privaite/api/dependencies.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from fastapi import Request
+
+if TYPE_CHECKING:
+    from privaite.config.schema import PrivAiTeConfig
+    from privaite.pii.engine import PIIEngine
+    from privaite.providers.router import ProviderRouter
+
+
+def get_config(request: Request) -> PrivAiTeConfig:
+    return request.app.state.config
+
+
+def get_pii_engine(request: Request) -> PIIEngine | None:
+    return getattr(request.app.state, "pii_engine", None)
+
+
+def get_provider_router(request: Request) -> ProviderRouter:
+    return request.app.state.provider_router

privaite/api/embeddings.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from fastapi import APIRouter, Depends, Request
+
+from privaite.api.dependencies import get_config, get_pii_engine, get_provider_router
+from privaite.config.schema import PrivAiTeConfig
+from privaite.pii.engine import UnsupportedContentError
+from privaite.providers.router import ProviderRouter
+from privaite.utils.errors import openai_error, provider_error_response
+
+logger = logging.getLogger("privaite.api.embeddings")
+
+router = APIRouter(prefix="/v1")
+
+
+@router.post("/embeddings", response_model=None)
+async def embeddings(
+    request: Request,
+    config: PrivAiTeConfig = Depends(get_config),
+    pii_engine: Any = Depends(get_pii_engine),
+    provider_router: ProviderRouter = Depends(get_provider_router),
+):
+    body = await request.json()
+    model = body.get("model")
+    input_text = body.get("input", "")
+
+    if not model:
+        return openai_error("model is required", "invalid_request_error", 400)
+
+    if not provider_router.has_model(model):
+        return openai_error(f"Model '{model}' not found", "not_found_error", 404)
+
+    if config.pii.enabled and pii_engine is not None:
+        try:
+            if isinstance(input_text, str):
+                msgs = [{"role": "user", "content": input_text}]
+                msgs, _ = await pii_engine.process_request(msgs)
+                input_text = msgs[0]["content"]
+            elif isinstance(input_text, list):
+                anonymized = []
+                for text in input_text:
+                    msgs = [{"role": "user", "content": text}]
+                    msgs, _ = await pii_engine.process_request(msgs)
+                    anonymized.append(msgs[0]["content"])
+                input_text = anonymized
+        except UnsupportedContentError as exc:
+            return openai_error(str(exc), "invalid_request_error", 400)
+        except Exception:
+            logger.exception("PII processing failed")
+            if config.pii.on_error == "block":
+                return openai_error(
+                    "PII anonymization failed. Request blocked for privacy.",
+                    "server_error", 500, "pii_error",
+                )
+
+    kwargs = {k: v for k, v in body.items() if k not in ("model", "input")}
+
+    try:
+        response = await provider_router.embedding(
+            model_alias=model, input_text=input_text, **kwargs
+        )
+    except Exception as exc:
+        logger.exception("Provider error for model %s", model)
+        return provider_error_response(exc)
+
+    if hasattr(response, "model_dump"):
+        return response.model_dump()
+    return dict(response)

privaite/api/health.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Request
+
+router = APIRouter()
+
+
+@router.get("/health")
+async def health() -> dict:
+    return {"status": "ok"}
+
+
+@router.get("/ready")
+async def ready(request: Request) -> dict:
+    checks: dict[str, bool] = {}
+
+    provider_router = getattr(request.app.state, "provider_router", None)
+    checks["providers_configured"] = (
+        provider_router is not None and len(provider_router.models) > 0
+    )
+
+    pii_engine = getattr(request.app.state, "pii_engine", None)
+    if pii_engine is not None:
+        checks["pii_engine_ready"] = pii_engine.is_ready
+    else:
+        checks["pii_engine_ready"] = True
+
+    all_ready = all(checks.values())
+    return {"ready": all_ready, "checks": checks}
+
+
+@router.get("/stats")
+async def stats(request: Request) -> dict:
+    tracker = getattr(request.app.state, "pii_tracker", None)
+    if tracker is None:
+        return {"enabled": False}
+
+    sessions = {}
+    with tracker._lock:
+        for sid, s in tracker._sessions.items():
+            label = sid[:16] + "..." if len(sid) > 16 else sid
+            sessions[label] = {
+                "requests": s.request_count,
+                "total_pii": s.total_pii,
+                "by_type": dict(s.pii_count),
+            }
+
+    return {"enabled": True, "sessions": sessions}

privaite/api/models.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import time
+
+from fastapi import APIRouter, Depends
+
+from privaite.api.dependencies import get_provider_router
+from privaite.providers.router import ProviderRouter
+
+router = APIRouter(prefix="/v1")
+
+
+@router.get("/models")
+async def list_models(
+    provider_router: ProviderRouter = Depends(get_provider_router),
+) -> dict:
+    models = []
+    for model_name in provider_router.models:
+        models.append({
+            "id": model_name,
+            "object": "model",
+            "created": int(time.time()),
+            "owned_by": "privaite",
+        })
+    return {"object": "list", "data": models}

privaite/api/router.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from fastapi import APIRouter
+
+from privaite.api import chat, completions, embeddings, health, models
+
+api_router = APIRouter()
+api_router.include_router(health.router, tags=["health"])
+api_router.include_router(models.router, tags=["models"])
+api_router.include_router(chat.router, tags=["chat"])
+api_router.include_router(completions.router, tags=["completions"])
+api_router.include_router(embeddings.router, tags=["embeddings"])

privaite/app.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
+from privaite.api.router import api_router
+from privaite.config.schema import PrivAiTeConfig
+from privaite.middleware.auth import AuthMiddleware
+from privaite.middleware.limits import RequestSizeLimitMiddleware
+from privaite.providers.router import ProviderRouter
+from privaite.utils.logging import setup_logging
+from privaite.utils.security import get_api_keys
+
+logger = logging.getLogger("privaite.app")
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+    config: PrivAiTeConfig = app.state.config
+
+    app.state.provider_router = ProviderRouter(config.providers)
+    logger.info(
+        "Provider router ready with %d model(s)", len(config.providers)
+    )
+
+    if config.auth.enabled and not get_api_keys():
+        logger.warning(
+            "Auth is enabled but PRIVAITE_API_KEYS is empty; all requests will be "
+            "rejected with 401 until a key is set (or set auth.enabled=false)."
+        )
+
+    if config.pii.enabled:
+        from privaite.pii.engine import PIIEngine
+        from privaite.pii.tracker import PIITracker
+
+        engine = PIIEngine(config.pii)
+        await engine.initialize()
+        app.state.pii_engine = engine
+        app.state.pii_tracker = PIITracker()
+        logger.info("PII engine initialized")
+    else:
+        app.state.pii_engine = None
+        app.state.pii_tracker = None
+        logger.info("PII processing disabled")
+
+    yield
+
+    if app.state.pii_engine is not None:
+        await app.state.pii_engine.shutdown()
+    logger.info("PrivAiTe shutdown complete")
+
+
+def create_app(config: PrivAiTeConfig | None = None) -> FastAPI:
+    if config is None:
+        from privaite.config.loader import load_config
+        config = load_config()
+
+    setup_logging(level=config.logging.level, fmt=config.logging.format)
+
+    app = FastAPI(
+        title="PrivAiTe",
+        description="Privacy-first LLM proxy with transparent PII anonymization",
+        version="0.2.3",
+        lifespan=lifespan,
+    )
+
+    app.state.config = config
+
+    app.add_middleware(AuthMiddleware)
+    app.add_middleware(
+        RequestSizeLimitMiddleware, max_bytes=config.server.max_request_bytes
+    )
+    app.include_router(api_router)
+
+    return app

privaite/cli.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import os
+
+import click
+import uvicorn
+
+from privaite.config.loader import load_config
+
+
+@click.command()
+@click.option("--config", "config_path", default=None, help="Path to config YAML file")
+@click.option("--host", default=None, help="Override server host")
+@click.option("--port", default=None, type=int, help="Override server port")
+@click.option("--reload", is_flag=True, help="Auto-reload on file changes (dev mode)")
+def main(config_path: str | None, host: str | None, port: int | None, reload: bool) -> None:
+    # uvicorn re-imports create_app() in the worker process, so the path has to
+    # travel through the environment to reach the factory's load_config().
+    if config_path:
+        os.environ["PRIVAITE_CONFIG_PATH"] = config_path
+    config = load_config(config_path)
+
+    run_host = host or config.server.host
+    run_port = port or config.server.port
+
+    click.echo(f"Starting PrivAiTe on {run_host}:{run_port}")
+    click.echo(f"PII processing: {'enabled' if config.pii.enabled else 'disabled'}")
+    click.echo(f"Providers: {len(config.providers)} configured")
+    if reload:
+        click.echo("Auto-reload enabled (dev mode)")
+
+    uvicorn.run(
+        "privaite.app:create_app",
+        host=run_host,
+        port=run_port,
+        workers=config.server.workers,
+        log_level=config.server.log_level,
+        factory=True,
+        reload=reload,
+        reload_dirs=["privaite", "config"] if reload else None,
+    )
+
+
+if __name__ == "__main__":
+    main()

privaite/config/__init__.py

@@ -0,0 +1,4 @@
+from privaite.config.loader import load_config
+from privaite.config.schema import PrivAiTeConfig
+
+__all__ = ["load_config", "PrivAiTeConfig"]

privaite/config/loader.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import os
+import re
+from pathlib import Path
+
+import yaml
+from dotenv import load_dotenv
+
+from privaite.config.schema import PrivAiTeConfig
+
+_ENV_VAR_PATTERN = re.compile(r"\$\{([^}]+)}")
+
+
+def _interpolate_env_vars(obj: object) -> object:
+    if isinstance(obj, str):
+        def _replace(match: re.Match) -> str:
+            var_name = match.group(1)
+            value = os.environ.get(var_name)
+            if value is None:
+                raise ValueError(f"Environment variable '{var_name}' is not set")
+            return value
+
+        return _ENV_VAR_PATTERN.sub(_replace, obj)
+
+    if isinstance(obj, dict):
+        return {k: _interpolate_env_vars(v) for k, v in obj.items()}
+
+    if isinstance(obj, list):
+        return [_interpolate_env_vars(item) for item in obj]
+
+    return obj
+
+
+def load_config(path: str | Path | None = None) -> PrivAiTeConfig:
+    # override=False so the explicit environment (e.g. PRIVAITE_CONFIG_PATH set by
+    # the CLI from --config) wins over .env; .env only fills in variables that are
+    # not already set, such as OPENAI_API_KEY.
+    load_dotenv(override=False)
+
+    if path is None:
+        path = os.environ.get("PRIVAITE_CONFIG_PATH", "config/privaite.yaml")
+
+    path = Path(path)
+
+    if not path.exists():
+        return PrivAiTeConfig()
+
+    with open(path) as f:
+        raw = yaml.safe_load(f) or {}
+
+    interpolated = _interpolate_env_vars(raw)
+    return PrivAiTeConfig.model_validate(interpolated)