ai/entity.py

import openai
import os
from dotenv import load_dotenv

load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")

def gpt_extract_field(field_name: str, sentence: str) -> str | None:
    """
    Extract or validate fields like origin, destination, date, people, or days.
    Handles both direct single-word input and full sentences.
    """
    sentence = sentence.strip().lower()

    # Direct return if the input is just one word (likely the answer itself)
    if len(sentence.split()) == 1:
        return sentence

    examples = {
        "date": "'20 June 2025', 'June 20 2025', '2025/06/20', 'tomorrow', 'in 2 days'",
        "origin": "'Karachi', 'Lahore', 'Dubai'",
        "destination": "'New York', 'Istanbul'",
        "people": "'one', '1', 'two', '5', 'seven', 'a group of 4'",
        "days": "'ten', '5','1','one', 'fourteen'"
    }

    if field_name in ["people", "days"]:
        field_instruction = f"""
Your task is to extract the number of {field_name} from the input.
Rules:
- Return the value as a digit (e.g., 'six' -> 6, 'thirty two' -> 32).
- If the value is already a number, return it as is.
- If it's a phrase like 'a group of four', return 4.
- If it's unclear or no number is present, return 'none'.
- Respond ONLY with the number or 'none'. No extra text.
Examples:
Input: 'five' → Output: 5
Input: '32' → Output: 32
Input: 'a group of four' → Output: 4
Input: 'no idea' → Output: none
Input: 'me and one friend' → Output: 2
""".strip()
    else:
        field_instruction = f"""
If the input contains the {field_name}, extract it.
Rules:
- If it's written as words like 'five' or 'two', convert to digits like 5 or 2.
- If the input is directly the {field_name} itself, return it as is.
- For dates, preserve the format as given: {examples.get(field_name, '')}.
- If invalid or not present, respond 'none'.
Only reply with the {field_name} or 'none'.
""".strip()

    prompt = (
        f"Extract the {field_name} from the following input.\n"
        f"Input: '{sentence}'\n"
        f"{field_instruction}"
    )

    try:
        client = openai.OpenAI()
        response = client.chat.completions.create(
            model="gpt-4-turbo",
            messages=[
                {"role": "system", "content": "You are a field extractor. Be concise and return only the requested value."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=30,
            temperature=0,
        )
        extracted_value = response.choices[0].message.content.strip().lower()
        return extracted_value if extracted_value != "none" else None

    except Exception as e:
        print(f"OpenAI API Error: {e}")
        return None