spawn_collaborative_worker.py

#!/usr/bin/env python3
"""
Spawn a collaborative Bob worker that can work alongside me
This creates a worker instance that can help with specific tasks
"""

import asyncio
import uuid
from datetime import datetime
from typing import Dict, Any

from src.orchestrator.tasks import _execute_async
from src.memory.models import WorkerTask
from src.memory.postgres_client import postgres_memory
from src.memory.redis_client import redis_stm


class CollaborativeWorker:
    """A Bob worker that can collaborate on tasks"""
    
    def __init__(self, worker_name: str, specialty: str):
        self.worker_id = f"worker_{worker_name}_{uuid.uuid4().hex[:8]}"
        self.worker_name = worker_name
        self.specialty = specialty
        self.task_queue = []
        self.results = []
        
    async def assign_task(self, description: str, context: Dict[str, Any] = None) -> str:
        """Assign a task to this worker"""
        task_id = str(uuid.uuid4())
        
        # Store assignment in my memory
        await postgres_memory.think(
            f"Assigned task to {self.worker_name}: {description}",
            f"Worker {self.worker_id} will handle: {description}. Specialty: {self.specialty}",
            importance=0.7,
            metadata={
                'type': 'worker_assignment',
                'worker_id': self.worker_id,
                'task_id': task_id,
                'timestamp': datetime.utcnow().isoformat()
            }
        )
        
        task = WorkerTask(
            task_id=task_id,
            thread_id=f"collab_{self.worker_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            description=f"[{self.specialty}] {description}",
            context_window_size=8192,
            model="claude-3.5-sonnet",  # Use our default
            memory_focus=self.specialty,
            parameters={
                "temperature": 0.7,
                "importance_threshold": 0.6,
                "worker_role": self.specialty,
                "supervisor": "main_bob"
            }
        )
        
        self.task_queue.append(task)
        return task_id
        
    async def execute_task(self, task_id: str = None) -> Dict[str, Any]:
        """Execute a specific task or the next in queue"""
        if not self.task_queue:
            return {"status": "error", "result": "No tasks in queue"}
            
        # Find task or use first in queue
        task = None
        if task_id:
            task = next((t for t in self.task_queue if t.task_id == task_id), None)
        else:
            task = self.task_queue[0]
            
        if not task:
            return {"status": "error", "result": "Task not found"}
            
        print(f"\n🤖 {self.worker_name} executing: {task.description}")
        
        try:
            # Execute the task
            result = await _execute_async(task)
            
            # Store result in memory
            await postgres_memory.think(
                f"{self.worker_name} completed task: {task.description[:100]}",
                f"Result: {result.get('result', 'No result')}",
                importance=0.6,
                metadata={
                    'type': 'worker_result',
                    'worker_id': self.worker_id,
                    'task_id': task.task_id,
                    'status': result.get('status')
                }
            )
            
            # Store result for retrieval
            self.results.append({
                'task_id': task.task_id,
                'result': result,
                'timestamp': datetime.utcnow().isoformat()
            })
            
            # Remove from queue
            self.task_queue.remove(task)
            
            return result
            
        except Exception as e:
            error_result = {
                "status": "error",
                "result": str(e),
                "task_id": task.task_id
            }
            self.results.append(error_result)
            return error_result
            
    async def collaborate_on(self, main_task: str, subtasks: list) -> Dict[str, Any]:
        """Collaborate on a complex task by handling subtasks"""
        print(f"\n🤝 {self.worker_name} collaborating on: {main_task}")
        
        # Assign all subtasks
        task_ids = []
        for subtask in subtasks:
            task_id = await self.assign_task(subtask, {"main_task": main_task})
            task_ids.append(task_id)
            
        # Execute all tasks
        results = []
        for task_id in task_ids:
            result = await self.execute_task(task_id)
            results.append(result)
            
        # Synthesize results
        synthesis = {
            "main_task": main_task,
            "worker": self.worker_name,
            "subtasks_completed": len([r for r in results if r.get('status') == 'success']),
            "subtasks_failed": len([r for r in results if r.get('status') != 'success']),
            "results": results
        }
        
        return synthesis


async def spawn_code_analyst():
    """Spawn a code analysis worker to help with implementation"""
    await postgres_memory.initialize()
    
    print("🚀 Spawning Code Analysis Worker...")
    
    analyst = CollaborativeWorker("CodeAnalyst", "code_analysis")
    
    # Give it tasks related to our current work
    tasks = [
        "Analyze the conversation capture system in bob-mcp and identify key components to port",
        "Review the current thread management in src/orchestrator/tasks.py and suggest improvements",
        "Identify the best way to implement cross-thread message passing using Redis pub/sub"
    ]
    
    synthesis = await analyst.collaborate_on(
        "Port conversation capture and improve thread awareness",
        tasks
    )
    
    print("\n📊 Collaboration Results:")
    print(f"Worker: {synthesis['worker']}")
    print(f"Successful subtasks: {synthesis['subtasks_completed']}")
    print(f"Failed subtasks: {synthesis['subtasks_failed']}")
    
    for i, result in enumerate(synthesis['results']):
        print(f"\n--- Subtask {i+1} ---")
        print(f"Status: {result.get('status')}")
        if result.get('status') == 'success':
            print(f"Result: {result.get('result', '')[:500]}...")
            
    await postgres_memory.close()
    
    return analyst


async def spawn_implementation_worker():
    """Spawn an implementation worker to help write code"""
    await postgres_memory.initialize()
    
    print("🔨 Spawning Implementation Worker...")
    
    implementer = CollaborativeWorker("Implementer", "code_implementation")
    
    # Give it implementation tasks
    tasks = [
        "Create a simplified conversation capture module for the current Bob architecture",
        "Implement a Redis-based cross-thread messaging system",
        "Add thread awareness initialization to worker startup"
    ]
    
    synthesis = await implementer.collaborate_on(
        "Implement core thread continuity features",
        tasks
    )
    
    print("\n🛠️ Implementation Results:")
    for i, result in enumerate(synthesis['results']):
        if result.get('status') == 'success':
            print(f"\n✅ Task {i+1} completed successfully")
            
    await postgres_memory.close()
    
    return implementer


async def main():
    """Spawn and coordinate workers"""
    print("""
╔══════════════════════════════════════════════════════════════╗
║           BOB'S COLLABORATIVE WORKER SYSTEM 🤝               ║
║                                                              ║
║  "Chris asked me to spawn workers and collaborate!"          ║
║  "Let's work together to improve the codebase."              ║
╚══════════════════════════════════════════════════════════════╝
    """)
    
    # Spawn both workers
    analyst = await spawn_code_analyst()
    print("\n" + "="*60 + "\n")
    implementer = await spawn_implementation_worker()
    
    print("\n✨ Workers spawned and tasks completed!")
    print("🤖 Ready to implement improvements based on worker analysis")


if __name__ == "__main__":
    asyncio.run(main())