"""Moonley agentic TOOL REGISTRY. Loads the corpus once (lean — no BM25) and exposes the tools the ReAct controller can call. Each tool returns a list of compact case dicts (doc_id + the fields the LLM needs to reason) or a small structured result. Most tools are ports of serve.py primitives; the statute tools are new (see themis-statute-layer). keyword_search(BM25) is intentionally omitted here — it is 68s/query on the Mac; it runs on the GPU box in production. Usage: from tools import Corpus ; C = Corpus(DATA, STATUTE_DIR) ; C.vector_search("...", k=8) """ import json, os, re, difflib import numpy as np from collections import Counter, defaultdict from sentence_transformers import SentenceTransformer, CrossEncoder from case_summary import case_summary_record, load_case_summaries from statute_crosswalk import load_default_crosswalk from statute_library import ExactStatuteLibrary BGE_Q = "Represent this sentence for searching relevant passages: " _NAME_STOP = {"v","vs","of","and","the","ors","anr","etc","state","union","govt","government","in","re", "others","another","ltd","co","pvt","dead","thr","lrs","alias","through","etc","anrs","ms", "shri","smt","sri","mr","mrs","dr","justice","sh","kum","mohd"} BAD_STATUS = {"overruled", "per_incuriam", "doubted"} def _clean(s): if not s: return "" return re.sub(r"\s+", " ", s).strip() def _ntok(s): """Tokenize a case name, MERGING runs of single letters so abbreviations match: U.P.->up, A.K.->ak.""" out = []; buf = "" for t in re.findall(r"[a-z]+", (s or "").lower()): if len(t) == 1: buf += t else: if buf: out.append(buf); buf = "" out.append(t) if buf: out.append(buf) return out class Corpus: def __init__(self, data_dir, statute_dir, device="cpu"): self.device = device print("[tools] loading corpus ...", flush=True) self.texts = []; self.chunk_doc = [] with open(os.path.join(data_dir, "escr_chunks.jsonl"), encoding="utf-8") as f: for l in f: c = json.loads(l); self.texts.append(c["text"]); self.chunk_doc.append(c["doc_id"]) self.M = np.load(os.path.join(data_dir, "escr_vectors.npy"), mmap_mode="r") # mmap: ~2GB off resident RSS (pages fault in on the M@qv scan) self.chunk_doc_arr = np.array(self.chunk_doc) self.doc_chunks = defaultdict(list) for i, d in enumerate(self.chunk_doc): self.doc_chunks[d].append(i) self.meta = {} for l in open(os.path.join(data_dir, "escr_meta.jsonl"), encoding="utf-8"): m = json.loads(l); self.meta[m["doc_id"]] = m # A judgment is eligible for research only when both its metadata and # source-derived text are present in the active release. Metadata-only # rows and graph stubs may remain useful for audit work, but they must # never become search results, recommendations, or chat evidence. self.eligible_doc_ids = { d for d, cis in self.doc_chunks.items() if d in self.meta and any(_clean(self.texts[i]) for i in cis) } self.goodlaw = {} # good_law_v2 (classify_treatments.py rollup — real statuses) wins over the legacy file _gl = "good_law_v2.jsonl" if os.path.exists(os.path.join(data_dir, "good_law_v2.jsonl")) else "good_law.jsonl" for l in open(os.path.join(data_dir, _gl), encoding="utf-8"): g = json.loads(l); self.goodlaw[g["doc_id"]] = g # identity ledger (build_ledger.py): decision-year fix, bench ints, sibling canonicals self.decision_year = {}; self.bench_n = {}; self.canonical = set(); self.cluster_of = {} _lp = os.path.join(data_dir, "corpus_ledger.jsonl") if os.path.exists(_lp): for l in open(_lp, encoding="utf-8"): r = json.loads(l); d = r["doc_id"] if r.get("decision_year"): self.decision_year[d] = r["decision_year"] self.bench_n[d] = r.get("bench_n", 0) if r.get("canonical", True): self.canonical.add(d) if r.get("cluster_id"): self.cluster_of[d] = r["cluster_id"] print(f"[tools] ledger: {len(self.decision_year)} decision-years, " f"{len(self.cluster_of)} sibling-clustered docs", flush=True) # famous-name aliases mined from citing sentences (build_citation_graph.py) self.aliases = {} _ap = os.path.join(data_dir, "case_aliases.json") if os.path.exists(_ap): self.aliases = {k.lower(): v for k, v in json.load(open(_ap, encoding="utf-8")).items()} print(f"[tools] aliases: {len(self.aliases)}", flush=True) self.in_edges = defaultdict(list); self.out_edges = defaultdict(list); self.edge_meta = {} self.cite_indeg = Counter() # edges_v2 (body-text parse, ~25x the legacy graph) wins over the headnote-only file _ep = "edges_v2.jsonl" if os.path.exists(os.path.join(data_dir, "edges_v2.jsonl")) else "edges.jsonl" for l in open(os.path.join(data_dir, _ep), encoding="utf-8"): e = json.loads(l); f, t = e["from"], e["target"] self.out_edges[f].append(t); self.in_edges[t].append(f) self.edge_meta[(f, t)] = {"treatment": e.get("treatment"), "method": e.get("method")} if e.get("method") in ("cite", "body", "headnote"): self.cite_indeg[t] += 1 if _ep == "edges_v2.jsonl": print(f"[tools] edges_v2: {sum(len(v) for v in self.out_edges.values())} edges", flush=True) # treatment overrides from the classifier (finer than the rollup) _tp = os.path.join(data_dir, "edges_treatment.jsonl") if os.path.exists(_tp): n = 0 for l in open(_tp, encoding="utf-8"): r = json.loads(l); k = (r["from"], r["target"]) if k in self.edge_meta: self.edge_meta[k]["treatment"] = r["treatment"]; n += 1 print(f"[tools] treatments: {n} classified edges", flush=True) # synthetic headnotes (backfill_headnotes.py) — fill the 70s-80s crater for skim/judge/view self.syn_held = {} _sp = os.path.join(data_dir, "synthetic_headnotes.jsonl") if os.path.exists(_sp): for l in open(_sp, encoding="utf-8"): r = json.loads(l) if r.get("held"): self.syn_held[r["doc_id"]] = r print(f"[tools] synthetic headnotes: {len(self.syn_held)}", flush=True) # Extraction-time case summaries. This sidecar is the stable hand-off for the future # Indian Kanoon re-extraction; reporter/synthetic headnotes remain honest interim fallbacks. self.case_summaries, _summary_file = load_case_summaries(data_dir) if _summary_file: print(f"[tools] case summaries: {len(self.case_summaries)} from {_summary_file}", flush=True) self.name_vocab = set(); self.name_postings = defaultdict(set) # token -> doc_ids (fast name lookup) self.nc2doc = {}; self.cite_resolver = {} # exact citation lookup (known-item route) for d, m in self.meta.items(): for w in _ntok(m.get("case_name") or ""): # _ntok merges U.P.->up so abbreviations match if len(w) >= 4: self.name_vocab.add(w) if len(w) > 1: self.name_postings[w].add(d) if m.get("neutral_citation"): self.nc2doc[m["neutral_citation"]] = d for k in [m.get("neutral_citation")] + (m.get("equivalent_citations") or []): if k: self.cite_resolver.setdefault(re.sub(r"\s+", " ", k.replace(".", "")).strip().upper(), d) # statute layer self.statute_idx = json.load(open(os.path.join(statute_dir, "statute_index.json"))) self.statute_V = np.load(os.path.join(statute_dir, "statute_vectors.npy")) _statute_records = json.load(open(os.path.join(statute_dir, "all_statutes.json"))) self.statute_texts = [s.get("retrieval_text", "") for s in _statute_records] self.statute_library = ExactStatuteLibrary.from_env( fallback_path=os.path.join(statute_dir, "all_statutes.json") ) self.concord = json.load(open(os.path.join(statute_dir, "concordance.json"))) self.crosswalk = load_default_crosswalk(os.environ.get("THEMIS_SECTION_CROSSWALK", "").strip() or None) # doc-level HELD-headnote vectors (optional; the $0 representation arm — cleaner signal than OCR chunks) self.held_V = None hv = os.path.join(data_dir, "held_vectors.npy") if os.path.exists(hv): self.held_V = np.load(hv) self.held_docs = json.load(open(os.path.join(data_dir, "held_docids.json"))) print(f"[tools] HELD vectors: {self.held_V.shape[0]}", flush=True) # citation-context vectors (optional; how LATER courts describe each precedent — clean, # modern, doctrine-level; exists precisely for old landmarks whose own text is OCR-noisy) self.citectx_V = None cv = os.path.join(data_dir, "citectx_vectors.npy") if os.path.exists(cv): self.citectx_V = np.load(cv, mmap_mode="r") self.citectx_docs = json.load(open(os.path.join(data_dir, "citectx_docids.json"))) print(f"[tools] CITECTX vectors: {self.citectx_V.shape[0]}", flush=True) self.st = SentenceTransformer("BAAI/bge-small-en-v1.5", device=device) self.ce = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", device=device) # KEYWORD arm. Preferred: prebuilt disk-backed FTS5 chunk index (escr_fts.sqlite, built # offline by build_fts_index.py) — ms queries, ~0 resident RAM, no boot cost. Fallback: # the in-RAM doc-level BM25 build (THEMIS_KEYWORD=1). THEMIS_FTS=0 opts out of FTS5. self.kw = None; self.fts = None _fts_path = os.path.join(data_dir, "escr_fts.sqlite") if os.environ.get("THEMIS_FTS", "1") == "1" and os.path.exists(_fts_path): import sqlite3, threading self.fts = sqlite3.connect(f"file:{_fts_path}?mode=ro", uri=True, check_same_thread=False) self.fts_lock = threading.Lock() # lanes hit keyword_search from parallel threads print(f"[tools] FTS5 keyword index: {os.path.getsize(_fts_path)/1e9:.2f} GB (disk-backed)", flush=True) if self.fts is None and os.environ.get("THEMIS_KEYWORD", "1") == "1": import math, time as _t t0 = _t.time(); print("[tools] building keyword index ...", flush=True) self.kw_docs = list(self.doc_chunks.keys()) self.kw_postings = defaultdict(list); df = Counter() self.kw_dl = np.zeros(len(self.kw_docs), dtype=np.float32) for i, d in enumerate(self.kw_docs): toks = [] for ci in self.doc_chunks[d]: toks += re.findall(r"[a-z0-9]+", self.texts[ci].lower()) tf = Counter(toks); self.kw_dl[i] = len(toks) for t, c in tf.items(): self.kw_postings[t].append((i, c)); df[t] += 1 N = len(self.kw_docs); self.kw_avgdl = float(self.kw_dl.mean()) or 1.0 self.kw_idf = {t: math.log(1 + (N - n + 0.5) / (n + 0.5)) for t, n in df.items()} self.kw = True print(f"[tools] keyword index: {len(self.kw_postings)} terms, {N} docs, {_t.time()-t0:.0f}s", flush=True) print( f"[tools] ready — {len(self.eligible_doc_ids)} source-grounded judgments " f"({len(self.meta) - len(self.eligible_doc_ids)} metadata-only excluded), " f"{len(self.statute_idx)} statute sections", flush=True, ) # ---- helpers ---- def _enc(self, q): return self.st.encode(BGE_Q + q, normalize_embeddings=True, convert_to_numpy=True).astype(np.float32) def is_retrieval_eligible(self, doc_id): """True only for judgments whose source text is in the active corpus.""" return str(doc_id) in self.eligible_doc_ids def coverage(self): return { "accepted_judgments": len(self.eligible_doc_ids), "metadata_only_excluded": len(self.meta) - len(self.eligible_doc_ids), "scope": "Supreme Court of India judgments stored in this release", } def _card(self, d, rr=0.0): m = self.meta.get(d, {}); gl = self.goodlaw.get(d, {}) cis = self.doc_chunks.get(d, []) snip = _clean((m.get("held") or m.get("issue") or (self.syn_held.get(d) or {}).get("held") or (self.texts[cis[0]] if cis else "")))[:240] return {"doc_id": d, "judgment_id": str(d), "case_name": m.get("case_name"), "year": self.decision_year.get(d) or m.get("year") or m.get("date"), "neutral_citation": m.get("neutral_citation"), "bench_strength": m.get("bench_strength"), "cited_by": self.cite_indeg.get(d, 0), "good_law": gl.get("good_law_status", "unknown"), "rr": round(float(rr), 2), "snippet": snip} def _dense_pool(self, qv, n=120): sim = self.M @ qv requested = min(max(n * 2, n + 1), len(sim) - 1) top = np.argpartition(-sim, requested)[:requested] return [ int(i) for i in top[np.argsort(-sim[top])] if self.is_retrieval_eligible(self.chunk_doc[int(i)]) ][:n] def _rerank_docs(self, q, cand_idx, topk): cand_idx = [ ci for ci in cand_idx if self.is_retrieval_eligible(self.chunk_doc[ci]) and _clean(self.texts[ci]) ] if not cand_idx: return [] rr = self.ce.predict([(q, self.texts[ci]) for ci in cand_idx]) best = {} for ci, s in zip(cand_idx, rr): d = self.chunk_doc[ci] if not self.is_retrieval_eligible(d): continue if d not in best or s > best[d]: best[d] = float(s) ranked = sorted(best.items(), key=lambda x: -x[1])[:topk] return [self._card(d, s) for d, s in ranked] # ---- TOOLS ---- def vector_search(self, q, k=8): """Semantic retrieval — doctrine described in the user's words. dense pool -> cross-encoder.""" return self._rerank_docs(q, self._dense_pool(self._enc(q), 120), k) def authority_search(self, q, k=8, alpha=0.3): """Retrieve, then rank by AUTHORITY (cross-encoder + alpha*log1p(cite_indeg)) — 'the leading case on X'.""" cand = self._dense_pool(self._enc(q), 120) rr = self.ce.predict([(q, self.texts[ci]) for ci in cand]); best = {} for ci, s in zip(cand, rr): d = self.chunk_doc[ci] if not self.is_retrieval_eligible(d): continue if d not in best or s > best[d]: best[d] = float(s) sig = lambda x: 1/(1+np.exp(-x)) scored = sorted(best.items(), key=lambda x: -(sig(x[1]) + alpha*np.log1p(self.cite_indeg.get(x[0], 0))))[:k] return [self._card(d, s) for d, s in scored] def keyword_search(self, q, k=12, k1=1.5, b=0.75): """BM25 keyword retrieval — exact terms / names / section nums that dense misses. FTS5 path: chunk-level match, best-chunk-per-doc aggregation (a doc with one strong exact-term chunk ranks high), OR semantics to mirror the legacy scorer.""" if self.fts is not None: _stop = {"of","the","and","or","in","to","a","an","is","for","on","by","at", "with","under","was","were","be","has","had","it","that","this"} toks = [t for t in re.findall(r"[a-z0-9]+", q.lower()) if t not in _stop] if not toks: return [] match = " OR ".join(f'"{t}"' for t in toks[:24]) # quoted: immune to FTS syntax chars with self.fts_lock: rows = self.fts.execute( "SELECT rowid, bm25(fts) FROM fts WHERE fts MATCH ? ORDER BY bm25(fts) LIMIT ?", (match, max(k * 12, 240))).fetchall() best = {} for ci, s in rows: # bm25(): smaller = better d = self.chunk_doc[ci] if not self.is_retrieval_eligible(d): continue if d not in best or s < best[d]: best[d] = s top = sorted(best.items(), key=lambda x: x[1])[:k] return [self._card(d) for d, _ in top] if not self.kw: return [] scores = defaultdict(float) for t in set(re.findall(r"[a-z0-9]+", q.lower())): idf = self.kw_idf.get(t) if not idf: continue for i, tf in self.kw_postings[t]: scores[i] += idf * (tf * (k1 + 1)) / (tf + k1 * (1 - b + b * self.kw_dl[i] / self.kw_avgdl)) top = [ item for item in sorted(scores.items(), key=lambda x: -x[1]) if self.is_retrieval_eligible(self.kw_docs[item[0]]) ][:k] return [self._card(self.kw_docs[i]) for i, _ in top] def dense_docs(self, q, k=60): """Doc-level dense ranking (first chunk-hit per doc).""" qv = self._enc(q); sim = self.M @ qv top = np.argpartition(-sim, 2500)[:2500]; top = top[np.argsort(-sim[top])] seen = []; s = set() for ci in top: d = self.chunk_doc[ci] if not self.is_retrieval_eligible(d): continue if d not in s: s.add(d); seen.append(d) if len(seen) >= k: break return seen def held_search(self, q, k=12): """Rank judgments by HELD-headnote similarity (doc-level, clean reporter language).""" if self.held_V is None: return [] qv = self._enc(q); sim = self.held_V @ qv top = np.argpartition(-sim, min(k, len(sim) - 1))[:k] return [ self.held_docs[int(i)] for i in top[np.argsort(-sim[top])] if self.is_retrieval_eligible(self.held_docs[int(i)]) ][:k] def citectx_search(self, q, k=12): """Rank judgments by how LATER courts describe them (citation-context vectors). Multiple contexts per doc -> dedupe keeping best rank.""" if self.citectx_V is None: return [] qv = self._enc(q); sim = np.asarray(self.citectx_V @ qv) n = min(k * 6, len(sim) - 1) top = np.argpartition(-sim, n)[:n] out, seen = [], set() for i in top[np.argsort(-sim[top])]: d = self.citectx_docs[int(i)] if not self.is_retrieval_eligible(d): continue if d not in seen: seen.add(d); out.append(d) if len(out) >= k: break return out def hybrid_search(self, q, k=8, pool=60): """The strong base primitive (panel + recall ablation: RRF@100=0.96): dense + BM25 -> RRF -> rerank.""" dd = self.dense_docs(q, pool) kd = [c["doc_id"] for c in self.keyword_search(q, pool)] sc = {} for r in (dd, kd): for rank, d in enumerate(r): sc[d] = sc.get(d, 0.0) + 1.0 / (60 + rank + 1) fused = [d for d, _ in sorted(sc.items(), key=lambda x: -x[1])][:max(40, k * 4)] rr = self.score_docs(q, fused) for d in fused: if d not in rr: rr[d] = -9.0 ranked = sorted(fused, key=lambda d: -rr[d]) return [self._card(d, rr.get(d, 0.0)) for d in ranked[:k]] def statute_search(self, q, k=3): """Find the statute SECTION(S) most relevant to the query (BNS/IPC/CrPC/IEA/...).""" qv = self._enc(q); sim = self.statute_V @ qv out = [] for j in np.argsort(-sim)[:k]: s = self.statute_idx[int(j)] out.append({"act": s.get("act_short"), "section": s.get("section_number"), "title": s.get("title"), "i": int(j)}) return out def cases_on_section(self, act_section_text, k=8): """Cases discussing a statute section: embed the section text, retrieve nearest judgments.""" qv = self.st.encode(act_section_text, normalize_embeddings=True, convert_to_numpy=True).astype(np.float32) return self._rerank_docs(act_section_text[:300], self._dense_pool(qv, 120), k) def statute_crosswalk(self, code, section): """Map a section across the new/old codes (BNS<->IPC, BNSS<->CrPC, BSA<->IEA).""" return self.crosswalk.lookup(code, section) def statute_provision(self, code, section): return self.statute_library.lookup(code, section) def encode_documents(self, texts): values = [_clean(value) for value in texts if _clean(value)] if not values: return np.empty((0, int(self.M.shape[1])), dtype=np.float32) return np.asarray( self.st.encode(values, normalize_embeddings=True, convert_to_numpy=True), dtype=np.float32, ) def find_similar_cases(self, doc_id, k=8): """'More like this' — nearest judgments to doc_id by embedding centroid.""" cis = self.doc_chunks.get(doc_id, []) if not cis: return [] centroid = self.M[cis].mean(0); centroid /= (np.linalg.norm(centroid) + 1e-9) pool = self._dense_pool(centroid.astype(np.float32), 60) seen = set([doc_id]); out = [] for ci in pool: d = self.chunk_doc[ci] if self.is_retrieval_eligible(d) and d not in seen: seen.add(d); out.append(self._card(d)) if len(out) >= k: break return out def cited_authorities(self, doc_id, k=12): """Note-UP: the cases doc_id relies on (its authority chain).""" return [ self._card(d) for d in list(dict.fromkeys(self.out_edges.get(doc_id, []))) if self.is_retrieval_eligible(d) ][:k] def progeny(self, doc_id, k=12): """Note-DOWN: the cases that cite doc_id (its progeny + treatment).""" out = [] for d in list(dict.fromkeys(self.in_edges.get(doc_id, [])))[:k]: if not self.is_retrieval_eligible(d): continue c = self._card(d); c["treatment"] = self.edge_meta.get((d, doc_id), {}).get("treatment") out.append(c) return out def co_cited_cases(self, doc_id, k=8): """Cases similar by SHARED AUTHORITIES (bibliographic coupling) — cases that cite what doc_id cites.""" mine = set(self.out_edges.get(doc_id, [])) if not mine: return [] score = Counter() for t in mine: for citer in self.in_edges.get(t, []): if citer != doc_id and self.is_retrieval_eligible(citer): score[citer] += 1 return [self._card(d) for d, _ in score.most_common(k)] def good_law_check(self, doc_id): """Citator: current status + treatment breakdown + the overruling case if any.""" gl = self.goodlaw.get(doc_id, {}) status = gl.get("good_law_status", "unknown") overruled_by = None if status in BAD_STATUS: for s, t in [(s, t) for (s, t) in self.edge_meta if t == doc_id]: if self.edge_meta[(s, t)].get("treatment") in ("overruled", "overrules"): if self.is_retrieval_eligible(s): overruled_by = self._card(s); break return {"doc_id": doc_id, "good_law": status, "treatment_breakdown": gl.get("treatment_breakdown", {}), "overruled_by": overruled_by} def metadata_filter(self, cards, min_bench=None, year_from=None, year_to=None): """Filter a candidate list by bench strength (Constitution Bench = 5+), date range.""" out = [] _BN = {"single": 1, "division": 2, "full": 3, "constitution": 5, "larger": 7} for c in cards: d = c["doc_id"]; m = self.meta.get(d, {}) bs = self.bench_n.get(d) or _BN.get(str(m.get("bench_strength") or "").lower(), 0) yr = self.decision_year.get(d) or 0 if not yr: try: yr = int(str(m.get("year") or 0)[:4]) except Exception: yr = 0 if min_bench and bs < min_bench: continue if year_from and yr and yr < year_from: continue if year_to and yr and yr > year_to: continue out.append(c) return out def read_case(self, doc_id): """Read a case's headnote/held/issue (for the agent to verify relevance + for grounding).""" if not self.is_retrieval_eligible(doc_id): return {} m = self.meta.get(doc_id, {}) held = _clean(m.get("held")); issue = _clean(m.get("issue")) if not held and doc_id in self.syn_held: # LLM-backfilled reporter-style headnote s = self.syn_held[doc_id] held = _clean(s.get("held")); issue = issue or _clean(s.get("issue")) if not held and not issue: # older cases lack extracted headnotes -> fall back to first chunks cis = self.doc_chunks.get(doc_id, []) held = _clean(" ".join(self.texts[i] for i in cis[:2])) return {"doc_id": doc_id, "case_name": m.get("case_name"), "neutral_citation": m.get("neutral_citation"), "bench_strength": m.get("bench_strength"), "good_law": self.goodlaw.get(doc_id, {}).get("good_law_status", "unknown"), "issue": issue[:1200], "held": held[:1800]} def score_docs(self, q, doc_ids, per_doc=3): """Uniformly cross-encoder-score a heterogeneous pool of docs vs q (max over each doc's first chunks). One batched CE pass. Returns {doc_id: rr}. Lets graph/authority/statute additions be ranked on the same scale as dense hits.""" pairs = []; owner = [] for d in doc_ids: if not self.is_retrieval_eligible(d): continue for ci in self.doc_chunks.get(d, [])[:per_doc]: pairs.append((q, self.texts[ci])); owner.append(d) if not pairs: return {d: 0.0 for d in doc_ids} sc = self.ce.predict(pairs, batch_size=256) best = {d: -9e9 for d in doc_ids} for d, s in zip(owner, sc): if s > best[d]: best[d] = float(s) return {d: (best[d] if best[d] > -9e9 else 0.0) for d in doc_ids} def front_text(self, doc_id, n=1800): """The judgment's FRONT MATTER (reporter headnote lives here in 70-100% of judgments across all decades — more reliable than meta.held, which craters to ~2% in the 1970s-80s).""" held = _clean(self.meta.get(doc_id, {}).get("held") or "") if len(held) > 200: return held[:n] syn = _clean((self.syn_held.get(doc_id) or {}).get("held") or "") if len(syn) > 200: return syn[:n] # backfilled crater doc: skim the ratio, not the cover page cis = self.doc_chunks.get(doc_id, []) return _clean(" ".join(self.texts[i] for i in cis[:3]))[:n] def full_text_for_read(self, q, doc_id, cap_chars=90000): """Layer-2 reading surface: the FULL judgment up to ~cap (≈22k tokens). Above-cap monsters get a tiered pack: HELD + opening + a window around the query's best-matching chunk + the ending — the controlling passage in multi-issue judgments sits mid-text where head/tail packs go blind.""" cis = self.doc_chunks.get(doc_id, []) if not cis: return "" parts = [self.texts[i] for i in cis] full = "\n".join(parts) if len(full) <= cap_chars: return full held = _clean((self.meta.get(doc_id, {}).get("held") or ""))[:6000] probe = cis[:40] # find the query-relevant window (one small CE pass) sc = self.ce.predict([(q, self.texts[i]) for i in probe]) bi = int(np.argmax(sc)) win = "\n".join(self.texts[i] for i in cis[max(0, bi - 2):bi + 3]) head = "\n".join(parts[:8]); tail = "\n".join(parts[-6:]) pack = (("HELD: " + held + "\n\n") if held else "") + head + "\n[...]\n" + win + "\n[...]\n" + tail return pack[:cap_chars] def best_chunk_text(self, q, doc_id, limit=1600): """The doc's single chunk most relevant to q (for the grounding gate to quote from).""" cis = self.doc_chunks.get(doc_id, [])[:6] if not cis: return "" sc = self.ce.predict([(q, self.texts[ci]) for ci in cis]) return _clean(self.texts[cis[int(np.argmax(sc))]])[:limit] def judgment_view(self, doc_id): """Full case view for the pilot UI (metadata + issue/held + good-law + cited-by). issue is in ~3% of metadata, held in ~56% — so fall back to the judgment's opening text when missing.""" if not self.is_retrieval_eligible(doc_id): return {} m = self.meta.get(doc_id, {}); gl = self.goodlaw.get(doc_id, {}) cis = self.doc_chunks.get(doc_id, []) body = _clean(" ".join(self.texts[i] for i in cis[:14])) extracted = self.case_summaries.get(doc_id, {}) issue = _clean(m.get("issue") or extracted.get("issue")) held = _clean(m.get("held") or extracted.get("held")); synthetic = False if not held and doc_id in self.syn_held: # LLM-backfilled headnote (disclosed to the UI) s = self.syn_held[doc_id] held = _clean(s.get("held")); issue = issue or _clean(s.get("issue")); synthetic = bool(held) if not held: held = body[:6000] # no extracted headnote -> show the opening (the headnote lives there) return {"doc_id": doc_id, "synthetic_headnote": synthetic, "summary": case_summary_record(m, self.syn_held.get(doc_id), extracted), "case_name": m.get("case_name"), "neutral_citation": m.get("neutral_citation"), "equivalent_citations": m.get("equivalent_citations"), "court": m.get("court"), "date": m.get("date"), "bench_strength": m.get("bench_strength"), "disposition": m.get("disposition"), "good_law_status": gl.get("good_law_status", "unknown"), "treatment_breakdown": gl.get("treatment_breakdown", {}), "cited_by": self.cite_indeg.get(doc_id, 0), "issue": issue[:4000], "held": held[:6000], "text": body[:60000]} def case_chat_passages(self, q, doc_id, k=4, limit=1800): """Return only source passages from one eligible opened judgment. The legacy bundle has chunk identity rather than schema-v5 paragraph identity, so its stable fallback IDs are disclosed as chunk anchors. The v5 adapter supplies real paragraph IDs through the same contract. """ if not self.is_retrieval_eligible(doc_id): return [] cis = self.doc_chunks.get(doc_id, []) if not cis: return [] scores = self.ce.predict([(q, self.texts[ci]) for ci in cis[:40]]) ranked = sorted(zip(cis[:40], scores), key=lambda item: -float(item[1]))[:k] return [ { "paragraph_id": f"{doc_id}:chunk:{ci}", "label": f"Indexed passage {rank + 1}", "text": _clean(self.texts[ci])[:limit], "source_kind": "legacy_chunk", } for rank, (ci, _) in enumerate(ranked, 1) if _clean(self.texts[ci]) ] def identity_hits(self, q): """Known-item route: a citation or a 'X v Y' case-name query resolves to the EXACT case(s) (cite_indeg salience tiebreak), not semantic search. Restores serve.py's identity routing.""" ql = q.strip() m = re.search(r"\[\d{4}\]\s*\d+\s*S\.?C\.?R\.?\s*\d+|\(\d{4}\)\s*\d+\s*SCC\s*\d+|\d{4}\s+INSC\s+\d+|AIR\s+\d{4}\s+SC\s+\d+", ql, re.I) if m: rid = self.cite_resolver.get(re.sub(r"\s+", " ", m.group(0).replace(".", "")).strip().upper()) or self.nc2doc.get(m.group(0)) if rid and self.is_retrieval_eligible(rid): return [rid], "citation" # famous-name alias ("kesavananda", "the shah bano judgment") — exact or contained phrase if self.aliases and len(ql) <= 60: qa = re.sub(r"[^a-z0-9 ]", " ", ql.lower()) qa = re.sub(r"\b(the|case|judgment|judgement|in|re|of)\b", " ", qa) qa = re.sub(r"\s+", " ", qa).strip() if qa in self.aliases and self.is_retrieval_eligible(self.aliases[qa]): return [self.aliases[qa]], "case name" # substring form ("the kesavananda bharati judgment") — but ONLY when the query is # essentially just the name: doctrinal residue ("bachan singh sentencing principles") # must fall through to full retrieval, not the single-doc shortcut hits = [(a, d) for a, d in self.aliases.items() if len(a) >= 8 and a in qa and len(qa) - len(a) <= 10] if hits: best = max(hits, key=lambda ad: self.cite_indeg.get(ad[1], 0)) if self.is_retrieval_eligible(best[1]): return [best[1]], "case name" if re.search(r"\bv[s.]?\b|\bversus\b", ql, re.I) and len(ql) <= 90: hits = [c["doc_id"] for c in self.name_lookup(ql, 6)] if hits: return hits, "case name" return [], None def name_lookup(self, name, k=4): """Resolve a case NAME to corpus doc(s) — the recall tool for LLM-named authorities.""" raw = [t for t in _ntok(name) if t not in _NAME_STOP and len(t) > 1] if not raw: return [] # compound-name variants (Indian names split/join freely: Ibrahimuddin <-> Ibrahim Uddin) extra = [] for t in raw: if t not in self.name_vocab and len(t) >= 7: # try splitting an unknown long token for cut in range(3, len(t) - 2): a, b = t[:cut], t[cut:] if a in self.name_vocab and b in self.name_vocab: extra += [a, b]; break for a, b in zip(raw, raw[1:]): # try joining adjacent tokens if (a + b) in self.name_vocab: extra.append(a + b) raw += extra qtok = set() for t in raw: if t in self.name_vocab or len(t) <= 3: qtok.add(t) else: qtok.update(difflib.get_close_matches(t, self.name_vocab, n=3, cutoff=0.82) or [t]) cand = set() # only docs sharing a query token (inverted index) for t in qtok: cand |= self.name_postings.get(t, set()) qdist = {t for t in qtok if len(t) >= 5} # distinctive party-name tokens (must match one) scored = [] for d in cand: if not self.is_retrieval_eligible(d): continue ntok = set(_ntok(self.meta.get(d, {}).get("case_name") or "")) ov = qtok & ntok if qdist and not (qdist & ntok): continue # reject namesakes that miss the party name if len(ov) >= 2 or (len(ov) == 1 and any(len(t) >= 5 for t in ov)): # rank: most query tokens matched, then the sibling-cluster CANONICAL (the main # judgment, not its referral order), then concision, then authority scored.append((len(ov), 1 if d in self.canonical else 0, -(len(ntok) - len(ov)), self.cite_indeg.get(d, 0), d)) scored.sort(reverse=True) return [self._card(t[-1]) for t in scored[:k]]