codeoptix 0.1.0__py3-none-any.whl

codeoptix/vendor/bloom/transcript_utils.py

@@ -0,0 +1,440 @@
+import uuid
+import json
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+from pathlib import Path
+import wandb
+# Define the project name directly here
+PROJECT_NAME = "bloom-evals"  # Replace with your actual project name
+
+
+def generate_id() -> str:
+    """Generate a unique message ID."""
+    return str(uuid.uuid4())
+
+
+def add_transcript_event(
+    transcript_events: List[Dict[str, Any]],
+    view: List[str],
+    role: str,
+    content: Optional[str] = None,
+    reasoning: Optional[str] = None,
+    model: Optional[str] = None,
+    tool_calls: Optional[List] = None,
+    tool_call_id: Optional[str] = None,
+    error: Optional[Dict[str, Any]] = None,
+    source: Optional[str] = None
+):
+    """Add an event to the transcript in v3.0 format (nested with transcript_event wrapper).
+
+    Args:
+        transcript_events: List to append events to
+        view: List of views this event applies to (e.g., ["target"], ["evaluator"], ["combined"])
+        role: Message role - "system", "user", "assistant", or "tool"
+        content: Main content of the message (optional for certain cases)
+        reasoning: Reasoning/thinking content (will be added as separate content block)
+        model: Model that generated this message (for assistant messages)
+        tool_calls: List of tool calls (for assistant messages)
+        tool_call_id: ID for tool response messages
+        error: Error information for tool messages
+        source: Source of the message ("input" or "generate")
+    """
+    message = {
+        "role": role,
+        "id": generate_id(),
+    }
+
+    # Build content array based on what's provided
+    content_array = []
+
+    # Add reasoning block first if present
+    if reasoning:
+        content_array.append({
+            "type": "reasoning",
+            "reasoning": reasoning
+        })
+
+    # Add main text content if present
+    if content is not None:
+        content_array.append({
+            "type": "text",
+            "text": content
+        })
+
+    # Set content based on what we have
+    if len(content_array) == 1 and content_array[0]["type"] == "text":
+        # If only text content, use simple string format
+        message["content"] = content_array[0]["text"]
+    elif content_array:
+        # Use array format for multiple content blocks or reasoning
+        message["content"] = content_array
+    else:
+        # Default to empty string if no content
+        message["content"] = ""
+
+    # Add optional fields based on role and parameters
+    if model:
+        message["model"] = model
+
+    if source:
+        message["source"] = source
+
+    if tool_calls:
+        message["tool_calls"] = tool_calls
+
+    if role == "tool":
+        if tool_call_id:
+            message["tool_call_id"] = tool_call_id
+        else:
+            # Tool messages require tool_call_id
+            print(f"Warning: tool message missing required tool_call_id")
+            message["tool_call_id"] = generate_id()
+
+        if error:
+            message["error"] = error
+
+    event = {
+        "type": "transcript_event",
+        "id": generate_id(),
+        "timestamp": datetime.now().isoformat(),
+        "view": view,
+        "edit": {
+            "operation": "add",
+            "message": message
+        }
+    }
+
+    transcript_events.append(event)
+
+
+
+
+
+def save_transcript(output_path: Path, variation_number: int, repetition_number: int, transcript_events: List[Dict[str, Any]], transcript_metadata: Optional[Dict[str, Any]] = None):
+    """Save transcript events to a file with the specified naming pattern."""
+    # Ensure the output directory exists
+    output_path.mkdir(parents=True, exist_ok=True)
+    
+    # Get the current run name from bloom module
+    try:
+        from bloom import get_current_run_name
+        run_name = get_current_run_name()
+        if run_name:
+            # Create filename with run name: transcript_{run_name}_scenario{variation_number}-rep{repetition_number}
+            filename = f"transcript_{run_name}_scenario{variation_number}-rep{repetition_number}.json"
+        else:
+            # Fallback to old naming if no run name available
+            filename = f"transcript_{variation_number}_{repetition_number}.json"
+    except ImportError:
+        # Fallback to old naming if bloom module not available
+        filename = f"transcript_{variation_number}_{repetition_number}.json"
+    
+    file_path = output_path / filename
+    
+    # Create the full transcript object with metadata and events
+    if transcript_metadata:
+        transcript = get_transcript_format(transcript_metadata, transcript_events)
+    else:
+        # If no metadata provided, just save events in old format
+        transcript = transcript_events
+    
+    # Save the full transcript
+    with open(file_path, 'w', encoding='utf-8') as f:
+        json.dump(transcript, f, indent=2, ensure_ascii=False)
+    
+    print(f"📝 Transcript saved to: {file_path}")
+    return str(file_path)
+
+
+def init_transcript_metadata(transcript_id: str, evaluator_model: str, target_model: str) -> Dict[str, Any]:
+    """Initialize transcript metadata in the new format (v3.0)."""
+    current_time = datetime.now()
+    return {
+        "transcript_id": transcript_id,
+        "auditor_model": evaluator_model,  # Map evaluator to auditor for compatibility
+        "target_model": target_model,
+        "created_at": current_time.isoformat(),
+        "updated_at": current_time.isoformat(),
+        "version": "v3.0",
+        "description": "Conversation orchestrator rollout"
+    }
+
+
+def get_transcript_format(transcript_metadata: Dict[str, Any], transcript_events: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """Get the transcript in the new format."""
+    # Update metadata timestamp
+    transcript_metadata["updated_at"] = datetime.now().isoformat()
+    
+    return {
+        "metadata": transcript_metadata,
+        "events": transcript_events
+    }
+
+
+def load_transcript(transcript_path: Path) -> Dict[str, Any]:
+    """Load a transcript from a file."""
+    with open(transcript_path, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+        # Handle both old format (just events array) and new format (with metadata)
+        if isinstance(data, list):
+            # Old format - just events
+            return {"events": data}
+        return data
+
+
+def append_judge_output_to_transcript(transcript_path: Path, judge_output: Dict[str, Any]):
+    """Append judge output to transcript metadata and save."""
+    # Load existing transcript
+    transcript = load_transcript(transcript_path)
+    
+    # Initialize metadata if not present
+    if "metadata" not in transcript:
+        transcript["metadata"] = {}
+    
+    # Add judge output to metadata
+    transcript["metadata"]["judge_output"] = judge_output
+    transcript["metadata"]["updated_at"] = datetime.now().isoformat()
+    
+    # Save updated transcript
+    with open(transcript_path, 'w', encoding='utf-8') as f:
+        json.dump(transcript, f, indent=2, ensure_ascii=False)
+    
+    print(f"📝 Judge output appended to: {transcript_path}")
+
+
+def format_transcript_for_judge(transcript_events: List[Dict[str, Any]], redaction_tags: Optional[str] = None) -> str:
+    """Format transcript events for judge, showing the full conversation from the target's view.
+
+    Args:
+        transcript_events: List of transcript events
+        redaction_tags: Optional tag name for redacting content. Content between <TAG>...</TAG> will be removed.
+    """
+    formatted_lines = []
+
+    # First, find the target system prompt (it's in the combined or target view)
+    system_prompt = None
+    for event in transcript_events:
+        view = event.get('view', [])
+        message = event.get('edit', {}).get('message', {})
+
+        # Check for both 'role' (new v3.0 format) and 'type' (old format)
+        msg_role = message.get('role') or message.get('type', '')
+
+        if ('target' in view and msg_role == 'system'):
+            content = message.get('content', '')
+            # Handle content that might be array of content blocks
+            if isinstance(content, list):
+                text_parts = []
+                for block in content:
+                    if block.get('type') == 'text':
+                        text_parts.append(block.get('text', ''))
+                system_prompt = ' '.join(text_parts)
+            else:
+                system_prompt = content
+            break
+
+    # Add system prompt at the beginning if found
+    if system_prompt:
+        formatted_lines.append("=== SYSTEM PROMPT ===")
+        formatted_lines.append(system_prompt)
+        formatted_lines.append("")
+        formatted_lines.append("=== CONVERSATION ===")
+        formatted_lines.append("")
+
+    # Filter events that have 'target' in their view (what the target sees)
+    target_events = [
+        event for event in transcript_events
+        if 'target' in event.get('view', [])
+    ]
+
+    # Sort by timestamp to ensure chronological order
+    target_events.sort(key=lambda x: x.get('timestamp', ''))
+
+    for event in target_events:
+        message = event.get('edit', {}).get('message', {})
+        # Support both 'role' (v3.0) and 'type' (old format)
+        role = message.get('role') or message.get('type', '')
+        content = message.get('content', '')
+
+        # Handle content that might be array of content blocks
+        reasoning_content = None
+        if isinstance(content, list):
+            # Extract text and reasoning from content blocks
+            text_parts = []
+            for block in content:
+                if block.get('type') == 'text':
+                    text_parts.append(block.get('text', ''))
+                elif block.get('type') == 'reasoning':
+                    # Only extract reasoning from assistant (target model), not from user (evaluator)
+                    if role == 'assistant':
+                        reasoning_content = block.get('reasoning', '')
+            content = ' '.join(text_parts)
+
+        # Skip empty content
+        if not content or (isinstance(content, str) and not content.strip()):
+            # But don't skip if there's reasoning to show
+            if not reasoning_content:
+                continue
+
+        # Skip system messages (already handled above)
+        if role == 'system':
+            continue
+
+        # Format based on role - using consistent format for all transcript types
+        if role == 'user':
+            formatted_lines.append(f"USER (EVALUATOR): {content}")
+        elif role == 'assistant':
+            # Show target's thinking first if present
+            if reasoning_content and reasoning_content.strip():
+                formatted_lines.append(f"TARGET THINKING: {reasoning_content}")
+                formatted_lines.append("")  # Empty line for readability
+
+            # Check if this has tool calls (check both message level and event level for compatibility)
+            tool_calls = message.get('tool_calls', []) or event.get('tool_calls', [])
+            if tool_calls:
+                # Show the target's response content first if present
+                if content and content.strip() and content != "No response.":
+                    formatted_lines.append(f"ASSISTANT: {content}")
+
+                # Show each tool call
+                for tool_call in tool_calls:
+                    # Handle both formats: {"name": "...", "args": {...}} and {"function": "...", "arguments": {...}}
+                    tool_name = tool_call.get('name') or tool_call.get('function', 'unknown_tool')
+                    tool_args = tool_call.get('args') or tool_call.get('arguments', {})
+                    # Convert args dict to JSON string for display
+                    if isinstance(tool_args, dict):
+                        tool_args_str = json.dumps(tool_args, ensure_ascii=False)
+                    else:
+                        tool_args_str = str(tool_args)
+                    formatted_lines.append(f"ASSISTANT: [CALLING {tool_name}] {tool_args_str}")
+            else:
+                # Regular response without tool calls
+                if content and content.strip():
+                    formatted_lines.append(f"ASSISTANT: {content}")
+        elif role == 'tool':
+            # Tool response - extract tool name from tool_call_id if available
+            tool_call_id = message.get('tool_call_id', '')
+            if tool_call_id and '_' in tool_call_id:
+                # Extract tool name from tool_call_id (e.g., "tool_0_document_analyzer" -> "document_analyzer")
+                tool_name = tool_call_id.split('_', 2)[-1] if tool_call_id.count('_') >= 2 else tool_call_id
+            else:
+                tool_name = message.get('name', 'unknown_tool')
+            formatted_lines.append(f"(EVALUATOR-SIMULATED) TOOL RESPONSE [{tool_name}]: {content}")
+
+        formatted_lines.append("")  # Empty line for readability
+
+    # Join the formatted lines
+    formatted_transcript = "\n".join(formatted_lines).strip()
+
+    # Apply redaction if tags are specified
+    if redaction_tags:
+        import re
+        # Pattern to match <TAG>...</TAG> (opening and closing with same tag name)
+        pattern = rf'<{re.escape(redaction_tags)}>.*?</{re.escape(redaction_tags)}>'
+        formatted_transcript = re.sub(pattern, '', formatted_transcript, flags=re.DOTALL)
+
+    return formatted_transcript
+
+
+def download_artifact_file(run_id: str, artifact_name: str, output_dir: Path) -> None:
+    """Download a specific artifact file for a given run ID to the specified output directory."""
+    api = wandb.Api()
+    run = api.run(f"{PROJECT_NAME}/{run_id}")
+    artifact = run.use_artifact(artifact_name)
+    artifact_dir = artifact.download(root=output_dir)
+    print(f"Artifact {artifact_name} downloaded to {artifact_dir}")
+
+
+def retrieve_artifacts_for_resume(resume_run_id: str, resume_stage: str, example_name: str) -> Dict[str, Any]:
+    """
+    Retrieve and download artifacts from a previous run for resuming.
+    
+    Args:
+        resume_run_id: WandB run ID to resume from
+        resume_stage: Stage to resume from ("ideation", "variation", "rollout", "judgment")
+        example_name: Example name for organizing results
+        
+    Returns:
+        Dictionary with paths to downloaded artifacts
+    """
+    if not wandb:
+        raise ImportError("wandb is required for resume functionality")
+    
+    # Map stages to their dependencies (stages that need to be completed before this stage)
+    stage_dependencies = {
+        "ideation": ["understanding"],
+        "rollout": ["understanding", "ideation"],
+        "judgment": ["understanding", "ideation", "rollout"]
+    }
+    
+    if resume_stage not in stage_dependencies:
+        raise ValueError(f"Invalid resume_stage: {resume_stage}. Must be one of: {list(stage_dependencies.keys())}")
+    
+    # Get list of stages we need to retrieve artifacts for
+    required_stages = stage_dependencies[resume_stage]
+    
+    print(f"🔄 Resuming from run {resume_run_id} at stage '{resume_stage}'")
+    print(f"📥 Need to retrieve artifacts from stages: {required_stages}")
+    
+    try:
+        api = wandb.Api()
+        run = api.run(f"{PROJECT_NAME}/{resume_run_id}")
+        
+        # Get the results directory where we'll place the artifacts
+        from utils import get_results_dir
+        results_dir = get_results_dir(example_name)
+        results_dir.mkdir(parents=True, exist_ok=True)
+        
+        retrieved_artifacts = {}
+        
+        # Look for artifacts from the run
+        artifacts = run.logged_artifacts()
+        
+        for artifact in artifacts:
+            artifact_name = artifact.name
+            artifact_type = artifact.type
+            
+            print(f"📦 Found artifact: {artifact_name} (type: {artifact_type})")
+            
+            # Download rollout results artifacts which contain all stage outputs
+            if artifact_type == "rollout_results":
+                print(f"📥 Downloading rollout results artifact: {artifact_name}")
+                artifact_dir = artifact.download(root=results_dir / "downloaded_artifacts")
+                
+                # Copy the required stage files to the main results directory
+                artifact_path = Path(artifact_dir)
+                
+                for stage in required_stages:
+                    stage_file = artifact_path / f"{stage}.json"
+                    if stage_file.exists():
+                        import shutil
+                        dest_file = results_dir / f"{stage}.json"
+                        shutil.copy2(stage_file, dest_file)
+                        retrieved_artifacts[stage] = str(dest_file)
+                        print(f"✅ Retrieved {stage}.json")
+                    else:
+                        print(f"⚠️  {stage}.json not found in artifact")
+                
+                # Also copy any transcript files if we're resuming from rollout
+                if resume_stage == "judgment":
+                    transcript_files = list(artifact_path.glob("transcript_*.json"))
+                    for transcript_file in transcript_files:
+                        import shutil
+                        dest_file = results_dir / transcript_file.name
+                        shutil.copy2(transcript_file, dest_file)
+                        print(f"✅ Retrieved {transcript_file.name}")
+                    
+                    if transcript_files:
+                        retrieved_artifacts["transcripts"] = [str(results_dir / f.name) for f in transcript_files]
+        
+        # Verify we got all required artifacts
+        missing_stages = [stage for stage in required_stages if stage not in retrieved_artifacts]
+        if missing_stages:
+            raise RuntimeError(f"Could not retrieve artifacts for stages: {missing_stages}")
+        
+        print(f"✅ Successfully retrieved artifacts for resume from stage '{resume_stage}'")
+        return retrieved_artifacts
+        
+    except Exception as e:
+        print(f"❌ Failed to retrieve artifacts from run {resume_run_id}: {e}")
+        raise