univin-weights / main.py
Vnmrsharma's picture
Update main.py
a93dc66 verified
import openai
import re
import os
from flask import Flask, request, jsonify
app = Flask(__name__)
application = app
openai.api_key = os.environ.get("OPENAI_API_KEY")
SECRET_KEY_CONST = os.environ.get("SECRET_KEY")
def extract_weight(text):
"""
Extracts weight in grams from a given text.
"""
patterns = [
r'(\d+(?:\.\d+)?)\s*(?:g|grams)',
r'(\d+(?:\.\d+)?)\s*(?:kg|kilograms)'
]
for pattern in patterns:
match = re.search(pattern, text, re.IGNORECASE)
if match:
weight = float(match.group(1))
if 'kg' in pattern:
weight *= 1000
return weight
return None
def get_weight_from_api(product_name, prompt):
"""
Sends a prompt to the OpenAI API and returns the response content.
"""
response = openai.chat.completions.create(
model="gpt-4o-search-preview",
messages=[{"role": "user", "content": prompt}],
max_tokens=10,
stream=False
)
return response.choices[0].message.content.strip()
def get_product_weight(product_name):
"""
Retrieves the product weight by first searching for a numerical value using OpenAI API.
If not found, it makes an estimate. Finally, it adds packaging weight of 50 grams.
Returns the total weight in grams (or None if not computable).
"""
# Step 1: Search for product weight
search_prompt = (
f"What is the weight in grams of {product_name}? "
"Provide only the numerical value."
)
search_response = get_weight_from_api(product_name, search_prompt)
weight = extract_weight(search_response)
# Step 2: If not found, estimate weight
if weight is None:
estimate_prompt = (
f"Estimate the weight in grams of a {product_name}. "
"Provide only the numerical value."
)
estimate_response = get_weight_from_api(product_name, estimate_prompt)
try:
weight = float(estimate_response)
except ValueError:
# Unable to determine the weight, return None.
return None
# Step 3: Add packaging weight (assumed to be 50 grams)
packaging_weight = 50
total_weight = weight + packaging_weight
return total_weight
@app.route('/get_weight', methods=['POST'])
def weight_endpoint():
"""
Endpoint to retrieve product weight.
Expected JSON payload:
{
"product_name": "name of the product",
"secret_key": "your_secret_key"
}
"""
data = request.get_json()
if not data:
return jsonify({"error": "Missing JSON payload"}), 400
secret_key = data.get("secret_key")
product_name = data.get("product_name")
if not secret_key or not product_name:
return jsonify({"error": "Missing secret key or product name"}), 400
if secret_key != SECRET_KEY_CONST:
return jsonify({"error": "Unauthorized access"}), 403
total_weight = get_product_weight(product_name)
if total_weight is None:
return jsonify({"error": "Unable to determine the product weight."}), 500
return jsonify({
"product": product_name,
"total_weight": total_weight
})
if __name__ == "__main__":
app.run(debug=True, host="0.0.0.0")