import gradio as gr
import os
import openai
import pinecone

openai.api_key = os.environ["OPENAI-API-KEY"]

pinecone.init(
    api_key=os.environ["PINECONE-API-KEY"],
    environment="asia-southeast1-gcp-free",
)

limit = 5000
# 3750

embed_model = "text-embedding-ada-002"

index_name = 'consulting-docs'
index = pinecone.Index(index_name)

# retrieve relevant answers
def retrieve(query):
    res = openai.Embedding.create(
        input=[query],
        engine=embed_model,
    )

    # retrieve from Pinecone
    xq = res['data'][0]['embedding']

    # get relevant contexts
    res = index.query(xq, top_k=3, include_metadata=True)
    contexts = [
        x['metadata']['text'] for x in res['matches']
    ]

    # build our prompt with the retrieved contexts included
    prompt_start = (
        "Answer the question based on the context below.\n\n"+
        "Context:\n"
    )
    prompt_end = (
        f"\n\nQuestion: {query}\nAnswer:"
    )

    # append contexts until hitting limit
    for i in range(1, len(contexts)):
        if len("\n\n---\n\n".join(contexts[:i])) >= limit:
            prompt = (
                prompt_start +
                "\n\n---\n\n".join(contexts[:i-1]) +
                prompt_end
            )
            break
        elif i == len(contexts)-1:
            prompt = (
                prompt_start +
                "\n\n---\n\n".join(contexts) +
                prompt_end
            )
    return prompt

# then we complete the context-infused query
def complete(prompt):
    # query text-davinci-003
    res = openai.Completion.create(
        engine='text-davinci-003',
        prompt=prompt,
        temperature=0,
        max_tokens=500,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=None
    )
    return res['choices'][0]['text'].strip()

def greet(query):
    # first we retrieve relevant items from Pinecone
    query_with_contexts = retrieve(query)
    # return only the main answer
    result = complete(query_with_contexts)
    return result

iface = gr.Interface(fn=greet, inputs="text", outputs="text")
iface.launch()