from brain import Brain

# Example conversation history
conversation_history = [
    {"role": "system", "content": "LeaveBot"},
    {"role": "user", "content": "Hey"},
    {"role": "assistant", "content": "Hello there! How can I assist you today? 
    If you have any questions related to Aubergine Solutions company's policy or guidelines, feel free to ask!"},
]

# Instantiate the Brain
leave_bot = Brain()

# User interacts with the chatbot
user_query_1 = "How many leave days am I entitled to per year?"
user_query_2 = "Can I take a leave for more than 10 days consecutively?"

# Initial interaction
conversation_history.append({"role": "user", "content": user_query_1})
response_1 = " ".join(response for response in leave_bot.llm_call(chain=conversation_history))
conversation_history.append({"role": "assistant", "content": response_1})

# User's next query
conversation_history.append({"role": "user", "content": user_query_2})
response_2 = " ".join(response for response in leave_bot.llm_call(chain=conversation_history))
conversation_history.append({"role": "assistant", "content": response_2})

The odd numbers in this group add up to an even number: 4, 8, 9, 15, 12, 2, 1.
A: Adding all the odd numbers (9, 15, 1) gives 25. The answer is False.
The odd numbers in this group add up to an even number: 17,  10, 19, 4, 8, 12, 24.
A: Adding all the odd numbers (17, 19) gives 36. The answer is True.
The odd numbers in this group add up to an even number: 16,  11, 14, 4, 8, 13, 24.
A: Adding all the odd numbers (11, 13) gives 24. The answer is True.
The odd numbers in this group add up to an even number: 17,  9, 10, 12, 13, 4, 2.
A: Adding all the odd numbers (17, 9, 13) gives 39. The answer is False.
The odd numbers in this group add up to an even number: 15, 32, 5, 13, 82, 7, 1. 
A:

Adding all the odd numbers (15, 5, 13, 7, 1) gives 41. The answer is False.

Note: 'Let's Think step by step' <---- add this special prompt.
Question: A juggler can juggle 16 balls. Half of the balls are golf balls, and half of the golf balls are blue. How many blue golf balls are there?
Answer:

There are 16 balls in total. Half of the balls are golf balls. That means that there are 8 golf balls. Half of the golf balls are blue. That means that there are 4 blue golf balls.

import os
from decouple import config
from chromadb.utils import embedding_functions
from long-chain.text_splitter import RecursiveCharacterTextSplitter
from pdf import PdfReader

def data_loader():
    # Set the directory path where PDF files are stored
    data_dir = 'data/policy'
    
    # Check if the directory exists
    if not os.path.exists(data_dir) or not os. path.isdir(data_dir):
        print("The 'data' directory does not exist in the current working directory.")
        return
    
    # Initialize an empty list to store data
    all_data = []
    
    # Loop through each file in the directory
    for file in os.listdir(data_dir):
        if file.endswith(".pdf"):
            # Extract the file name without the extension
            file_name = file.split(".")[0]
            
            # Construct the full PDF file path
            pdf_url = os.path.join(data_dir, file)
            
            # Read the PDF file
            pdf_reader = PdfReader(pdf_url)
            
            # Extract text from each page, remove spaces, and join lines
            pdf_texts = [" ".join(page.extract_text().replace(" ", "").split('\n')) for page in pdf_reader.pages]
            
            # Remove empty texts
            pdf_texts = [text for text in pdf_texts if text]
            
            # Initialize a text splitter
            character_splitter = RecursiveCharacterTextSplitter(
                separators=["\n\n", "\n", ". ", " ", ""],
                chunk_size=1000,
                chunk_overlap=0
            )
            
            # Split text into chunks
            character_split_texts = character_splitter.split_text('\n\n'.join(pdf_texts))
            
            # Display the total number of chunks
            print(f"\nTotal chunks: {len(character_split_texts)}")
            
            # Append each chunk to the data list with associated topic
            for chunk in character_split_texts:
                all_data.append({"topic": file_name, "content": chunk})

    # Get the collection name and Chroma client host from environment variables
    collection_name = config('COMPANY_POLICY_DATA')
    chroma_client_host = config('CHROMA_CLIENT_HOST')
    
    # Initialize a ChromaDB HTTP client
    chroma_client = chromadb.HttpClient(host=chroma_client_host, port=8000)
    
    # Try to delete the existing collection (if any)
    try:
        chroma_client.delete_collection(name=collection_name)
    except:
        pass
    
    # Define distance functions for indexing
    distance_functions = ["l2", "ip", "cosine"]
    
    # Initialize a SentenceTransformer embedding function
    sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="multi-qa-mpnet-base-dot-v1")
    
    # Create a new ChromaDB collection with specified metadata
    collection = chroma_client.create_collection(
        name=collection_name, embedding_function=sentence_transformer_ef, metadata={"hnsw:space": distance_functions[0]})

    # Initialize lists to store documents, metadata, and document IDs
    documents = []
    metadata = []
    ids = []
    n = 1
    
    # Iterate through each data chunk and extract content and topic
    for data in all_data:
        documents.append(data['content'])
        metadata.append({"topic": data['topic']})
        ids.append(f"{data['topic']}_{n}")
        n += 1
    
    # Add documents, metadata, and IDs to the ChromaDB collection
    collection.add(
        documents=documents,
        metadatas=metadata,
        ids=ids
    )
    
    # Return True indicating successful data loading
    return True

import chromadb
import os
from decouple import config
from pypdf import PdfReader
from  chromadb.utils  import  embedding_functions
from langchain.text_splitter import RecursiveCharacterTextSplitter

def data_loader():
    data_dir = 'data/policy'
    if not os.path.exists(data_dir) or not os.path.isdir(data_dir):
        print("The 'data' directory does not exist in the current working directory.")
        return
    all_data = []
    for file in os.listdir(data_dir):
        if file.endswith(".pdf"):
            file_name = file.split(".")[0]
            pdf_url = os.path.join(data_dir, file)
            pdf_reader = PdfReader(pdf_url)
            pdf_texts = [" ".join(page.extract_text().replace(" ", "").split('\n')) for page in pdf_reader.pages]
            pdf_texts = [text for text in pdf_texts if text]
            character_splitter = RecursiveCharacterTextSplitter(
                separators=["\n\n", "\n", ". ", " ", ""],
                chunk_size=1000,
                chunk_overlap=0
            )
            character_split_texts = character_splitter.split_text('\n\n'.join(pdf_texts))
            print(f"\nTotal chunks: {len(character_split_texts)}")
            for chunk in character_split_texts:
                all_data.append({"topic": file_name, "content": chunk})

    collection_name = config('COMPANY_POLICY_DATA')
    chroma_client_host = config('CHROMA_CLIENT_HOST')
    chroma_client = chromadb.HttpClient(host=chroma_client_host, port=8000)
    try:
        chroma_client.delete_collection(name=collection_name)
    except:
        pass
    distance_functions = ["l2", "ip", "cosine"]
    sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="multi-qa-mpnet-base-dot-v1")
    collection = chroma_client.create_collection(
        name=collection_name, embedding_function=sentence_transformer_ef, metadata={"hnsw:space": distance_functions[0]})
    documents = []
    metadata = []
    ids = []
    n=1
    for data in all_data:
        documents.append(data['content'])
        metadata.append({"topic": data['topic']})
        ids.append(f"{data['topic']}_{n}")
        n+=1
    collection.add(
        documents=documents,
        metadatas=metadata,
        ids=ids
    )
    return True