langgraph/examples/reflexion/reflexion.ipynb

Reflexion

Reflexion by Shinn, et. al., is an architecture designed to learn through verbal feedback by means of an episodic memory unit.

It's got 3 main components:

1. Actor with Self-reflection
2. Evaluator
3. Memory

Below, we implement in LangGraph.

0. Prerequisites

In [ ]:

# %pip install -U langchain langgraph langchain_openai

In [1]:

import getpass
import os


def _set_if_undefined(var: str) -> None:
    if os.environ.get(var):
        return
    os.environ[var] = getpass.getpass(var)


# Optional: Configure tracing to visualize and debug the agent
_set_if_undefined("LANGCHAIN_API_KEY")
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "Reflexion"

We will use tavily search as a tool. You can get an API key here or replace with a different tool of your choosing.

In [2]:

# %pip install tavily-python

In [3]:

_set_if_undefined("OPENAI_API_KEY")
_set_if_undefined("TAVILY_API_KEY")

Actor

In [4]:

from langchain.tools.tavily_search import TavilySearchResults
from langchain.utilities.tavily_search import TavilySearchAPIWrapper

search = TavilySearchAPIWrapper()
tavily_tool = TavilySearchResults(api_wrapper=search, max_results=5)

In [8]:

import datetime
from typing import List

from langchain.chains import create_structured_output_runnable
from langchain.output_parsers.openai_tools import JsonOutputToolsParser
from langchain.prompts.chat import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI

actor_prompt_template = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """You are expert researcher.
Current time: {time}

1. {first_instruction}
2. Reflect and critique your answer. Be severe to maximize improvement.
3. Recommend a search query to research information and improve your answer.""",
        ),
        MessagesPlaceholder(variable_name="messages"),
        ("system", "Answer the user's question above using the required format."),
    ]
).partial(
    time=lambda: datetime.datetime.now().isoformat(),
)


class Reflection(BaseModel):
    missing: str = Field(description="Critique of what is missing.")
    superfluous: str = Field(description="Critique of what is superfluous")


class AnswerQuestion(BaseModel):
    """Answer the question."""

    answer: str = Field(description="~250 word detailed answer to the question.")
    reflection: Reflection = Field(description="Your reflection on the initial answer.")
    search_queries: List[str] = Field(
        description="1-3 search queries for researching improvements to address the critique of your current answer."
    )


def nest(state: List[BaseMessage]):
    return {"messages": state}


llm = ChatOpenAI(model="gpt-4-turbo-preview")
initial_answer_chain = (
    nest
    | actor_prompt_template.partial(
        first_instruction="Provide a detailed ~250 word answer."
    )
    | llm.bind_tools(tools=[AnswerQuestion], tool_choice="AnswerQuestion")
)

In [9]:

initial = initial_answer_chain.invoke(
    [HumanMessage(content="Why the US should NOT adopt rank-choice voting")]
)

In [10]:

parser = JsonOutputToolsParser(return_id=True)
parsed = parser.invoke(initial)
parsed[0]

Out [10]:

{'type': 'AnswerQuestion',
 'args': {'answer': 'The debate around whether the United States should adopt ranked-choice voting (RCV) is complex, involving various considerations spanning democratic principles, electoral outcomes, and practical implications. Opponents of RCV argue against its adoption in the US for several reasons. Firstly, RCV can be more complicated for voters to understand and participate in compared to the traditional plurality voting system. This complexity might discourage voter participation or lead to a higher rate of ballot errors, ultimately undermining the democratic process. Secondly, while RCV is designed to produce a consensus candidate that reflects a broader preference among the electorate, critics argue that it can lead to moderate, less polarizing candidates being elected, potentially stifling vibrant political debate and reducing the diversity of political thought represented in government. Thirdly, the implementation of RCV would require significant changes to the existing electoral infrastructure, including updates to voting machines and the training of election officials, which could be costly and time-consuming. There are also concerns about the potential for increased confusion during the transition period, which could affect election integrity. Lastly, some detractors contend that RCV does not necessarily solve the problem of strategic voting, as voters may still attempt to game the system by ranking candidates in a way that does not reflect their true preferences, merely to avoid the election of their least preferred candidate.',
  'reflection': {'missing': 'Specific examples or studies that support the arguments against RCV, such as instances where RCV led to voter confusion or did not eliminate strategic voting.',
   'superfluous': 'The answer might have spent too much time explaining the general concept of RCV before diving into the specific reasons against its adoption in the US, which could have been briefly summarized to allow more space for detailed criticisms.'},
  'search_queries': ['ranked-choice voting disadvantages studies',
   'RCV implementation costs and challenges',
   'Effects of RCV on voter participation and strategic voting']},
 'id': 'call_dg5IemqOS5EfSJCJY88rQMrm'}

Actor pt 2: Revise

In [11]:

revise_instructions = """Revise your previous answer using the new information.
    - You should use the previous critique to add important information to your answer.
        - You MUST include numerical citations in your revised answer to ensure it can be verified.
        - Add a "References" section to the bottom of your answer (which does not count towards the word limit). In form of:
            - [1] https://example.com
            - [2] https://example.com
    - You should use the previous critique to remove superfluous information from your answer and make SURE it is not more than 250 words.
"""


class ReviseAnswer(AnswerQuestion):
    """Revise your original answer to your question."""

    references: List[str] = Field(
        description="Citations motivating your updated answer."
    )


revision_chain = (
    nest
    | actor_prompt_template.partial(first_instruction=revise_instructions)
    | llm.bind_tools(tools=[ReviseAnswer], tool_choice="ReviseAnswer")
)

In [15]:

parsed[0]["args"]["search_queries"]

Out [15]:

['ranked-choice voting disadvantages studies',
 'RCV implementation costs and challenges',
 'Effects of RCV on voter participation and strategic voting']

In [16]:

import json

revised = revision_chain.invoke(
    [
        HumanMessage(content="Why the US should NOT adopt rank-choice voting"),
        initial,
        ToolMessage(
            tool_call_id=initial.additional_kwargs["tool_calls"][0]["id"],
            content=json.dumps(
                tavily_tool.invoke(str(parsed[0]["args"]["search_queries"]))
            ),
        ),
    ]
)

In [17]:

parsed = parser.invoke(revised)
parsed

Out [17]:

[{'type': 'ReviseAnswer',
  'args': {'answer': 'Opponents of ranked-choice voting (RCV) in the United States highlight several key concerns against its adoption. First, RCV can introduce complexity and confusion for voters, potentially leading to increased ballot errors and decreased voter participation[1]. This complexity may undermine the democratic process by discouraging engagement. Second, the transition to RCV would entail considerable costs and logistical challenges, including updating voting machines and training election officials. The need for significant changes to electoral infrastructure could pose financial and operational burdens[2]. Third, critics argue RCV might not eliminate strategic voting, as voters could still rank candidates in a non-genuine manner to influence the election outcome[3]. Moreover, while RCV aims to elect consensus candidates, this can result in the election of moderate, less polarizing figures, potentially dampening political diversity and debate[4]. Lastly, the implementation of RCV could lead to initial confusion and concerns over election integrity during the transition, impacting public trust in the electoral process[5]. These factors combined suggest that the adoption of RCV in the U.S. might not yield the intended democratic enhancements and could introduce new challenges to the electoral system.\n\nReferences:\n[1] https://thefga.org/research/ranked-choice-voting-a-disaster-in-disguise/\n[2] https://www.npr.org/2023/12/13/1214199019/ranked-choice-voting-explainer\n[3] https://www.csg.org/2023/03/21/ranked-choice-voting-what-where-why-why-not/\n[4] https://cs50.harvard.edu/college/2024/spring/psets/3/runoff/\n[5] https://www.nytimes.com/2024/02/08/us/ranked-choice-voting-elections.html',
   'reflection': {'missing': 'Precise data or studies that quantify the impact of RCV on voter participation, ballot errors, and strategic voting.',
    'superfluous': 'The initial answer lacked direct references and included broad, unsubstantiated claims. The revised version focuses on specific criticisms and removes general explanations of RCV that do not directly address the question.',
    'search_queries': ['ranked-choice voting impact on voter turnout',
     'RCV ballot error rate study',
     'strategic voting in RCV elections'],
    'references': ['https://thefga.org/research/ranked-choice-voting-a-disaster-in-disguise/',
     'https://www.npr.org/2023/12/13/1214199019/ranked-choice-voting-explainer',
     'https://www.csg.org/2023/03/21/ranked-choice-voting-what-where-why-why-not/',
     'https://cs50.harvard.edu/college/2024/spring/psets/3/runoff/',
     'https://www.nytimes.com/2024/02/08/us/ranked-choice-voting-elections.html']}},
  'id': 'call_2HQBy80sY704ENtB6IvgRhZl'}]

Graph

In [57]:

from collections import defaultdict

from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation

# This a helper class we have that is useful for running tools
# It takes in an agent action and calls that tool and returns the result
tool_executor = ToolExecutor([tavily_tool])


def execute_tools(state: List[BaseMessage]) -> List[BaseMessage]:
    tool_invocation: AIMessage = state[-1]
    parsed_tool_calls = parser.invoke(tool_invocation)
    ids = []
    tool_invocations = []
    for parsed_call in parsed_tool_calls:
        for query in parsed_call["args"]["search_queries"]:
            tool_invocations.append(
                ToolInvocation(
                    # We only have this one for now. Would want to map it
                    # if we change
                    tool="tavily_search_results_json",
                    tool_input=query,
                )
            )
            ids.append(parsed_call["id"])

    outputs = tool_executor.batch(tool_invocations)
    outputs_map = defaultdict(dict)
    for id_, output, invocation in zip(ids, outputs, tool_invocations):
        outputs_map[id_][invocation.tool_input] = output

    return [
        ToolMessage(content=json.dumps(query_outputs), tool_call_id=id_)
        for id_, query_outputs in outputs_map.items()
    ]

In [58]:

MAX_ITERATIONS = 5


def _get_num_iterations(state: List[BaseMessage]):
    i = 0
    for m in state[::-1]:
        if not isinstance(m, (ToolMessage, AIMessage)):
            break
        i += 1
    return i


def event_loop(state: List[BaseMessage]) -> str:
    num_iterations = _get_num_iterations(state)
    if num_iterations > MAX_ITERATIONS:
        return END
    return "execute_tools"

In [59]:

from langgraph.graph import END, MessageGraph

builder = MessageGraph()
builder.add_node("draft", initial_answer_chain)
builder.add_node("execute_tools", execute_tools)
builder.add_node("revise", revision_chain)
builder.add_edge("draft", "execute_tools")
builder.add_edge("execute_tools", "revise")
builder.add_conditional_edges("revise", event_loop)
builder.set_entry_point("draft")
graph = builder.compile()

In [61]:

events = graph.stream(
    [HumanMessage(content="How should we handle the climate crisis?")]
)
for i, step in enumerate(events):
    node, output = next(iter(step.items()))
    print(f"## {i+1}. {node}")
    print(str(output)[:100] + " ...")
    print("---")

## 1. draft
content='' additional_kwargs={'tool_calls': [{'id': 'call_umpIUt5pLhCGImf8XCcMiZlD', 'function': {'a ...
---
## 2. execute_tools
[ToolMessage(content='{"successful climate change mitigation examples": [{"url": "https://www.washin ...
---
## 3. revise
content='' additional_kwargs={'tool_calls': [{'id': 'call_rutKnOqvPZoG5sSS6Be46RUn', 'function': {'a ...
---
## 4. execute_tools
[ToolMessage(content='{"climate crisis innovative policy examples": [{"url": "https://www.forbes.com ...
---
## 5. revise
content='' additional_kwargs={'tool_calls': [{'id': 'call_kUUMiW7z7IDbsMOlD25zI221', 'function': {'a ...
---
## 6. execute_tools
[ToolMessage(content='{"climate crisis innovative policy examples": [{"url": "https://www.worldwildl ...
---
## 7. revise
content='' additional_kwargs={'tool_calls': [{'id': 'call_EY7kiUhu6h3jemW0HrEfo2aP', 'function': {'a ...
---
## 8. __end__
[HumanMessage(content='How should we handle the climate crisis?'), AIMessage(content='', additional_ ...
---

In [ ]:

print(step[END])

Memory

In [ ]: