Mini Legal Casebase Search Engine

from dataclasses import dataclass, field
from datetime import date

@dataclass
class Paragraph:
    number: int          # Paragraph number in judgment
    text: str            # Raw paragraph text
    embedding: list[float]  # 384-dim vector (all-MiniLM-L6-v2)

@dataclass
class Case:
    id: str              # SHA-256 of normalised citation (stable, unique)
    name: str            # "Donoghue v Stevenson"
    citation: str        # "[1932] AC 562"
    court: str           # "House of Lords"
    jurisdiction: str    # "United Kingdom"
    date: date           # Date of judgment
    judges: list[str]    # Presiding judges
    catchwords: list[str]   # Issue labels (manual or NLP-extracted)
    full_text: str       # Complete judgment text (UTF-8)
    paragraphs: list[Paragraph]  # Segmented, each independently embedded
    citations_out: list[str]  # Case IDs this case cites
    citations_in: list[str]   # Case IDs that cite this one (reverse-indexed)

from whoosh.fields import Schema, TEXT, ID, STORED, KEYWORD, DATETIME
from whoosh.analysis import StemmingAnalyzer, KeywordAnalyzer

# Field boosts express the relative importance of a match in each field.
# A citation hit is worth twice a body hit; a name hit worth 1.5×.
case_schema = Schema(
    id=ID(stored=True, unique=True),
    name=TEXT(
        stored=True,
        analyzer=StemmingAnalyzer(),
        field_boost=1.5,    # "Donoghue" matches are highly relevant
    ),
    citation=TEXT(
        stored=True,
        analyzer=KeywordAnalyzer(),   # Never stem "[1932] AC 562"
        field_boost=2.0,    # Exact citation lookup is the most precise query
    ),
    court=TEXT(stored=True, field_boost=1.2),
    catchwords=TEXT(
        stored=True,
        analyzer=StemmingAnalyzer(),
        field_boost=1.3,    # Issue labels are denser signal than body text
    ),
    jurisdiction=KEYWORD(stored=True, commas=True),
    date=DATETIME(stored=True),
    full_text=TEXT(analyzer=StemmingAnalyzer(), field_boost=1.0),
)

# Boolean / phrase search is parsed with Whoosh's MultifieldParser:
#   negligence AND "duty of care"     → AND operator
#   Donoghue OR Stevenson             → OR operator
#   "neighbour principle"             → phrase search
#   court:Lords negligence            → field-scoped term

def search(query: str, alpha: float = 0.7, top_k: int = 10):
    # ── 1. BM25 via Whoosh ────────────────────────────────────────────
    with ix.searcher() as searcher:
        hits = searcher.search(qp.parse(query), limit=top_k * 2)
        bm25_raw = {h["id"]: h.score for h in hits}

    # ── 2. Query routing ──────────────────────────────────────────────
    # If BM25 is already confident (exact citation or case name match),
    # bypass vector search entirely — keyword wins.
    if bm25_raw and max(bm25_raw.values()) > HIGH_CONFIDENCE_THRESHOLD:
        return build_results(bm25_raw, mode="keyword")

    # ── 3. Vector search (paragraph-level) ────────────────────────────
    # Embed the query using the same model used at index time.
    q_vec = model.encode(query)          # shape: (384,)

    cosine_raw = {}
    for case_id, paragraphs in paragraph_store.items():
        # A case's score = the best paragraph match (max pooling).
        # Paragraph-level search finds relevant passages inside long judgments.
        scores = [1 - cosine(q_vec, p.embedding) for p in paragraphs]
        cosine_raw[case_id] = max(scores)

    # ── 4. Normalise both distributions to [0, 1] ──────────────────────
    def norm(d):
        lo, hi = min(d.values()), max(d.values())
        return {k: (v - lo) / (hi - lo + 1e-9) for k, v in d.items()}

    bm25_n   = norm(bm25_raw)
    cosine_n = norm(cosine_raw)

    # ── 5. Weighted linear combination ────────────────────────────────
    #   score = α · norm(BM25) + (1 − α) · norm(cosine)
    #   Default α = 0.7 favours keyword precision;
    #   lower α shifts weight toward semantic recall.
    all_ids = set(bm25_n) | set(cosine_n)
    hybrid = {
        cid: alpha * bm25_n.get(cid, 0) + (1 - alpha) * cosine_n.get(cid, 0)
        for cid in all_ids
    }
    ranked = sorted(hybrid, key=hybrid.get, reverse=True)[:top_k]
    return build_results({cid: hybrid[cid] for cid in ranked}, mode="hybrid")

def generate_snippet(text: str, terms: list[str], window: int = 50) -> str:
    words = text.split()
    term_set = {t.lower() for t in terms}

    # Slide a window across the text, scoring each position by term density.
    # Higher density = more query terms in this region = more relevant excerpt.
    best_start, best_score = 0, 0
    for i in range(max(1, len(words) - window)):
        score = sum(
            1 for w in words[i : i + window]
            if w.lower().strip(".,;:()") in term_set
        )
        if score > best_score:
            best_start, best_score = i, score

    # Expand the window slightly and snap to sentence boundaries.
    excerpt = words[max(0, best_start - 8) : best_start + window + 8]
    text = " ".join(excerpt)

    # Wrap first sentence boundary at each end to keep the snippet coherent.
    text = re.sub(r'^[^.!?]*[.!?]\s*', '', text)  # trim leading partial sentence
    m = re.search(r'[.!?]', text[::-1])             # find last sentence end
    if m:
        text = text[: len(text) - m.start()]

    # Highlight matched terms with <mark> so the UI can render them.
    for term in terms:
        text = re.sub(
            rf'\b({re.escape(term)})\b', r'<mark>\1</mark>',
            text, flags=re.IGNORECASE
        )
    return f"…{text}…"

# Citations appear in multiple jurisdiction-specific formats.
# The regex union below covers the major Australian and UK patterns.
PATTERNS = [
    # Australian neutral citations   → [2001] HCA 14
    r'\[(?P<year>\d{4})\]\s*(?P<court>HCA|FCA|FCAFC|NSWCA|VSCA|QCA)\s*(?P<num>\d+)',

    # English series citations        → [1990] 2 AC 605  /  (1932) AC 562
    r'[\[\(](?P<year>\d{4})[\]\)]\s*(?:\d+\s*)?'
    r'(?P<series>AC|QB|Ch|WLR|All ER|EWCA Civ|EWCA Crim|UKSC)\s*(?P<page>\d+)',

    # Commonwealth Law Reports        → (2001) 207 CLR 562
    r'\((?P<year>\d{4})\)\s*(?P<vol>\d+)\s*CLR\s*(?P<page>\d+)',
]

def extract_citations(text: str) -> list[str]:
    found = []
    for pat in PATTERNS:
        for m in re.finditer(pat, text):
            found.append(normalise(m.group(0)))
    return list(dict.fromkeys(found))   # deduplicate, preserve order

# Edges are stored in Postgres for O(1) reverse lookup ("cited by"):
# CREATE TABLE citation_edges (
#     src_id TEXT REFERENCES cases(id),
#     tgt_citation TEXT,          -- raw citation string
#     tgt_id TEXT,                -- NULL if target not in corpus
#     PRIMARY KEY (src_id, tgt_citation)
# );

from fastapi import FastAPI
from pydantic import BaseModel
from typing import Literal

class SearchRequest(BaseModel):
    q: str
    mode: Literal["keyword", "semantic", "hybrid"] = "hybrid"
    courts: list[str] = []
    jurisdictions: list[str] = []
    year_from: int | None = None
    year_to: int | None = None
    page: int = 1
    per_page: int = 10

class CaseResult(BaseModel):
    id: str
    name: str
    citation: str
    court: str
    jurisdiction: str
    date: str
    snippet: str           # HTML with <mark> highlights
    score: float           # Normalised to [0, 1]
    score_breakdown: dict  # {"bm25": 0.82, "vector": 0.71, "final": 0.78}

class SearchResponse(BaseModel):
    total: int
    page: int
    results: list[CaseResult]
    facet_counts: dict     # {"court": {"House of Lords": 8, ...}, ...}
    mode_used: str         # "keyword" | "hybrid" — which path fired

@app.post("/search", response_model=SearchResponse)
async def search(req: SearchRequest):
    results = engine.search(
        query=req.q,
        mode=req.mode,
        filters={"courts": req.courts, "jurisdictions": req.jurisdictions,
                 "year_range": (req.year_from, req.year_to)},
    )
    return paginate(results, req.page, req.per_page)

@app.get("/cases/{case_id}")
async def get_case(case_id: str) -> FullCase:
    # Returns full text + metadata + citation graph for one case.
    return db.get_case(case_id)