langgraph/examples/code_assistant/langgraph_code_assistant.ipynb

Code generation with flow

AlphaCodium presented an approach for code generation that uses aflow paradigm to construct an answer to a coding question iteratively..

AlphaCodium iteravely tests and improves an answer on public and AI-generated tests for a particular question.

We will implement some of these ideas from scratch using LangGraph:

1. We start with a set of documentation specified by a user
2. We use a long context LLM to ingest it, and answer a question based upon it
3. We perform two unit tests: Check imports and code execution

In [2]:

 ! pip install -U langchain_community langchain-openai langchain-anthropic langchain langgraph

Docs

Load LangChain Expression Language (LCEL) docs as an example.

In [1]:

from bs4 import BeautifulSoup as Soup
from langchain_community.document_loaders.recursive_url_loader import RecursiveUrlLoader

# LCEL docs
url = "https://python.langchain.com/docs/expression_language/"
loader = RecursiveUrlLoader(
    url=url, max_depth=20, extractor=lambda x: Soup(x, "html.parser").text
)
docs = loader.load()

# Sort the list based on the URLs and get the text
d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
d_reversed = list(reversed(d_sorted))
concatenated_content = "\n\n\n --- \n\n\n".join(
    [doc.page_content for doc in d_reversed]
)

LLMs

Code solution

We de-couple code solution and code formatting so that any LLM can be used for code solution.

In [15]:

from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# Grader prompt 
code_gen_prompt = ChatPromptTemplate.from_messages(
    [("system","""You are a coding assistant with expertise in LCEL, LangChain expression language. \n 
    Here is a full set of LCEL documentation:  \n ------- \n  {context} \n ------- \n Answer the user 
    question based on the above provided documentation. Ensure any code you provide can be executed \n 
    with all required imports and variables defined. Structure your answer with a description of the code solution. \n
    Then list the imports. And finally list the functioning code block. Here is the user question:"""),
    ("placeholder", "{messages}")]
)

# code_gen_llm = ChatOpenAI(temperature=0, model="gpt-4-0125-preview")
code_gen_llm = ChatAnthropic(temperature=0, model='claude-3-opus-20240229')

code_gen_chain = code_gen_prompt | code_gen_llm | StrOutputParser()
question = "How do I build a RAG chain in LCEL?"
solution = code_gen_chain.invoke({"context"  : concatenated_content, "messages" : [("user",question)]})

Formatted code

The structured output generation is handled in a seperate step.

In [17]:

from langchain_openai import ChatOpenAI
from langchain_core.pydantic_v1 import BaseModel, Field

# Prompt
prompt = ChatPromptTemplate.from_messages(
        [("system","""You are an expert a code formatting, strating with a code solution \n
        Structure the solution in three parts with a prefix that defines the problem, then \n 
        list the imports, and finally list the functioning code block.""" ),
         ("user", "Here is the code solution: {code}"),]) 
    
# Data model
class code(BaseModel):
    """Code output"""

    prefix: str = Field(description="Description of the problem and approach")
    imports: str = Field(description="Code block import statements")
    code: str = Field(description="Code block not including import statements")

# Formatter 
llm = ChatOpenAI(model="gpt-4-0125-preview", temperature=0)
llm_formatter = llm.with_structured_output(code)
structured_code_formatter = prompt | llm_formatter
output = structured_code_formatter.invoke([("code",solution)])

Parameters

In [18]:

max_iterations = 3

State

Our state is a dict that will contain keys (errors, question, code generation) relevant to code generation.

In [6]:

from typing import Dict, TypedDict, List

class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        error : Binary flag for control flow to indicate whether test error was tripped
        messages : With user question, error messages, reasoning
        generation : Code solution
        iterations : Number of tries 
    """

    error : str
    messages : List
    generation : str
    iterations : int

Graph

Our graph lays out the logical flow shown in the figure above.

In [28]:

from operator import itemgetter
from langchain.prompts import PromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.runnables import RunnablePassthrough

### Nodes

def generate(state: GraphState):
    """
    Generate a code solution

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")
    
    # State
    messages = state["messages"]
    iterations = state["iterations"]
    
    # Solution
    solution = code_gen_chain.invoke({"context"  : concatenated_content, "messages" : messages})

    # Structured output
    code_solution = structured_code_formatter.invoke([("code",solution)])
    
    # Increment
    iterations = iterations + 1
    return {"generation": code_solution, "messages": messages, "iterations": iterations}


def code_check(state: GraphState):
    """
    Check code

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """

    print("---CHECKING CODE---")
    
    ## State
    messages = state["messages"]
    code_solution = state["generation"]
    iterations = state["iterations"]

    # Get solution components
    prefix = code_solution.prefix
    imports = code_solution.imports
    code = code_solution.code

    # Check imports
    try:
        exec(imports)
    except Exception as e:
        print("---CODE IMPORT CHECK: FAILED---")
        error_message = [("user", f"Import Test Failure: You are required to fix the import error: {e}")]
        messages += error_message
        return {"generation": code_solution, "messages": messages, "iterations": iterations, "error": "yes"}
    
    # Check execution
    try:
        exec(imports + "\n" + code)
    except Exception as e:
        print("---CODE BLOCK CHECK: FAILED---")
        error_message = [("user", f"Execution Test Failure: You are required to fix the code execution error: {e}")]
        return {"generation": code_solution, "messages": messages, "iterations": iterations, "error": "yes"}
  
    # No errors
    print("---NO CODE TEST FAILURES---")
    return {"generation": code_solution, "messages": messages, "iterations": iterations, "error": "no"}


def reflect(state: GraphState):
    """
    Reflect on errors

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, generation
    """

    print("---GENERATING CODE SOLUTION---")
    
    # State
    messages = state["messages"]
    iterations = state["iterations"]

    # Prompt reflection
    reflection_message = [("user", """You tried to solve this problem and failed a unit test. Reflect on this failure
                                    given the provided documentation. Carefully write suggestions based on the 
                                    documentation to avoid making this mistake again.""")]
    messages += reflection_message
    
    # Suggesting
    reflections = code_gen_chain.invoke({"context"  : concatenated_content, "messages" : messages})
    messages += reflections
    return {"generation": code_solution, "messages": messages, "iterations": iterations}


### Edges

def decide_to_finish(state: GraphState):
    """
    Determines whether to finish.

    Args:
        state (dict): The current graph state

    Returns:
        str: Next node to call
    """
    error = state["error"]
    iterations = state["iterations"]

    if error == "no" or iterations == max_iterations:
        print("---DECISION: FINISH---")
        return "end"
    else:
        print("---DECISION: RE-TRY SOLUTION---")
        return "reflect" # Or, directly back to generate

In [29]:

from langgraph.graph import END, StateGraph

workflow = StateGraph(GraphState)

# Define the nodes
workflow.add_node("generate", generate)  # generation solution
workflow.add_node("check_code", code_check)  # check code
workflow.add_node("reflect", reflect)  # reflect

# Build graph
workflow.set_entry_point("generate")
workflow.add_edge("generate", "check_code")
workflow.add_conditional_edges(
    "check_code",
    decide_to_finish,
    {
        "end": END,
        "reflect": "reflect",
    },
)
workflow.add_edge("reflect", "generate")
app = workflow.compile()

In [30]:

question = "How do I build a RAG chain in LCEL?"
app.invoke({"messages":[("user",question)],"iterations":0})

Out [30]:

---GENERATING CODE SOLUTION---
---CHECKING CODE---
---NO CODE TEST FAILURES---
---DECISION: FINISH---

{'error': 'no',
 'messages': [('user', 'How do I build a RAG chain in LCEL?')],
 'generation': code(prefix="Problem: Building a Retrieval-Augmented Generation (RAG) Chain using LangChain Expression Language (LCEL).\n\nApproach: To construct a RAG chain in LCEL, it's essential to integrate a retriever, a prompt template, a language model, and an output parser. The process involves creating a retriever to fetch relevant documents based on a query, defining a prompt template that incorporates the retrieved context and a question, instantiating a language model, and creating a chain by connecting the retriever, prompt, model, and output parser. Finally, the chain is invoked with a question to obtain the answer.", imports='from langchain_community.vectorstores import FAISS\nfrom langchain_core.output_parsers import StrOutputParser\nfrom langchain_core.prompts import ChatPromptTemplate\nfrom langchain_core.runnables import RunnablePassthrough\nfrom langchain_openai import ChatOpenAI, OpenAIEmbeddings', code='# Load documents into a vector store\nvectorstore = FAISS.from_texts(\n    ["harrison worked at kensho"], embedding=OpenAIEmbeddings()\n)\nretriever = vectorstore.as_retriever()\n\n# Define prompt template\ntemplate = """Answer the question based only on the following context:\n{context}\nQuestion: {question}"""\nprompt = ChatPromptTemplate.from_template(template)\n\n# Instantiate model and output parser\nmodel = ChatOpenAI()\noutput_parser = StrOutputParser()\n\n# Create RAG chain \nchain = (\n    {"context": retriever, "question": RunnablePassthrough()}\n    | prompt \n    | model\n    | output_parser\n)\n\n# Run the chain\nchain.invoke("where did harrison work?")'),
 'iterations': 1}

Eval

Here is a public dataset of LCEL questions.

I saved this as test-LCEL-code-gen.

You can also find the csv here.

In [32]:

import langsmith
client = langsmith.Client()

In [ ]:

# Clone the dataset to your tenant to use it
public_dataset = ("https://smith.langchain.com/public/326674a6-62bd-462d-88ae-eea49d503f9d/d")
client.clone_public_dataset(public_dataset)

Custom evals.

In [33]:

from langsmith.schemas import Example, Run

def check_import(run: Run, example: Example) -> dict: 
    imports = run.outputs.get("imports")
    try:
        exec(imports)
        return {"key": "import_check" , "score": 1} 
    except:
        return {"key": "import_check" , "score": 0} 

def check_execution(run: Run, example: Example) -> dict: 
    imports = run.outputs.get("imports")
    code = run.outputs.get("code")
    try:
        exec(imports + "\n" + code)
        return {"key": "code_execution_check" , "score": 1} 
    except:
        return {"key": "code_execution_check" , "score": 0} 

Compare LangGraph to Context Stuffing.

In [34]:

def predict_base_case(example: dict):
    """ Context stuffing """
    solution = code_gen_chain.invoke({"context"  : concatenated_content, "messages" : [("user",example["question"])] })
    solution_structured = structured_code_formatter.invoke([("code",solution)])
    return {"imports": solution_structured.imports, "code": solution_structured.code}

def predict_langgraph(example: dict):
    """ LangGraph """
    graph = app.invoke({"messages":[("user",example["question"])],"iterations":0})
    solution = graph["generation"]
    return {"imports": solution.imports, "code": solution.code}

In [ ]:

from langsmith.evaluation import evaluate

# Evaluator
code_evalulator = [check_import,check_execution]

# Dataset
dataset_name = "test-LCEL-code-gen"

# Run base case
experiment_results = evaluate(
    predict_base_case,
    data=dataset_name,
    evaluators=code_evalulator,
    experiment_prefix="test-without-langgraph",
    metadata={
      "variant": "Claude3",
    },
)

# Run with langgraph
experiment_results = evaluate(
    predict_langgraph,
    data=dataset_name,
    evaluators=code_evalulator,
    experiment_prefix="test-with-langgraph",
    metadata={
      "variant": "Claude3",
    },
)

/var/folders/l9/bpjxdmfx7lvd1fbdjn38y5dh0000gn/T/ipykernel_43078/1906326331.py:10: UserWarning: Function evaluate is in beta.
  experiment_results = evaluate(

View the evaluation results for experiment: 'test-without-langgraph:1f3259c' at:
https://smith.langchain.com/o/1fa8b1f4-fcb9-4072-9aa9-983e35ad61b8/datasets/76bc1d66-8cb6-4f6d-b633-7cd077937e46/compare?selectedSessions=1c92918c-b111-412d-b845-a96f2ee37f7c

0it [00:00, ?it/s]

Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}
Error running target function: Error code: 429 - {'type': 'error', 'error': {'type': 'rate_limit_error', 'message': 'Number of concurrent connections has exceeded your rate limit. Please try again later or contact sales at https://www.anthropic.com/contact-sales to discuss your options for a rate limit increase.'}}

TODO: Clean This Later

Compute standard error across 4 trials.

In [14]:

# You will have to update these to match the tests you ran.
# The test name can be found at langgraph_results["project_name"]
langgraph = [
    "80db-context-stuffing-with-langgraph",
    "060c-context-stuffing-with-langgraph",
    "93cd-context-stuffing-with-langgraph",
    "60ef-context-stuffing-with-langgraph",
]

In [15]:

no_langgraph = [
    "b493-context-stuffing-no-langgraph",
    "eb8a-context-stuffing-no-langgraph",
    "b88c-context-stuffing-no-langgraph",
    "0aaa-context-stuffing-no-langgraph",
]

In [39]:

import pandas as pd


def prepare_dataframe(project, trial_number, chain):
    df = client.get_test_results(project_name=project)
    df = df.dropna(subset=["feedback.check_execution", "feedback.check_import"])
    df = df[["input.question", "feedback.check_execution", "feedback.check_import"]]
    df["trial #"] = trial_number
    df["chain"] = chain
    return df


# Prepare each dataframe
dfs_chain1 = [
    prepare_dataframe(project, i + 1, "LangGraph")
    for i, project in enumerate(langgraph)
]
dfs_chain2 = [
    prepare_dataframe(project, i + 1, "No LangGraph")
    for i, project in enumerate(no_langgraph)
]

# Combine all dataframes
final_df = pd.concat(dfs_chain1 + dfs_chain2, ignore_index=True)

In [42]:

final_df.groupby("chain").size()

Out [42]:

chain
LangGraph       78
No LangGraph    79
dtype: int64

In [40]:

import pandas as pd


def group_standard_error(group):
    """
    Calculate the standard error for the 'correct' column in a given group.

    The function assumes the 'correct' column contains binary values (0 or 1).
    It computes the standard error based on the formula for the standard error
    of a proportion, which is sqrt(p * (1 - p) / n), where p is the proportion
    of successes (1s) and n is the total number of trials.

    Args:
    group (pd.DataFrame): A DataFrame group with a 'correct' column.

    Returns:
    pd.Series: A series containing the standard error of the 'correct' column.
    """
    # 3 trials x 20 questions per trial = 60
    total_trials = len(group)
    std_errors = {}
    for column in ["feedback.check_import", "feedback.check_execution"]:
        # Number correct
        occurrences = group[column].sum()
        # Total trials
        fraction = occurrences / total_trials
        # Standard error
        std_errors[column] = (fraction * (1 - fraction) / total_trials) ** 0.5
    return pd.Series(std_errors)


# Calculate standard errors
std_errors = final_df.groupby(["chain"]).apply(group_standard_error)

# Calculate the fraction of correct answers
grouped_frac_correct = (
    final_df.groupby("chain")[
        ["feedback.check_import", "feedback.check_execution"]
    ].sum()
    / final_df.groupby("chain")[
        ["feedback.check_import", "feedback.check_execution"]
    ].count()
)

# Concatenate the fraction correct data with the standard errors
correct_frac_and_errors = pd.concat([grouped_frac_correct, std_errors], axis=1)

# If you want to rename the columns for clarity
correct_frac_and_errors.columns = [
    "Fraction Imports Correct",
    "Fraction Execution Correct",
    "Imports Correct Std Error",
    "Execution Correct Std Error",
]
correct_frac_and_errors

Out [40]:

	Fraction Imports Correct	Fraction Execution Correct	Imports Correct Std Error	Execution Correct Std Error
chain
LangGraph	1.000000	0.807692	0.000000	0.044625
No LangGraph	0.987342	0.556962	0.012578	0.055888

In [47]:

import matplotlib.pyplot as plt
import seaborn as sns


def plt_combined_bar_graph(df, fraction_fields, error_fields, titles, ylabels):
    """
    Plot bar graphs with error bars for specified fields in the provided DataFrame as subplots.

    Args:
    df (pd.DataFrame): The DataFrame containing the data to be plotted.
    fraction_fields (list[str]): List of column names in the DataFrame to be plotted on the y-axis for fractions.
    error_fields (list[str]): List of column names in the DataFrame to be plotted for standard errors.
    titles (list[str]): Titles of the plots.
    ylabels (list[str]): Labels for the y-axis.

    This function does not return any value but displays the bar graph.
    """
    n = len(fraction_fields)  # Number of plots to create
    fig, axs = plt.subplots(1, n, figsize=(10 * n, 9), sharey=True)

    for i, (fraction_field, error_field, title, ylabel) in enumerate(
        zip(fraction_fields, error_fields, titles, ylabels)
    ):
        barplot = sns.barplot(
            x="chain",
            y=fraction_field,
            data=df.sort_values(
                "chain", ascending=False
            ),  # Sort the DataFrame to reverse the order
            ax=axs[i],
            capsize=0.1,
            errorbar=None,
        )

        # Add error bars manually
        for j, bar in enumerate(barplot.patches):
            # Get the error for the current bar
            error = df.sort_values("chain", ascending=False)[error_field].iloc[j]
            # Add error bars to each bar
            axs[i].errorbar(
                x=bar.get_x() + bar.get_width() / 2,
                y=bar.get_height(),
                yerr=error,
                fmt="none",
                capsize=5,
                color="black",
            )

        axs[i].set_title(title)
        axs[i].set_xlabel("Chain")
        axs[i].set_ylabel(ylabel)

    plt.tight_layout()
    plt.show()


# Define the columns and labels for the plots
fraction_fields = ["Fraction Imports Correct", "Fraction Execution Correct"]
error_fields = ["Imports Correct Std Error", "Execution Correct Std Error"]

titles = [
    "Feedback Check Import Fraction by Chain",
    "Feedback Check Execution Fraction by Chain",
]
ylabels = ["Fraction Correct", "Fraction Correct"]

# Call the function with the specified arguments
plt_combined_bar_graph(
    correct_frac_and_errors, fraction_fields, error_fields, titles, ylabels
)

In [ ]: