# extractor_gpt4.py
import os
import json
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

SYSTEM_PROMPT = """You are a research document metadata extractor.
Given the text of a research document, extract the specified metadata fields.
Be precise. If you are unsure about a field, use your best judgment based on context."""

EXTRACTION_SCHEMA = {
    "name": "extract_document_metadata",
    "parameters": {
        "type": "object",
        "properties": {
            "methodology_type": {
                "type": "string",
                "enum": ["experimental", "observational", "review", "simulation", "mixed"]
            },
            "dataset_source": {"type": "string"},
            "primary_metric": {"type": "string"},
            "year": {"type": "integer"},
            "confidence_score": {"type": "number"}
        },
        "required": ["methodology_type", "dataset_source", "year"]
    }
}

def extract_metadata_gpt4(document_text: str) -> dict:
    response = client.chat.completions.create(
        model="gpt-4-turbo",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"Extract metadata from this document:\n\n{document_text[:8000]}"}
        ],
        functions=[EXTRACTION_SCHEMA],
        function_call={"name": "extract_document_metadata"},
        temperature=0,
        timeout=30
    )
    try:
        args = response.choices[0].message.function_call.arguments
        return json.loads(args)
    except (AttributeError, json.JSONDecodeError) as e:
        raise ValueError(f"Failed to parse GPT-4 response: {e}")