gemini_api.py

import google.generativeai as genai
import json
import os

genai.configure(api_key="AIzaSyAT-eqTQUIpoVpmM-PL5hR5el1n6YqZRV4")
# 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."

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)
    obj = json.loads(response.text)
    ret_num = obj.get('order number', NO_ORDER)
    if ret_num is None:
      return NO_ORDER
    else:
      return ret_num