# Entity extraction function that provide uses the query and provide stock name and quantity in dictionary format
import openai
from src.constants import *
openai.api_key = os.getenv("OPENAI_API_KEY")
import json
def system_message(labels):
return f"""
You are an expert in Natural Language Processing. Your task is to identify common Named Entities (NER) in a given text.
The possible common Named Entities (NER) types are exclusively: ({", ".join(labels)})."""
def assisstant_message():
return f"""
EXAMPLE:
Text: 'I want to buy 10 shares of MSFT. Can you help me with that? I also want information about last month stock prices of MSFT. Also tell me USD to INR conversion rate. What is the price of Nvidia stocks?'
{{"currency": ["USD", "INR"],"stock": ["MSFT", "Nvidia"],"quantity": ["10"],"date": ["last month"]
}}
--"""
def user_message(text):
return f"""
TASK:
Text: {text}
"""
def run_openai_task(labels, text):
messages = [
{"role": "system", "content": system_message(labels=labels)},
{"role": "assistant", "content": assisstant_message()},
{"role": "user", "content": user_message(text=text)}]
try:
response = openai.chat.completions.create(
model = "gpt-3.5-turbo-0613",
messages = messages,
temperature=0,
frequency_penalty=0,
presence_penalty=0
)
response_message = response.choices[0].message
return response_message.content
except Exception as e:
return str(e)
def extract_entities(query, user_intent=None):
labels = [
"currency", # Any country currency
"stock", # Any stock name or abbreviation
"quantity", # Quantity of stock and currency
"date" # absolute or relative dates or periods(12 Feb 2023, yesterday, last month, next week, 1 month, 2 weeks)
]
response = run_openai_task(labels, query)
# Try to convert the response to json
try:
response = json.loads(response)
# If response is a dictionary, return it else return an empty dictionary
labels = ["stock", "currency", "quantity", "date"]
if isinstance(response, dict):
# Check the available keys in the response and if not present, add them with empty list
for label in labels:
if label not in response:
response[label] = []
return response
else:
return {"stock": [], "currency": [], "quantity": [], "date": []}
except:
return {"stock": [], "currency": [], "quantity": [], "date": []}