testStreamlit.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

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")
modelNormalRouter = llm("llama3.1:8b", 0, "json")
modelSmallRouter = llm("llama3.2:3b", 0, "json")


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

RAG_TEMPLATE = """
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, just say that you don't know.
Use three sentences maximum and keep the answer concise.

<context>
{context}
</context>

Answer the following question:

{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 personal data:
<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 Toms' personal data, parents, hobbies ect.
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}"""

###
##########  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 test_router(router_chain):
    test_router1 = router_chain.invoke(
        "Who is favored to win the NFC Championship game in the 2024 season?"
    )
    test_router2 = router_chain.invoke(
        "Who Is Jona?"
    )
    test_router3 = router_chain.invoke(
        "What Tom studies?"
    )
    print(
        json.loads((test_router1)),
        json.loads((test_router2)),
        json.loads((test_router3)),
    )
###     Test of router chain    ###
#test_router()
###                             ###


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)
    directionOfQ = json.loads(directionOfQJSON)["datasource"]
    print(f"Choosen way is: {directionOfQ}")
    return directionOfQ

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





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)






pdf_path = "/home/tom/תואר/4thYearProject/PythonProjects/testMathSyb.pdf"


#
path = "/home/tom/ollamaProjects/RAGtextFile"
# path = "/home/tom/תואר/4thYearProject/PythonProjects/syllabusCalculus1TomTranslated"
# path = "/home/tom/תואר/4thYearProject/PythonProjects/testMathSymb"

FileType="txt"
LanguageType="eng"

webPath ="https://bedtimestorieskd.com/small-bear//"

data= Retriver(path,LanguageType,FileType)


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

#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
chainRAG = chain(RAG_TEMPLATE, modelSmall, data)





# question = "What the course about?"
# print(chainRAG.chainRunner.invoke(question)+"\n")

# question = "What is the kids name? and who are his parents?"
# print(chainRAG.chainRunner.invoke(question)+"\n")
# question = "מה השם של הילד?"
# print(reverse_Heb(chainRAG.invoke(question)))

# # question = "Tell me about Toms favorite place and his parents names"
# # print(chainRAG.invoke(question))
# question = "ספר לי על המקום אהוב על תום ומה שם הוריו."
# print(reverse_Heb(chainRAG.invoke(question)))

# question = "What is special about potatoes?"
# print(chainRAG.invoke(question)+"\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():
    history=""
    print("Welcome to ChatRAG(Routing), type exit to quit.")
    while True:
        userInput = input("Input: ")
        if userInput.lower() =="exit":
            break
        if(RouterDirection(userInput,router_chain1)=="vectorstore"):
            #If need RAG so use chainRAG
            result = chainRAG.chainRunner.invoke(userInput)
        else:
            #If RAG not needed use simple chainChat
            result = chainChat.invoke({
                "history": history,
                "question": userInput
            })
        print("AI: ", result)
        history = history + f"\nUser: {userInput}\nAI: {result}"


# handle_conv()

def handle_conv_steamlitTut():
    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
            if RouterDirection(userInput, router_chain2) == "vectorstore":
                result = chainRAG.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(result)
        st.session_state.messages.append({"role": "assistant", "content": result})


handle_conv_steamlitTut()


def basicStreamlit():
    prompt = st.text_input("Say something")
    if prompt:
        st.write(f"User prompt: {prompt}")


# 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


def testLaTexSteamlit():
    st.title("LaTeX Equation Display")
    st.markdown("This is an inline equation: $E = mc^2$.")
    st.latex(r"""
    \begin{align*}
    a^2 + b^2 &= c^2 \\
    x &= \frac{-b \pm \sqrt{b^2 - 4ac}}{2a}
    \end{align*}
    """)

    text_from_trainMATH = "Solving for $x$ and $y,$ we find\n\\[x = \\frac{11t - 1}{3}, \\quad y = \\frac{5t + 5}{3}.\\]From the first equation,\n\\[t = \\frac{3x + 1}{11}.\\]Then\n\\begin{align*}\ny &= \\frac{5t + 5}{3} \\\\\n&= \\frac{5 \\cdot \\frac{3x + 1}{11} + 5}{3} \\\\\n&= \\frac{5}{11} x + \\frac{20}{11},\n\\end{align*}Thus, the slope of the line is $\\boxed{\\frac{5}{11}}.$"
    st.latex(text_from_trainMATH)
    st.markdown(text_from_trainMATH)

    st.markdown("\n\n")
    st.markdown("\n\n")
    st.markdown("This is question number: $57$.")
    Question_num_57 = "What is the volume of the region in three-dimensional space defined by the inequalities $|x|+|y|+|z|\\le1$ and $|x|+|y|+|z-1|\\le1$?"
    st.markdown(clean_latex(Question_num_57))
    st.markdown("This is answer number: $57$.")
    Question_num_57 = "In the octant where $x \\ge 0,$ $y \\ge 0,$ and $z \\ge 0,$ the inequality $|x| + |y| + |z| \\le 1$ becomes\n\\[x + y + z \\le 1.\\]Thus, the region in this octant is the tetrahedron with vertices $(0,0,0),$ $(1,0,0),$ $(0,1,0),$ and $(1,0,0).$  By symmetry, the region defined by $|x| + |y| + |z| \\le 1$ is the octahedron with vertices $(\\pm 1,0,0),$ $(0,\\pm 1,0),$ and $(0,0,\\pm 1).$  Let the base of the upper-half of the octahedron be $ABCD,$ and let $E = (0,0,1).$\n\nSimilarly, the region defined by $|x| + |y| + |z - 1| \\le 1$ is also an octahedron, centered at $(0,0,1).$  Let the base of the lower-half of the octahedron be $A'B'C'D',$ and let $E' = (0,0,0).$\n\n[asy]\nimport three;\n\nsize(250);\ncurrentprojection = perspective(6,3,2);\n\ntriple A, B, C, D, E, Ap, Bp, Cp, Dp, Ep, M, N, P, Q;\n\nA = (1,0,0);\nB = (0,1,0);\nC = (-1,0,0);\nD = (0,-1,0);\nE = (0,0,1);\nAp = (1,0,1);\nBp = (0,1,1);\nCp = (-1,0,1);\nDp = (0,-1,1);\nEp = (0,0,0);\nM = (A + E)/2;\nN = (B + E)/2;\nP = (C + E)/2;\nQ = (D + E)/2;\n\ndraw(D--A--B);\ndraw(D--C--B,dashed);\ndraw(C--E,dashed);\ndraw(A--M);\ndraw(M--E,dashed);\ndraw(B--N);\ndraw(N--E,dashed);\ndraw(D--Q);\ndraw(Q--E,dashed);\ndraw(Ap--Bp--Cp--Dp--cycle);\ndraw(Ap--M);\ndraw(M--Ep,dashed);\ndraw(Bp--N);\ndraw(N--Ep,dashed);\ndraw(Cp--Ep,dashed);\ndraw(Dp--Q);\ndraw(Q--Ep,dashed);\ndraw(Q--M--N);\ndraw(Q--P--N,dashed);\n\nlabel(\"$A$\", A, SW);\nlabel(\"$B$\", B, dir(0));\nlabel(\"$C$\", C, S);\nlabel(\"$D$\", D, W);\nlabel(\"$E$\", E, dir(90));\nlabel(\"$A'$\", Ap, dir(90));\nlabel(\"$B'$\", Bp, dir(0));\nlabel(\"$C'$\", Cp, dir(90));\nlabel(\"$D'$\", Dp, W);\nlabel(\"$E'$\", Ep, S);\nlabel(\"$M$\", M, SW);\nlabel(\"$N$\", N, dir(0));\nlabel(\"$P$\", P, NE);\nlabel(\"$Q$\", Q, W);\n[/asy]\n\nFaces $ABE$ and $A'B'E'$ intersect in line segment $\\overline{MN},$ where $M$ is the midpoint of $\\overline{AE},$ and $N$ is the midpoint of $\\overline{BE}.$  Thus, the intersection of the two octahedra is another octahedra, consisting of the upper-half of pyramid $ABCDE,$ and the lower-half of pyramid $A'B'C'D'E'.$\n\nThe volume of pyramid $ABCDE$ is\n\\[\\frac{1}{3} \\cdot (\\sqrt{2})^2 \\cdot 1 = \\frac{2}{3},\\]so the volume of its upper half is $\\left( \\frac{1}{2} \\right)^3 \\cdot \\frac{2}{3} = \\frac{1}{12}.$  Then the volume of the smaller octahedron is $\\frac{2}{12} = \\boxed{\\frac{1}{6}}.$"
    st.markdown(clean_latex(Question_num_57))



####     For LaTex Testing   ####
#testLaTexSteamlit()
####     For Chat Testing   ####
#handle_conv_steamlitTut()