langgraph/examples/rag/langgraph_agentic_rag.ipynb

In [ ]:

! pip install langchain_community tiktoken langchain-openai langchainhub chromadb langchain langgraph

LangGraph Retrieval Agent

We can implement Retrieval Agents in LangGraph.

Retriever

In [14]:

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain_openai import 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=100, chunk_overlap=50
)
doc_splits = text_splitter.split_documents(docs_list)

# Add to vectorDB
vectorstore = Chroma.from_documents(
    documents=doc_splits,
    collection_name="rag-chroma",
    embedding=OpenAIEmbeddings(),
)
retriever = vectorstore.as_retriever()

In [17]:

from langchain.tools.retriever import create_retriever_tool

tool = create_retriever_tool(
    retriever,
    "retrieve_blog_posts",
    "Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.",
)

tools = [tool]

from langgraph.prebuilt import ToolExecutor

tool_executor = ToolExecutor(tools)

Agent state

We will defined a graph.

A state object that it passes around to each node.

Our state will be a list of messages.

Each node in our graph will append to it.

In [18]:

import operator
from typing import Annotated, Sequence, TypedDict

from langchain_core.messages import BaseMessage


class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]

Nodes and Edges

Each node will -

1/ Either be a function or a runnable.

2/ Modify the state.

The edges choose which node to call next.

We can lay out an agentic RAG graph like this:

In [19]:

import json
import operator
from typing import Annotated, Sequence, TypedDict

from langchain.output_parsers import PydanticOutputParser
from langchain.prompts import PromptTemplate
from langchain.tools.render import format_tool_to_openai_function
from langchain_core.messages import BaseMessage, FunctionMessage
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import ToolInvocation

### Edges


def should_retrieve(state):
    """
    Decides whether the agent should retrieve more information or end the process.

    This function checks the last message in the state for a function call. If a function call is
    present, the process continues to retrieve information. Otherwise, it ends the process.

    Args:
        state (messages): The current state of the agent, including all messages.

    Returns:
        str: A decision to either "continue" the retrieval process or "end" it.
    """
    print("---DECIDE TO RETRIEVE---")
    messages = state["messages"]
    last_message = messages[-1]
    # If there is no function call, then we finish
    if "function_call" not in last_message.additional_kwargs:
        print("---DECISION: DO NOT RETRIEVE / DONE---")
        return "end"
    # Otherwise there is a function call, so we continue
    else:
        print("---DECISION: RETRIEVE---")
        return "continue"


def check_relevance(state):
    """
    Determines whether the Agent should continue based on the relevance of retrieved documents.

    This function checks if the last message in the conversation is of type FunctionMessage, indicating
    that document retrieval has been performed. It then evaluates the relevance of these documents to the user's
    initial question using a predefined model and output parser. If the documents are relevant, the conversation
    is considered complete. Otherwise, the retrieval process is continued.

    Args:
        state messages: The current state of the conversation, including all messages.

    Returns:
        str: A directive to either "end" the conversation if relevant documents are found, or "continue" the retrieval process.
    """

    print("---CHECK RELEVANCE---")

    # Output
    class FunctionOutput(BaseModel):
        binary_score: str = Field(description="Relevance score 'yes' or 'no'")

    # Create an instance of the PydanticOutputParser
    parser = PydanticOutputParser(pydantic_object=FunctionOutput)

    # Get the format instructions from the output parser
    format_instructions = parser.get_format_instructions()

    # Create a prompt template with format instructions and the query
    prompt = PromptTemplate(
        template="""You are a grader assessing relevance of retrieved docs to a user question. \n 
        Here are the retrieved docs:
        \n ------- \n
        {context} 
        \n ------- \n
        Here is the user question: {question}
        If the docs contain keyword(s) in the user question, then score them as relevant. \n
        Give a binary score 'yes' or 'no' score to indicate whether the docs are relevant to the question. \n 
        Output format instructions: \n {format_instructions}""",
        input_variables=["question"],
        partial_variables={"format_instructions": format_instructions},
    )

    model = ChatOpenAI(temperature=0, model="gpt-4-0125-preview")

    chain = prompt | model | parser

    messages = state["messages"]
    last_message = messages[-1]
    score = chain.invoke(
        {"question": messages[0].content, "context": last_message.content}
    )

    # If relevant
    if score.binary_score == "yes":
        print("---DECISION: DOCS RELEVANT---")
        return "yes"

    else:
        print("---DECISION: DOCS NOT RELEVANT---")
        print(score.binary_score)
        return "no"


### Nodes


# Define the function that calls the model
def call_model(state):
    """
    Invokes the agent model to generate a response based on the current state.

    This function calls the agent model to generate a response to the current conversation state.
    The response is added to the state's messages.

    Args:
        state (messages): The current state of the agent, including all messages.

    Returns:
        dict: The updated state with the new message added to the list of messages.
    """
    print("---CALL AGENT---")
    messages = state["messages"]
    model = ChatOpenAI(temperature=0, streaming=True, model="gpt-4-0125-preview")
    functions = [format_tool_to_openai_function(t) for t in tools]
    model = model.bind_functions(functions)
    response = model.invoke(messages)
    # We return a list, because this will get added to the existing list
    return {"messages": [response]}


# Define the function to execute tools
def retrieve(state):
    """
    Executes a tool based on the last message's function call.

    This function is responsible for executing a tool invocation based on the function call
    specified in the last message. The result from the tool execution is added to the conversation
    state as a new message.

    Args:
        state (messages): The current state of the agent, including all messages.

    Returns:
        dict: The updated state with the new function message added to the list of messages.
    """
    print("---EXECUTE RETRIEVAL---")
    messages = state["messages"]
    # Based on the continue condition
    # we know the last message involves a function call
    last_message = messages[-1]
    # We construct an ToolInvocation from the function_call
    action = ToolInvocation(
        tool=last_message.additional_kwargs["function_call"]["name"],
        tool_input=json.loads(
            last_message.additional_kwargs["function_call"]["arguments"]
        ),
    )
    # We call the tool_executor and get back a response
    response = tool_executor.invoke(action)
    # print(type(response))
    # We use the response to create a FunctionMessage
    function_message = FunctionMessage(content=str(response), name=action.tool)

    # We return a list, because this will get added to the existing list
    return {"messages": [function_message]}

# Rewrite query
def rewrite(state):
    
    """
    Transform the query to produce a better question.
    
    Args:
        state (messages): The current state of the agent, including all messages.
    
    Returns:
        dict: The updated state with the new function message added to the list of messages.
    """
    
    print("---TRANSFORM QUERY---")
    # we know the first message involves a user question
    question = messages[0]

    msg = HumanMessage(
        content=f""" \n 
    Look at the input and try to reason about the underlying semantic intent / meaning. \n 
    Here is the initial question:
    \n ------- \n
    {question} 
    \n ------- \n
    Formulate an improved question: """,
    )

    # Grader
    model = ChatOpenAI(temperature=0, model="gpt-4-0125-preview", streaming=True)
    response = model.invoke(msg)
    return {"messages": [response]}

Graph

• Start with an agent, call_model
• Agent make a decision to call a function
• If so, then action to call tool (retriever)
• Then call agent with the tool output added to messages (state)

In [28]:

from langgraph.graph import END, StateGraph

# Define a new graph
workflow = StateGraph(AgentState)

# Define the nodes we will cycle between
workflow.add_node("agent", call_model)  # agent
workflow.add_node("retrieve", retrieve)  # retrieval
workflow.add_node("rewrite", rewrite)  # retrieval

In [29]:

# Call agent node to decide to retrieve or not
workflow.set_entry_point("agent")

# Decide whether to retrieve
workflow.add_conditional_edges(
    "agent",
    # Assess agent decision
    should_retrieve,
    {
        # Call tool node
        "continue": "retrieve",
        "end": END,
    },
)

# Edges taken after the `action` node is called.
workflow.add_conditional_edges(
    "retrieve",
    # Assess agent decision
    check_relevance,
    {
        # Call agent node
        "yes": "agent",
        "no": "rewrite",  
    },
)
workflow.add_edge("agent", END)
workflow.add_edge("rewrite", "agent")

# Compile
app = workflow.compile()

In [30]:

inputs = {
    "messages": [
        HumanMessage(
            content="What does Lilian Weng say about the types of agent memory?"
        )
    ]
}

app.invoke(inputs)

---CALL AGENT---
---DECIDE TO RETRIEVE---
---DECISION: RETRIEVE---
---EXECUTE RETRIEVAL---
---CHECK RELEVANCE---
---DECISION: DOCS RELEVANT---
---CALL AGENT---
---DECIDE TO RETRIEVE---
---DECISION: DO NOT RETRIEVE / DONE---

---------------------------------------------------------------------------
InvalidUpdateError                        Traceback (most recent call last)
File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:736, in _apply_writes(checkpoint, channels, pending_writes, config, for_step)
    735 try:
--> 736     channels[chan].update(vals)
    737 except InvalidUpdateError as e:

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/channels/last_value.py:47, in LastValue.update(self, values)
     46 if len(values) != 1:
---> 47     raise InvalidUpdateError("LastValue can only receive one value per step.")
     49 self.value = values[-1]

InvalidUpdateError: LastValue can only receive one value per step.

The above exception was the direct cause of the following exception:

InvalidUpdateError                        Traceback (most recent call last)
Cell In[30], line 9
      1 inputs = {
      2     "messages": [
      3         HumanMessage(
   (...)
      6     ]
      7 }
----> 9 app.invoke(inputs)

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:569, in Pregel.invoke(self, input, config, output_keys, input_keys, **kwargs)
    559 def invoke(
    560     self,
    561     input: Union[dict[str, Any], Any],
   (...)
    566     **kwargs: Any,
    567 ) -> Union[dict[str, Any], Any]:
    568     latest: Union[dict[str, Any], Any] = None
--> 569     for chunk in self.stream(
    570         input,
    571         config,
    572         output_keys=output_keys if output_keys is not None else self.output,
    573         input_keys=input_keys,
    574         **kwargs,
    575     ):
    576         latest = chunk
    577     return latest

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:605, in Pregel.transform(self, input, config, output_keys, input_keys, **kwargs)
    596 def transform(
    597     self,
    598     input: Iterator[Union[dict[str, Any], Any]],
   (...)
    603     **kwargs: Any,
    604 ) -> Iterator[Union[dict[str, Any], Any]]:
--> 605     for chunk in self._transform_stream_with_config(
    606         input,
    607         self._transform,
    608         config,
    609         output_keys=output_keys,
    610         input_keys=input_keys,
    611         **kwargs,
    612     ):
    613         yield chunk

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langchain_core/runnables/base.py:1497, in Runnable._transform_stream_with_config(self, input, transformer, config, run_type, **kwargs)
   1495 try:
   1496     while True:
-> 1497         chunk: Output = context.run(next, iterator)  # type: ignore
   1498         yield chunk
   1499         if final_output_supported:

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:350, in Pregel._transform(self, input, run_manager, config, input_keys, output_keys)
    347 _interrupt_or_proceed(done, inflight, step)
    349 # apply writes to channels
--> 350 _apply_writes(
    351     checkpoint, channels, pending_writes, config, step + 1
    352 )
    354 if self.debug:
    355     print_checkpoint(step, channels)

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:738, in _apply_writes(checkpoint, channels, pending_writes, config, for_step)
    736     channels[chan].update(vals)
    737 except InvalidUpdateError as e:
--> 738     raise InvalidUpdateError(
    739         f"Invalid update for channel {chan}: {e}"
    740     ) from e
    741 checkpoint["channel_versions"][chan] += 1
    742 updated_channels.add(chan)

InvalidUpdateError: Invalid update for channel __end__: LastValue can only receive one value per step.

In [22]:

import pprint

from langchain_core.messages import HumanMessage

inputs = {
    "messages": [
        HumanMessage(
            content="What does Lilian Weng say about the types of agent memory?"
        )
    ]
}
for output in app.stream(inputs):
    for key, value in output.items():
        pprint.pprint(f"Output from node '{key}':")
        pprint.pprint("---")
        pprint.pprint(value, indent=2, width=80, depth=None)
    pprint.pprint("\n---\n")

---CALL AGENT---
"Output from node 'agent':"
'---'
{ 'messages': [ AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{"query":"types of agent memory"}', 'name': 'retrieve_blog_posts'}})]}
'\n---\n'
---DECIDE TO RETRIEVE---
---DECISION: RETRIEVE---
"Output from node '__end__':"
'---'
{ 'messages': [ HumanMessage(content='What does Lilian Weng say about the types of agent memory?'),
                AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{"query":"types of agent memory"}', 'name': 'retrieve_blog_posts'}})]}
'\n---\n'
---EXECUTE RETRIEVAL---
"Output from node 'retrieve':"
'---'
{ 'messages': [ FunctionMessage(content='Table of Contents\n\n\n\nAgent System Overview\n\nComponent One: Planning\n\nTask Decomposition\n\nSelf-Reflection\n\n\nComponent Two: Memory\n\nTypes of Memory\n\nMaximum Inner Product Search (MIPS)\n\n\nComponent Three: Tool Use\n\nCase Studies\n\nScientific Discovery Agent\n\nGenerative Agents Simulation\n\nProof-of-Concept Examples\n\n\nChallenges\n\nCitation\n\nReferences\n\nTable of Contents\n\n\n\nAgent System Overview\n\nComponent One: Planning\n\nTask Decomposition\n\nSelf-Reflection\n\n\nComponent Two: Memory\n\nTypes of Memory\n\nMaximum Inner Product Search (MIPS)\n\n\nComponent Three: Tool Use\n\nCase Studies\n\nScientific Discovery Agent\n\nGenerative Agents Simulation\n\nProof-of-Concept Examples\n\n\nChallenges\n\nCitation\n\nReferences\n\nPlanning\n\nSubgoal and decomposition: The agent breaks down large tasks into smaller, manageable subgoals, enabling efficient handling of complex tasks.\nReflection and refinement: The agent can do self-criticism and self-reflection over past actions, learn from mistakes and refine them for future steps, thereby improving the quality of final results.\n\n\nMemory\n\nPlanning\n\nSubgoal and decomposition: The agent breaks down large tasks into smaller, manageable subgoals, enabling efficient handling of complex tasks.\nReflection and refinement: The agent can do self-criticism and self-reflection over past actions, learn from mistakes and refine them for future steps, thereby improving the quality of final results.\n\n\nMemory', name='retrieve_blog_posts')]}
'\n---\n'
---CHECK RELEVANCE---
---DECISION: DOCS RELEVANT---
---CALL AGENT---
"Output from node 'agent':"
'---'
{ 'messages': [ AIMessage(content='Lilian Weng discusses the concept of memory within agent systems, highlighting its importance but does not provide specific details on the types of agent memory in the provided excerpt. The discussion on memory is part of a broader overview of agent systems, which also includes planning and tool use. In the context of planning, agents are capable of breaking down large tasks into smaller, manageable subgoals (task decomposition) and engaging in self-reflection and refinement based on past actions to improve future outcomes.\n\nWhile the excerpt mentions a section titled "Types of Memory," specific details or descriptions of these types are not provided in the provided content. Additionally, there\'s a mention of Maximum Inner Product Search (MIPS) in the context of memory, suggesting it might be a technique or tool related to how agents access or utilize their memory, but again, specific details are not given.\n\nFor a more detailed understanding of the types of agent memory Lilian Weng discusses, it would be necessary to access the full content of her blog post or publication.')]}
'\n---\n'
---DECIDE TO RETRIEVE---
---DECISION: DO NOT RETRIEVE / DONE---

---------------------------------------------------------------------------
InvalidUpdateError                        Traceback (most recent call last)
File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:736, in _apply_writes(checkpoint, channels, pending_writes, config, for_step)
    735 try:
--> 736     channels[chan].update(vals)
    737 except InvalidUpdateError as e:

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/channels/last_value.py:47, in LastValue.update(self, values)
     46 if len(values) != 1:
---> 47     raise InvalidUpdateError("LastValue can only receive one value per step.")
     49 self.value = values[-1]

InvalidUpdateError: LastValue can only receive one value per step.

The above exception was the direct cause of the following exception:

InvalidUpdateError                        Traceback (most recent call last)
Cell In[22], line 12
      3 from langchain_core.messages import HumanMessage
      5 inputs = {
      6     "messages": [
      7         HumanMessage(
   (...)
     10     ]
     11 }
---> 12 for output in app.stream(inputs):
     13     for key, value in output.items():
     14         pprint.pprint(f"Output from node '{key}':")

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:605, in Pregel.transform(self, input, config, output_keys, input_keys, **kwargs)
    596 def transform(
    597     self,
    598     input: Iterator[Union[dict[str, Any], Any]],
   (...)
    603     **kwargs: Any,
    604 ) -> Iterator[Union[dict[str, Any], Any]]:
--> 605     for chunk in self._transform_stream_with_config(
    606         input,
    607         self._transform,
    608         config,
    609         output_keys=output_keys,
    610         input_keys=input_keys,
    611         **kwargs,
    612     ):
    613         yield chunk

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langchain_core/runnables/base.py:1497, in Runnable._transform_stream_with_config(self, input, transformer, config, run_type, **kwargs)
   1495 try:
   1496     while True:
-> 1497         chunk: Output = context.run(next, iterator)  # type: ignore
   1498         yield chunk
   1499         if final_output_supported:

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:350, in Pregel._transform(self, input, run_manager, config, input_keys, output_keys)
    347 _interrupt_or_proceed(done, inflight, step)
    349 # apply writes to channels
--> 350 _apply_writes(
    351     checkpoint, channels, pending_writes, config, step + 1
    352 )
    354 if self.debug:
    355     print_checkpoint(step, channels)

File ~/miniforge3/envs/llama2/lib/python3.9/site-packages/langgraph/pregel/__init__.py:738, in _apply_writes(checkpoint, channels, pending_writes, config, for_step)
    736     channels[chan].update(vals)
    737 except InvalidUpdateError as e:
--> 738     raise InvalidUpdateError(
    739         f"Invalid update for channel {chan}: {e}"
    740     ) from e
    741 checkpoint["channel_versions"][chan] += 1
    742 updated_channels.add(chan)

InvalidUpdateError: Invalid update for channel __end__: LastValue can only receive one value per step.

In [ ]: