testUserDataSaver.py

###     Use: streamlit run <file.py>    ###

import streamlit as st
from Retrieval import Retriver
from llm import llm
from chain import chain

import json
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import OllamaLLM
import pdfplumber
from langchain.schema import Document

import chromadb
import os

chromadb.api.client.SharedSystemClient.clear_system_cache()

#   https://docs.streamlit.io/

# {
# Using VirtualchatbotTest memory
# (venv) python3 -m venv chatbotTest
# (venv) source chatbotTest/bin/activate
# (chatbotTest) (venv)
# }



modelPirate = llm("ollamapirate")
modelNormal = llm("llama3.1:8b")
modelSmall = llm("llama3.2:3b", 0.8) # temperature = 0.8
modelSmall2 = llm("llama3.2:3b") # temperature = 0

modelNormalRouter = llm("llama3.1:8b", 0, "json")
modelSmallRouter = llm("llama3.2:3b", 0, "json")


###
##########  TEMPLATES START  ########
###

userDataSaver_instructions = """
You are personal assistant.
You are an expert at saving data about the user from given chat input and data context file.
Use the following retrieved context data.
Overwrite existing information ONLY if new one is more relevant.
You will not save new data as 'null'.
The data will be saved in JSON format with 3 keys.

'Name':
'Age':
'Preferred Language':

If you don't know, or want to save 'null', do not save.

<context>
{context}
</context>

User's input:
{question}

"""

userDataExtract_instructions = """
You are a personal assistant.
You are an expert at extracting data about the user from given context file.
Use user request to extract relevant data.
Use the following piece of retrieved context data.

<context>
{context}
</context>

User request: 
{question}
"""

RAG_TEMPLATE_MATH = """
You are an expert at question generation using a given file.
Use the following pieces of retrieved context to generate a new question based on user's request.
The generated question MUST be similar to the original AND on the same level.
If you don't know, just say that you don't know.
You MUST NOT return answer in correct Latex format.

<context>
{context}
</context>


User's request:
{question}"""

### Works, slowly, reads all the data each time. ###
router_instructions1 = """
You are an expert at routing a user question to a vectorstore or self answer.
The vectorstore contains documents related to math questions:
<context>
{context}
</context>
Use the vectorstore for questions on these topics. For all else, and especially for current events, generate independent answer.
Return JSON with single key, datasource, that is 'independent' or 'vectorstore' depending on the question:
{question}"""

### Also Works, given what data is stored in vector. ###
router_instructions2 = """
You are an expert at routing a user question to a vectorstore or self answer.
The vectorstore contains documents related to math questions in calculus 1.
Use the vectorstore for questions on these topics. For all else, and especially for current events, and given information generate independent answer.
Return JSON with single key, datasource, that is 'independent' or 'vectorstore' depending on the question:
{question}"""



###
##########  TEMPLATES END  ##########
###
def extract_hebrew_text(pdf_path):
    documents = []
    with pdfplumber.open(pdf_path) as pdf:
        for i, page in enumerate(pdf.pages):
            text = page.extract_text()
            if text:  # Only include pages with text
                documents.append(Document(page_content=text,
                                          metadata={"page": i + 1}))
    return documents


def extract_hebrew_text_toString(pdf_path):
    with pdfplumber.open(pdf_path) as pdf:
        text = ""
        for page in pdf.pages:
            text += page.extract_text()
    return text




def RouterDirection(question,router_chain):
    ##### Use router_chain2 for simple check.
    ##### User router_chain1 for heavy check. - runs through files.
    print("*** Routing QUESTION ****")
    directionOfQJSON = router_chain.chainRunner.invoke(question)
    # print(directionOfQJSON)
    directionOfQ = json.loads(directionOfQJSON)["datasource"]
    print(f"Choosen way is: {directionOfQ}")
    return directionOfQ







def reverse_Heb(sentence):
    words = sentence.split()
    reversed_words = [word[::-1] for word in words] #reverse each word
    reversed_words = reversed_words[::-1] #reverse order of words
    return ' '.join(reversed_words)







path_userData = "/home/tom/תואר/4thYearProject/PythonProjects/userInfo/userData_eng.txt"
path = "/home/tom/תואר/4thYearProject/PythonProjects/calculusJsonDataShort_eng.txt"
if os.path.exists(path_userData):
    noData=0
    # File exists, open in read and write mode ('r+')
    open(path_userData, 'r+')  # 'r+' allows reading and writing
else:
    noData=1
    # File doesn't exist, create and open in write mode ('w')
    with open(path_userData, 'w') as f: # 'w' creates and opens for writing
         f.write("EMPTY")

print("--Start Data")
data= Retriver(path)
dataUser= Retriver(path_userData)
print("--End Data")

router_chain1 = chain(router_instructions1, modelSmallRouter, data)
router_chain2 =chain(router_instructions2,modelSmallRouter)

def test_router(router_chain):
    test_router1 = router_chain.chainRunner.invoke(
        "Who is favored to win the NFC Championship game in the 2024 season?"
    )
    test_router2 = router_chain.chainRunner.invoke(
        "Who Is Jona?"
    )
    test_router3 = router_chain.chainRunner.invoke(
        "I need a math question?"
    )
    print(
        json.loads((test_router1)),
        json.loads((test_router2)),
        json.loads((test_router3)),
    )
#test_router(router_chain2)

## TEST RouterDirection ###
# print(RouterDirection(
#     "Who is favored to win the NFC Championship game in the 2024 season?",
#     router_chain=router_chain2))






#chainRAG ---> can select which model to serach with, best is modelSmall
###         Story PDF-Eng = 1:48 minutes --> llama3.1:8b
###         Story Txt-Eng = 0:54 minutes --> llama3.1:8b
###         Story PDF-Eng = 0:24 minutes --> llama3.2:3b
###         Story Txt-Eng = 0:24 minutes --> llama3.2:3b
userDataSAVER_agent = chain(userDataSaver_instructions, modelSmallRouter, dataUser)
userDataEXTRACTOR_agent = chain(userDataExtract_instructions, modelSmall, dataUser)
chainRAG = chain(RAG_TEMPLATE_MATH, modelSmall, data)










####      TEST RAG CHAIN      ####
# question = "generate a question similar to Question number 5"
# ans = (chainRAG.chainRunner.invoke(question)+"\n")
# print(ans+"\n")





templateChat="""
Answer the question below, use less than 200 words.

Here is the conversation history: {history}

Question: {question}

Answer:
"""
modelChatting = OllamaLLM(model="llama3.1:8b")
## Notice - OllamaLLM is !NOT! ChatOllama (used in the start)
promptChat = ChatPromptTemplate.from_template(templateChat)
chainChat = promptChat | modelChatting
#chainRAG



def handle_conv_steamlit():
    if noData:
        get_user_info()

    if "messages" not in st.session_state:
        st.session_state.messages = []


    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if userInput := st.chat_input("What is up?"):
        st.session_state.messages.append({"role": "user", "content": userInput})
        with st.chat_message("user"):
            st.markdown(userInput)

        with st.chat_message("assistant"):
            if userInput.lower() == "exit":
                st.write("**Exiting the chat. Refresh the page to start again.**")
                return

            # Determine route and invoke appropriate chain
            routerWay = RouterDirection(userInput, router_chain2)
            if routerWay == "math":
                #Need RAG
                #Give a new question from RAG
                result = chainRAG.chainRunner.invoke(userInput)
            # elif routerWay == "personal":
            #     result = userDataEXTRACTOR_agent.chainRunner.invoke(userInput)
            else:
                result = chainChat.invoke({
                    "history": [
                                {"role": m["role"], "content": m["content"]}
                                for m in st.session_state.messages
                                ],
                    "question": userInput
                })
            st.markdown(clean_latex(result))

            # newData = userDataSAVER_agent.chainRunner.invoke(userInput)
            # print(newData+"\n")
            # with open(path_userData, 'w') as file:
            #     file.write(newData)

        st.session_state.messages.append({"role": "assistant", "content": result})




DATA_FILE = path_userData

def save_user_data(name, age,degree, taskComplete, topic1):
    """Saves user name and age to a text file."""
    with open(DATA_FILE, "w") as file:
        file.write(f"{name}\n{age}\n{degree}\n{taskComplete}\n{topic1}")

def load_user_data():
    """Loads user name and age from the text file."""
    if os.path.exists(DATA_FILE):
        with open(DATA_FILE, "r") as file:
            data = file.read().strip().split("\n")
            if (data[0] != "EMPTY"):
                return data
    return None, None, None, None, None
def get_user_info():
    """Collects user's Name and Age before allowing chat interaction."""
    st.title("Welcome! Please enter your details to continue.")

    # Load user data from file
    name, age, degree, taskComplete, topic1 = load_user_data()

    if name and age and degree:
        st.session_state.name = name
        st.session_state.age = age
        st.session_state.degree = degree
        st.session_state.taskComplete = taskComplete
        st.session_state.topic1 = topic1
        st.session_state.info_collected = True
        st.rerun()  # Refresh page to move to chat

    with st.form("user_info_form"):
        name = st.text_input("Enter your Name:")
        age = st.text_input("Enter your Age:")
        degree = st.text_input("Enter your Degree:")
        submitted = st.form_submit_button("Continue")

        if submitted:
            if name.strip() and age.strip().isdigit():
                taskComplete =0
                topic1 =0
                st.session_state.name = name
                st.session_state.age = age
                st.session_state.degree = degree
                st.session_state.taskComplete = taskComplete
                st.session_state.topic1 = topic1
                st.session_state.info_collected = True
                save_user_data(name, age, degree, taskComplete, topic1)
                st.rerun()  # Refresh page to move to chat
            else:
                st.warning("Please enter a valid Name and Age (numeric).")

# Chat function
def handle_chat():
    # st.title(f"Welcome back, {st.session_state.name}! Let's Chat.")

    user_data = {"Name": st.session_state.name,
                 "Age": st.session_state.age,
                 "Degree": st.session_state.degree,
                 "Task Complete": st.session_state.taskComplete,
                 "Topic1 Level": st.session_state.topic1}

    st.sidebar.title("User Progress")
    st.sidebar.metric("Name", user_data["Name"])
    st.sidebar.metric("Age", user_data["Age"])
    st.sidebar.metric("Degree", user_data["Degree"])
    st.sidebar.metric("Task Complete", user_data["Task Complete"])
    st.sidebar.header("Topic 1 Progress")
    progress = st.sidebar.progress(int(user_data["Topic1 Level"]))


    if "messages" not in st.session_state:
        st.session_state.messages = []

    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if userInput := st.chat_input("What is up?"):
        st.session_state.messages.append({"role": "user", "content": userInput})

        with st.chat_message("user"):
            st.markdown(userInput)

        with st.chat_message("assistant"):
            response = f"Hello {st.session_state.name}, you said: {userInput}"  # Replace with LLM response
            st.markdown(response)

        st.session_state.messages.append({"role": "assistant", "content": response})

# Run the app
if "info_collected" not in st.session_state or not st.session_state.info_collected:
    get_user_info()
else:
    handle_chat()


import time

def basicStreamlit():
    prompt = st.text_input("Say something")
    if prompt:
        st.write(f"User prompt: {prompt}")
    user_data = {"Score": 85, "Tasks Completed": 7, "Level": 3}

    # Sidebar content
    st.sidebar.title("User Progress")
    st.sidebar.metric("Score", user_data["Score"])
    st.sidebar.metric("Tasks Completed", user_data["Tasks Completed"])
    st.sidebar.metric("Level", user_data["Level"])

    # Progress bar simulation
    progress = st.sidebar.progress(0)
    for i in range(101):  # Simulating progress
        time.sleep(0.05)
        progress.progress(i)

    # Main content
    st.title("Main Application")
    st.write("This is your main content area.")


# basicStreamlit()


import re
def clean_latex(input_string):
    # Remove [asy] environments and their contents
    cleaned_string = re.sub(r'\[asy\](.*?)\[/asy\]', '', input_string, flags=re.DOTALL)

    # Remove extra newlines and trim whitespace
    #cleaned_string = re.sub(r'\n\s*\n', '\n', cleaned_string).strip()

    return cleaned_string




####     For Chat Testing   ####
#handle_conv_steamlit()