src/models/qwen_client.py

#!/usr/bin/env python3
"""
HuggingFace Qwen 2.5 Model Client
Handles inference for router, main, and complex models with cost tracking
"""

import os
import time
import logging
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from enum import Enum

from huggingface_hub import InferenceClient
from langchain_huggingface import HuggingFaceEndpoint
from langchain_core.language_models.llms import LLM

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ModelTier(Enum):
    """Model complexity tiers for cost optimization"""
    ROUTER = "router"     # 3B - Fast, cheap routing decisions
    MAIN = "main"         # 14B - Balanced performance  
    COMPLEX = "complex"   # 32B - Best performance for hard tasks

@dataclass
class ModelConfig:
    """Configuration for each Qwen model"""
    name: str
    tier: ModelTier
    max_tokens: int
    temperature: float
    cost_per_token: float  # Estimated cost per token
    timeout: int

@dataclass
class InferenceResult:
    """Result of model inference with metadata"""
    response: str
    model_used: str
    tokens_used: int
    cost_estimate: float
    response_time: float
    success: bool
    error: Optional[str] = None

class QwenClient:
    """HuggingFace client for Qwen 2.5 model family"""
    
    def __init__(self, hf_token: Optional[str] = None):
        """Initialize the Qwen client with HuggingFace token"""
        self.hf_token = hf_token or os.getenv("HUGGINGFACE_TOKEN")
        if not self.hf_token:
            logger.warning("No HuggingFace token provided. API access may be limited.")
            
        # Define model configurations - Updated with best available models
        self.models = {
            ModelTier.ROUTER: ModelConfig(
                name="Qwen/Qwen2.5-7B-Instruct",  # Fast router for classification
                tier=ModelTier.ROUTER,
                max_tokens=512,
                temperature=0.1,
                cost_per_token=0.0003,  # 7B model
                timeout=15
            ),
            ModelTier.MAIN: ModelConfig(
                name="Qwen/Qwen2.5-32B-Instruct",  # 4.5x more powerful for main tasks
                tier=ModelTier.MAIN,
                max_tokens=1024,
                temperature=0.1,
                cost_per_token=0.0008,  # Higher cost for 32B
                timeout=25
            ),
            ModelTier.COMPLEX: ModelConfig(
                name="Qwen/Qwen2.5-72B-Instruct",  # 10x more powerful for complex reasoning!
                tier=ModelTier.COMPLEX,
                max_tokens=2048,
                temperature=0.1,
                cost_per_token=0.0015,  # Premium for 72B model
                timeout=35
            )
        }
        
        # Initialize clients
        self.inference_clients = {}
        self.langchain_clients = {}
        self._initialize_clients()
        
        # Cost tracking
        self.total_cost = 0.0
        self.request_count = 0
        self.budget_limit = 0.10  # $0.10 total budget
        
    def _initialize_clients(self):
        """Initialize HuggingFace clients for each model"""
        for tier, config in self.models.items():
            try:
                # HuggingFace InferenceClient for direct API calls
                self.inference_clients[tier] = InferenceClient(
                    model=config.name,
                    token=self.hf_token
                )
                
                # LangChain wrapper for integration
                self.langchain_clients[tier] = HuggingFaceEndpoint(
                    repo_id=config.name,
                    max_new_tokens=config.max_tokens,
                    temperature=config.temperature,
                    huggingfacehub_api_token=self.hf_token,
                    timeout=config.timeout
                )
                
                logger.info(f"✅ Initialized {tier.value} model: {config.name}")
                
            except Exception as e:
                logger.error(f"❌ Failed to initialize {tier.value} model: {e}")
                self.inference_clients[tier] = None
                self.langchain_clients[tier] = None
    
    def get_model_status(self) -> Dict[str, bool]:
        """Check which models are available"""
        status = {}
        for tier in ModelTier:
            status[tier.value] = (
                self.inference_clients.get(tier) is not None and 
                self.langchain_clients.get(tier) is not None
            )
        return status
    
    def select_model_tier(self, complexity: str = "medium", budget_conscious: bool = True, question_text: str = "") -> ModelTier:
        """Smart model selection based on task complexity, budget, and question analysis"""
        
        # Check budget constraints
        budget_used_percent = (self.total_cost / self.budget_limit) * 100
        
        if budget_conscious and budget_used_percent > 80:
            logger.warning(f"Budget critical ({budget_used_percent:.1f}% used), forcing router model")
            return ModelTier.ROUTER
        elif budget_conscious and budget_used_percent > 60:
            logger.warning(f"Budget warning ({budget_used_percent:.1f}% used), limiting complex model usage")
            complexity = "simple" if complexity == "complex" else complexity
            
        # Enhanced complexity analysis based on question content
        if question_text:
            question_lower = question_text.lower()
            
            # Indicators for complex reasoning (use 72B model)
            complex_indicators = [
                "analyze", "explain why", "reasoning", "logic", "complex", "difficult",
                "multi-step", "calculate and explain", "compare and contrast",
                "what is the relationship", "how does", "why is", "prove that",
                "step by step", "detailed analysis", "comprehensive"
            ]
            
            # Indicators for simple tasks (use 7B model)  
            simple_indicators = [
                "what is", "who is", "when", "where", "simple", "quick",
                "yes or no", "true or false", "list", "name", "find"
            ]
            
            # Math and coding indicators (use 32B model - good balance)
            math_indicators = [
                "calculate", "compute", "solve", "equation", "formula", "math",
                "number", "total", "sum", "average", "percentage", "code", "program"
            ]
            
            # File processing indicators (use 32B+ models)
            file_indicators = [
                "image", "picture", "photo", "audio", "sound", "video", "file",
                "document", "excel", "csv", "data", "chart", "graph"
            ]
            
            # Count indicators
            complex_score = sum(1 for indicator in complex_indicators if indicator in question_lower)
            simple_score = sum(1 for indicator in simple_indicators if indicator in question_lower)
            math_score = sum(1 for indicator in math_indicators if indicator in question_lower)
            file_score = sum(1 for indicator in file_indicators if indicator in question_lower)
            
            # Auto-detect complexity based on content
            if complex_score >= 2 or len(question_text) > 200:
                complexity = "complex"
            elif file_score >= 1 or math_score >= 2:
                complexity = "medium"
            elif simple_score >= 2 and complex_score == 0:
                complexity = "simple"
                
        # Select based on complexity with budget awareness
        if complexity == "complex" and budget_used_percent < 70:
            selected_tier = ModelTier.COMPLEX
        elif complexity == "simple" or budget_used_percent > 75:
            selected_tier = ModelTier.ROUTER
        else:
            selected_tier = ModelTier.MAIN
            
        # Fallback if selected model unavailable
        if not self.inference_clients.get(selected_tier):
            logger.warning(f"Selected model {selected_tier.value} unavailable, falling back")
            for fallback in [ModelTier.MAIN, ModelTier.ROUTER, ModelTier.COMPLEX]:
                if self.inference_clients.get(fallback):
                    selected_tier = fallback
                    break
            else:
                raise RuntimeError("No models available")
        
        # Log selection reasoning
        logger.info(f"Selected {selected_tier.value} model (complexity: {complexity}, budget: {budget_used_percent:.1f}%)")
        return selected_tier
    
    async def generate_async(self, 
                           prompt: str, 
                           tier: Optional[ModelTier] = None,
                           max_tokens: Optional[int] = None) -> InferenceResult:
        """Async text generation with the specified model tier"""
        
        if tier is None:
            tier = self.select_model_tier()
            
        config = self.models[tier]
        client = self.inference_clients.get(tier)
        
        if not client:
            return InferenceResult(
                response="",
                model_used=config.name,
                tokens_used=0,
                cost_estimate=0.0,
                response_time=0.0,
                success=False,
                error=f"Model {tier.value} not available"
            )
        
        start_time = time.time()
        
        try:
            # Use specified max_tokens or model default
            tokens = max_tokens or config.max_tokens
            
            # Use chat completion API for conversational models
            messages = [{"role": "user", "content": prompt}]
            
            response = client.chat_completion(
                messages=messages,
                model=config.name,
                max_tokens=tokens,
                temperature=config.temperature
            )
            
            response_time = time.time() - start_time
            
            # Extract response from chat completion
            if response and response.choices:
                response_text = response.choices[0].message.content
            else:
                raise ValueError("No response received from model")
            
            # Estimate tokens used (rough approximation)
            estimated_tokens = len(prompt.split()) + len(response_text.split())
            cost_estimate = estimated_tokens * config.cost_per_token
            
            # Update tracking
            self.total_cost += cost_estimate
            self.request_count += 1
            
            logger.info(f"✅ Generated response using {tier.value} model in {response_time:.2f}s")
            
            return InferenceResult(
                response=response_text,
                model_used=config.name,
                tokens_used=estimated_tokens,
                cost_estimate=cost_estimate,
                response_time=response_time,
                success=True
            )
            
        except Exception as e:
            response_time = time.time() - start_time
            logger.error(f"❌ Generation failed with {tier.value} model: {e}")
            
            return InferenceResult(
                response="",
                model_used=config.name, 
                tokens_used=0,
                cost_estimate=0.0,
                response_time=response_time,
                success=False,
                error=str(e)
            )
    
    def generate(self, 
                prompt: str, 
                tier: Optional[ModelTier] = None,
                max_tokens: Optional[int] = None) -> InferenceResult:
        """Synchronous text generation (wrapper for async)"""
        import asyncio
        
        # Create event loop if needed
        try:
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
        return loop.run_until_complete(
            self.generate_async(prompt, tier, max_tokens)
        )
    
    def get_langchain_llm(self, tier: ModelTier) -> Optional[LLM]:
        """Get LangChain LLM instance for agent integration"""
        return self.langchain_clients.get(tier)
    
    def get_usage_stats(self) -> Dict[str, Any]:
        """Get current usage and cost statistics"""
        return {
            "total_cost": self.total_cost,
            "request_count": self.request_count,
            "budget_limit": self.budget_limit,
            "budget_remaining": self.budget_limit - self.total_cost,
            "budget_used_percent": (self.total_cost / self.budget_limit) * 100,
            "average_cost_per_request": self.total_cost / max(self.request_count, 1),
            "models_available": self.get_model_status()
        }
    
    def reset_usage_tracking(self):
        """Reset usage statistics (for testing/development)"""
        self.total_cost = 0.0
        self.request_count = 0
        logger.info("Usage tracking reset")

# Test functions
def test_model_connection(client: QwenClient, tier: ModelTier):
    """Test connection to a specific model tier"""
    test_prompt = "Hello! Please respond with 'Connection successful' if you can read this."
    
    logger.info(f"Testing {tier.value} model...")
    result = client.generate(test_prompt, tier=tier, max_tokens=50)
    
    if result.success:
        logger.info(f"✅ {tier.value} model test successful: {result.response[:50]}...")
        logger.info(f"   Response time: {result.response_time:.2f}s")
        logger.info(f"   Cost estimate: ${result.cost_estimate:.6f}")
    else:
        logger.error(f"❌ {tier.value} model test failed: {result.error}")
    
    return result.success

def test_all_models():
    """Test all available models"""
    logger.info("🧪 Testing all Qwen models...")
    
    client = QwenClient()
    
    results = {}
    for tier in ModelTier:
        results[tier] = test_model_connection(client, tier)
    
    logger.info("📊 Test Results Summary:")
    for tier, success in results.items():
        status = "✅ PASS" if success else "❌ FAIL"
        logger.info(f"   {tier.value:8}: {status}")
    
    logger.info("💰 Usage Statistics:")
    stats = client.get_usage_stats()
    for key, value in stats.items():
        if key != "models_available":
            logger.info(f"   {key}: {value}")
    
    return results

if __name__ == "__main__":
    # Load environment variables for testing
    from dotenv import load_dotenv
    load_dotenv()
    
    # Run tests when script executed directly
    test_all_models()