reasoning_classifier.py

"""
Reasoning Classifier Module
Uses DeepSeek-R1 with Chain of Thought for AI intelligence classification
"""

import json
import asyncio
from typing import Dict, List, Optional
from datetime import datetime, timedelta
from functools import lru_cache
import hashlib
import re


class ReasoningClassifier:
    """Reasoning-based classifier using DeepSeek-R1 with Chain of Thought"""
    
    def __init__(self, cache_ttl: int = 86400):  # 24 hour cache
        self.cache_ttl = cache_ttl
        self.cache = {}
        self.system_instruction = (
            "You are an AI Research Scientist specializing in AGI, ASI, and ACI taxonomies. "
            "Your task is to analyze research paper abstracts and categorize them.\n\n"
            "CRITERIA:\n"
            "- AGI (Artificial General Intelligence): Focus on cross-domain reasoning, System 2 thinking, and 'generality.'\n"
            "- ASI (Artificial Superintelligence): Focus on recursive self-improvement, alignment at scale, and superhuman capabilities.\n"
            "- ACI (Artificial Collective Intelligence): Focus on multi-agent systems, swarm intelligence, and human-AI collaboration.\n"
            "- Narrow AI: Focus on specific, single-domain optimizations (e.g., just 'faster vision' or 'better LLM weights').\n\n"
            "OUTPUT FORMAT (JSON):\n"
            "{\n"
            '  "category": "AGI | ASI | ACI | Narrow AI",\n'
            '  "confidence_score": 0-100,\n'
            '  "analysis": "A brief technical justification of why this fits the category based on architectural depth.",\n'
            '  "aci_potential": "High/Low"\n'
            "}\n"
        )
    
    def _construct_fallback_response(self, content: str) -> Optional[Dict]:
        """
        Construct a fallback response from incomplete JSON by extracting available information
        
        Args:
            content: The content that failed JSON parsing
            
        Returns:
            Dictionary with extracted information or None if extraction fails
        """
        try:
            # Try to extract category
            category_match = re.search(r'"category"\s*:\s*"([^"]+)"', content)
            category = category_match.group(1) if category_match else "Unknown"
            
            # Try to extract confidence score
            confidence_match = re.search(r'"confidence_score"\s*:\s*(\d+)', content)
            confidence = int(confidence_match.group(1)) if confidence_match else 50
            
            # Try to extract analysis (partial)
            analysis_match = re.search(r'"analysis"\s*:\s*"([^"]*)', content)
            analysis = analysis_match.group(1) if analysis_match else "Could not extract full analysis"
            
            # Try to extract ACI potential
            aci_match = re.search(r'"aci_potential"\s*:\s*"([^"]+)"', content)
            aci_potential = aci_match.group(1) if aci_match else "Unknown"
            
            print(f"DEBUG: Fallback extraction - category: {category}, confidence: {confidence}")
            
            return {
                'category': category,
                'confidence_score': confidence,
                'analysis': analysis + " (Extracted from incomplete response)",
                'aci_potential': aci_potential,
                'model_used': 'deepseek-r1',
                'classification_timestamp': datetime.now().isoformat(),
                'fallback': True
            }
        except Exception as e:
            print(f"DEBUG: Fallback extraction failed: {e}")
            return None
    
    def _get_cache_key(self, title: str, summary: str) -> str:
        """Generate cache key from paper content"""
        content = f"{title}:{summary}"
        return hashlib.md5(content.encode()).hexdigest()
    
    def _get_from_cache(self, cache_key: str) -> Optional[Dict]:
        """Get classification result from cache if available and not expired"""
        if cache_key in self.cache:
            result, timestamp = self.cache[cache_key]
            if (datetime.now() - timestamp).seconds < self.cache_ttl:
                return result
        return None
    
    def _set_cache(self, cache_key: str, result: Dict):
        """Store classification result in cache"""
        self.cache[cache_key] = (result, datetime.now())
    
    def classify_paper(self, paper_data: Dict, use_cache: bool = True) -> Dict:
        """
        Classify a paper using reasoning-based approach
        
        Args:
            paper_data: Dictionary containing paper information (title, summary, etc.)
            use_cache: Whether to use cached results if available
            
        Returns:
            Dictionary with classification results
        """
        title = paper_data.get('title', '')
        summary = paper_data.get('summary', '')
        
        # Check cache first
        cache_key = self._get_cache_key(title, summary)
        if use_cache:
            cached_result = self._get_from_cache(cache_key)
            if cached_result:
                cached_result['cached'] = True
                return cached_result
        
        # Perform reasoning classification
        result = self._classify_with_reasoning(title, summary)
        
        # Cache the result
        if use_cache:
            self._set_cache(cache_key, result)
        
        result['cached'] = False
        return result
    
    def _classify_with_reasoning(self, title: str, summary: str) -> Dict:
        """
        Perform actual reasoning classification using DeepSeek-R1
        
        Args:
            title: Paper title
            summary: Paper summary
            
        Returns:
            Classification result dictionary
        """
        try:
            from huggingface_hub import InferenceClient
            import os
            
            # Get API key from environment variable
            api_key = os.getenv("HUGGINGFACE_API_KEY")
            
            # Debug: Check if API key is available
            print(f"DEBUG: HUGGINGFACE_API_KEY found: {bool(api_key)}")
            if not api_key:
                print("DEBUG: HUGGINGFACE_API_KEY not set - will try without authentication (may fail)")
            
            # Initialize client with API key if available
            if api_key:
                client = InferenceClient("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", token=api_key)
            else:
                client = InferenceClient("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B")
            
            user_input = f"Analyze this abstract: {title}. {summary}"
            
            response = client.chat_completion(
                messages=[
                    {"role": "system", "content": self.system_instruction},
                    {"role": "user", "content": user_input}
                ],
                max_tokens=1000,
                temperature=0.1
            )
            
            content = response.choices[0].message.content
            
            # Debug: Log raw response
            print(f"DEBUG: Raw model response length: {len(content)}")
            print(f"DEBUG: Raw response preview: {content[:200]}")
            
            # Strip CoT tags if present (DeepSeek-R1 uses specific tags)
            try:
                # Try to extract content between reasoning tags if present
                if "<|begin_of_thought|>" in content and "<|end_of_thought|>" in content:
                    # DeepSeek specific tags
                    content = content.split("<|end_of_thought|>")[-1].strip()
                elif "```json" in content:
                    # Extract JSON from code block
                    content = content.split("```json")[-1].split("```")[0].strip()
                elif "```" in content:
                    # Extract from any code block
                    content = content.split("```")[-1].split("```")[0].strip()
                # Handle Chinese/Unicode CoT tags by looking for JSON pattern
                elif "{" in content and "}" in content:
                    # Try to extract just the JSON part
                    start_idx = content.find("{")
                    end_idx = content.rfind("}") + 1
                    if start_idx != -1 and end_idx > start_idx:
                        content = content[start_idx:end_idx]
                
                print(f"DEBUG: Content after tag stripping length: {len(content)}")
                print(f"DEBUG: Content preview: {content[:200]}")
            except Exception as e:
                # If tag stripping fails, use content as-is
                print(f"DEBUG: Tag stripping failed: {e}, using raw content")
                pass
            
            # Try to parse JSON response
            try:
                result = json.loads(content)
            except json.JSONDecodeError as e:
                # Try to extract JSON from mixed content
                print(f"DEBUG: Initial JSON parse failed: {e}")
                print(f"DEBUG: Attempting to extract JSON from mixed content")
                
                # Try to find JSON-like structure
                # Improved regex pattern to handle nested structures
                json_pattern = r'\{(?:[^{}]|(?:\{[^{}]*\}))*\}'
                matches = re.findall(json_pattern, content, re.DOTALL)
                
                if matches:
                    # Try the largest match first (most likely to be complete)
                    matches.sort(key=len, reverse=True)
                    for match in matches:
                        try:
                            result = json.loads(match)
                            print(f"DEBUG: Successfully extracted JSON from mixed content")
                            print(f"DEBUG: Extracted JSON length: {len(match)}")
                            break
                        except:
                            continue
                    else:
                        # If all matches fail, try to construct a minimal valid response
                        print(f"DEBUG: All JSON matches failed, attempting fallback")
                        # Try to extract category and confidence from text
                        fallback_result = self._construct_fallback_response(content)
                        if fallback_result:
                            return fallback_result
                        raise e
                else:
                    print(f"DEBUG: No JSON patterns found in content")
                    # Try to construct a minimal valid response
                    fallback_result = self._construct_fallback_response(content)
                    if fallback_result:
                        return fallback_result
                    raise e
            
            # Add metadata
            result['model_used'] = 'deepseek-r1'
            result['classification_timestamp'] = datetime.now().isoformat()
            
            return result
            
        except json.JSONDecodeError as e:
            # JSON parsing failed, return error result
            print(f"DEBUG: Final JSON parsing error: {e}")
            return {
                'category': 'Error',
                'confidence_score': 0,
                'analysis': f'JSON parsing error: {str(e)}',
                'aci_potential': 'Unknown',
                'model_used': 'deepseek-r1',
                'classification_timestamp': datetime.now().isoformat(),
                'error': str(e)
            }
        except Exception as e:
            # API call failed, return error result
            error_msg = str(e)
            print(f"DEBUG: General exception: {error_msg}")
            if "api_key" in error_msg.lower() or "api key" in error_msg.lower():
                error_msg = "HUGGINGFACE_API_KEY not configured or invalid. Please configure it in Space Settings or use keyword mode instead."
            return {
                'category': 'Error',
                'confidence_score': 0,
                'analysis': f'API error: {error_msg}',
                'aci_potential': 'Unknown',
                'model_used': 'deepseek-r1',
                'classification_timestamp': datetime.now().isoformat(),
                'error': error_msg
            }
    
    async def classify_paper_async(self, paper_data: Dict, use_cache: bool = True) -> Dict:
        """
        Async version of classify_paper for batch processing
        
        Args:
            paper_data: Dictionary containing paper information
            use_cache: Whether to use cached results if available
            
        Returns:
            Classification result dictionary
        """
        # For now, use synchronous version with asyncio wrapper
        # In future, can implement true async with aiohttp
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.classify_paper, paper_data, use_cache)
    
    async def batch_classify_async(self, papers: List[Dict], use_cache: bool = True, 
                                   max_concurrent: int = 5) -> List[Dict]:
        """
        Classify multiple papers asynchronously with concurrency control
        
        Args:
            papers: List of paper dictionaries
            use_cache: Whether to use cached results
            max_concurrent: Maximum number of concurrent API calls
            
        Returns:
            List of classification results
        """
        semaphore = asyncio.Semaphore(max_concurrent)
        
        async def classify_with_semaphore(paper):
            async with semaphore:
                return await self.classify_paper_async(paper, use_cache)
        
        tasks = [classify_with_semaphore(paper) for paper in papers]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Handle exceptions
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                processed_results.append({
                    'category': 'Error',
                    'confidence_score': 0,
                    'analysis': f'Batch processing error: {str(result)}',
                    'aci_potential': 'Unknown',
                    'model_used': 'deepseek-r1',
                    'classification_timestamp': datetime.now().isoformat(),
                    'error': str(result)
                })
            else:
                processed_results.append(result)
        
        return processed_results
    
    def batch_classify(self, papers: List[Dict], use_cache: bool = True) -> List[Dict]:
        """
        Synchronous batch classification
        
        Args:
            papers: List of paper dictionaries
            use_cache: Whether to use cached results
            
        Returns:
            List of classification results
        """
        results = []
        for paper in papers:
            result = self.classify_paper(paper, use_cache)
            results.append(result)
        return results
    
    def clear_cache(self):
        """Clear the classification cache"""
        self.cache.clear()
    
    def get_cache_stats(self) -> Dict:
        """Get cache statistics"""
        total_entries = len(self.cache)
        valid_entries = 0
        expired_entries = 0
        
        for cache_key, (result, timestamp) in self.cache.items():
            if (datetime.now() - timestamp).seconds < self.cache_ttl:
                valid_entries += 1
            else:
                expired_entries += 1
        
        return {
            'total_entries': total_entries,
            'valid_entries': valid_entries,
            'expired_entries': expired_entries,
            'cache_ttl_hours': self.cache_ttl / 3600
        }
    
    def get_supported_categories(self) -> List[str]:
        """Get list of supported classification categories"""
        return ['AGI', 'ASI', 'ACI', 'Narrow AI', 'Not Related', 'Error']


# Test the reasoning classifier
if __name__ == "__main__":
    classifier = ReasoningClassifier()
    
    # Test with sample papers
    test_papers = [
        {
            'title': 'Neural Computers: A New Computing Paradigm',
            'summary': 'Researchers propose Neural Computers that unify computation, memory, and I/O in a single learned runtime state, potentially leading to artificial general intelligence.'
        },
        {
            'title': 'Multi-Agent Reinforcement Learning for Swarm Coordination',
            'summary': 'A novel approach to coordinating large swarms of autonomous agents using decentralized reinforcement learning and emergent collective intelligence.'
        },
        {
            'title': 'Image Classification with Deep Learning',
            'summary': 'A new approach to image classification using convolutional neural networks.'
        }
    ]
    
    print("Testing Reasoning Classifier...")
    for i, paper in enumerate(test_papers, 1):
        print(f"\nPaper {i}: {paper['title']}")
        result = classifier.classify_paper(paper)
        print(f"Category: {result['category']}")
        print(f"Confidence: {result['confidence_score']}")
        print(f"Analysis: {result['analysis']}")
        print(f"ACI Potential: {result['aci_potential']}")
        print(f"Cached: {result['cached']}")