codepp 0.0.437__py3-none-any.whl

code_puppy/round_robin_model.py

@@ -0,0 +1,150 @@
+import threading
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass, field
+from typing import Any, AsyncIterator, List
+
+from pydantic_ai._run_context import RunContext
+from pydantic_ai.models import (
+    Model,
+    ModelMessage,
+    ModelRequestParameters,
+    ModelResponse,
+    ModelSettings,
+    StreamedResponse,
+)
+
+try:
+    from opentelemetry.context import get_current_span
+except ImportError:
+    # If opentelemetry is not installed, provide a dummy implementation
+    def get_current_span():
+        class DummySpan:
+            def is_recording(self):
+                return False
+
+            def set_attributes(self, attributes):
+                pass
+
+        return DummySpan()
+
+
+@dataclass(init=False)
+class RoundRobinModel(Model):
+    """A model that cycles through multiple models in a round-robin fashion.
+
+    This model distributes requests across multiple candidate models to help
+    overcome rate limits or distribute load.
+    """
+
+    models: List[Model]
+    _current_index: int = field(default=0, repr=False)
+    _model_name: str = field(repr=False)
+    _rotate_every: int = field(default=1, repr=False)
+    _request_count: int = field(default=0, repr=False)
+    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
+
+    def __init__(
+        self,
+        *models: Model,
+        rotate_every: int = 1,
+        settings: ModelSettings | None = None,
+    ):
+        """Initialize a round-robin model instance.
+
+        Args:
+            models: The model instances to cycle through.
+            rotate_every: Number of requests before rotating to the next model (default: 1).
+            settings: Model settings that will be used as defaults for this model.
+        """
+        super().__init__(settings=settings)
+        if not models:
+            raise ValueError("At least one model must be provided")
+        if rotate_every < 1:
+            raise ValueError("rotate_every must be at least 1")
+        self.models = list(models)
+        self._current_index = 0
+        self._request_count = 0
+        self._rotate_every = rotate_every
+        self._lock = threading.Lock()
+
+    @property
+    def model_name(self) -> str:
+        """The model name showing this is a round-robin model with its candidates."""
+        base_name = f"round_robin:{','.join(model.model_name for model in self.models)}"
+        if self._rotate_every != 1:
+            return f"{base_name}:rotate_every={self._rotate_every}"
+        return base_name
+
+    @property
+    def system(self) -> str:
+        """System prompt from the current model."""
+        return self.models[self._current_index].system
+
+    @property
+    def base_url(self) -> str | None:
+        """Base URL from the current model."""
+        return self.models[self._current_index].base_url
+
+    def _get_next_model(self) -> Model:
+        """Get the next model in the round-robin sequence and update the index."""
+        with self._lock:
+            model = self.models[self._current_index]
+            self._request_count += 1
+            if self._request_count >= self._rotate_every:
+                self._current_index = (self._current_index + 1) % len(self.models)
+                self._request_count = 0
+            return model
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        """Make a request using the next model in the round-robin sequence."""
+        current_model = self._get_next_model()
+        # Use prepare_request to merge settings and customize parameters
+        merged_settings, prepared_params = current_model.prepare_request(
+            model_settings, model_request_parameters
+        )
+
+        try:
+            response = await current_model.request(
+                messages, merged_settings, prepared_params
+            )
+            self._set_span_attributes(current_model)
+            return response
+        except Exception:
+            # Unlike FallbackModel, we don't try other models here
+            # The round-robin strategy is about distribution, not failover
+            raise
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        """Make a streaming request using the next model in the round-robin sequence."""
+        current_model = self._get_next_model()
+        # Use prepare_request to merge settings and customize parameters
+        merged_settings, prepared_params = current_model.prepare_request(
+            model_settings, model_request_parameters
+        )
+
+        async with current_model.request_stream(
+            messages, merged_settings, prepared_params, run_context
+        ) as response:
+            self._set_span_attributes(current_model)
+            yield response
+
+    def _set_span_attributes(self, model: Model):
+        """Set span attributes for observability."""
+        with suppress(Exception):
+            span = get_current_span()
+            if span.is_recording():
+                attributes = getattr(span, "attributes", {})
+                if attributes.get("gen_ai.request.model") == self.model_name:
+                    span.set_attributes(model.model_attributes(model))

code_puppy/scheduler/__init__.py

@@ -0,0 +1,41 @@
+"""Code Puppy Scheduler - Run scheduled prompts automatically.
+
+This module provides a cross-platform scheduler daemon that executes
+Code Puppy prompts on configurable schedules (intervals, cron expressions).
+
+Components:
+    - config: Task definitions and JSON persistence
+    - daemon: Background scheduler process
+    - executor: Task execution logic
+    - platform: Cross-platform daemon management
+"""
+
+from code_puppy.scheduler.config import (
+    SCHEDULER_LOG_DIR,
+    SCHEDULER_PID_FILE,
+    SCHEDULES_FILE,
+    ScheduledTask,
+    add_task,
+    delete_task,
+    get_task,
+    load_tasks,
+    save_tasks,
+    toggle_task,
+    update_task,
+)
+from code_puppy.scheduler.daemon import start_daemon_background
+
+__all__ = [
+    "ScheduledTask",
+    "load_tasks",
+    "save_tasks",
+    "add_task",
+    "update_task",
+    "delete_task",
+    "get_task",
+    "toggle_task",
+    "start_daemon_background",
+    "SCHEDULES_FILE",
+    "SCHEDULER_PID_FILE",
+    "SCHEDULER_LOG_DIR",
+]

code_puppy/scheduler/__main__.py

@@ -0,0 +1,9 @@
+"""Entry point for running scheduler daemon directly.
+
+Usage: python -m code_puppy.scheduler
+"""
+
+from code_puppy.scheduler.daemon import start_daemon
+
+if __name__ == "__main__":
+    start_daemon(foreground=True)

code_puppy/scheduler/cli.py

@@ -0,0 +1,118 @@
+"""CLI subcommands for the scheduler.
+
+Handles command-line operations like starting/stopping the daemon,
+listing tasks, and running tasks immediately.
+"""
+
+from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning
+
+
+def handle_scheduler_start() -> bool:
+    """Start the scheduler daemon in background."""
+    from code_puppy.scheduler.daemon import get_daemon_pid, start_daemon_background
+
+    pid = get_daemon_pid()
+    if pid:
+        emit_warning(f"Scheduler daemon already running (PID {pid})")
+        return True
+
+    emit_info("Starting scheduler daemon...")
+
+    if start_daemon_background():
+        pid = get_daemon_pid()
+        emit_success(f"Scheduler daemon started (PID {pid})")
+        return True
+    else:
+        emit_error("Failed to start scheduler daemon")
+        return False
+
+
+def handle_scheduler_stop() -> bool:
+    """Stop the scheduler daemon."""
+    from code_puppy.scheduler.daemon import get_daemon_pid, stop_daemon
+
+    pid = get_daemon_pid()
+    if not pid:
+        emit_info("Scheduler daemon is not running")
+        return True
+
+    emit_info(f"Stopping scheduler daemon (PID {pid})...")
+
+    if stop_daemon():
+        emit_success("Scheduler daemon stopped")
+        return True
+    else:
+        emit_error("Failed to stop scheduler daemon")
+        return False
+
+
+def handle_scheduler_status() -> bool:
+    """Show scheduler daemon status."""
+    from code_puppy.scheduler.config import load_tasks
+    from code_puppy.scheduler.daemon import get_daemon_pid
+
+    pid = get_daemon_pid()
+    if pid:
+        emit_success(f"🐕 Scheduler daemon: RUNNING (PID {pid})")
+    else:
+        emit_warning("🐕 Scheduler daemon: STOPPED")
+
+    tasks = load_tasks()
+    enabled_count = sum(1 for t in tasks if t.enabled)
+
+    emit_info(f"\n📅 Scheduled tasks: {len(tasks)} total, {enabled_count} enabled")
+
+    if tasks:
+        emit_info("\nTasks:")
+        for task in tasks:
+            status_icon = "🟢" if task.enabled else "🔴"
+            last_run = task.last_run[:19] if task.last_run else "never"
+            emit_info(
+                f"  {status_icon} {task.name} ({task.schedule_type}: {task.schedule_value})"
+            )
+            emit_info(
+                f"      Last run: {last_run}, Status: {task.last_status or 'pending'}"
+            )
+
+    return True
+
+
+def handle_scheduler_list() -> bool:
+    """List all scheduled tasks."""
+    from code_puppy.scheduler.config import load_tasks
+
+    tasks = load_tasks()
+
+    if not tasks:
+        emit_info("No scheduled tasks configured.")
+        emit_info("Use '/scheduler' to create one.")
+        return True
+
+    emit_info(f"📅 Scheduled Tasks ({len(tasks)}):\n")
+
+    for task in tasks:
+        status = "🟢 enabled" if task.enabled else "🔴 disabled"
+        emit_info(f"  [{task.id}] {task.name}")
+        emit_info(f"      Status: {status}")
+        emit_info(f"      Schedule: {task.schedule_type} ({task.schedule_value})")
+        emit_info(f"      Agent: {task.agent}, Model: {task.model or 'default'}")
+        if task.last_run:
+            emit_info(f"      Last run: {task.last_run[:19]} ({task.last_status})")
+        emit_info("")
+
+    return True
+
+
+def handle_scheduler_run(task_id: str) -> bool:
+    """Run a specific task immediately."""
+    from code_puppy.scheduler.executor import run_task_by_id
+
+    emit_info(f"Running task {task_id}...")
+    success, message = run_task_by_id(task_id)
+
+    if success:
+        emit_success(message)
+    else:
+        emit_error(message)
+
+    return success

code_puppy/scheduler/config.py

@@ -0,0 +1,126 @@
+"""Scheduler configuration and task management.
+
+Handles ScheduledTask dataclass definition and JSON persistence
+for scheduled Code Puppy tasks.
+"""
+
+import json
+import os
+import uuid
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional
+
+# Import from existing config
+from code_puppy.config import DATA_DIR
+
+SCHEDULES_FILE = os.path.join(DATA_DIR, "scheduled_tasks.json")
+SCHEDULER_PID_FILE = os.path.join(DATA_DIR, "scheduler.pid")
+SCHEDULER_LOG_DIR = os.path.join(DATA_DIR, "scheduler_logs")
+
+
+@dataclass
+class ScheduledTask:
+    """A scheduled Code Puppy task."""
+
+    id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
+    name: str = ""
+    prompt: str = ""
+    agent: str = "code-puppy"
+    model: str = ""  # Uses default if empty
+    schedule_type: str = "interval"  # "interval", "cron", "daily", "hourly"
+    schedule_value: str = "1h"  # e.g., "30m", "1h", "0 9 * * *" for cron
+    working_directory: str = "."
+    log_file: str = ""  # Auto-generated if empty
+    enabled: bool = True
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    last_run: Optional[str] = None
+    last_status: Optional[str] = None  # "success", "failed", "running"
+    last_exit_code: Optional[int] = None
+
+    def __post_init__(self):
+        if not self.log_file:
+            self.log_file = os.path.join(SCHEDULER_LOG_DIR, f"{self.id}.log")
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ScheduledTask":
+        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
+
+
+def ensure_scheduler_dirs() -> None:
+    """Create scheduler directories if they don't exist."""
+    os.makedirs(SCHEDULER_LOG_DIR, mode=0o700, exist_ok=True)
+
+
+def load_tasks() -> List[ScheduledTask]:
+    """Load all scheduled tasks from JSON file."""
+    ensure_scheduler_dirs()
+    if not os.path.exists(SCHEDULES_FILE):
+        return []
+    try:
+        with open(SCHEDULES_FILE, "r") as f:
+            data = json.load(f)
+            return [ScheduledTask.from_dict(t) for t in data]
+    except (json.JSONDecodeError, IOError):
+        return []
+
+
+def save_tasks(tasks: List[ScheduledTask]) -> None:
+    """Save all scheduled tasks to JSON file."""
+    ensure_scheduler_dirs()
+    temp_path = Path(SCHEDULES_FILE).with_suffix(".tmp")
+    with open(temp_path, "w", encoding="utf-8") as f:
+        json.dump([t.to_dict() for t in tasks], f, indent=2, ensure_ascii=False)
+    temp_path.replace(SCHEDULES_FILE)
+
+
+def add_task(task: ScheduledTask) -> None:
+    """Add a new scheduled task."""
+    tasks = load_tasks()
+    tasks.append(task)
+    save_tasks(tasks)
+
+
+def update_task(task: ScheduledTask) -> bool:
+    """Update an existing task. Returns True if found and updated."""
+    tasks = load_tasks()
+    for i, t in enumerate(tasks):
+        if t.id == task.id:
+            tasks[i] = task
+            save_tasks(tasks)
+            return True
+    return False
+
+
+def delete_task(task_id: str) -> bool:
+    """Delete a task by ID. Returns True if found and deleted."""
+    tasks = load_tasks()
+    original_len = len(tasks)
+    tasks = [t for t in tasks if t.id != task_id]
+    if len(tasks) < original_len:
+        save_tasks(tasks)
+        return True
+    return False
+
+
+def get_task(task_id: str) -> Optional[ScheduledTask]:
+    """Get a task by ID."""
+    tasks = load_tasks()
+    for t in tasks:
+        if t.id == task_id:
+            return t
+    return None
+
+
+def toggle_task(task_id: str) -> Optional[bool]:
+    """Toggle a task's enabled state. Returns new state or None if not found."""
+    task = get_task(task_id)
+    if task:
+        task.enabled = not task.enabled
+        update_task(task)
+        return task.enabled
+    return None