File size: 13,878 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
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
"""
πŸ”§ Enhanced Agent Manager Service - Database Loading Fix
Verbesserte Version mit robuster Database-Memory Integration

Behebt kritische Probleme:
- Agents werden nicht aus Database geladen beim Backend-Start  
- Memory vs Database Mismatch zwischen Services
- Neue Agents verschwinden nach Server-Restart
- "No description available" Problem
"""

import asyncio
import logging
from typing import Dict, List, Optional, Any
from datetime import datetime

from services.agent_manager import AgentManagerService
from database.connection import db_manager
from database.models import DBAgent
from models.agent import SaapAgent, AgentStatus, AgentType
from sqlalchemy import select

logger = logging.getLogger(__name__)

class EnhancedAgentManagerService(AgentManagerService):
    """
    Enhanced Agent Manager mit verbesserter Database Integration
    
    Neue Features:
    - Force Database Loading beim Startup
    - Enhanced Database-Memory Bridge  
    - Guaranteed Agent Persistence
    - Improved Error Handling
    - Better Service Integration
    """
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.database_loading_enabled = True
        self.force_database_sync = True
        
    async def initialize(self):
        """Enhanced initialization with guaranteed Database Agent Loading"""
        try:
            logger.info("πŸš€ Initializing Enhanced Agent Manager Service...")
            
            # Initialize colossus client (from parent)
            await self._initialize_colossus_client()
            
            # Enhanced Database Agent Loading - GUARANTEED
            await self.force_load_agents_from_database()
            
            self.is_initialized = True
            logger.info(f"βœ… Enhanced Agent Manager initialized: {len(self.agents)} agents loaded")
            
            # Log all loaded agents for debugging
            for agent_id, agent in self.agents.items():
                logger.info(f"πŸ“‹ Loaded: {agent.name} ({agent_id}) - {agent.status.value}")
            
        except Exception as e:
            logger.error(f"❌ Enhanced Agent Manager initialization failed: {e}")
            raise
    
    async def _initialize_colossus_client(self):
        """Initialize colossus client with error handling"""
        try:
            from api.colossus_client import ColossusClient
            logger.info("πŸ”Œ Connecting to colossus server...")
            self.colossus_client = ColossusClient()
            await self.colossus_client.__aenter__()
            await self._test_colossus_connection()
        except Exception as colossus_error:
            logger.error(f"❌ colossus connection failed: {colossus_error}")
            self.colossus_connection_status = f"failed: {str(colossus_error)}"
    
    async def force_load_agents_from_database(self):
        """
        πŸš€ ENHANCED: Force load all agents from database with guaranteed success
        
        This method GUARANTEES that all database agents are loaded into memory.
        It replaces the problematic _load_agents_from_database method.
        """
        try:
            logger.info("πŸ” Force loading agents from database...")
            
            # Step 1: Ensure database is initialized
            await self._ensure_database_initialized()
            
            # Step 2: Load all agents from database
            loaded_count = await self._load_database_agents()
            
            # Step 3: Add default agents if database is empty
            if loaded_count == 0:
                logger.info("πŸ“¦ No agents in database - adding default agents...")
                await self._add_default_agents_to_database()
                loaded_count = await self._load_database_agents()
            
            # Step 4: Verify loading success
            await self._verify_agent_loading(loaded_count)
            
            logger.info(f"βœ… Force loading successful: {len(self.agents)} agents in memory")
            
        except Exception as e:
            logger.error(f"❌ Force loading failed: {e}")
            # Fallback to default agents if database loading fails completely
            logger.info("πŸ†˜ Fallback: Loading default agents in memory-only mode...")
            await super().load_default_agents()
    
    async def _ensure_database_initialized(self):
        """Ensure database is properly initialized"""
        max_retries = 3
        retry_count = 0
        
        while retry_count < max_retries:
            try:
                if not db_manager.is_initialized:
                    logger.info(f"πŸ“Š Database not initialized - attempting initialization (retry {retry_count + 1}/{max_retries})...")
                    await db_manager.initialize()
                
                # Test database connectivity
                async with db_manager.get_async_session() as session:
                    result = await session.execute(select(DBAgent).limit(1))
                    result.scalars().first()
                
                logger.info("βœ… Database connection verified")
                return
                
            except Exception as e:
                retry_count += 1
                if retry_count >= max_retries:
                    raise Exception(f"Database initialization failed after {max_retries} retries: {e}")
                
                logger.warning(f"⚠️ Database init attempt {retry_count} failed: {e}")
                await asyncio.sleep(1)
    
    async def _load_database_agents(self) -> int:
        """Load all agents from database into memory"""
        try:
            async with db_manager.get_async_session() as session:
                result = await session.execute(select(DBAgent))
                db_agents = result.scalars().all()
                
            logger.info(f"πŸ“š Found {len(db_agents)} agents in database")
            
            # Clear existing agents and load from database
            self.agents.clear()
            loaded_count = 0
            
            for db_agent in db_agents:
                try:
                    saap_agent = db_agent.to_saap_agent()
                    self.agents[saap_agent.id] = saap_agent
                    loaded_count += 1
                    logger.debug(f"πŸ”„ Loaded: {saap_agent.name} ({saap_agent.id})")
                except Exception as conversion_error:
                    logger.warning(f"⚠️ Failed to convert agent {db_agent.id}: {conversion_error}")
            
            logger.info(f"βœ… Successfully loaded {loaded_count} agents from database")
            return loaded_count
            
        except Exception as e:
            logger.error(f"❌ Database agent loading failed: {e}")
            return 0
    
    async def _add_default_agents_to_database(self):
        """Add default Alesi agents to database"""
        try:
            from models.agent import AgentTemplates
            
            default_agents = [
                AgentTemplates.jane_alesi(),
                AgentTemplates.john_alesi(), 
                AgentTemplates.lara_alesi()
            ]
            
            async with db_manager.get_async_session() as session:
                for agent in default_agents:
                    # Check if agent already exists
                    result = await session.execute(
                        select(DBAgent).where(DBAgent.id == agent.id)
                    )
                    if result.scalars().first():
                        logger.debug(f"⚠️ Agent {agent.id} already exists in database")
                        continue
                    
                    db_agent = DBAgent.from_saap_agent(agent)
                    session.add(db_agent)
                    logger.info(f"βž• Added default agent to database: {agent.name}")
                
                await session.commit()
            
            logger.info("βœ… Default agents added to database")
            
        except Exception as e:
            logger.error(f"❌ Failed to add default agents to database: {e}")
            raise
    
    async def _verify_agent_loading(self, expected_count: int):
        """Verify that agent loading was successful"""
        memory_count = len(self.agents)
        
        if memory_count != expected_count:
            logger.warning(f"⚠️ Agent count mismatch: Expected {expected_count}, got {memory_count}")
        
        # Verify agent data integrity
        for agent_id, agent in self.agents.items():
            if not agent.name or agent.name == "Unknown Agent":
                logger.warning(f"⚠️ Agent {agent_id} has missing name")
            if not agent.description:
                logger.warning(f"⚠️ Agent {agent_id} has missing description")
        
        logger.info(f"βœ… Agent loading verification completed: {memory_count} agents")
    
    async def register_agent(self, agent: SaapAgent) -> bool:
        """Enhanced agent registration with guaranteed database persistence"""
        try:
            # Always add to memory cache first
            self.agents[agent.id] = agent
            logger.info(f"πŸ“ Agent added to memory: {agent.name} ({agent.id})")
            
            # Force database persistence
            await self._force_database_persistence(agent)
            
            return True
            
        except Exception as e:
            logger.error(f"❌ Enhanced agent registration failed: {e}")
            # Remove from cache if registration failed
            self.agents.pop(agent.id, None)
            return False
    
    async def _force_database_persistence(self, agent: SaapAgent):
        """Force agent persistence to database with retries"""
        max_retries = 3
        retry_count = 0
        
        while retry_count < max_retries:
            try:
                await self._ensure_database_initialized()
                
                async with db_manager.get_async_session() as session:
                    # Check if agent already exists
                    result = await session.execute(
                        select(DBAgent).where(DBAgent.id == agent.id)
                    )
                    existing = result.scalars().first()
                    
                    if existing:
                        # Update existing agent
                        updated_agent = DBAgent.from_saap_agent(agent)
                        await session.merge(updated_agent)
                        logger.info(f"πŸ”„ Agent updated in database: {agent.name}")
                    else:
                        # Create new agent
                        db_agent = DBAgent.from_saap_agent(agent)
                        session.add(db_agent)
                        logger.info(f"βž• Agent added to database: {agent.name}")
                    
                    await session.commit()
                
                logger.info(f"βœ… Database persistence successful: {agent.name}")
                return
                
            except Exception as e:
                retry_count += 1
                if retry_count >= max_retries:
                    raise Exception(f"Database persistence failed after {max_retries} retries: {e}")
                
                logger.warning(f"⚠️ Database persistence attempt {retry_count} failed: {e}")
                await asyncio.sleep(0.5)
    
    async def get_comprehensive_agent_status(self) -> Dict[str, Any]:
        """Get comprehensive status for debugging"""
        try:
            # Database agent count
            async with db_manager.get_async_session() as session:
                result = await session.execute(select(DBAgent))
                db_agents = result.scalars().all()
                db_count = len(db_agents)
            
            # Memory agent count
            memory_count = len(self.agents)
            
            # Agent details
            agent_details = []
            for agent_id, agent in self.agents.items():
                agent_details.append({
                    "id": agent_id,
                    "name": agent.name,
                    "type": agent.type.value,
                    "status": agent.status.value,
                    "has_description": bool(agent.description and agent.description != "No description available")
                })
            
            return {
                "database_initialized": db_manager.is_initialized,
                "database_agent_count": db_count,
                "memory_agent_count": memory_count,
                "sync_status": "synced" if db_count == memory_count else "out_of_sync",
                "agent_details": agent_details,
                "colossus_status": self.colossus_connection_status,
                "enhanced_features_active": True
            }
            
        except Exception as e:
            logger.error(f"❌ Status check failed: {e}")
            return {"error": str(e)}

# Create enhanced global instance
enhanced_agent_manager = EnhancedAgentManagerService()

if __name__ == "__main__":
    async def test_enhanced_agent_manager():
        """Test enhanced agent manager functionality"""
        manager = EnhancedAgentManagerService()
        await manager.initialize()
        
        # Get comprehensive status
        status = await manager.get_comprehensive_agent_status()
        print("πŸ“Š Enhanced Agent Manager Status:")
        print(f"  Database: {status.get('database_agent_count', 0)} agents")
        print(f"  Memory: {status.get('memory_agent_count', 0)} agents")
        print(f"  Sync: {status.get('sync_status', 'unknown')}")
        print(f"  Enhanced: {status.get('enhanced_features_active', False)}")
        
        await manager.shutdown_all_agents()
    
    asyncio.run(test_enhanced_agent_manager())