import gradio as gr
from openai import OpenAI
import tiktoken

# Set your OpenAI API key
openai = OpenAI(api_key='sk-proj-VbcCZjqEbB0ah7tKm37wT3BlbkFJbLfFYNDdbR8hiKXhTMIB')

def count_tokens(messages):
    encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
    num_tokens = 0
    for message in messages:
        num_tokens += len(encoding.encode(message["content"]))
    return num_tokens

def chatgpt(prompt):
    messages = [
        {
            "role": "system",
            "content": (
                "You are an AI assistant for a soldier. "
                "Do not share personal information. "
                "Provide survival tips, medical responses, and identify threats based on camera input. "
                "Keep responses short to save tokens."
            )
        },
        {"role": "user", "content": prompt},
    ]
    
    token_usage = count_tokens(messages)
    print(f"Token usage: {token_usage} tokens")

    response = openai.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=messages,
        max_tokens=50  # Limit the number of tokens in the response
    )
    return response.choices[0].message.content

iface = gr.Interface(fn=chatgpt, inputs="text", outputs="text", title="ShadowWatch001")
iface.launch()