File size: 2,197 Bytes
07caf53
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
#!/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()