removed asyncio

Changes: Cleaned imports, removed nest_asyncio, fixed output_type argument, added robust result parsing.

agents/extractor.py

@@ -1,29 +1,28 @@
 import os
-import nest_asyncio
 from pydantic_ai import Agent
 from .models import InvoiceSchema, ResumeSchema, GenericSchema
 from utils.client import get_api_key
 from utils.prompt_loader import load_prompt
 
-# Fix loop for some environments
-nest_asyncio.apply()
-
-# Ensure key is available
+# Ensure API key is loaded into environment for the library
 os.environ["GOOGLE_API_KEY"] = get_api_key()
 
-# Load the SOTA Master Prompt from YAML
+# Load Externalized Prompt
 master_system_prompt = load_prompt("system.yaml")
 
-# Initialize Agent with GEMINI 2.5 FLASH (The 2026 Standard)
+# Initialize Agent (Gemini 2.5 Flash)
+# Note: Using the model ID compatible with Google's GenAI SDK
 agent = Agent(
     'google-gla:models/gemini-2.5-flash',
-    system_prompt=master_system_prompt 
+    system_prompt=master_system_prompt
 )
 
 async def process_data(text: str, schema_type: str):
     """
-    Orchestrates the extraction process using the Externalized Prompt.
+    Orchestrates the extraction process.
+    Selects the correct Pydantic schema and enforces type-safety.
     """
+    # 1. Router Logic
     if schema_type == "invoice":
         target_model = InvoiceSchema
     elif schema_type == "resume":
@@ -32,14 +31,17 @@ async def process_data(text: str, schema_type: str):
         target_model = GenericSchema
 
     try:
-        # Run Agent with Type-Safe Output
+        # 2. Execution (With Schema Enforcement)
+        # Using 'output_type' as required by PydanticAI v1.40+
         result = await agent.run(text, output_type=target_model)
         
-        # Access Data (Universal Access)
-        try:
+        # 3. Result Parsing (Handle version differences in return object)
+        if hasattr(result, 'data'):
             return result.data
-        except:
+        elif hasattr(result, 'output'):
             return result.output
+        else:
+            return result
             
     except Exception as e:
-        return {"error": f"Agent Failure: {str(e)}"}
+        return {"error": f"Extraction Logic Failed: {str(e)}"}