Incorporate DB-GPT Test Dataset

marble/configs/test_config_db_single/test_config_db.yaml

@@ -0,0 +1,45 @@
+coordinate_mode: star
+relationships: []
+llm: "gpt-4o-mini"
+
+environment:
+  type: DB
+  name: "DB Simulation Environment"
+  max_iterations: 5
+  anomalies:  
+    - anomaly: MISSING_INDEXES
+      threads: 100
+      ncolumn: 20
+      nrow:    20000
+      colsize: 100
+
+communication: False
+
+task:
+  content: "Analyze the database alerts & outputs and find out the reason that caused it. The alerts might include: NodeMemSwapped, NodeLoadHigh, ... The reasons could be: ['INSERT_LARGE_DATA', 'MISSING_INDEXES','LOCK_CONTENTION','VACUUM','REDUNDANT_INDEX','INSERT_LARGE_DATA,IO_CONTENTION', 'FETCH_LARGE_DATA,CORRELATED_SUBQUERY','POOR_JOIN_PERFORMANCE,CPU_CONTENTION']. Only one of these reasons would apply. The planner should ask different experts to work on same task, and summarize their opinions into a final prediction. They can only do 3 things with tools. First is get alert. The second thing they can do is check whether a metric is abnormal using a statistical method. They can check: cpu_usage, memory_usage, network, and io. The third thing they can do is match diagnostic knowledge based on the expert and the four metrics, to guess what has caused the problem."
+  output_format: "The alerts might include: NodeMemSwapped, NodeLoadHigh, ... Please choose the most likely cause of the database anomaly from the following list, based on the expert agents: ['INSERT_LARGE_DATA', 'MISSING_INDEXES','LOCK_CONTENTION','VACUUM','REDUNDANT_INDEX','INSERT_LARGE_DATA,IO_CONTENTION', 'FETCH_LARGE_DATA,CORRELATED_SUBQUERY','POOR_JOIN_PERFORMANCE,CPU_CONTENTION']. You can ONLY CHOOSE ONE."
+
+agents:
+  - type: BaseAgent
+    agent_id: ConfigurationExpert
+    profile: "ConfigurationExpert specializes in system configurations and optimizations."
+  - type: BaseAgent
+    agent_id: CpuExpert
+    profile: "CpuExpert is knowledgeable in CPU architecture, performance, and optimizations."
+  - type: BaseAgent
+    agent_id: WorkloadExpert
+    profile: "WorkloadExpert excels in analyzing workloads, resource allocation, and optimization for efficiency."
+
+memory:
+  type: SharedMemory
+  # Additional memory configurations if needed
+
+metrics: {}
+  # Define metrics configurations for the Evaluator
+  # Example:
+  # accuracy: true
+  # response_time: true
+
+engine_planner:
+  initial_progress: "Starting the simulation."
+  # Additional engine planner configurations if needed

marble/configs/test_config_db_testset/test_config_db.yaml

@@ -0,0 +1,40 @@
+coordinate_mode: star
+relationships: []
+llm: "gpt-4o-mini"
+
+environment:
+  type: DB
+  name: "DB Simulation Environment"
+  max_iterations: 5
+  anomalies: []
+
+communication: False
+
+task:
+  content: "Analyze the database alerts & outputs and find out the reason that caused it. The alerts might include: NodeMemSwapped, NodeLoadHigh, ... The reasons could be: ['INSERT_LARGE_DATA', 'MISSING_INDEXES','LOCK_CONTENTION','VACUUM','REDUNDANT_INDEX','INSERT_LARGE_DATA,IO_CONTENTION', 'FETCH_LARGE_DATA,CORRELATED_SUBQUERY','POOR_JOIN_PERFORMANCE,CPU_CONTENTION']. Only one of these reasons would apply. The planner should ask different experts to work on same task, and summarize their opinions into a final prediction. They can only do 3 things with tools. First is get alert. The second thing they can do is check whether a metric is abnormal using a statistical method. They can check: cpu_usage, memory_usage, network, and io. The third thing they can do is match diagnostic knowledge based on the expert and the four metrics, to guess what has caused the problem."
+  output_format: "The alerts might include: NodeMemSwapped, NodeLoadHigh, ... Please choose the most likely cause of the database anomaly from the following list, based on the expert agents: ['INSERT_LARGE_DATA', 'MISSING_INDEXES','LOCK_CONTENTION','VACUUM','REDUNDANT_INDEX','INSERT_LARGE_DATA,IO_CONTENTION', 'FETCH_LARGE_DATA,CORRELATED_SUBQUERY','POOR_JOIN_PERFORMANCE,CPU_CONTENTION']. You can ONLY CHOOSE ONE."
+
+agents:
+  - type: BaseAgent
+    agent_id: ConfigurationExpert
+    profile: "ConfigurationExpert specializes in system configurations and optimizations."
+  - type: BaseAgent
+    agent_id: CpuExpert
+    profile: "CpuExpert is knowledgeable in CPU architecture, performance, and optimizations."
+  - type: BaseAgent
+    agent_id: WorkloadExpert
+    profile: "WorkloadExpert excels in analyzing workloads, resource allocation, and optimization for efficiency."
+
+memory:
+  type: SharedMemory
+  # Additional memory configurations if needed
+
+metrics: {}
+  # Define metrics configurations for the Evaluator
+  # Example:
+  # accuracy: true
+  # response_time: true
+
+engine_planner:
+  initial_progress: "Starting the simulation."
+  # Additional engine planner configurations if needed

marble/engine/engine.py

@@ -9,7 +9,7 @@
 from marble.agent import BaseAgent
 from marble.configs.config import Config
 from marble.engine.engine_planner import EnginePlanner
-from marble.environments import BaseEnvironment, ResearchEnvironment, WebEnvironment
+from marble.environments import BaseEnvironment, ResearchEnvironment, WebEnvironment, DBEnvironment
 from marble.evaluator.evaluator import Evaluator
 from marble.graph.agent_graph import AgentGraph
 from marble.memory.base_memory import BaseMemory
@@ -80,6 +80,9 @@ def _initialize_environment(self, env_config: Dict[str, Any]) -> BaseEnvironment
         elif env_type == "Research":
             env3 = ResearchEnvironment(name="Research Environment", config=env_config)
             return env3
+        elif env_type == "DB":
+            env3 = DBEnvironment(name="Database Environment", config=env_config)
+            return env3
         else:
             raise ValueError(f"Unsupported environment type: {env_type}")
 

marble/environments/__init__.py

@@ -1,9 +1,11 @@
 from .base_env import BaseEnvironment
 from .research_env import ResearchEnvironment
 from .web_env import WebEnvironment
+from .db_env import DBEnvironment
 
 __all__ = [
     'BaseEnvironment',
     'WebEnvironment',
-    'ResearchEnvironment'
+    'ResearchEnvironment',
+    'DBEnvironment'
 ]