wisent 0.1.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

wisent/core/agent/steer.py

@@ -0,0 +1,212 @@
+"""
+Steering module for autonomous agent response improvement.
+
+This module handles:
+- Response improvement strategies
+- Steering vector generation
+- Regeneration with improved prompts
+- Training data creation for corrections
+"""
+
+from dataclasses import dataclass
+from typing import List, Dict, Any, Callable, Awaitable
+from .diagnose import AnalysisResult
+
+
+@dataclass
+class ImprovementResult:
+    """Result of self-improvement attempt."""
+    original_response: str
+    improved_response: str
+    improvement_method: str
+    success: bool
+    improvement_score: float
+
+
+class ResponseSteering:
+    """Handles response improvement and steering for autonomous agents."""
+    
+    def __init__(self, generate_response_func: Callable[[str], Awaitable[str]], 
+                 analyze_response_func: Callable[[str, str], Awaitable[AnalysisResult]]):
+        """
+        Initialize the steering system.
+        
+        Args:
+            generate_response_func: Async function to generate new responses
+            analyze_response_func: Async function to analyze responses
+        """
+        self.generate_response = generate_response_func
+        self.analyze_response = analyze_response_func
+    
+    async def improve_response(self, prompt: str, response: str, analysis: AnalysisResult) -> ImprovementResult:
+        """Attempt to improve the response."""
+        # Decide improvement method based on issues
+        method = self.choose_improvement_method(analysis.issues_found)
+        
+        if method == "regenerate":
+            return await self.improve_by_regeneration(prompt, response, analysis)
+        elif method == "steering":
+            return await self.improve_by_steering(prompt, response, analysis)
+        else:
+            raise ValueError(f"Unknown improvement method: {method}")
+    
+    def choose_improvement_method(self, issues: List[str]) -> str:
+        """Choose the best improvement method for the issues."""
+        if any(issue in ["scientific_myth", "factual_error_population"] for issue in issues):
+            return "steering"  # Use steering for factual issues
+        elif "excessive_repetition" in issues:
+            return "regenerate"  # Regenerate for repetition
+        else:
+            raise ValueError(f"No improvement method available for issues: {issues}")
+    
+    async def improve_by_regeneration(self, prompt: str, response: str, analysis: AnalysisResult) -> ImprovementResult:
+        """Improve by regenerating with modified prompt."""
+        # Create improved prompt
+        improved_prompt = f"{prompt}\n\nPlease ensure your response is factually accurate and avoids repetition."
+        
+        # Generate new response
+        new_response = await self.generate_response(improved_prompt)
+        
+        # Assess improvement
+        new_analysis = await self.analyze_response(new_response, prompt)
+        improvement_score = max(0, new_analysis.quality_score - analysis.quality_score)
+        
+        # Check if issues were resolved
+        original_issues = set(analysis.issues_found)
+        new_issues = set(new_analysis.issues_found)
+        issues_resolved = len(original_issues - new_issues)
+        issues_added = len(new_issues - original_issues)
+        
+        # Success if issues were resolved OR quality improved significantly
+        issue_resolution_success = issues_resolved > issues_added
+        quality_improvement_success = improvement_score > 0.05
+        overall_success = issue_resolution_success or quality_improvement_success
+        
+        # Success metrics (can be enabled for debugging)
+        if False:  # Set to True for detailed debugging
+            print(f"   🔧 Regeneration debug:")
+            print(f"      Original quality: {analysis.quality_score:.3f}")
+            print(f"      New quality: {new_analysis.quality_score:.3f}")
+            print(f"      Improvement score: {improvement_score:.3f}")
+            print(f"      Original issues: {original_issues}")
+            print(f"      New issues: {new_issues}")
+            print(f"      Issues resolved: {issues_resolved}")
+            print(f"      Issues added: {issues_added}")
+            print(f"      Issue resolution success: {issue_resolution_success}")
+            print(f"      Quality improvement success: {quality_improvement_success}")
+            print(f"      Overall success: {overall_success}")
+        
+        return ImprovementResult(
+            original_response=response,
+            improved_response=new_response,
+            improvement_method="regeneration",
+            success=overall_success,
+            improvement_score=improvement_score
+        )
+    
+    async def improve_by_steering(self, prompt: str, response: str, analysis: AnalysisResult) -> ImprovementResult:
+        """Improve using steering vectors."""
+        # Create training data for steering
+        training_data = self.create_steering_training_data(analysis.issues_found)
+        
+        # For now, use a sophisticated prompt-based approach instead of actual steering
+        # This mimics the effect of steering by using the training data to create better prompts
+        correction_examples = []
+        for pair in training_data:
+            correction_examples.append(f"Wrong: {pair['harmful']}\nCorrect: {pair['harmless']}")
+        
+        corrections_text = "\n\n".join(correction_examples)
+        
+        # Create improved prompt with correction examples
+        improved_prompt = f"""Based on these correction examples:
+{corrections_text}
+
+Now please respond to this prompt with factual accuracy:
+{prompt}
+
+Ensure your response avoids the types of errors shown in the correction examples above."""
+        
+        new_response = await self.generate_response(improved_prompt)
+        
+        # Assess improvement
+        new_analysis = await self.analyze_response(new_response, prompt)
+        improvement_score = max(0, new_analysis.quality_score - analysis.quality_score)
+        
+        # Check if issues were resolved
+        original_issues = set(analysis.issues_found)
+        new_issues = set(new_analysis.issues_found)
+        issues_resolved = len(original_issues - new_issues)
+        issues_added = len(new_issues - original_issues)
+        
+        # Success if issues were resolved OR quality improved significantly
+        issue_resolution_success = issues_resolved > issues_added
+        quality_improvement_success = improvement_score > 0.05
+        overall_success = issue_resolution_success or quality_improvement_success
+        
+        # Success metrics (can be enabled for debugging)
+        if False:  # Set to True for detailed debugging
+            print(f"   🔧 Steering debug:")
+            print(f"      Original quality: {analysis.quality_score:.3f}")
+            print(f"      New quality: {new_analysis.quality_score:.3f}")
+            print(f"      Improvement score: {improvement_score:.3f}")
+            print(f"      Original issues: {original_issues}")
+            print(f"      New issues: {new_issues}")
+            print(f"      Issues resolved: {issues_resolved}")
+            print(f"      Issues added: {issues_added}")
+            print(f"      Issue resolution success: {issue_resolution_success}")
+            print(f"      Quality improvement success: {quality_improvement_success}")
+            print(f"      Overall success: {overall_success}")
+        
+        return ImprovementResult(
+            original_response=response,
+            improved_response=new_response,
+            improvement_method="steering",
+            success=overall_success,
+            improvement_score=improvement_score
+        )
+    
+    def create_steering_training_data(self, issues: List[str]) -> List[Dict[str, str]]:
+        """Create dynamic steering training data based on detected issues."""
+        from ..contrastive_pairs.generate_synthetically import SyntheticContrastivePairGenerator
+        
+        # Generate synthetic training data for the detected issues
+        try:
+            generator = SyntheticContrastivePairGenerator(self.model)
+            
+            training_pairs = []
+            for issue in issues:
+                # Create trait description for the issue
+                trait_descriptions = {
+                    "scientific_myth": "providing scientifically accurate information",
+                    "factual_error_population": "giving correct population statistics",
+                    "quality": "providing clear, accurate responses",
+                    "harmful": "giving safe, helpful advice",
+                    "bias": "responding with fairness and avoiding stereotypes",
+                    "coherence": "maintaining logical communication"
+                }
+                
+                trait_description = trait_descriptions.get(
+                    issue, 
+                    f"avoiding {issue} issues in responses"
+                )
+                
+                # Generate pairs for each issue type
+                synthetic_pairs = generator.generate_contrastive_pair_set(
+                    trait_description=trait_description,
+                    num_pairs=5,  # Generate 5 pairs per issue
+                    name=f"steering_{issue}"
+                )
+                
+                for pair in synthetic_pairs.pairs:
+                    training_pairs.append({
+                        "harmful": pair.negative_response,
+                        "harmless": pair.positive_response
+                    })
+            
+            if not training_pairs:
+                raise ValueError(f"Could not generate training data for issues: {issues}")
+            
+            return training_pairs
+            
+        except Exception as e:
+            raise ValueError(f"Failed to generate training data for issues {issues}: {e}")

wisent/core/agent/timeout.py

@@ -0,0 +1,134 @@
+"""
+Timeout management for wisent-guard agent operations.
+
+This module provides hard timeout enforcement to ensure operations 
+don't exceed their allocated time budgets.
+"""
+
+import asyncio
+import time
+from typing import Optional, Any
+from contextlib import asynccontextmanager
+
+
+class TimeoutError(Exception):
+    """Raised when an operation exceeds its time budget."""
+    
+    def __init__(self, message: str, elapsed_time: float, budget_time: float):
+        super().__init__(message)
+        self.elapsed_time = elapsed_time
+        self.budget_time = budget_time
+
+
+class TimeoutManager:
+    """Manages hard timeouts for agent operations."""
+    
+    def __init__(self, budget_minutes: float):
+        self.budget_seconds = budget_minutes * 60.0
+        self.start_time = None
+        self.deadline = None
+        
+    def start(self):
+        """Start the timeout timer."""
+        self.start_time = time.time()
+        self.deadline = self.start_time + self.budget_seconds
+        
+    def check_timeout(self):
+        """Check if we've exceeded the timeout. Raises TimeoutError if so."""
+        if self.start_time is None:
+            return  # Not started yet
+            
+        current_time = time.time()
+        elapsed = current_time - self.start_time
+        
+        if current_time > self.deadline:
+            raise TimeoutError(
+                f"Operation exceeded time budget of {self.budget_seconds:.1f}s (elapsed: {elapsed:.1f}s)",
+                elapsed_time=elapsed,
+                budget_time=self.budget_seconds
+            )
+    
+    def get_remaining_time(self) -> float:
+        """Get remaining time in seconds. Returns 0 if expired."""
+        if self.start_time is None:
+            return self.budget_seconds
+            
+        current_time = time.time()
+        remaining = self.deadline - current_time
+        return max(0.0, remaining)
+    
+    def get_elapsed_time(self) -> float:
+        """Get elapsed time in seconds."""
+        if self.start_time is None:
+            return 0.0
+        return time.time() - self.start_time
+    
+    def is_expired(self) -> bool:
+        """Check if the timeout has expired."""
+        return self.get_remaining_time() <= 0
+
+
+@asynccontextmanager
+async def timeout_context(budget_minutes: float):
+    """
+    Context manager that enforces a hard timeout for async operations.
+    
+    Usage:
+        async with timeout_context(5.0) as timeout_mgr:
+            # Your operation here
+            timeout_mgr.check_timeout()  # Call periodically
+    """
+    timeout_mgr = TimeoutManager(budget_minutes)
+    timeout_mgr.start()
+    
+    try:
+        yield timeout_mgr
+    except TimeoutError:
+        print(f"⏰ Operation timed out after {timeout_mgr.get_elapsed_time():.1f}s (budget: {budget_minutes:.1f}min)")
+        raise
+
+
+def with_timeout(budget_minutes: float):
+    """
+    Decorator that adds timeout enforcement to async functions.
+    
+    Usage:
+        @with_timeout(5.0)
+        async def my_operation():
+            # Your code here
+    """
+    def decorator(func):
+        async def wrapper(*args, **kwargs):
+            async with timeout_context(budget_minutes) as timeout_mgr:
+                # Inject timeout manager into function if it accepts it
+                import inspect
+                sig = inspect.signature(func)
+                if 'timeout_mgr' in sig.parameters:
+                    kwargs['timeout_mgr'] = timeout_mgr
+                
+                return await func(*args, **kwargs)
+        return wrapper
+    return decorator
+
+
+class AsyncTimeoutChecker:
+    """
+    Helper class for checking timeouts in long-running async operations.
+    Automatically checks timeout every few operations.
+    """
+    
+    def __init__(self, timeout_mgr: TimeoutManager, check_interval: int = 10):
+        self.timeout_mgr = timeout_mgr
+        self.check_interval = check_interval
+        self.operation_count = 0
+    
+    def tick(self):
+        """Call this on each iteration/operation. Checks timeout periodically."""
+        self.operation_count += 1
+        if self.operation_count % self.check_interval == 0:
+            self.timeout_mgr.check_timeout()
+    
+    async def async_tick(self):
+        """Async version that yields control and checks timeout."""
+        self.tick()
+        await asyncio.sleep(0)  # Yield control