genxai-framework 0.1.0__py3-none-any.whl

genxai/performance/pooling.py

@@ -0,0 +1,289 @@
+"""Connection pooling for GenXAI performance optimization."""
+
+import asyncio
+from typing import Optional, Any, Dict
+from contextlib import asynccontextmanager, contextmanager
+import queue
+import threading
+
+
+class ConnectionPool:
+    """Generic connection pool."""
+    
+    def __init__(
+        self,
+        create_connection,
+        min_size: int = 5,
+        max_size: int = 20,
+        timeout: float = 30.0
+    ):
+        """Initialize connection pool.
+        
+        Args:
+            create_connection: Function to create new connection
+            min_size: Minimum pool size
+            max_size: Maximum pool size
+            timeout: Connection timeout
+        """
+        self.create_connection = create_connection
+        self.min_size = min_size
+        self.max_size = max_size
+        self.timeout = timeout
+        
+        self._pool = asyncio.Queue(maxsize=max_size)
+        self._size = 0
+        self._lock = asyncio.Lock()
+    
+    async def initialize(self):
+        """Initialize pool with minimum connections."""
+        for _ in range(self.min_size):
+            conn = await self.create_connection()
+            await self._pool.put(conn)
+            self._size += 1
+    
+    async def acquire(self):
+        """Acquire connection from pool."""
+        try:
+            # Try to get existing connection
+            conn = await asyncio.wait_for(
+                self._pool.get(),
+                timeout=self.timeout
+            )
+            return conn
+        except asyncio.TimeoutError:
+            # Create new connection if under max size
+            async with self._lock:
+                if self._size < self.max_size:
+                    conn = await self.create_connection()
+                    self._size += 1
+                    return conn
+            
+            # Wait for available connection
+            conn = await self._pool.get()
+            return conn
+    
+    async def release(self, conn):
+        """Release connection back to pool."""
+        await self._pool.put(conn)
+    
+    async def close(self):
+        """Close all connections."""
+        while not self._pool.empty():
+            conn = await self._pool.get()
+            if hasattr(conn, 'close'):
+                await conn.close()
+        self._size = 0
+    
+    @asynccontextmanager
+    async def connection(self):
+        """Context manager for connection."""
+        conn = await self.acquire()
+        try:
+            yield conn
+        finally:
+            await self.release(conn)
+
+
+class DatabaseConnectionPool:
+    """Database connection pool."""
+    
+    def __init__(
+        self,
+        db_url: str,
+        min_size: int = 5,
+        max_size: int = 20
+    ):
+        """Initialize database connection pool.
+        
+        Args:
+            db_url: Database URL
+            min_size: Minimum pool size
+            max_size: Maximum pool size
+        """
+        self.db_url = db_url
+        self.min_size = min_size
+        self.max_size = max_size
+        self._pool = None
+    
+    async def initialize(self):
+        """Initialize pool."""
+        try:
+            import asyncpg
+            self._pool = await asyncpg.create_pool(
+                self.db_url,
+                min_size=self.min_size,
+                max_size=self.max_size
+            )
+        except ImportError:
+            raise ImportError("asyncpg not installed. Install with: pip install asyncpg")
+    
+    @asynccontextmanager
+    async def connection(self):
+        """Get connection from pool."""
+        async with self._pool.acquire() as conn:
+            yield conn
+    
+    async def close(self):
+        """Close pool."""
+        if self._pool:
+            await self._pool.close()
+
+
+class HTTPConnectionPool:
+    """HTTP connection pool using aiohttp."""
+    
+    def __init__(
+        self,
+        connector_limit: int = 100,
+        connector_limit_per_host: int = 30
+    ):
+        """Initialize HTTP connection pool.
+        
+        Args:
+            connector_limit: Total connection limit
+            connector_limit_per_host: Per-host connection limit
+        """
+        self.connector_limit = connector_limit
+        self.connector_limit_per_host = connector_limit_per_host
+        self._session = None
+    
+    async def initialize(self):
+        """Initialize session."""
+        try:
+            import aiohttp
+            connector = aiohttp.TCPConnector(
+                limit=self.connector_limit,
+                limit_per_host=self.connector_limit_per_host
+            )
+            self._session = aiohttp.ClientSession(connector=connector)
+        except ImportError:
+            raise ImportError("aiohttp not installed. Install with: pip install aiohttp")
+    
+    async def get(self, url: str, **kwargs):
+        """GET request."""
+        async with self._session.get(url, **kwargs) as response:
+            return await response.json()
+    
+    async def post(self, url: str, **kwargs):
+        """POST request."""
+        async with self._session.post(url, **kwargs) as response:
+            return await response.json()
+    
+    async def close(self):
+        """Close session."""
+        if self._session:
+            await self._session.close()
+
+
+class VectorStoreConnectionPool:
+    """Vector store connection pool."""
+    
+    def __init__(
+        self,
+        store_type: str,
+        config: Dict[str, Any],
+        pool_size: int = 10
+    ):
+        """Initialize vector store connection pool.
+        
+        Args:
+            store_type: Vector store type (chromadb, pinecone, weaviate)
+            config: Store configuration
+            pool_size: Pool size
+        """
+        self.store_type = store_type
+        self.config = config
+        self.pool_size = pool_size
+        self._connections = queue.Queue(maxsize=pool_size)
+        self._lock = threading.Lock()
+    
+    def initialize(self):
+        """Initialize pool."""
+        for _ in range(self.pool_size):
+            conn = self._create_connection()
+            self._connections.put(conn)
+    
+    def _create_connection(self):
+        """Create vector store connection."""
+        if self.store_type == "chromadb":
+            import chromadb
+            return chromadb.Client(self.config)
+        elif self.store_type == "pinecone":
+            import pinecone
+            pinecone.init(**self.config)
+            return pinecone
+        elif self.store_type == "weaviate":
+            import weaviate
+            return weaviate.Client(**self.config)
+        else:
+            raise ValueError(f"Unknown store type: {self.store_type}")
+    
+    @contextmanager
+    def connection(self):
+        """Get connection from pool."""
+        conn = self._connections.get(timeout=30)
+        try:
+            yield conn
+        finally:
+            self._connections.put(conn)
+    
+    def close(self):
+        """Close all connections."""
+        while not self._connections.empty():
+            conn = self._connections.get()
+            if hasattr(conn, 'close'):
+                conn.close()
+
+
+# Global pools
+_db_pool = None
+_http_pool = None
+_vector_pool = None
+
+
+async def get_db_pool() -> DatabaseConnectionPool:
+    """Get global database pool."""
+    global _db_pool
+    
+    if _db_pool is None:
+        import os
+        db_url = os.getenv("POSTGRES_URL", "postgresql://localhost/genxai")
+        _db_pool = DatabaseConnectionPool(db_url)
+        await _db_pool.initialize()
+    
+    return _db_pool
+
+
+async def get_http_pool() -> HTTPConnectionPool:
+    """Get global HTTP pool."""
+    global _http_pool
+    
+    if _http_pool is None:
+        _http_pool = HTTPConnectionPool()
+        await _http_pool.initialize()
+    
+    return _http_pool
+
+
+def get_vector_pool(store_type: str = "chromadb") -> VectorStoreConnectionPool:
+    """Get global vector store pool."""
+    global _vector_pool
+    
+    if _vector_pool is None:
+        import os
+        config = {}
+        
+        if store_type == "pinecone":
+            config = {
+                "api_key": os.getenv("PINECONE_API_KEY"),
+                "environment": os.getenv("PINECONE_ENV", "us-west1-gcp")
+            }
+        elif store_type == "weaviate":
+            config = {
+                "url": os.getenv("WEAVIATE_URL", "http://localhost:8080")
+            }
+        
+        _vector_pool = VectorStoreConnectionPool(store_type, config)
+        _vector_pool.initialize()
+    
+    return _vector_pool

genxai/security/audit.py

@@ -0,0 +1,304 @@
+"""Audit logging and approvals for governance."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+import json
+import os
+import sqlite3
+
+
+@dataclass
+class AuditEvent:
+    action: str
+    actor_id: str
+    resource_id: str
+    status: str
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    timestamp: datetime = field(default_factory=datetime.utcnow)
+
+
+@dataclass
+class ApprovalRequest:
+    request_id: str
+    action: str
+    resource_id: str
+    actor_id: str
+    status: str = "pending"
+    created_at: datetime = field(default_factory=datetime.utcnow)
+
+
+class AuditStore:
+    """SQLite-backed store for audit logs and approvals."""
+
+    def __init__(self, db_path: Path) -> None:
+        self.db_path = db_path
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._ensure_db()
+
+    def _connect(self) -> sqlite3.Connection:
+        return sqlite3.connect(self.db_path)
+
+    def _with_connection(self, handler):
+        conn = self._connect()
+        try:
+            return handler(conn)
+        finally:
+            conn.close()
+
+    def _ensure_db(self) -> None:
+        def _init(conn: sqlite3.Connection) -> None:
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                CREATE TABLE IF NOT EXISTS audit_events (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    action TEXT NOT NULL,
+                    actor_id TEXT NOT NULL,
+                    resource_id TEXT NOT NULL,
+                    status TEXT NOT NULL,
+                    metadata TEXT NOT NULL,
+                    timestamp TEXT NOT NULL
+                )
+                """
+            )
+            cursor.execute(
+                """
+                CREATE TABLE IF NOT EXISTS approval_requests (
+                    request_id TEXT PRIMARY KEY,
+                    action TEXT NOT NULL,
+                    resource_id TEXT NOT NULL,
+                    actor_id TEXT NOT NULL,
+                    status TEXT NOT NULL,
+                    created_at TEXT NOT NULL
+                )
+                """
+            )
+            conn.commit()
+
+        self._with_connection(_init)
+
+    def load_events(self) -> List[AuditEvent]:
+        def _load(conn: sqlite3.Connection):
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT action, actor_id, resource_id, status, metadata, timestamp FROM audit_events"
+            )
+            return cursor.fetchall()
+
+        rows = self._with_connection(_load)
+        events: List[AuditEvent] = []
+        for action, actor_id, resource_id, status, metadata, timestamp in rows:
+            events.append(
+                AuditEvent(
+                    action=action,
+                    actor_id=actor_id,
+                    resource_id=resource_id,
+                    status=status,
+                    metadata=json.loads(metadata),
+                    timestamp=datetime.fromisoformat(timestamp),
+                )
+            )
+        return events
+
+    def save_event(self, event: AuditEvent) -> None:
+        def _save(conn: sqlite3.Connection) -> None:
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                INSERT INTO audit_events (action, actor_id, resource_id, status, metadata, timestamp)
+                VALUES (?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    event.action,
+                    event.actor_id,
+                    event.resource_id,
+                    event.status,
+                    json.dumps(event.metadata, default=str),
+                    event.timestamp.isoformat(),
+                ),
+            )
+            conn.commit()
+
+        self._with_connection(_save)
+
+    def clear_events(self) -> None:
+        def _clear(conn: sqlite3.Connection) -> None:
+            cursor = conn.cursor()
+            cursor.execute("DELETE FROM audit_events")
+            conn.commit()
+
+        self._with_connection(_clear)
+
+    def load_requests(self) -> Dict[str, ApprovalRequest]:
+        def _load(conn: sqlite3.Connection):
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT request_id, action, resource_id, actor_id, status, created_at FROM approval_requests"
+            )
+            return cursor.fetchall()
+
+        rows = self._with_connection(_load)
+        requests: Dict[str, ApprovalRequest] = {}
+        for request_id, action, resource_id, actor_id, status, created_at in rows:
+            requests[request_id] = ApprovalRequest(
+                request_id=request_id,
+                action=action,
+                resource_id=resource_id,
+                actor_id=actor_id,
+                status=status,
+                created_at=datetime.fromisoformat(created_at),
+            )
+        return requests
+
+    def save_request(self, request: ApprovalRequest) -> None:
+        def _save(conn: sqlite3.Connection) -> None:
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                INSERT OR REPLACE INTO approval_requests
+                (request_id, action, resource_id, actor_id, status, created_at)
+                VALUES (?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    request.request_id,
+                    request.action,
+                    request.resource_id,
+                    request.actor_id,
+                    request.status,
+                    request.created_at.isoformat(),
+                ),
+            )
+            conn.commit()
+
+        self._with_connection(_save)
+
+    def clear_requests(self) -> None:
+        def _clear(conn: sqlite3.Connection) -> None:
+            cursor = conn.cursor()
+            cursor.execute("DELETE FROM approval_requests")
+            conn.commit()
+
+        self._with_connection(_clear)
+
+    def vacuum(self) -> None:
+        def _vacuum(conn: sqlite3.Connection) -> None:
+            conn.execute("VACUUM")
+
+        self._with_connection(_vacuum)
+
+
+class AuditLog:
+    """In-memory audit log."""
+
+    def __init__(self) -> None:
+        self._store = _get_audit_store()
+        self._events: List[AuditEvent] = self._store.load_events()
+
+    def record(self, event: AuditEvent) -> None:
+        self._events.append(event)
+        self._store.save_event(event)
+
+    def list_events(self) -> List[AuditEvent]:
+        self._events = self._store.load_events()
+        return list(self._events)
+
+    def clear(self) -> None:
+        self._events.clear()
+        self._store.clear_events()
+
+
+class ApprovalService:
+    """Simple in-memory approval workflow."""
+
+    def __init__(self) -> None:
+        self._store = _get_audit_store()
+        self._requests: Dict[str, ApprovalRequest] = self._store.load_requests()
+        self._counter = self._infer_counter()
+
+    def _infer_counter(self) -> int:
+        max_counter = 0
+        for request_id in self._requests:
+            if request_id.startswith("approval_"):
+                try:
+                    max_counter = max(max_counter, int(request_id.split("_", 1)[1]))
+                except ValueError:
+                    continue
+        return max_counter
+
+    def submit(self, action: str, resource_id: str, actor_id: str) -> ApprovalRequest:
+        self._counter += 1
+        request_id = f"approval_{self._counter}"
+        request = ApprovalRequest(
+            request_id=request_id,
+            action=action,
+            resource_id=resource_id,
+            actor_id=actor_id,
+        )
+        self._requests[request_id] = request
+        self._store.save_request(request)
+        return request
+
+    def approve(self, request_id: str) -> Optional[ApprovalRequest]:
+        request = self._requests.get(request_id)
+        if request:
+            request.status = "approved"
+            self._store.save_request(request)
+        return request
+
+    def reject(self, request_id: str) -> Optional[ApprovalRequest]:
+        request = self._requests.get(request_id)
+        if request:
+            request.status = "rejected"
+            self._store.save_request(request)
+        return request
+
+    def get(self, request_id: str) -> Optional[ApprovalRequest]:
+        return self._requests.get(request_id)
+
+    def clear(self) -> None:
+        self._requests.clear()
+        self._counter = 0
+        self._store.clear_requests()
+
+
+_audit_log: Optional[AuditLog] = None
+_approval_service: Optional[ApprovalService] = None
+_audit_store: Optional[AuditStore] = None
+
+
+def _get_default_db_path() -> Path:
+    return Path(__file__).resolve().parents[1] / "data" / "audit.db"
+
+
+def _get_audit_store() -> AuditStore:
+    global _audit_store
+    if _audit_store is None:
+        db_path = Path(os.getenv("GENXAI_AUDIT_DB", str(_get_default_db_path())))
+        _audit_store = AuditStore(db_path)
+    return _audit_store
+
+
+def get_audit_log() -> AuditLog:
+    global _audit_log
+    if _audit_log is None:
+        _audit_log = AuditLog()
+    return _audit_log
+
+
+def get_approval_service() -> ApprovalService:
+    global _approval_service
+    if _approval_service is None:
+        _approval_service = ApprovalService()
+    return _approval_service
+
+
+def reset_audit_services() -> None:
+    """Reset audit services (useful for tests)."""
+    global _audit_log, _approval_service, _audit_store
+    _audit_log = None
+    _approval_service = None
+    _audit_store = None