# Install required packages
%pip install --upgrade --quiet langchain langchain-community langchainhub chromadb bs4
%pip install --quiet langchain-openai
%pip install --quiet gradio_client==0.02.10 gradio==3.38.0

import getpass
import os
import warnings
warnings.filterwarnings("ignore")

# Set up OpenAI API key
if not os.environ.get("OPENAI_API_KEY"):
    os.environ['OPENAI_API_KEY'] = getpass.getpass("Enter your OpenAI API key: ")

# Import required libraries
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
import bs4
from langchain import hub
from langchain.chains import create_retrieval_chain, create_history_aware_retriever
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.messages import AIMessage, HumanMessage

# Initialize the language model
llm = ChatOpenAI(model="gpt-3.5-turbo-0125")

# Step 1: Load, Chunk and Index the contents to create a retriever
def create_retriever_from_url(url):
    """Create a retriever from a web URL"""
    # Configure BeautifulSoup parser
    bs4_parser = bs4.SoupStrainer(
        class_=("posts-content", "post-title", "post-header")
    )
    
    # Load documents
    loader = WebBaseLoader(
        web_paths=[url],
        bs_kwargs={"parse_only": bs4_parser}
    )
    
    docs = loader.load()
    
    # Split documents into chunks
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )
    splits = text_splitter.split_documents(docs)
    
    # Create vector store and retriever
    vectorstore = Chroma.from_documents(
        documents=splits,
        embedding=OpenAIEmbeddings()
    )
    retriever = vectorstore.as_retriever()
    
    return retriever

# Step 2: Create the RAG chain with history awareness
def create_rag_system(retriever):
    """Create a complete RAG system with chat history"""
    
    # System prompt for the final answer generation
    system_prompt = (
        "You are an assistant for question-answering tasks. "
        "Use the following pieces of retrieved context to answer "
        "the question. If you don't know the answer, say that you "
        "don't know. Use three sentences maximum and keep the "
        "answer concise.\n\n"
        "{context}"
    )
    
    # Prompt for history-aware query reformulation
    contextualize_q_system_prompt = (
        "Given a chat history and the latest user question "
        "which might reference context in the chat history, "
        "formulate a standalone question that can be understood "
        "without the chat history. Do NOT answer the question, "
        "just reformulate it if needed and otherwise return it as is."
    )
    
    # Create history-aware retriever
    contextualize_q_prompt = ChatPromptTemplate.from_messages([
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])
    
    history_aware_retriever = create_history_aware_retriever(
        llm, retriever, contextualize_q_prompt
    )
    
    # Create question-answering chain
    qa_prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])
    
    question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
    
    # Combine into final RAG chain
    rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
    
    return rag_chain

# Step 3: Initialize the system
def initialize_rag_system(url="https://lilianweng.github.io/posts/2023-06-23-agent/"):
    """Initialize the complete RAG system"""
    print("Creating retriever from URL...")
    retriever = create_retriever_from_url(url)
    print("Creating RAG chain...")
    rag_chain = create_rag_system(retriever)
    print("RAG system initialized successfully!")
    return rag_chain

# Step 4: Chat interface
class RAGChatBot:
    def __init__(self, rag_chain):
        self.rag_chain = rag_chain
        self.chat_history = []
    
    def ask_question(self, question):
        """Ask a question and get a response"""
        response = self.rag_chain.invoke({
            "input": question,
            "chat_history": self.chat_history
        })
        
        # Update chat history
        self.chat_history.extend([
            HumanMessage(content=question),
            AIMessage(content=response["answer"])
        ])
        
        return response

# Example usage
if __name__ == "__main__":
    # Initialize the RAG system
    rag_chain = initialize_rag_system()
    chatbot = RAGChatBot(rag_chain)
    
    # Example conversation
    print("\n=== RAG System Demo ===")
    
    # First question
    question1 = "What is Task Decomposition?"
    print(f"\nQ: {question1}")
    response1 = chatbot.ask_question(question1)
    print(f"A: {response1['answer']}")
    
    # Follow-up question (uses chat history)
    question2 = "What are common ways of doing it?"
    print(f"\nQ: {question2}")
    response2 = chatbot.ask_question(question2)
    print(f"A: {response2['answer']}")
    
    # Another follow-up
    question3 = "Can you give examples of each method?"
    print(f"\nQ: {question3}")
    response3 = chatbot.ask_question(question3)
    print(f"A: {response3['answer']}")

    # Print context used for the last response
    print(f"\nContext used: {len(response3['context'])} documents")