testRAG.py

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

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

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_eng.txt"
# path = "/home/tom/תואר/4thYearProject/PythonProjects/syllabusCalculus1TomTranslated"
# path = "/home/tom/תואר/4thYearProject/PythonProjects/testMathSymb"

# FileType="txt"
# LanguageType="eng"
# data= Retriver(path,LanguageType,FileType)
webPath ="https://bedtimestorieskd.com/small-bear/"

data= Retriver(path)


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()