scripts/build_index.py

"""Build the Chroma index the triage agent retrieves from.

Two collections:
  - policy rules  (one citable chunk per rule, from data/policy.md)
  - past cases    (precedent adjudications, from data/past_cases.json)

Run once locally; commit data/chroma/ for the Hugging Face Space (Spaces can't ingest).

    python -m scripts.build_index
"""
from __future__ import annotations

import json
from pathlib import Path

from src.config import CONFIG
from src.policy import parse_policy_rules
from src.store import count, index_documents


def build_policy() -> int:
    rules = parse_policy_rules()
    if not rules:
        print("WARNING: no policy rules parsed — check data/policy.md formatting.")
        return 0
    index_documents(
        CONFIG.policy_collection,
        ids=[r.rule_id for r in rules],
        texts=[r.text for r in rules],  # r.text already begins with the rule ID + name
        metadatas=[{"rule_id": r.rule_id, "section": r.section} for r in rules],
    )
    return len(rules)


def build_cases() -> int:
    path = Path("data/past_cases.json")
    if not path.exists():
        return 0
    cases = json.loads(path.read_text(encoding="utf-8"))
    index_documents(
        CONFIG.cases_collection,
        ids=[c["case_id"] for c in cases],
        texts=[
            f"{c['title']}. {c['summary']} Outcome: {c['outcome']}"
            + (f" (rules: {', '.join(c['rule_ids'])})." if c.get("rule_ids") else ".")
            for c in cases
        ],
        metadatas=[{"case_id": c["case_id"], "outcome": c["outcome"]} for c in cases],
    )
    return len(cases)


def main() -> None:
    nr = build_policy()
    nc = build_cases()
    print(f"Indexed {nr} policy rules  -> '{CONFIG.policy_collection}' (total {count(CONFIG.policy_collection)})")
    print(f"Indexed {nc} past cases    -> '{CONFIG.cases_collection}' (total {count(CONFIG.cases_collection)})")
    print("Done.")


if __name__ == "__main__":
    main()