wizelit-sdk 0.1.23__py3-none-any.whl

wizelit_sdk/agent_wrapper/job.py

@@ -0,0 +1,379 @@
+"""
+Job class for managing execution context and logging in Wizelit Agent Wrapper.
+"""
+import logging
+import asyncio
+import uuid
+import time
+from datetime import datetime, UTC
+from typing import List, Optional, Awaitable, Any, TYPE_CHECKING
+from fastmcp import Context
+
+if TYPE_CHECKING:
+    from wizelit_sdk.database import DatabaseManager
+    from wizelit_sdk.agent_wrapper.streaming import LogStreamer
+
+
+class MemoryLogHandler(logging.Handler):
+    """
+    Custom logging handler that stores log messages in a list.
+    """
+
+    def __init__(self, logs_list: List[str]):
+        super().__init__()
+        self.logs_list = logs_list
+        self.setFormatter(logging.Formatter('%(message)s'))
+
+    def emit(self, record: logging.LogRecord) -> None:
+        """
+        Emit a log record by appending it to the logs list.
+        """
+        try:
+            # Format timestamp
+            ts = time.strftime("%H:%M:%S")
+
+            # Format message with level and timestamp
+            formatted_message = f"[{record.levelname}] [{ts}] {record.getMessage()}"
+
+            # Append to logs list
+            self.logs_list.append(formatted_message)
+        except Exception:
+            # Prevent exceptions in logging handler from breaking execution
+            self.handleError(record)
+
+
+class DatabaseLogHandler(logging.Handler):
+    """
+    Logging handler that persists log messages to PostgreSQL database.
+    Writes asynchronously to avoid blocking the logging thread.
+    """
+
+    def __init__(self, job_id: str, db_manager: 'DatabaseManager'):
+        super().__init__()
+        self.job_id = job_id
+        self.db_manager = db_manager
+        self.setFormatter(logging.Formatter('%(message)s'))
+
+    def emit(self, record: logging.LogRecord) -> None:
+        """
+        Emit a log record by writing it to the database asynchronously.
+        """
+        try:
+            # Import here to avoid circular dependency
+            from models.job import JobLogModel
+            from datetime import datetime
+
+            async def write_log():
+                try:
+                    async with self.db_manager.get_session() as session:
+                        log = JobLogModel(
+                            job_id=self.job_id,
+                            message=record.getMessage(),
+                            level=record.levelname,
+                            timestamp=datetime.now(UTC).replace(tzinfo=None)
+                        )
+                        session.add(log)
+                        await session.commit()
+                except Exception as e:
+                    # Log to stderr but don't break execution
+                    print(f"Error writing log to database: {e}", flush=True)
+
+            # Schedule async write without awaiting
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(write_log())
+            except RuntimeError:
+                # No event loop running - log warning
+                print("Warning: No event loop running, cannot write log to database", flush=True)
+        except Exception:
+            # Prevent exceptions in logging handler from breaking execution
+            self.handleError(record)
+
+
+class StreamingLogHandler(logging.Handler):
+    """
+    Logging handler that publishes log messages to Redis for real-time streaming.
+    Enables push-based log delivery without polling.
+    """
+
+    def __init__(self, job_id: str, log_streamer: 'LogStreamer'):
+        super().__init__()
+        self.job_id = job_id
+        self.log_streamer = log_streamer
+        self.setFormatter(logging.Formatter('%(message)s'))
+
+    def emit(self, record: logging.LogRecord) -> None:
+        """
+        Emit a log record by publishing it to Redis Pub/Sub.
+        """
+        try:
+            async def publish_log():
+                try:
+                    await self.log_streamer.publish_log(
+                        self.job_id,
+                        record.getMessage(),
+                        record.levelname
+                    )
+                except Exception as e:
+                    # Log to stderr but don't break execution
+                    print(f"Error streaming log: {e}", flush=True)
+
+            # Schedule async publish without awaiting
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(publish_log())
+            except RuntimeError:
+                # No event loop running - log warning
+                print("Warning: No event loop running, cannot stream log to Redis", flush=True)
+        except Exception as e:
+            # Prevent exceptions in logging handler from breaking execution
+            print(f"Error in StreamingLogHandler.emit: {e}", flush=True)
+            self.handleError(record)
+
+
+class Job:
+    """
+    Job instance that provides logging capabilities and execution context.
+    Each decorated function execution gets a Job instance injected.
+    """
+
+    def __init__(
+        self,
+        ctx: Context,
+        job_id: Optional[str] = None,
+        db_manager: Optional['DatabaseManager'] = None,
+        log_streamer: Optional['LogStreamer'] = None
+    ):
+        """
+        Initialize a Job instance.
+
+        Args:
+            ctx: FastMCP Context for progress reporting
+            job_id: Optional job identifier (generates UUID if not provided)
+            db_manager: Optional DatabaseManager for persisting logs
+            log_streamer: Optional LogStreamer for real-time log streaming
+        """
+        self._ctx = ctx
+        self._id = job_id or f"JOB-{str(uuid.uuid4())[:8]}"
+        self._status = "running"
+        self._logs: List[str] = []
+        self._result: Optional[str] = None
+        self._error: Optional[str] = None
+        self._db_manager = db_manager
+        self._log_streamer = log_streamer
+
+        # Set up logger
+        self._setup_logger(ctx)
+
+    @property
+    def id(self) -> str:
+        """Unique job identifier."""
+        return self._id
+
+    @property
+    def logger(self) -> logging.Logger:
+        """Python Logger instance configured with MCP streaming handler."""
+        return self._logger
+
+    @property
+    def logs(self) -> List[str]:
+        """List of log messages (timestamped strings)."""
+        return self._logs
+
+    @property
+    def status(self) -> str:
+        """Job status: 'running', 'completed', or 'failed'."""
+        return self._status
+
+    @status.setter
+    def status(self, value: str) -> None:
+        """Set job status and publish status change event."""
+        self._status = value
+        # Publish status change to Redis if streamer is available
+        if self._log_streamer:
+            print(f"[DEBUG] Publishing status change for job {self._id}: {value}", flush=True)
+            async def publish_status():
+                try:
+                    await self._log_streamer.publish_status_change(
+                        self._id,
+                        value,
+                        result=self._result,
+                        error=self._error
+                    )
+                    print(f"[DEBUG] Status change published successfully for job {self._id}", flush=True)
+                except Exception as e:
+                    print(f"Error publishing status change: {e}", flush=True)
+
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(publish_status())
+            except RuntimeError:
+                print("Warning: No event loop running, cannot publish status change to Redis", flush=True)
+        else:
+            print(f"[DEBUG] No log_streamer available for job {self._id}, skipping Redis publish", flush=True)
+
+    @property
+    def result(self) -> Optional[str | dict[str, Any]]:
+        """Job result (if completed successfully)."""
+        return self._result
+
+    @result.setter
+    def result(self, value: Optional[str | dict[str, Any]]) -> None:
+        """Set job result."""
+        self._result = value
+
+    @property
+    def error(self) -> Optional[str]:
+        """Job error message (if failed)."""
+        return self._error
+
+    @error.setter
+    def error(self, value: Optional[str]) -> None:
+        """Set job error message."""
+        self._error = value
+
+    async def _heartbeat(self, interval_seconds: float = 5.0) -> None:
+        """
+        Periodically append a heartbeat log while a job is running so the UI
+        has visible progress even during long operations.
+        """
+        start = time.monotonic()
+        while self._status == "running":
+            await asyncio.sleep(interval_seconds)
+            # Re-check in case status changed while sleeping
+            if self._status != "running":
+                break
+            elapsed = int(time.monotonic() - start)
+            # Use logger so logs are captured in memory and streamed if enabled
+            self.logger.info(f"⏳ Still working... ({elapsed}s)")
+
+    def run(
+        self,
+        coro: Awaitable[Any],
+        *,
+        heartbeat_interval: float = 5.0,
+    ) -> "asyncio.Task[Any]":
+        """
+        Run a coroutine in the background, managing heartbeat, status, result, and error.
+
+        This is intended for long-running jobs. It:
+        - Marks the job as running
+        - Starts a heartbeat logger
+        - Awaits the provided coroutine
+        - On success: stores the result (if string) and marks status 'completed'
+        - On failure: stores the error message and marks status 'failed'
+        """
+        import asyncio
+
+        async def _runner() -> Any:
+            self.status = "running"
+            # Persist initial job state
+            await self.persist_to_db()
+
+            heartbeat_task = asyncio.create_task(self._heartbeat(heartbeat_interval))
+            try:
+                result = await coro
+                # Store string results for convenience
+                if isinstance(result, str|dict):
+                    self.result = result
+                if self._status == "running":
+                    self.status = "completed"
+                    # Persist completion state
+                    await self.persist_to_db()
+                return result
+            except Exception as e:  # noqa: BLE001 - we deliberately capture all
+                self.error = str(e)
+                self.status = "failed"
+                # Persist failure state
+                await self.persist_to_db()
+                # Also log the error so it shows up in logs UI
+                self.logger.error(f"❌ [System] Error: {e}")
+                raise
+            finally:
+                # Stop heartbeat
+                heartbeat_task.cancel()
+                try:
+                    import contextlib
+
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await heartbeat_task
+                except Exception:
+                    # Ignore heartbeat shutdown errors
+                    pass
+
+        # Schedule the runner in the current event loop and return the Task
+        return asyncio.create_task(_runner())
+
+    def _setup_logger(self, ctx: Context) -> None:
+        """
+        Configure logger with custom handlers for streaming and storage.
+
+        Args:
+            ctx: FastMCP Context for progress reporting
+        """
+        _ = ctx  # ctx reserved for potential streaming handler setup
+        # Create logger with unique name per job
+        logger_name = f"wizelit.job.{self._id}"
+        self._logger = logging.getLogger(logger_name)
+
+        # Set level to INFO by default
+        self._logger.setLevel(logging.INFO)
+
+        # Remove any existing handlers to avoid duplicates
+        self._logger.handlers.clear()
+
+        # Add MemoryLogHandler for internal storage (backward compatibility)
+        memory_handler = MemoryLogHandler(self._logs)
+        memory_handler.setLevel(logging.INFO)
+        self._logger.addHandler(memory_handler)
+
+        # Add DatabaseLogHandler if db_manager provided
+        if self._db_manager:
+            db_handler = DatabaseLogHandler(self._id, self._db_manager)
+            db_handler.setLevel(logging.INFO)
+            self._logger.addHandler(db_handler)
+
+        # Add StreamingLogHandler if log_streamer provided
+        if self._log_streamer:
+            streaming_handler = StreamingLogHandler(self._id, self._log_streamer)
+            streaming_handler.setLevel(logging.INFO)
+            self._logger.addHandler(streaming_handler)
+
+        # Prevent propagation to root logger
+        self._logger.propagate = False
+
+    async def persist_to_db(self) -> None:
+        """
+        Persist the job state to the database.
+        Creates or updates the job record.
+        """
+        if not self._db_manager:
+            return
+
+        try:
+            from models.job import JobModel
+
+            async with self._db_manager.get_session() as session:
+                # Check if job already exists
+                existing_job = await session.get(JobModel, self._id)
+
+                if existing_job:
+                    # Update existing job
+                    existing_job.status = self._status
+                    existing_job.result = self._result
+                    existing_job.error = self._error
+                    existing_job.updated_at = datetime.now(UTC).replace(tzinfo=None)
+                else:
+                    # Create new job
+                    job = JobModel(
+                        id=self._id,
+                        status=self._status,
+                        result=self._result,
+                        error=self._error
+                    )
+                    session.add(job)
+
+                await session.commit()
+        except Exception as e:
+            # Log error but don't break execution
+            print(f"Error persisting job to database: {e}", flush=True)

wizelit_sdk/agent_wrapper/streaming.py

@@ -0,0 +1,206 @@
+"""
+Real-time log streaming using Redis Pub/Sub.
+Enables push-based log delivery from workers to hub without polling.
+"""
+import json
+import logging
+from datetime import datetime, UTC
+from typing import AsyncGenerator, Optional, Dict, Any
+import asyncio
+
+try:
+    import redis.asyncio as redis
+except ImportError:
+    redis = None
+
+logger = logging.getLogger(__name__)
+
+
+class LogStreamer:
+    """
+    Manages real-time log streaming via Redis Pub/Sub.
+
+    Workers publish log events to job-specific channels.
+    Hub subscribes to these channels for real-time updates.
+    """
+
+    def __init__(self, redis_url: str = "redis://localhost:6379"):
+        """
+        Initialize the log streamer.
+
+        Args:
+            redis_url: Redis connection URL
+        """
+        if redis is None:
+            raise ImportError(
+                "redis package is required for log streaming. "
+                "Install it with: pip install redis"
+            )
+
+        self.redis_url = redis_url
+        self._redis: Optional[redis.Redis] = None
+        self._pubsub: Optional[redis.client.PubSub] = None
+
+    async def _ensure_connected(self) -> redis.Redis:
+        """Ensure Redis connection is established."""
+        if self._redis is None:
+            try:
+                self._redis = redis.from_url(
+                    self.redis_url,
+                    decode_responses=True,
+                    socket_connect_timeout=5,
+                    socket_keepalive=True,
+                )
+                # Test connection
+                await self._redis.ping()
+                logger.info(f"Connected to Redis at {self.redis_url}")
+            except Exception as e:
+                logger.error(f"Failed to connect to Redis: {e}")
+                raise
+        return self._redis
+
+    async def publish_log(
+        self,
+        job_id: str,
+        message: str,
+        level: str = "INFO",
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """
+        Publish a log event to Redis.
+
+        Args:
+            job_id: Job identifier
+            message: Log message
+            level: Log level (INFO, ERROR, WARNING, DEBUG)
+            metadata: Additional metadata to include
+        """
+        try:
+            redis_client = await self._ensure_connected()
+
+            event = {
+                "job_id": job_id,
+                "message": message,
+                "level": level,
+                "timestamp": datetime.now(UTC).isoformat(),
+            }
+
+            if metadata:
+                event["metadata"] = metadata
+
+            channel = f"job:{job_id}:logs"
+            await redis_client.publish(channel, json.dumps(event))
+
+        except Exception as e:
+            # Don't let streaming errors break the main execution
+            logger.error(f"Failed to publish log for job {job_id}: {e}")
+
+    async def publish_status_change(
+        self,
+        job_id: str,
+        status: str,
+        result: Optional[Any] = None,
+        error: Optional[str] = None
+    ) -> None:
+        """
+        Publish a job status change event.
+
+        Args:
+            job_id: Job identifier
+            status: New status (running, completed, failed)
+            result: Job result (for completed jobs)
+            error: Error message (for failed jobs)
+        """
+        try:
+            redis_client = await self._ensure_connected()
+
+            event = {
+                "job_id": job_id,
+                "status": status,
+                "timestamp": datetime.now(UTC).isoformat(),
+            }
+
+            if result is not None:
+                event["result"] = result
+            if error is not None:
+                event["error"] = error
+
+            channel = f"job:{job_id}:status"
+            await redis_client.publish(channel, json.dumps(event))
+            logger.info(f"Published status change for job {job_id}: {status}")
+
+        except Exception as e:
+            logger.error(f"Failed to publish status change for job {job_id}: {e}")
+
+    async def subscribe_logs(
+        self,
+        job_id: str,
+        timeout: Optional[float] = None
+    ) -> AsyncGenerator[Dict[str, Any], None]:
+        """
+        Subscribe to log stream for a specific job.
+
+        Args:
+            job_id: Job identifier
+            timeout: Optional timeout in seconds
+
+        Yields:
+            Dict containing log event data
+        """
+        redis_client = await self._ensure_connected()
+        pubsub = redis_client.pubsub()
+
+        try:
+            # Subscribe to both logs and status channels
+            log_channel = f"job:{job_id}:logs"
+            status_channel = f"job:{job_id}:status"
+            await pubsub.subscribe(log_channel, status_channel)
+            logger.info(f"Subscribed to channels for job {job_id}")
+
+            start_time = asyncio.get_event_loop().time() if timeout else None
+
+            async for message in pubsub.listen():
+                # Check timeout
+                if timeout and start_time:
+                    elapsed = asyncio.get_event_loop().time() - start_time
+                    if elapsed > timeout:
+                        logger.info(f"Subscription timeout for job {job_id}")
+                        break
+
+                if message["type"] == "message":
+                    try:
+                        event = json.loads(message["data"])
+                        yield event
+
+                        # Stop listening if job is completed or failed
+                        if event.get("status") in ["completed", "failed"]:
+                            logger.info(f"Job {job_id} finished with status: {event.get('status')}")
+                            break
+
+                    except json.JSONDecodeError as e:
+                        logger.error(f"Failed to decode message: {e}")
+
+        except asyncio.CancelledError:
+            logger.info(f"Subscription cancelled for job {job_id}")
+        except Exception as e:
+            logger.error(f"Error in subscription for job {job_id}: {e}")
+        finally:
+            await pubsub.unsubscribe()
+            await pubsub.close()
+
+    async def close(self) -> None:
+        """Close Redis connection."""
+        if self._redis:
+            await self._redis.close()
+            await self._redis.connection_pool.disconnect()
+            self._redis = None
+            logger.info("Redis connection closed")
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self._ensure_connected()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()

wizelit_sdk/agent_wrapper/utils.py

@@ -0,0 +1,45 @@
+"""Utility functions for Wizelit Agent Wrapper."""
+
+
+def greet(name: str = "World") -> str:
+    """
+    Generate a greeting message.
+    
+    Args:
+        name: Name to greet. Defaults to "World".
+        
+    Returns:
+        A greeting string.
+        
+    Example:
+        >>> greet("Alice")
+        'Hello, Alice!'
+        >>> greet()
+        'Hello, World!'
+    """
+    return f"Hello, {name}!"
+
+def greet_many(names: list[str]) -> list[str]:
+    """
+    Generate a greeting message for each name.
+    
+    Args:
+        names: List of names to greet.
+        
+    Returns:
+        A list of greeting strings.
+    """
+    return [f"Hello, {name}!" for name in names]
+
+
+def greet_many3(names: list[str]) -> str:
+    """
+    Generate a greeting message for each name.
+    
+    Args:
+        names: List of names to greet.
+        
+    Returns:
+        A list of greeting strings.
+    """
+    return ', '.join([f"Hello, {name}!" for name in names])