import os
import gradio as gr

# Langchain
from langchain.chains import RetrievalQA

# HuggingFace
from langchain_community.embeddings import HuggingFaceEmbeddings

# GeminiPro
from langchain_google_genai import ChatGoogleGenerativeAI

# Pinecone vector database
from pinecone import Pinecone, ServerlessSpec
from langchain_pinecone import PineconeVectorStore


setid = "global"

embeddings = HuggingFaceEmbeddings(model_name=os.getenv("EMBEDDINGS_MODEL"))

# model = ChatOpenAI(temperature=0.0)
model = ChatGoogleGenerativeAI(
    model="gemini-pro", temperature=0.1, convert_system_message_to_human=True
)

pc = Pinecone(
        api_key=os.getenv("PINECONE_API_KEY")
    )
index = pc.Index(setid)
vectorstore = PineconeVectorStore(index, embeddings, "text")

def greet(query):
    retrieval_chain = RetrievalQA.from_llm(
        llm=model,
        retriever=vectorstore.as_retriever(kwargs={"k": 5}),
        return_source_documents=True,
    )
    output = retrieval_chain.invoke(query)
    return output


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