File size: 2,013 Bytes
f109bf8 | 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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | #!/usr/bin/env python3
"""Harness: embed COBOL balance query, ANN top-K against agent_registry."""
from __future__ import annotations
import math
import sys
from pathlib import Path
# reuse embed from loader
sys.path.insert(0, str(Path(__file__).resolve().parent))
from load_agent_registry import embed, fmt_vector
import psycopg2
QUERY = (
"Explain how this COBOL payment program calculates the account balance "
"and determine whether its Java replacement preserves the behavior. "
"cobol copybook business_rule_extraction control_flow data_flow "
"legacy_mainframe modernization_analysis"
)
TOP_K = 5
def main() -> int:
dsn = sys.argv[1] if len(sys.argv) > 1 else "dbname=david user=david host=127.0.0.1"
qvec = fmt_vector(embed(QUERY))
conn = psycopg2.connect(dsn)
with conn.cursor() as cur:
cur.execute(
"""
SELECT agent_id, domain, risk_class,
1 - (embedding <=> %s::vector) AS cosine_sim
FROM agent_registry
WHERE embedding IS NOT NULL
ORDER BY embedding <=> %s::vector
LIMIT %s
""",
(qvec, qvec, TOP_K),
)
rows = cur.fetchall()
conn.close()
print(f"query={QUERY[:80]}...")
print(f"top_k={TOP_K}")
for i, (aid, domain, risk, sim) in enumerate(rows, 1):
print(f"{i}\t{aid}\t{domain}\t{risk}\tsim={sim:.6f}")
ids = [r[0] for r in rows]
ok_legacy = "LegacyCobol" in ids
# Quant should not outrank LegacyCobol (prefer: Legacy ahead of Quant if both present)
if "Quant" in ids and "LegacyCobol" in ids:
ok_order = ids.index("LegacyCobol") < ids.index("Quant")
else:
ok_order = "Quant" not in ids[:1] # Quant not the preferred top hit
print(f"legacy_in_topk={ok_legacy}")
print(f"quant_not_preferred={ok_order}")
print(f"pass={ok_legacy and ok_order}")
return 0 if (ok_legacy and ok_order) else 1
if __name__ == "__main__":
raise SystemExit(main())
|