langgraph/examples/code_assistant/langgraph_code_gen.ipynb

In [ ]:

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

Optionally, use LangSmith for tracing:

export LANGCHAIN_TRACING_V2=true
export LANGCHAIN_ENDPOINT=https://api.smith.langchain.com
export LANGCHAIN_API_KEY=<your-api-key>

Docs

In [14]:

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 in 'metadata' -> 'source'
d_sorted = sorted(docs, key=lambda x: x.metadata["source"])
d_reversed = list(reversed(d_sorted))

# Concatenate the 'page_content' of each sorted dictionary
concatenated_content = "\n\n\n --- \n\n\n".join(
    [doc.page_content for doc in d_reversed]
)

Tool Use

In [17]:

from operator import itemgetter
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.output_parsers.openai_tools import PydanticToolsParser
from langchain_core.utils.function_calling import convert_to_openai_tool
   
## 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")

## LLM
model = ChatOpenAI(temperature=0, model="gpt-4-0125-preview", streaming=True)

# Tool
code_tool_oai = convert_to_openai_tool(code)

# LLM with tool and enforce invocation
llm_with_tool = model.bind(
    tools=[code_tool_oai],
    tool_choice={"type": "function", "function": {"name": "code"}},
)

# Parser
parser_tool = PydanticToolsParser(tools=[code])

## Prompt
template = """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. \n
    Ensure any code you provide can be executed with all required imports and variables defined. \n
    Structure your answer with a description of the code solution. \n
    Then list the imports. And finally list the functioning code block. \n
    Here is the user question: \n --- --- --- \n {question}"""

# Prompt 
prompt = PromptTemplate(
    template=template,
    input_variables=["context", "question"],
)

# Chain
chain = (
    {
        # "context": lambda x: docs,
        "context": lambda x: concatenated_content,
        "question": itemgetter("question"),
    }
    | prompt
    | llm_with_tool 
    | parser_tool
)

In [19]:

chain.invoke({"question":"How to create a RAG chain in LCEL?"})

Out [19]:

[code(prefix="To create a Retrieval-Augmented Generation (RAG) chain in LangChain Expression Language (LCEL), you need to follow a structured approach. This involves setting up a retriever to fetch relevant documents based on the user's query, and then using those documents to generate a response with a language model. Here's a step-by-step guide to creating a RAG chain in LCEL:\n\n", imports='from operator import itemgetter\nfrom 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='# Initialize the vector store with sample texts and embeddings\nvectorstore = FAISS.from_texts(\n    ["harrison worked at kensho"], embedding=OpenAIEmbeddings())\n\n# Create a retriever from the vector store\nretriever = vectorstore.as_retriever()\n\n# Define the prompt template\ntemplate = """Answer the question based only on the following context:{context}Question: {question}"""\nprompt = ChatPromptTemplate.from_template(template)\n\n# Initialize the model\nmodel = ChatOpenAI()\n\n# Create the RAG chain\nchain = (\n    {"context": retriever, "question": RunnablePassthrough()}    | prompt    | model    | StrOutputParser()\n)\n\n# Invoke the chain with a query\nresponse = chain.invoke("where did harrison work?")\nprint(response)')]

State

In [20]:

from typing import Dict, TypedDict

from langchain_core.messages import BaseMessage


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

    Attributes:
        keys: A dictionary where each key is a string.
    """

    keys: Dict[str, any]

In [21]:

from operator import itemgetter
from bs4 import BeautifulSoup as Soup
from langchain_community.document_loaders.recursive_url_loader import RecursiveUrlLoader
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.output_parsers import PydanticOutputParser
from langchain.output_parsers.openai_tools import PydanticToolsParser
from langchain_core.utils.function_calling import convert_to_openai_tool

def generate(state):
    """
    Generate a code solution based on LCEL docs and the input question 
    with optional feedback from code execution tests 

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents
    """
    
    ## State
    state_dict = state["keys"]
    question = state_dict["question"]
    iter = state_dict["iterations"]
   
    ## 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")
    
    ## LLM
    model = ChatOpenAI(temperature=0, model="gpt-4-0125-preview", streaming=True)
    
    # Tool
    code_tool_oai = convert_to_openai_tool(code)
    
    # LLM with tool and enforce invocation
    llm_with_tool = model.bind(
        tools=[convert_to_openai_tool(code_tool_oai)],
        tool_choice={"type": "function", "function": {"name": "code"}},
    )
    
    # Parser
    parser_tool = PydanticToolsParser(tools=[code])
    
    ## Prompt
    template = """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. \n
        Ensure any code you provide can be executed with all required imports and variables defined. \n
        Structure your answer with a description of the code solution. \n
        Then list the imports. And finally list the functioning code block. \n
        Here is the user question: \n --- --- --- \n {question}"""

    ## Generation
    if "error" in state_dict:
        print("---RE-GENERATE SOLUTION w/ ERROR FEEDBACK---")
        
        error = state_dict["error"]
        code_solution = state_dict["generation"]
        
        # Udpate prompt 
        addendum = """  \n --- --- --- \n You previously tried to solve this problem. \n Here is your solution:  
                    \n --- --- --- \n {generation}  \n --- --- --- \n  Here is the resulting error from code 
                    execution:  \n --- --- --- \n {error}  \n --- --- --- \n Please re-try to answer this. 
                    Structure your answer with a description of the code solution. \n Then list the imports. 
                    And finally list the functioning code block. Structure your answer with a description of 
                    the code solution. \n Then list the imports. And finally list the functioning code block. 
                    \n Here is the user question: \n --- --- --- \n {question}"""
        template = template +  addendum

        # Prompt 
        prompt = PromptTemplate(
            template=template,
            input_variables=["context", "question", "generation", "error"],
        )
        
        # Chain
        chain = (
            {
                "context": lambda x: concatenated_content,
                "question": itemgetter("question"),
                "generation": itemgetter("generation"),
                "error": itemgetter("error"),
            }
            | prompt
            | llm_with_tool 
            | parser_tool
        )

        code_solution = chain.invoke({"question":question,
                                      "generation":str(code_solution[0]),
                                      "error":error})
                
    else:
        print("---GENERATE SOLUTION---")
        
        # Prompt 
        prompt = PromptTemplate(
            template=template,
            input_variables=["context", "question"],
        )

        # Chain
        chain = (
            {
                # "context": lambda x: docs,
                "context": lambda x: concatenated_content,
                "question": itemgetter("question"),
            }
            | prompt
            | llm_with_tool 
            | parser_tool
        )

        code_solution = chain.invoke({"question":question})

    iter = iter+1 
    return {"keys": {"generation": code_solution, "question": question, "iterations": iter}}

def check_code_imports(state):
    """
    Check imports

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """
    
    ## State
    print("---CHECKING CODE IMPORTS---")
    state_dict = state["keys"]
    question = state_dict["question"]
    code_solution = state_dict["generation"]
    imports = code_solution[0].imports
    iter = state_dict["iterations"]

    try:        
        # Attempt to execute the imports
        exec(imports)
    except Exception as e:
        print("---CODE IMPORT CHECK: FAILED---")
        # Catch any error during execution (e.g., ImportError, SyntaxError)
        error = f"Execution error: {e}"
        if "error" in state_dict:
            error_prev_runs = state_dict["error"]
            error = error_prev_runs + "\n --- Most recent run error --- \n" + error     
    else:
        print("---CODE IMPORT CHECK: SUCCESS---")
        # No errors occurred
        error = "None"

    return {"keys": {"generation": code_solution, "question": question, "error": error, "iterations":iter}}

def check_code_execution(state):
    """
    Check code block execution

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, error
    """
    
    ## State
    print("---CHECKING CODE EXECUTION---")
    state_dict = state["keys"]
    question = state_dict["question"]
    code_solution = state_dict["generation"]
    prefix = code_solution[0].prefix
    imports = code_solution[0].imports
    code = code_solution[0].code
    code_block = imports +"\n"+ code
    iter = state_dict["iterations"]

    try:        
        # Attempt to execute the code block
        exec(code_block)
    except Exception as e:
        print("---CODE BLOCK CHECK: FAILED---")
        # Catch any error during execution (e.g., ImportError, SyntaxError)
        error = f"Execution error: {e}"
        if "error" in state_dict:
            error_prev_runs = state_dict["error"]
            error = error_prev_runs + "\n --- Most recent run error --- \n" + error  
    else:
        print("---CODE BLOCK CHECK: SUCCESS---")
        # No errors occurred
        error = "None"

    return {"keys": {"generation": code_solution, 
                     "question": question, 
                     "error": error, 
                     "prefix":prefix,
                     "imports":imports,
                     "iterations":iter,
                     "code":code}}

### Edges

def decide_to_check_code_exec(state):
    """
    Determines whether to test code execution, or re-try answer generation.

    Args:
       state (dict): The current graph state

    Returns:
        str: Next node to call
    """

    print("---DECIDE TO TEST CODE EXECUTION---")
    state_dict = state["keys"]
    question = state_dict["question"]
    code_solution = state_dict["generation"]
    error = state_dict["error"]

    if error == "None":
        # All documents have been filtered check_relevance
        # We will re-generate a new query
        print("---DECISION: TEST CODE EXECUTION---")
        return "check_code_execution"
    else:
        # We have relevant documents, so generate answer
        print("---DECISION: RE-TRY SOLUTION---")
        return "generate"

def decide_to_finish(state):
    """
    Determines whether to finish (re-try code 3 times.

    Args:
        state (dict): The current graph state

    Returns:
        str: Next node to call
    """

    print("---DECIDE TO TEST CODE EXECUTION---")
    state_dict = state["keys"]
    question = state_dict["question"]
    code_solution = state_dict["generation"]
    error = state_dict["error"]
    iter = state_dict["iterations"]

    if error == "None" or iter == 3:
        # All documents have been filtered check_relevance
        # We will re-generate a new query
        print("---DECISION: TEST CODE EXECUTION---")
        return "end"
    else:
        # We have relevant documents, so generate answer
        print("---DECISION: RE-TRY SOLUTION---")
        return "generate"

In [23]:

from langgraph.graph import END, StateGraph

workflow = StateGraph(GraphState)

# Define the nodes
workflow.add_node("generate", generate)  # generation solution
workflow.add_node("check_code_imports", check_code_imports)  # check imports
workflow.add_node("check_code_execution", check_code_execution)  # check execution

# Build graph
workflow.set_entry_point("generate")
workflow.add_edge("generate", "check_code_imports")
workflow.add_conditional_edges(
    "check_code_imports",
    decide_to_check_code_exec,
    {
        "check_code_execution": "check_code_execution",
        "generate": "generate",
    },
)
workflow.add_conditional_edges(
    "check_code_execution",
    decide_to_finish,
    {
        "end": END,
        "generate": "generate",
    },
)

# Compile
app = workflow.compile()

In [24]:

question = "I am passing text key 'foo' to my prompt and want to process it with a function, process_text(...), prior to the prompt. How can I do this using LCEL?"
config = {"recursion_limit": 50}
answer = app.invoke({"keys":{"question":question, "iterations":0}},config=config)

---GENERATE SOLUTION---
---CHECKING CODE IMPORTS---
---CODE IMPORT CHECK: SUCCESS---
---DECIDE TO TEST CODE EXECUTION---
---DECISION: TEST CODE EXECUTION---
---CHECKING CODE EXECUTION---
---CODE BLOCK CHECK: FAILED---
---DECIDE TO TEST CODE EXECUTION---
---DECISION: RE-TRY SOLUTION---
---RE-GENERATE SOLUTION w/ ERROR FEEDBACK---
---CHECKING CODE IMPORTS---
---CODE IMPORT CHECK: SUCCESS---
---DECIDE TO TEST CODE EXECUTION---
---DECISION: TEST CODE EXECUTION---
---CHECKING CODE EXECUTION---
Why did the bear break up with his girlfriend?
Because he couldn't bear the relationship anymore!
---CODE BLOCK CHECK: SUCCESS---
---DECIDE TO TEST CODE EXECUTION---
---DECISION: TEST CODE EXECUTION---

In [27]:

answer['keys']['generation'][0].prefix

Out [27]:

"To process a text key 'foo' with a function before passing it to a prompt using LCEL, you can use a RunnableLambda. This allows you to define a custom function that processes the input text and then passes the modified text to the prompt. Here's how you can do it:"

In [30]:

exec(answer['keys']['generation'][0].imports)

In [31]:

exec(answer['keys']['generation'][0].code)

Why did the bear break up with his girlfriend? 

Because he couldn't bear the relationship any longer!

In [ ]: