import gradio as gr
import sqlite_utils
import llm
import pandas as pd
from openai import OpenAI
import os

if not os.getenv("OPENAI_API_KEY"):
    raise ValueError("OPENAI_API_KEY must be set")

def retrieval(question: str):
    db = sqlite_utils.Database("metropolen.db")
    embedding_model = llm.get_embedding_model("ada-002")
    collection = llm.Collection("sections", db, model=embedding_model)

    data = [(entry.id, entry.content, entry.score) for entry in collection.similar(question, number=3)]
    df = pd.DataFrame(data, columns=['section_heading', 'content', 'score'])
    documents = (df.content).str.cat(sep='\n')

    return documents

def make_prompt(documents, question):
    prompt = f"""Given the context below, answer the question at the end of the text.
    ---
    {documents}
    ---
    {question}
    """
    return prompt

def generation(prompt, api_key=None, model="gpt-4o"):
    api_key = api_key if api_key else os.getenv('OPENAI_API_KEY')
    client = OpenAI(api_key=api_key)
    c = client.chat.completions.create(model=model,
                                       messages=[{"role": "user", "content": prompt}])
    return c.choices[0].message.content

def rag(frage: str) -> str:
    documents = retrieval(question=frage)
    prompt = make_prompt(documents=documents, question=frage)
    antwort = generation(prompt=prompt, api_key=os.getenv('OPENAI_API_KEY'), model="gpt-4o")
    return antwort

demo = gr.Interface(fn=rag, inputs="text", outputs="text")
demo.launch()