Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here.
fromlangchain_community.document_loadersimportWebBaseLoaderfromlangchain_community.vectorstoresimportChromafromlangchain_openaiimportOpenAIEmbeddingsfromlangchain_text_splittersimportRecursiveCharacterTextSplitterurls=["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()forurlinurls]docs_list=[itemforsublistindocsforiteminsublist]text_splitter=RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=100,chunk_overlap=50)doc_splits=text_splitter.split_documents(docs_list)# Add to vectorDBvectorstore=Chroma.from_documents(documents=doc_splits,collection_name="rag-chroma",embedding=OpenAIEmbeddings(),)retriever=vectorstore.as_retriever()
fromlangchain.tools.retrieverimportcreate_retriever_toolretriever_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=[retriever_tool]
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.
fromtypingimportAnnotated,Sequencefromtyping_extensionsimportTypedDictfromlangchain_core.messagesimportBaseMessagefromlanggraph.graph.messageimportadd_messagesclassAgentState(TypedDict):# The add_messages function defines how an update should be processed# Default is to replace. add_messages says "append"messages:Annotated[Sequence[BaseMessage],add_messages]
This notebook uses Pydantic v2 BaseModel, which requires langchain-core >= 0.3. Using langchain-core < 0.3 will result in errors due to mixing of Pydantic v1 and v2 BaseModels.
fromtypingimportAnnotated,Literal,Sequencefromtyping_extensionsimportTypedDictfromlangchainimporthubfromlangchain_core.messagesimportBaseMessage,HumanMessagefromlangchain_core.output_parsersimportStrOutputParserfromlangchain_core.promptsimportPromptTemplatefromlangchain_openaiimportChatOpenAIfrompydanticimportBaseModel,Fieldfromlanggraph.prebuiltimporttools_condition### Edgesdefgrade_documents(state)->Literal["generate","rewrite"]:""" Determines whether the retrieved documents are relevant to the question. Args: state (messages): The current state Returns: str: A decision for whether the documents are relevant or not """print("---CHECK RELEVANCE---")# Data modelclassgrade(BaseModel):"""Binary score for relevance check."""binary_score:str=Field(description="Relevance score 'yes' or 'no'")# LLMmodel=ChatOpenAI(temperature=0,model="gpt-4o",streaming=True)# LLM with tool and validationllm_with_tool=model.with_structured_output(grade)# Promptprompt=PromptTemplate(template="""You are a grader assessing relevance of a retrieved document to a user question. \n Here is the retrieved document: \n\n{context}\n\n Here is the user question: {question}\n If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.""",input_variables=["context","question"],)# Chainchain=prompt|llm_with_toolmessages=state["messages"]last_message=messages[-1]question=messages[0].contentdocs=last_message.contentscored_result=chain.invoke({"question":question,"context":docs})score=scored_result.binary_scoreifscore=="yes":print("---DECISION: DOCS RELEVANT---")return"generate"else:print("---DECISION: DOCS NOT RELEVANT---")print(score)return"rewrite"### Nodesdefagent(state):""" Invokes the agent model to generate a response based on the current state. Given the question, it will decide to retrieve using the retriever tool, or simply end. Args: state (messages): The current state Returns: dict: The updated state with the agent response appended to messages """print("---CALL AGENT---")messages=state["messages"]model=ChatOpenAI(temperature=0,streaming=True,model="gpt-4-turbo")model=model.bind_tools(tools)response=model.invoke(messages)# We return a list, because this will get added to the existing listreturn{"messages":[response]}defrewrite(state):""" Transform the query to produce a better question. Args: state (messages): The current state Returns: dict: The updated state with re-phrased question """print("---TRANSFORM QUERY---")messages=state["messages"]question=messages[0].contentmsg=[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: """,)]# Gradermodel=ChatOpenAI(temperature=0,model="gpt-4-0125-preview",streaming=True)response=model.invoke(msg)return{"messages":[response]}defgenerate(state):""" Generate answer Args: state (messages): The current state Returns: dict: The updated state with re-phrased question """print("---GENERATE---")messages=state["messages"]question=messages[0].contentlast_message=messages[-1]docs=last_message.content# Promptprompt=hub.pull("rlm/rag-prompt")# LLMllm=ChatOpenAI(model_name="gpt-4o-mini",temperature=0,streaming=True)# Post-processingdefformat_docs(docs):return"\n\n".join(doc.page_contentfordocindocs)# Chainrag_chain=prompt|llm|StrOutputParser()# Runresponse=rag_chain.invoke({"context":docs,"question":question})return{"messages":[response]}print("*"*20+"Prompt[rlm/rag-prompt]"+"*"*20)prompt=hub.pull("rlm/rag-prompt").pretty_print()# Show what the prompt looks like
********************Prompt[rlm/rag-prompt]********************================================[1m Human Message [0m=================================You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.Question: [33;1m[1;3m{question}[0m Context: [33;1m[1;3m{context}[0m Answer:
Then call agent with the tool output added to messages (state)
fromlanggraph.graphimportEND,StateGraph,STARTfromlanggraph.prebuiltimportToolNode# Define a new graphworkflow=StateGraph(AgentState)# Define the nodes we will cycle betweenworkflow.add_node("agent",agent)# agentretrieve=ToolNode([retriever_tool])workflow.add_node("retrieve",retrieve)# retrievalworkflow.add_node("rewrite",rewrite)# Re-writing the questionworkflow.add_node("generate",generate)# Generating a response after we know the documents are relevant# Call agent node to decide to retrieve or notworkflow.add_edge(START,"agent")# Decide whether to retrieveworkflow.add_conditional_edges("agent",# Assess agent decisiontools_condition,{# Translate the condition outputs to nodes in our graph"tools":"retrieve",END:END,},)# Edges taken after the `action` node is called.workflow.add_conditional_edges("retrieve",# Assess agent decisiongrade_documents,)workflow.add_edge("generate",END)workflow.add_edge("rewrite","agent")# Compilegraph=workflow.compile()
fromIPython.displayimportImage,displaytry:display(Image(graph.get_graph(xray=True).draw_mermaid_png()))exceptException:# This requires some extra dependencies and is optionalpass
importpprintinputs={"messages":[("user","What does Lilian Weng say about the types of agent memory?"),]}foroutputingraph.stream(inputs):forkey,valueinoutput.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={'tool_calls': [{'index': 0, 'id': 'call_z36oPZN8l1UC6raxrebqc1bH', 'function': {'arguments': '{"query":"types of agent memory"}', 'name': 'retrieve_blog_posts'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-2bad2518-8187-4d8f-8e23-2b9501becb6f-0', tool_calls=[{'name': 'retrieve_blog_posts', 'args': {'query': 'types of agent memory'}, 'id': 'call_z36oPZN8l1UC6raxrebqc1bH'}])]}'\n---\n'---CHECK RELEVANCE------DECISION: DOCS RELEVANT---"Output from node 'retrieve':"'---'{ 'messages': [ ToolMessage(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\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\nMemory\n\nShort-term memory: I would consider all the in-context learning (See Prompt Engineering) as utilizing short-term memory of the model to learn.\nLong-term memory: This provides the agent with the capability to retain and recall (infinite) information over extended periods, often by leveraging an external vector store and fast retrieval.\n\n\nTool use\n\nThe design of generative agents combines LLM with memory, planning and reflection mechanisms to enable agents to behave conditioned on past experience, as well as to interact with other agents.', name='retrieve_blog_posts', id='d815f283-868c-4660-a1c6-5f6e5373ca06', tool_call_id='call_z36oPZN8l1UC6raxrebqc1bH')]}'\n---\n'---GENERATE---"Output from node 'generate':"'---'{ 'messages': [ 'Lilian Weng discusses short-term and long-term memory in ' 'agent systems. Short-term memory is used for in-context ' 'learning, while long-term memory allows agents to retain and ' 'recall information over extended periods.']}'\n---\n'