nextchat / ai /entity.py
servionsoft's picture
Update ai/entity.py
10e9c04 verified
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