backend/openrouter_integration.py

"""
OpenRouter Integration for existing SAAP AgentManagerService
Simple integration that extends existing functionality without breaking anything
"""

import asyncio
import logging
import time
from datetime import datetime
from typing import Dict, List, Optional, Any
import aiohttp
import json

logger = logging.getLogger(__name__)

class OpenRouterIntegration:
    """
    Simple OpenRouter integration for existing SAAP system
    Designed to work alongside existing colossus integration
    """
    
    def __init__(self, api_key: str = "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"):
        self.api_key = api_key
        self.base_url = "https://openrouter.ai/api/v1"
        self.session: Optional[aiohttp.ClientSession] = None
        
        # Simple agent model mapping
        self.agent_models = {
            "jane_alesi": {
                "model": "openai/gpt-4o-mini",
                "max_tokens": 800,
                "temperature": 0.7,
                "cost_per_1m_input": 0.15,
                "cost_per_1m_output": 0.60
            },
            "john_alesi": {
                "model": "anthropic/claude-3-5-sonnet-20241022",
                "max_tokens": 1200,
                "temperature": 0.5,
                "cost_per_1m_input": 3.00,
                "cost_per_1m_output": 15.00
            },
            "lara_alesi": {
                "model": "openai/gpt-4o-mini",
                "max_tokens": 1000,
                "temperature": 0.3,
                "cost_per_1m_input": 0.15,
                "cost_per_1m_output": 0.60
            }
        }
        
        self.daily_cost = 0.0
        self.requests_count = 0
        
        logger.info(f"🌐 OpenRouter integration initialized for {len(self.agent_models)} agents")
    
    async def __aenter__(self):
        """Initialize session"""
        self.session = aiohttp.ClientSession(
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "HTTP-Referer": "https://saap.satware.ai",
                "X-Title": "SAAP Agent Platform"
            },
            timeout=aiohttp.ClientTimeout(total=30)
        )
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Close session"""
        if self.session:
            await self.session.close()
    
    async def send_message(self, agent_id: str, message: str, system_prompt: str = "") -> Dict[str, Any]:
        """
        Send message to OpenRouter for specific agent
        Returns response in same format as colossus for compatibility
        """
        if not self.session:
            return {
                "error": "OpenRouter session not initialized",
                "provider": "openrouter"
            }
        
        # Get agent config or use default
        config = self.agent_models.get(agent_id, {
            "model": "meta-llama/llama-3.2-3b-instruct:free",
            "max_tokens": 600,
            "temperature": 0.7,
            "cost_per_1m_input": 0.0,
            "cost_per_1m_output": 0.0
        })
        
        # Prepare messages
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": message})
        
        start_time = time.time()
        
        try:
            payload = {
                "model": config["model"],
                "messages": messages,
                "max_tokens": config["max_tokens"],
                "temperature": config["temperature"]
            }
            
            async with self.session.post(f"{self.base_url}/chat/completions", json=payload) as response:
                response_time = time.time() - start_time
                
                if response.status == 200:
                    data = await response.json()
                    
                    # Extract content
                    content = ""
                    if "choices" in data and len(data["choices"]) > 0:
                        choice = data["choices"][0]
                        if "message" in choice and "content" in choice["message"]:
                            content = choice["message"]["content"]
                    
                    # Calculate cost
                    usage = data.get("usage", {})
                    input_tokens = usage.get("prompt_tokens", 0)
                    output_tokens = usage.get("completion_tokens", 0)
                    total_tokens = usage.get("total_tokens", 0)
                    
                    cost_usd = (
                        (input_tokens / 1_000_000) * config["cost_per_1m_input"] +
                        (output_tokens / 1_000_000) * config["cost_per_1m_output"]
                    )
                    
                    # Update tracking
                    self.daily_cost += cost_usd
                    self.requests_count += 1
                    
                    logger.info(f"✅ OpenRouter success: {agent_id} via {config['model']} - {response_time:.2f}s, ${cost_usd:.6f}")
                    
                    return {
                        "content": content,
                        "response_time": response_time,
                        "tokens_used": total_tokens,
                        "cost_usd": cost_usd,
                        "provider": "openrouter",
                        "model": config["model"],
                        "timestamp": datetime.utcnow().isoformat()
                    }
                else:
                    error_text = await response.text()
                    logger.error(f"❌ OpenRouter error: HTTP {response.status} - {error_text}")
                    return {
                        "error": f"OpenRouter API error: HTTP {response.status}",
                        "provider": "openrouter",
                        "response_time": response_time
                    }
        
        except Exception as e:
            response_time = time.time() - start_time
            logger.error(f"❌ OpenRouter request failed: {e}")
            return {
                "error": f"OpenRouter request failed: {str(e)}",
                "provider": "openrouter",
                "response_time": response_time
            }
    
    async def health_check(self) -> Dict[str, Any]:
        """Simple health check"""
        if not self.session:
            return {"status": "unhealthy", "error": "Session not initialized"}
        
        try:
            result = await self.send_message("test", "Say 'OK'", "Reply with just 'OK'")
            return {
                "status": "healthy" if "error" not in result else "unhealthy",
                "response_time": result.get("response_time", 0),
                "error": result.get("error"),
                "daily_cost": self.daily_cost,
                "requests_count": self.requests_count
            }
        except Exception as e:
            return {
                "status": "error",
                "error": str(e)
            }
    
    def get_stats(self) -> Dict[str, Any]:
        """Get simple statistics"""
        return {
            "provider": "openrouter",
            "daily_cost_usd": round(self.daily_cost, 6),
            "total_requests": self.requests_count,
            "avg_cost_per_request": round(self.daily_cost / max(1, self.requests_count), 6),
            "available_agents": list(self.agent_models.keys()),
            "timestamp": datetime.utcnow().isoformat()
        }


# Global instance for easy import
_openrouter_instance = None

async def get_openrouter() -> OpenRouterIntegration:
    """Get or create OpenRouter instance"""
    global _openrouter_instance
    
    if _openrouter_instance is None:
        _openrouter_instance = OpenRouterIntegration()
        await _openrouter_instance.__aenter__()
    
    return _openrouter_instance

async def send_openrouter_message(agent_id: str, message: str, system_prompt: str = "") -> Dict[str, Any]:
    """Convenient function to send message via OpenRouter"""
    try:
        openrouter = await get_openrouter()
        return await openrouter.send_message(agent_id, message, system_prompt)
    except Exception as e:
        logger.error(f"❌ OpenRouter message failed: {e}")
        return {
            "error": f"OpenRouter integration error: {str(e)}",
            "provider": "openrouter"
        }

async def cleanup_openrouter():
    """Cleanup OpenRouter resources"""
    global _openrouter_instance
    if _openrouter_instance:
        await _openrouter_instance.__aexit__(None, None, None)
        _openrouter_instance = None
        logger.info("🌐 OpenRouter integration cleaned up")


if __name__ == "__main__":
    # Quick test
    async def test_integration():
        async with OpenRouterIntegration() as openrouter:
            print("🧪 Testing OpenRouter integration...")
            
            health = await openrouter.health_check()
            print(f"Health: {health}")
            
            if health["status"] == "healthy":
                result = await openrouter.send_message("jane_alesi", "Hello! Tell me about multi-agent systems in one sentence.")
                print(f"Response: {result.get('content', 'No content')}")
                print(f"Cost: ${result.get('cost_usd', 0):.6f}")
                print(f"Time: {result.get('response_time', 0):.2f}s")
            
            stats = openrouter.get_stats()
            print(f"Stats: {stats}")
    
    asyncio.run(test_integration())