%%capture --no-stderr
! pip install -U langchain_community tiktoken langchain-openai langchain-cohere langchainhub chromadb langchain langgraph  tavily-python

import getpass
import os
 
def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")
 
 
_set_env("OPENAI_API_KEY")
_set_env("COHERE_API_KEY")
_set_env("TAVILY_API_KEY")

### Build Index
 
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
 
### from langchain_cohere import CohereEmbeddings
 
# Set embeddings
embd = OpenAIEmbeddings()
 
# Docs to index
urls = [
    "https://lilianweng.github.io/posts/2023-06-23-agent/",
    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
    "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
]
 
# Load
docs = [WebBaseLoader(url).load() for url in urls]
docs_list = [item for sublist in docs for item in sublist]
 
# Split
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=500, chunk_overlap=0
)
doc_splits = text_splitter.split_documents(docs_list)
 
# Add to vectorstore
vectorstore = Chroma.from_documents(
    documents=doc_splits,
    collection_name="rag-chroma",
    embedding=embd,
)
retriever = vectorstore.as_retriever()

embd = OpenAIEmbeddings()

urls = [
    "https://lilianweng.github.io/posts/2023-06-23-agent/",
    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
    "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
]

docs = [WebBaseLoader(url).load() for url in urls]
docs_list = [item for sublist in docs for item in sublist]

text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=500, chunk_overlap=0
)
doc_splits = text_splitter.split_documents(docs_list)

vectorstore = Chroma.from_documents(
    documents=doc_splits,
    collection_name="rag-chroma",
    embedding=embd,
)

retriever = vectorstore.as_retriever()

### Router
from typing import Literal
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
 
 
# Data model
class RouteQuery(BaseModel):
    """Route a user query to the most relevant datasource."""
 
    datasource: Literal["vectorstore", "web_search"] = Field(
        ...,
        description="Given a user question choose to route it to web search or a vectorstore.",
    )
 
 
# LLM with function call
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
structured_llm_router = llm.with_structured_output(RouteQuery)
 
# Prompt
system = """You are an expert at routing a user question to a vectorstore or web search.
The vectorstore contains documents related to agents, prompt engineering, and adversarial attacks.
Use the vectorstore for questions on these topics. Otherwise, use web-search."""
route_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "{question}"),
    ]
)
 
question_router = route_prompt | structured_llm_router
 
print(question_router.invoke({"question": "Who will the Bears draft first in the NFL draft?"}))
print(question_router.invoke({"question": "What are the types of agent memory?"}))

class RouteQuery(BaseModel):
    """Route a user query to the most relevant datasource."""
 
    datasource: Literal["vectorstore", "web_search"] = Field(
        ...,
        description="Given a user question choose to route it to web search or a vectorstore.",
    )

llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
structured_llm_router = llm.with_structured_output(RouteQuery)

system = """You are an expert at routing a user question to a vectorstore or web search.
The vectorstore contains documents related to agents, prompt engineering, and adversarial attacks.
Use the vectorstore for questions on these topics. Otherwise, use web-search."""
route_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "{question}"),
    ]
)

question_router = route_prompt | structured_llm_router

print(question_router.invoke({"question": "Who will the Bears draft first in the NFL draft?"}))
print(question_router.invoke({"question": "What are the types of agent memory?"}))

### Retrieval Grader
 
# Data model
class GradeDocuments(BaseModel):
    """Binary score for relevance check on retrieved documents."""
 
    binary_score: str = Field(
        description="Documents are relevant to the question, 'yes' or 'no'"
    )
 
# LLM with function call
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
structured_llm_grader = llm.with_structured_output(GradeDocuments)
 
# Prompt
system = """You are a grader assessing relevance of a retrieved document to a user question. \n 
    If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
    It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""
grade_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Retrieved document: \n\n {document} \n\n User question: {question}"),
    ]
)
 
retrieval_grader = grade_prompt | structured_llm_grader
question = "agent memory"
docs = retriever.invoke(question)
doc_txt = docs[1].page_content
print(retrieval_grader.invoke({"question": question, "document": doc_txt}))

### Generate
 
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
 
# Prompt
prompt = hub.pull("rlm/rag-prompt")
 
# LLM
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
 
# Post-processing
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)
 
# Chain
rag_chain = prompt | llm | StrOutputParser()
 
# Run
generation = rag_chain.invoke({"context": docs, "question": question})
print(generation)

### Hallucination Grader
 
# Data model
class GradeHallucinations(BaseModel):
    """Binary score for hallucination present in generation answer."""
 
    binary_score: str = Field(
        description="Answer is grounded in the facts, 'yes' or 'no'"
    )
 
# LLM with function call
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
structured_llm_grader = llm.with_structured_output(GradeHallucinations)
 
# Prompt
system = """You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \n 
     Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts."""
hallucination_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Set of facts: \n\n {documents} \n\n LLM generation: {generation}"),
    ]
)
 
hallucination_grader = hallucination_prompt | structured_llm_grader
hallucination_grader.invoke({"documents": docs, "generation": generation})

### Answer Grader
 
# Data model
class GradeAnswer(BaseModel):
    """Binary score to assess answer addresses question."""
 
    binary_score: str = Field(
        description="Answer addresses the question, 'yes' or 'no'"
    )
 
 
# LLM with function call
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
structured_llm_grader = llm.with_structured_output(GradeAnswer)
 
# Prompt
system = """You are a grader assessing whether an answer addresses / resolves a question \n 
     Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question."""
answer_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "User question: \n\n {question} \n\n LLM generation: {generation}"),
    ]
)
 
answer_grader = answer_prompt | structured_llm_grader
answer_grader.invoke({"question": question, "generation": generation})

### Question Re-writer
 
# LLM
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
 
# Prompt
system = """You a question re-writer that converts an input question to a better version that is optimized \n 
     for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning."""
re_write_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        (
            "human",
            "Here is the initial question: \n\n {question} \n Formulate an improved question.",
        ),
    ]
)
 
question_rewriter = re_write_prompt | llm | StrOutputParser()
question_rewriter.invoke({"question": question})

### Search
 
from langchain_community.tools.tavily_search import TavilySearchResults
web_search_tool = TavilySearchResults(k=3)

from langchain.schema import Document
 
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, "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 generation
    generation = 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 doc
    filtered_docs = []
    for d in documents:
        score = retrieval_grader.invoke(
            {"question": question, "document": d.page_content}
        )
        grade = score.binary_score
        if grade == "yes":
            print("---GRADE: DOCUMENT RELEVANT---")
            filtered_docs.append(d)
        else:
            print("---GRADE: DOCUMENT NOT RELEVANT---")
            continue
    return {"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 question
    better_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 search
    docs = 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"]
    source = question_router.invoke({"question": question})
    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---")
    state["question"]
    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, TRANSFORM QUERY---"
        )
        return "transform_query"
    else:
        # We have relevant documents, so generate answer
        print("---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.binary_score
 
    # Check hallucination
    if grade == "yes":
        print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
        # Check question-answering
        print("---GRADE GENERATION vs QUESTION---")
        score = answer_grader.invoke({"question": question, "generation": generation})
        grade = score.binary_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"

from typing import List
from typing_extensions import TypedDict
 
class GraphState(TypedDict):
    """
    Represents the state of our graph.
 
    Attributes:
        question: question
        generation: LLM generation
        documents: list of documents
    """
 
    question: str
    generation: str
    documents: List[str]

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"]
 source = question_router.invoke({"question": question})
 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---")
 state["question"]
 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, TRANSFORM QUERY---"
     )
     return "transform_query"
 else:
     # We have relevant documents, so generate answer
     print("---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.binary_score
 
 # Check hallucination
 if grade == "yes":
     print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
     # Check question-answering
     print("---GRADE GENERATION vs QUESTION---")
     score = answer_grader.invoke({"question": question, "generation": generation})
     grade = score.binary_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"

from langgraph.graph import END, StateGraph, START
 
workflow = StateGraph(GraphState)
 
# Define the nodes
workflow.add_node("web_search", web_search)  # web search
workflow.add_node("retrieve", retrieve)  # retrieve
workflow.add_node("grade_documents", grade_documents)  # grade documents
workflow.add_node("generate", generate)  # generatae
workflow.add_node("transform_query", transform_query)  # transform_query
 
# Build graph
workflow.add_conditional_edges(
    START,
    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",
    },
)
 
# Compile
app = workflow.compile()

workflow = StateGraph(GraphState)

workflow.add_node("web_search", web_search)  # web search
workflow.add_node("retrieve", retrieve)  # retrieve
workflow.add_node("grade_documents", grade_documents)  # grade documents
workflow.add_node("generate", generate)  # generatae
workflow.add_node("transform_query", transform_query)  # transform_query

workflow.add_conditional_edges(
    START,
    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",
    },
)

from pprint import pprint
 
# Run
inputs = {
    "question": "ベアーズが2024年のNFLドラフトで最初に指名すると予想される選手は誰ですか？"
}
for output in app.stream(inputs):
    for key, value in output.items():
        # Node
        pprint(f"Node '{key}':")
        # Optional: print full state at each node
        # pprint.pprint(value["keys"], indent=2, width=80, depth=None)
    pprint("\n---\n")
 
# Final generation
pprint(value["generation"])

# Run
inputs = {"question": "エージェントの記憶の種類にはどのようなものがありますか？"}
for output in app.stream(inputs):
    for key, value in output.items():
        # Node
        pprint(f"Node '{key}':")
        # Optional: print full state at each node
        # pprint.pprint(value["keys"], indent=2, width=80, depth=None)
    pprint("\n---\n")
 
# Final generation
pprint(value["generation"])