| """CPU runtime for an immutable schema-v5 Moonley serving release.""" |
|
|
| from __future__ import annotations |
|
|
| import difflib |
| import json |
| import math |
| import os |
| import re |
| import sqlite3 |
| import threading |
| from collections import Counter, OrderedDict, defaultdict |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
|
|
| from statute_crosswalk import load_default_crosswalk |
| from statute_library import ExactStatuteLibrary |
|
|
|
|
| QUERY_TASK = ( |
| "Given a legal research query, retrieve relevant passages from judgments " |
| "of the Supreme Court of India that answer the query" |
| ) |
| NAME_STOP = { |
| "v", "vs", "of", "and", "the", "ors", "anr", "etc", "state", "union", |
| "govt", "government", "in", "re", "ltd", "co", "pvt", "through", "another", |
| } |
| NAME_QUERY_NOISE = NAME_STOP | { |
| "about", "case", "court", "decision", "did", "give", "held", "holding", |
| "for", "is", "judgement", "judgment", "know", "me", "on", "passed", "please", |
| "say", "tell", "was", "what", "which", |
| } |
| BAD_STATUS = {"overruled", "partly_overruled", "per_incuriam", "doubted"} |
|
|
| ACT_ALIASES = { |
| "tpa": "transfer property act", |
| "transfer of property act": "transfer property act", |
| "transfer property act": "transfer property act", |
| "ipc": "indian penal code", |
| "crpc": "code criminal procedure", |
| "cpc": "code civil procedure", |
| "iea": "indian evidence act", |
| "ni act": "negotiable instruments act", |
| "bns": "bharatiya nyaya sanhita", |
| "bnss": "bharatiya nagarik suraksha sanhita", |
| "bsa": "bharatiya sakshya adhiniyam", |
| } |
|
|
|
|
| def _clean(value: object) -> str: |
| return re.sub(r"\s+", " ", str(value or "")).strip() |
|
|
|
|
| def _json(value: object, default: object) -> object: |
| try: |
| return json.loads(str(value)) if value not in (None, "") else default |
| except (TypeError, ValueError, json.JSONDecodeError): |
| return default |
|
|
|
|
| def _ntok(value: object) -> list[str]: |
| out, single = [], "" |
| for token in re.findall(r"[a-z]+", str(value or "").lower()): |
| if len(token) == 1: |
| single += token |
| else: |
| if single: |
| out.append(single); single = "" |
| out.append(token) |
| if single: |
| out.append(single) |
| return out |
|
|
|
|
| def _norm_identity(value: object) -> str: |
| text = str(value or "").lower().replace("versus", " v ").replace("vs.", " v ") |
| text = re.sub(r"\bvs?\b", " v ", text) |
| return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9]+", " ", text)).strip() |
|
|
|
|
| def _norm_citation(value: object) -> str: |
| return re.sub(r"\s+", " ", re.sub(r"[^A-Z0-9]+", " ", str(value or "").upper())).strip() |
|
|
|
|
| def _norm_act(value: object) -> str: |
| """Collapse common abbreviations and harmless title/year variants.""" |
| text = re.sub(r"\b(?:18|19|20)\d{2}\b", " ", str(value or "").lower()) |
| text = re.sub(r"[^a-z0-9]+", " ", text) |
| text = re.sub(r"\s+", " ", text).strip() |
| if text in ACT_ALIASES: |
| return ACT_ALIASES[text] |
| tokens = [token for token in text.split() if token not in {"the", "of", "india"}] |
| normalized = " ".join(tokens) |
| return ACT_ALIASES.get(normalized, normalized) |
|
|
|
|
| def _norm_section(value: object) -> str: |
| text = re.sub(r"^\s*(?:sections?|ss?\.?)[\s:-]*", "", str(value or ""), flags=re.I) |
| return re.sub(r"\s+", "", text).strip(".,;:") |
|
|
|
|
| class CorpusV5: |
| """Expose the legacy agent tool contract over the Qwen/schema-v5 bundle. |
| |
| Every public judgment method starts at ``eligible_doc_ids``. The release |
| builder has already proved that each member has metadata, stored paragraphs, |
| and authoritative search units; the runtime rechecks those counts at boot. |
| """ |
|
|
| def __init__( |
| self, |
| data_dir: str | os.PathLike[str], |
| statute_dir: str | os.PathLike[str] | None = None, |
| device: str = "cpu", |
| *, |
| index: Any | None = None, |
| query_encoder: Any | None = None, |
| ): |
| self.data_dir = Path(data_dir) |
| self.device = device |
| self.manifest = json.loads((self.data_dir / "release_manifest.json").read_text(encoding="utf-8")) |
| if self.manifest.get("status") != "complete": |
| raise RuntimeError("schema-v5 serving release is not complete") |
| self.model_config = self.manifest.get("model") or {} |
| self.dimension = int(self.model_config.get("dimension") or 2560) |
| self._db_path = self.data_dir / str((self.manifest.get("artifacts") or {}).get("database", {}).get("name") or "corpus.sqlite3") |
| if not self._db_path.exists(): |
| raise RuntimeError(f"serving database missing: {self._db_path}") |
| self._local = threading.local() |
| self._encoder_lock = threading.Lock() |
| self._model_load_lock = threading.Lock() |
| self._query_encoder = query_encoder |
| self._query_cache: OrderedDict[str, np.ndarray] = OrderedDict() |
| self._query_cache_lock = threading.Lock() |
| self._query_cache_size = max(8, int(os.environ.get("THEMIS_QUERY_CACHE", "64"))) |
|
|
| if index is None: |
| import faiss |
|
|
| index_path = self.data_dir / str((self.manifest.get("artifacts") or {}).get("faiss_index", {}).get("name") or "index.faiss") |
| flags = getattr(faiss, "IO_FLAG_MMAP", 0) | getattr(faiss, "IO_FLAG_READ_ONLY", 0) |
| self.index = faiss.read_index(str(index_path), flags) |
| else: |
| self.index = index |
| if int(getattr(self.index, "d", self.dimension)) != self.dimension: |
| raise RuntimeError("FAISS dimension does not match release manifest") |
|
|
| self.meta: dict[str, dict[str, Any]] = {} |
| self.goodlaw: dict[str, dict[str, Any]] = {} |
| self.decision_year: dict[str, int] = {} |
| self.bench_n: dict[str, int] = {} |
| self.canonical: set[str] = set() |
| self.name_vocab: set[str] = set() |
| self.name_postings: dict[str, set[str]] = defaultdict(set) |
| self.aliases: dict[str, str] = {} |
| self.nc2doc: dict[str, str] = {} |
| self.cite_resolver: dict[str, str] = {} |
| self._load_metadata() |
| self.eligible_doc_ids = set(self.meta) |
| self.canonical = set(self.eligible_doc_ids) |
|
|
| self._doc_rows: dict[str, list[int]] = defaultdict(list) |
| self._row_doc: list[str] = [] |
| self._row_type: list[str] = [] |
| self._unit_cache: OrderedDict[int, dict[str, Any]] = OrderedDict() |
| self._load_unit_map() |
| expected_units = int((self.manifest.get("corpus") or {}).get("units") or 0) |
| if expected_units and (len(self._row_doc) != expected_units or int(getattr(self.index, "ntotal", expected_units)) != expected_units): |
| raise RuntimeError("unit-table, manifest, and FAISS row counts diverge") |
| missing = self.eligible_doc_ids - set(self._doc_rows) |
| if missing: |
| raise RuntimeError(f"{len(missing)} accepted judgments have no serving units") |
|
|
| self.out_edges: dict[str, list[str]] = defaultdict(list) |
| self.in_edges: dict[str, list[str]] = defaultdict(list) |
| self.edge_meta: dict[tuple[str, str], dict[str, Any]] = {} |
| self.cite_indeg: Counter[str] = Counter() |
| self._load_graph() |
| self.statute_idx = self._load_statute_index() |
| self._provision_act_names: dict[str, list[str]] = defaultdict(list) |
| for row in self._connection().execute( |
| "SELECT DISTINCT act_name FROM provisions WHERE act_name IS NOT NULL AND act_name != ''" |
| ): |
| name = str(row["act_name"]) |
| self._provision_act_names[_norm_act(name)].append(name) |
| self.concord = {} |
| concordance = Path(statute_dir or "") / "concordance.json" if statute_dir else None |
| if concordance and concordance.exists(): |
| self.concord = json.loads(concordance.read_text(encoding="utf-8")) |
| crosswalk_path = os.environ.get("THEMIS_SECTION_CROSSWALK", "").strip() or None |
| self.crosswalk = load_default_crosswalk(crosswalk_path) |
| provisions_path = Path(statute_dir or "") / "all_statutes.json" if statute_dir else None |
| self.statute_library = ExactStatuteLibrary.from_env(fallback_path=provisions_path) |
| print( |
| f"[corpus-v5] ready — {len(self.eligible_doc_ids)} accepted judgments, " |
| f"{len(self._row_doc)} Qwen units, {sum(self.cite_indeg.values())} resolved internal edges", |
| flush=True, |
| ) |
|
|
| def _connection(self) -> sqlite3.Connection: |
| connection = getattr(self._local, "connection", None) |
| if connection is None: |
| connection = sqlite3.connect(f"file:{self._db_path}?mode=ro", uri=True, check_same_thread=False, timeout=30) |
| connection.row_factory = sqlite3.Row |
| self._local.connection = connection |
| return connection |
|
|
| def _load_metadata(self) -> None: |
| connection = sqlite3.connect(f"file:{self._db_path}?mode=ro", uri=True) |
| connection.row_factory = sqlite3.Row |
| aliases_by_doc: dict[str, list[str]] = defaultdict(list) |
| alias_owners: dict[str, set[str]] = defaultdict(set) |
| for row in connection.execute("SELECT alias,normalized_alias,judgment_id FROM aliases"): |
| aliases_by_doc[str(row["judgment_id"])].append(str(row["alias"])) |
| alias_owners[str(row["normalized_alias"])].add(str(row["judgment_id"])) |
| for key, owners in alias_owners.items(): |
| if len(owners) == 1: |
| self.aliases[key] = next(iter(owners)) |
| citation_owners: dict[str, set[str]] = defaultdict(set) |
| for row in connection.execute("SELECT * FROM judgments"): |
| d = str(row["judgment_id"]) |
| equivalents = list(_json(row["equivalent_citations_json"], [])) |
| bench = list(_json(row["bench_json"], [])) |
| acts_records = list(_json(row["acts_json"], [])) |
| provisions = list(_json(row["provisions_json"], [])) |
| summary = dict(_json(row["summary_json"], {})) |
| graph_metrics = dict(_json(row["graph_metrics_json"], {})) |
| case_numbers = list(_json(row["case_numbers_json"], [])) |
| case_number = next((item.get("raw") for item in case_numbers if isinstance(item, dict) and item.get("raw")), None) |
| m = { |
| "doc_id": d, "judgment_id": d, "case_name": row["case_name"], |
| "neutral_citation": row["neutral_citation"], "equivalent_citations": equivalents, |
| "date": row["decision_date"], "year": row["year"], "court": row["court"], |
| "case_number": case_number, "bench_strength": row["bench_size"] or row["bench_bucket"], |
| "bench": bench, "author_judge": None, "disposition": row["disposition"], |
| "acts": [item.get("name") for item in acts_records if isinstance(item, dict) and item.get("name")], |
| "acts_records": acts_records, "provisions": provisions, "issue": row["issue"], "held": row["held"], |
| "summary": summary, "source_url": row["source_url"], "source_provider": row["source_provider"], |
| "source_ik_tid": row["source_ik_tid"], "review_status": row["review_status"], |
| "aliases": aliases_by_doc.get(d, []), "graph_metrics": graph_metrics, |
| } |
| self.meta[d] = m |
| status = str(row["good_law_status"] or "unknown") |
| good_law = dict(_json(row["good_law_json"], {})) |
| self.goodlaw[d] = { |
| **good_law, "good_law_status": status, |
| "display_state": row["display_state"] or "grey", |
| "treatment_breakdown": graph_metrics.get("treatment_breakdown") or {}, |
| } |
| try: |
| self.decision_year[d] = int(row["year"]) |
| except (TypeError, ValueError): |
| pass |
| self.bench_n[d] = int(row["bench_size"] or 0) |
| for token in _ntok(row["case_name"]): |
| if len(token) >= 4: |
| self.name_vocab.add(token) |
| if len(token) > 1: |
| self.name_postings[token].add(d) |
| if row["neutral_citation"]: |
| self.nc2doc[str(row["neutral_citation"])] = d |
| for citation in [row["neutral_citation"], *equivalents]: |
| normalized = _norm_citation(citation) |
| if normalized: |
| citation_owners[normalized].add(d) |
| for key, owners in citation_owners.items(): |
| if len(owners) == 1: |
| self.cite_resolver[key] = next(iter(owners)) |
| connection.close() |
|
|
| def _load_unit_map(self) -> None: |
| for row in self._connection().execute("SELECT row_id,judgment_id,unit_type FROM units ORDER BY row_id"): |
| row_id = int(row["row_id"]) |
| if row_id != len(self._row_doc): |
| raise RuntimeError("serving unit rows are not contiguous") |
| judgment_id = str(row["judgment_id"]) |
| self._row_doc.append(judgment_id) |
| self._row_type.append(str(row["unit_type"])) |
| self._doc_rows[judgment_id].append(row_id) |
|
|
| def _load_graph(self) -> None: |
| query = "SELECT * FROM graph_edges WHERE target_id IS NOT NULL" |
| for row in self._connection().execute(query): |
| source, target = str(row["source_id"]), str(row["target_id"]) |
| if source not in self.eligible_doc_ids or target not in self.eligible_doc_ids: |
| continue |
| self.out_edges[source].append(target); self.in_edges[target].append(source) |
| self.edge_meta[(source, target)] = { |
| "treatment": row["relation"] or "referred_to", "scope": row["scope"], |
| "confidence": row["confidence"], "evidence": list(_json(row["evidence_json"], [])), |
| "method": row["resolution_method"], |
| } |
| self.cite_indeg[target] += 1 |
|
|
| def _load_statute_index(self) -> list[dict[str, Any]]: |
| rows = self._connection().execute( |
| "SELECT act_name,number,MIN(raw_mention) title,COUNT(DISTINCT judgment_id) cases " |
| "FROM provisions WHERE act_name IS NOT NULL GROUP BY act_name,number ORDER BY cases DESC LIMIT 25000" |
| ) |
| return [ |
| {"act_short": row["act_name"], "section_number": row["number"], "title": row["title"], "cases": row["cases"]} |
| for row in rows |
| ] |
|
|
| def _load_encoder(self) -> Any: |
| if self._query_encoder is not None: |
| return self._query_encoder |
| with self._model_load_lock: |
| if self._query_encoder is not None: |
| return self._query_encoder |
| import torch |
| from sentence_transformers import SentenceTransformer |
|
|
| model_path = os.environ.get("THEMIS_QWEN_MODEL") or self.model_config.get("model_id") or "Qwen/Qwen3-Embedding-4B" |
| dtype_name = os.environ.get("THEMIS_QWEN_DTYPE", "bfloat16").lower() |
| dtype = torch.bfloat16 if dtype_name == "bfloat16" else torch.float32 |
| torch.set_num_threads(max(1, int(os.environ.get("THEMIS_TORCH_THREADS", str(os.cpu_count() or 4))))) |
| kwargs: dict[str, Any] = { |
| "device": "cpu", "trust_remote_code": True, |
| "model_kwargs": {"dtype": dtype, "low_cpu_mem_usage": True}, |
| } |
| if Path(str(model_path)).exists(): |
| kwargs["local_files_only"] = True |
| else: |
| kwargs["revision"] = self.model_config.get("revision") |
| model = SentenceTransformer(str(model_path), **kwargs) |
| model.max_seq_length = int(self.model_config.get("max_seq_length") or 2048) |
| self._query_encoder = model |
| return model |
|
|
| def warmup(self) -> None: |
| self._enc("Supreme Court legal research") |
|
|
| def _enc(self, query: str) -> np.ndarray: |
| normalized = _clean(query) |
| with self._query_cache_lock: |
| cached = self._query_cache.get(normalized) |
| if cached is not None: |
| self._query_cache.move_to_end(normalized) |
| return cached.copy() |
| prompt = f"Instruct: {self.model_config.get('query_task') or QUERY_TASK}\nQuery: {normalized}" |
| encoder = self._load_encoder() |
| with self._encoder_lock: |
| if callable(encoder) and not hasattr(encoder, "encode"): |
| vector = encoder(prompt) |
| else: |
| vector = encoder.encode(prompt, normalize_embeddings=True, convert_to_numpy=True) |
| vector = np.asarray(vector, dtype=np.float32).reshape(-1) |
| if vector.shape != (self.dimension,): |
| raise RuntimeError(f"query encoder returned {vector.shape}; expected {(self.dimension,)}") |
| vector /= max(float(np.linalg.norm(vector)), 1e-12) |
| with self._query_cache_lock: |
| self._query_cache[normalized] = vector.copy() |
| self._query_cache.move_to_end(normalized) |
| while len(self._query_cache) > self._query_cache_size: |
| self._query_cache.popitem(last=False) |
| return vector |
|
|
| def encode_documents(self, texts: list[str]) -> np.ndarray: |
| """Encode private knowledge chunks in the Qwen document space.""" |
| values = [_clean(text) for text in texts if _clean(text)] |
| if not values: |
| return np.empty((0, self.dimension), dtype=np.float32) |
| encoder = self._load_encoder() |
| with self._encoder_lock: |
| if callable(encoder) and not hasattr(encoder, "encode"): |
| matrix = np.vstack([encoder(value) for value in values]) |
| else: |
| matrix = encoder.encode( |
| values, |
| batch_size=max(1, int(os.environ.get("THEMIS_KNOWLEDGE_BATCH", "4"))), |
| normalize_embeddings=True, |
| convert_to_numpy=True, |
| show_progress_bar=False, |
| ) |
| matrix = np.asarray(matrix, dtype=np.float32) |
| if matrix.shape != (len(values), self.dimension): |
| raise RuntimeError( |
| f"document encoder returned {matrix.shape}; expected {(len(values), self.dimension)}" |
| ) |
| return matrix |
|
|
| def _unit(self, row_id: int) -> dict[str, Any]: |
| cached = self._unit_cache.get(int(row_id)) |
| if cached is not None: |
| self._unit_cache.move_to_end(int(row_id)); return cached |
| row = self._connection().execute("SELECT * FROM units WHERE row_id=?", (int(row_id),)).fetchone() |
| if row is None: |
| return {} |
| unit = dict(row); unit["paragraph_ids"] = list(_json(unit.pop("paragraph_ids_json", "[]"), [])) |
| self._unit_cache[int(row_id)] = unit |
| if len(self._unit_cache) > 4096: |
| self._unit_cache.popitem(last=False) |
| return unit |
|
|
| def _dense_units(self, query: str, n: int = 1500) -> list[tuple[int, float]]: |
| limit = max(1, min(int(n), len(self._row_doc))) |
| scores, rows = self.index.search(self._enc(query).reshape(1, -1), limit) |
| return [(int(row), float(score)) for row, score in zip(rows[0], scores[0]) if int(row) >= 0] |
|
|
| def _card(self, judgment_id: str, rr: float = 0.0, passage: str | None = None) -> dict[str, Any]: |
| d = str(judgment_id); m = self.meta.get(d, {}); gl = self.goodlaw.get(d, {}) |
| snippet = _clean(passage or m.get("held") or (m.get("summary") or {}).get("one_line") or (m.get("summary") or {}).get("text"))[:420] |
| return { |
| "doc_id": d, "judgment_id": d, "case_name": m.get("case_name"), "year": m.get("year"), |
| "date": m.get("date"), "neutral_citation": m.get("neutral_citation"), |
| "equivalent_citations": m.get("equivalent_citations") or [], "court": m.get("court"), |
| "bench_strength": m.get("bench_strength"), "disposition": m.get("disposition"), |
| "cited_by": self.cite_indeg.get(d, 0), "good_law": gl.get("good_law_status", "unknown"), |
| "good_law_status": gl.get("good_law_status", "unknown"), "rr": round(float(rr), 6), |
| "snippet": snippet, "passage": snippet, |
| } |
|
|
| def is_retrieval_eligible(self, doc_id: object) -> bool: |
| return str(doc_id) in self.eligible_doc_ids |
|
|
| def coverage(self) -> dict[str, Any]: |
| corpus = self.manifest.get("corpus") or {} |
| return { |
| "accepted_judgments": len(self.eligible_doc_ids), "metadata_only_excluded": 0, |
| "units": len(self._row_doc), "paragraphs": int(corpus.get("paragraphs") or 0), |
| "scope": corpus.get("source_scope") or "Supreme Court of India judgments stored in this release", |
| "release_version": self.manifest.get("release_version"), |
| } |
|
|
| def vector_search(self, query: str, k: int = 8) -> list[dict[str, Any]]: |
| out, seen = [], set() |
| for row_id, score in self._dense_units(query, max(800, k * 80)): |
| d = self._row_doc[row_id] |
| if d in seen or not self.is_retrieval_eligible(d): |
| continue |
| seen.add(d); out.append(self._card(d, score, self._unit(row_id).get("text"))) |
| if len(out) >= k: |
| break |
| return out |
|
|
| def dense_docs(self, query: str, k: int = 60) -> list[str]: |
| return [card["doc_id"] for card in self.vector_search(query, k)] |
|
|
| def keyword_search(self, query: str, k: int = 12, **_: Any) -> list[dict[str, Any]]: |
| stop = {"of", "the", "and", "or", "in", "to", "a", "an", "is", "for", "on", "by", "with"} |
| tokens = [token for token in re.findall(r"[a-z0-9]+", query.lower()) if token not in stop] |
| if not tokens: |
| return [] |
| match = " OR ".join(f'"{token}"' for token in tokens[:24]) |
| rows = self._connection().execute( |
| "SELECT rowid,bm25(unit_fts) score FROM unit_fts WHERE unit_fts MATCH ? ORDER BY score LIMIT ?", |
| (match, max(300, k * 40)), |
| ).fetchall() |
| out, seen = [], set() |
| for row in rows: |
| row_id = int(row["rowid"]); d = self._row_doc[row_id] |
| if d in seen or not self.is_retrieval_eligible(d): |
| continue |
| seen.add(d); out.append(self._card(d, -float(row["score"]), self._unit(row_id).get("text"))) |
| if len(out) >= k: |
| break |
| return out |
|
|
| def score_docs(self, query: str, doc_ids: list[str], per_doc: int = 8) -> dict[str, float]: |
| qv = self._enc(query); scores: dict[str, float] = {} |
| for value in doc_ids: |
| d = str(value) |
| if not self.is_retrieval_eligible(d): |
| continue |
| rows = self._doc_rows[d][: max(1, per_doc)] |
| if not rows: |
| continue |
| vectors = np.vstack([np.asarray(self.index.reconstruct(int(row)), dtype=np.float32) for row in rows]) |
| scores[d] = float(np.max(vectors @ qv)) |
| return scores |
|
|
| def hybrid_search(self, query: str, k: int = 8, pool: int = 60) -> list[dict[str, Any]]: |
| dense = self.vector_search(query, pool); keyword = self.keyword_search(query, pool) |
| fused: dict[str, float] = defaultdict(float) |
| cards: dict[str, dict[str, Any]] = {} |
| for lane in (dense, keyword): |
| for rank, card in enumerate(lane, 1): |
| d = card["doc_id"]; fused[d] += 1.0 / (60 + rank); cards.setdefault(d, card) |
| ranked = sorted(fused, key=lambda d: -fused[d])[: max(k * 5, 40)] |
| refined = self.score_docs(query, ranked) |
| ranked.sort(key=lambda d: -(refined.get(d, 0.0) + 8 * fused[d])) |
| return [self._card(d, refined.get(d, fused[d]), cards[d].get("passage")) for d in ranked[:k]] |
|
|
| def authority_search(self, query: str, k: int = 8, alpha: float = 0.3) -> list[dict[str, Any]]: |
| cards = self.vector_search(query, max(80, k * 10)) |
| cards.sort(key=lambda card: -(float(card.get("rr") or 0) + alpha * math.log1p(self.cite_indeg.get(card["doc_id"], 0)))) |
| return cards[:k] |
|
|
| def search_lanes(self, query: str, frame: dict[str, Any], lane_n: int = 6) -> dict[str, list[dict[str, Any]]]: |
| """Build all protected lanes from one Qwen query encoding/index scan. |
| |
| CPU serving cannot afford to encode every LLM paraphrase independently. |
| The approved lawyer query is the semantic anchor; frame variants shape |
| deterministic lane ordering and FTS lookups without another 4B-model pass. |
| """ |
| base = self.hybrid_search(query, max(24, lane_n * 4)) |
| factual = base[:lane_n] |
| doctrine = sorted( |
| base, |
| key=lambda card: -( |
| float(card.get("rr") or 0) |
| + 0.18 * math.log1p(self.cite_indeg.get(card["doc_id"], 0)) |
| + 0.04 * self.bench_n.get(card["doc_id"], 0) |
| ), |
| )[: max(lane_n, 10)] |
| seen = {card["doc_id"] for card in doctrine} |
| for authority in frame.get("authorities") or []: |
| for card in self.name_lookup(str(authority), 2): |
| if card["doc_id"] not in seen: |
| card["named"] = True; card["rr"] = max(1.0, float(card.get("rr") or 0)) |
| doctrine.append(card); seen.add(card["doc_id"]) |
| |
| |
| |
| statute_runs: list[list[dict[str, Any]]] = [] |
| for section in (frame.get("sections") or [])[:3]: |
| act = str(section.get("act") or "") |
| number = str(section.get("section") or "") |
| provisions = [{"act": act, "section": number}] |
| corresponding = self.statute_crosswalk(act, number).get("corresponding") or [] |
| provisions.extend(corresponding[:5]) |
| for provision in provisions: |
| mapped_act = str(provision.get("act") or "") |
| mapped_number = str(provision.get("section") or "") |
| exact = self.provision_cases(mapped_act, mapped_number, max(3, lane_n), query=query) |
| if not exact: |
| exact = self.keyword_search( |
| f"{mapped_act} section {mapped_number}", max(3, lane_n) |
| ) |
| statute_runs.append(exact) |
| statute: list[dict[str, Any]] = [] |
| statute_seen: set[str] = set() |
| for rank in range(max((len(run) for run in statute_runs), default=0)): |
| for run in statute_runs: |
| if rank >= len(run): |
| continue |
| card = run[rank] |
| if card["doc_id"] not in statute_seen: |
| statute_seen.add(card["doc_id"]); statute.append(card) |
| if len(statute) >= lane_n: |
| break |
| if len(statute) >= lane_n: |
| break |
| known = [] |
| for value in frame.get("known_citations") or []: |
| ids, _ = self.identity_hits(str(value)) |
| for d in ids: |
| if all(card["doc_id"] != d for card in known): |
| known.append(self._card(d)) |
| return {"factual": factual, "doctrine": doctrine, "statute": statute[:lane_n], "known": known[:5]} |
|
|
| def provision_cases( |
| self, |
| act: object, |
| section: object, |
| k: int = 8, |
| *, |
| query: str | None = None, |
| ) -> list[dict[str, Any]]: |
| """Return judgments carrying an exact structured act/section match. |
| |
| The query embedding remains the Qwen judgment embedding. Bare-act BGE |
| vectors, when enabled, are a separate retrieval space and never enter |
| this score calculation. |
| """ |
| act_key, number = _norm_act(act), _norm_section(section) |
| names = self._provision_act_names.get(act_key, []) |
| if not names or not number: |
| return [] |
| placeholders = ",".join("?" for _ in names) |
| rows = self._connection().execute( |
| f"SELECT judgment_id,GROUP_CONCAT(DISTINCT salience) saliences " |
| f"FROM provisions WHERE act_name IN ({placeholders}) AND number=? GROUP BY judgment_id", |
| (*names, number), |
| ).fetchall() |
| doc_ids = [str(row["judgment_id"]) for row in rows if self.is_retrieval_eligible(row["judgment_id"])] |
| if not doc_ids: |
| return [] |
| salience_by_doc = {str(row["judgment_id"]): str(row["saliences"] or "") for row in rows} |
| topical = self.score_docs(query, doc_ids) if query else {} |
|
|
| def rank_signals(doc_id: str) -> tuple[int, int, float]: |
| metrics = self.meta.get(doc_id, {}).get("graph_metrics") or {} |
| try: |
| external = max(0, int(metrics.get("cited_by_count") or 0)) |
| except (TypeError, ValueError): |
| external = 0 |
| saliences = {value.strip().lower() for value in salience_by_doc.get(doc_id, "").split(",")} |
| salience = 2 if saliences & {"core", "ratio", "primary"} else 1 if "supporting" in saliences else 0 |
| blended = ( |
| float(topical.get(doc_id, 0.0)) |
| + 0.35 * math.log1p(external) |
| + 0.10 * math.log1p(self.cite_indeg.get(doc_id, 0)) |
| + 0.025 * self.bench_n.get(doc_id, 0) |
| ) |
| return salience, external, blended |
|
|
| |
| |
| |
| signals = {doc_id: rank_signals(doc_id) for doc_id in doc_ids} |
| doc_ids.sort(key=lambda doc_id: (-signals[doc_id][0], -signals[doc_id][1], -signals[doc_id][2], str(doc_id))) |
| out = [] |
| for doc_id in doc_ids[: max(1, int(k))]: |
| card = self._card(doc_id, topical.get(doc_id, 0.0)) |
| card["provision_match"] = {"act": act, "section": number, "exact": True} |
| card["provision_salience"] = salience_by_doc.get(doc_id, "") |
| card["native_cited_by"] = signals[doc_id][1] |
| card["protected"] = True |
| out.append(card) |
| return out |
|
|
| def held_search(self, query: str, k: int = 12) -> list[str]: |
| out, seen = [], set() |
| for row_id, _ in self._dense_units(query, max(1600, k * 100)): |
| if self._row_type[row_id] not in {"holdings_ratio", "summary_overview", "issues_facts"}: |
| continue |
| d = self._row_doc[row_id] |
| if d not in seen: |
| seen.add(d); out.append(d) |
| if len(out) >= k: |
| break |
| return out |
|
|
| def citectx_search(self, query: str, k: int = 12) -> list[str]: |
| out, seen = [], set() |
| for row_id, _ in self._dense_units(query, max(2500, k * 140)): |
| if "citation" not in self._row_type[row_id]: |
| continue |
| d = self._row_doc[row_id] |
| if d not in seen: |
| seen.add(d); out.append(d) |
| if len(out) >= k: |
| break |
| return out |
|
|
| def identity_hits(self, query: str) -> tuple[list[str], str | None]: |
| cite_match = re.search(r"\b\d{4}\s+INSC\s+\d+\b|\[\d{4}\]\s*\d+\s*S\.?C\.?R\.?\s*\d+|\(\d{4}\)\s*\d+\s*SCC\s*\d+|AIR\s+\d{4}\s+SC\s+\d+", query, re.I) |
| if cite_match: |
| d = self.cite_resolver.get(_norm_citation(cite_match.group(0))) |
| if d: |
| return [d], "citation" |
| normalized = _norm_identity(query) |
| if normalized in self.aliases: |
| return [self.aliases[normalized]], "case name" |
| candidates = [(alias, d) for alias, d in self.aliases.items() if len(alias) >= 8 and alias in normalized and len(normalized) - len(alias) <= 12] |
| if candidates: |
| candidates.sort(key=lambda item: (-len(item[0]), -self.cite_indeg.get(item[1], 0))) |
| return [candidates[0][1]], "case name" |
| |
| |
| |
| |
| words = _ntok(query) |
| explicit_lookup = bool(re.search(r"\b(?:case|judg(?:e)?ment|holding|held|decision)\b", str(query or ""), re.I)) |
| name_tokens = [token for token in words if token not in NAME_QUERY_NOISE and len(token) > 1] |
| if 1 <= len(name_tokens) <= 5 and (explicit_lookup or len(name_tokens) >= 2): |
| candidate_ids: set[str] = set() |
| for token in name_tokens: |
| vocabulary = [token] |
| if token not in self.name_vocab and len(token) >= 5: |
| vocabulary.extend(difflib.get_close_matches(token, self.name_vocab, n=3, cutoff=0.78)) |
| for value in vocabulary: |
| candidate_ids.update(self.name_postings.get(value, set())) |
| scored = [] |
| for d in candidate_ids: |
| case_tokens = [token for token in _ntok(self.meta[d].get("case_name")) if token not in NAME_STOP] |
| ratios = [max((difflib.SequenceMatcher(None, token, other, autojunk=False).ratio() for other in case_tokens), default=0.0) for token in name_tokens] |
| if ratios and all(score >= 0.78 for score in ratios): |
| score = sum(ratios) / len(ratios) |
| scored.append((score, sum(value == 1.0 for value in ratios), self.cite_indeg.get(d, 0), d)) |
| scored.sort(reverse=True) |
| if scored and scored[0][0] >= 0.88: |
| if len(scored) > 1 and scored[1][0] >= 0.88 and scored[0][0] - scored[1][0] <= 0.015: |
| return [item[-1] for item in scored[:3]], "ambiguous case name" |
| kind = "case name" if scored[0][0] >= 0.999 else "close case name" |
| return [scored[0][-1]], kind |
| explicit_named_case = bool(re.search(r"\b(?:case|v(?:s)?\.?|versus)\b", str(query or ""), re.I)) |
| if explicit_named_case and any(token in self.name_vocab for token in name_tokens): |
| return [], "unresolved case name" |
| return [], None |
|
|
| def name_lookup(self, name: str, k: int = 4) -> list[dict[str, Any]]: |
| normalized = _norm_identity(name) |
| exact = self.aliases.get(normalized) |
| if exact: |
| return [self._card(exact)] |
| raw = [token for token in _ntok(name) if token not in NAME_STOP and len(token) > 1] |
| if not raw: |
| return [] |
| expanded = list(raw) |
| for token in raw: |
| if token not in self.name_vocab and len(token) >= 7: |
| expanded.extend(difflib.get_close_matches(token, self.name_vocab, n=2, cutoff=0.84)) |
| query_tokens = set(expanded) |
| candidates: set[str] = set() |
| for token in query_tokens: |
| candidates.update(self.name_postings.get(token, set())) |
| scored = [] |
| for d in candidates: |
| case_tokens = set(_ntok(self.meta[d].get("case_name"))) |
| overlap = query_tokens & case_tokens |
| if overlap: |
| scored.append((len(overlap), -abs(len(case_tokens) - len(query_tokens)), self.cite_indeg.get(d, 0), d)) |
| scored.sort(reverse=True) |
| return [self._card(item[-1]) for item in scored[:k]] |
|
|
| def statute_search(self, query: str, k: int = 3) -> list[dict[str, Any]]: |
| tokens = {token for token in re.findall(r"[a-z0-9]+", query.lower()) if len(token) > 1} |
| scored = [] |
| for item in self.statute_idx: |
| value = f"{item.get('act_short')} {item.get('section_number')} {item.get('title')}".lower() |
| overlap = sum(1 for token in tokens if token in value) |
| if overlap: |
| scored.append((overlap, int(item.get("cases") or 0), item)) |
| scored.sort(key=lambda item: (-item[0], -item[1])) |
| return [ |
| {"act": item[2].get("act_short"), "section": item[2].get("section_number"), "title": item[2].get("title"), "i": index} |
| for index, item in enumerate(scored[:k]) |
| ] |
|
|
| def cases_on_section(self, text: str, k: int = 8) -> list[dict[str, Any]]: |
| return self.hybrid_search(text, k) |
|
|
| def statute_crosswalk(self, code: str, section: object) -> dict[str, Any]: |
| return self.crosswalk.lookup(code, section) |
|
|
| def statute_provision(self, code: str, section: object) -> dict[str, Any] | None: |
| """Return exact bare-act text; never substitute a semantic neighbour.""" |
| return self.statute_library.lookup(code, section) |
|
|
| def cited_authorities(self, doc_id: str, k: int = 12) -> list[dict[str, Any]]: |
| return [self._card(d) for d in dict.fromkeys(self.out_edges.get(str(doc_id), [])) if self.is_retrieval_eligible(d)][:k] |
|
|
| def progeny(self, doc_id: str, k: int = 12) -> list[dict[str, Any]]: |
| values = sorted(set(self.in_edges.get(str(doc_id), [])), key=lambda d: -self.cite_indeg.get(d, 0)) |
| return [self._card(d) for d in values if self.is_retrieval_eligible(d)][:k] |
|
|
| def co_cited_cases(self, doc_id: str, k: int = 8) -> list[dict[str, Any]]: |
| score: Counter[str] = Counter() |
| for target in set(self.out_edges.get(str(doc_id), [])): |
| for citer in self.in_edges.get(target, []): |
| if citer != str(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: str) -> dict[str, Any]: |
| d = str(doc_id); gl = self.goodlaw.get(d, {}); status = gl.get("good_law_status", "unknown") |
| overruled_by = None |
| if status in BAD_STATUS: |
| for source in self.in_edges.get(d, []): |
| if self.edge_meta.get((source, d), {}).get("treatment") in {"overruled", "overrules"}: |
| overruled_by = self._card(source); break |
| return {"doc_id": d, "good_law": status, "treatment_breakdown": gl.get("treatment_breakdown", {}), "overruled_by": overruled_by} |
|
|
| def metadata_filter(self, cards: list[dict[str, Any]], min_bench: int | None = None, year_from: int | None = None, year_to: int | None = None) -> list[dict[str, Any]]: |
| out = [] |
| for card in cards: |
| d = card["doc_id"]; bench, year = self.bench_n.get(d, 0), self.decision_year.get(d, 0) |
| if min_bench and bench < min_bench or year_from and year and year < year_from or year_to and year and year > year_to: |
| continue |
| out.append(card) |
| return out |
|
|
| def read_case(self, doc_id: str) -> dict[str, Any]: |
| d = str(doc_id) |
| if not self.is_retrieval_eligible(d): |
| return {} |
| m = self.meta[d] |
| return {"doc_id": d, "case_name": m.get("case_name"), "neutral_citation": m.get("neutral_citation"), "bench_strength": m.get("bench_strength"), "good_law": self.goodlaw[d].get("good_law_status", "unknown"), "issue": _clean(m.get("issue"))[:1600], "held": _clean(m.get("held") or (m.get("summary") or {}).get("text"))[:2600]} |
|
|
| def front_text(self, doc_id: str, n: int = 1800) -> str: |
| m = self.meta.get(str(doc_id), {}); summary = m.get("summary") or {} |
| return _clean(m.get("held") or summary.get("text") or summary.get("one_line"))[:n] |
|
|
| def _rank_doc_rows(self, query: str, doc_id: str, limit: int = 6) -> list[tuple[int, float]]: |
| rows = self._doc_rows.get(str(doc_id), []) |
| if not rows: |
| return [] |
| qv = self._enc(query) |
| vectors = np.vstack([np.asarray(self.index.reconstruct(int(row)), dtype=np.float32) for row in rows]) |
| scores = vectors @ qv |
| order = np.argsort(-scores)[:limit] |
| return [(rows[int(i)], float(scores[int(i)])) for i in order] |
|
|
| def best_chunk_text(self, query: str, doc_id: str, limit: int = 1600) -> str: |
| ranked = self._rank_doc_rows(query, str(doc_id), 1) |
| return _clean(self._unit(ranked[0][0]).get("text"))[:limit] if ranked else "" |
|
|
| def full_text_for_read(self, query: str, doc_id: str, cap_chars: int = 90000) -> str: |
| d = str(doc_id) |
| rows = self._connection().execute("SELECT text FROM paragraphs WHERE judgment_id=? ORDER BY sequence", (d,)).fetchall() |
| full = "\n".join(str(row["text"]) for row in rows) |
| if len(full) <= cap_chars: |
| return full |
| relevant = "\n".join(self._unit(row_id).get("text", "") for row_id, _ in self._rank_doc_rows(query, d, 5)) |
| return (self.front_text(d, 6000) + "\n[...]\n" + relevant + "\n[...]\n" + "\n".join(str(row["text"]) for row in rows[-8:]))[:cap_chars] |
|
|
| def judgment_paragraphs(self, doc_id: str, offset: int = 0, limit: int = 100) -> dict[str, Any]: |
| d = str(doc_id); offset, limit = max(0, int(offset)), max(1, min(int(limit), 2000)) |
| total = int(self._connection().execute("SELECT COUNT(*) FROM paragraphs WHERE judgment_id=?", (d,)).fetchone()[0]) |
| rows = self._connection().execute( |
| "SELECT * FROM paragraphs WHERE judgment_id=? ORDER BY sequence LIMIT ? OFFSET ?", (d, limit, offset) |
| ).fetchall() |
| paragraphs = [ |
| { |
| "paragraph_id": row["paragraph_id"], "sequence": row["sequence"], |
| "paragraph_number": row["paragraph_number"], "page_number": row["page_number"], |
| "label": (f"¶ {row['paragraph_number']}" if row["paragraph_number"] else None) or row["citation_label"] or f"¶ {row['sequence']}", |
| "coordinate_status": row["coordinate_status"], "text": row["text"], |
| "html_anchor": "paragraph-" + re.sub(r"[^A-Za-z0-9_-]", "-", str(row["paragraph_id"])), |
| "source_kind": "stored_paragraph", |
| } |
| for row in rows |
| ] |
| return {"judgment_id": d, "paragraphs": paragraphs, "offset": offset, "limit": limit, "total": total, "next_offset": offset + len(paragraphs) if offset + len(paragraphs) < total else None} |
|
|
| def judgment_view(self, doc_id: str) -> dict[str, Any]: |
| d = str(doc_id) |
| if not self.is_retrieval_eligible(d): |
| return {} |
| m = self.meta[d]; page = self.judgment_paragraphs(d, 0, 2000); paragraphs = page["paragraphs"] |
| text = "\n\n".join(f"{p.get('paragraph_number') or p['sequence']}. {p['text']}" for p in paragraphs) |
| return { |
| "doc_id": d, "judgment_id": d, "summary": m.get("summary") or {}, |
| "case_name": m.get("case_name"), "neutral_citation": m.get("neutral_citation"), |
| "equivalent_citations": m.get("equivalent_citations") or [], "court": m.get("court"), |
| "date": m.get("date"), "bench_strength": m.get("bench_strength"), "disposition": m.get("disposition"), |
| "good_law_status": self.goodlaw[d].get("good_law_status", "unknown"), |
| "treatment_breakdown": self.goodlaw[d].get("treatment_breakdown", {}), |
| "cited_by": self.cite_indeg.get(d, 0), "issue": _clean(m.get("issue"))[:8000], |
| "held": _clean(m.get("held"))[:10000], "text": text, "paragraphs": paragraphs, |
| "paragraph_count": page["total"], "paragraphs_truncated": page["next_offset"] is not None, |
| "source_url": m.get("source_url"), "source_provider": m.get("source_provider"), |
| } |
|
|
| def case_chat_passages(self, query: str, doc_id: str, k: int = 5, limit: int = 1800) -> list[dict[str, Any]]: |
| d = str(doc_id) |
| if not self.is_retrieval_eligible(d): |
| return [] |
| paragraph_ids: list[str] = [] |
| for row_id, _ in self._rank_doc_rows(query, d, max(8, k * 2)): |
| for paragraph_id in self._unit(row_id).get("paragraph_ids") or []: |
| if paragraph_id not in paragraph_ids: |
| paragraph_ids.append(str(paragraph_id)) |
| if len(paragraph_ids) >= k: |
| break |
| if len(paragraph_ids) >= k: |
| break |
| out = [] |
| for paragraph_id in paragraph_ids: |
| row = self._connection().execute("SELECT * FROM paragraphs WHERE paragraph_id=? AND judgment_id=?", (paragraph_id, d)).fetchone() |
| if row is None: |
| continue |
| out.append({ |
| "paragraph_id": paragraph_id, "label": (f"¶ {row['paragraph_number']}" if row["paragraph_number"] else None) or row["citation_label"] or f"¶ {row['sequence']}", |
| "text": _clean(row["text"])[:limit], "source_kind": "stored_paragraph", |
| "sequence": row["sequence"], |
| "html_anchor": "paragraph-" + re.sub(r"[^A-Za-z0-9_-]", "-", paragraph_id), |
| }) |
| return out |
|
|
| def relevant_passages(self, query: str, doc_id: str, k: int = 6) -> list[dict[str, Any]]: |
| """Case-local semantic pinpoints for the query, resolved to stored paragraphs.""" |
| return [ |
| {**item, "highlight_kind": "query_relevance"} |
| for item in self.case_chat_passages(query, doc_id, k=k, limit=6000) |
| ] |
|
|
|
|
| __all__ = ["CorpusV5"] |
|
|
