12345678910111213141516171819202122232425 |
- from typing import Dict, Optional
- from pqai_agent.data_models.agent_configuration import AgentConfiguration
- from pqai_agent.logging_service import logger
- class AgentConfigManager:
- def __init__(self, session_maker):
- self.session_maker = session_maker
- self.agent_configs: Dict[int, AgentConfiguration] = {}
- self.refresh_configs()
- def refresh_configs(self):
- try:
- with self.session_maker() as session:
- data = session.query(AgentConfiguration).filter_by(is_delete=False).all()
- agent_configs = {}
- for config in data:
- agent_configs[config.id] = config
- self.agent_configs = agent_configs
- logger.debug(f"Refreshed agent configurations: {list(self.agent_configs.keys())}")
- except Exception as e:
- logger.error(f"Failed to refresh agent configurations: {e}")
- def get_config(self, agent_id: int) -> Optional[AgentConfiguration]:
- return self.agent_configs.get(agent_id, None)
|