Update main.py

main.py

@@ -44,9 +44,6 @@ class StudentRecord(BaseModel):
     section: str = Field(default="", description="Section letter")
     mobile: str = Field(default="", description="Mobile number")
 
-class ExtractResponse(BaseModel):
-    students: List[StudentRecord] = Field(default_factory=list)
-
 class PDFRequest(BaseModel):
     pdfUrl: str
 
@@ -76,10 +73,9 @@ Return ONLY this JSON format:
   ]
 }
 
-IMPORTANT: Ensure ALL students are included. If you see 200 students, return all 200. No truncation allowed.
+IMPORTANT: Ensure ALL students are included. No truncation allowed. Close all JSON arrays and objects properly.
 """,
-    output_type=ExtractResponse,
-    model_settings=ModelSettings(temperature=0.05, top_p=0.7, max_tokens=4000)
+    model_settings=ModelSettings(temperature=0.05, top_p=0.7, max_tokens=8000)
 )
 
 runner = Runner()
@@ -103,6 +99,41 @@ def download_and_extract_text(pdf_url: str) -> str:
     return text
 
 
+def parse_json_from_output(output: str) -> dict:
+    """Parse JSON from agent output, handling truncation and errors"""
+    if not output:
+        return {"students": []}
+    
+    # Try direct JSON parse first
+    try:
+        return json.loads(output)
+    except json.JSONDecodeError:
+        pass
+    
+    # Try to extract JSON from markdown or other formatting
+    json_match = re.search(r'\{[\s\S]*\}', output)
+    if not json_match:
+        return {"students": []}
+    
+    json_str = json_match.group(0)
+    
+    # Try to fix incomplete JSON
+    try:
+        return json.loads(json_str)
+    except json.JSONDecodeError:
+        # Try closing the JSON if it's truncated
+        open_braces = json_str.count('{') - json_str.count('}')
+        open_brackets = json_str.count('[') - json_str.count(']')
+        
+        json_str = json_str.rstrip().rstrip(',') + ']' * open_brackets + '}' * open_braces
+        
+        try:
+            return json.loads(json_str)
+        except json.JSONDecodeError as e:
+            print(f"⚠️ Failed to parse JSON even after fixing: {e}")
+            return {"students": []}
+
+
 def regex_fallback_extraction(text: str) -> dict:
     """Robust regex-based extraction for when agent fails"""
     print("🔄 Using regex fallback for extraction...")
@@ -111,17 +142,17 @@ def regex_fallback_extraction(text: str) -> dict:
     # Try multiple regex patterns for flexibility
     patterns = [
         # Pattern 1: Name | Roll | Class | Section | Mobile
-        r'^([A-Za-z\s]+?)\s*\|\s*(\d+)\s*\|\s*([\w\d\s\.]+?)\s*\|\s*([A-Z])\s*\|\s*(\d+)',
+        r'^([A-Za-z\s]+?)\s*\|\s*(\d+)\s*\|\s*([\w\d\s\.,-]+?)\s*\|\s*([A-Za-z0-9\s,.-]+?)\s*\|\s*(\d+)',
         # Pattern 2: Space-separated format
-        r'^([A-Za-z\s]+?)\s+(\d+)\s+([\w\d\s\.]+?)\s+([A-Z])\s+(\d{10,})',
-        # Pattern 3: Tab-separated or multiple spaces
-        r'^([A-Za-z\s]+?)\t+(\d+)\t+([\w\d\s\.]+?)\t+([A-Z])\t+(\d{10,})',
+        r'^([A-Za-z\s]+?)\s+(\d{8,})\s+([\w\d\s\.,-]+?)\s+([A-Za-z0-9\s,.-]+?)\s+(\d{11,})',
+        # Pattern 3: Tab-separated
+        r'^([A-Za-z\s]+?)\t+(\d{8,})\t+([\w\d\s\.,-]+?)\t+([A-Za-z0-9\s,.-]+?)\t+(\d{11,})',
     ]
     
     seen = set()
     for line in text.splitlines():
         line = line.strip()
-        if not line or "name" in line.lower() or "roll" in line.lower():
+        if not line or "name" in line.lower() or "roll" in line.lower() or "generated" in line.lower():
             continue
         
         for pattern in patterns:
@@ -130,11 +161,11 @@ def regex_fallback_extraction(text: str) -> dict:
                 name = match.group(1).strip()
                 roll_no = match.group(2)
                 class_name = match.group(3).strip()
-                section = match.group(4)
+                section = match.group(4).strip()
                 mobile = match.group(5)
                 
                 key = (name, roll_no)
-                if key not in seen:
+                if key not in seen and name and roll_no:
                     seen.add(key)
                     students.append({
                         "name": name,
@@ -150,10 +181,9 @@ def regex_fallback_extraction(text: str) -> dict:
 
 
 async def extract_from_text_chunked(text: str) -> dict:
-    """Runs the agent with chunking for large datasets"""
+    """Runs the agent with flexible JSON parsing for large datasets"""
     print(f"📄 Extracting from {len(text)} characters...")
     
-    # Try direct extraction first
     try:
         resp = await runner.run(
             student_agent,
@@ -168,14 +198,23 @@ async def extract_from_text_chunked(text: str) -> dict:
             output = resp.final_output
         
         if output:
-            result = output.model_dump() if hasattr(output, "model_dump") else output
+            # Convert to string if needed
+            if isinstance(output, str):
+                output_str = output
+            else:
+                output_str = str(output)
+            
+            # Parse JSON flexibly
+            result = parse_json_from_output(output_str)
             student_count = len(result.get("students", []))
             print(f"✅ Agent extracted {student_count} students")
             
             if student_count > 0:
                 return result
+            else:
+                print("⚠️ Agent returned empty results")
     except Exception as e:
-        print(f"⚠️ Agent extraction failed: {e}")
+        print(f"⚠️ Agent extraction error: {e}")
     
     # Fallback to regex if agent fails or returns empty
     return regex_fallback_extraction(text)
@@ -187,8 +226,11 @@ def clean_and_deduplicate(students: List[dict]) -> List[dict]:
     unique = []
     
     for s in students:
-        key = (s.get("name", "").strip(), s.get("roll_no", "").strip())
-        if key and key not in seen:
+        name = str(s.get("name", "")).strip()
+        roll_no = str(s.get("roll_no", "")).strip()
+        key = (name, roll_no)
+        
+        if key and key[0] and key[1] and key not in seen:
             seen.add(key)
             unique.append(s)
     
@@ -204,7 +246,7 @@ async def extract_student(req: PDFRequest):
     
     Features:
     - Handles large datasets (200+ students)
-    - Agent-based extraction with fallback
+    - No strict validation - flexible JSON parsing
     - Regex backup for comprehensive coverage
     - Automatic deduplication
     """
@@ -220,11 +262,15 @@ async def extract_student(req: PDFRequest):
             "success": True,
             "pdfUrl": req.pdfUrl,
             "total_students": len(cleaned),
-            "students": cleaned,
-            "raw_text_preview": text[:1000]
+            "students": cleaned
         }
     except Exception as e:
         print(f"❌ Error: {e}")
         import traceback
         traceback.print_exc()
-        return {"success": False, "error": str(e)}
+        return {
+            "success": False,
+            "error": str(e),
+            "total_students": 0,
+            "students": []
+        }