Replaced brittle prompt with pydantic schema validation

Author: Nossks

GSoC25/NEF/NEF.py

@@ -9,6 +9,8 @@
 from urllib.parse import quote
 from getpass import getpass
 
+from output_schema import DisambiguationResult
+
 # =============== Gemini client bootstrap ===============
 try:
     from google import genai
@@ -280,13 +282,6 @@ def _safe_idx(lst: List[Any], value: Any) -> Optional[int]:
             i = 0
         return max(0, min(i, len(lst) - 1))
 
-    @staticmethod
-    def _get_json(text: Optional[str]) -> Dict[str, Any]:
-        try:
-            return json.loads((text or "").strip() or "{}")
-        except Exception:
-            return {}
-
     # --------------------- main entrypoint ---------------------
 
     def disambiguate_triple(
@@ -332,46 +327,45 @@ def disambiguate_triple(
         subj_list_text = self._fmt_indexed(subject_candidates)
         obj_list_text  = self._fmt_indexed(object_candidates)
 
-        prompt = f"""Pick the best RDF triple using ONLY these options.
+        prompt = f"""
+        Analyze the context and disambiguate the triple by selecting the correct Subject, Predicate, and Object.
 
-Allowed predicate URIs:
-{pred_list_text}
+        Context: {context}
 
-Subject candidates (choose by INDEX):
-{subj_list_text}
+        Allowed Predicates (Select one URI):
+        {pred_list_text}
 
-Object candidates (choose by INDEX):
-{obj_list_text}
+        Subject Candidates (Select by Index):
+        {subj_list_text}
 
-Context (helps decide, but does NOT add new options):
-{context}
+        Object Candidates (Select by Index):
+        {obj_list_text}
+        """
 
-Return ONLY strict JSON on one line (no prose):
-{{"subject_index": 0, "predicate": "URI", "object_index": 0}}
-Rules:
-- "predicate" MUST be exactly one URI from Allowed predicate URIs.
-- "subject_index" MUST be an integer index from Subject candidates.
-- "object_index" MUST be an integer index from Object candidates.
-- Do not invent or modify URIs. Do not swap roles.
-"""
+        generation_config = {
+            "response_mime_type": "application/json",
+            "response_schema": DisambiguationResult  # <--- The Pydantic class output
+        }
 
         # Call the model (Gemini client style)
         resp = self.client.models.generate_content(
             model=self.model_name,
             contents=prompt,
-            config={"response_mime_type": "application/json"},
+            config=generation_config
         )
 
-        data = self._get_json(getattr(resp, "text", None))
+        data = DisambiguationResult.model_validate_json(resp.text)
+        pred_uri = data.predicate_uri
+        s_idx = data.subject_index
+        o_idx = data.object_index
 
         # Validate predicate
-        pred_uri = data.get("predicate", "")
         if pred_uri not in allowed:
-            pred_uri = allowed[0]
+            pred_uri = None      ## if none
 
         # Clamp indices and map to URIs
-        si = self._safe_idx(subject_candidates, data.get("subject_index", 0))
-        oi = self._safe_idx(object_candidates,  data.get("object_index", 0))
+        si = self._safe_idx(subject_candidates, s_idx)
+        oi = self._safe_idx(object_candidates,  o_idx)
 
         s_uri = subject_candidates[si][0] if (subject_candidates and si is not None) else ""
         o_uri = object_candidates[oi][0] if (object_candidates and oi is not None) else ""
@@ -392,7 +386,7 @@ def disambiguate_triple(
         # Build meta (compatible with your previous code)
         chosen_sim, rank0 = sim_map.get(pred_uri, (None, None))
         meta = {
-            "label": "candidate",
+            "label": "candidate" if pred_uri else "hallucinaiton_rejected",
             "chosen_similarity": float(chosen_sim) if chosen_sim is not None else None,
             "rank_in_topk": (rank0 + 1) if rank0 is not None else None,
             "topk": total_k,

GSoC25/NEF/output_schema.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel, Field
+
+class DisambiguationResult(BaseModel):
+    subject_index: int = Field(description="The integer index of the correct subject from the provided list.")
+    predicate_uri: str = Field(description="The exact URI of the predicate selected from the allowed list.")
+    object_index: int = Field(description="The integer index of the correct object from the provided list.")