Hugging Face's logo

Spaces:
harshrawat18
/
govbridge-api
Sleeping

App Files Community
govbridge-api
File size: 9,804 Bytes
aa68ee9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
"""
GovBridge India — Unit Tests for Knowledge Extraction Pipeline
Sprint 30: PROJECT INDRA Phase 1.5

Tests cover:
  1. Prompt construction
  2. Pydantic validation (valid predicates, invalid predicates)
  3. Text chunking with overlap
  4. Triplet deduplication
  5. Entity normalization

Run: pytest tests/test_knowledge_extractor.py -v
"""

import sys
import os
import pytest

# Ensure gov_backend and ingestion are on sys.path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ingestion"))


# ═══════════════════════════════════════════════════════════════
# TEST 1: Predicate Ontology
# ═══════════════════════════════════════════════════════════════

class TestGovPredicate:
    def test_all_10_predicates_exist(self):
        from prompts import VALID_PREDICATES
        expected = {
            "AMENDS", "SUPERSEDES", "REPEALS", "ENACTS",
            "DELEGATES_TO", "FUNDS", "MANDATES", "APPLIES_TO",
            "DEFINES", "PENALIZES"
        }
        assert VALID_PREDICATES == expected

    def test_predicate_enum_values(self):
        from prompts import GovPredicate
        assert GovPredicate.AMENDS.value == "AMENDS"
        assert GovPredicate.DELEGATES_TO.value == "DELEGATES_TO"


# ═══════════════════════════════════════════════════════════════
# TEST 2: Pydantic Validation
# ═══════════════════════════════════════════════════════════════

class TestKnowledgeTriplet:
    def test_valid_triplet(self):
        from knowledge_extractor import KnowledgeTriplet
        t = KnowledgeTriplet(
            subject="Finance Act 2024",
            predicate="AMENDS",
            object="Income Tax Act 1961"
        )
        assert t.predicate == "AMENDS"
        assert t.subject == "Finance Act 2024"

    def test_predicate_normalized_to_uppercase(self):
        from knowledge_extractor import KnowledgeTriplet
        t = KnowledgeTriplet(
            subject="Act A",
            predicate="amends",  # lowercase
            object="Act B"
        )
        assert t.predicate == "AMENDS"

    def test_invalid_predicate_rejected(self):
        from knowledge_extractor import KnowledgeTriplet
        from pydantic import ValidationError
        with pytest.raises(ValidationError):
            KnowledgeTriplet(
                subject="Act A",
                predicate="MODIFIES",  # Not in ontology
                object="Act B"
            )

    def test_empty_subject_rejected(self):
        from knowledge_extractor import KnowledgeTriplet
        from pydantic import ValidationError
        with pytest.raises(ValidationError):
            KnowledgeTriplet(
                subject="",
                predicate="AMENDS",
                object="Act B"
            )

    def test_whitespace_trimmed(self):
        from knowledge_extractor import KnowledgeTriplet
        t = KnowledgeTriplet(
            subject="  Finance Act 2024  ",
            predicate="  REPEALS  ",
            object="  Old Act  "
        )
        assert t.subject == "Finance Act 2024"
        assert t.object == "Old Act"


# ═══════════════════════════════════════════════════════════════
# TEST 3: Text Chunking
# ═══════════════════════════════════════════════════════════════

class TestTextChunking:
    def test_short_text_single_chunk(self):
        from knowledge_extractor import chunk_text_for_extraction
        chunks = chunk_text_for_extraction("Short text", chunk_size=1000)
        assert len(chunks) == 1
        assert chunks[0] == "Short text"

    def test_long_text_produces_multiple_chunks(self):
        from knowledge_extractor import chunk_text_for_extraction
        # Create text that's clearly > 1000 chars
        text = "A" * 500 + "\n\n" + "B" * 500 + "\n\n" + "C" * 500
        chunks = chunk_text_for_extraction(text, chunk_size=600, overlap=100)
        assert len(chunks) >= 2

    def test_overlap_is_applied(self):
        from knowledge_extractor import chunk_text_for_extraction
        # Create text with clear paragraph boundaries
        text = ("Para one. " * 50) + "\n\n" + ("Para two. " * 50) + "\n\n" + ("Para three. " * 50)
        chunks = chunk_text_for_extraction(text, chunk_size=300, overlap=50)
        # Verify overlap: last N chars of chunk[0] should appear in chunk[1]
        if len(chunks) > 1:
            tail = chunks[0][-50:]
            assert tail in chunks[1] or len(chunks[1]) > 50


# ═══════════════════════════════════════════════════════════════
# TEST 4: Deduplication
# ═══════════════════════════════════════════════════════════════

class TestDeduplication:
    def test_exact_duplicates_removed(self):
        from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
        triplets = [
            KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
            KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
        ]
        result = deduplicate_triplets(triplets)
        assert len(result) == 1

    def test_case_insensitive_dedup(self):
        from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
        triplets = [
            KnowledgeTriplet(subject="Finance Act", predicate="AMENDS", object="IT Act"),
            KnowledgeTriplet(subject="finance act", predicate="AMENDS", object="it act"),
        ]
        result = deduplicate_triplets(triplets)
        assert len(result) == 1

    def test_different_predicates_kept(self):
        from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
        triplets = [
            KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
            KnowledgeTriplet(subject="Act A", predicate="SUPERSEDES", object="Act B"),
        ]
        result = deduplicate_triplets(triplets)
        assert len(result) == 2

    def test_article_noise_dedup(self):
        from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
        triplets = [
            KnowledgeTriplet(subject="The Finance Act", predicate="AMENDS", object="The IT Act"),
            KnowledgeTriplet(subject="Finance Act", predicate="AMENDS", object="IT Act"),
        ]
        result = deduplicate_triplets(triplets)
        assert len(result) == 1


# ═══════════════════════════════════════════════════════════════
# TEST 5: Entity Normalization
# ═══════════════════════════════════════════════════════════════

class TestEntityNormalization:
    def test_normalize_strips_whitespace(self):
        from knowledge_extractor import _normalize_entity
        assert _normalize_entity("  Finance Act  ") == "finance act"

    def test_normalize_removes_articles(self):
        from knowledge_extractor import _normalize_entity
        assert _normalize_entity("The Finance Act") == "finance act"

    def test_normalize_collapses_spaces(self):
        from knowledge_extractor import _normalize_entity
        assert _normalize_entity("Finance   Act   2024") == "finance act 2024"

    def test_normalize_removes_legal_filler(self):
        from knowledge_extractor import _normalize_entity
        assert _normalize_entity("said Ministry of Finance") == "ministry of finance"


# ═══════════════════════════════════════════════════════════════
# TEST 6: Prompt Construction
# ═══════════════════════════════════════════════════════════════

class TestPromptConstruction:
    def test_build_chunk_user_prompt_includes_context(self):
        from prompts import build_chunk_user_prompt
        ctx = {
            "document_title": "Finance Act 2024",
            "primary_body": "Ministry of Finance",
            "key_definitions": ["tax", "income"],
            "document_date": "2024-01-15"
        }
        prompt = build_chunk_user_prompt(ctx, "Some chunk text here")
        assert "Finance Act 2024" in prompt
        assert "Ministry of Finance" in prompt
        assert "Some chunk text here" in prompt
        assert "tax, income" in prompt

    def test_build_chunk_user_prompt_handles_empty_context(self):
        from prompts import build_chunk_user_prompt
        ctx = {}
        prompt = build_chunk_user_prompt(ctx, "Chunk text")
        assert "Unknown" in prompt
        assert "Chunk text" in prompt