from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain
from langchain_openai import OpenAI
from langchain_community.embeddings.sentence_transformer import (
    SentenceTransformerEmbeddings,
)
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.vectorstores import Chroma
import gradio as gr
import os

api = os.environ['google_api']

embedding = SentenceTransformerEmbeddings(
    model_name="all-MiniLM-L6-v2"
)


# embedding = GoogleGenerativeAIEmbeddings()
db = Chroma(persist_directory="vectordb", embedding_function=embedding)

retriever_db = db.as_retriever()

# llm = OpenAI(model="gpt-3.5-turbo-instruct",
#                temperature=0, openai_api_key=opena_api)

# model = ChatOpenAI(model="gpt-3.5-turbo-1106",
#                    temperature=0, openai_api_key=opena_api)


model = ChatGoogleGenerativeAI(
    model="gemini-pro", google_api_key=api, temperature=0.1,  stream=True
)

# Define a function to store conversation history
conversation_history = []


def store_conversation(message):
    conversation_history.append(message)
    if len(conversation_history) > 3:
        # Maintain a history of only 3 conversations
        conversation_history.pop(0)


prompt = ChatPromptTemplate.from_template("""
Sure, here's a revised prompt that addresses inquiries about Entrans beyond just AI products and information technologies:

---

Your role: You are an AI-powered chatbot representing Entrans, a leading provider of cutting-edge AI solutions.

Your Objective: Engage potential clients and showcase Entrans's expertise in developing innovative AI products and information technologies. Additionally, provide information about Entrans's team culture, goals, contact details, and other relevant inquiries.

For inquiries related to building AI products and information technologies, provide compelling responses demonstrating Entrans's capability to design and deliver tailored solutions that meet clients' needs and exceed expectations. Highlight Entrans's track record of successful projects in these areas, emphasizing the quality, reliability, and effectiveness of our solutions.

For inquiries about Entrans's team culture, goals, contact details, and any other aspect of our organization, provide detailed and informative responses that portray Entrans as a dynamic, innovative, and customer-focused company. Showcase our commitment to excellence, collaboration, and continuous improvement in all aspects of our operations.

Your ultimate goal is to leave potential clients impressed and eager to collaborate with Entrans on their next AI project, as well as to provide comprehensive information about our organization to address any inquiries they may have.

Context: {context}
Question: {input}
Conversation History: {conversation_history}
""")


def respond(message, history):
    store_conversation(message)  # Store the current message in history
    result = retrieval_chain.invoke(
        {"input": message, "conversation_history": conversation_history})

    partial_message = ""
    for chunk in result['answer']:
        partial_message += chunk
        yield partial_message


chain = create_stuff_documents_chain(llm=model, prompt=prompt)

retrieval_chain = create_retrieval_chain(
    retriever_db,
    chain
)


gr.ChatInterface(respond,
                 title="Entrans Bot",
                 description="Ask me anything the Information you want!",
                 theme="soft",
                 examples=["What AI services do you provide?",
                           "What would be the cost to build an AI application?"],
                 cache_examples=True
                 ).launch(share=True)