agent_config_manager.py 1.0 KB

12345678910111213141516171819202122232425
  1. from typing import Dict, Optional
  2. from pqai_agent.data_models.agent_configuration import AgentConfiguration
  3. from pqai_agent.logging_service import logger
  4. class AgentConfigManager:
  5. def __init__(self, session_maker):
  6. self.session_maker = session_maker
  7. self.agent_configs: Dict[int, AgentConfiguration] = {}
  8. self.refresh_configs()
  9. def refresh_configs(self):
  10. try:
  11. with self.session_maker() as session:
  12. data = session.query(AgentConfiguration).filter_by(is_delete=False).all()
  13. agent_configs = {}
  14. for config in data:
  15. agent_configs[config.id] = config
  16. self.agent_configs = agent_configs
  17. logger.debug(f"Refreshed agent configurations: {list(self.agent_configs.keys())}")
  18. except Exception as e:
  19. logger.error(f"Failed to refresh agent configurations: {e}")
  20. def get_config(self, agent_id: int) -> Optional[AgentConfiguration]:
  21. return self.agent_configs.get(agent_id, None)