Building an Autonomous Data and Infrastructure Management System with Lightweight LLMs
This guide walks you through constructing a dynamic, agent-driven data infrastructure strategy using the compact yet powerful Qwen2.5-0.5B-Instruct model. We start by establishing a versatile large language model (LLM) agent framework, then craft dedicated agents responsible for distinct facets of data handling-from ingestion and quality evaluation to infrastructure tuning. These agents are orchestrated to collaborate seamlessly, enabling efficient multi-agent workflows across complex data pipelines. Practical scenarios, including retail and IoT data streams, illustrate how autonomous agents can simplify and optimize data operations.
Setting Up the Core Lightweight LLM Agent
Our foundation is a lightweight LLM agent built on the Qwen2.5-0.5B-Instruct model, optimized for resource-efficient deployment. This base agent loads the tokenizer and model, supports contextual dialogue, and generates intelligent responses tailored to its assigned role. This modular design allows us to extend functionality easily and run the system effectively in environments like Google Colab.
!pip install -q transformers torch accelerate datasets huggingface_hub
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import json, time
from typing import List, Dict, Any
from dataclasses import dataclass
from datetime import datetime
import pandas as pd
class LightweightLLMAgent:
def __init__(self, role: str, model_name: str = "Qwen/Qwen2.5-0.5B-Instruct"):
self.role = role
self.model_name = model_name
self.device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading {model_name} for {role} agent on {self.device}...")
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
device_map="auto"
)
self.conversation_history = []
def generate_response(self, prompt: str, max_tokens: int = 150) -> str:
messages = [
{"role": "system", "content": f"You are a {self.role} agent in a data infrastructure system."},
{"role": "user", "content": prompt}
]
text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
model_inputs = self.tokenizer([text], return_tensors="pt").to(self.device)
with torch.no_grad():
generated_ids = self.model.generate(
model_inputs.input_ids,
max_new_tokens=max_tokens,
temperature=0.7,
do_sample=True,
top_p=0.95
)
generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
self.conversation_history.append({"prompt": prompt, "response": response})
return response
Specialized Agents for Data Pipeline Management
We develop three distinct agents, each focusing on a critical layer of data pipeline management:
Data Ingestion Agent
This agent analyzes incoming data sources and recommends optimal ingestion strategies based on source type, volume, and frequency. For example, it might suggest batch processing for daily CSV uploads or streaming ingestion for real-time APIs.
class DataIngestionAgent(LightweightLLMAgent):
def __init__(self):
super().__init__(role="Data Ingestion Specialist")
def analyze_data_source(self, source_info: Dict) -> Dict:
prompt = f"""Evaluate this data source and propose an ingestion plan:
Source Type: {source_info.get('type', 'unknown')}
Volume: {source_info.get('volume', 'unknown')}
Frequency: {source_info.get('frequency', 'unknown')}
Outline: 1) Recommended ingestion method, 2) Key considerations."""
strategy = self.generate_response(prompt, max_tokens=100)
return {"source": source_info, "strategy": strategy, "timestamp": datetime.now().isoformat()}
Data Quality Agent
Focused on assessing data integrity, this agent reviews metrics like completeness and consistency, identifies issues, and offers prioritized recommendations to improve data reliability.
class DataQualityAgent(LightweightLLMAgent):
def __init__(self):
super().__init__(role="Data Quality Analyst")
def assess_data_quality(self, data_sample: Dict) -> Dict:
prompt = f"""Evaluate the quality of this data sample:
Completeness: {data_sample.get('completeness', 'N/A')}%
Consistency: {data_sample.get('consistency', 'N/A')}%
Detected Issues: {data_sample.get('issues', 0)}
Provide a concise quality assessment and top two improvement suggestions."""
assessment = self.generate_response(prompt, max_tokens=100)
return {"assessment": assessment, "severity": self._calculate_severity(data_sample), "timestamp": datetime.now().isoformat()}
def _calculate_severity(self, data_sample: Dict) -> str:
completeness = data_sample.get('completeness', 100)
consistency = data_sample.get('consistency', 100)
avg_score = (completeness + consistency) / 2
if avg_score >= 90:
return "LOW"
elif avg_score >= 70:
return "MEDIUM"
else:
return "HIGH"
Infrastructure Optimization Agent
This agent continuously monitors system metrics such as CPU load, memory consumption, storage usage, and query latency. It then generates actionable recommendations to enhance infrastructure efficiency and scalability.
class InfrastructureOptimizationAgent(LightweightLLMAgent):
def __init__(self):
super().__init__(role="Infrastructure Optimization Specialist")
def optimize_resources(self, metrics: Dict) -> Dict:
prompt = f"""Review these infrastructure metrics and suggest improvements:
CPU Usage: {metrics.get('cpu_usage', 0)}%
Memory Usage: {metrics.get('memory_usage', 0)}%
Storage: {metrics.get('storage_used', 0)}GB / {metrics.get('storage_total', 0)}GB
Query Latency: {metrics.get('query_latency', 0)}ms
Provide two optimization recommendations."""
recommendations = self.generate_response(prompt, max_tokens=100)
return {
"current_metrics": metrics,
"recommendations": recommendations,
"priority": self._calculate_priority(metrics),
"timestamp": datetime.now().isoformat()
}
def _calculate_priority(self, metrics: Dict) -> str:
cpu = metrics.get('cpu_usage', 0)
memory = metrics.get('memory_usage', 0)
if cpu > 85 or memory > 85:
return "CRITICAL"
elif cpu > 70 or memory > 70:
return "HIGH"
else:
return "NORMAL"
Coordinating Agents with an Orchestrator
To unify the agents into a cohesive system, we implement an orchestrator that manages the sequential execution of ingestion analysis, quality assessment, and infrastructure optimization. This centralized controller logs each stage’s output and compiles comprehensive reports, enabling transparent monitoring and streamlined automation.
class AgenticDataOrchestrator:
def __init__(self):
print("n" + "="*70)
print("Initializing Agentic Data Infrastructure System")
print("="*70 + "n")
self.ingestion_agent = DataIngestionAgent()
self.quality_agent = DataQualityAgent()
self.optimization_agent = InfrastructureOptimizationAgent()
self.execution_log = []
def process_data_pipeline(self, pipeline_config: Dict) -> Dict:
results = {
"pipeline_id": pipeline_config.get("id", "unknown"),
"start_time": datetime.now().isoformat(),
"stages": []
}
print("n[Stage 1] Data Ingestion Analysis")
ingestion_result = self.ingestion_agent.analyze_data_source(pipeline_config.get("source", {}))
print(f"Strategy: {ingestion_result['strategy'][:150]}...")
results["stages"].append({"stage": "ingestion", "result": ingestion_result})
print("n[Stage 2] Data Quality Assessment")
quality_result = self.quality_agent.assess_data_quality(pipeline_config.get("quality_metrics", {}))
print(f"Assessment: {quality_result['assessment'][:150]}...")
print(f"Severity: {quality_result['severity']}")
results["stages"].append({"stage": "quality", "result": quality_result})
print("n[Stage 3] Infrastructure Optimization")
optimization_result = self.optimization_agent.optimize_resources(pipeline_config.get("infrastructure_metrics", {}))
print(f"Recommendations: {optimization_result['recommendations'][:150]}...")
print(f"Priority: {optimization_result['priority']}")
results["stages"].append({"stage": "optimization", "result": optimization_result})
results["end_time"] = datetime.now().isoformat()
results["status"] = "completed"
self.execution_log.append(results)
return results
def generate_summary_report(self) -> pd.DataFrame:
if not self.execution_log:
return pd.DataFrame()
summary_data = []
for log in self.execution_log:
summary_data.append({
"Pipeline ID": log["pipeline_id"],
"Start Time": log["start_time"],
"Status": log["status"],
"Stages Completed": len(log["stages"])
})
return pd.DataFrame(summary_data)
Real-World Applications: E-commerce and IoT Pipelines
To showcase the system’s capabilities, we simulate two distinct data pipelines:
- E-commerce Pipeline: Handles real-time REST API data with moderate volume and frequent updates, focusing on customer transactions and inventory data.
- IoT Sensor Pipeline: Processes high-volume streaming data from Kafka message queues, typical of sensor telemetry in smart environments.
Each pipeline is autonomously managed by the orchestrator, with agents performing their specialized tasks and contributing to an integrated workflow. The final summary report consolidates execution details, demonstrating the system’s effectiveness in diverse scenarios.
def main():
orchestrator = AgenticDataOrchestrator()
print("n" + "="*70)
print("EXAMPLE 1: E-commerce Data Pipeline")
print("="*70)
ecommerce_pipeline = {
"id": "ecommerce_pipeline_001",
"source": {"type": "REST API", "volume": "10GB/day", "frequency": "real-time"},
"quality_metrics": {"completeness": 87, "consistency": 92, "issues": 15},
"infrastructure_metrics": {"cpu_usage": 78, "memory_usage": 82, "storage_used": 450, "storage_total": 1000, "query_latency": 250}
}
result1 = orchestrator.process_data_pipeline(ecommerce_pipeline)
print("n" + "="*70)
print("EXAMPLE 2: IoT Sensor Data Pipeline")
print("="*70)
iot_pipeline = {
"id": "iot_pipeline_002",
"source": {"type": "Message Queue (Kafka)", "volume": "50GB/day", "frequency": "streaming"},
"quality_metrics": {"completeness": 95, "consistency": 88, "issues": 8},
"infrastructure_metrics": {"cpu_usage": 65, "memory_usage": 71, "storage_used": 780, "storage_total": 2000, "query_latency": 180}
}
result2 = orchestrator.process_data_pipeline(iot_pipeline)
print("n" + "="*70)
print("EXECUTION SUMMARY REPORT")
print("="*70 + "n")
summary_df = orchestrator.generate_summary_report()
print(summary_df.to_string(index=False))
print("n" + "="*70)
print("Tutorial Complete!")
print("="*70)
print("nKey Highlights:")
print("✔ Modular lightweight LLM agent design")
print("✔ Dedicated agents for ingestion, quality, and optimization")
print("✔ Coordinated multi-agent orchestration")
print("✔ Real-time infrastructure monitoring and tuning")
print("✔ Autonomous decision-making in complex data workflows")
if __name__ == "__main__":
main()
Summary and Future Outlook
This project demonstrates how a compact open-source LLM can power a multi-agent data infrastructure system capable of autonomous analysis, evaluation, and optimization. By distributing responsibilities across specialized agents and orchestrating their collaboration, we transform traditional data pipelines into adaptive, self-regulating ecosystems. Such frameworks are increasingly vital for scalable enterprise applications, where agility and efficiency in data management are paramount.
