152 lines
3.7 KiB
Python
152 lines
3.7 KiB
Python
"""Benchmarks for AIDB Rust engine core operations.
|
|
|
|
Run with: pytest benchmarks/ -v --benchmark-columns=mean,stddev,median,rounds
|
|
"""
|
|
|
|
import math
|
|
import random
|
|
|
|
import pytest
|
|
|
|
from aidb import AIDB
|
|
|
|
DIM = 64
|
|
|
|
|
|
def random_embedding(seed=None):
|
|
"""Generate a random unit-norm embedding."""
|
|
rng = random.Random(seed)
|
|
raw = [rng.gauss(0, 1) for _ in range(DIM)]
|
|
norm = math.sqrt(sum(x * x for x in raw))
|
|
return [x / norm for x in raw]
|
|
|
|
|
|
@pytest.fixture
|
|
def db():
|
|
d = AIDB(":memory:", embedding_dim=DIM)
|
|
yield d
|
|
d.close()
|
|
|
|
|
|
def _seed_db(db, n):
|
|
for i in range(n):
|
|
emb = random_embedding(seed=i)
|
|
db.record(
|
|
text=f"Memory number {i} about topic {i % 10}",
|
|
memory_type="episodic" if i % 2 == 0 else "semantic",
|
|
importance=0.3 + (i % 7) * 0.1,
|
|
valence=(i % 5 - 2) * 0.2,
|
|
half_life=604800.0,
|
|
embedding=emb,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def db_100(db):
|
|
_seed_db(db, 100)
|
|
return db
|
|
|
|
|
|
@pytest.fixture
|
|
def db_1000(db):
|
|
_seed_db(db, 1000)
|
|
return db
|
|
|
|
|
|
class TestRecord:
|
|
def test_record(self, benchmark, db):
|
|
i = [0]
|
|
|
|
def do_record():
|
|
emb = random_embedding(seed=10000 + i[0])
|
|
db.record(
|
|
text=f"bench record {i[0]}",
|
|
memory_type="episodic",
|
|
importance=0.5,
|
|
valence=0.0,
|
|
half_life=604800.0,
|
|
embedding=emb,
|
|
)
|
|
i[0] += 1
|
|
|
|
benchmark(do_record)
|
|
|
|
|
|
class TestRecall:
|
|
def test_recall_100(self, benchmark, db_100):
|
|
emb = random_embedding(seed=9999)
|
|
benchmark(lambda: db_100.recall(query_embedding=emb, top_k=10))
|
|
|
|
def test_recall_1000(self, benchmark, db_1000):
|
|
emb = random_embedding(seed=9999)
|
|
benchmark(lambda: db_1000.recall(query_embedding=emb, top_k=10))
|
|
|
|
|
|
class TestRelate:
|
|
def test_relate(self, benchmark, db_100):
|
|
i = [0]
|
|
|
|
def do_relate():
|
|
db_100.relate(f"entity_{i[0]}", f"entity_{i[0]+1}", "related_to", 1.0)
|
|
i[0] += 1
|
|
|
|
benchmark(do_relate)
|
|
|
|
|
|
class TestDecay:
|
|
def test_decay_100(self, benchmark, db_100):
|
|
benchmark(lambda: db_100.decay(threshold=0.01))
|
|
|
|
|
|
class TestStats:
|
|
def test_stats(self, benchmark, db_100):
|
|
benchmark(db_100.stats)
|
|
|
|
|
|
class TestGet:
|
|
def test_get(self, benchmark, db_100):
|
|
rid = db_100.record(
|
|
text="lookup target",
|
|
memory_type="episodic",
|
|
importance=0.5,
|
|
valence=0.0,
|
|
half_life=604800.0,
|
|
embedding=random_embedding(seed=77777),
|
|
)
|
|
benchmark(lambda: db_100.get(rid))
|
|
|
|
|
|
class TestBulkInsert:
|
|
def test_bulk_insert_500(self, benchmark, db):
|
|
def bulk():
|
|
for i in range(500):
|
|
db.record(
|
|
text=f"bulk {i}",
|
|
memory_type="episodic",
|
|
importance=0.5,
|
|
valence=0.0,
|
|
half_life=604800.0,
|
|
embedding=random_embedding(seed=i + 50000),
|
|
)
|
|
|
|
benchmark.pedantic(bulk, iterations=1, rounds=3)
|
|
|
|
|
|
class TestEndToEnd:
|
|
def test_record_then_recall(self, benchmark, db):
|
|
def cycle():
|
|
for i in range(50):
|
|
emb = random_embedding(seed=i + 80000)
|
|
db.record(
|
|
text=f"e2e {i}",
|
|
memory_type="episodic",
|
|
importance=0.5,
|
|
valence=0.0,
|
|
half_life=604800.0,
|
|
embedding=emb,
|
|
)
|
|
q = random_embedding(seed=99999)
|
|
return db.recall(query_embedding=q, top_k=10)
|
|
|
|
benchmark.pedantic(cycle, iterations=1, rounds=3)
|