from flask import Flask , render_template , jsonify, request
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv
from src.prompt_template import system_prompt
import os
load_dotenv()
app = Flask(__name__)
PINECONE_API_KEY = os.environ['PINECONE_API_KEY']
GOOGLE_API_KEY = os.environ['GOOGLE_API_KEY']

llm = ChatGoogleGenerativeAI(
    model="gemini-1.5-pro",
    temperature=0.4,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key=GOOGLE_API_KEY
)
embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001",google_api_key=GOOGLE_API_KEY)

doc_search = PineconeVectorStore.from_existing_index(
    index_name='customer-support',
    embedding=embeddings
)
retriever = doc_search.as_retriever(searh_type = 'similarity', search_kwards={'k':3})
prompt = ChatPromptTemplate(
    [
        ("system",system_prompt),
        ('human',"{input}")
    ]
)

question_answer_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)

@app.route("/")
def index():
    return render_template("index.html")

@app.route("/get",methods = ['GET','POST'])
def chat():
    text = request.form['text']
    try:
        response = rag_chain.invoke({'input':text})
        return str(response['answer'])

    except Exception as e:
        print(e)
        return "Some error occurred !!"