wehub-resource-sync
229 lines
9.0 KiB
Python
229 lines
9.0 KiB
Python
"""Deterministic corpus seeder for the performance benchmark.
|
||
|
||
The v1 harness booted an empty DB, so the read journeys measured a best-case
|
||
near-empty table. This seeds a sizeable, realistic corpus directly through the
|
||
store API (no HTTP, no runner) so ``list_sessions`` / ``get_session`` /
|
||
``load_conversation_history`` read a production-shaped volume.
|
||
|
||
Writes to the same DB URI the server later boots against; startup migrations
|
||
are an idempotent no-op on an at-head DB. The seed is deterministic (fixed RNG,
|
||
fixed counts) so the same config always yields the same corpus — which is what
|
||
makes "seed once, reuse" sound. The reuse marker records the Alembic head read
|
||
at seed time, so a corpus from an older schema is auto-reseeded (no manual
|
||
revision bookkeeping).
|
||
|
||
Listable-corpus recipe, per session (the permission grant is the gotcha — the
|
||
loopback server resolves every request to user ``"local"`` and
|
||
``list_sessions`` filters by it):
|
||
|
||
1. ``create_session_with_agent`` — conversation + session-scoped agent row.
|
||
2. ``permission_store.grant("local", sid, LEVEL_OWNER)`` — makes it listable.
|
||
3. one batched ``append(sid, items)`` — user-role message items.
|
||
|
||
Run standalone::
|
||
|
||
uv run --no-sync dev/benchmarks/omnigent/seed.py \
|
||
--database-uri sqlite:///tmp/bench.db --sessions 5000 --items-per-session 50
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import random
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# Allow ``uv run <path>`` (no package context) to import omnigent + siblings.
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
|
||
|
||
from omnigent.db.utils import _get_head_db_revision, generate_agent_id
|
||
from omnigent.entities import MessageData, NewConversationItem
|
||
from omnigent.server.auth import LEVEL_OWNER, RESERVED_USER_LOCAL
|
||
from omnigent.stores.conversation_store.sqlalchemy_store import SqlAlchemyConversationStore
|
||
from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore
|
||
|
||
# Label key stamped on the first seeded session recording the corpus config, so
|
||
# a later run can detect an existing (and matching) seed and skip re-seeding.
|
||
_SEED_META_LABEL = "omni_bench_seed"
|
||
|
||
# Fixed identifiers so the corpus is byte-stable across runs at a given config.
|
||
_AGENT_NAME = "bench-agent"
|
||
_DEFAULT_SESSIONS = 5000
|
||
_DEFAULT_ITEMS = 50
|
||
_DEFAULT_RNG_SEED = 1234
|
||
|
||
# A pool of realistic-ish message fragments; the RNG assembles item text from
|
||
# these so search_text has lexical variety without external data.
|
||
_FRAGMENTS = (
|
||
"investigate the failing migration",
|
||
"the runner keeps disconnecting under load",
|
||
"add pagination to the sessions endpoint",
|
||
"why does the policy classifier time out",
|
||
"refactor the conversation store append path",
|
||
"benchmark the list endpoints against postgres",
|
||
"the web UI drops the last streamed token",
|
||
"trace the tunnel handshake for this runner id",
|
||
"summarize the changes in this pull request",
|
||
"reproduce the elicitation race on reconnect",
|
||
)
|
||
|
||
|
||
def _meta_value(sessions: int, items_per_session: int, rng_seed: int, head: str) -> str:
|
||
"""Serialize the corpus config into the seed-marker label value.
|
||
|
||
Includes the Alembic *head* read at seed time, so a corpus seeded under an
|
||
older schema auto-mismatches the current head and is reseeded — no
|
||
hand-maintained revision constant.
|
||
"""
|
||
return f"sessions={sessions};items={items_per_session};rng={rng_seed};rev={head}"
|
||
|
||
|
||
def _existing_seed_meta(conv: SqlAlchemyConversationStore) -> str | None:
|
||
"""Return the seed-marker label value if a bench corpus already exists.
|
||
|
||
Looks up the most recent ``bench-agent`` session and reads its
|
||
``omni_bench_seed`` label. ``None`` means no (recognizable) seed present.
|
||
"""
|
||
listing = conv.list_conversations(limit=1, agent_name=_AGENT_NAME)
|
||
if not listing.data:
|
||
return None
|
||
marked = conv.get_conversation(listing.data[0].id)
|
||
return marked.labels.get(_SEED_META_LABEL) if marked is not None else None
|
||
|
||
|
||
def _make_items(rng: random.Random, count: int) -> list[NewConversationItem]:
|
||
"""Build *count* deterministic user-role message items.
|
||
|
||
User-role only: assistant messages require an ``agent`` field the store
|
||
only assigns after a real turn, and the seeded read path is role-agnostic.
|
||
"""
|
||
items: list[NewConversationItem] = []
|
||
for i in range(count):
|
||
text = f"{rng.choice(_FRAGMENTS)} (item {i})"
|
||
items.append(
|
||
NewConversationItem(
|
||
type="message",
|
||
response_id=f"resp_seed_{i}",
|
||
data=MessageData(role="user", content=[{"type": "input_text", "text": text}]),
|
||
)
|
||
)
|
||
return items
|
||
|
||
|
||
def seed(
|
||
db_uri: str,
|
||
*,
|
||
sessions: int = _DEFAULT_SESSIONS,
|
||
items_per_session: int = _DEFAULT_ITEMS,
|
||
rng_seed: int = _DEFAULT_RNG_SEED,
|
||
reseed: bool = False,
|
||
) -> int:
|
||
"""Seed *sessions* sessions × *items_per_session* items into *db_uri*.
|
||
|
||
Idempotent: if a matching seed already exists (same config + schema
|
||
revision) it is left untouched unless *reseed* is set. Constructing the
|
||
store runs migrations to head on first init, so *db_uri* need not
|
||
pre-exist.
|
||
|
||
:param db_uri: SQLAlchemy URI the server will also boot against, e.g.
|
||
``"sqlite:///abs/bench.db"`` or ``"postgresql+psycopg://…"``.
|
||
:param sessions: Number of listable sessions to create.
|
||
:param items_per_session: Conversation items appended to each session.
|
||
:param rng_seed: Seed for the deterministic text RNG.
|
||
:param reseed: Seed even when a matching corpus is already present.
|
||
:returns: The number of sessions created (0 when a matching seed is reused).
|
||
"""
|
||
conv = SqlAlchemyConversationStore(db_uri)
|
||
perms = SqlAlchemyPermissionStore(db_uri)
|
||
|
||
# Read the current schema head at runtime (no DB contacted) and fold it into
|
||
# the reuse marker, so a corpus from an older schema is auto-reseeded.
|
||
head = _get_head_db_revision("sqlite:///:memory:")
|
||
want = _meta_value(sessions, items_per_session, rng_seed, head)
|
||
if not reseed:
|
||
existing = _existing_seed_meta(conv)
|
||
if existing == want:
|
||
print(f"seed: matching corpus already present ({want}); skipping")
|
||
return 0
|
||
if existing is not None:
|
||
print(f"seed: existing corpus differs ({existing!r} != {want!r}); pass --reseed")
|
||
return 0
|
||
|
||
perms.ensure_user(RESERVED_USER_LOCAL)
|
||
rng = random.Random(rng_seed)
|
||
|
||
last_sid = ""
|
||
for s in range(sessions):
|
||
created = conv.create_session_with_agent(
|
||
agent_id=generate_agent_id(),
|
||
agent_name=_AGENT_NAME,
|
||
agent_bundle_location="bench/seed", # never validated on the read path
|
||
agent_description=None,
|
||
title=f"bench session {s}: {rng.choice(_FRAGMENTS)}",
|
||
)
|
||
sid = created.conversation.id
|
||
last_sid = sid
|
||
perms.grant(RESERVED_USER_LOCAL, sid, LEVEL_OWNER)
|
||
if items_per_session:
|
||
conv.append(sid, _make_items(rng, items_per_session))
|
||
if sessions >= 100 and s % (sessions // 10) == 0 and s:
|
||
print(f"seed: {s}/{sessions} sessions")
|
||
|
||
# Stamp the corpus config on the LAST (newest) session — that's the one
|
||
# ``_existing_seed_meta``'s default desc listing returns, so the reuse
|
||
# check finds it regardless of corpus size.
|
||
if last_sid:
|
||
conv.set_labels(last_sid, {_SEED_META_LABEL: want})
|
||
|
||
print(f"seed: created {sessions} sessions × {items_per_session} items ({want})")
|
||
return sessions
|
||
|
||
|
||
def _parse_args(argv: list[str]) -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
prog="omnigent-benchmark-seed",
|
||
description=__doc__,
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
)
|
||
parser.add_argument(
|
||
"--database-uri",
|
||
metavar="URI",
|
||
help="DB to seed. Required unless --print-head.",
|
||
)
|
||
parser.add_argument("--sessions", type=int, default=_DEFAULT_SESSIONS, metavar="N")
|
||
parser.add_argument("--items-per-session", type=int, default=_DEFAULT_ITEMS, metavar="N")
|
||
parser.add_argument("--rng-seed", type=int, default=_DEFAULT_RNG_SEED, metavar="N")
|
||
parser.add_argument(
|
||
"--reseed",
|
||
action="store_true",
|
||
help="Seed even if a matching corpus is already present.",
|
||
)
|
||
parser.add_argument(
|
||
"--print-head",
|
||
action="store_true",
|
||
help="Print the repo's Alembic head revision and exit (drift-check helper).",
|
||
)
|
||
return parser.parse_args(argv)
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> int:
|
||
args = _parse_args(argv if argv is not None else sys.argv[1:])
|
||
if args.print_head:
|
||
print(_get_head_db_revision("sqlite:///:memory:"))
|
||
return 0
|
||
if not args.database_uri:
|
||
print("seed: --database-uri is required (unless --print-head)", file=sys.stderr)
|
||
return 2
|
||
seed(
|
||
args.database_uri,
|
||
sessions=args.sessions,
|
||
items_per_session=args.items_per_session,
|
||
rng_seed=args.rng_seed,
|
||
reseed=args.reseed,
|
||
)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|