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:
            """,
        },
    },
    "embedder": {
        "provider": "google",
        "config": {
            "model": "models/embedding-001",
        },
    },
})

app.add("https://www.forbes.com/profile/elon-musk")

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


demo = gr.ChatInterface(query)

demo.launch()