# Import Langchain modules
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS 
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI


# Import Environment Modules
import os
from dotenv import load_dotenv

# Import API Modules
from fastapi import FastAPI
from fastapi.responses import HTMLResponse, JSONResponse
import uvicorn

# Import Other Modules
import json
import logging
import warnings
warnings.filterwarnings("ignore")

# Load configuration
with open('config.json', 'r') as f:
    config = json.load(f)

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def environment_setup() -> None:
    """
    Load environment variables and set OpenAI API key.
    """
    load_dotenv()
    os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

def load_documents(document_path: str) -> list:
    """
    Load the pdf file and split it into pages.
    """
    try:
        loader = PyPDFLoader(document_path)
        pages = loader.load_and_split()
        return pages
    except Exception as e:
        logging.error(f"Error loading documents from {document_path}: {e}")
        return []

def split_documents(pages: list) -> list:
    """
    Split the pages into chunks.
    """
    try:
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=200,
            chunk_overlap=0,
            length_function=len,
            is_separator_regex=True,
        )
        docs = text_splitter.split_documents(pages)
        return docs
    except Exception as e:
        logging.error(f"Error splitting documents: {e}")
        return []

def process_documents() -> list:
    """
    Process all documents in the specified path.
    """
    document_paths = [os.path.join(config['DOCUMENTS_PATH'], f) for f in os.listdir(config['DOCUMENTS_PATH']) if f.endswith(".pdf")]

    all_docs = []
    for document_path in document_paths:
        pages = load_documents(document_path)
        docs = split_documents(pages)
        all_docs.extend(docs)

    return all_docs

def embeddings(docs: list) -> FAISS:
    """
    Load the embeddings and store them in a vector store.
    """
    try:
        embeddings = OpenAIEmbeddings()
        db = FAISS.from_documents(docs, embeddings)
        return db
    except Exception as e:
        logging.error(f"Error creating embeddings: {e}")
        return None

def initialize_model() -> OpenAI:
    """
    Initialize the model.
    """
    llm = OpenAI()
    return llm

def LLM_chain(llm: OpenAI, db: FAISS) -> RetrievalQA:
    """
    Create a retrieval chain with the LLM and vector store.
    """
    chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=db.as_retriever(search_kwargs={"k": 5}))
    return chain

def initialize_all() -> tuple:
    """
    Initialize all components.
    """
    environment_setup()
    docs = process_documents()
    db = embeddings(docs)
    llm = initialize_model()
    llm_chain = LLM_chain(llm, db)
    return llm_chain, db

def process_message(chain: RetrievalQA, user_message: str, db: FAISS) -> str:
    """
    Process the user's message and return the bot's response.
    """
    try:
        query = user_message
        docs = db.similarity_search(query)
        result = chain.run(input_documents=docs, query=query)
        return result
    except Exception as e:
        logging.error(f"Error generating response: {e}", exc_info=True)
        return "Sorry, I couldn't understand your message."


def setup_fastapi(llm_chain: RetrievalQA, db: FAISS) -> FastAPI:
    """
    Setup FastAPI with routes.
    """
    app = FastAPI()

    @app.get("/", response_class=HTMLResponse)
    def read_root() -> HTMLResponse:
        """
        Serve the chatbot HTML page.
        """
        try:
            with open('templates/chatbot.html', 'r') as f:
                html_content = f.read()
            return HTMLResponse(content=html_content, status_code=200)
        except Exception as e:
            logging.error(f"Error reading HTML file: {e}", exc_info=True)
            return HTMLResponse(content="Sorry, something went wrong.", status_code=500)

    @app.get("/chatbot/{user_message}")
    def get_bot_response(user_message: str) -> JSONResponse:
        """
        Process the user's message and return the bot's response.
        """
        try:
            bot_response = process_message(llm_chain, user_message, db)
            return JSONResponse(content={"answer": bot_response})
        except Exception as e:
            logging.error(f"Error processing message: {e}", exc_info=True)
            return JSONResponse(content={"answer": "Sorry, something went wrong."})

    return app

if __name__ == "__main__":
    try:
        llm_chain, db = initialize_all()
        fastapi_app = setup_fastapi(llm_chain, db)
        uvicorn.run(fastapi_app, host="0.0.0.0", port=8000)
    except Exception as e:
        logging.error(f"Error during initialization: {e}", exc_info=True)
langchain.text_splitter.RecursiveCharacterTextSplitter langchain.llms.OpenAI langchain.vectorstores.FAISS.from_documents langchain.document_loaders.PyPDFLoader langchain.embeddings.OpenAIEmbeddings
[((573, 606), 'warnings.filterwarnings', 'warnings.filterwarnings', (['"""ignore"""'], {}), "('ignore')\n", (596, 606), False, 'import warnings\n'), ((712, 808), 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.INFO', 'format': '"""%(asctime)s - %(levelname)s - %(message)s"""'}), "(level=logging.INFO, format=\n    '%(asctime)s - %(levelname)s - %(message)s')\n", (731, 808), False, 'import logging\n'), ((678, 690), 'json.load', 'json.load', (['f'], {}), '(f)\n', (687, 690), False, 'import json\n'), ((913, 926), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (924, 926), False, 'from dotenv import load_dotenv\n'), ((962, 989), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (971, 989), False, 'import os\n'), ((2654, 2662), 'langchain.llms.OpenAI', 'OpenAI', ([], {}), '()\n', (2660, 2662), False, 'from langchain.llms import OpenAI\n'), ((3800, 3809), 'fastapi.FastAPI', 'FastAPI', ([], {}), '()\n', (3807, 3809), False, 'from fastapi import FastAPI\n'), ((1128, 1154), 'langchain.document_loaders.PyPDFLoader', 'PyPDFLoader', (['document_path'], {}), '(document_path)\n', (1139, 1154), False, 'from langchain.document_loaders import PyPDFLoader\n'), ((1462, 1575), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(200)', 'chunk_overlap': '(0)', 'length_function': 'len', 'is_separator_regex': '(True)'}), '(chunk_size=200, chunk_overlap=0,\n    length_function=len, is_separator_regex=True)\n', (1492, 1575), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((1926, 1967), 'os.path.join', 'os.path.join', (["config['DOCUMENTS_PATH']", 'f'], {}), "(config['DOCUMENTS_PATH'], f)\n", (1938, 1967), False, 'import os\n'), ((2374, 2392), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (2390, 2392), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((2406, 2444), 'langchain.vectorstores.FAISS.from_documents', 'FAISS.from_documents', (['docs', 'embeddings'], {}), '(docs, embeddings)\n', (2426, 2444), False, 'from langchain.vectorstores import FAISS\n'), ((5019, 5070), 'uvicorn.run', 'uvicorn.run', (['fastapi_app'], {'host': '"""0.0.0.0"""', 'port': '(8000)'}), "(fastapi_app, host='0.0.0.0', port=8000)\n", (5030, 5070), False, 'import uvicorn\n'), ((1251, 1318), 'logging.error', 'logging.error', (['f"""Error loading documents from {document_path}: {e}"""'], {}), "(f'Error loading documents from {document_path}: {e}')\n", (1264, 1318), False, 'import logging\n'), ((1738, 1786), 'logging.error', 'logging.error', (['f"""Error splitting documents: {e}"""'], {}), "(f'Error splitting documents: {e}')\n", (1751, 1786), False, 'import logging\n'), ((1977, 2013), 'os.listdir', 'os.listdir', (["config['DOCUMENTS_PATH']"], {}), "(config['DOCUMENTS_PATH'])\n", (1987, 2013), False, 'import os\n'), ((2498, 2546), 'logging.error', 'logging.error', (['f"""Error creating embeddings: {e}"""'], {}), "(f'Error creating embeddings: {e}')\n", (2511, 2546), False, 'import logging\n'), ((3552, 3615), 'logging.error', 'logging.error', (['f"""Error generating response: {e}"""'], {'exc_info': '(True)'}), "(f'Error generating response: {e}', exc_info=True)\n", (3565, 3615), False, 'import logging\n'), ((4087, 4138), 'fastapi.responses.HTMLResponse', 'HTMLResponse', ([], {'content': 'html_content', 'status_code': '(200)'}), '(content=html_content, status_code=200)\n', (4099, 4138), False, 'from fastapi.responses import HTMLResponse, JSONResponse\n'), ((4629, 4675), 'fastapi.responses.JSONResponse', 'JSONResponse', ([], {'content': "{'answer': bot_response}"}), "(content={'answer': bot_response})\n", (4641, 4675), False, 'from fastapi.responses import HTMLResponse, JSONResponse\n'), ((5106, 5171), 'logging.error', 'logging.error', (['f"""Error during initialization: {e}"""'], {'exc_info': '(True)'}), "(f'Error during initialization: {e}', exc_info=True)\n", (5119, 5171), False, 'import logging\n'), ((4182, 4243), 'logging.error', 'logging.error', (['f"""Error reading HTML file: {e}"""'], {'exc_info': '(True)'}), "(f'Error reading HTML file: {e}', exc_info=True)\n", (4195, 4243), False, 'import logging\n'), ((4263, 4332), 'fastapi.responses.HTMLResponse', 'HTMLResponse', ([], {'content': '"""Sorry, something went wrong."""', 'status_code': '(500)'}), "(content='Sorry, something went wrong.', status_code=500)\n", (4275, 4332), False, 'from fastapi.responses import HTMLResponse, JSONResponse\n'), ((4719, 4781), 'logging.error', 'logging.error', (['f"""Error processing message: {e}"""'], {'exc_info': '(True)'}), "(f'Error processing message: {e}', exc_info=True)\n", (4732, 4781), False, 'import logging\n'), ((4801, 4865), 'fastapi.responses.JSONResponse', 'JSONResponse', ([], {'content': "{'answer': 'Sorry, something went wrong.'}"}), "(content={'answer': 'Sorry, something went wrong.'})\n", (4813, 4865), False, 'from fastapi.responses import HTMLResponse, JSONResponse\n')]