memoryweave/scripts/benchmark.py

#!/usr/bin/env python3
import json
import urllib.request
import time
import statistics
import concurrent.futures

API_KEY = "zhiyi-dev-key-2026"

def api(method, path, data=None):
    req = urllib.request.Request(
        f"http://localhost:7821{path}",
        data=json.dumps(data).encode() if data else None,
        headers={"X-API-Key": API_KEY, "Content-Type": "application/json"},
        method=method
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            return json.loads(resp.read())
    except Exception as e:
            if "429" in str(e):
                time.sleep(0.5)
                return api(method, path, data)  # retry once
            raise

def time_request(method, path, data=None):
    start = time.perf_counter()
    api(method, path, data)
    return (time.perf_counter() - start) * 1000  # ms

# Warmup
time.sleep(1)
for _ in range(3):
    api("POST", "/api/v1/recall", {"query": "test", "limit": 5, "namespace": "shared"})
    time.sleep(0.1)

# === 1. Recall ===
print("=== Semantic Recall Test ===")
recall_tests = [
    ("牧尘的系统是什么", "Arch"),
    ("牧尘的内存多大", "GB"),
    ("Docker", "容器"),
    ("织忆图谱", "图谱"),
    ("测试时间", "2026"),
    ("编译", "编译"),
    ("接口", "API"),
    ("memoryweave", "织忆"),
    ("MiniMax", "M2"),
    ("GPU", "RTX"),
]
recall_hits_5 = recall_hits_10 = 0
for query, expected in recall_tests:
    r5 = api("POST", "/api/v1/recall", {"query": query, "limit": 5, "namespace": "shared"})
    r10 = api("POST", "/api/v1/recall", {"query": query, "limit": 10, "namespace": "shared"})
    hit_5 = any(expected in str(r.get("content", "")) for r in r5.get("results", []))
    hit_10 = any(expected in str(r.get("content", "")) for r in r10.get("results", []))
    if hit_5:
        recall_hits_5 += 1
    if hit_10:
        recall_hits_10 += 1
    print(f"  '{query}' expected='{expected}': @5={'hit' if hit_5 else 'miss'}, @10={'hit' if hit_10 else 'miss'}")

recall_at_5 = recall_hits_5 / len(recall_tests)
recall_at_10 = recall_hits_10 / len(recall_tests)
print(f"Recall@5={recall_at_5:.4f}, Recall@10={recall_at_10:.4f}")

# === 2. Latency ===
print("\n=== Latency Test (100 reqs) ===")
endpoints = [
    ("POST", "/api/v1/recall", {"query": "牧尘", "limit": 5, "namespace": "shared"}),
    ("GET", "/api/v1/stats", None),
    ("GET", "/api/v1/graph/stats", None),
    ("POST", "/api/v1/graph/navigate", {"entity": "Docker", "max_hops": 2}),
]
lat_stats = {}
for method, path, data in endpoints:
    # warmup
    for _ in range(3):
        time_request(method, path, data)
        time.sleep(0.05)
    lats = [time_request(method, path, data) for _ in range(30)]
    sorted_lats = sorted(lats)
    p50_idx = int(len(sorted_lats) * 0.5)
    p95_idx = int(len(sorted_lats) * 0.95)
    p99_idx = int(len(sorted_lats) * 0.99)
    lat_stats[path] = {
        "p50": sorted_lats[p50_idx],
        "p95": sorted_lats[p95_idx],
        "p99": sorted_lats[p99_idx]
    }
    print(f"  {method} {path}: p50={lat_stats[path]['p50']:.1f}ms p95={lat_stats[path]['p95']:.1f}ms p99={lat_stats[path]['p99']:.1f}ms")

# === 3. Graph Navigation ===
print("\n=== Graph Navigation Test ===")
gstats = api("GET", "/api/v1/graph/stats")
nodes, edges = gstats.get("node_count", 0), gstats.get("edge_count", 0)
print(f"Graph: {nodes} nodes, {edges} edges")
entities = ["Docker", "Go", "Linux", "Arch", "Deepin", "RTX", "GPU", "Memory", "API", "Model"]
h1 = h2 = h3 = 0
for e in entities:
    for hops in [1, 2, 3]:
        r = api("POST", "/api/v1/graph/navigate", {"entity": e, "max_hops": hops})
        time.sleep(0.05)
        if r.get("paths"):
            if hops == 1:
                h1 += 1
            elif hops == 2:
                h2 += 1
            else:
                h3 += 1
trigger_rate = (h1 + h2 + h3) * 100 / 30
print(f"1-hop: {h1}/10, 2-hop: {h2}/10, 3-hop: {h3}/10, trigger rate: {trigger_rate:.1f}%")

# === 4. Concurrent Stability ===
print("\n=== Concurrent Stability Test (50 workers x 10 reqs) ===")


def worker():
    ok = fail = 0
    for i in range(10):
        try:
            api("POST", "/api/v1/recall", {"query": f"concurrent{i}", "limit": 5, "namespace": "shared"})
            ok += 1
        except:
            fail += 1
    return ok, fail


t0 = time.time()
with concurrent.futures.ThreadPoolExecutor(max_workers=50) as ex:
    results = list(ex.map(lambda _: worker(), range(50)))
t1 = time.time()
total_ok = sum(r[0] for r in results)
total_fail = sum(r[1] for r in results)
total_req = 500
success_rate = total_ok * 100 / total_req
throughput = total_req / (t1 - t0)
print(f"Success: {total_ok}/{total_req} ({success_rate:.2f}%), Fail: {total_fail}")
print(f"Duration: {t1-t0:.2f}s, Throughput: {throughput:.1f} req/s")

# Print results for extraction
print("\n=== RESULTS ===")
print(f"RECALL_5={recall_at_5:.4f}")
print(f"RECALL_10={recall_at_10:.4f}")
print(f"LAT_P50={lat_stats['/api/v1/recall']['p50']:.1f}")
print(f"LAT_P95={lat_stats['/api/v1/recall']['p95']:.1f}")
print(f"LAT_P99={lat_stats['/api/v1/recall']['p99']:.1f}")
print(f"GRAPH_NODES={nodes}")
print(f"GRAPH_EDGES={edges}")
print(f"GRAPH_TRIGGER={trigger_rate:.1f}")
print(f"CONCURRENT_OK={total_ok}")
print(f"CONCURRENT_TOTAL={total_req}")
print(f"CONCURRENT_RATE={success_rate:.2f}")
print(f"CONCURRENT_TP={throughput:.1f}")