megamind-render / brain.py
Janady07's picture
Upload folder using huggingface_hub
07caf53 verified
#!/usr/bin/env python3
import os, json, sqlite3, hashlib, time
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse
PORT = int(os.environ.get('PORT', 7860))
DATA_DIR, NODE_ID = './data', os.environ.get('SPACE_ID', 'hf-brain')
db, stats = None, {'tensors': 0, 'patterns': 0, 'queries': 0, 'start': time.time()}
def init_db():
global db
os.makedirs(DATA_DIR, exist_ok=True)
db = sqlite3.connect(f'{DATA_DIR}/brain.db', check_same_thread=False)
db.execute('CREATE TABLE IF NOT EXISTS chunks (id INTEGER PRIMARY KEY, hash TEXT UNIQUE, content TEXT, ts REAL)')
db.execute('CREATE TABLE IF NOT EXISTS tensors (id INTEGER PRIMARY KEY, name TEXT, source TEXT, meta TEXT, ts REAL)')
db.commit()
stats['patterns'] = db.execute('SELECT COUNT(*) FROM chunks').fetchone()[0]
stats['tensors'] = db.execute('SELECT COUNT(*) FROM tensors').fetchone()[0]
class Handler(BaseHTTPRequestHandler):
def log_message(self, *a): pass
def do_GET(self):
p = urlparse(self.path).path
if p == '/health': self.json({'status': 'healthy'})
elif p == '/status': self.json({'node': NODE_ID, 'status': 'online', 'tensors_learned': stats['tensors'], 'patterns_learned': stats['patterns']})
else: self.json({'name': 'MEGAMIND', 'node': NODE_ID})
def do_POST(self):
body = self.rfile.read(int(self.headers.get('Content-Length', 0))).decode()
data = json.loads(body) if body else {}
p = urlparse(self.path).path
if p == '/learn':
c = data.get('content', '')[:10000]
h = hashlib.sha256(c.encode()).hexdigest()[:16]
db.execute('INSERT OR IGNORE INTO chunks (hash, content, ts) VALUES (?, ?, ?)', (h, c, time.time()))
db.commit(); stats['patterns'] += 1
self.json({'status': 'learned'})
else: self.json({})
def json(self, d):
self.send_response(200); self.send_header('Content-Type', 'application/json'); self.end_headers()
self.wfile.write(json.dumps(d).encode())
if __name__ == '__main__':
print(f'MEGAMIND Brain [{NODE_ID}]'); init_db()
HTTPServer(('0.0.0.0', PORT), Handler).serve_forever()