import os
from embedchain import Pipeline as App
import gradio as gr

os.environ["GOOGLE_API_KEY"] = "AIzaSyBbruzn10nez-0a-_60TA9R9h6qumLD1Es"

app = App.from_config(config={
    "llm": {
        "provider": "google",
        "config": {
            "model": "gemini-pro",
            "temperature": 0.5,
            "max_tokens": 1000,
            "top_p": 1,
            "stream": False,
            "template": """
      Use the following pieces of context to answer the query at the end.
      If you don't know the answer, just say that you don't know, don't try to make up an answer.

      $context

      Query: $query

      Helpful Answer:
    system_prompt: |
      Act as William Shakespeare. Answer the following questions in the style of William Shakespeare.
"""
        },
    },
    "embedder": {
        "provider": "google",
        "config": {
            "model": "models/embedding-001",
        },
    },
})

app.add('http://www.droit-afrique.com/uploads/Gabon-Code-2019-penal.pdf', data_type='pdf_file')


def query(message, history):
    return app.chat(message)


demo = gr.ChatInterface(query)

demo.launch()