from langchain.text_splitter import RecursiveCharacterTextSplitterfrom langchain_community.vectorstores import Chromafrom langchain_community.embeddings import GPT4AllEmbeddingsfrom langchain.prompts import PromptTemplatefrom langchain_community.chat_models import ChatOllamafrom langchain_core.output_parsers import JsonOutputParserfrom langchain.prompts import PromptTemplatefrom langchain_community.chat_models import ChatOllamafrom langchain_core.output_parsers import JsonOutputParserfrom langchain import hubfrom langchain_community.chat_models import ChatOllamafrom langchain_core.output_parsers import StrOutputParserfrom typing_extensions import TypedDictfrom typing import Listfrom langchain.schema import Documentfrom langgraph.graph import END, StateGraphfrom langchain_community.tools.tavily_search import TavilySearchResultsfrom langchain.schema import Documentfrom langchain_community.document_loaders import PyPDFLoaderimport streamlit as stimport os

local_llm = "llama3"tavily_api_key = os.environ['TAVILY_API_KEY'] = 'API_KEY'st.title("Multi-PDF ChatBot using LLAMA3 & Adaptive RAG")user_input = st.text_input("Question:", placeholder="Ask about your PDF", key='input')
with st.sidebar:uploaded_files = st.file_uploader("Upload your file", type=['pdf'], accept_multiple_files=True)process = st.button("Process")if process:if not uploaded_files:st.warning("Please upload at least one PDF file.")st.stop()

# Ensure the temp directory existstemp_dir = 'Users/home/temp/'if not os.path.exists(temp_dir):os.makedirs(temp_dir)
# Process each uploaded filefor uploaded_file in uploaded_files:temp_file_path = os.path.join(temp_dir, uploaded_file.name)
# Save the file to diskwith open(temp_file_path, "wb") as file:file.write(uploaded_file.getbuffer())# Use getbuffer() for Streamlit's UploadedFile
# Load the PDF using PyPDFLoadertry:loader = PyPDFLoader(temp_file_path)data = loader.load()# Assuming loader.load() is the correct method callst.write(f"Data loaded for {uploaded_file.name}")except Exception as e:st.error(f"Failed to load {uploaded_file.name}: {str(e)}")

text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=250, chunk_overlap=0)text_chunks = text_splitter.split_documents(data)
# Add to vectorDBvectorstore = Chroma.from_documents(documents=text_chunks,collection_name="rag-chroma",embedding=GPT4AllEmbeddings(),)retriever = vectorstore.as_retriever()llm = ChatOllama(model=local_llm, format="json", temperature=0)

 prompt = PromptTemplate(template="""You are an expert at routing a user question to a vectorstore or web search. \\nUse the vectorstore for questions on LLMagents, prompt engineering, and adversarial attacks. \\nYou do not need to be stringent with the keywords in the question related to these topics. \\nOtherwise, use web-search. Give a binary choice 'web_search' or 'vectorstore' based on the question. \\nReturn the a JSON with a single key 'datasource' and no premable or explaination. \\nQuestion to route: {question}""",input_variables=["question"],)
question_router = prompt | llm | JsonOutputParser()question = "llm agent memory"docs = retriever.get_relevant_documents(question)doc_txt = docs[1].page_contentquestion_router.invoke({"question": question})llm = ChatOllama(model=local_llm, format="json", temperature=0)

prompt = PromptTemplate(template="""You are a grader assessing relevance of a retrieved document to a user question. \\nHere is the retrieved document: \\n\\n {document} \\n\\nHere is the user question: {question} \\nIf the document contains keywords related to the user question, grade it as relevant. \\nIt does not need to be a stringent test. The goal is to filter out erroneous retrievals. \\nGive a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question. \\nProvide the binary score as a JSON with a single key 'score' and no premable or explaination.""",input_variables=["question", "document"],)
retrieval_grader = prompt | llm | JsonOutputParser()question = "agent memory"docs = retriever.get_relevant_documents(question)doc_txt = docs[1].page_contentst.write(retrieval_grader.invoke({"question": question, "document": doc_txt}))

### Generateprompt = hub.pull("rlm/rag-prompt")
# LLMllm = ChatOllama(model=local_llm, temperature=0)
# Post-processingdef format_docs(docs):return "\\n\\n".join(doc.page_content for doc in docs)
# Chainrag_chain = prompt | llm | StrOutputParser()
# Runquestion = "agent memory"generation = rag_chain.invoke({"context": docs, "question": question})print(generation)

### Hallucination Grader# LLMllm = ChatOllama(model=local_llm, format="json", temperature=0)
# Promptprompt = PromptTemplate(template="""You are a grader assessing whether an answer is grounded in / supported by a set of facts. \\nHere are the facts:\\n ------- \\n{documents}\\n ------- \\nHere is the answer: {generation}Give a binary score 'yes' or 'no' score to indicate whether the answer is grounded in / supported by a set of facts. \\nProvide the binary score as a JSON with a single key 'score' and no preamble or explanation.""",input_variables=["generation", "documents"],)
hallucination_grader = prompt | llm | JsonOutputParser()hallucination_grader.invoke({"documents": docs, "generation": generation})

### Answer Grader
# LLMllm = ChatOllama(model=local_llm, format="json", temperature=0)
# Promptprompt = PromptTemplate(template="""You are a grader assessing whether an answer is useful to resolve a question. \\nHere is the answer:\\n ------- \\n{generation}\\n ------- \\nHere is the question: {question}Give a binary score 'yes' or 'no' to indicate whether the answer is useful to resolve a question. \\nProvide the binary score as a JSON with a single key 'score' and no preamble or explanation.""",input_variables=["generation", "question"],)
answer_grader = prompt | llm | JsonOutputParser()answer_grader.invoke({"question": question,"generation": generation})

### Question Re-writer
# LLMllm = ChatOllama(model=local_llm, temperature=0)
# Promptre_write_prompt = PromptTemplate(template="""You a question re-writer that converts an input question to a better version that is optimized \\n for vectorstore retrieval. Look at the initial and formulate an improved question. \\n Here is the initial question: \\n\\n {question}. Improved question with no preamble: \\n """,input_variables=["generation", "question"],)
question_rewriter = re_write_prompt | llm | StrOutputParser()question_rewriter.invoke({"question": question})

web_search_tool = TavilySearchResults(k=3,tavily_api_key=tavily_api_key)

class GraphState(TypedDict):"""Represents the state of our graph.
Attributes:question: questiongeneration: LLM generationdocuments: list of documents"""question : strgeneration : strdocuments : List[str]

def retrieve(state):"""Retrieve documents
Args:state (dict): The current graph state
Returns:state (dict): New key added to state, documents, that contains retrieved documents"""print("---RETRIEVE---")question = state["question"]
# Retrievaldocuments = retriever.get_relevant_documents(question)return {"documents": documents, "question": question}
def generate(state):"""Generate answer
Args:state (dict): The current graph state
Returns:state (dict): New key added to state, generation, that contains LLM generation"""print("---GENERATE---")question = state["question"]documents = state["documents"]
# RAG generationgeneration = rag_chain.invoke({"context": documents, "question": question})return {"documents": documents, "question": question, "generation": generation}
def grade_documents(state):"""Determines whether the retrieved documents are relevant to the question.
Args:state (dict): The current graph state
Returns:state (dict): Updates documents key with only filtered relevant documents"""
print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")question = state["question"]documents = state["documents"]
# Score each docfiltered_docs = []for d in documents:score = retrieval_grader.invoke({"question": question, "document": d.page_content})grade = score['score']if grade == "yes":print("---GRADE: DOCUMENT RELEVANT---")filtered_docs.append(d)else:print("---GRADE: DOCUMENT NOT RELEVANT---")continuereturn {"documents": filtered_docs, "question": question}
def transform_query(state):"""Transform the query to produce a better question.
Args:state (dict): The current graph state
Returns:state (dict): Updates question key with a re-phrased question"""
print("---TRANSFORM QUERY---")question = state["question"]documents = state["documents"]
# Re-write questionbetter_question = question_rewriter.invoke({"question": question})return {"documents": documents, "question": better_question}

def web_search(state):"""Web search based on the re-phrased question.
Args:state (dict): The current graph state
Returns:state (dict): Updates documents key with appended web results"""
print("---WEB SEARCH---")question = state["question"]
# Web searchdocs = web_search_tool.invoke({"query": question})web_results = "\\n".join([d["content"] for d in docs])web_results = Document(page_content=web_results)
return {"documents": web_results, "question": question}
### Edges ###
def route_question(state):"""Route question to web search or RAG.
Args:state (dict): The current graph state
Returns:str: Next node to call"""
print("---ROUTE QUESTION---")question = state["question"]print(question)source = question_router.invoke({"question": question})print(source)print(source['datasource'])if source['datasource'] == 'web_search':print("---ROUTE QUESTION TO WEB SEARCH---")return "web_search"elif source['datasource'] == 'vectorstore':print("---ROUTE QUESTION TO RAG---")return "vectorstore"
def decide_to_generate(state):"""Determines whether to generate an answer, or re-generate a question.
Args:state (dict): The current graph state
Returns:str: Binary decision for next node to call"""
print("---ASSESS GRADED DOCUMENTS---")question = state["question"]filtered_documents = state["documents"]
if not filtered_documents:# All documents have been filtered check_relevance# We will re-generate a new queryprint("---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---")return "transform_query"else:# We have relevant documents, so generate answerprint("---DECISION: GENERATE---")return "generate"
def grade_generation_v_documents_and_question(state):"""Determines whether the generation is grounded in the document and answers question.
Args:state (dict): The current graph state
Returns:str: Decision for next node to call"""
print("---CHECK HALLUCINATIONS---")question = state["question"]documents = state["documents"]generation = state["generation"]
score = hallucination_grader.invoke({"documents": documents, "generation": generation})grade = score['score']
# Check hallucinationif grade == "yes":print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")# Check question-answeringprint("---GRADE GENERATION vs QUESTION---")score = answer_grader.invoke({"question": question,"generation": generation})grade = score['score']if grade == "yes":print("---DECISION: GENERATION ADDRESSES QUESTION---")return "useful"else:print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")return "not useful"else:pprint("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")return "not supported"

workflow = StateGraph(GraphState)
# Define the nodesworkflow.add_node("web_search", web_search) # web searchworkflow.add_node("retrieve", retrieve) # retrieveworkflow.add_node("grade_documents", grade_documents) # grade documentsworkflow.add_node("generate", generate) # generataeworkflow.add_node("transform_query", transform_query) # transform_query
# Build graphworkflow.set_conditional_entry_point(route_question,{"web_search": "web_search","vectorstore": "retrieve",},)workflow.add_edge("web_search", "generate")workflow.add_edge("retrieve", "grade_documents")workflow.add_conditional_edges("grade_documents",decide_to_generate,{"transform_query": "transform_query","generate": "generate",},)workflow.add_edge("transform_query", "retrieve")workflow.add_conditional_edges("generate",grade_generation_v_documents_and_question,{"not supported": "generate","useful": END,"not useful": "transform_query",},)
# Compileapp = workflow.compile()

inputs = {"question": user_input}for output in app.stream(inputs):for key, value in output.items():# Nodest.write(f"Node '{key}':")# Optional: print full state at each node# pprint.pprint(value["keys"], indent=2, width=80, depth=None)print("\\n---\\n")
# Final generationst.write(value["generation"])