# Required modules
import os
from pinecone import Pinecone
from transformers import AutoModel
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq
from dotenv import load_dotenv
load_dotenv()


# Initialize clients, indexes, models etc.
pc_client = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
pc_index = pc_client.Index(os.getenv("PINECONE_INDEX"))
embedding_model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
groq_llm=ChatGroq(
    groq_api_key=os.getenv("GROQ_API_KEY"),
    model_name="Llama3-8b-8192"
)


#context retrivel
def retrive_context(user_query:str) -> str:
    """Retrives the context for asked query from vector database

    Args:
        user_query (str): Questions asked by user to bot

    Returns:
        context (str): Question's context
    """
    
    context = ""
    try:
        embedded_query = embedding_model.encode(user_query).tolist()
    except Exception as e:
        return 500
    
    try:
        res = pc_index.query(
            vector=embedded_query,
            top_k=5,
            include_values=True,
            include_metadata = True
        )
    except Exception as e:
        return 500
    
    for match in res['matches']:
        context = context + match['metadata']['text'] + " "
        
    print(context)
    return context


# Prompt Engineering for LLM
prompt = ChatPromptTemplate.from_template(
    """
    Hello! As a RAG agent for Biskane, your task is to answer the user's question using the provided context. Please keep your responses brief and straightforward.

    <context>
    {context}
    <context>
    Question: {query}
    """
)


# Response generator
def generate_response(query:str, context:str) -> str:
    """Generates the response for asked question from given context

    Args:
        query (str): Query asked by user to bot
        context (str): Context, retrived from vector database

    Returns:
        answer (str): Generated response
    """
    try: 
        chain = prompt | groq_llm
        llm_response = chain.invoke({
            "context": context,
            "query": query
        })
        return llm_response.content
    except Exception as e:
        return 500