Spaces:
Configuration error
Configuration error
File size: 4,317 Bytes
6733714 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
LETA Context Compressor — Stage 4 of TITAN architecture.
Reduces injected LLM context from the raw 80 KB chunk dump to a
targeted ~8 KB focused-excerpt block. The compressor:
1. Scores every chunk for query relevance (token overlap, no extra model call)
2. Extracts the best sentence window from each chunk
3. Caps total output at MAX_COMPRESSED_CHARS
This sits between retriever and synthesizer. build_context() is still
used for source-citation metadata; compress_context() provides the
dense factual block the LLM actually reasons over.
"""
import os
import re
from typing import List
MAX_COMPRESSED_CHARS_QA = 18000 # ~4 500 tokens — Q&A mode
MAX_COMPRESSED_CHARS_DRAFT = 20000 # ~5 000 tokens — draft mode needs full statutory text
MAX_EXCERPT_CHARS_QA = 1200 # per-chunk window Q&A — full statutory provisions fit
MAX_EXCERPT_CHARS_DRAFT = 1400 # per-chunk window draft — needs verbatim statutory paragraphs
MAX_CHUNKS_USED_QA = 20 # Q&A cap
MAX_CHUNKS_USED_DRAFT = 25 # draft uses all retrieved chunks
def _query_tokens(query: str) -> set:
"""Lower-cased word stems ≥ 3 chars from the query."""
return set(re.findall(r'\b[a-z]{3,}\b', query.lower()))
def _overlap_score(text: str, q_tokens: set) -> float:
"""Jaccard-like token overlap — fast, no embeddings needed."""
if not text or not q_tokens:
return 0.0
chunk_tokens = set(re.findall(r'\b[a-z]{3,}\b', text.lower()))
return len(chunk_tokens & q_tokens) / max(len(q_tokens), 1)
def _best_window(text: str, q_tokens: set, max_chars: int = MAX_EXCERPT_CHARS_QA) -> str:
"""
Split text into sentences and return the highest-scoring 2-sentence
window (± 1 sentence context around the best match).
Falls back to a straight head-truncation if no sentence found.
"""
sentences = re.split(r'(?<=[.!?])\s+', text.strip())
if len(sentences) <= 2:
return text[:max_chars]
best_idx = max(range(len(sentences)), key=lambda i: _overlap_score(sentences[i], q_tokens))
start = max(0, best_idx - 1)
end = min(len(sentences), best_idx + 2)
window = " ".join(sentences[start:end])
if len(window) > max_chars:
window = window[:max_chars] + "…"
return window
def compress_context(chunks: List[dict], query: str, is_draft: bool = False) -> str:
"""
Build a focused, compressed context string from retrieved chunks.
Draft mode uses larger limits so verbatim statutory text and full judgment
paragraphs fit — needed for 5000-word SCN replies.
Args:
chunks: Retrieved chunk dicts (as returned by retriever.search)
query: Original user query string
is_draft: True for SCN replies / appeals / drafting mode
Returns:
Compressed context string.
"""
if not chunks:
return "No relevant statutory documents retrieved."
max_total = MAX_COMPRESSED_CHARS_DRAFT if is_draft else MAX_COMPRESSED_CHARS_QA
max_excerpt = MAX_EXCERPT_CHARS_DRAFT if is_draft else MAX_EXCERPT_CHARS_QA
max_chunks = MAX_CHUNKS_USED_DRAFT if is_draft else MAX_CHUNKS_USED_QA
q_tokens = _query_tokens(query)
# Score each chunk: combine reranker score with query-overlap score
scored = []
for c in chunks:
rerank = float(c.get("_final_legal_score", c.get("_rerank_score", 0.0)) or 0.0)
text = c.get("text", "")
overlap = _overlap_score(text, q_tokens)
combined = 0.6 * rerank + 0.4 * overlap
scored.append((combined, c))
scored.sort(key=lambda x: x[0], reverse=True)
top = scored[:max_chunks]
parts: List[str] = []
total = 0
for rank, (score, chunk) in enumerate(top, 1):
if total >= max_total:
break
text = chunk.get("text", "").strip()
source = os.path.basename(chunk.get("source", "Unknown"))
page = chunk.get("page", "N/A")
excerpt = _best_window(text, q_tokens, max_chars=max_excerpt)
remaining = max_total - total
if len(excerpt) > remaining:
excerpt = excerpt[:remaining] + "…"
entry = f"[{rank}] {source} p.{page}\n{excerpt}"
parts.append(entry)
total += len(entry) + 4 # +4 for separator
return "\n\n---\n\n".join(parts)
|