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:
    """
    Use OpenAI GPT to extract only the specified field from the given sentence.
    Args:
        field_name (str): The field to extract, e.g., 'days', 'destination', 'origin', 'adults', 'date'.
        sentence (str): The user input sentence to extract the data from.

    Returns:
        str | None: The extracted field value as a string, or None if not found.
    """
    prompt = (
        f"Extract only the {field_name} from this sentence:\n"
        f"Sentence: '{sentence}'\n\n"
        f"If the {field_name} is a number given as a word (e.g., 'ten', 'five'), convert it to digits (e.g., 10, 5).\n"
        f"If the {field_name} is not present or invalid, respond with 'none'.\n"
        f"Respond ONLY with the extracted {field_name} value (no extra text)."
    )

    try:
        client = openai.OpenAI()
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "You are an assistant that extracts specific information from sentences."},
                {"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