a

step5_assign_colors.py

@@ -0,0 +1,250 @@
+#!/usr/bin/env python3
+"""
+Step 5: Assign colors to speakers based on their characteristics.
+
+Input: Speaker files in "_speakers/" folder
+Output: _colors.json with speaker-color mappings
+
+Output format:
+{
+  "Malabar": "golden",
+  "Moon": "silver",
+  "Earth": "green",
+  ...
+}
+
+Usage:
+    uv run step5_assign_colors.py
+
+Environment Variables:
+    OPENAI_API_KEY - Required
+    OPENAI_BASE_URL - Optional (for Kimi/GLM APIs)
+    LLM_MODEL - Optional (e.g., "glm-4.5-air")
+"""
+
+import os
+import re
+import sys
+import json
+from pathlib import Path
+from typing import List, Dict, Tuple, Optional, Set
+from openai import OpenAI
+
+# ============== Configuration ==============
+
+INPUT_DIR = Path("_speakers")
+OUTPUT_FILE = Path("_colors.json")
+
+# Fixed color assignments
+FIXED_COLORS = {
+    "Malabar": "#FFD700"  # Gold
+}
+
+# Default configurations for different providers
+DEFAULT_CONFIGS = {
+    "openai": {
+        "base_url": None,
+        "model": "gpt-4o-mini"
+    },
+    "moonshot": {
+        "base_url": "https://api.moonshot.cn/v1",
+        "model": "kimi-latest"
+    },
+    "bigmodel": {  # Zhipu AI (GLM)
+        "base_url": "https://open.bigmodel.cn/api/paas/v4",
+        "model": "glm-4.5-air"
+    }
+}
+
+
+def get_llm_config() -> Tuple[str, str]:
+    """Get LLM configuration from environment."""
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        raise ValueError("OPENAI_API_KEY environment variable is required")
+    
+    base_url = os.getenv("OPENAI_BASE_URL")
+    model = os.getenv("LLM_MODEL")
+    
+    if base_url:
+        if model:
+            return base_url, model
+        if "bigmodel" in base_url:
+            return base_url, DEFAULT_CONFIGS["bigmodel"]["model"]
+        elif "moonshot" in base_url or "kimi" in base_url:
+            return base_url, DEFAULT_CONFIGS["moonshot"]["model"]
+        else:
+            return base_url, DEFAULT_CONFIGS["openai"]["model"]
+    else:
+        return None, model or DEFAULT_CONFIGS["openai"]["model"]
+
+
+def collect_speakers(input_dir: Path) -> Set[str]:
+    """Collect all unique speakers from speaker files."""
+    speakers = set()
+    
+    for file_path in input_dir.glob("*_speakers.txt"):
+        with open(file_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                # Parse line: [timestamp](Speaker) text
+                match = re.match(r'^\[\d{2}:\d{2}\]\(([^)]+)\)', line)
+                if match:
+                    speakers.add(match.group(1))
+    
+    return speakers
+
+
+def assign_colors(speakers: Set[str], client: OpenAI, model: str) -> Dict[str, str]:
+    """Assign colors to speakers using LLM."""
+    # Start with fixed colors
+    color_mapping = FIXED_COLORS.copy()
+    
+    # Filter out speakers that already have fixed colors
+    remaining_speakers = [s for s in speakers if s not in color_mapping]
+    
+    if not remaining_speakers:
+        return color_mapping
+    
+    # Build prompt
+    speakers_list = ", ".join(remaining_speakers)
+    
+    prompt = f"""Assign CSS hex color codes to each speaker from "Little Malabar" based on their characteristics.
+
+Speakers to assign colors:
+{speakers_list}
+
+Color assignment guidelines (use hex codes like #FF0000):
+- Mars → #CD5C5C (red planet) or #FF4500
+- Earth → #228B22 (forest green) or #4169E1 (royal blue)
+- Moon → #C0C0C0 (silver) or #A9A9A9 (dark gray)
+- Sun → #FFD700 (gold) or #FFA500 (orange)
+- Jupiter → #D2691E (chocolate/orange)
+- Galaxy → #9370DB (medium purple) or #FF69B4 (hot pink)
+- Star → #FFFFFF (white) or #FFFACD (lemon chiffon)
+- Volcano → #8B0000 (dark red) or #FF4500 (orange red)
+- Kangaroo/Giraffe → #D2B48C (tan) or #F4A460 (sandy brown)
+- Song → #87CEEB (sky blue) or #DDA0DD (plum)
+
+Fixed assignment:
+- Malabar → #FFD700 (gold, already set)
+
+Reply with ONLY a JSON object mapping speaker names to hex color codes:
+{{"SpeakerName": "#RRGGBB", ...}}
+
+JSON:"""
+    
+    try:
+        response = client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": "You assign colors to characters. Reply with ONLY valid JSON."},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.3,
+            max_tokens=500
+        )
+        
+        message = response.choices[0].message
+        result = message.content or ""
+        
+        # GLM models may put response in reasoning_content
+        if not result and hasattr(message, 'reasoning_content') and message.reasoning_content:
+            result = message.reasoning_content
+        
+        # Try to parse JSON
+        json_match = re.search(r'\{[^}]+\}', result)
+        if json_match:
+            try:
+                parsed = json.loads(json_match.group())
+                for speaker, color in parsed.items():
+                    if speaker in remaining_speakers:
+                        color_mapping[speaker] = color
+            except json.JSONDecodeError:
+                print(f"    Warning: Could not parse JSON response")
+        
+        # Assign default hex colors for any remaining speakers
+        for speaker in remaining_speakers:
+            if speaker not in color_mapping:
+                # Simple fallback based on name (using hex codes)
+                name_lower = speaker.lower()
+                if 'mars' in name_lower:
+                    color_mapping[speaker] = "#FF4500"  # Orange red
+                elif 'earth' in name_lower:
+                    color_mapping[speaker] = "#228B22"  # Forest green
+                elif 'moon' in name_lower:
+                    color_mapping[speaker] = "#C0C0C0"  # Silver
+                elif 'sun' in name_lower:
+                    color_mapping[speaker] = "#FFD700"  # Gold
+                elif 'jupiter' in name_lower:
+                    color_mapping[speaker] = "#D2691E"  # Chocolate/orange
+                elif 'star' in name_lower:
+                    color_mapping[speaker] = "#FFFFFF"  # White
+                elif 'galaxy' in name_lower:
+                    color_mapping[speaker] = "#9370DB"  # Medium purple
+                elif 'volcano' in name_lower:
+                    color_mapping[speaker] = "#8B0000"  # Dark red
+                elif 'song' in name_lower:
+                    color_mapping[speaker] = "#87CEEB"  # Sky blue
+                else:
+                    color_mapping[speaker] = "#808080"  # Gray
+        
+        return color_mapping
+        
+    except Exception as e:
+        print(f"    Error assigning colors: {e}")
+        # Return defaults for all remaining speakers
+        for speaker in remaining_speakers:
+            color_mapping[speaker] = "gray"
+        return color_mapping
+
+
+def main():
+    # Get LLM config
+    base_url, model = get_llm_config()
+    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=base_url)
+    
+    print(f"Using model: {model}")
+    print(f"Endpoint: {base_url or 'OpenAI default'}")
+    
+    # Check input directory
+    if not INPUT_DIR.exists():
+        print(f"Error: Input directory {INPUT_DIR}/ not found")
+        sys.exit(1)
+    
+    # Collect all speakers
+    print(f"\nCollecting speakers from {INPUT_DIR}/...")
+    speakers = collect_speakers(INPUT_DIR)
+    
+    if not speakers:
+        print("Error: No speakers found")
+        sys.exit(1)
+    
+    print(f"Found {len(speakers)} unique speakers:")
+    for speaker in sorted(speakers):
+        if speaker in FIXED_COLORS:
+            print(f"  - {speaker}: {FIXED_COLORS[speaker]} (fixed)")
+        else:
+            print(f"  - {speaker}")
+    
+    # Assign colors
+    print(f"\nAssigning colors...")
+    color_mapping = assign_colors(speakers, client, model)
+    
+    print(f"\nFinal color assignments:")
+    for speaker, color in sorted(color_mapping.items()):
+        fixed = " (fixed)" if speaker in FIXED_COLORS else ""
+        print(f"  - {speaker}: {color}{fixed}")
+    
+    # Save to JSON
+    with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
+        json.dump(color_mapping, f, ensure_ascii=False, indent=2)
+    
+    print(f"\nSaved to: {OUTPUT_FILE}")
+    print(f"\nStep 5 Complete!")
+
+
+if __name__ == "__main__":
+    main()