import pandas as pd
df = pd.read_csv('./anime.csv')

context_data = []
for i in range(min(len(df), 100)):  # Loop over rows
    context = ""
    for j in range(7):  # Loop over the first 8 columns
        context += df.columns[j]  # Add column name
        context += ": "
        context += str(df.iloc[i][j])  # Convert value to string
        context += " "
    context_data.append(context)

import os

# Get the secret key from the environment
groq_key = os.environ.get('Animepedia')

## LLM used for RAG
from langchain_groq import ChatGroq

llm = ChatGroq(model="llama-3.1-70b-versatile",api_key=groq_key)

## Embedding model
from langchain_huggingface import HuggingFaceEmbeddings
embed_model = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")

# create vector store
from langchain_chroma import Chroma

vectorstore = Chroma(
    collection_name="Anime_dataset_store",
    embedding_function=embed_model,
    persist_directory="./",
)

vectorstore.get().keys()

# add data to vector nstore
vectorstore.add_texts(context_data)

retriever = vectorstore.as_retriever()

from langchain_core.prompts import PromptTemplate

# Modified template for anime dataset
template = ("""You are an anime expert.
    Use the provided context to answer the question.
    If you don't know the answer, say so. Explain your answer in detail.
    Do not discuss the context in your response; just provide the answer directly.

    Context: {context}

    Question: {question}

    Answer:""")

# Create the prompt
rag_prompt = PromptTemplate.from_template(template)

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | rag_prompt
    | llm
    | StrOutputParser()
)

import gradio as gr

# Function to handle chat input and generate responses using rag_chain
def animepedia_memory_stream(message, history):
    partial_text = ""
    for new_text in rag_chain.stream(message):  # Assuming rag_chain is configured for Animepedia
        partial_text += new_text
        yield partial_text

# Examples of user queries for Animepedia
examples = [
    "What is the highest-rated action anime?",
    "Can you recommend an anime with less than 12 episodes?",
    "Tell me about a family-friendly anime.",
]

# Description and title for the Animepedia chatbot
description = "Real-time Anime Companion to Answer Questions and Provide Recommendations About Your Favorite Shows."
title = "Animepedia: Your Ultimate Anime Guide"

# Creating the Gradio Chat Interface
demo = gr.ChatInterface(
    fn=animepedia_memory_stream,
    type="messages",
    title=title,
    description=description,
    fill_height=True,
    examples=examples,
    theme="glass",
)

# Launching the chatbot interface
if __name__ == "__main__":
    demo.launch()