feat: add basic schema validation layer for extraction outputs

Author: Lochit-Vinay

src/llm.py

@@ -2,6 +2,7 @@
 import os
 import requests
 from api.services.prompt_builder import build_extraction_prompt
+from src.validation import validate_extraction
 
 def safe_extract_value(response: str):
     if not response:
@@ -160,4 +161,9 @@ def handle_plural_values(self, plural_value):
         return values
 
     def get_data(self):
-        return self._json
+        validated_data, errors = validate_extraction(self._json)
+
+        if errors:
+            print("[Validation Warning]", errors)
+
+        return validated_data

src/main.py

@@ -2,7 +2,8 @@
 # from backend import Fill  
 from commonforms import prepare_form 
 from pypdf import PdfReader
-from controller import Controller
+from typing import Union
+from src.controller import Controller
 
 def input_fields(num_fields: int):
     fields = []

src/validation.py

@@ -0,0 +1,29 @@
+from pydantic import BaseModel, ValidationError
+from typing import Optional
+
+
+class ExtractionSchema(BaseModel):
+    name: Optional[str]
+    location: Optional[str]
+    date: Optional[str]
+    incident_type: Optional[str]
+    description: Optional[str]
+
+    class Config:
+        extra = "allow"
+
+
+def validate_extraction(data: dict):
+    try:
+        validated = ExtractionSchema(**data)
+        return validated.dict(), None
+    except ValidationError as e:
+        formatted_errors = []
+
+        for err in e.errors():
+            formatted_errors.append({
+                "field": err.get("loc", ["unknown"])[0],
+                "issue": err.get("msg", "Invalid value")
+            })
+
+        return data, formatted_errors