# Here is one of the many custom scripts i build. 
# Costs to use it is exactly 0
# Even runs with llama3.1 70B or 405B..and few more...

import streamlit as st
from llm_chatbot import LLMChatBot
from streamlit_option_menu import option_menu
import speech_recognition as sr
import pyttsx3
import os
import getpass
from uuid import uuid4
import faiss
import numpy as np
import requests
import io
import warnings
import torch
import pickle
import asyncio
import json
from git import Repo
from rich import print as rp
from typing import Union, List, Generator, Any, Mapping, Optional, Dict
from requests.sessions import RequestsCookieJar
from dotenv import load_dotenv, find_dotenv
from langchain import hub
from langchain_core.documents import Document
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain
from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter, Language
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma, FAISS
from langchain.vectorstores.base import VectorStore
from langchain.retrievers import MultiQueryRetriever
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.llms import BaseLLM
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain.retrievers.document_compressors import DocumentCompressorPipeline
from langchain_community.document_transformers import EmbeddingsRedundantFilter
from langchain_text_splitters import CharacterTextSplitter
from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain.memory.buffer import ConversationBufferMemory
from langchain.chains import StuffDocumentsChain, LLMChain, ConversationalRetrievalChain
from uber_toolkit_class import UberToolkit
from glob import glob
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.io as pio
from sklearn.decomposition import PCA
from sklearn.preprocessing import MinMaxScaler
from langchain_core.documents import Document
from scipy.stats import gaussian_kde
from huggingface_hub import InferenceClient
from hugchat import hugchat
from hugchat.login import Login
from hugchat.message import Message
from hugchat.types.assistant import Assistant
from hugchat.types.model import Model
from hugchat.types.message import MessageNode, Conversation
from langchain_community.document_loaders import TextLoader
from TTS.api import TTS
import time
from playsound import playsound
from system_prompts import __all__ as prompts
from profiler import VoiceProfileManager, VoiceProfile

# Load environment variables
load_dotenv(find_dotenv())

class ChatbotApp:

  def __init__(self, email, password, default_llm=1):

    self.email = email

    self.password = password

    self.default_llm = default_llm

    self.embeddings = HuggingFaceEmbeddings(

      model_name="all-MiniLM-L6-v2",

      model_kwargs={'device': 'cpu'},

      encode_kwargs={'normalize_embeddings': True}

    )

    self.vectorstore = None


  def create_vectorstore_from_github(self):

    repo_url = "YOUR_REPO_URL"

    local_repo_path = self.clone_github_repo(repo_url)

    loader = DirectoryLoader(path=local_repo_path, glob=f"**/*", show_progress=True, recursive=True)

    loaded_files = loader.load()

    documents = [Document(page_content=file_content) for file_content in loaded_files]

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)

    split_documents = text_splitter.split_documents(documents)

    texts = [doc.page_content for doc in split_documents]

    print(f"Texts for embedding: {texts}") # Debug print

    self.vectorstore = FAISS.from_texts(texts, self.embeddings)


  def create_vectorstore(self, docs):

     

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)

    # Wrap text content in Document objects

    documents = [Document(page_content=doc) for doc in docs]

    # Split documents using the text splitter

    split_documents = text_splitter.split_documents(documents)

    # Convert split documents back to plain text

    texts = [doc.page_content for doc in split_documents]

    vectorstore = FAISS.from_texts(texts, self.setup_embeddings())

    return vectorstore

   

  def setup_session_state(self):

    if 'chat_history' not in st.session_state:

      st.session_state.chat_history = []

    if 'voice_mode' not in st.session_state:

      st.session_state.voice_mode = False

    if 'vectorstore' not in st.session_state:

      st.session_state.vectorstore = None

    if 'retriever' not in st.session_state:

      st.session_state.retriever = None

    if 'compression_retriever' not in st.session_state:

      st.session_state.compression_retriever = None



  def text_to_speech(self, text):

    self.engine.say(text)

    self.engine.runAndWait()



  def speech_to_text(self):

    r = sr.Recognizer()

    with sr.Microphone() as source:

      st.write("Listening...")

      audio = r.listen(source)

      try:

        text = r.recognize_google(audio)

        return text

      except:

        return "Sorry, I didn't catch that."


  def setup_embeddings(self):

    return HuggingFaceEmbeddings(

      model_name="all-MiniLM-L6-v2",

      model_kwargs={'device': 'cpu'},

      encode_kwargs={'normalize_embeddings': True}

    )

 
  def create_vector_store(self, docs):

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)

    # Wrap text content in Document objects

    documents = [Document(page_content=doc) for doc in docs]

    # Split documents using the text splitter

    split_documents = text_splitter.split_documents(documents)

    print(f"Split documents: {split_documents}") # Debug print

    # Convert split documents back to plain text

    texts = [doc.page_content for doc in split_documents]

    print(f"Texts: {texts}") # Debug print

    if not texts:

      print("No valid texts found for embedding. Check your repository content.")

      return


    try:

      self.vectorstore = FAISS.from_texts(texts, self.embeddings)

      print("Vector store created successfully")

    except Exception as e:

      print(f"Error creating vector store: {str(e)}")


  def setup_retriever(self, k=5, similarity_threshold=0.76):

    self.retriever = st.session_state.vectorstore.as_retriever(k=k)

    splitter = CharacterTextSplitter(chunk_size=300, chunk_overlap=0, separator=". ")

    redundant_filter = EmbeddingsRedundantFilter(embeddings=self.setup_embeddings())

    relevant_filter = EmbeddingsFilter(embeddings=self.setup_embeddings(), similarity_threshold=similarity_threshold)

    pipeline_compressor = DocumentCompressorPipeline(

      transformers=[splitter, redundant_filter, relevant_filter]

    )

    st.session_state.compression_retriever = ContextualCompressionRetriever(base_compressor=pipeline_compressor, base_retriever=self.retriever)


  def create_retrieval_chain(self):

    rag_prompt = hub.pull("langchain-ai/retrieval-qa-chat")

    combine_docs_chain = create_stuff_documents_chain(self.llm, rag_prompt)

    self.high_retrieval_chain = create_retrieval_chain(st.session_state.compression_retriever, combine_docs_chain)

    self.low_retrieval_chain = create_retrieval_chain(self.retriever, combine_docs_chain)



  def setup_tts(self, model_name="tts_models/en/ljspeech/fast_pitch"):

    self.tts = TTS(model_name=model_name, progress_bar=False, vocoder_path='vocoder_models/en/ljspeech/univnet')


  def setup_speech_recognition(self):

    self.recognizer = sr.Recognizer()


  def setup_folders(self):

    self.dirs = ["test_input", "vectorstore", "test"]

    for d in self.dirs:

      os.makedirs(d, exist_ok=True)


  def send_message(self, message, web=False):

    message_result = self.llm.chat(message, web_search=web)

    return message_result.wait_until_done()


  def stream_response(self, message, web=False, stream=False):

    responses = []

    for resp in self.llm.query(message, stream=stream, web_search=web):

      responses.append(resp['token'])

    return ' '.join(responses)


  def web_search(self, text):

    result = self.send_message(text, web=True)

    return result


  def retrieve_context(self, query: str):

    context = []

    lowres = self.retriever._get_relevant_documents(query)

    highres = st.session_state.compression_retriever.get_relevant_documents(query)

    context = "\n".join([doc.page_content for doc in lowres + highres])

    return context


  def get_conversation_chain(self):

    EMAIL = os.getenv("EMAIL")

    PASSWD = os.getenv("PASSWD")

    model = 1

    self.llm = LLMChatBot(EMAIL, PASSWD, default_llm=model)
    self.llm.create_new_conversation(system_prompt=self.llm.default_system_prompt, switch_to=True)

    memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)

    conversation_chain = ConversationalRetrievalChain.from_llm(

      llm=self.llm,

      retriever=st.session_state.vectorstore.as_retriever(),

      memory=memory

    )

    return conversation_chain

  async def handle_user_input(self, user_input):

    response = st.session_state.conversation({'question': user_input})

    st.session_state.chat_history = response['chat_history']



    for i, message in enumerate(st.session_state.chat_history):

      if i % 2 == 0:

        st.write(f"Human: {message.content}")

      else:

        st.write(f"AI: {message.content}")

        if st.session_state.voice_mode:

          self.text_to_speech(message.content)

  def clone_github_repo(self, repo_url, local_path='./repo'):

    if os.path.exists(local_path):

      st.write("Repository already cloned.")

      return local_path

    Repo.clone_from(repo_url, local_path)

    return local_path


  def glob_recursive_multiple_extensions(base_dir, extensions):

    all_files = []

    for ext in extensions:

      pattern = os.path.join(base_dir, '**', f'*.{ext}')

      files = glob(pattern, recursive=True)

      all_files.extend(files)

    return all_files


  def load_documents_from_github(self, repo_url, file_types=['*.py', '*.md', '*.txt', '*.html']):

    local_repo_path = self.clone_github_repo(repo_url)

    globber=f"**/*/{{{','.join(file_types)}}}"

    rp(globber)

    loader = DirectoryLoader(path=local_repo_path, glob=globber, show_progress=True, recursive=True,loader_cls=TextLoader)

    loaded_files = loader.load()

    st.write(f"Nr. files loaded: {len(loaded_files)}")

    print(f"Loaded files: {len(loaded_files)}") # Debug print

    # Convert the loaded files to Document objects

    documents = [Document(page_content=file_content) for file_content in loaded_files]

    print(f"Documents: {documents}") # Debug print

    return documents


  def split_documents(self, documents, chunk_s=512, chunk_o=0):

    split_docs = []

    splitter=None

    for doc in documents:

      ext = os.path.splitext(getattr(doc, 'source', '') or getattr(doc, 'filename', ''))[1].lower()

      if ext == '.py':

        splitter = RecursiveCharacterTextSplitter.from_language(language=Language.PYTHON, chunk_size=chunk_s, chunk_overlap=chunk_o)

      elif ext in ['.md', '.markdown']:

        splitter = RecursiveCharacterTextSplitter.from_language(language=Language.MARKDOWN, chunk_size=chunk_s, chunk_overlap=chunk_o)

      elif ext in ['.html', '.htm']:

        splitter = RecursiveCharacterTextSplitter.from_language(language=Language.HTML, chunk_size=chunk_s, chunk_overlap=chunk_o)

      else:

        splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_s, chunk_overlap=chunk_o)

      split_docs.extend(splitter.split_documents([doc]))

    return split_docs, splitter


  def visualize_vectorstore(self):

    if st.session_state.vectorstore is None:

      st.write("Vectorstore is not initialized.")

      return

    documents = st.session_state.vectorstore.get_all_documents()

    embeddings = [doc.embedding for doc in documents]

    pca = PCA(n_components=3)

    embeddings_3d = pca.fit_transform(embeddings)

    scaler = MinMaxScaler()

    embeddings_3d_normalized = scaler.fit_transform(embeddings_3d)

    colors = embeddings_3d_normalized[:, 0]

    hover_text = [f"Document {i}:<br>{doc.page_content[:100]}..." for i, doc in enumerate(documents)]

    fig = go.Figure(data=[go.Scatter3d(

      x=embeddings_3d_normalized[:, 0],

      y=embeddings_3d_normalized[:, 1],

      z=embeddings_3d_normalized[:, 2],

      mode='markers',

      marker=dict(

        size=5,

        color=colors,

        colorscale='Viridis',

        opacity=0.8

      ),

      text=hover_text,

      hoverinfo='text'

    )])


    fig.update_layout(

      title="Interactive 3D Vectorstore Document Distribution",

      scene=dict(

        xaxis_title="PCA Component 1",

        yaxis_title="PCA Component 2",

        zaxis_title="PCA Component 3"

      ),

      width=800,

      height=600,

    )

    st.plotly_chart(fig)


  def chatbot_page(self):

    st.title("Chatbot")

    # Toggle for voice mode

    st.session_state.voice_mode = st.toggle("Voice Mode")

    # File uploader for context injection

    uploaded_file = st.file_uploader("Choose a file for context injection")

    if uploaded_file is not None:

      documents = [uploaded_file.read().decode()]

      st.session_state.vectorstore = self.create_vector_store(documents)

      st.session_state.conversation = self.get_conversation_chain()

    # GitHub repository URL input

    repo_url = st.text_input("Enter GitHub repository URL")

    if repo_url:

      documents = self.load_documents_from_github(repo_url)

      split_docs, _ = self.split_documents(documents)

      st.session_state.vectorstore = self.create_vector_store(split_docs)

      st.session_state.conversation = self.get_conversation_chain()

    # Chat interface

    user_input = st.text_input("You: ", key="user_input")

    if user_input:

      asyncio.run(self.handle_user_input(user_input))

    if st.session_state.voice_mode:

      if st.button("Speak"):

        user_speech = self.speech_to_text()

        st.text_input("You: ", value=user_speech, key="user_speech_input")

        if user_speech != "Sorry, I didn't catch that.":

          asyncio.run(self.handle_user_input(user_speech))


  def dashboard_page(self):

    st.title("Dashboard")


    if st.session_state.vectorstore is not None:

      st.write("Vectorstore Visualization")

      self.visualize_vectorstore()

    else:

      st.write("Vectorstore is not initialized. Please add documents in the Chatbot page.")


  def main(self):

    st.set_page_config(page_title="Enhanced Multi-page Chatbot App", layout="wide")

    # Sidebar navigation

    with st.sidebar:

      selected = option_menu(

        menu_title="Navigation",

        options=["Chatbot", "Dashboard"],

        icons=["chat", "bar-chart"],

        menu_icon="cast",

        default_index=0,

      )

    if selected == "Chatbot":

      self.chatbot_page()

    elif selected == "Dashboard":

      self.dashboard_page()


if __name__ == "__main__":

  app = ChatbotApp(os.getenv("EMAIL"),os.getenv("PASSWD"))

  app.main()
#https://www.linkedin.com/pulse/multi-type-ragollama31-405b-chatbot-boudewijn-kooy-t5lue/?trackingId=Q5pqCmYoQYGWkbViMWtqLQ%3D%3D