check-bge-knowledge.py 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. """Read-only migration checks. Output only counts/assertions, never source bodies."""
  2. from collections import Counter
  3. import json
  4. from pathlib import Path
  5. import sys
  6. ROOT = Path(__file__).resolve().parents[1]
  7. sys.path.insert(0, str(ROOT))
  8. from openpyxl import load_workbook
  9. from src.step2_knowledge_ingest.ingest import load
  10. from src.step3_question_answer.retrieval import Retriever
  11. def main():
  12. prepared = ROOT / 'knowledge/prepared'
  13. records = load(prepared)
  14. excel = [r for r in records if r['raw_record'].get('workbook_source')]
  15. counts = Counter(r['raw_record']['workbook_source']['sheet'] for r in excel)
  16. assert sorted(counts.values()) == [11, 16, 26]
  17. workbook = load_workbook(ROOT / excel[0]['source']['file'], data_only=False)
  18. try:
  19. for sheet_name in counts:
  20. sheet = workbook[sheet_name]
  21. represented = {}
  22. for record in excel:
  23. source = record['raw_record']['workbook_source']
  24. if source['sheet'] == sheet_name:
  25. for row in source['rows']:
  26. represented.update(row['cells'])
  27. expected = {cell.coordinate: cell.value for row in sheet.iter_rows(min_row=4)
  28. for cell in row if cell.value is not None}
  29. assert all(represented.get(k) == v for k, v in expected.items())
  30. finally:
  31. workbook.close()
  32. retriever = Retriever(ROOT / '.env', prepared / 'ingestion/manifest.json')
  33. checks = []
  34. try:
  35. selected = [next(r for r in excel if r['raw_record']['workbook_source']['sheet'] == sheet) for sheet in counts]
  36. selected += [next(r for r in records if r['kind'] == kind) for kind in ('policy', 'finance', 'qa')]
  37. for record in selected:
  38. query = record['data']['question'] if record['kind'] == 'qa' else record['title']
  39. result = retriever.search(query)
  40. checks.append({'complete': result['complete'], 'expected_record_found': any(
  41. candidate['record']['id'] == record['id'] for candidate in result['candidates'])})
  42. assert all(c['complete'] and c['expected_record_found'] for c in checks)
  43. finally:
  44. retriever.close()
  45. report = {'records': len(records), 'workbook_counts': dict(counts), 'nonempty_cells_preserved': True,
  46. 'retrieval_checks': checks, 'scope': 'exact-title retrieval smoke; not recall-quality evaluation'}
  47. path = ROOT / 'output/checks/bge-knowledge.json'
  48. path.parent.mkdir(parents=True, exist_ok=True)
  49. path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + '\n', encoding='utf-8')
  50. print(json.dumps(report, ensure_ascii=False))
  51. if __name__ == '__main__':
  52. main()