File size: 9,609 Bytes
4343907
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
"""
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())