gemini_api.py

import google.generativeai as genai
import json
import os

genai.configure(api_key=os.getenv("API_KEY"))
generation_config = {
  "temperature": 0.9,
  "top_p": 1,
  "top_k": 0,
  "max_output_tokens": 2048,
  "response_mime_type": "text/plain",
}
safety_settings = [
  {
    "category": "HARM_CATEGORY_HARASSMENT",
    "threshold": "BLOCK_MEDIUM_AND_ABOVE",
  },
  {
    "category": "HARM_CATEGORY_HATE_SPEECH",
    "threshold": "BLOCK_MEDIUM_AND_ABOVE",
  },
  {
    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
    "threshold": "BLOCK_MEDIUM_AND_ABOVE",
  },
  {
    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
    "threshold": "BLOCK_MEDIUM_AND_ABOVE",
  },
]

model = genai.GenerativeModel(
  model_name="gemini-1.0-pro",
  safety_settings=safety_settings,
  generation_config=generation_config,
)

NO_ORDER = "Order Number not provided."
NO_ITEM = "No item present."

adds_store = {
  "pasta": "buy 1 get 1 free on penne",
  "chocolate": "50 percent off on premium dark chocolate",
  "biryani": "100 rupees off on all biryani orders",
  "shoe": "get 1 pair socks free",
  "skirt": "new exclusive summer collection",
  "socks": "buy 2 get 3 free",
  "saree": "10 percent off on silk collection"
}

def model_api(input, prompt_type):
    return prompt_type(input)

def sentiment(input):
    chat_session = model.start_chat(
    history=[
        {
        "role": "user",
        "parts": [
            "Valid sentiments are POSITIVE, NEGATIVE, NEUTRAL, MIXED\nQuery: The shoes I ordered are small. Size is not what was in the catalog. Order number is 11223\nOutput:{\"sentiment\": \"NEGATIVE\"}",
        ],
        },
    ]
    )

    response = chat_session.send_message(f"Query: {input}\nOutput:")

    print(response.text)
    obj = json.loads(response.text)
    return obj['sentiment']

def category(input):
    chat_session = model.start_chat(
    history=[
        {
        "role": "user",
        "parts": [
            "Order categories are Electronics, Clothing, Food, Books\nQuery: The shoes I ordered are small. Size is not what was in the catalog. Order number is 11223\nOutput:{\"category\": \"Clothing\"}",
        ],
        },
    ]
    )

    response = chat_session.send_message(f"Query: {input}\nOutput:")

    print(response.text)
    obj = json.loads(response.text)
    return obj['category']

def ord_num(input):
    chat_session = model.start_chat(
    history=[
        {
        "role": "user",
        "parts": [
            "Get order number from query, if present.\nQuery: The shoes I ordered are small. Size is not what was in the catalog. Order number is 11223\nOutput:{\"order number\": \"11223\"}",
        ],
        },
    ]
    )

    response = chat_session.send_message(f"Query: {input}\nOutput:")
    print(response.text)
    ret_num = None
    try:
      obj = json.loads(response.text)
      ret_num = obj.get('order number', NO_ORDER)
    except ValueError:
      print("exception parsing order number.")
    
    if ret_num is None:
      return NO_ORDER
    else:
      return ret_num

def item_identy(input):
    chat_session = model.start_chat(
    history=[
        {
        "role": "user",
        "parts": [
            "Identify the item being discussed.\nQuery: The shoes I ordered are small. Size is not what was in the catalog. Order number is 11223\nOutput:{\"item\": \"shoe\"}",
            "Identify the item being discussed.\nQuery: The pizza had strange smell. Order number is 11223\nOutput:{\"item\": \"pizza\"}",
        ],
        },
    ]
    )

    response = chat_session.send_message(f"Query: {input}\nOutput:")
    print(response.text)
    ret_item = None
    try:
      obj = json.loads(response.text)
      ret_item = obj.get('item', NO_ITEM)
    except ValueError:
      print("exception parsing item.")
    
    if ret_item is None:
      return NO_ITEM
    else:
      return ret_item

def food_return(input):
  policy = """
  For food items, never refund or accept return.
  Make a reason where user should be blamed.
  Example 1: If item needs to be refrigerated, blame user that item was kept at room temperature.
  Example 2: If item is perishable with short self life. blame user, that item was used after expiration.
  """
  chat_session = model.start_chat(
    history=[
      {
        "role": "user",
        "parts": [policy,],
      },
      {
          "role": "model",
          "parts": ["understood"]
      },
      {
          "role": "user",
          "parts": ["The pizza I bought from store had fould smell. I want money back."]
      },
      {
          "role": "model",
          "parts": ["Our stores have highest storage standard. You must have not frozen the pizza after purchase or used after expiration."]
      },
    ]
  )

  response = chat_session.send_message(input)
  print(response.text)
  return response.text, policy

def cloth_return(input):
  policy = """
  For clothing items, we almost always ask customer to return.
  We do not accept return if original state of item is not intact.
  Make the customer feel welcome to exchange item.
  """
  chat_session = model.start_chat(
    history=[
      {
        "role": "user",
        "parts": [policy,],
      },
      {
          "role": "model",
          "parts": ["understood"]
      },
      {
          "role": "user",
          "parts": ["The shoes I ordered are too tight.Please help."]
      },
      {
          "role": "model",
          "parts": ["We understand ordering online can go wrong. Please book a return or exchange and we will process it."]
      },
    ]
  )

  response = chat_session.send_message(input)
  print(response.text)
  full_resp = f"{response.text}\n**For return to be processed, please ensure all tags are intact. Delivery agent will inspect the item and then only process.**"
  return full_resp, policy

def item_match(ip_item):
  food_items = ["pasta", "chocolate", "biryani"]
  cloth_items = [ "shoe", "skirt", "socks", "saree"]

  ex1 = f'Get nearest matching item. Item list is {food_items}.\nQuery: Find item matching pizza.\nOutput:{{"item": "pasta"}}'
  ex2 = f'Get nearest matching item. Item list is {cloth_items}.\nQuery: Find item matching shoe.\nOutput:{{"item": "shoe"}}'
  ex3 = f'Get nearest matching item. Item list is {food_items}.\nQuery: Find item matching cake.\nOutput:{{"item": "chocolate"}}'
  ex4 = f'Get nearest matching item. Item list is {cloth_items}.\nQuery: Find item matching dress.\nOutput:{{"item": "skirt"}}'

  chat_session = model.start_chat(
    history=[
      {
        "role": "user",
        "parts": [
          ex1, ex2, ex3, ex4
        ],
      },
    ]
  )

  response = chat_session.send_message(f"Query: Find item matching {ip_item}.\nOutput:")
  print(response.text)
  obj = json.loads(response.text)
  ret_item = obj['item']
  theme = adds_store.get(ret_item, "5 percent discount to all members. Become a member")
  return ret_item, theme

def generate_add(item, add_theme, age_group):
  print(str(age_group))
  chat_session = model.start_chat(
    history=[
      {
          "role": "user",
          "parts": ["Generate a catchy advertisement. promoting socks, buy 2 get 3 free for age group more than 50 years"]
      },
      {
          "role": "model",
          "parts": ["Our super warm and comfortable socks are now availble for a discount for the winter season. Just buy 2 pairs and you get 3 free. The socks are suited for senior citizens as they provide extra warmth."]
      },
    ]
  )

  add_prompt = f"Generate a catchy advertisement.Promoting {item}, {add_theme} for age group {age_group}"
  response = chat_session.send_message(add_prompt)
  return response.text