core.py

Setup
from dotenv import load_dotenv
load_dotenv()

LLM
from langchain_ollama import ChatOllama
local_llm = "llama2:7b"
llm = ChatOllama(model=local_llm, temperature=0, base_url="http://localhost:11434")
llm_json_mode = ChatOllama(model=local_llm, temperature=0, format="json", base_url="http://localhost:11434")

Retriever
from langchain_pinecone import PineconeVectorStore
from langchain_huggingface import HuggingFaceEmbeddings
index_name= "tactical-edge-rag-index"
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = PineconeVectorStore(index_name=index_name, embedding=embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

Router
from langchain_core.prompts import ChatPromptTemplate
system = """You are an expert at routing a user question to a vectorstore or web search.
The vectorstore contains documents related to 2022 Annual report of Mobily and Operation and
Maintenance Manual of Caterpillar.
Use the vectorstore for questions on these topics. For all else, and especially for current events, use web-search.
Return JSON with single key, datasource, that is 'websearch' or 'vectorstore' depending on the question."""
human_msg = """    User question:     {question}"""
route_prompt = ChatPromptTemplate.from_messages(
[
("system", system),
("human", human_msg),
])
question_router = route_prompt | llm_json_mode

Documents Grader
system = """You are a grader assessing relevance of a retrieved document to a user question. If the document contains keyword(s) or semantic meaning related to the question, grade it as relevant.
It does not need to be a stringent test. The goal is to filter out erroneous retrievals.
Return JSON with single key, binary_score, that is 'yes' or 'no' score to indicate whether the document contains at least some information that is relevant to the question."""
human_msg = """    Retrieved documents:     {documents}         User question:     {question}"""
grade_prompt = ChatPromptTemplate.from_messages(
[
("system", system),
("human", human_msg),
])
retrieval_grader = grade_prompt | llm_json_mode

Generate Answer
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
prompt = hub.pull("rlm/rag-prompt") #prompt has context and question parameter

Chain
rag_chain = prompt | llm | StrOutputParser()

Hallucination Grader
system = """You are a teacher grading a quiz. You will be given FACTS and a STUDENT ANSWER. Here is the grade criteria to follow:
(1) Ensure the STUDENT ANSWER is grounded in the FACTS. (2) Ensure the STUDENT ANSWER does not contain "hallucinated" information outside the scope of the FACTS.
Score:
A score of yes means that the student's answer meets all of the criteria. This is the highest (best) score. A score of no means that the student's answer does not meet all of the criteria. This is the lowest possible score you can give.
Explain your reasoning in a step-by-step manner to ensure your reasoning and conclusion are correct. Avoid simply stating the correct answer at the outset.
Return JSON with two two keys, binary_score is 'yes' or 'no' score to indicate whether the STUDENT ANSWER is grounded in the FACTS. And a key, explanation, that contains an explanation of the score."""
human_msg = """    Set of facts:     {documents}         Student answer:     {generation}"""

hallucination_prompt = ChatPromptTemplate.from_messages(
[
("system", system),
("human", human_msg),
])
hallucination_grader = hallucination_prompt | llm_json_mode

Answer Grader
system = """"You are a teacher grading a quiz. You will be given a QUESTION and a STUDENT ANSWER. Here is the grade criteria to follow:
(1) The STUDENT ANSWER helps to answer the QUESTION
Score:
A score of yes means that the student's answer meets all of the criteria. This is the highest (best) score. The student can receive a score of yes if the answer contains extra information that is not explicitly asked for in the question.
A score of no means that the student's answer does not meet all of the criteria. This is the lowest possible score you can give.
Explain your reasoning in a step-by-step manner to ensure your reasoning and conclusion are correct. Avoid simply stating the correct answer at the outset.
Return JSON with two two keys, binary_score is 'yes' or 'no' score to indicate whether the STUDENT ANSWER meets the criteria. And a key, explanation, that contains an explanation of the score."""
human_msg = """    User question:    {question}         LLM generation:    {generation}"""
answer_prompt = ChatPromptTemplate.from_messages(
[
("system", system),
("human", human_msg),
])
answer_grader = answer_prompt | llm_json_mode

Question Re-writer
system = """You a question re-writer that converts an input question to a better version that is optimized for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning.
Return only the rewritten question without any explanation."""
human_msg = """    Here is the initial question:     {question}        Formulate an improved question."""
re_write_prompt = ChatPromptTemplate.from_messages(
[
("system", system),
("human", human_msg)
])
question_rewriter = re_write_prompt | llm | StrOutputParser()

Web Search Tool
from langchain_community.tools.tavily_search import TavilySearchResults
web_search_tool = TavilySearchResults(max_results=3)

Graph State
from typing import List
from typing_extensions import TypedDict
from langchain.schema import Document
import json
class State(TypedDict):
"""
Represents the state of our graph.
Attributes:
question: question
generation: LLM generation
max_retries: Max number of retries for answer generation
loop_step: number of loops for answer generation
documents: list of documents

"""
question: str
generation: str
max_retries: int
loop_step: int
documents: List[Document]

Retriever Node<
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"]
# Retrieval
documents = retriever.invoke(question)

return {"documents": documents}

Generate Node
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"]
loop_step = state.get("loop_step", 0)
# RAG generation
generation = rag_chain.invoke({"context": documents, "question": question})

return {"generation": generation, "loop_step": loop_step + 1}

Documents Grader Node
If any docs are relevant, we can proceed with generating answer
def grade(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 doc
filtered_docs = []
for doc in documents:
score = retrieval_grader.invoke(
{"question": question, "documents": doc}
)
grade = json.loads(score.content)["binary_score"]
if grade.lower() == "yes":
print("---GRADE: DOCUMENT RELEVANT---")
filtered_docs.append(doc)
else:
print("---GRADE: DOCUMENT NOT RELEVANT---")

return {"documents": filtered_docs}

Question Re-writer Node
def rewrite(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("---Rewrite---")
question = state["question"]
# Re-write question
better_question = question_rewriter.invoke({"question": question})

return {"question": better_question}

Web Search Node
def 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 search
docs = web_search_tool.invoke(question)
web_results = "\n".join([doc["content"] for doc in docs])
documents = Document(page_content=web_results)
return {"documents": documents}

Conditional Edge
from typing import Literal
def route_question(state) -> Literal["vectorstore", "web_search"]:
"""
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"]
answer  = question_router.invoke({"question": question})
source = json.loads(answer.content)["datasource"]
if source == "web_search":
print("---ROUTE QUESTION TO WEB SEARCH---")
return "web_search"
elif source == "vectorstore":
print("---ROUTE QUESTION TO RAG---")
return "vectorstore"
def decide_to_generate(state) -> Literal["generate", "rewrite"]:
"""
Determines whether to generate an answer, or rewrite a question.
Args:
state (dict): The current graph state
Returns:
str: Binary decision for next node to call
"""
print("---ASSESS GRADED DOCUMENTS---")
filtered_documents = state["documents"]
if not filtered_documents:
# All documents have been filtered check_relevance
# We will re-generate a new query
print(
"---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, REWRITE QUESTION---"
)
return "rewrite"
else:
# We have relevant documents, so generate answer
print("---DECISION: GENERATE---")
return "generate"
def decide_to_answer(state) -> Literal["useful", "not useful", "not supported", "max retries"]:
"""
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"]
max_retries = state.get("max_retries", 3)
hallucination_score = hallucination_grader.invoke(
{"documents": documents, "generation": generation}
)
hallucination_grade = json.loads(hallucination_score.content)["binary_score"]
# Check hallucination
if hallucination_grade.lower() == "yes":
print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
# Check question-answering
print("---GRADE GENERATION vs QUESTION---")
answer_score = answer_grader.invoke({"question": question, "generation": generation})
answer_grade = json.loads(answer_score.content)["binary_score"]
if answer_grade.lower() == "yes":
print("---DECISION: GENERATION ADDRESSES QUESTION---")
return "useful"
elif state["loop_step"] <= max_retries:
print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
return "not useful"
else:
print("---DECISION: MAX RETRIES REACHED---")
return "max retries"
elif state["loop_step"] <= max_retries:
print("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
return "not supported"
else:
print("---DECISION: MAX RETRIES REACHED---")
return "max retries"

Compile Graph
from IPython.display import Image, display
from langgraph.graph import StateGraph, START, END
workflow = StateGraph(State)

Define the nodes
workflow.add_node("retrieve", retrieve)  # retrieve
workflow.add_node("grade", grade)  # grade
workflow.add_node("generate", generate)  # generatae
workflow.add_node("rewrite", rewrite)  # rewrite
workflow.add_node("search", search)  # web search

Build graph
workflow.add_conditional_edges(
START,
route_question,
{
"web_search": "search",
"vectorstore": "retrieve",
},
)
workflow.add_edge("search", "generate")
workflow.add_edge("retrieve", "grade")
workflow.add_conditional_edges(
"grade",
decide_to_generate,
{
"rewrite": "rewrite",
"generate": "generate",
},
)
workflow.add_edge("rewrite", "retrieve")
workflow.add_conditional_edges(
"generate",
decide_to_answer,
{
"useful": END,
"not useful": "rewrite",
"not supported": "generate",
"max retries": END,
},
)

Compile
graph = workflow.compile()