rajkumarrawal commited on
Commit
45d2fdb
·
1 Parent(s): 91a4bb5

Refactor AgentManager to lazily initialize agents to avoid pickling on startup

Browse files

Remove eager initialization in __init__ and call _initialize_default_agents on first get_agent call. Update Gradio dropdown to static list of agent types to avoid referencing not-yet-created agent dict.

Files changed (1) hide show
  1. app.py +20 -13
app.py CHANGED
@@ -206,7 +206,7 @@ class SpacesConfig:
206
 
207
  class AgentManager:
208
  """Manages agent instances and resources for the Spaces deployment."""
209
-
210
  def __init__(self, config: SpacesConfig):
211
  self.config = config
212
  self.agents = {}
@@ -221,24 +221,26 @@ class AgentManager:
221
  "cpu_usage": [],
222
  "timestamps": []
223
  }
224
-
225
- # Initialize default agents
226
- self._initialize_default_agents()
227
-
228
- logger.info("Agent manager initialized with default agents")
229
 
230
  def _initialize_default_agents(self):
231
- """Initialize default agent instances."""
 
 
 
232
  try:
233
  # Initialize different agent types for different use cases
234
  agent_types = [
235
  "EnterpriseAgent",
236
- "ConsumerAgent",
237
  "CreativeAgent",
238
  "VoiceAgent",
239
  "GeneralAgent"
240
  ]
241
-
242
  for agent_type in agent_types:
243
  try:
244
  agent = RefactoredAutonomousAgent(agent_type)
@@ -252,7 +254,7 @@ class AgentManager:
252
  logger.info(f"Initialized {agent_type} agent")
253
  except Exception as e:
254
  logger.error(f"Failed to initialize {agent_type} agent: {e}")
255
-
256
  # Add a general-purpose agent as fallback
257
  if "GeneralAgent" not in self.agents:
258
  self.agents["GeneralAgent"] = RefactoredAutonomousAgent("GeneralAgent")
@@ -263,7 +265,7 @@ class AgentManager:
263
  "last_used": None
264
  }
265
  logger.info("Initialized fallback GeneralAgent")
266
-
267
  except Exception as e:
268
  logger.error(f"Failed to initialize agents: {e}")
269
  # Create a minimal fallback agent
@@ -276,7 +278,12 @@ class AgentManager:
276
  }
277
 
278
  def get_agent(self, agent_type: str = "GeneralAgent") -> Optional[RefactoredAutonomousAgent]:
279
- """Get an agent instance by type."""
 
 
 
 
 
280
  return self.agents.get(agent_type, self.agents.get("GeneralAgent"))
281
 
282
  def update_performance_metrics(self, agent_type: str, response_time: float, success: bool):
@@ -487,7 +494,7 @@ class SpacesApp:
487
  with gr.Column(scale=2):
488
  gr.HTML("<h3>🎯 Select Agent Type</h3>")
489
  agent_type = gr.Dropdown(
490
- choices=list(self.agent_manager.agents.keys()),
491
  value="GeneralAgent",
492
  label="Choose Agent Type"
493
  )
 
206
 
207
  class AgentManager:
208
  """Manages agent instances and resources for the Spaces deployment."""
209
+
210
  def __init__(self, config: SpacesConfig):
211
  self.config = config
212
  self.agents = {}
 
221
  "cpu_usage": [],
222
  "timestamps": []
223
  }
224
+
225
+ # Don't initialize agents during construction to avoid pickling issues
226
+ # Agents will be created lazily when first requested
227
+ logger.info("Agent manager initialized (agents will be created lazily)")
 
228
 
229
  def _initialize_default_agents(self):
230
+ """Initialize default agent instances lazily."""
231
+ if self.agents: # Already initialized
232
+ return
233
+
234
  try:
235
  # Initialize different agent types for different use cases
236
  agent_types = [
237
  "EnterpriseAgent",
238
+ "ConsumerAgent",
239
  "CreativeAgent",
240
  "VoiceAgent",
241
  "GeneralAgent"
242
  ]
243
+
244
  for agent_type in agent_types:
245
  try:
246
  agent = RefactoredAutonomousAgent(agent_type)
 
254
  logger.info(f"Initialized {agent_type} agent")
255
  except Exception as e:
256
  logger.error(f"Failed to initialize {agent_type} agent: {e}")
257
+
258
  # Add a general-purpose agent as fallback
259
  if "GeneralAgent" not in self.agents:
260
  self.agents["GeneralAgent"] = RefactoredAutonomousAgent("GeneralAgent")
 
265
  "last_used": None
266
  }
267
  logger.info("Initialized fallback GeneralAgent")
268
+
269
  except Exception as e:
270
  logger.error(f"Failed to initialize agents: {e}")
271
  # Create a minimal fallback agent
 
278
  }
279
 
280
  def get_agent(self, agent_type: str = "GeneralAgent") -> Optional[RefactoredAutonomousAgent]:
281
+ """Get an agent instance by type, creating agents lazily."""
282
+ # Initialize agents on first request
283
+ if not self.agents:
284
+ self._initialize_default_agents()
285
+
286
+ # Return requested agent or fallback
287
  return self.agents.get(agent_type, self.agents.get("GeneralAgent"))
288
 
289
  def update_performance_metrics(self, agent_type: str, response_time: float, success: bool):
 
494
  with gr.Column(scale=2):
495
  gr.HTML("<h3>🎯 Select Agent Type</h3>")
496
  agent_type = gr.Dropdown(
497
+ choices=["EnterpriseAgent", "ConsumerAgent", "CreativeAgent", "VoiceAgent", "GeneralAgent"],
498
  value="GeneralAgent",
499
  label="Choose Agent Type"
500
  )