caption-flow 0.3.4__py3-none-any.whl → 0.4.1__py3-none-any.whl

caption_flow/utils/json_utils.py

@@ -1,37 +1,41 @@
 """JSON serialization utilities for handling special types like datetime."""
 
 import json
-from datetime import datetime, date
-from decimal import Decimal
-from pathlib import Path
-from typing import Any, Dict, List, Union
 from dataclasses import asdict, is_dataclass
+from datetime import date, datetime
+from decimal import Decimal
 from enum import Enum
+from pathlib import Path
+from typing import Any, Dict, Union
 
 
 def safe_json_dumps(obj: Any, **kwargs) -> str:
-    """
-    Safely serialize objects to JSON, handling special types.
+    """Safely serialize objects to JSON, handling special types.
 
     Args:
+    ----
         obj: Object to serialize
         **kwargs: Additional arguments to pass to json.dumps
 
     Returns:
+    -------
         JSON string representation
+
     """
     return json.dumps(obj, default=json_serializer, **kwargs)
 
 
 def safe_dict(obj: Any) -> Dict[str, Any]:
-    """
-    Convert an object to a dictionary, handling special types.
+    """Convert an object to a dictionary, handling special types.
 
     Args:
+    ----
         obj: Object to convert (dataclass, dict, etc.)
 
     Returns:
+    -------
         Dictionary with JSON-serializable values
+
     """
     if is_dataclass(obj):
         data = asdict(obj)
@@ -46,14 +50,16 @@ def safe_dict(obj: Any) -> Dict[str, Any]:
 
 
 def sanitize_dict(data: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Recursively sanitize a dictionary to ensure all values are JSON-serializable.
+    """Recursively sanitize a dictionary to ensure all values are JSON-serializable.
 
     Args:
+    ----
         data: Dictionary to sanitize
 
     Returns:
+    -------
         Sanitized dictionary
+
     """
     result = {}
 
@@ -83,14 +89,16 @@ def sanitize_dict(data: Dict[str, Any]) -> Dict[str, Any]:
 
 
 def sanitize_value(value: Any) -> Any:
-    """
-    Sanitize a single value for JSON serialization.
+    """Sanitize a single value for JSON serialization.
 
     Args:
+    ----
         value: Value to sanitize
 
     Returns:
+    -------
         JSON-serializable value
+
     """
     if value is None:
         return None
@@ -115,17 +123,20 @@ def sanitize_value(value: Any) -> Any:
 
 
 def json_serializer(obj: Any) -> Any:
-    """
-    Default JSON serializer for special types.
+    """Default JSON serializer for special types.
 
     Args:
+    ----
         obj: Object to serialize
 
     Returns:
+    -------
         JSON-serializable representation
 
     Raises:
+    ------
         TypeError: If object type is not supported
+
     """
     if isinstance(obj, (datetime, date)):
         return obj.isoformat()
@@ -146,14 +157,16 @@ def json_serializer(obj: Any) -> Any:
 
 
 def parse_datetime(dt_string: Union[str, datetime, None]) -> Union[datetime, None]:
-    """
-    Parse a datetime string or return existing datetime.
+    """Parse a datetime string or return existing datetime.
 
     Args:
+    ----
         dt_string: ISO format datetime string, datetime object, or None
 
     Returns:
+    -------
         datetime object or None
+
     """
     if dt_string is None:
         return None
@@ -171,31 +184,35 @@ def parse_datetime(dt_string: Union[str, datetime, None]) -> Union[datetime, Non
 
 # Convenience functions for common use cases
 def to_json_dict(obj: Any) -> Dict[str, Any]:
-    """
-    Convert any object to a JSON-serializable dictionary.
+    """Convert any object to a JSON-serializable dictionary.
 
     This is a convenience wrapper around safe_dict.
 
     Args:
+    ----
         obj: Object to convert
 
     Returns:
+    -------
         JSON-serializable dictionary
+
     """
     return safe_dict(obj)
 
 
 def to_json_string(obj: Any, indent: int = None) -> str:
-    """
-    Convert any object to a JSON string.
+    """Convert any object to a JSON string.
 
     This is a convenience wrapper around safe_json_dumps.
 
     Args:
+    ----
         obj: Object to convert
         indent: Number of spaces for indentation (None for compact)
 
     Returns:
+    -------
         JSON string
+
     """
     return safe_json_dumps(obj, indent=indent)

caption_flow/utils/prompt_template.py

@@ -1,8 +1,8 @@
 """Prompt template system for dynamic column substitution."""
 
-import re
 import logging
-from typing import Dict, Any, List, Optional
+import re
+from typing import Any, Dict, List
 
 logger = logging.getLogger(__name__)
 
@@ -14,12 +14,13 @@ class PromptTemplate:
     COLUMN_PATTERN = re.compile(r"\{(?:column|col):([\w-]+)\}")
 
     def __init__(self, template: str):
-        """
-        Initialize with a prompt template.
+        """Initialize with a prompt template.
 
         Args:
+        ----
             template: Prompt template string, e.g.
                      "describe this image. tags: {column:user_tags}"
+
         """
         self.template = template
         self.required_columns = self._extract_columns()
@@ -30,14 +31,16 @@ class PromptTemplate:
         return list(set(matches))  # Remove duplicates
 
     def format(self, item_data: Dict[str, Any]) -> str:
-        """
-        Format the template with actual column values.
+        """Format the template with actual column values.
 
         Args:
+        ----
             item_data: Dictionary containing column values from dataset
 
         Returns:
+        -------
             Formatted prompt string
+
         """
         prompt = self.template
 
@@ -64,11 +67,12 @@ class PromptTemplate:
         return prompt.strip()
 
     def validate_columns(self, available_columns: List[str]) -> List[str]:
-        """
-        Validate that required columns are available.
+        """Validate that required columns are available.
 
-        Returns:
+        Returns
+        -------
             List of missing column names
+
         """
         missing = []
         for col in self.required_columns:
@@ -81,11 +85,12 @@ class PromptTemplateManager:
     """Manages multiple prompt templates."""
 
     def __init__(self, prompts: List[str]):
-        """
-        Initialize with list of prompt strings (which may contain templates).
+        """Initialize with list of prompt strings (which may contain templates).
 
         Args:
+        ----
             prompts: List of prompt strings
+
         """
         self.templates = [PromptTemplate(p) for p in prompts]
         self._all_required_columns = None
@@ -101,14 +106,16 @@ class PromptTemplateManager:
         return self._all_required_columns
 
     def format_all(self, item_data: Dict[str, Any]) -> List[str]:
-        """
-        Format all templates with item data.
+        """Format all templates with item data.
 
         Args:
+        ----
             item_data: Dictionary containing column values
 
         Returns:
+        -------
             List of formatted prompts
+
         """
         formatted = []
         for template in self.templates:
@@ -123,11 +130,12 @@ class PromptTemplateManager:
         return formatted
 
     def validate_all(self, available_columns: List[str]) -> Dict[str, List[str]]:
-        """
-        Validate all templates against available columns.
+        """Validate all templates against available columns.
 
-        Returns:
+        Returns
+        -------
             Dict mapping template string to list of missing columns
+
         """
         issues = {}
         for template in self.templates:

caption_flow/utils/vllm_config.py

@@ -1,8 +1,8 @@
 """vLLM configuration management utilities."""
 
 import logging
-from typing import Dict, Any, Optional, Tuple, List
 from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
@@ -140,11 +140,12 @@ class VLLMConfigManager:
     def update_runtime_config(
         self, vllm_instance, old_config: Dict[str, Any], new_config: Dict[str, Any]
     ) -> Tuple[bool, Optional[Any]]:
-        """
-        Update vLLM configuration at runtime without reload.
+        """Update vLLM configuration at runtime without reload.
 
-        Returns:
+        Returns
+        -------
             Tuple of (success, new_sampling_params)
+
         """
         change = self.analyze_config_change(old_config, new_config)
 

caption_flow/viewer.py

@@ -1,24 +1,16 @@
 """TUI viewer for browsing CaptionFlow datasets with image preview using Urwid."""
 
-import asyncio
-import io
-import json
 import logging
 import os
-import sys
-from datetime import datetime
+import tempfile
 from pathlib import Path
-from typing import Dict, Any, List, Optional, Tuple
 from urllib.parse import urlparse
-import tempfile
 
-import aiohttp
 import pandas as pd
-import pyarrow.parquet as pq
 import urwid
 
 try:
-    from term_image.image import from_file, from_url, BaseImage
+    from term_image.image import BaseImage, from_file, from_url
     from term_image.widget import UrwidImage, UrwidImageScreen
 
     TERM_IMAGE_AVAILABLE = True
@@ -232,7 +224,7 @@ class DatasetViewer:
                 filename = filename[:17] + "..."
 
             # Count captions
-            caption_count = self._count_captions(item)
+            self._count_captions(item)
 
             text = f"{idx:3d}. {filename}"
             if idx == self.current_item_idx:
@@ -411,8 +403,8 @@ class DatasetViewer:
 
         try:
             # Download image
-            import urllib.request
             import urllib.error
+            import urllib.request
 
             # Create request with user agent to avoid 403 errors
             request = urllib.request.Request(

caption_flow/workers/base.py

@@ -4,13 +4,12 @@ import asyncio
 import json
 import logging
 import ssl
-import time
 from abc import ABC, abstractmethod
-from typing import Dict, Any, Optional
 from threading import Event
+from typing import Any, Dict, Optional
 
 import websockets
-from websockets.client import WebSocketClientProtocol
+from websockets.asyncio.client import ClientConnection
 
 logger = logging.getLogger(__name__)
 
@@ -18,6 +17,8 @@ logger = logging.getLogger(__name__)
 class BaseWorker(ABC):
     """Base class for all WebSocket-based workers with common connection logic."""
 
+    gpu_id: Optional[int] = None
+
     def __init__(self, config: Dict[str, Any]):
         self.config = config
         self.server_url = config["server"]
@@ -29,7 +30,7 @@ class BaseWorker(ABC):
 
         # State
         self.worker_id: Optional[str] = None
-        self.websocket: Optional[WebSocketClientProtocol] = None
+        self.websocket: Optional[ClientConnection] = None
         self.running = False
         self.connected = Event()
         self.main_loop: Optional[asyncio.AbstractEventLoop] = None