Browse Source

Add agent_config_manager and service_module_manager

StrayWarrior 1 month ago
parent
commit
49a6dd7c43
2 changed files with 52 additions and 0 deletions
  1. 25 0
      pqai_agent/agent_config_manager.py
  2. 27 0
      pqai_agent/service_module_manager.py

+ 25 - 0
pqai_agent/agent_config_manager.py

@@ -0,0 +1,25 @@
+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)

+ 27 - 0
pqai_agent/service_module_manager.py

@@ -0,0 +1,27 @@
+from pqai_agent.data_models.service_module import ServiceModule, ModuleAgentType
+from pqai_agent.logging_service import logger
+
+class ServiceModuleManager:
+    def __init__(self, session_maker):
+        self.session_maker = session_maker
+        self.module_configs = {}
+        self.refresh_configs()
+
+    def refresh_configs(self):
+        try:
+            with self.session_maker() as session:
+                data = session.query(ServiceModule).filter_by(is_delete=False).all()
+                module_configs = {}
+                for module in data:
+                    module_configs[module.name] = {
+                        'display_name': module.display_name,
+                        'default_agent_type': ModuleAgentType(module.default_agent_type),
+                        'default_agent_id': module.default_agent_id
+                    }
+                self.module_configs = module_configs
+                logger.debug(f"Refreshed module configurations: {module_configs}")
+        except Exception as e:
+            logger.error(f"Error refreshing module configs: {e}")
+
+    def get_module_config(self, module_name: str):
+        return self.module_configs.get(module_name)