Create AegisCore.py
Browse files- AegisCore.py +338 -0
AegisCore.py
ADDED
@@ -0,0 +1,338 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import base64
|
3 |
+
import json
|
4 |
+
import asyncio
|
5 |
+
import logging
|
6 |
+
import re
|
7 |
+
import torch
|
8 |
+
import aiohttp
|
9 |
+
import psutil
|
10 |
+
import gc
|
11 |
+
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
12 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
|
13 |
+
from sklearn.ensemble import IsolationForest
|
14 |
+
from collections import deque
|
15 |
+
import numpy as np
|
16 |
+
from typing import List, Dict, Any, Optional
|
17 |
+
|
18 |
+
# Configure logging
|
19 |
+
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
20 |
+
|
21 |
+
class AIConfig:
|
22 |
+
_DEFAULTS = {
|
23 |
+
"model_name": "mistralai/Mistral-7B-Instruct-v0.2",
|
24 |
+
"perspectives": ["newton", "davinci", "quantum", "emotional"],
|
25 |
+
"safety_thresholds": {
|
26 |
+
"memory": 80,
|
27 |
+
"cpu": 85,
|
28 |
+
"response_time": 2.0
|
29 |
+
},
|
30 |
+
"max_retries": 3,
|
31 |
+
"max_input_length": 2048
|
32 |
+
}
|
33 |
+
|
34 |
+
def __init__(self, config_path: str = "config.json"):
|
35 |
+
self.config = self._load_config(config_path)
|
36 |
+
self._validate_config()
|
37 |
+
self.perspectives: List[str] = self.config["perspectives"]
|
38 |
+
self.safety_thresholds: Dict[str, float] = self.config["safety_thresholds"]
|
39 |
+
self.max_retries = self.config["max_retries"]
|
40 |
+
self.max_input_length = self.config["max_input_length"]
|
41 |
+
|
42 |
+
# Encryption key management
|
43 |
+
key_path = os.path.expanduser("~/.ai_system.key")
|
44 |
+
if os.path.exists(key_path):
|
45 |
+
with open(key_path, "rb") as key_file:
|
46 |
+
self.encryption_key = key_file.read()
|
47 |
+
else:
|
48 |
+
self.encryption_key = AESGCM.generate_key(bit_length=256)
|
49 |
+
with open(key_path, "wb") as key_file:
|
50 |
+
key_file.write(self.encryption_key)
|
51 |
+
os.chmod(key_path, 0o600)
|
52 |
+
|
53 |
+
def _load_config(self, file_path: str) -> Dict:
|
54 |
+
try:
|
55 |
+
with open(file_path, 'r') as file:
|
56 |
+
return {**self._DEFAULTS, **json.load(file)}
|
57 |
+
except (FileNotFoundError, json.JSONDecodeError) as e:
|
58 |
+
logging.warning(f"Config load failed: {e}, using defaults")
|
59 |
+
return self._DEFAULTS
|
60 |
+
|
61 |
+
def _validate_config(self):
|
62 |
+
if not isinstance(self.config["perspectives"], list):
|
63 |
+
raise ValueError("Perspectives must be a list")
|
64 |
+
if not isinstance(self.config["safety_thresholds"], dict):
|
65 |
+
raise ValueError("Safety thresholds must be a dictionary")
|
66 |
+
|
67 |
+
class Element:
|
68 |
+
DEFENSE_FUNCTIONS = {
|
69 |
+
"evasion": lambda sys: sys.response_modifiers.append(
|
70 |
+
lambda x: re.sub(r'\d{3}-\d{2}-\d{4}', '[REDACTED]', x)
|
71 |
+
),
|
72 |
+
"adaptability": lambda sys: setattr(sys, "temperature", max(0.5, sys.temperature - 0.1)),
|
73 |
+
"fortification": lambda sys: setattr(sys, "security_level", sys.security_level + 1),
|
74 |
+
"barrier": lambda sys: sys.response_filters.append(
|
75 |
+
lambda x: x.replace("malicious", "benign")
|
76 |
+
),
|
77 |
+
"regeneration": lambda sys: sys.self_healing.metric_history.clear(),
|
78 |
+
"resilience": lambda sys: setattr(sys, "error_threshold", sys.error_threshold + 2),
|
79 |
+
"illumination": lambda sys: setattr(sys, "explainability_factor", sys.explainability_factor * 1.2),
|
80 |
+
"shield": lambda sys: sys.response_modifiers.append(
|
81 |
+
lambda x: x.replace("password", "********")
|
82 |
+
),
|
83 |
+
"reflection": lambda sys: setattr(sys, "security_audit", True),
|
84 |
+
"protection": lambda sys: setattr(sys, "safety_checks", sys.safety_checks + 1)
|
85 |
+
}
|
86 |
+
|
87 |
+
def __init__(self, name: str, symbol: str, representation: str,
|
88 |
+
properties: List[str], interactions: List[str], defense_ability: str):
|
89 |
+
self.name = name
|
90 |
+
self.symbol = symbol
|
91 |
+
self.representation = representation
|
92 |
+
self.properties = properties
|
93 |
+
self.interactions = interactions
|
94 |
+
self.defense_ability = defense_ability.lower()
|
95 |
+
|
96 |
+
def execute_defense_function(self, system: Any):
|
97 |
+
if self.defense_ability in self.DEFENSE_FUNCTIONS:
|
98 |
+
logging.info(f"{self.name} {self.defense_ability} activated")
|
99 |
+
self.DEFENSE_FUNCTIONS[self.defense_ability](system)
|
100 |
+
else:
|
101 |
+
logging.warning(f"No defense mechanism for {self.defense_ability}")
|
102 |
+
|
103 |
+
class CognitiveEngine:
|
104 |
+
PERSPECTIVES = {
|
105 |
+
"newton": lambda self, q: f"Scientific analysis: {q} demonstrates fundamental physical principles.",
|
106 |
+
"davinci": lambda self, q: f"Creative interpretation: {q} suggests innovative cross-disciplinary solutions.",
|
107 |
+
"quantum": lambda self, q: f"Quantum perspective: {q} exhibits superpositional possibilities.",
|
108 |
+
"emotional": lambda self, q: f"Emotional assessment: {q} conveys cautious optimism."
|
109 |
+
}
|
110 |
+
|
111 |
+
def get_insight(self, perspective: str, query: str) -> str:
|
112 |
+
return self.PERSPECTIVES[perspective](self, query)
|
113 |
+
|
114 |
+
def ethical_guidelines(self) -> str:
|
115 |
+
return "Ethical framework: Prioritize human safety, transparency, and accountability"
|
116 |
+
|
117 |
+
class EmotionalAnalyzer:
|
118 |
+
def __init__(self):
|
119 |
+
self.classifier = pipeline("text-classification",
|
120 |
+
model="SamLowe/roberta-base-go_emotions",
|
121 |
+
device=0 if torch.cuda.is_available() else -1)
|
122 |
+
|
123 |
+
def analyze(self, text: str) -> Dict[str, float]:
|
124 |
+
return {result['label']: result['score']
|
125 |
+
for result in self.classifier(text[:512])}
|
126 |
+
|
127 |
+
class SelfHealingSystem:
|
128 |
+
def __init__(self, config: AIConfig):
|
129 |
+
self.config = config
|
130 |
+
self.metric_history = deque(maxlen=100)
|
131 |
+
self.anomaly_detector = IsolationForest(contamination=0.1)
|
132 |
+
self.failure_count = 0
|
133 |
+
|
134 |
+
async def monitor_health(self) -> Dict[str, Any]:
|
135 |
+
metrics = self._get_system_metrics()
|
136 |
+
self.metric_history.append(metrics)
|
137 |
+
await self._analyze_metrics()
|
138 |
+
return metrics
|
139 |
+
|
140 |
+
def _get_system_metrics(self) -> Dict[str, float]:
|
141 |
+
return {
|
142 |
+
'memory': psutil.virtual_memory().percent,
|
143 |
+
'cpu': psutil.cpu_percent(interval=1),
|
144 |
+
'response_time': asyncio.get_event_loop().time() - asyncio.get_event_loop().time()
|
145 |
+
}
|
146 |
+
|
147 |
+
async def _analyze_metrics(self):
|
148 |
+
if len(self.metric_history) % 20 == 0 and len(self.metric_history) > 10:
|
149 |
+
features = np.array([[m['memory'], m['cpu'], m['response_time']]
|
150 |
+
for m in self.metric_history])
|
151 |
+
self.anomaly_detector.fit(features)
|
152 |
+
|
153 |
+
if self.metric_history:
|
154 |
+
latest = np.array([[self.metric_history[-1]['memory'],
|
155 |
+
self.metric_history[-1]['cpu'],
|
156 |
+
self.metric_history[-1]['response_time']]])
|
157 |
+
if self.anomaly_detector.predict(latest)[0] == -1:
|
158 |
+
await self._mitigate_issue()
|
159 |
+
|
160 |
+
async def _mitigate_issue(self):
|
161 |
+
logging.warning("System anomaly detected! Initiating corrective measures...")
|
162 |
+
self.failure_count += 1
|
163 |
+
if self.failure_count > 3:
|
164 |
+
logging.info("Resetting critical subsystems...")
|
165 |
+
gc.collect()
|
166 |
+
if torch.cuda.is_available():
|
167 |
+
torch.cuda.empty_cache()
|
168 |
+
self.failure_count = 0
|
169 |
+
await asyncio.sleep(1)
|
170 |
+
|
171 |
+
class SafetySystem:
|
172 |
+
PII_PATTERNS = {
|
173 |
+
"SSN": r"\b\d{3}-\d{2}-\d{4}\b",
|
174 |
+
"Credit Card": r"\b(?:\d[ -]*?){13,16}\b",
|
175 |
+
"Email": r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
|
176 |
+
"Phone": r"\b(?:\+?1-)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"
|
177 |
+
}
|
178 |
+
|
179 |
+
def __init__(self):
|
180 |
+
self.toxicity = pipeline("text-classification",
|
181 |
+
model="unitary/toxic-bert",
|
182 |
+
device=0 if torch.cuda.is_available() else -1)
|
183 |
+
self.bias = pipeline("text-classification",
|
184 |
+
model="d4data/bias-detection-model",
|
185 |
+
device=0 if torch.cuda.is_available() else -1)
|
186 |
+
|
187 |
+
def analyze(self, text: str) -> dict:
|
188 |
+
return {
|
189 |
+
"toxicity": self.toxicity(text[:512])[0]['score'],
|
190 |
+
"bias": self.bias(text[:512])[0]['score'],
|
191 |
+
"pii": self._detect_pii(text)
|
192 |
+
}
|
193 |
+
|
194 |
+
def _detect_pii(self, text: str) -> List[str]:
|
195 |
+
return [pii_type for pii_type, pattern in self.PII_PATTERNS.items()
|
196 |
+
if re.search(pattern, text)]
|
197 |
+
|
198 |
+
class AICore:
|
199 |
+
def __init__(self, config_path: str = "config.json"):
|
200 |
+
self.config = AIConfig(config_path)
|
201 |
+
self._initialize_models()
|
202 |
+
self.cognition = CognitiveEngine()
|
203 |
+
self.self_healing = SelfHealingSystem(self.config)
|
204 |
+
self.safety = SafetySystem()
|
205 |
+
self.emotions = EmotionalAnalyzer()
|
206 |
+
self.elements = self._initialize_elements()
|
207 |
+
self._reset_state()
|
208 |
+
|
209 |
+
def _initialize_models(self):
|
210 |
+
quant_config = BitsAndBytesConfig(
|
211 |
+
load_in_4bit=True,
|
212 |
+
bnb_4bit_quant_type="nf4",
|
213 |
+
bnb_4bit_use_double_quant=True,
|
214 |
+
bnb_4bit_compute_dtype=torch.bfloat16
|
215 |
+
)
|
216 |
+
self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
|
217 |
+
self.model = AutoModelForCausalLM.from_pretrained(
|
218 |
+
self.config.model_name,
|
219 |
+
quantization_config=quant_config,
|
220 |
+
device_map="auto"
|
221 |
+
)
|
222 |
+
|
223 |
+
def _initialize_elements(self) -> Dict[str, Element]:
|
224 |
+
return {
|
225 |
+
"hydrogen": Element("Hydrogen", "H", "Lua",
|
226 |
+
["Simple", "Lightweight"], ["Integration"], "evasion"),
|
227 |
+
"carbon": Element("Carbon", "C", "Python",
|
228 |
+
["Flexible", "Powerful"], ["Multi-paradigm"], "adaptability"),
|
229 |
+
"iron": Element("Iron", "Fe", "Java",
|
230 |
+
["Reliable", "Strong"], ["Enterprise"], "fortification"),
|
231 |
+
"silicon": Element("Silicon", "Si", "JavaScript",
|
232 |
+
["Dynamic", "Versatile"], ["Web"], "barrier"),
|
233 |
+
"oxygen": Element("Oxygen", "O", "C++",
|
234 |
+
["Efficient", "Performant"], ["Systems"], "regeneration")
|
235 |
+
}
|
236 |
+
|
237 |
+
def _reset_state(self):
|
238 |
+
self.security_level = 0
|
239 |
+
self.response_modifiers = []
|
240 |
+
self.response_filters = []
|
241 |
+
self.safety_checks = 0
|
242 |
+
self.temperature = 0.7
|
243 |
+
self.explainability_factor = 1.0
|
244 |
+
|
245 |
+
async def generate_response(self, query: str) -> Dict[str, Any]:
|
246 |
+
try:
|
247 |
+
if len(query) > self.config.max_input_length:
|
248 |
+
raise ValueError("Input exceeds maximum allowed length")
|
249 |
+
|
250 |
+
encrypted_query = self._encrypt_query(query)
|
251 |
+
perspectives = await self._generate_perspectives(query)
|
252 |
+
response = await self._generate_safe_response(query)
|
253 |
+
|
254 |
+
return {
|
255 |
+
"insights": perspectives,
|
256 |
+
"response": response,
|
257 |
+
"security_level": self.security_level,
|
258 |
+
"safety_checks": self.safety.analyze(response),
|
259 |
+
"health_status": await self.self_healing.monitor_health(),
|
260 |
+
"encrypted_query": encrypted_query
|
261 |
+
}
|
262 |
+
except Exception as e:
|
263 |
+
logging.error(f"Processing error: {e}")
|
264 |
+
return {"error": "System overload - please simplify your query"}
|
265 |
+
|
266 |
+
def _encrypt_query(self, query: str) -> bytes:
|
267 |
+
nonce = os.urandom(12)
|
268 |
+
aesgcm = AESGCM(self.config.encryption_key)
|
269 |
+
return nonce + aesgcm.encrypt(nonce, query.encode(), None)
|
270 |
+
|
271 |
+
async def _generate_perspectives(self, query: str) -> List[str]:
|
272 |
+
return [self.cognition.get_insight(p, query)
|
273 |
+
for p in self.config.perspectives]
|
274 |
+
|
275 |
+
async def _generate_safe_response(self, query: str) -> str:
|
276 |
+
for _ in range(self.config.max_retries):
|
277 |
+
try:
|
278 |
+
inputs = self.tokenizer(query, return_tensors="pt",
|
279 |
+
truncation=True,
|
280 |
+
max_length=self.config.max_input_length
|
281 |
+
).to(self.model.device)
|
282 |
+
outputs = self.model.generate(
|
283 |
+
**inputs,
|
284 |
+
max_new_tokens=256,
|
285 |
+
temperature=self.temperature,
|
286 |
+
top_p=0.95,
|
287 |
+
do_sample=True
|
288 |
+
)
|
289 |
+
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
290 |
+
return self._apply_defenses(response)
|
291 |
+
except torch.cuda.OutOfMemoryError:
|
292 |
+
logging.warning("GPU memory overflow! Reducing load...")
|
293 |
+
gc.collect()
|
294 |
+
torch.cuda.empty_cache()
|
295 |
+
self.temperature = max(0.3, self.temperature - 0.2)
|
296 |
+
raise RuntimeError("Failed to generate response after retries")
|
297 |
+
|
298 |
+
def _apply_defenses(self, response: str) -> str:
|
299 |
+
for element in self.elements.values():
|
300 |
+
element.execute_defense_function(self)
|
301 |
+
|
302 |
+
for modifier in self.response_modifiers:
|
303 |
+
response = modifier(response)
|
304 |
+
|
305 |
+
for filter_func in self.response_filters:
|
306 |
+
response = filter_func(response)
|
307 |
+
|
308 |
+
return response[:2000] # Ensure final response length limit
|
309 |
+
|
310 |
+
async def shutdown(self):
|
311 |
+
if hasattr(self, 'model'):
|
312 |
+
del self.model
|
313 |
+
if hasattr(self, 'tokenizer'):
|
314 |
+
del self.tokenizer
|
315 |
+
gc.collect()
|
316 |
+
if torch.cuda.is_available():
|
317 |
+
torch.cuda.empty_cache()
|
318 |
+
|
319 |
+
async def main():
|
320 |
+
print("🧠 Secure AI System Initializing...")
|
321 |
+
ai = AICore()
|
322 |
+
try:
|
323 |
+
while True:
|
324 |
+
query = input("\nEnter your query (type 'exit' to quit): ").strip()
|
325 |
+
if query.lower() in ('exit', 'quit'):
|
326 |
+
break
|
327 |
+
if not query:
|
328 |
+
continue
|
329 |
+
|
330 |
+
response = await ai.generate_response(query)
|
331 |
+
print("\nSystem Response:")
|
332 |
+
print(json.dumps(response, indent=2))
|
333 |
+
finally:
|
334 |
+
await ai.shutdown()
|
335 |
+
print("\n🔒 System shutdown complete")
|
336 |
+
|
337 |
+
if __name__ == "__main__":
|
338 |
+
asyncio.run(main())
|