caption-flow 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

caption_flow/__init__.py

@@ -3,7 +3,8 @@
 __version__ = "0.1.0"
 
 from .orchestrator import Orchestrator
-from .worker import Worker
+from .workers.data import DataWorker
+from .workers.caption import CaptionWorker
 from .monitor import Monitor
 
-__all__ = ["Orchestrator", "Worker", "Monitor"]
+__all__ = ["Orchestrator", "DataWorker", "CaptionWorker", "Monitor"]

caption_flow/cli.py

@@ -15,7 +15,6 @@ from rich.logging import RichHandler
 from datetime import datetime
 
 from .orchestrator import Orchestrator
-from .worker import Worker
 from .monitor import Monitor
 from .utils.certificates import CertificateManager
 
@@ -168,8 +167,6 @@ def orchestrator(ctx, config: Optional[str], **kwargs):
     else:
         config_data = base_config
 
-    console.print(f"Config contents: {config_data}")
-
     # Apply CLI overrides
     if kwargs.get("port"):
         config_data["port"] = kwargs["port"]
@@ -242,11 +239,11 @@ def worker(ctx, config: Optional[str], **kwargs):
 
     # Choose worker type
     if kwargs.get("vllm") or config_data.get("vllm"):
-        from .worker_vllm import VLLMWorker
+        from .workers.caption import CaptionWorker
 
-        worker_instance = VLLMWorker(config_data)
+        worker_instance = CaptionWorker(config_data)
     else:
-        worker_instance = Worker(config_data)
+        raise ValueError(f"Not sure how to handle worker for {config_data.get('type')} type setup.")
 
     try:
         asyncio.run(worker_instance.start())
@@ -262,84 +259,92 @@ def worker(ctx, config: Optional[str], **kwargs):
 @click.option("--no-verify-ssl", is_flag=True, help="Skip SSL verification")
 @click.option("--debug", is_flag=True, help="Enable debug output")
 @click.pass_context
-def monitor(ctx, config: Optional[str], server: Optional[str], token: Optional[str], 
-           no_verify_ssl: bool, debug: bool):
+def monitor(
+    ctx,
+    config: Optional[str],
+    server: Optional[str],
+    token: Optional[str],
+    no_verify_ssl: bool,
+    debug: bool,
+):
     """Start the monitoring TUI."""
-    
+
     # Enable debug logging if requested
     if debug:
         setup_logging(verbose=True)
         console.print("[yellow]Debug mode enabled[/yellow]")
-    
+
     # Load configuration
-    base_config = ConfigManager.find_config('monitor', config)
-    
+    base_config = ConfigManager.find_config("monitor", config)
+
     if not base_config:
         # Try to find monitor config in orchestrator config as fallback
-        orch_config = ConfigManager.find_config('orchestrator')
-        if orch_config and 'monitor' in orch_config:
-            base_config = {'monitor': orch_config['monitor']}
+        orch_config = ConfigManager.find_config("orchestrator")
+        if orch_config and "monitor" in orch_config:
+            base_config = {"monitor": orch_config["monitor"]}
             console.print("[dim]Using monitor config from orchestrator.yaml[/dim]")
         else:
             base_config = {}
             if not server or not token:
                 console.print("[yellow]No monitor config found, using CLI args[/yellow]")
-    
+
     # Handle different config structures
     # Case 1: Config has top-level 'monitor' section
-    if 'monitor' in base_config:
-        config_data = base_config['monitor']
+    if "monitor" in base_config:
+        config_data = base_config["monitor"]
     # Case 2: Config IS the monitor config (no wrapper)
     else:
         config_data = base_config
-    
+
     # Apply CLI overrides (CLI always wins)
     if server:
-        config_data['server'] = server
+        config_data["server"] = server
     if token:
-        config_data['token'] = token
+        config_data["token"] = token
     if no_verify_ssl:
-        config_data['verify_ssl'] = False
-    
+        config_data["verify_ssl"] = False
+
     # Debug output
     if debug:
         console.print("\n[cyan]Final monitor configuration:[/cyan]")
         console.print(f"  Server: {config_data.get('server', 'NOT SET')}")
-        console.print(f"  Token: {'***' + config_data.get('token', '')[-4:] if config_data.get('token') else 'NOT SET'}")
+        console.print(
+            f"  Token: {'***' + config_data.get('token', '')[-4:] if config_data.get('token') else 'NOT SET'}"
+        )
         console.print(f"  Verify SSL: {config_data.get('verify_ssl', True)}")
         console.print()
-    
+
     # Validate required fields
-    if not config_data.get('server'):
+    if not config_data.get("server"):
         console.print("[red]Error: --server required (or set 'server' in monitor.yaml)[/red]")
         console.print("\n[dim]Example monitor.yaml:[/dim]")
         console.print("server: wss://localhost:8765")
         console.print("token: your-token-here")
         sys.exit(1)
-        
-    if not config_data.get('token'):
+
+    if not config_data.get("token"):
         console.print("[red]Error: --token required (or set 'token' in monitor.yaml)[/red]")
         console.print("\n[dim]Example monitor.yaml:[/dim]")
         console.print("server: wss://localhost:8765")
         console.print("token: your-token-here")
         sys.exit(1)
-    
+
     # Set defaults for optional settings
-    config_data.setdefault('refresh_interval', 1.0)
-    config_data.setdefault('show_inactive_workers', False)
-    config_data.setdefault('max_log_lines', 100)
-    
+    config_data.setdefault("refresh_interval", 1.0)
+    config_data.setdefault("show_inactive_workers", False)
+    config_data.setdefault("max_log_lines", 100)
+
     # Create and start monitor
     try:
         monitor_instance = Monitor(config_data)
-        
+
         if debug:
             console.print("[green]Starting monitor...[/green]")
             console.print(f"[dim]Connecting to: {config_data['server']}[/dim]")
             sys.exit(1)
-        
+
         asyncio.run(monitor_instance.start())
-        
+
     except KeyboardInterrupt:
         console.print("\n[yellow]Closing monitor...[/yellow]")
     except ConnectionRefusedError:
@@ -350,9 +355,11 @@ def monitor(ctx, config: Optional[str], server: Optional[str], token: Optional[s
         console.print(f"\n[red]Error starting monitor: {e}[/red]")
         if debug:
             import traceback
+
             traceback.print_exc()
         sys.exit(1)
 
+
 @main.command()
 @click.option("--config", type=click.Path(exists=True), help="Configuration file")
 @click.option("--server", help="Orchestrator WebSocket URL")
@@ -376,11 +383,21 @@ def reload_config(
     if not server or not token:
         base_config = ConfigManager.find_config("orchestrator", config) or {}
         admin_config = base_config.get("admin", {})
+        admin_tokens = base_config.get("orchestrator", {}).get("auth", {}).get("admin_tokens", [])
+        has_admin_tokens = False
+        if len(admin_tokens) > 0:
+            has_admin_tokens = True
+            first_admin_token = admin_tokens[0].get("token", None)
+        # Do not print sensitive admin token to console.
 
         if not server:
-            server = admin_config.get("server")
+            server = admin_config.get("server", "ws://localhost:8765")
         if not token:
-            token = admin_config.get("token")
+            token = admin_config.get("token", None)
+            if token is None and has_admin_tokens:
+                # grab the first one, we'll just assume we're localhost.
+                console.print("Using first admin token.")
+                token = first_admin_token
 
     if not server:
         console.print("[red]Error: --server required (or set in config)[/red]")
@@ -447,7 +464,7 @@ def reload_config(
                     return True
                 else:
                     error = reload_response.get("error", "Unknown error")
-                    console.print(f"[red]Reload failed: {error}[/red]")
+                    console.print(f"[red]Reload failed: {error} ({reload_response=})[/red]")
                     return False
 
         except Exception as e:

caption_flow/models.py

@@ -1,9 +1,9 @@
 """Data models for CaptionFlow."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
-from typing import Optional
+from typing import Any, Dict, List, Optional
 
 
 class JobStatus(Enum):
@@ -51,9 +51,10 @@ class Caption:
     timestamp: datetime
     caption_count: int = 1  # Number of captions generated for this item
     caption: Optional[str] = None
-    captions: Optional[list] = None
+    captions: Optional[List[str]] = None
+    outputs: Dict[str, List[str]] = field(default_factory=dict)
     quality_score: Optional[float] = None
-    quality_scores: Optional[list] = None
+    quality_scores: Optional[List[float]] = None
 
     # Image metadata
     image_width: Optional[int] = None
@@ -66,6 +67,7 @@ class Caption:
     total_captions: Optional[int] = None  # Total captions for this image
     processing_time_ms: Optional[float] = None
     chunk_id: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self):
         if self.caption is None and self.captions is None:

caption_flow/monitor.py

@@ -166,6 +166,8 @@ class Monitor:
         # Filter out rate stats (already shown in rates panel)
         for key, value in self.stats.items():
             if key not in ["current_rate", "average_rate", "expected_rate"]:
+                if isinstance(value, dict):
+                    value = json.dumps(value, indent=2)
                 stats_table.add_row(key.replace("_", " ").title(), str(value))
 
         layout["stats"].update(Panel(stats_table, title="System Statistics", border_style="green"))
@@ -178,9 +180,18 @@ class Monitor:
         leaderboard_table.add_column("Trust", style="green")
 
         for i, contributor in enumerate(self.leaderboard[:10], 1):
+            # Format name with active worker count
+            name = contributor.get("name", "Unknown")
+            active_workers = contributor.get("active_workers", 0)
+
+            if active_workers > 0:
+                name_display = f"{name} [bright_green](x{active_workers})[/bright_green]"
+            else:
+                name_display = f"{name} [dim](offline)[/dim]"
+
             leaderboard_table.add_row(
                 str(i),
-                contributor.get("name", "Unknown"),
+                name_display,
                 str(contributor.get("total_captions", 0)),
                 "⭐" * contributor.get("trust_level", 0),
             )
@@ -188,7 +199,6 @@ class Monitor:
         layout["leaderboard"].update(
             Panel(leaderboard_table, title="Top Contributors", border_style="yellow")
         )
-
         # Activity panel
         activity_text = Text()
         for activity in self.recent_activity[-10:]: