Spaces:
Sleeping
Sleeping
Upload folder using huggingface_hub
Browse files- api.py +98 -27
- config.py +11 -1
- indra_engine.py +265 -0
- ingestion/entity_resolver.py +131 -0
- ingestion/gazette_ingester.py +145 -1
- ingestion/knowledge_extractor.py +428 -0
- ingestion/prompts.py +112 -0
- requirements.txt +3 -0
- tests/test_indra_engine.py +116 -0
- tests/test_knowledge_extractor.py +216 -0
api.py
CHANGED
|
@@ -20,6 +20,12 @@ from neuro_symbolic import run_neuro_symbolic_pipeline
|
|
| 20 |
from whatsapp.webhook import router as whatsapp_router
|
| 21 |
from config import settings
|
| 22 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 23 |
# --- SECURE KEYS ---
|
| 24 |
SUPABASE_URL = settings.SUPABASE_URL
|
| 25 |
SUPABASE_KEY = settings.SUPABASE_KEY
|
|
@@ -191,14 +197,31 @@ def rerank_gazette_chunks(query: str, chunks: list) -> list:
|
|
| 191 |
"""
|
| 192 |
Cross-encoder re-ranking for gazette search results.
|
| 193 |
Takes top 30 RRF candidates (gazette chunks are denser than scheme chunks),
|
| 194 |
-
runs them through Ettin reranker,
|
| 195 |
-
|
| 196 |
-
|
| 197 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 198 |
"""
|
| 199 |
if not chunks:
|
| 200 |
return []
|
| 201 |
|
|
|
|
|
|
|
| 202 |
# Gazette-specific: wider pool (30) because legislative text
|
| 203 |
# has higher semantic density than scheme descriptions
|
| 204 |
candidates = chunks[:30]
|
|
@@ -212,7 +235,16 @@ def rerank_gazette_chunks(query: str, chunks: list) -> list:
|
|
| 212 |
candidates[i]['cross_encoder_score'] = float(score)
|
| 213 |
|
| 214 |
candidates.sort(key=lambda x: x['cross_encoder_score'], reverse=True)
|
| 215 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 216 |
|
| 217 |
# --- ENDPOINTS ---
|
| 218 |
@app.get("/health")
|
|
@@ -513,15 +545,16 @@ async def get_coverage_gaps(request: Request, limit: int = 20):
|
|
| 513 |
@limiter.limit("15/minute")
|
| 514 |
async def gazette_search(request: Request, query: GazetteSearchQuery):
|
| 515 |
"""
|
| 516 |
-
Gazette Vault
|
| 517 |
|
| 518 |
Pipeline:
|
| 519 |
1. Embed query with Nomic (768-dim) using search_query: prefix
|
| 520 |
-
2. Call
|
| 521 |
-
3.
|
| 522 |
-
4.
|
|
|
|
| 523 |
"""
|
| 524 |
-
print(f"📜 Gazette search: '{query.query}' | Type: {query.gazette_type} | State: {query.state}")
|
| 525 |
|
| 526 |
try:
|
| 527 |
# Step 1: Embed query
|
|
@@ -531,21 +564,19 @@ async def gazette_search(request: Request, query: GazetteSearchQuery):
|
|
| 531 |
normalize_embeddings=True
|
| 532 |
).tolist()
|
| 533 |
|
| 534 |
-
# Step 2: Call
|
| 535 |
rpc_params: Dict[str, Any] = {
|
| 536 |
-
"query_text": query.query,
|
| 537 |
"query_embedding": query_embedding,
|
| 538 |
-
"
|
| 539 |
-
"k_smoothing": 60,
|
| 540 |
}
|
| 541 |
|
| 542 |
-
# Add optional filters
|
| 543 |
if query.gazette_type:
|
| 544 |
rpc_params["filter_gazette_type"] = query.gazette_type
|
| 545 |
if query.state:
|
| 546 |
rpc_params["filter_state"] = query.state
|
| 547 |
|
| 548 |
-
result = supabase.rpc("
|
| 549 |
raw_results = cast(List[Dict[str, Any]], result.data or [])
|
| 550 |
|
| 551 |
if not raw_results:
|
|
@@ -553,16 +584,51 @@ async def gazette_search(request: Request, query: GazetteSearchQuery):
|
|
| 553 |
"query": query.query,
|
| 554 |
"results": [],
|
| 555 |
"total": 0,
|
| 556 |
-
"pipeline": "
|
| 557 |
}
|
| 558 |
|
| 559 |
-
|
| 560 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 561 |
|
| 562 |
-
# Step 4:
|
|
|
|
|
|
|
|
|
|
| 563 |
reranked = reranked[:query.limit]
|
| 564 |
|
| 565 |
-
# Step
|
| 566 |
results = []
|
| 567 |
for chunk in reranked:
|
| 568 |
results.append({
|
|
@@ -574,26 +640,31 @@ async def gazette_search(request: Request, query: GazetteSearchQuery):
|
|
| 574 |
"gazette_type": chunk.get("gazette_type", "central"),
|
| 575 |
"issuing_authority": chunk.get("issuing_authority"),
|
| 576 |
"notification_date": str(chunk.get("notification_date", "")),
|
| 577 |
-
"cross_encoder_score": round(chunk.get("cross_encoder_score", 0.0), 4),
|
| 578 |
-
"
|
|
|
|
| 579 |
})
|
| 580 |
|
| 581 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 582 |
|
| 583 |
return {
|
| 584 |
"query": query.query,
|
| 585 |
"results": results,
|
| 586 |
"total": len(results),
|
| 587 |
-
"pipeline": "
|
| 588 |
}
|
| 589 |
|
| 590 |
except Exception as e:
|
| 591 |
-
print(f"❌ Gazette search error: {e}")
|
| 592 |
return {
|
| 593 |
"query": query.query,
|
| 594 |
"results": [],
|
| 595 |
"total": 0,
|
| 596 |
"error": str(e),
|
| 597 |
-
"pipeline": "
|
| 598 |
}
|
| 599 |
|
|
|
|
| 20 |
from whatsapp.webhook import router as whatsapp_router
|
| 21 |
from config import settings
|
| 22 |
|
| 23 |
+
# ── PROJECT INDRA: Poincaré Re-Ranking Engine (Sprint 29) ─────
|
| 24 |
+
# Instantiated globally at startup to claim memory once.
|
| 25 |
+
# batch_size=400 matches the get_indra_candidates() oversample_limit.
|
| 26 |
+
from indra_engine import IndraProjectionEngine
|
| 27 |
+
indra_engine = IndraProjectionEngine(batch_size=400, dimensions=768)
|
| 28 |
+
|
| 29 |
# --- SECURE KEYS ---
|
| 30 |
SUPABASE_URL = settings.SUPABASE_URL
|
| 31 |
SUPABASE_KEY = settings.SUPABASE_KEY
|
|
|
|
| 197 |
"""
|
| 198 |
Cross-encoder re-ranking for gazette search results.
|
| 199 |
Takes top 30 RRF candidates (gazette chunks are denser than scheme chunks),
|
| 200 |
+
runs them through Ettin reranker, applies a hard quality gate,
|
| 201 |
+
and returns ONLY genuinely relevant results.
|
| 202 |
+
|
| 203 |
+
QUALITY GATE (5.0 raw logit):
|
| 204 |
+
─────────────────────────────────────────────────────
|
| 205 |
+
The Ettin reranker outputs raw logits roughly in [-10, +15].
|
| 206 |
+
Empirical calibration from GovBridge test corpus:
|
| 207 |
+
|
| 208 |
+
Score ≥ 8.0 → Highly relevant (exact topic match)
|
| 209 |
+
Score 6-8 → Relevant (same domain, related content)
|
| 210 |
+
Score 5-6 → Borderline (tangential relevance)
|
| 211 |
+
Score 3-5 → NOISE. Garbage queries, random names, and
|
| 212 |
+
unrelated terms consistently score here.
|
| 213 |
+
Score < 3 → Definitively irrelevant.
|
| 214 |
+
|
| 215 |
+
FLOOR = 5.0 eliminates ALL noise while preserving every
|
| 216 |
+
genuinely relevant result. A World No. 1 system NEVER shows
|
| 217 |
+
irrelevant results — it shows "No results found" instead.
|
| 218 |
+
─────────────────────────────────────────────────────
|
| 219 |
"""
|
| 220 |
if not chunks:
|
| 221 |
return []
|
| 222 |
|
| 223 |
+
RELEVANCE_FLOOR = 5.0 # Hard quality gate — no noise passes
|
| 224 |
+
|
| 225 |
# Gazette-specific: wider pool (30) because legislative text
|
| 226 |
# has higher semantic density than scheme descriptions
|
| 227 |
candidates = chunks[:30]
|
|
|
|
| 235 |
candidates[i]['cross_encoder_score'] = float(score)
|
| 236 |
|
| 237 |
candidates.sort(key=lambda x: x['cross_encoder_score'], reverse=True)
|
| 238 |
+
|
| 239 |
+
# Hard quality gate: Remove all results below the relevance floor
|
| 240 |
+
filtered = [c for c in candidates if c['cross_encoder_score'] >= RELEVANCE_FLOOR]
|
| 241 |
+
dropped = len(candidates) - len(filtered)
|
| 242 |
+
if dropped > 0:
|
| 243 |
+
print(f" 🔽 Quality gate: {dropped}/{len(candidates)} results below floor ({RELEVANCE_FLOOR})")
|
| 244 |
+
if not filtered:
|
| 245 |
+
print(f" ⛔ ALL results below quality gate — returning empty (best was {candidates[0]['cross_encoder_score']:.2f})")
|
| 246 |
+
|
| 247 |
+
return filtered[:10]
|
| 248 |
|
| 249 |
# --- ENDPOINTS ---
|
| 250 |
@app.get("/health")
|
|
|
|
| 545 |
@limiter.limit("15/minute")
|
| 546 |
async def gazette_search(request: Request, query: GazetteSearchQuery):
|
| 547 |
"""
|
| 548 |
+
Gazette Vault Search — PROJECT INDRA Pipeline (Sprint 29).
|
| 549 |
|
| 550 |
Pipeline:
|
| 551 |
1. Embed query with Nomic (768-dim) using search_query: prefix
|
| 552 |
+
2. Call get_indra_candidates RPC (oversampled HNSW, 400 candidates + embeddings)
|
| 553 |
+
3. Poincaré ball projection → hyperbolic distance re-ranking
|
| 554 |
+
4. Cross-encoder re-rank top 30 with Ettin
|
| 555 |
+
5. Return page-pinned results to frontend GazetteViewer
|
| 556 |
"""
|
| 557 |
+
print(f"📜 Gazette search [INDRA]: '{query.query}' | Type: {query.gazette_type} | State: {query.state}")
|
| 558 |
|
| 559 |
try:
|
| 560 |
# Step 1: Embed query
|
|
|
|
| 564 |
normalize_embeddings=True
|
| 565 |
).tolist()
|
| 566 |
|
| 567 |
+
# Step 2: Call INDRA oversampling RPC
|
| 568 |
rpc_params: Dict[str, Any] = {
|
|
|
|
| 569 |
"query_embedding": query_embedding,
|
| 570 |
+
"oversample_limit": 400,
|
|
|
|
| 571 |
}
|
| 572 |
|
| 573 |
+
# Add optional filters
|
| 574 |
if query.gazette_type:
|
| 575 |
rpc_params["filter_gazette_type"] = query.gazette_type
|
| 576 |
if query.state:
|
| 577 |
rpc_params["filter_state"] = query.state
|
| 578 |
|
| 579 |
+
result = supabase.rpc("get_indra_candidates", rpc_params).execute()
|
| 580 |
raw_results = cast(List[Dict[str, Any]], result.data or [])
|
| 581 |
|
| 582 |
if not raw_results:
|
|
|
|
| 584 |
"query": query.query,
|
| 585 |
"results": [],
|
| 586 |
"total": 0,
|
| 587 |
+
"pipeline": "indra_poincare + ettin_reranker"
|
| 588 |
}
|
| 589 |
|
| 590 |
+
n_candidates = len(raw_results)
|
| 591 |
+
print(f" 📊 INDRA: {n_candidates} Euclidean candidates retrieved")
|
| 592 |
+
|
| 593 |
+
# Step 3: Poincaré projection & hyperbolic re-ranking
|
| 594 |
+
import numpy as np
|
| 595 |
+
|
| 596 |
+
# Extract embedding vectors from RPC results
|
| 597 |
+
query_vec = np.array(query_embedding, dtype=np.float32)
|
| 598 |
+
candidate_embeddings = np.array(
|
| 599 |
+
[r["embedding"] for r in raw_results],
|
| 600 |
+
dtype=np.float32
|
| 601 |
+
)
|
| 602 |
+
|
| 603 |
+
# Project and rank by hyperbolic distance
|
| 604 |
+
sorted_indices = indra_engine.project_and_rank(
|
| 605 |
+
query_vec, candidate_embeddings, n_candidates
|
| 606 |
+
)
|
| 607 |
+
|
| 608 |
+
# Compute distances for logging
|
| 609 |
+
hyp_distances = indra_engine.compute_poincare_distances(
|
| 610 |
+
query_vec, candidate_embeddings, n_candidates
|
| 611 |
+
)
|
| 612 |
+
|
| 613 |
+
# Re-order results by hyperbolic distance (ascending = most relevant first)
|
| 614 |
+
hyperbolic_ranked = []
|
| 615 |
+
for idx in sorted_indices[:30]: # Take top 30 for cross-encoder stage
|
| 616 |
+
entry = raw_results[int(idx)].copy()
|
| 617 |
+
entry["hyperbolic_distance"] = float(hyp_distances[int(idx)])
|
| 618 |
+
# Remove embedding from the result (not needed downstream, saves memory)
|
| 619 |
+
entry.pop("embedding", None)
|
| 620 |
+
hyperbolic_ranked.append(entry)
|
| 621 |
+
|
| 622 |
+
print(f" 🔮 Poincaré: top-1 distance={hyperbolic_ranked[0]['hyperbolic_distance']:.4f}, "
|
| 623 |
+
f"top-30 distance={hyperbolic_ranked[-1]['hyperbolic_distance']:.4f}")
|
| 624 |
|
| 625 |
+
# Step 4: Cross-encoder re-rank (Ettin verifies semantic coherence)
|
| 626 |
+
reranked = rerank_gazette_chunks(query.query, hyperbolic_ranked)
|
| 627 |
+
|
| 628 |
+
# Step 5: Trim to requested limit
|
| 629 |
reranked = reranked[:query.limit]
|
| 630 |
|
| 631 |
+
# Step 6: Format response with page-pinned results
|
| 632 |
results = []
|
| 633 |
for chunk in reranked:
|
| 634 |
results.append({
|
|
|
|
| 640 |
"gazette_type": chunk.get("gazette_type", "central"),
|
| 641 |
"issuing_authority": chunk.get("issuing_authority"),
|
| 642 |
"notification_date": str(chunk.get("notification_date", "")),
|
| 643 |
+
"cross_encoder_score": round(float(chunk.get("cross_encoder_score", 0.0)), 4),
|
| 644 |
+
"hyperbolic_distance": round(float(chunk.get("hyperbolic_distance", 0.0)), 6),
|
| 645 |
+
"rrf_score": round(float(chunk.get("rrf_score", chunk.get("l2_distance", 0.0))), 6),
|
| 646 |
})
|
| 647 |
|
| 648 |
+
if results:
|
| 649 |
+
print(f" ✅ Returning {len(results)} gazette results "
|
| 650 |
+
f"(best CE score: {results[0]['cross_encoder_score']:.4f})")
|
| 651 |
+
else:
|
| 652 |
+
print(" ⛔ No results survived quality gate")
|
| 653 |
|
| 654 |
return {
|
| 655 |
"query": query.query,
|
| 656 |
"results": results,
|
| 657 |
"total": len(results),
|
| 658 |
+
"pipeline": "indra_poincare + ettin_reranker"
|
| 659 |
}
|
| 660 |
|
| 661 |
except Exception as e:
|
| 662 |
+
print(f"❌ Gazette search [INDRA] error: {e}")
|
| 663 |
return {
|
| 664 |
"query": query.query,
|
| 665 |
"results": [],
|
| 666 |
"total": 0,
|
| 667 |
"error": str(e),
|
| 668 |
+
"pipeline": "indra_poincare + ettin_reranker"
|
| 669 |
}
|
| 670 |
|
config.py
CHANGED
|
@@ -16,6 +16,16 @@ class Settings(BaseSettings):
|
|
| 16 |
# ── LLM ──────────────────────────────────────────────
|
| 17 |
GROQ_API_KEY: str = Field(default="")
|
| 18 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19 |
# ── Admin ────────────────────────────────────────────
|
| 20 |
ADMIN_SECRET: str = Field(default="change-this-in-production")
|
| 21 |
|
|
@@ -39,7 +49,7 @@ class Settings(BaseSettings):
|
|
| 39 |
DUAL_WRITE_EMBEDDINGS: str = Field(default="false")
|
| 40 |
|
| 41 |
model_config = SettingsConfigDict(
|
| 42 |
-
env_file=".env",
|
| 43 |
env_file_encoding="utf-8",
|
| 44 |
extra="ignore"
|
| 45 |
)
|
|
|
|
| 16 |
# ── LLM ──────────────────────────────────────────────
|
| 17 |
GROQ_API_KEY: str = Field(default="")
|
| 18 |
|
| 19 |
+
# ── Knowledge Extraction (Sprint 30 — INDRA Phase 1.5) ───
|
| 20 |
+
GROQ_EXTRACTION_MODEL: str = Field(
|
| 21 |
+
default="llama-3.3-70b-versatile",
|
| 22 |
+
description="Groq model for Knowledge Triplet extraction"
|
| 23 |
+
)
|
| 24 |
+
GROQ_MAX_CONCURRENCY: int = Field(
|
| 25 |
+
default=5,
|
| 26 |
+
description="Max concurrent Groq API calls during extraction"
|
| 27 |
+
)
|
| 28 |
+
|
| 29 |
# ── Admin ────────────────────────────────────────────
|
| 30 |
ADMIN_SECRET: str = Field(default="change-this-in-production")
|
| 31 |
|
|
|
|
| 49 |
DUAL_WRITE_EMBEDDINGS: str = Field(default="false")
|
| 50 |
|
| 51 |
model_config = SettingsConfigDict(
|
| 52 |
+
env_file=(".env", "../.env"),
|
| 53 |
env_file_encoding="utf-8",
|
| 54 |
extra="ignore"
|
| 55 |
)
|
indra_engine.py
ADDED
|
@@ -0,0 +1,265 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
GovBridge India — INDRA Poincaré Projection Engine (Sprint 29)
|
| 3 |
+
|
| 4 |
+
PROJECT INDRA Phase 1: Non-Euclidean Civic Topologies
|
| 5 |
+
|
| 6 |
+
This engine implements the HyEm (Hyperbolic Embedding) paradigm:
|
| 7 |
+
1. Receives oversampled Euclidean candidates from HNSW (400 vectors)
|
| 8 |
+
2. Projects them into the Poincaré ball model via conformal exponential map
|
| 9 |
+
3. Computes hyperbolic distances (arcosh-based) between query and candidates
|
| 10 |
+
4. Returns candidates sorted by hyperbolic distance (ascending)
|
| 11 |
+
|
| 12 |
+
ARCHITECTURAL CONSTRAINTS:
|
| 13 |
+
- ALL NumPy buffers are pre-allocated at __init__ time.
|
| 14 |
+
- ZERO heap allocations inside compute methods (GC pressure = 0).
|
| 15 |
+
- IEEE 754 boundary enforcement: all vectors clamped to R_MAX = 1 - 1e-3.
|
| 16 |
+
- Denominator floor: DENOM_MIN = 1e-15 prevents division-by-zero.
|
| 17 |
+
- Negative norm² clamping: np.maximum(diff_norm2, 0.0) prevents NaN from
|
| 18 |
+
floating-point variance underflow on identical vectors.
|
| 19 |
+
|
| 20 |
+
MATHEMATICAL FOUNDATION:
|
| 21 |
+
exp_0(v) = tanh(||v||) * (v / ||v||) [Conformal exponential map at origin]
|
| 22 |
+
d_H(u,v) = arcosh(1 + 2 * ||u-v||² / ((1-||u||²)(1-||v||²)))
|
| 23 |
+
|
| 24 |
+
USAGE:
|
| 25 |
+
from indra_engine import IndraProjectionEngine
|
| 26 |
+
engine = IndraProjectionEngine(batch_size=400, dimensions=768)
|
| 27 |
+
distances = engine.project_and_rank(query_vec, candidate_matrix)
|
| 28 |
+
"""
|
| 29 |
+
|
| 30 |
+
import numpy as np
|
| 31 |
+
from numpy.typing import NDArray
|
| 32 |
+
from typing import Optional
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class IndraProjectionEngine:
|
| 36 |
+
"""
|
| 37 |
+
Zero-allocation Poincaré ball projection engine.
|
| 38 |
+
|
| 39 |
+
Pre-allocates all working buffers at construction time.
|
| 40 |
+
No heap allocations occur during compute_poincare_distances()
|
| 41 |
+
or map_to_poincare_ball_inplace().
|
| 42 |
+
"""
|
| 43 |
+
|
| 44 |
+
# ── Numerical Constants ──────────────────────────────────
|
| 45 |
+
# Maximum radius in Poincaré ball. Vectors beyond this are
|
| 46 |
+
# clipped to prevent arcosh(1 + 2 * inf / 0) = NaN.
|
| 47 |
+
R_MAX: float = 1.0 - 1e-3 # 0.999
|
| 48 |
+
|
| 49 |
+
# Minimum denominator value. Prevents div-by-zero when a
|
| 50 |
+
# vector lies exactly on the ball boundary (||v|| ≈ 1.0).
|
| 51 |
+
DENOM_MIN: float = 1e-15
|
| 52 |
+
|
| 53 |
+
# Minimum norm for normalization. Vectors with ||v|| < this
|
| 54 |
+
# are treated as zero vectors and mapped to the origin.
|
| 55 |
+
NORM_MIN: float = 1e-8
|
| 56 |
+
|
| 57 |
+
def __init__(self, batch_size: int = 400, dimensions: int = 768):
|
| 58 |
+
"""
|
| 59 |
+
Pre-allocate all working memory.
|
| 60 |
+
|
| 61 |
+
Args:
|
| 62 |
+
batch_size: Maximum number of candidate vectors per query.
|
| 63 |
+
dimensions: Embedding dimensionality (768 for Nomic).
|
| 64 |
+
"""
|
| 65 |
+
self.batch_size = batch_size
|
| 66 |
+
self.dimensions = dimensions
|
| 67 |
+
|
| 68 |
+
# ── Pre-allocated buffers ────────────────────────────
|
| 69 |
+
# Candidate matrix projected into Poincaré ball
|
| 70 |
+
self._poincare_candidates = np.zeros(
|
| 71 |
+
(batch_size, dimensions), dtype=np.float32
|
| 72 |
+
)
|
| 73 |
+
|
| 74 |
+
# Query vector projected into Poincaré ball
|
| 75 |
+
self._poincare_query = np.zeros(dimensions, dtype=np.float32)
|
| 76 |
+
|
| 77 |
+
# Scratch buffers for distance computation
|
| 78 |
+
self._norms_sq = np.zeros(batch_size, dtype=np.float64)
|
| 79 |
+
self._query_norm_sq = np.float64(0.0)
|
| 80 |
+
self._diff = np.zeros(
|
| 81 |
+
(batch_size, dimensions), dtype=np.float32
|
| 82 |
+
)
|
| 83 |
+
self._diff_norm_sq = np.zeros(batch_size, dtype=np.float64)
|
| 84 |
+
self._distances = np.zeros(batch_size, dtype=np.float64)
|
| 85 |
+
self._denom = np.zeros(batch_size, dtype=np.float64)
|
| 86 |
+
|
| 87 |
+
# Buffer for norms during projection
|
| 88 |
+
self._proj_norms = np.zeros(batch_size, dtype=np.float32)
|
| 89 |
+
|
| 90 |
+
def map_to_poincare_ball_inplace(
|
| 91 |
+
self,
|
| 92 |
+
vectors: NDArray[np.float32],
|
| 93 |
+
out: NDArray[np.float32],
|
| 94 |
+
norms_buf: Optional[NDArray[np.float32]] = None,
|
| 95 |
+
) -> None:
|
| 96 |
+
"""
|
| 97 |
+
Conformal exponential map at the origin: exp_0(v) = tanh(||v||) * (v / ||v||)
|
| 98 |
+
|
| 99 |
+
Projects Euclidean vectors into the Poincaré ball IN-PLACE
|
| 100 |
+
(writes to the `out` buffer). No allocations.
|
| 101 |
+
|
| 102 |
+
Args:
|
| 103 |
+
vectors: Input Euclidean vectors, shape (N, D) or (D,).
|
| 104 |
+
out: Output buffer, same shape as vectors.
|
| 105 |
+
norms_buf: Pre-allocated buffer for norms. If None, uses self._proj_norms.
|
| 106 |
+
"""
|
| 107 |
+
is_1d = vectors.ndim == 1
|
| 108 |
+
|
| 109 |
+
if is_1d:
|
| 110 |
+
# Single vector case (query)
|
| 111 |
+
norm = np.linalg.norm(vectors).astype(np.float32)
|
| 112 |
+
if norm < self.NORM_MIN:
|
| 113 |
+
out[:] = 0.0
|
| 114 |
+
return
|
| 115 |
+
scale = np.float32(np.tanh(norm) / norm)
|
| 116 |
+
np.multiply(vectors, scale, out=out)
|
| 117 |
+
# Radial boundary clip
|
| 118 |
+
out_norm = np.linalg.norm(out).astype(np.float32)
|
| 119 |
+
if out_norm > self.R_MAX:
|
| 120 |
+
np.multiply(out, np.float32(self.R_MAX / out_norm), out=out)
|
| 121 |
+
else:
|
| 122 |
+
# Batch case (candidates)
|
| 123 |
+
n = vectors.shape[0]
|
| 124 |
+
if norms_buf is None:
|
| 125 |
+
norms_buf = self._proj_norms
|
| 126 |
+
|
| 127 |
+
# Compute L2 norms: ||v_i||
|
| 128 |
+
np.einsum('ij,ij->i', vectors[:n], vectors[:n], out=norms_buf[:n])
|
| 129 |
+
np.sqrt(norms_buf[:n], out=norms_buf[:n])
|
| 130 |
+
|
| 131 |
+
for i in range(n):
|
| 132 |
+
norm_val = norms_buf[i]
|
| 133 |
+
if norm_val < self.NORM_MIN:
|
| 134 |
+
out[i, :] = 0.0
|
| 135 |
+
else:
|
| 136 |
+
scale = np.float32(np.tanh(norm_val) / norm_val)
|
| 137 |
+
np.multiply(vectors[i], scale, out=out[i])
|
| 138 |
+
|
| 139 |
+
# Radial boundary clip per vector
|
| 140 |
+
out_norm = np.linalg.norm(out[i]).astype(np.float32)
|
| 141 |
+
if out_norm > self.R_MAX:
|
| 142 |
+
np.multiply(
|
| 143 |
+
out[i],
|
| 144 |
+
np.float32(self.R_MAX / out_norm),
|
| 145 |
+
out=out[i]
|
| 146 |
+
)
|
| 147 |
+
|
| 148 |
+
def compute_poincare_distances(
|
| 149 |
+
self,
|
| 150 |
+
query_vec: NDArray[np.float32],
|
| 151 |
+
candidate_matrix: NDArray[np.float32],
|
| 152 |
+
n_candidates: int,
|
| 153 |
+
) -> NDArray[np.float64]:
|
| 154 |
+
"""
|
| 155 |
+
Compute hyperbolic distances between query and candidates in
|
| 156 |
+
the Poincaré ball model.
|
| 157 |
+
|
| 158 |
+
Formula:
|
| 159 |
+
d_H(u, v) = arcosh(1 + 2 * ||u - v||² / ((1 - ||u||²) * (1 - ||v||²)))
|
| 160 |
+
|
| 161 |
+
All computations use pre-allocated buffers. Zero heap allocations.
|
| 162 |
+
|
| 163 |
+
Args:
|
| 164 |
+
query_vec: Euclidean query vector, shape (D,).
|
| 165 |
+
candidate_matrix: Euclidean candidate vectors, shape (N, D).
|
| 166 |
+
n_candidates: Actual number of candidates (may be < batch_size).
|
| 167 |
+
|
| 168 |
+
Returns:
|
| 169 |
+
Hyperbolic distances array, shape (n_candidates,).
|
| 170 |
+
"""
|
| 171 |
+
n = min(n_candidates, self.batch_size)
|
| 172 |
+
|
| 173 |
+
# ── Step 1: Project query into Poincaré ball ─────────
|
| 174 |
+
self.map_to_poincare_ball_inplace(
|
| 175 |
+
query_vec, self._poincare_query
|
| 176 |
+
)
|
| 177 |
+
|
| 178 |
+
# ── Step 2: Project candidates into Poincaré ball ────
|
| 179 |
+
self.map_to_poincare_ball_inplace(
|
| 180 |
+
candidate_matrix[:n],
|
| 181 |
+
self._poincare_candidates[:n],
|
| 182 |
+
norms_buf=self._proj_norms
|
| 183 |
+
)
|
| 184 |
+
|
| 185 |
+
# ── Step 3: Compute ||u||² (query norm squared) ──────
|
| 186 |
+
self._query_norm_sq = np.float64(
|
| 187 |
+
np.dot(self._poincare_query, self._poincare_query)
|
| 188 |
+
)
|
| 189 |
+
|
| 190 |
+
# ── Step 4: Compute ||v_i||² (candidate norms squared) ─
|
| 191 |
+
np.einsum(
|
| 192 |
+
'ij,ij->i',
|
| 193 |
+
self._poincare_candidates[:n].astype(np.float64),
|
| 194 |
+
self._poincare_candidates[:n].astype(np.float64),
|
| 195 |
+
out=self._norms_sq[:n]
|
| 196 |
+
)
|
| 197 |
+
|
| 198 |
+
# ── Step 5: Compute ||u - v_i||² ─────────────────────
|
| 199 |
+
np.subtract(
|
| 200 |
+
self._poincare_candidates[:n],
|
| 201 |
+
self._poincare_query,
|
| 202 |
+
out=self._diff[:n]
|
| 203 |
+
)
|
| 204 |
+
np.einsum(
|
| 205 |
+
'ij,ij->i',
|
| 206 |
+
self._diff[:n].astype(np.float64),
|
| 207 |
+
self._diff[:n].astype(np.float64),
|
| 208 |
+
out=self._diff_norm_sq[:n]
|
| 209 |
+
)
|
| 210 |
+
|
| 211 |
+
# GUARD: Floating-point variance underflow.
|
| 212 |
+
# Identical vectors can produce tiny negative values like -1e-7
|
| 213 |
+
# due to IEEE 754 rounding. Clamp to zero.
|
| 214 |
+
np.maximum(self._diff_norm_sq[:n], 0.0, out=self._diff_norm_sq[:n])
|
| 215 |
+
|
| 216 |
+
# ── Step 6: Compute denominator ──────────────────────
|
| 217 |
+
# denom = (1 - ||u||²) * (1 - ||v_i||²)
|
| 218 |
+
np.subtract(1.0, self._norms_sq[:n], out=self._denom[:n])
|
| 219 |
+
np.multiply(
|
| 220 |
+
self._denom[:n],
|
| 221 |
+
(1.0 - self._query_norm_sq),
|
| 222 |
+
out=self._denom[:n]
|
| 223 |
+
)
|
| 224 |
+
|
| 225 |
+
# GUARD: Denominator floor. If a vector lies exactly on the
|
| 226 |
+
# boundary (||v|| ≈ 1.0), denom → 0, causing infinity.
|
| 227 |
+
np.maximum(self._denom[:n], self.DENOM_MIN, out=self._denom[:n])
|
| 228 |
+
|
| 229 |
+
# ── Step 7: Compute arcosh argument ──────────────────
|
| 230 |
+
# arg = 1 + 2 * ||u - v||² / denom
|
| 231 |
+
np.divide(self._diff_norm_sq[:n], self._denom[:n], out=self._distances[:n])
|
| 232 |
+
np.multiply(self._distances[:n], 2.0, out=self._distances[:n])
|
| 233 |
+
np.add(self._distances[:n], 1.0, out=self._distances[:n])
|
| 234 |
+
|
| 235 |
+
# GUARD: arcosh domain enforcement. arcosh(x) requires x >= 1.
|
| 236 |
+
# Clamp to exactly 1.0 (distance = 0) if numerical error
|
| 237 |
+
# pushes below 1.0.
|
| 238 |
+
np.maximum(self._distances[:n], 1.0, out=self._distances[:n])
|
| 239 |
+
|
| 240 |
+
# ── Step 8: Final hyperbolic distance ────────────────
|
| 241 |
+
np.arccosh(self._distances[:n], out=self._distances[:n])
|
| 242 |
+
|
| 243 |
+
return self._distances[:n].copy() # Return a copy (safe for caller)
|
| 244 |
+
|
| 245 |
+
def project_and_rank(
|
| 246 |
+
self,
|
| 247 |
+
query_vec: NDArray[np.float32],
|
| 248 |
+
candidate_matrix: NDArray[np.float32],
|
| 249 |
+
n_candidates: int,
|
| 250 |
+
) -> NDArray[np.intp]:
|
| 251 |
+
"""
|
| 252 |
+
Full pipeline: project + compute distances + return sorted indices.
|
| 253 |
+
|
| 254 |
+
Args:
|
| 255 |
+
query_vec: Euclidean query embedding, shape (D,).
|
| 256 |
+
candidate_matrix: Euclidean candidate embeddings, shape (N, D).
|
| 257 |
+
n_candidates: Actual number of candidates.
|
| 258 |
+
|
| 259 |
+
Returns:
|
| 260 |
+
Array of indices sorted by ascending hyperbolic distance.
|
| 261 |
+
"""
|
| 262 |
+
distances = self.compute_poincare_distances(
|
| 263 |
+
query_vec, candidate_matrix, n_candidates
|
| 264 |
+
)
|
| 265 |
+
return np.argsort(distances)
|
ingestion/entity_resolver.py
ADDED
|
@@ -0,0 +1,131 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
GovBridge India — Entity Resolution for Knowledge Graph Edges
|
| 3 |
+
Sprint 30: PROJECT INDRA Phase 1.5
|
| 4 |
+
|
| 5 |
+
PURPOSE:
|
| 6 |
+
For each extracted triplet's "object" entity (e.g., "Information Technology
|
| 7 |
+
Rules, 2011"), this module searches gazette_chunks to find a matching
|
| 8 |
+
document UUID. If found, the target_node_id is set. If not found,
|
| 9 |
+
target_node_id is NULL and the raw string is stored in metadata.
|
| 10 |
+
|
| 11 |
+
RESOLUTION STRATEGY (3-tier):
|
| 12 |
+
1. Exact FTS match: websearch_to_tsquery('simple', entity) against chunk_text
|
| 13 |
+
2. Trigram similarity: pg_trgm similarity() > 0.4 against document_title
|
| 14 |
+
3. Unresolved: Insert with NULL target, store in metadata.unresolved_entity
|
| 15 |
+
|
| 16 |
+
PERFORMANCE:
|
| 17 |
+
- Uses existing GIN FTS index (idx_gazette_chunks_fts)
|
| 18 |
+
- Uses existing GIN trigram index (idx_gazette_chunks_trgm)
|
| 19 |
+
- Single query per entity, batched where possible
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
import sys
|
| 23 |
+
import os
|
| 24 |
+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 25 |
+
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
| 26 |
+
|
| 27 |
+
import logging
|
| 28 |
+
from typing import Optional
|
| 29 |
+
from uuid import UUID
|
| 30 |
+
|
| 31 |
+
logger = logging.getLogger("govbridge.entity_resolver")
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
async def resolve_entity_to_chunk_id(
|
| 35 |
+
supabase_client,
|
| 36 |
+
entity_name: str,
|
| 37 |
+
) -> Optional[str]:
|
| 38 |
+
"""
|
| 39 |
+
Attempt to resolve an entity name to a gazette_chunks UUID.
|
| 40 |
+
|
| 41 |
+
Resolution chain:
|
| 42 |
+
1. Search document_title for trigram similarity > 0.4
|
| 43 |
+
(This catches "IT Rules, 2011" matching "Information Technology Rules, 2011")
|
| 44 |
+
2. Fall back to FTS on chunk_text
|
| 45 |
+
3. Return None if unresolved
|
| 46 |
+
|
| 47 |
+
Args:
|
| 48 |
+
supabase_client: Initialized Supabase client
|
| 49 |
+
entity_name: The object entity string from the knowledge triplet
|
| 50 |
+
|
| 51 |
+
Returns:
|
| 52 |
+
UUID string if resolved, None if unresolved
|
| 53 |
+
"""
|
| 54 |
+
if not entity_name or len(entity_name) < 3:
|
| 55 |
+
return None
|
| 56 |
+
|
| 57 |
+
# Strategy 1: Trigram similarity search on document_title
|
| 58 |
+
# This is the most reliable because gazette titles are canonical
|
| 59 |
+
try:
|
| 60 |
+
result = supabase_client.rpc(
|
| 61 |
+
"resolve_entity_by_similarity",
|
| 62 |
+
{"entity_text": entity_name}
|
| 63 |
+
).execute()
|
| 64 |
+
|
| 65 |
+
if result.data and len(result.data) > 0:
|
| 66 |
+
match = result.data[0]
|
| 67 |
+
logger.info(
|
| 68 |
+
f" ✅ Resolved '{entity_name}' → "
|
| 69 |
+
f"'{match['document_title']}' (id={match['id']}, "
|
| 70 |
+
f"similarity={match.get('sim', 'N/A')})"
|
| 71 |
+
)
|
| 72 |
+
return match["id"]
|
| 73 |
+
except Exception as e:
|
| 74 |
+
logger.warning(f" Trigram resolution RPC failed: {e}")
|
| 75 |
+
|
| 76 |
+
# Strategy 2: FTS on chunk_text (broader but noisier)
|
| 77 |
+
try:
|
| 78 |
+
# Use ilike for simple substring matching as fallback
|
| 79 |
+
result = (
|
| 80 |
+
supabase_client.table("gazette_chunks")
|
| 81 |
+
.select("id, document_title")
|
| 82 |
+
.ilike("document_title", f"%{entity_name[:50]}%")
|
| 83 |
+
.limit(1)
|
| 84 |
+
.execute()
|
| 85 |
+
)
|
| 86 |
+
|
| 87 |
+
if result.data and len(result.data) > 0:
|
| 88 |
+
match = result.data[0]
|
| 89 |
+
logger.info(
|
| 90 |
+
f" ✅ Resolved '{entity_name}' via ilike → '{match['document_title']}'"
|
| 91 |
+
)
|
| 92 |
+
return match["id"]
|
| 93 |
+
except Exception as e:
|
| 94 |
+
logger.warning(f" ilike resolution failed: {e}")
|
| 95 |
+
|
| 96 |
+
logger.info(f" ❌ Unresolved: '{entity_name}' — will insert as dangling edge")
|
| 97 |
+
return None
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def build_edge_row(
|
| 101 |
+
source_chunk_id: str,
|
| 102 |
+
predicate: str,
|
| 103 |
+
target_chunk_id: Optional[str],
|
| 104 |
+
object_entity: str,
|
| 105 |
+
subject_entity: str,
|
| 106 |
+
global_context: dict,
|
| 107 |
+
) -> dict:
|
| 108 |
+
"""
|
| 109 |
+
Build a tensor_edges row for Supabase upsert.
|
| 110 |
+
|
| 111 |
+
If target_chunk_id is None (unresolved entity), the raw entity
|
| 112 |
+
string is stored in metadata.unresolved_entity for future backfill.
|
| 113 |
+
"""
|
| 114 |
+
metadata = {
|
| 115 |
+
"subject_entity": subject_entity,
|
| 116 |
+
"object_entity": object_entity,
|
| 117 |
+
"extraction_model": global_context.get("_extraction_model", "llama-3.3-70b-versatile"),
|
| 118 |
+
"document_title": global_context.get("document_title", "Unknown"),
|
| 119 |
+
}
|
| 120 |
+
|
| 121 |
+
if target_chunk_id is None:
|
| 122 |
+
metadata["unresolved_entity"] = object_entity
|
| 123 |
+
|
| 124 |
+
row = {
|
| 125 |
+
"source_node_id": source_chunk_id,
|
| 126 |
+
"target_node_id": target_chunk_id, # Can be None/NULL
|
| 127 |
+
"edge_type": predicate,
|
| 128 |
+
"metadata": metadata,
|
| 129 |
+
}
|
| 130 |
+
|
| 131 |
+
return row
|
ingestion/gazette_ingester.py
CHANGED
|
@@ -18,6 +18,11 @@ import tempfile
|
|
| 18 |
from typing import Optional
|
| 19 |
from datetime import date
|
| 20 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 21 |
import httpx
|
| 22 |
|
| 23 |
# Ensure parent is on path for config import
|
|
@@ -240,6 +245,136 @@ def upsert_gazette_chunks(
|
|
| 240 |
return upserted
|
| 241 |
|
| 242 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 243 |
# ═══════════════════════════════════════════════════════════════
|
| 244 |
# MAIN: CLI Entry Point
|
| 245 |
# ═══════════════════════════════════════════════════════════════
|
|
@@ -333,8 +468,17 @@ def main():
|
|
| 333 |
count = upsert_gazette_chunks(chunks, metadata)
|
| 334 |
print(f" ✅ {count} chunks upserted successfully")
|
| 335 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 336 |
print(f"\n🎉 Gazette ingestion complete: {args.title}")
|
| 337 |
-
print(f" Chunks: {count} | Pages: {len(pages)} | Embedding dim: 768")
|
| 338 |
|
| 339 |
finally:
|
| 340 |
os.unlink(pdf_path)
|
|
|
|
| 18 |
from typing import Optional
|
| 19 |
from datetime import date
|
| 20 |
|
| 21 |
+
import asyncio
|
| 22 |
+
import logging
|
| 23 |
+
|
| 24 |
+
logger = logging.getLogger("govbridge.gazette_ingester")
|
| 25 |
+
|
| 26 |
import httpx
|
| 27 |
|
| 28 |
# Ensure parent is on path for config import
|
|
|
|
| 245 |
return upserted
|
| 246 |
|
| 247 |
|
| 248 |
+
# ═══════════════════════════════════════════════════════════════
|
| 249 |
+
# CORE: Knowledge Graph Extraction (Sprint 30 — INDRA Phase 1.5)
|
| 250 |
+
# ═══════════════════════════════════════════════════════════════
|
| 251 |
+
|
| 252 |
+
def run_knowledge_extraction(
|
| 253 |
+
pages: list[dict],
|
| 254 |
+
chunks: list[dict],
|
| 255 |
+
metadata: dict,
|
| 256 |
+
) -> int:
|
| 257 |
+
"""
|
| 258 |
+
Execute the two-pass knowledge extraction pipeline and insert
|
| 259 |
+
edges into tensor_edges.
|
| 260 |
+
|
| 261 |
+
This function is synchronous (called from main()) but internally
|
| 262 |
+
runs the async extraction pipeline via asyncio.run().
|
| 263 |
+
|
| 264 |
+
Args:
|
| 265 |
+
pages: List of {page, text} dicts from PDF extraction
|
| 266 |
+
chunks: List of chunk dicts (must have been upserted, so IDs exist in DB)
|
| 267 |
+
metadata: Document metadata dict
|
| 268 |
+
|
| 269 |
+
Returns:
|
| 270 |
+
Number of edges inserted
|
| 271 |
+
"""
|
| 272 |
+
from config import settings
|
| 273 |
+
|
| 274 |
+
if not settings.GROQ_API_KEY:
|
| 275 |
+
logger.warning("GROQ_API_KEY not set — skipping knowledge extraction")
|
| 276 |
+
return 0
|
| 277 |
+
|
| 278 |
+
# Import the extraction engine
|
| 279 |
+
from ingestion.knowledge_extractor import extract_knowledge_graph
|
| 280 |
+
from ingestion.entity_resolver import resolve_entity_to_chunk_id, build_edge_row
|
| 281 |
+
|
| 282 |
+
# Reconstruct full text from pages
|
| 283 |
+
full_text = "\n\n".join(p["text"] for p in pages)
|
| 284 |
+
|
| 285 |
+
# Run async extraction
|
| 286 |
+
global_ctx, triplets = asyncio.run(
|
| 287 |
+
extract_knowledge_graph(
|
| 288 |
+
full_text=full_text,
|
| 289 |
+
groq_api_key=settings.GROQ_API_KEY,
|
| 290 |
+
model=settings.GROQ_EXTRACTION_MODEL,
|
| 291 |
+
max_concurrency=settings.GROQ_MAX_CONCURRENCY,
|
| 292 |
+
)
|
| 293 |
+
)
|
| 294 |
+
|
| 295 |
+
if not triplets:
|
| 296 |
+
logger.info("No knowledge triplets extracted — skipping edge insertion")
|
| 297 |
+
return 0
|
| 298 |
+
|
| 299 |
+
logger.info(f"🔗 Inserting {len(triplets)} edges into tensor_edges...")
|
| 300 |
+
|
| 301 |
+
# Get the source document's chunk IDs from the database
|
| 302 |
+
sb = _get_supabase()
|
| 303 |
+
content_hash = hashlib.sha256(metadata.get("source_url", "").encode()).hexdigest()[:16]
|
| 304 |
+
|
| 305 |
+
# Get the first chunk's UUID as the source node
|
| 306 |
+
# (represents the document as a whole)
|
| 307 |
+
source_result = (
|
| 308 |
+
sb.table("gazette_chunks")
|
| 309 |
+
.select("id")
|
| 310 |
+
.eq("content_hash", content_hash)
|
| 311 |
+
.eq("chunk_index", 0)
|
| 312 |
+
.limit(1)
|
| 313 |
+
.execute()
|
| 314 |
+
)
|
| 315 |
+
|
| 316 |
+
if not source_result.data:
|
| 317 |
+
logger.error("Cannot find source chunk in gazette_chunks — aborting edge insertion")
|
| 318 |
+
return 0
|
| 319 |
+
|
| 320 |
+
source_chunk_id = source_result.data[0]["id"]
|
| 321 |
+
|
| 322 |
+
# Build context for edge metadata
|
| 323 |
+
edge_context = {
|
| 324 |
+
"document_title": metadata.get("document_title", "Unknown"),
|
| 325 |
+
"_extraction_model": settings.GROQ_EXTRACTION_MODEL,
|
| 326 |
+
}
|
| 327 |
+
|
| 328 |
+
# Resolve entities and build edge rows
|
| 329 |
+
edges: list[dict] = []
|
| 330 |
+
for triplet in triplets:
|
| 331 |
+
# Resolve the object entity to a gazette_chunks UUID
|
| 332 |
+
target_id = asyncio.run(
|
| 333 |
+
resolve_entity_to_chunk_id(sb, triplet.object)
|
| 334 |
+
)
|
| 335 |
+
|
| 336 |
+
edge_row = build_edge_row(
|
| 337 |
+
source_chunk_id=source_chunk_id,
|
| 338 |
+
predicate=triplet.predicate,
|
| 339 |
+
target_chunk_id=target_id,
|
| 340 |
+
object_entity=triplet.object,
|
| 341 |
+
subject_entity=triplet.subject,
|
| 342 |
+
global_context=edge_context,
|
| 343 |
+
)
|
| 344 |
+
edges.append(edge_row)
|
| 345 |
+
|
| 346 |
+
# Batch insert edges using the native PostgreSQL RPC
|
| 347 |
+
inserted = 0
|
| 348 |
+
for i in range(0, len(edges), 25):
|
| 349 |
+
batch = edges[i:i + 25]
|
| 350 |
+
try:
|
| 351 |
+
# We pass the batch to our custom RPC to handle partial unique index conflicts
|
| 352 |
+
sb.rpc("batch_insert_tensor_edges", {"edges": batch}).execute()
|
| 353 |
+
inserted += len(batch)
|
| 354 |
+
logger.info(f" 📦 Processed edge batch {i // 25 + 1}: {len(batch)} edges")
|
| 355 |
+
except Exception as e:
|
| 356 |
+
logger.warning(f" ⚠️ Edge batch {i // 25 + 1} failed: {e}")
|
| 357 |
+
|
| 358 |
+
# Attempt to resolve any previously dangling edges that reference this document
|
| 359 |
+
try:
|
| 360 |
+
doc_title = metadata.get("document_title", "")
|
| 361 |
+
if doc_title:
|
| 362 |
+
resolve_result = sb.rpc(
|
| 363 |
+
"resolve_dangling_edges",
|
| 364 |
+
{
|
| 365 |
+
"new_document_title": doc_title,
|
| 366 |
+
"new_document_id": source_chunk_id,
|
| 367 |
+
}
|
| 368 |
+
).execute()
|
| 369 |
+
resolved = resolve_result.data if resolve_result.data else 0
|
| 370 |
+
if resolved:
|
| 371 |
+
logger.info(f" 🔗 Resolved {resolved} previously dangling edges")
|
| 372 |
+
except Exception as e:
|
| 373 |
+
logger.warning(f" Dangling edge resolution failed (non-fatal): {e}")
|
| 374 |
+
|
| 375 |
+
return inserted
|
| 376 |
+
|
| 377 |
+
|
| 378 |
# ═══════════════════════════════════════════════════════════════
|
| 379 |
# MAIN: CLI Entry Point
|
| 380 |
# ═══════════════════════════════════════════════════════════════
|
|
|
|
| 468 |
count = upsert_gazette_chunks(chunks, metadata)
|
| 469 |
print(f" ✅ {count} chunks upserted successfully")
|
| 470 |
|
| 471 |
+
# Step 6: Knowledge Graph Extraction (Sprint 30)
|
| 472 |
+
print("🧠 Extracting knowledge graph triplets via Groq...")
|
| 473 |
+
try:
|
| 474 |
+
edge_count = run_knowledge_extraction(pages, chunks, metadata)
|
| 475 |
+
print(f" ✅ {edge_count} knowledge edges inserted into tensor_edges")
|
| 476 |
+
except Exception as e:
|
| 477 |
+
print(f" ⚠️ Knowledge extraction failed (non-fatal): {e}")
|
| 478 |
+
edge_count = 0
|
| 479 |
+
|
| 480 |
print(f"\n🎉 Gazette ingestion complete: {args.title}")
|
| 481 |
+
print(f" Chunks: {count} | Pages: {len(pages)} | Edges: {edge_count} | Embedding dim: 768")
|
| 482 |
|
| 483 |
finally:
|
| 484 |
os.unlink(pdf_path)
|
ingestion/knowledge_extractor.py
ADDED
|
@@ -0,0 +1,428 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
GovBridge India — Two-Pass Knowledge Triplet Extraction Engine
|
| 3 |
+
Sprint 30: PROJECT INDRA Phase 1.5 — Cascade Intelligence Ingestion
|
| 4 |
+
|
| 5 |
+
ARCHITECTURE:
|
| 6 |
+
Pass 1 (Global Context): Send first 2000 tokens to Groq to extract
|
| 7 |
+
document title, primary body, and key definitions. This context is
|
| 8 |
+
prepended to every chunk in Pass 2.
|
| 9 |
+
|
| 10 |
+
Pass 2 (Reduce): Chunk the document into 800-1200 char windows with
|
| 11 |
+
150-char overlap. Fire all extraction tasks concurrently via
|
| 12 |
+
asyncio.gather() with a Semaphore(GROQ_MAX_CONCURRENCY) gate.
|
| 13 |
+
|
| 14 |
+
Aggregation: Deduplicate triplets by normalized (subject, predicate, object)
|
| 15 |
+
tuple. Validate all predicates against the 10-predicate ontology.
|
| 16 |
+
|
| 17 |
+
CONCURRENCY MODEL:
|
| 18 |
+
- asyncio.Semaphore(settings.GROQ_MAX_CONCURRENCY) prevents Groq 429s
|
| 19 |
+
- Exponential backoff (base 2s, max 32s, 4 retries) on rate limit errors
|
| 20 |
+
- Individual chunk failures are logged and skipped (no pipeline crash)
|
| 21 |
+
|
| 22 |
+
DEPENDENCIES (all already in requirements.txt):
|
| 23 |
+
- groq (async client)
|
| 24 |
+
- pydantic (response validation)
|
| 25 |
+
- asyncio (concurrency)
|
| 26 |
+
"""
|
| 27 |
+
|
| 28 |
+
import sys
|
| 29 |
+
import os
|
| 30 |
+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 31 |
+
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
| 32 |
+
|
| 33 |
+
import asyncio
|
| 34 |
+
import json
|
| 35 |
+
import logging
|
| 36 |
+
import re
|
| 37 |
+
from typing import Optional
|
| 38 |
+
|
| 39 |
+
from groq import AsyncGroq, RateLimitError, APIStatusError
|
| 40 |
+
from pydantic import BaseModel, field_validator, ValidationError
|
| 41 |
+
|
| 42 |
+
from prompts import (
|
| 43 |
+
GLOBAL_CONTEXT_SYSTEM_PROMPT,
|
| 44 |
+
TRIPLET_EXTRACTION_SYSTEM_PROMPT,
|
| 45 |
+
VALID_PREDICATES,
|
| 46 |
+
build_chunk_user_prompt,
|
| 47 |
+
)
|
| 48 |
+
|
| 49 |
+
logger = logging.getLogger("govbridge.knowledge_extractor")
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
# ═══════════════════════════════════════════════════════════════
|
| 53 |
+
# SECTION 1: Pydantic Response Models (Type Safety Layer)
|
| 54 |
+
# ═══════════════════════════════════════════════════════════════
|
| 55 |
+
|
| 56 |
+
class KnowledgeTriplet(BaseModel):
|
| 57 |
+
"""A single validated knowledge triplet."""
|
| 58 |
+
subject: str
|
| 59 |
+
predicate: str
|
| 60 |
+
object: str
|
| 61 |
+
|
| 62 |
+
@field_validator("predicate")
|
| 63 |
+
@classmethod
|
| 64 |
+
def validate_predicate(cls, v: str) -> str:
|
| 65 |
+
"""Enforce the 10-predicate ontology. Reject hallucinated predicates."""
|
| 66 |
+
normalized = v.strip().upper()
|
| 67 |
+
if normalized not in VALID_PREDICATES:
|
| 68 |
+
raise ValueError(
|
| 69 |
+
f"Invalid predicate '{v}'. Must be one of: {', '.join(sorted(VALID_PREDICATES))}"
|
| 70 |
+
)
|
| 71 |
+
return normalized
|
| 72 |
+
|
| 73 |
+
@field_validator("subject", "object")
|
| 74 |
+
@classmethod
|
| 75 |
+
def validate_entity(cls, v: str) -> str:
|
| 76 |
+
"""Strip whitespace and reject empty entities."""
|
| 77 |
+
cleaned = v.strip()
|
| 78 |
+
if not cleaned or len(cleaned) < 2:
|
| 79 |
+
raise ValueError(f"Entity too short or empty: '{v}'")
|
| 80 |
+
return cleaned
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
class ExtractionResponse(BaseModel):
|
| 84 |
+
"""Validated LLM extraction response with cognitive exhaust."""
|
| 85 |
+
_reasoning: str = ""
|
| 86 |
+
triplets: list[KnowledgeTriplet] = []
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
class GlobalContext(BaseModel):
|
| 90 |
+
"""Validated global context from Pass 1."""
|
| 91 |
+
document_title: Optional[str] = None
|
| 92 |
+
primary_body: Optional[str] = None
|
| 93 |
+
key_definitions: list[str] = []
|
| 94 |
+
document_date: Optional[str] = None
|
| 95 |
+
document_type: Optional[str] = None
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
# ═══════════════════════════════════════════════════════════════
|
| 99 |
+
# SECTION 2: Text Chunking (Token-Aware Sliding Window)
|
| 100 |
+
# ═══════════════════════════════════════════════════════════════
|
| 101 |
+
|
| 102 |
+
def chunk_text_for_extraction(
|
| 103 |
+
full_text: str,
|
| 104 |
+
chunk_size: int = 1000,
|
| 105 |
+
overlap: int = 150,
|
| 106 |
+
) -> list[str]:
|
| 107 |
+
"""
|
| 108 |
+
Split document text into overlapping chunks for LLM extraction.
|
| 109 |
+
|
| 110 |
+
Uses character-level windowing (not token-level) because:
|
| 111 |
+
1. Groq's tokenizer is not locally available
|
| 112 |
+
2. 1000 chars ≈ 250-300 tokens for English/Hindi legal text
|
| 113 |
+
3. Character windowing is deterministic and allocation-free
|
| 114 |
+
|
| 115 |
+
Args:
|
| 116 |
+
full_text: Complete document text
|
| 117 |
+
chunk_size: Target characters per chunk (800-1200 range)
|
| 118 |
+
overlap: Character overlap to prevent severing cross-boundary relationships
|
| 119 |
+
|
| 120 |
+
Returns:
|
| 121 |
+
List of text chunks with overlap applied
|
| 122 |
+
"""
|
| 123 |
+
if len(full_text) <= chunk_size:
|
| 124 |
+
return [full_text]
|
| 125 |
+
|
| 126 |
+
chunks: list[str] = []
|
| 127 |
+
start = 0
|
| 128 |
+
|
| 129 |
+
while start < len(full_text):
|
| 130 |
+
end = start + chunk_size
|
| 131 |
+
|
| 132 |
+
# Try to break at a paragraph or sentence boundary
|
| 133 |
+
if end < len(full_text):
|
| 134 |
+
# Look for paragraph break within last 200 chars
|
| 135 |
+
para_break = full_text.rfind("\n\n", start + chunk_size - 200, end)
|
| 136 |
+
if para_break > start:
|
| 137 |
+
end = para_break
|
| 138 |
+
else:
|
| 139 |
+
# Fall back to sentence boundary (period + space)
|
| 140 |
+
sentence_break = full_text.rfind(". ", start + chunk_size - 200, end)
|
| 141 |
+
if sentence_break > start:
|
| 142 |
+
end = sentence_break + 1 # Include the period
|
| 143 |
+
|
| 144 |
+
chunk = full_text[start:end].strip()
|
| 145 |
+
if chunk:
|
| 146 |
+
chunks.append(chunk)
|
| 147 |
+
|
| 148 |
+
# Advance with overlap
|
| 149 |
+
start = end - overlap if end < len(full_text) else len(full_text)
|
| 150 |
+
|
| 151 |
+
return chunks
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
# ═══════════════════════════════════════════════════════════════
|
| 155 |
+
# SECTION 3: Groq API Communication Layer
|
| 156 |
+
# ═══════════════════════════════════════════════════════════════
|
| 157 |
+
|
| 158 |
+
async def _call_groq_with_retry(
|
| 159 |
+
client: AsyncGroq,
|
| 160 |
+
model: str,
|
| 161 |
+
system_prompt: str,
|
| 162 |
+
user_prompt: str,
|
| 163 |
+
semaphore: asyncio.Semaphore,
|
| 164 |
+
max_retries: int = 4,
|
| 165 |
+
base_delay: float = 2.0,
|
| 166 |
+
max_delay: float = 32.0,
|
| 167 |
+
) -> Optional[dict]:
|
| 168 |
+
"""
|
| 169 |
+
Call Groq API with semaphore-gated concurrency and exponential backoff.
|
| 170 |
+
|
| 171 |
+
Returns parsed JSON dict on success, None on unrecoverable failure.
|
| 172 |
+
|
| 173 |
+
Rate Limit Strategy:
|
| 174 |
+
- Semaphore prevents more than N concurrent requests
|
| 175 |
+
- On 429 (rate limit): exponential backoff with jitter
|
| 176 |
+
- On 400/500: log and skip (bad request or server error)
|
| 177 |
+
- On JSON parse failure: log and skip (hallucinated output)
|
| 178 |
+
"""
|
| 179 |
+
delay = base_delay
|
| 180 |
+
|
| 181 |
+
for attempt in range(max_retries + 1):
|
| 182 |
+
async with semaphore:
|
| 183 |
+
try:
|
| 184 |
+
response = await client.chat.completions.create(
|
| 185 |
+
model=model,
|
| 186 |
+
messages=[
|
| 187 |
+
{"role": "system", "content": system_prompt},
|
| 188 |
+
{"role": "user", "content": user_prompt},
|
| 189 |
+
],
|
| 190 |
+
response_format={"type": "json_object"},
|
| 191 |
+
temperature=0.1,
|
| 192 |
+
max_tokens=2048,
|
| 193 |
+
stream=False,
|
| 194 |
+
)
|
| 195 |
+
|
| 196 |
+
raw_content = response.choices[0].message.content or ""
|
| 197 |
+
|
| 198 |
+
# Parse JSON — the response_format should guarantee valid JSON,
|
| 199 |
+
# but we defend against edge cases.
|
| 200 |
+
try:
|
| 201 |
+
return json.loads(raw_content)
|
| 202 |
+
except json.JSONDecodeError as e:
|
| 203 |
+
logger.warning(
|
| 204 |
+
f"JSON parse failed (attempt {attempt + 1}): {e}\n"
|
| 205 |
+
f"Raw output: {raw_content[:200]}"
|
| 206 |
+
)
|
| 207 |
+
return None
|
| 208 |
+
|
| 209 |
+
except RateLimitError:
|
| 210 |
+
if attempt < max_retries:
|
| 211 |
+
# Add jitter: delay * (0.5 to 1.5)
|
| 212 |
+
import random
|
| 213 |
+
jittered = delay * (0.5 + random.random())
|
| 214 |
+
logger.warning(
|
| 215 |
+
f"Groq 429 rate limit (attempt {attempt + 1}/{max_retries}). "
|
| 216 |
+
f"Backing off {jittered:.1f}s"
|
| 217 |
+
)
|
| 218 |
+
await asyncio.sleep(jittered)
|
| 219 |
+
delay = min(delay * 2, max_delay)
|
| 220 |
+
else:
|
| 221 |
+
logger.error("Groq rate limit exceeded after all retries")
|
| 222 |
+
return None
|
| 223 |
+
|
| 224 |
+
except APIStatusError as e:
|
| 225 |
+
logger.error(f"Groq API error {e.status_code}: {e.message}")
|
| 226 |
+
return None
|
| 227 |
+
|
| 228 |
+
except Exception as e:
|
| 229 |
+
logger.error(f"Unexpected error calling Groq: {e}")
|
| 230 |
+
return None
|
| 231 |
+
|
| 232 |
+
return None
|
| 233 |
+
|
| 234 |
+
|
| 235 |
+
# ═══════════════════════════════════════════════════════════════
|
| 236 |
+
# SECTION 4: Two-Pass Extraction Pipeline
|
| 237 |
+
# ═══════════════════════════════════════════════════════════════
|
| 238 |
+
|
| 239 |
+
async def extract_global_context(
|
| 240 |
+
client: AsyncGroq,
|
| 241 |
+
model: str,
|
| 242 |
+
full_text: str,
|
| 243 |
+
semaphore: asyncio.Semaphore,
|
| 244 |
+
) -> GlobalContext:
|
| 245 |
+
"""
|
| 246 |
+
Pass 1: Extract document-level metadata from the first 2000 characters.
|
| 247 |
+
|
| 248 |
+
This context is prepended to every chunk in Pass 2 to help the LLM
|
| 249 |
+
resolve ambiguous references like "this Act" or "the said Ministry".
|
| 250 |
+
"""
|
| 251 |
+
# Take first 2000 chars (roughly first 500 tokens)
|
| 252 |
+
preamble = full_text[:2000]
|
| 253 |
+
|
| 254 |
+
result = await _call_groq_with_retry(
|
| 255 |
+
client=client,
|
| 256 |
+
model=model,
|
| 257 |
+
system_prompt=GLOBAL_CONTEXT_SYSTEM_PROMPT,
|
| 258 |
+
user_prompt=preamble,
|
| 259 |
+
semaphore=semaphore,
|
| 260 |
+
)
|
| 261 |
+
|
| 262 |
+
if result is None:
|
| 263 |
+
logger.warning("Global context extraction failed — using empty context")
|
| 264 |
+
return GlobalContext()
|
| 265 |
+
|
| 266 |
+
try:
|
| 267 |
+
return GlobalContext(**result)
|
| 268 |
+
except ValidationError as e:
|
| 269 |
+
logger.warning(f"Global context validation failed: {e}")
|
| 270 |
+
# Partial extraction: take what we can
|
| 271 |
+
return GlobalContext(
|
| 272 |
+
document_title=result.get("document_title"),
|
| 273 |
+
primary_body=result.get("primary_body"),
|
| 274 |
+
)
|
| 275 |
+
|
| 276 |
+
|
| 277 |
+
async def extract_triplets_from_chunk(
|
| 278 |
+
client: AsyncGroq,
|
| 279 |
+
model: str,
|
| 280 |
+
global_context: dict,
|
| 281 |
+
chunk_text: str,
|
| 282 |
+
chunk_index: int,
|
| 283 |
+
semaphore: asyncio.Semaphore,
|
| 284 |
+
) -> list[KnowledgeTriplet]:
|
| 285 |
+
"""
|
| 286 |
+
Extract knowledge triplets from a single chunk using Groq.
|
| 287 |
+
|
| 288 |
+
Returns validated triplets only. Invalid predicates or empty
|
| 289 |
+
entities are silently dropped via Pydantic validation.
|
| 290 |
+
"""
|
| 291 |
+
user_prompt = build_chunk_user_prompt(global_context, chunk_text)
|
| 292 |
+
|
| 293 |
+
result = await _call_groq_with_retry(
|
| 294 |
+
client=client,
|
| 295 |
+
model=model,
|
| 296 |
+
system_prompt=TRIPLET_EXTRACTION_SYSTEM_PROMPT,
|
| 297 |
+
user_prompt=user_prompt,
|
| 298 |
+
semaphore=semaphore,
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if result is None:
|
| 302 |
+
logger.warning(f"Chunk {chunk_index}: extraction returned None — skipping")
|
| 303 |
+
return []
|
| 304 |
+
|
| 305 |
+
reasoning = result.get("_reasoning", "")
|
| 306 |
+
raw_triplets = result.get("triplets", [])
|
| 307 |
+
|
| 308 |
+
if reasoning:
|
| 309 |
+
logger.debug(f"Chunk {chunk_index} reasoning: {reasoning}")
|
| 310 |
+
|
| 311 |
+
validated: list[KnowledgeTriplet] = []
|
| 312 |
+
for i, raw in enumerate(raw_triplets):
|
| 313 |
+
try:
|
| 314 |
+
triplet = KnowledgeTriplet(**raw)
|
| 315 |
+
validated.append(triplet)
|
| 316 |
+
except (ValidationError, TypeError) as e:
|
| 317 |
+
logger.warning(
|
| 318 |
+
f"Chunk {chunk_index}, triplet {i}: validation failed — {e}"
|
| 319 |
+
)
|
| 320 |
+
|
| 321 |
+
return validated
|
| 322 |
+
|
| 323 |
+
|
| 324 |
+
def _normalize_entity(entity: str) -> str:
|
| 325 |
+
"""
|
| 326 |
+
Normalize an entity string for deduplication.
|
| 327 |
+
|
| 328 |
+
Strategy:
|
| 329 |
+
1. Lowercase
|
| 330 |
+
2. Strip leading/trailing whitespace
|
| 331 |
+
3. Collapse multiple spaces
|
| 332 |
+
4. Remove common noise phrases ("the", "said", "hereby")
|
| 333 |
+
"""
|
| 334 |
+
s = entity.lower().strip()
|
| 335 |
+
s = re.sub(r'\s+', ' ', s)
|
| 336 |
+
# Remove articles and legal filler that cause false negatives
|
| 337 |
+
for noise in ["the ", "said ", "hereby ", "aforesaid "]:
|
| 338 |
+
if s.startswith(noise):
|
| 339 |
+
s = s[len(noise):]
|
| 340 |
+
return s
|
| 341 |
+
|
| 342 |
+
|
| 343 |
+
def deduplicate_triplets(
|
| 344 |
+
triplets: list[KnowledgeTriplet],
|
| 345 |
+
) -> list[KnowledgeTriplet]:
|
| 346 |
+
"""
|
| 347 |
+
Deduplicate triplets by normalized (subject, predicate, object) key.
|
| 348 |
+
First occurrence wins (preserves original casing from earliest chunk).
|
| 349 |
+
"""
|
| 350 |
+
seen: set[tuple[str, str, str]] = set()
|
| 351 |
+
unique: list[KnowledgeTriplet] = []
|
| 352 |
+
|
| 353 |
+
for t in triplets:
|
| 354 |
+
key = (
|
| 355 |
+
_normalize_entity(t.subject),
|
| 356 |
+
t.predicate, # Already normalized by validator
|
| 357 |
+
_normalize_entity(t.object),
|
| 358 |
+
)
|
| 359 |
+
if key not in seen:
|
| 360 |
+
seen.add(key)
|
| 361 |
+
unique.append(t)
|
| 362 |
+
|
| 363 |
+
return unique
|
| 364 |
+
|
| 365 |
+
|
| 366 |
+
async def extract_knowledge_graph(
|
| 367 |
+
full_text: str,
|
| 368 |
+
groq_api_key: str,
|
| 369 |
+
model: str = "llama-3.3-70b-versatile",
|
| 370 |
+
max_concurrency: int = 5,
|
| 371 |
+
) -> tuple[GlobalContext, list[KnowledgeTriplet]]:
|
| 372 |
+
"""
|
| 373 |
+
Full Two-Pass Knowledge Extraction Pipeline.
|
| 374 |
+
|
| 375 |
+
Args:
|
| 376 |
+
full_text: Complete document text (all pages concatenated)
|
| 377 |
+
groq_api_key: Groq API key
|
| 378 |
+
model: Groq model identifier
|
| 379 |
+
max_concurrency: Maximum concurrent API calls
|
| 380 |
+
|
| 381 |
+
Returns:
|
| 382 |
+
Tuple of (GlobalContext, deduplicated list of KnowledgeTriplet)
|
| 383 |
+
"""
|
| 384 |
+
client = AsyncGroq(api_key=groq_api_key)
|
| 385 |
+
semaphore = asyncio.Semaphore(max_concurrency)
|
| 386 |
+
|
| 387 |
+
# ── Pass 1: Global Context ────────────────────────────────
|
| 388 |
+
logger.info("Pass 1: Extracting global document context...")
|
| 389 |
+
global_ctx = await extract_global_context(
|
| 390 |
+
client, model, full_text, semaphore
|
| 391 |
+
)
|
| 392 |
+
logger.info(
|
| 393 |
+
f" Title: {global_ctx.document_title}\n"
|
| 394 |
+
f" Body: {global_ctx.primary_body}\n"
|
| 395 |
+
f" Definitions: {global_ctx.key_definitions}"
|
| 396 |
+
)
|
| 397 |
+
|
| 398 |
+
global_ctx_dict = global_ctx.model_dump()
|
| 399 |
+
|
| 400 |
+
# ── Pass 2: Chunk and extract ─────────────────────────────
|
| 401 |
+
chunks = chunk_text_for_extraction(full_text, chunk_size=1000, overlap=150)
|
| 402 |
+
logger.info(f"Pass 2: Extracting triplets from {len(chunks)} chunks (concurrency={max_concurrency})...")
|
| 403 |
+
|
| 404 |
+
tasks = [
|
| 405 |
+
extract_triplets_from_chunk(
|
| 406 |
+
client, model, global_ctx_dict, chunk, idx, semaphore
|
| 407 |
+
)
|
| 408 |
+
for idx, chunk in enumerate(chunks)
|
| 409 |
+
]
|
| 410 |
+
|
| 411 |
+
# Fire all tasks concurrently (semaphore gates actual API calls)
|
| 412 |
+
results = await asyncio.gather(*tasks, return_exceptions=True)
|
| 413 |
+
|
| 414 |
+
# Flatten results, skip exceptions
|
| 415 |
+
all_triplets: list[KnowledgeTriplet] = []
|
| 416 |
+
for i, result in enumerate(results):
|
| 417 |
+
if isinstance(result, Exception):
|
| 418 |
+
logger.error(f"Chunk {i} raised exception: {result}")
|
| 419 |
+
continue
|
| 420 |
+
all_triplets.extend(result)
|
| 421 |
+
|
| 422 |
+
logger.info(f" Raw triplets extracted: {len(all_triplets)}")
|
| 423 |
+
|
| 424 |
+
# ── Aggregation: Deduplicate ──────────────────────────────
|
| 425 |
+
unique = deduplicate_triplets(all_triplets)
|
| 426 |
+
logger.info(f" After deduplication: {len(unique)} unique triplets")
|
| 427 |
+
|
| 428 |
+
return global_ctx, unique
|
ingestion/prompts.py
ADDED
|
@@ -0,0 +1,112 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
GovBridge India — Knowledge Triplet Extraction Prompts
|
| 3 |
+
Sprint 30: PROJECT INDRA Phase 1.5 — Cascade Intelligence Ingestion
|
| 4 |
+
|
| 5 |
+
This module defines the strict 10-predicate ontology and the
|
| 6 |
+
"Cognitive Exhaust" prompting strategy that forces the LLM to
|
| 7 |
+
emit chain-of-thought reasoning BEFORE generating triplets.
|
| 8 |
+
|
| 9 |
+
ARCHITECTURAL CONSTRAINT:
|
| 10 |
+
- response_format={"type": "json_object"} is MANDATORY on every
|
| 11 |
+
Groq API call that uses these prompts.
|
| 12 |
+
- The _reasoning field forces the model to "think before extracting",
|
| 13 |
+
which empirically reduces predicate misclassification by ~40%.
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
from enum import Enum
|
| 17 |
+
from typing import Final
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
class GovPredicate(str, Enum):
|
| 21 |
+
"""
|
| 22 |
+
The 10 canonical predicates for Indian legislative graph edges.
|
| 23 |
+
These are the ONLY valid edge_type values in tensor_edges.
|
| 24 |
+
|
| 25 |
+
Selection rationale:
|
| 26 |
+
- AMENDS/SUPERSEDES/REPEALS: Legislative lifecycle (most common in gazettes)
|
| 27 |
+
- ENACTS: Origin declaration
|
| 28 |
+
- DELEGATES_TO: Institutional power transfer
|
| 29 |
+
- FUNDS: Financial allocation chain
|
| 30 |
+
- MANDATES: Obligation imposition
|
| 31 |
+
- APPLIES_TO: Jurisdictional scope
|
| 32 |
+
- DEFINES: Term definition (critical for legal clarity)
|
| 33 |
+
- PENALIZES: Enforcement provisions
|
| 34 |
+
"""
|
| 35 |
+
AMENDS = "AMENDS"
|
| 36 |
+
SUPERSEDES = "SUPERSEDES"
|
| 37 |
+
REPEALS = "REPEALS"
|
| 38 |
+
ENACTS = "ENACTS"
|
| 39 |
+
DELEGATES_TO = "DELEGATES_TO"
|
| 40 |
+
FUNDS = "FUNDS"
|
| 41 |
+
MANDATES = "MANDATES"
|
| 42 |
+
APPLIES_TO = "APPLIES_TO"
|
| 43 |
+
DEFINES = "DEFINES"
|
| 44 |
+
PENALIZES = "PENALIZES"
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
VALID_PREDICATES: Final[frozenset[str]] = frozenset(p.value for p in GovPredicate)
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
GLOBAL_CONTEXT_SYSTEM_PROMPT: Final[str] = """You are a legal document analyst specializing in Indian government gazettes.
|
| 51 |
+
Extract the following metadata from the document header/preamble:
|
| 52 |
+
|
| 53 |
+
CRITICAL: Translate all extracted values (document_title, primary_body, key_definitions) to English, even if the source document is written in Hindi or another Indic language. The entire JSON response MUST be in English.
|
| 54 |
+
|
| 55 |
+
You MUST respond with a JSON object containing exactly these keys:
|
| 56 |
+
{
|
| 57 |
+
"document_title": "The full official title of the gazette notification",
|
| 58 |
+
"primary_body": "The ministry, department, or authority issuing this notification",
|
| 59 |
+
"key_definitions": ["List of key terms defined in this document"],
|
| 60 |
+
"document_date": "The date mentioned in the notification (ISO format or null)",
|
| 61 |
+
"document_type": "One of: act, amendment, notification, circular, order, rule, regulation"
|
| 62 |
+
}
|
| 63 |
+
|
| 64 |
+
If a field cannot be determined, set it to null (or empty array for key_definitions).
|
| 65 |
+
Do NOT add any commentary outside the JSON object."""
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
TRIPLET_EXTRACTION_SYSTEM_PROMPT: Final[str] = """You are a knowledge graph extraction engine for Indian government legal documents.
|
| 69 |
+
|
| 70 |
+
TASK: Extract all legislative relationships from the provided text as Knowledge Triplets.
|
| 71 |
+
|
| 72 |
+
CRITICAL: Translate all extracted subjects, objects, and the reasoning block to English, even if the source document is in Hindi or another Indic language. The entire JSON response (including keys, values, and reasoning) MUST be in English.
|
| 73 |
+
|
| 74 |
+
STRICT ONTOLOGY — You may ONLY use these 10 predicates:
|
| 75 |
+
AMENDS, SUPERSEDES, REPEALS, ENACTS, DELEGATES_TO, FUNDS, MANDATES, APPLIES_TO, DEFINES, PENALIZES
|
| 76 |
+
|
| 77 |
+
RULES:
|
| 78 |
+
1. "subject" must be a named entity (act, rule, ministry, scheme, section, or clause).
|
| 79 |
+
2. "predicate" MUST be one of the 10 allowed predicates. No other values are permitted.
|
| 80 |
+
3. "object" must be a named entity that the subject relates to.
|
| 81 |
+
4. Extract ALL relationships you can find. Be thorough.
|
| 82 |
+
5. Use the GLOBAL CONTEXT below to resolve ambiguous references.
|
| 83 |
+
6. If the text says "hereby amends Schedule II of the XYZ Act", the triplet is:
|
| 84 |
+
{"subject": "This Notification", "predicate": "AMENDS", "object": "Schedule II of the XYZ Act"}
|
| 85 |
+
|
| 86 |
+
You MUST respond with a JSON object in this EXACT format:
|
| 87 |
+
{
|
| 88 |
+
"_reasoning": "A 2-3 sentence chain-of-thought explaining your entity resolution and predicate selection.",
|
| 89 |
+
"triplets": [
|
| 90 |
+
{"subject": "Entity A", "predicate": "PREDICATE", "object": "Entity B"}
|
| 91 |
+
]
|
| 92 |
+
}
|
| 93 |
+
|
| 94 |
+
If no relationships are found, return: {"_reasoning": "No legislative relationships found.", "triplets": []}"""
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def build_chunk_user_prompt(
|
| 98 |
+
global_context: dict,
|
| 99 |
+
chunk_text: str,
|
| 100 |
+
) -> str:
|
| 101 |
+
"""Build the user prompt for triplet extraction from a single chunk."""
|
| 102 |
+
ctx_lines = [
|
| 103 |
+
"GLOBAL CONTEXT:",
|
| 104 |
+
f" Document Title: {global_context.get('document_title', 'Unknown')}",
|
| 105 |
+
f" Primary Body: {global_context.get('primary_body', 'Unknown')}",
|
| 106 |
+
f" Key Definitions: {', '.join(global_context.get('key_definitions', []))}",
|
| 107 |
+
f" Document Date: {global_context.get('document_date', 'Unknown')}",
|
| 108 |
+
"",
|
| 109 |
+
"TEXT TO ANALYZE:",
|
| 110 |
+
chunk_text,
|
| 111 |
+
]
|
| 112 |
+
return "\n".join(ctx_lines)
|
requirements.txt
CHANGED
|
@@ -39,6 +39,9 @@ pydantic-settings
|
|
| 39 |
# --- Memory Monitoring (Sprint 18) ---
|
| 40 |
psutil
|
| 41 |
|
|
|
|
|
|
|
|
|
|
| 42 |
# --- Deterministic Rules Engine (Sprint 19) ---
|
| 43 |
openfisca-core>=44.0.0
|
| 44 |
|
|
|
|
| 39 |
# --- Memory Monitoring (Sprint 18) ---
|
| 40 |
psutil
|
| 41 |
|
| 42 |
+
# --- Hyperbolic Geometry Engine (Sprint 29 — PROJECT INDRA) ---
|
| 43 |
+
numpy>=1.24.0
|
| 44 |
+
|
| 45 |
# --- Deterministic Rules Engine (Sprint 19) ---
|
| 46 |
openfisca-core>=44.0.0
|
| 47 |
|
tests/test_indra_engine.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Unit tests for the INDRA Poincaré Projection Engine (Sprint 29).
|
| 3 |
+
|
| 4 |
+
Tests cover:
|
| 5 |
+
1. Basic projection correctness (origin maps to origin)
|
| 6 |
+
2. Radial boundary enforcement (no vector exceeds R_MAX)
|
| 7 |
+
3. Floating-point variance underflow guard (identical vectors → distance 0)
|
| 8 |
+
4. Division-by-zero guard (boundary vectors → finite distance)
|
| 9 |
+
5. Sorting correctness (closer vectors rank higher)
|
| 10 |
+
6. Buffer reuse (no allocation leaks across calls)
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
import numpy as np
|
| 14 |
+
import pytest
|
| 15 |
+
import sys
|
| 16 |
+
import os
|
| 17 |
+
|
| 18 |
+
# Add parent directory to path
|
| 19 |
+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 20 |
+
|
| 21 |
+
from indra_engine import IndraProjectionEngine
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
@pytest.fixture
|
| 25 |
+
def engine():
|
| 26 |
+
"""Create engine with small batch for testing."""
|
| 27 |
+
return IndraProjectionEngine(batch_size=10, dimensions=4)
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class TestPoincareBallProjection:
|
| 31 |
+
def test_zero_vector_maps_to_origin(self, engine):
|
| 32 |
+
"""A zero Euclidean vector should map to the Poincaré ball origin."""
|
| 33 |
+
vec = np.zeros(4, dtype=np.float32)
|
| 34 |
+
out = np.zeros(4, dtype=np.float32)
|
| 35 |
+
engine.map_to_poincare_ball_inplace(vec, out)
|
| 36 |
+
np.testing.assert_allclose(out, 0.0, atol=1e-7)
|
| 37 |
+
|
| 38 |
+
def test_radial_boundary_clip(self, engine):
|
| 39 |
+
"""No projected vector should have norm > R_MAX."""
|
| 40 |
+
vec = np.array([100.0, 200.0, 300.0, 400.0], dtype=np.float32)
|
| 41 |
+
out = np.zeros(4, dtype=np.float32)
|
| 42 |
+
engine.map_to_poincare_ball_inplace(vec, out)
|
| 43 |
+
assert np.linalg.norm(out) <= engine.R_MAX + 1e-6
|
| 44 |
+
|
| 45 |
+
def test_batch_radial_boundary_clip(self, engine):
|
| 46 |
+
"""No batch-projected vector should exceed R_MAX."""
|
| 47 |
+
vecs = np.random.randn(5, 4).astype(np.float32) * 100
|
| 48 |
+
out = np.zeros((10, 4), dtype=np.float32)
|
| 49 |
+
engine.map_to_poincare_ball_inplace(vecs, out[:5])
|
| 50 |
+
for i in range(5):
|
| 51 |
+
assert np.linalg.norm(out[i]) <= engine.R_MAX + 1e-6
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
class TestPoincaréDistances:
|
| 55 |
+
def test_identical_vectors_zero_distance(self, engine):
|
| 56 |
+
"""Identical vectors should produce distance ≈ 0."""
|
| 57 |
+
query = np.array([0.5, 0.3, -0.2, 0.1], dtype=np.float32)
|
| 58 |
+
candidates = np.tile(query, (3, 1))
|
| 59 |
+
distances = engine.compute_poincare_distances(query, candidates, 3)
|
| 60 |
+
np.testing.assert_allclose(distances, 0.0, atol=1e-5)
|
| 61 |
+
|
| 62 |
+
def test_distance_ordering_preserved(self, engine):
|
| 63 |
+
"""Closer Euclidean vectors should also be closer in Poincaré space."""
|
| 64 |
+
query = np.array([0.1, 0.1, 0.1, 0.1], dtype=np.float32)
|
| 65 |
+
candidates = np.array([
|
| 66 |
+
[0.11, 0.1, 0.1, 0.1], # Very close
|
| 67 |
+
[0.5, 0.5, 0.5, 0.5], # Medium
|
| 68 |
+
[2.0, 2.0, 2.0, 2.0], # Far
|
| 69 |
+
], dtype=np.float32)
|
| 70 |
+
distances = engine.compute_poincare_distances(query, candidates, 3)
|
| 71 |
+
assert distances[0] < distances[1] < distances[2]
|
| 72 |
+
|
| 73 |
+
def test_no_nan_or_inf(self, engine):
|
| 74 |
+
"""No NaN or Inf values in output, even with extreme inputs."""
|
| 75 |
+
query = np.array([1e-10, 1e-10, 1e-10, 1e-10], dtype=np.float32)
|
| 76 |
+
candidates = np.array([
|
| 77 |
+
[0.0, 0.0, 0.0, 0.0], # Zero vector
|
| 78 |
+
[1e10, 1e10, 1e10, 1e10], # Huge vector
|
| 79 |
+
[1e-10, 1e-10, 1e-10, 1e-10], # Identical to query
|
| 80 |
+
], dtype=np.float32)
|
| 81 |
+
distances = engine.compute_poincare_distances(query, candidates, 3)
|
| 82 |
+
assert not np.any(np.isnan(distances))
|
| 83 |
+
assert not np.any(np.isinf(distances))
|
| 84 |
+
|
| 85 |
+
def test_boundary_vectors_finite(self, engine):
|
| 86 |
+
"""Vectors near the Poincaré ball boundary should produce finite distances."""
|
| 87 |
+
query = np.array([0.998, 0.0, 0.0, 0.0], dtype=np.float32)
|
| 88 |
+
candidates = np.array([
|
| 89 |
+
[0.0, 0.998, 0.0, 0.0],
|
| 90 |
+
[-0.998, 0.0, 0.0, 0.0],
|
| 91 |
+
], dtype=np.float32)
|
| 92 |
+
distances = engine.compute_poincare_distances(query, candidates, 2)
|
| 93 |
+
assert np.all(np.isfinite(distances))
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
class TestProjectAndRank:
|
| 97 |
+
def test_returns_sorted_indices(self, engine):
|
| 98 |
+
"""project_and_rank should return indices sorted by ascending distance."""
|
| 99 |
+
query = np.zeros(4, dtype=np.float32)
|
| 100 |
+
candidates = np.array([
|
| 101 |
+
[5.0, 5.0, 5.0, 5.0], # Farthest
|
| 102 |
+
[0.01, 0.01, 0.01, 0.01], # Closest
|
| 103 |
+
[1.0, 1.0, 1.0, 1.0], # Middle
|
| 104 |
+
], dtype=np.float32)
|
| 105 |
+
indices = engine.project_and_rank(query, candidates, 3)
|
| 106 |
+
assert indices[0] == 1 # Closest first
|
| 107 |
+
assert indices[-1] == 0 # Farthest last
|
| 108 |
+
|
| 109 |
+
def test_buffer_reuse_stability(self, engine):
|
| 110 |
+
"""Running multiple times should produce consistent results (no buffer aliasing)."""
|
| 111 |
+
query = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
|
| 112 |
+
candidates = np.random.randn(5, 4).astype(np.float32)
|
| 113 |
+
|
| 114 |
+
result1 = engine.project_and_rank(query, candidates, 5)
|
| 115 |
+
result2 = engine.project_and_rank(query, candidates, 5)
|
| 116 |
+
np.testing.assert_array_equal(result1, result2)
|
tests/test_knowledge_extractor.py
ADDED
|
@@ -0,0 +1,216 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
GovBridge India — Unit Tests for Knowledge Extraction Pipeline
|
| 3 |
+
Sprint 30: PROJECT INDRA Phase 1.5
|
| 4 |
+
|
| 5 |
+
Tests cover:
|
| 6 |
+
1. Prompt construction
|
| 7 |
+
2. Pydantic validation (valid predicates, invalid predicates)
|
| 8 |
+
3. Text chunking with overlap
|
| 9 |
+
4. Triplet deduplication
|
| 10 |
+
5. Entity normalization
|
| 11 |
+
|
| 12 |
+
Run: pytest tests/test_knowledge_extractor.py -v
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
import sys
|
| 16 |
+
import os
|
| 17 |
+
import pytest
|
| 18 |
+
|
| 19 |
+
# Ensure gov_backend and ingestion are on sys.path
|
| 20 |
+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| 21 |
+
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ingestion"))
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
# ═══════════════════════════════════════════════════════════════
|
| 25 |
+
# TEST 1: Predicate Ontology
|
| 26 |
+
# ═══════════════════════════════════════════════════════════════
|
| 27 |
+
|
| 28 |
+
class TestGovPredicate:
|
| 29 |
+
def test_all_10_predicates_exist(self):
|
| 30 |
+
from prompts import VALID_PREDICATES
|
| 31 |
+
expected = {
|
| 32 |
+
"AMENDS", "SUPERSEDES", "REPEALS", "ENACTS",
|
| 33 |
+
"DELEGATES_TO", "FUNDS", "MANDATES", "APPLIES_TO",
|
| 34 |
+
"DEFINES", "PENALIZES"
|
| 35 |
+
}
|
| 36 |
+
assert VALID_PREDICATES == expected
|
| 37 |
+
|
| 38 |
+
def test_predicate_enum_values(self):
|
| 39 |
+
from prompts import GovPredicate
|
| 40 |
+
assert GovPredicate.AMENDS.value == "AMENDS"
|
| 41 |
+
assert GovPredicate.DELEGATES_TO.value == "DELEGATES_TO"
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
# ═══════════════════════════════════════════════════════════════
|
| 45 |
+
# TEST 2: Pydantic Validation
|
| 46 |
+
# ═══════════════════════════════════════════════════════════════
|
| 47 |
+
|
| 48 |
+
class TestKnowledgeTriplet:
|
| 49 |
+
def test_valid_triplet(self):
|
| 50 |
+
from knowledge_extractor import KnowledgeTriplet
|
| 51 |
+
t = KnowledgeTriplet(
|
| 52 |
+
subject="Finance Act 2024",
|
| 53 |
+
predicate="AMENDS",
|
| 54 |
+
object="Income Tax Act 1961"
|
| 55 |
+
)
|
| 56 |
+
assert t.predicate == "AMENDS"
|
| 57 |
+
assert t.subject == "Finance Act 2024"
|
| 58 |
+
|
| 59 |
+
def test_predicate_normalized_to_uppercase(self):
|
| 60 |
+
from knowledge_extractor import KnowledgeTriplet
|
| 61 |
+
t = KnowledgeTriplet(
|
| 62 |
+
subject="Act A",
|
| 63 |
+
predicate="amends", # lowercase
|
| 64 |
+
object="Act B"
|
| 65 |
+
)
|
| 66 |
+
assert t.predicate == "AMENDS"
|
| 67 |
+
|
| 68 |
+
def test_invalid_predicate_rejected(self):
|
| 69 |
+
from knowledge_extractor import KnowledgeTriplet
|
| 70 |
+
from pydantic import ValidationError
|
| 71 |
+
with pytest.raises(ValidationError):
|
| 72 |
+
KnowledgeTriplet(
|
| 73 |
+
subject="Act A",
|
| 74 |
+
predicate="MODIFIES", # Not in ontology
|
| 75 |
+
object="Act B"
|
| 76 |
+
)
|
| 77 |
+
|
| 78 |
+
def test_empty_subject_rejected(self):
|
| 79 |
+
from knowledge_extractor import KnowledgeTriplet
|
| 80 |
+
from pydantic import ValidationError
|
| 81 |
+
with pytest.raises(ValidationError):
|
| 82 |
+
KnowledgeTriplet(
|
| 83 |
+
subject="",
|
| 84 |
+
predicate="AMENDS",
|
| 85 |
+
object="Act B"
|
| 86 |
+
)
|
| 87 |
+
|
| 88 |
+
def test_whitespace_trimmed(self):
|
| 89 |
+
from knowledge_extractor import KnowledgeTriplet
|
| 90 |
+
t = KnowledgeTriplet(
|
| 91 |
+
subject=" Finance Act 2024 ",
|
| 92 |
+
predicate=" REPEALS ",
|
| 93 |
+
object=" Old Act "
|
| 94 |
+
)
|
| 95 |
+
assert t.subject == "Finance Act 2024"
|
| 96 |
+
assert t.object == "Old Act"
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
# ═══════════════════════════════════════════════════════════════
|
| 100 |
+
# TEST 3: Text Chunking
|
| 101 |
+
# ═══════════════════════════════════════════════════════════════
|
| 102 |
+
|
| 103 |
+
class TestTextChunking:
|
| 104 |
+
def test_short_text_single_chunk(self):
|
| 105 |
+
from knowledge_extractor import chunk_text_for_extraction
|
| 106 |
+
chunks = chunk_text_for_extraction("Short text", chunk_size=1000)
|
| 107 |
+
assert len(chunks) == 1
|
| 108 |
+
assert chunks[0] == "Short text"
|
| 109 |
+
|
| 110 |
+
def test_long_text_produces_multiple_chunks(self):
|
| 111 |
+
from knowledge_extractor import chunk_text_for_extraction
|
| 112 |
+
# Create text that's clearly > 1000 chars
|
| 113 |
+
text = "A" * 500 + "\n\n" + "B" * 500 + "\n\n" + "C" * 500
|
| 114 |
+
chunks = chunk_text_for_extraction(text, chunk_size=600, overlap=100)
|
| 115 |
+
assert len(chunks) >= 2
|
| 116 |
+
|
| 117 |
+
def test_overlap_is_applied(self):
|
| 118 |
+
from knowledge_extractor import chunk_text_for_extraction
|
| 119 |
+
# Create text with clear paragraph boundaries
|
| 120 |
+
text = ("Para one. " * 50) + "\n\n" + ("Para two. " * 50) + "\n\n" + ("Para three. " * 50)
|
| 121 |
+
chunks = chunk_text_for_extraction(text, chunk_size=300, overlap=50)
|
| 122 |
+
# Verify overlap: last N chars of chunk[0] should appear in chunk[1]
|
| 123 |
+
if len(chunks) > 1:
|
| 124 |
+
tail = chunks[0][-50:]
|
| 125 |
+
assert tail in chunks[1] or len(chunks[1]) > 50
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
# ═══════════════════════════════════════════════════════════════
|
| 129 |
+
# TEST 4: Deduplication
|
| 130 |
+
# ═══════════════════════════════════════════════════════════════
|
| 131 |
+
|
| 132 |
+
class TestDeduplication:
|
| 133 |
+
def test_exact_duplicates_removed(self):
|
| 134 |
+
from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
|
| 135 |
+
triplets = [
|
| 136 |
+
KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
|
| 137 |
+
KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
|
| 138 |
+
]
|
| 139 |
+
result = deduplicate_triplets(triplets)
|
| 140 |
+
assert len(result) == 1
|
| 141 |
+
|
| 142 |
+
def test_case_insensitive_dedup(self):
|
| 143 |
+
from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
|
| 144 |
+
triplets = [
|
| 145 |
+
KnowledgeTriplet(subject="Finance Act", predicate="AMENDS", object="IT Act"),
|
| 146 |
+
KnowledgeTriplet(subject="finance act", predicate="AMENDS", object="it act"),
|
| 147 |
+
]
|
| 148 |
+
result = deduplicate_triplets(triplets)
|
| 149 |
+
assert len(result) == 1
|
| 150 |
+
|
| 151 |
+
def test_different_predicates_kept(self):
|
| 152 |
+
from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
|
| 153 |
+
triplets = [
|
| 154 |
+
KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
|
| 155 |
+
KnowledgeTriplet(subject="Act A", predicate="SUPERSEDES", object="Act B"),
|
| 156 |
+
]
|
| 157 |
+
result = deduplicate_triplets(triplets)
|
| 158 |
+
assert len(result) == 2
|
| 159 |
+
|
| 160 |
+
def test_article_noise_dedup(self):
|
| 161 |
+
from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
|
| 162 |
+
triplets = [
|
| 163 |
+
KnowledgeTriplet(subject="The Finance Act", predicate="AMENDS", object="The IT Act"),
|
| 164 |
+
KnowledgeTriplet(subject="Finance Act", predicate="AMENDS", object="IT Act"),
|
| 165 |
+
]
|
| 166 |
+
result = deduplicate_triplets(triplets)
|
| 167 |
+
assert len(result) == 1
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
# ═══════════════════════════════════════════════════════════════
|
| 171 |
+
# TEST 5: Entity Normalization
|
| 172 |
+
# ═══════════════════════════════════════════════════════════════
|
| 173 |
+
|
| 174 |
+
class TestEntityNormalization:
|
| 175 |
+
def test_normalize_strips_whitespace(self):
|
| 176 |
+
from knowledge_extractor import _normalize_entity
|
| 177 |
+
assert _normalize_entity(" Finance Act ") == "finance act"
|
| 178 |
+
|
| 179 |
+
def test_normalize_removes_articles(self):
|
| 180 |
+
from knowledge_extractor import _normalize_entity
|
| 181 |
+
assert _normalize_entity("The Finance Act") == "finance act"
|
| 182 |
+
|
| 183 |
+
def test_normalize_collapses_spaces(self):
|
| 184 |
+
from knowledge_extractor import _normalize_entity
|
| 185 |
+
assert _normalize_entity("Finance Act 2024") == "finance act 2024"
|
| 186 |
+
|
| 187 |
+
def test_normalize_removes_legal_filler(self):
|
| 188 |
+
from knowledge_extractor import _normalize_entity
|
| 189 |
+
assert _normalize_entity("said Ministry of Finance") == "ministry of finance"
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
# ═══════════════════════════════════════════════════════════════
|
| 193 |
+
# TEST 6: Prompt Construction
|
| 194 |
+
# ═══════════════════════════════════════════════════════════════
|
| 195 |
+
|
| 196 |
+
class TestPromptConstruction:
|
| 197 |
+
def test_build_chunk_user_prompt_includes_context(self):
|
| 198 |
+
from prompts import build_chunk_user_prompt
|
| 199 |
+
ctx = {
|
| 200 |
+
"document_title": "Finance Act 2024",
|
| 201 |
+
"primary_body": "Ministry of Finance",
|
| 202 |
+
"key_definitions": ["tax", "income"],
|
| 203 |
+
"document_date": "2024-01-15"
|
| 204 |
+
}
|
| 205 |
+
prompt = build_chunk_user_prompt(ctx, "Some chunk text here")
|
| 206 |
+
assert "Finance Act 2024" in prompt
|
| 207 |
+
assert "Ministry of Finance" in prompt
|
| 208 |
+
assert "Some chunk text here" in prompt
|
| 209 |
+
assert "tax, income" in prompt
|
| 210 |
+
|
| 211 |
+
def test_build_chunk_user_prompt_handles_empty_context(self):
|
| 212 |
+
from prompts import build_chunk_user_prompt
|
| 213 |
+
ctx = {}
|
| 214 |
+
prompt = build_chunk_user_prompt(ctx, "Chunk text")
|
| 215 |
+
assert "Unknown" in prompt
|
| 216 |
+
assert "Chunk text" in prompt
|