├── utils
    ├── __init__.py
    └── utils.py
├── env.tmp
├── assets
    ├── lab1_1.png
    ├── lab1_2.png
    ├── lab2_1.png
    ├── lab3_1.png
    ├── lab3_2.png
    ├── lab5_1.png
    ├── lab5_2.png
    ├── lab6_1.png
    ├── lab6_2.png
    ├── lab6_3.png
    ├── lab6_4.png
    ├── lab6_5.png
    └── lab6_6.png
├── .gitignore
├── CODE_OF_CONDUCT.md
├── Lab_1
    ├── README.md
    └── Lab_1.ipynb
├── Lab_5
    ├── README.md
    └── Lab_5.ipynb
├── LICENSE
├── Lab_4
    ├── README.md
    └── Lab_4.ipynb
├── Lab_2
    ├── README.md
    └── Lab_2.ipynb
├── Lab_3
    ├── README.md
    └── Lab_3.ipynb
├── pyproject.toml
├── CONTRIBUTING.md
├── Lab_6
    ├── README.md
    ├── helper.py
    └── Lab_6.ipynb
└── README.md


/utils/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/env.tmp:
--------------------------------------------------------------------------------
1 | TAVILY_API_KEY=<paste your api key here>
2 | AWS_REGION=us-east-1


--------------------------------------------------------------------------------
/assets/lab1_1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab1_1.png


--------------------------------------------------------------------------------
/assets/lab1_2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab1_2.png


--------------------------------------------------------------------------------
/assets/lab2_1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab2_1.png


--------------------------------------------------------------------------------
/assets/lab3_1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab3_1.png


--------------------------------------------------------------------------------
/assets/lab3_2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab3_2.png


--------------------------------------------------------------------------------
/assets/lab5_1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab5_1.png


--------------------------------------------------------------------------------
/assets/lab5_2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab5_2.png


--------------------------------------------------------------------------------
/assets/lab6_1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab6_1.png


--------------------------------------------------------------------------------
/assets/lab6_2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab6_2.png


--------------------------------------------------------------------------------
/assets/lab6_3.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab6_3.png


--------------------------------------------------------------------------------
/assets/lab6_4.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab6_4.png


--------------------------------------------------------------------------------
/assets/lab6_5.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab6_5.png


--------------------------------------------------------------------------------
/assets/lab6_6.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/langgraph-agents-with-amazon-bedrock/HEAD/assets/lab6_6.png


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | /.DS_Store
2 | /__pycache__
3 | *__pycache__*
4 | *checkpoints.db
5 | .ipynb_checkpoints/
6 | .virtual_documents/
7 | .env


--------------------------------------------------------------------------------
/CODE_OF_CONDUCT.md:
--------------------------------------------------------------------------------
1 | ## Code of Conduct
2 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
3 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
4 | opensource-codeofconduct@amazon.com with any additional questions or comments.
5 | 


--------------------------------------------------------------------------------
/Lab_1/README.md:
--------------------------------------------------------------------------------
 1 | # Lab 1: Building a ReAct Agent from Scratch
 2 | 
 3 | Welcome to this section on building a ReAct (Reasoning and Acting) agent from scratch using Amazon Bedrock and Anthropic's Claude model. 
 4 | 
 5 | The ReAct approach aims to combine reasoning (e.g. chain-of-thought prompting) and acting (e.g. action plan generation) capabilities of large language models (LLMs) in an interleaved manner. It was proposed in the paper [ReAct: Synergizing Reasoning and Acting in Language Models](https://arxiv.org/abs/2210.03629).
 6 | 
 7 | ![ReAct approach](../assets/lab1_1.png)
 8 | 
 9 | The structure of the agent within the section looks like the picture below, with components:
10 | - User Input
11 | - System prompt
12 | - Language Model
13 | - Identified action
14 | - Tool to be executed by the action
15 | 
16 | ![Agent structure](../assets/lab1_2.png)
17 | 
18 | Let's begin!
19 | 


--------------------------------------------------------------------------------
/Lab_5/README.md:
--------------------------------------------------------------------------------
 1 | # Lab 5: Human in the Loop
 2 | 
 3 | Welcome to this advanced section on implementing human-in-the-loop interactions with AI agents using LangGraph, Anthropic's Claude model on Amazon Bedrock. 
 4 | This section is designed for solution architects and data scientists looking to build sophisticated AI systems with human oversight capabilities.
 5 | 
 6 | Before we dive into the code, let's understand the big picture. Imagine you're architecting a mission-critical AI system for a large enterprise. This system needs to make complex decisions, but also requires human oversight for accountability and error correction. This is where human-in-the-loop AI comes in.
 7 | 
 8 | Think of our system as a collaborative decision-making pipeline between AI and human experts. LangGraph provides the framework for this pipeline, the foundation models on Amazon Bedrock bring the AI capabilities, which also brings the scalable cloud infrastructure.
 9 | 
10 | Let's dive in!
11 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT No Attribution
 2 | 
 3 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of
 6 | this software and associated documentation files (the "Software"), to deal in
 7 | the Software without restriction, including without limitation the rights to
 8 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 9 | the Software, and to permit persons to whom the Software is furnished to do so.
10 | 
11 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
12 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
13 | FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
14 | COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
15 | IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
16 | CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
17 | 
18 | 


--------------------------------------------------------------------------------
/Lab_4/README.md:
--------------------------------------------------------------------------------
 1 | # Lab 4: Persistence and Streaming
 2 | 
 3 | Welcome to this section, where we will explore two critical concepts emerging in the development of sophisticated AI agents, particularly those designed for extended operations: persistence and streaming. These concepts are fundamental to creating robust, responsive, and scalable AI solutions.
 4 | 
 5 | Persistence enables the preservation of an agent's state at specific points in time, facilitating seamless resumption of operations in subsequent interactions. This capability is particularly crucial for long-running applications where continuity is paramount. Streaming, on the other hand, provides real-time visibility into the agent's operations, emitting a series of signals that indicate the agent's current actions and thought processes. This feature is invaluable for monitoring and understanding the behavior of long-running applications.
 6 | 
 7 | In this section, we will explore these concepts in depth utilizing Anthropic's Claude model on Amazon Bedrock as our foundation.
 8 | 
 9 | 
10 | Let's dive in!


--------------------------------------------------------------------------------
/Lab_2/README.md:
--------------------------------------------------------------------------------
 1 | # Lab 2: LangGraph Components
 2 | 
 3 | Welcome to this section on implementing intelligent agents using LangGraph, Anthropic's Claude model, and Amazon Bedrock. In our previous session, we took a hands-on approach by implementing a REACT (Reasoning and Acting) agent from scratch. This exercise gave us a fundamental understanding of how these agents operate and make decisions.
 4 | 
 5 | Now, we're taking a significant step forward. We'll be leveraging LangGraph, a powerful framework for creating structured conversations and decision-making flows with language models. Instead of building everything from the ground up, we'll use LangGraph to create a more sophisticated and flexible agent. We'll also be upgrading our language model, moving from a basic implementation to using Anthropic's Claude model via Amazon Bedrock.
 6 | 
 7 | This section will demonstrate how to create a robust, scalable agent capable of handling complex, multi-step queries. You'll see how the structure provided by LangGraph enhances our agent's capabilities, and how the advanced features of Claude and the scalability of Amazon Bedrock take our implementation to the next level.
 8 | By the end of this section, you'll have a clear understanding of how to use these tools to create state-of-the-art AI agents for real-world applications. 
 9 | 
10 | In essence, LangGraph is an extension of LangChain that supports graphs. Single and Multi-agent workflows are described and represented as graphs. The tool allows for well controlled "flows".
11 | 
12 | The components of a LangGraph graph are depicted in the image below.
13 | 
14 | ![LangGraph components](../assets/lab2_1.png)
15 | 
16 | Let's dive in!
17 | 


--------------------------------------------------------------------------------
/Lab_3/README.md:
--------------------------------------------------------------------------------
 1 | # Lab 3: Agentic Search Tools
 2 | 
 3 | Welcome to this section introducing the concept of agentic search, a more advanced form of search used by AI agents. Unlike standard search or zero-shot learning, agentic search allows agents to access dynamic data, provide sources for information, and handle complex queries by breaking them down into sub-questions. The process involves understanding the query, selecting the best source for information, extracting relevant data, and filtering out less important details. This approach helps reduce hallucinations and improves human-computer interaction.
 4 | 
 5 | The section demonstrates the difference between regular search tools and agentic search tools through practical examples. Using a weather query for San Francisco, it shows how a regular search (using `DuckDuckGo`) provides links that require further processing to extract useful information. In contrast, the agentic search tool (using `Tavily`) returns structured data in JSON format, which is ideal for AI agents to process. 
 6 | 
 7 | A basic example of an agentic search tool implementation is shown in the picture below.
 8 | 
 9 | ![Agentic Search Tool](../assets/lab3_1.png)
10 | 
11 | When the agent decides to send the query to the search tool, it would first work on understanding the question and divide it to sub-questions if needed. This is an important point because it allows agents to handle complex queries. Then, for each subquery, the search tool will have to find the best source, choosing from multiple integrations. For example, if an agent would ask "How is the weather in San Francisco?" The search tool should use the weather API for best results. The job doesn't end with finding the correct source. The search tool would then have to extract only the relevant information to the subquery. 
12 | 
13 | A basic implementation of this can be achieved through a process of chunking the source, and run a quick vector search to retrieve the top-K chunks. After retrieving the data from its source, the search tool would then score the results and filter out the less relevant information. This implementation is shown in the picture below.
14 | 
15 | ![Agentic Search Tool](../assets/lab3_2.png)
16 | 
17 | Let's dive in!


--------------------------------------------------------------------------------
/utils/utils.py:
--------------------------------------------------------------------------------
 1 | import logging
 2 | import os
 3 | import pprint
 4 | 
 5 | import boto3
 6 | 
 7 | 
 8 | def set_logger(log_level: str = "INFO") -> object:
 9 |     log_level = os.environ.get("LOG_LEVEL", log_level).strip().upper()
10 |     logging.basicConfig(format="[%(asctime)s] p%(process)s {%(filename)s:%(lineno)d} %(levelname)s - %(message)s")
11 |     logger = logging.getLogger(__name__)
12 |     logger.setLevel(log_level)
13 |     return logger
14 | 
15 | 
16 | logger = set_logger()
17 | 
18 | 
19 | def set_pretty_printer():
20 |     return pprint.PrettyPrinter(indent=2, width=100)
21 | 
22 | 
23 | def get_tavily_api(key: str, region_name: str) -> str:
24 | 
25 |     if not os.path.isfile("../.env"):
26 |         raise Exception('Local environment variable file .env not existing! Please create with the command: cp .env.tmp .env')
27 |     
28 |     tavily_api_prefix = "tvly-"
29 |     if not os.environ[key].startswith(tavily_api_prefix):
30 |         logger.info(f"{key} value not correctly set in the .env file, expected a key to start with \"{tavily_api_prefix}\" but got it starting with \"{os.environ[key][:5]}\". Trying from AWS Secrets Manager.")
31 |         session = boto3.session.Session()
32 |         secrets_manager = session.client(service_name="secretsmanager", region_name=region_name)
33 |         try:
34 |             secret_value = secrets_manager.get_secret_value(SecretId=key)
35 |         except Exception as e:
36 |             logger.error(f"{key} secret couldn't be retrieved correctly from AWS Secrets Manager either! Received error message:\n{e}")
37 |             raise e
38 | 
39 |         logger.info(f"{key} variable correctly retrieved from the AWS Secret Manager.")
40 |         secret_string = secret_value["SecretString"]
41 |         secret = eval(secret_string, {"__builtins__": {}}, {})[key]
42 |         if not secret:
43 |             raise Exception(f"{key} value not correctly set in the AWS Secrets Manager, expected a key to start with \"{tavily_api_prefix}\" but got it starting with \"{os.environ[key][:5]}\".")
44 |         os.environ[key] = secret
45 |     else:
46 |         logger.info(f"{key} variable correctly retrieved from the .env file.")
47 |         secret = os.environ[key]
48 | 
49 |     return secret
50 | 


--------------------------------------------------------------------------------
/pyproject.toml:
--------------------------------------------------------------------------------
  1 | [tool.poetry]
  2 | name = "deeplearningai-langgraph-agents-with-amazon-bedrock"
  3 | version = "0.1.0"
  4 | description = ""
  5 | authors = [
  6 |     "Philipp Kaindl <philikai@amazon.de>",
  7 |     "Luca Perrozzi <lperroz@amazon.ch>",
  8 |     "Gabriel Rodriguez Garcia <garoga@amazon.com>",
  9 |     "Markus Rollwagen <rollwag@amazon.ch>",
 10 | ]
 11 | readme = "README.md"
 12 | package-mode = false
 13 | 
 14 | [virtualenvs]
 15 | create = true
 16 | in-project = true
 17 | 
 18 | [tool.poetry.dependencies]
 19 | python = "^3.10"
 20 | aiofiles = "23.2.1"
 21 | aiohappyeyeballs = "2.3.5"
 22 | aiohttp = "3.10.2"
 23 | aiosignal = "1.3.1"
 24 | altair = "5.3.0"
 25 | annotated-types = "0.7.0"
 26 | anyio = "4.4.0"
 27 | attrs = "24.2.0"
 28 | boto3 = "1.34.155"
 29 | botocore = "1.34.155"
 30 | certifi = "2024.7.4"
 31 | charset-normalizer = "3.3.2"
 32 | click = "8.1.7"
 33 | contourpy = "1.2.1"
 34 | cycler = "0.12.1"
 35 | distro = "1.9.0"
 36 | fastapi = "0.115.2"
 37 | ffmpy = "0.4.0"
 38 | filelock = "3.15.4"
 39 | fonttools = "4.53.1"
 40 | frozenlist = "1.4.1"
 41 | fsspec = "2024.6.1"
 42 | gradio = "^5.23.0"                      # allow minor and patch version upgrades
 43 | gradio-client = "^1.8.0"
 44 | greenlet = "3.0.3"
 45 | h11 = "0.14.0"
 46 | httpcore = "1.0.5"
 47 | httpx = "0.27.0"
 48 | huggingface-hub = "0.28.1"
 49 | idna = "3.7"
 50 | importlib-resources = "6.4.0"
 51 | jinja2 = "3.1.6"
 52 | jiter = "0.5.0"
 53 | jmespath = "1.0.1"
 54 | jsonpatch = "1.33"
 55 | jsonpointer = "3.0.0"
 56 | jsonschema = "4.23.0"
 57 | jsonschema-specifications = "2023.12.1"
 58 | kiwisolver = "1.4.5"
 59 | langchain = "0.2.12"
 60 | langchain-aws = "0.1.15"
 61 | langchain-core = "0.2.28"
 62 | langchain-openai = "0.1.20"
 63 | langchain-text-splitters = "0.2.2"
 64 | langgraph = "0.0.53"
 65 | langsmith = "0.1.98"
 66 | markdown-it-py = "3.0.0"
 67 | markupsafe = "2.1.5"
 68 | matplotlib = "3.9.1.post1"
 69 | mdurl = "0.1.2"
 70 | multidict = "6.0.5"
 71 | numpy = "1.26.4"
 72 | openai = "1.40.0"
 73 | orjson = "3.10.6"
 74 | packaging = "24.1"
 75 | pandas = "2.2.2"
 76 | pillow = "10.4.0"
 77 | pydantic = "2.8.2"
 78 | pydantic-core = "2.20.1"
 79 | pydub = "0.25.1"
 80 | pygments = "2.18.0"
 81 | pyparsing = "3.1.2"
 82 | python-dateutil = "2.9.0.post0"
 83 | python-dotenv = "1.0.1"
 84 | python-multipart = "0.0.18"
 85 | pytz = "2024.1"
 86 | pyyaml = "6.0.2"
 87 | referencing = "0.35.1"
 88 | regex = "2024.7.24"
 89 | requests = "2.32.3"
 90 | rich = "13.7.1"
 91 | rpds-py = "0.20.0"
 92 | ruff = "0.11.1"
 93 | s3transfer = "0.10.2"
 94 | semantic-version = "2.10.0"
 95 | shellingham = "1.5.4"
 96 | six = "1.16.0"
 97 | sniffio = "1.3.1"
 98 | sqlalchemy = "2.0.32"
 99 | starlette = "0.40.0"
100 | tavily-python = "0.3.5"
101 | tenacity = "8.5.0"
102 | tiktoken = "0.7.0"
103 | tomlkit = "0.12.0"
104 | toolz = "0.12.1"
105 | tqdm = "4.66.5"
106 | typer = "0.12.3"
107 | typing-extensions = "4.12.2"
108 | tzdata = "2024.1"
109 | urllib3 = "2.2.2"
110 | uuid6 = "2024.7.10"
111 | uvicorn = "0.30.5"
112 | websockets = "11.0.3"
113 | yarl = "1.9.4"
114 | ipykernel = "6.29.5"
115 | jupyter = "1.0.0"
116 | langchain-community = "0.2.11"
117 | pygraphviz = ">=1.13,<2.0"
118 | duckduckgo-search = "^6.2.6"
119 | aiosqlite = "^0.20.0"
120 | 
121 | 
122 | [build-system]
123 | requires = ["poetry-core"]
124 | build-backend = "poetry.core.masonry.api"
125 | 


--------------------------------------------------------------------------------
/CONTRIBUTING.md:
--------------------------------------------------------------------------------
 1 | # Contributing Guidelines
 2 | 
 3 | Thank you for your interest in contributing to our project. Whether it's a bug report, new feature, correction, or additional
 4 | documentation, we greatly value feedback and contributions from our community.
 5 | 
 6 | Please read through this document before submitting any issues or pull requests to ensure we have all the necessary
 7 | information to effectively respond to your bug report or contribution.
 8 | 
 9 | 
10 | ## Reporting Bugs/Feature Requests
11 | 
12 | We welcome you to use the GitHub issue tracker to report bugs or suggest features.
13 | 
14 | When filing an issue, please check existing open, or recently closed, issues to make sure somebody else hasn't already
15 | reported the issue. Please try to include as much information as you can. Details like these are incredibly useful:
16 | 
17 | * A reproducible test case or series of steps
18 | * The version of our code being used
19 | * Any modifications you've made relevant to the bug
20 | * Anything unusual about your environment or deployment
21 | 
22 | 
23 | ## Contributing via Pull Requests
24 | Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that:
25 | 
26 | 1. You are working against the latest source on the *main* branch.
27 | 2. You check existing open, and recently merged, pull requests to make sure someone else hasn't addressed the problem already.
28 | 3. You open an issue to discuss any significant work - we would hate for your time to be wasted.
29 | 
30 | To send us a pull request, please:
31 | 
32 | 1. Fork the repository.
33 | 2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change.
34 | 3. Ensure local tests pass.
35 | 4. Commit to your fork using clear commit messages.
36 | 5. Send us a pull request, answering any default questions in the pull request interface.
37 | 6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation.
38 | 
39 | GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and
40 | [creating a pull request](https://help.github.com/articles/creating-a-pull-request/).
41 | 
42 | 
43 | ## Finding contributions to work on
44 | Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' issues is a great place to start.
45 | 
46 | 
47 | ## Code of Conduct
48 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
49 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
50 | opensource-codeofconduct@amazon.com with any additional questions or comments.
51 | 
52 | 
53 | ## Security issue notifications
54 | If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue.
55 | 
56 | 
57 | ## Licensing
58 | 
59 | See the [LICENSE](LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution.
60 | 


--------------------------------------------------------------------------------
/Lab_6/README.md:
--------------------------------------------------------------------------------
 1 | # Lab 6: Essay Writer
 2 | 
 3 | In this section, we'll embark on an exciting project that brings together various concepts we've learned so far: we're going to build an AI-powered Essay Writer. This project will demonstrate how to create a more complex, multi-step AI agent using Amazon Bedrock and Anthropic's Claude model.
 4 | 
 5 | ## What You'll Learn
 6 | 
 7 | 1. How to structure a multi-agent system for a complex task
 8 | 2. Implementing a state machine using LangGraph
 9 | 3. Integrating external research capabilities with the Tavily API
10 | 4. Using Claude on Amazon Bedrock for various subtasks within the system
11 | 5. Building a simple GUI to interact with your AI Essay Writer
12 | 
13 | ## Project Overview
14 | 
15 | Our Essay Writer will work through several stages:
16 | 
17 | 1. Planning: Generate an outline for the essay
18 | 2. Research: Use Tavily to gather relevant information
19 | 3. Writing: Generate a draft based on the plan and research
20 | 4. Reflection: Critique the current draft
21 | 5. Iteration: Revise the essay based on the critique
22 | 
23 | The flow of the agent is represented in the picture below
24 | 
25 | ![Essay Writer Flow](../assets/lab6_1.png)
26 | 
27 | 
28 | By the end of this section, you'll have a functional AI Essay Writer that can generate, critique, and refine essays on various topics. This project will give you hands-on experience in building a practical, multi-step AI application using state-of-the-art language models and tools.
29 | 
30 | Let's dive in and start building our AI Essay Writer!
31 | 
32 | # Conclusion
33 | 
34 | At the end of this workshop, you should have a good idea on how to build your own agents. Before to conclude, we will cover some other agent architectures that you should know about. 
35 | 
36 | - [Multi-Agent Collaboration](https://github.com/langchain-ai/langgraph/blob/main/docs/docs/tutorials/multi_agent/multi-agent-collaboration.ipynb):
37 | In a multi-agent architecture, multiple agents work collaboratively on a shared state. These agents could be a combination of prompts, language models, and tools, each contributing their capabilities. The key aspect is that they all operate on and pass around the same shared state, allowing them to build upon each other's work iteratively.
38 | 
39 | ![Multi-Agent Architecture](../assets/lab6_2.png)
40 | 
41 | - [Supervisor Agent](https://github.com/langchain-ai/langgraph/blob/main/docs/docs/tutorials/multi_agent/agent_supervisor.ipynb):
42 | A supervisor agent architecture involves a central supervisor agent that coordinates and manages the inputs and outputs of various sub-agents. The supervisor determines the specific inputs and tasks for each sub-agent, which can have their own internal states and processes. Unlike the multi-agent approach, there is no single shared state, but rather the supervisor orchestrates the flow of information between the sub-agents.
43 | 
44 | ![Supervisor Agent Architecture](../assets/lab6_3.png)
45 | 
46 | - [Flow Engineering](https://arxiv.org/abs/2401.08500):
47 | Flow engineering, as described in the AlphaCode paper, refers to designing architectures with a directed flow of information, punctuated by iterative loops at key points. This approach combines a linear pipeline with cyclical iterations, tailored to specific problem domains like coding. The goal is to engineer the optimal flow of information and decision-making for the task at hand.
48 | 
49 | ![Flow Engineering Architecture](../assets/lab6_4.png)
50 | 
51 | - [Plan and Execute](https://github.com/langchain-ai/langgraph/blob/main/docs/docs/tutorials/plan-and-execute/plan-and-execute.ipynb):
52 | The plan and execute paradigm involves an explicit planning phase followed by an execution phase. First, a plan or set of steps is generated, which sub-agents then execute sequentially. After each step, the plan may be updated or revised based on the results, and the process continues until the plan is completed or a replanning is required.
53 | 
54 | ![Plan and Execute Architecture](../assets/lab6_5.png)
55 | 
56 | - [Language Agent Tree Search](https://github.com/langchain-ai/langgraph/blob/main/docs/docs/tutorials/lats/lats.ipynb):
57 | This approach performs a tree search over possible action states, generating actions, reflecting on them, and backpropagating information to update parent nodes. It allows for jumping back to previous states in the tree, enabling the agent to explore different action paths while leveraging information from earlier reflections and updates.
58 | 
59 | ![Language Agent Tree Search Architecture](../assets/lab6_6.png)
60 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # LangGraph Agents with Amazon Bedrock
  2 | 
  3 | This repository contains a workshop adapted from the course [AI Agents in LangGraph](https://www.deeplearning.ai/short-courses/ai-agents-in-langgraph/) 
  4 | created by [Harrison Chase](https://www.linkedin.com/in/harrison-chase-961287118) (Co-Founder and CEO of [LangChain](https://www.langchain.com/)) and [Rotem Weiss](https://www.linkedin.com/in/rotem-weiss) (Co-founder and CEO of [Tavily](https://tavily.com/)), and hosted on [DeepLearning.AI](https://www.deeplearning.ai/).
  5 | The original content is used with the consent of the authors.
  6 | 
  7 | This workshop is also avalailable in AWS Workshop Studio [here](https://catalog.us-east-1.prod.workshops.aws/workshops/9bc28f51-d7c3-468b-ba41-72667f3273f1/en-US).
  8 | 
  9 | Make sure to read and follow this README before you go through the material to ensure a smooth experience.
 10 | 
 11 | ## Outline
 12 | 
 13 | The workshop:
 14 | - explores the latest advancements in AI agents and agentic workflows, leveraging improvements in function calling LLMs and specialized tools like agentic search
 15 | - utilizes LangChain's updated support for agentic workflows and introduces LangGraph, an extension for building complex agent behaviors
 16 | - provides insights into key design patterns in agentic workflows including *planning, tool use, reflection, multi-agent communication, memory*
 17 | 
 18 | The material is divided in six Jupyter Notebooks Labs that will help you understand the LangGraph framework, its underlying concepts, and how to use it with Amazon Bedrock:
 19 | 
 20 | - Lab 1: [Building a ReAct Agent from Scratch](Lab_1/)
 21 |     - Build a basic ReAct agent from scratch using Python and an LLM, implementing a loop of reasoning and acting to solve tasks through tool usage and observation
 22 | - Lab 2: [LangGraph Components](Lab_2/)
 23 |     - Introduction to LangGraph, a tool for implementing agents with cyclic graphs, demonstrating how to create a more structured and controllable agent using components like nodes, edges, and state management
 24 | - Lab 3: [Agentic Search Tools](Lab_3/)
 25 |     - Introduction to Agentic search tools, enhancing AI agents' capabilities by providing structured, relevant data from dynamic sources, improving accuracy and reducing hallucinations
 26 | - Lab 4: [Persistence and Streaming](Lab_4/)
 27 |     - Persistence and streaming are crucial for long-running agent tasks, enabling state preservation, resumption of conversations, and real-time visibility into agent actions and outputs
 28 | - Lab 5: [Human in the Loop](Lab_5/)
 29 |     - Advanced human-in-the-loop interaction patterns in LangGraph, including adding breaks, modifying states, time travel, and manual state updates for better control and interaction with AI agents
 30 | - Lab 6: [Essay Writer](Lab_6/)
 31 |     - Build an AI essay writer using a multi-step process involving planning, research, writing, reflection, and revision, implemented as a graph of interconnected agents
 32 | 
 33 | If this is your first time working with LangGraph, we recommend to refer to the [original course](https://www.deeplearning.ai/short-courses/ai-agents-in-langgraph/) for detailed video explanations.
 34 | 
 35 | Let's get started with the setup of the environment.
 36 | 
 37 | ## Setup your virtual environment
 38 | 
 39 | This instructions are meant to be used locally with [AWS authentication](https://docs.aws.amazon.com/cli/v1/userguide/cli-authentication-short-term.html), as well as within an [Amazon SageMaker JupyterLab](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated-jl.html) or [Amazon SageMaker Code Editor](https://docs.aws.amazon.com/sagemaker/latest/dg/code-editor.html) instance.
 40 | 
 41 | The course requires `Python >=3.10` (to install, visit this link: https://www.python.org/downloads/)
 42 | 
 43 | ### 1. Download the repository
 44 | 
 45 | ```
 46 | git clone https://github.com/aws-samples/langgraph-agents-with-amazon-bedrock.git
 47 | ```
 48 | 
 49 | ### 2. Install OS dependencies (Ubuntu/Debian)
 50 | 
 51 | ```
 52 | sudo apt update
 53 | sudo apt-get install graphviz graphviz-dev python3-dev
 54 | pip install pipx
 55 | pipx install poetry
 56 | pipx ensurepath
 57 | source ~/.bashrc
 58 | ```
 59 | 
 60 | Installation commands for other OS can be found here: https://pygraphviz.github.io/documentation/stable/install.html
 61 | 
 62 | ### 3. Create a virtual environment and install python dependencies
 63 | 
 64 | ```
 65 | cd langgraph-agents-with-amazon-bedrock
 66 | export POETRY_VIRTUALENVS_PATH="$PWD/.venv"
 67 | export INITIAL_WORKING_DIRECTORY=$(pwd)
 68 | poetry shell
 69 | ```
 70 | 
 71 | ```
 72 | cd $INITIAL_WORKING_DIRECTORY
 73 | poetry install
 74 | ```
 75 | 
 76 | ### 4. Add the kernel to the Jupyter Notebook server
 77 | The newly created python environment needs to be added to the list of available kernels of the Jupyter Notebook server.
 78 | This is possible from within the poetry environment with the command:
 79 | ```
 80 | poetry run python -m ipykernel install --user --name agents-dev-env
 81 | ```
 82 | The kernel might not appear right away in the list, in that case a refresh of the list will be needed.
 83 | 
 84 | ### 5. Create and set your Tavily API key
 85 | 
 86 | Head over to https://app.tavily.com/home and create an API KEY for free. 
 87 | 
 88 | ### 6. Setup the local environment variables
 89 | 
 90 | Create a personal copy of the temporary environment file [env.tmp](env.tmp) with the name `.env`, which is already added to the [.gitignore](.gitignore) to avoid committing personal information.
 91 | ```
 92 | cp env.tmp .env
 93 | ```
 94 | You can edit the preferred region to use Amazon Bedrock inside the `.env` file, if needed (the default is `us-east-1`).
 95 | 
 96 | ### 7. Store the Tavily API key
 97 | You have two options to store the Tavily API key: 
 98 | 
 99 | 1. Copy the Tavily API key inside the `.env` file. This option will be always checked first.
100 | 
101 | 2. [Create a new secret in AWS Secrets Manager](https://docs.aws.amazon.com/secretsmanager/latest/userguide/create_secret.html) with the name "TAVILY_API_KEY", retrieve the secret `arn` by clicking on it, and [add an inline policy](https://docs.aws.amazon.com/IAM/latest/UserGuide/access_policies_manage-attach-detach.html#add-policies-console) with the [permission to read the secret](https://docs.aws.amazon.com/secretsmanager/latest/userguide/auth-and-access_examples.html#auth-and-access_examples_read) to your [SageMaker execution role](https://docs.aws.amazon.com/sagemaker/latest/dg/domain-user-profile-view-describe.html) replacing the copied `arn` in the example below.
102 | ```
103 | {
104 |     "Version": "2012-10-17",
105 |     "Statement": [
106 |         {
107 |             "Effect": "Allow",
108 |             "Action": "secretsmanager:GetSecretValue",
109 |             "Resource": "arn:aws:secretsmanager:<Region>:<AccountId>:secret:SecretName-6RandomCharacters"
110 |         }
111 |     ]
112 | }
113 | ```
114 | 
115 | You are all set! Make sure to select the freshly created `agents-dev-env` kernel for each notebook.
116 | 
117 | # Additional resources
118 | 
119 | - [Amazon Bedrock User Guide](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html)
120 | - [LangChain documentation](https://python.langchain.com/v0.2/docs/introduction/)
121 | - [LangGraph github repository](https://github.com/langchain-ai/langgraph)
122 | - [LangSmith Prompt hub](https://smith.langchain.com/hub)
123 | 


--------------------------------------------------------------------------------
/Lab_3/Lab_3.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |  "cells": [
  3 |   {
  4 |    "cell_type": "markdown",
  5 |    "id": "2544dd96-a13c-42df-ab63-c427e68fd41e",
  6 |    "metadata": {},
  7 |    "source": [
  8 |     "# Lab 3: Agentic Search Tools\n"
  9 |    ]
 10 |   },
 11 |   {
 12 |    "cell_type": "markdown",
 13 |    "id": "575640a8",
 14 |    "metadata": {},
 15 |    "source": [
 16 |     "## Setting Up the Environment"
 17 |    ]
 18 |   },
 19 |   {
 20 |    "cell_type": "code",
 21 |    "execution_count": null,
 22 |    "id": "d0168aee-bce9-4d60-b827-f86a88187e31",
 23 |    "metadata": {
 24 |     "height": 198
 25 |    },
 26 |    "outputs": [],
 27 |    "source": [
 28 |     "from dotenv import load_dotenv\n",
 29 |     "import os\n",
 30 |     "import sys\n",
 31 |     "import json, re\n",
 32 |     "import pprint\n",
 33 |     "import boto3\n",
 34 |     "from botocore.client import Config\n",
 35 |     "import warnings\n",
 36 |     "\n",
 37 |     "warnings.filterwarnings(\"ignore\")\n",
 38 |     "import logging\n",
 39 |     "\n",
 40 |     "# import local modules\n",
 41 |     "dir_current = os.path.abspath(\"\")\n",
 42 |     "dir_parent = os.path.dirname(dir_current)\n",
 43 |     "if dir_parent not in sys.path:\n",
 44 |     "    sys.path.append(dir_parent)\n",
 45 |     "from utils import utils\n",
 46 |     "\n",
 47 |     "bedrock_config = Config(\n",
 48 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
 49 |     ")\n",
 50 |     "\n",
 51 |     "# Set basic configs\n",
 52 |     "logger = utils.set_logger()\n",
 53 |     "pp = utils.set_pretty_printer()\n",
 54 |     "\n",
 55 |     "# Load environment variables from .env file or Secret Manager\n",
 56 |     "_ = load_dotenv(\"../.env\")\n",
 57 |     "aws_region = os.getenv(\"AWS_REGION\")\n",
 58 |     "tavily_ai_api_key = utils.get_tavily_api(\"TAVILY_API_KEY\", aws_region)\n",
 59 |     "\n",
 60 |     "# Set bedrock configs\n",
 61 |     "bedrock_config = Config(\n",
 62 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
 63 |     ")\n",
 64 |     "\n",
 65 |     "# Create a bedrock runtime client\n",
 66 |     "bedrock_rt = boto3.client(\n",
 67 |     "    \"bedrock-runtime\", region_name=aws_region, config=bedrock_config\n",
 68 |     ")\n",
 69 |     "\n",
 70 |     "# Create a bedrock client to check available models\n",
 71 |     "bedrock = boto3.client(\"bedrock\", region_name=aws_region, config=bedrock_config)\n"
 72 |    ]
 73 |   },
 74 |   {
 75 |    "cell_type": "code",
 76 |    "execution_count": null,
 77 |    "id": "a2ba84ec-c172-4de7-ac55-e3158a531b23",
 78 |    "metadata": {
 79 |     "height": 147
 80 |    },
 81 |    "outputs": [],
 82 |    "source": [
 83 |     "from tavily import TavilyClient\n",
 84 |     "\n",
 85 |     "client = TavilyClient(api_key=tavily_ai_api_key)\n",
 86 |     "# run search\n",
 87 |     "result = client.search(\"What is in Nvidia's new Blackwell GPU?\", include_answer=True)\n",
 88 |     "\n",
 89 |     "# print the answer\n",
 90 |     "result[\"answer\"]"
 91 |    ]
 92 |   },
 93 |   {
 94 |    "cell_type": "markdown",
 95 |    "id": "d7f3f33c-c7bd-4b7f-9616-b65eef104514",
 96 |    "metadata": {},
 97 |    "source": [
 98 |     "## Regular search\n"
 99 |    ]
100 |   },
101 |   {
102 |    "cell_type": "code",
103 |    "execution_count": null,
104 |    "id": "876d5092-b8ef-4e38-b4d7-0e80c609bf7a",
105 |    "metadata": {
106 |     "height": 166
107 |    },
108 |    "outputs": [],
109 |    "source": [
110 |     "# choose location (try to change to your own city!)\n",
111 |     "\n",
112 |     "city = \"San Francisco\"\n",
113 |     "\n",
114 |     "query = f\"\"\"\n",
115 |     "    what is the current weather in {city}?\n",
116 |     "    Should I travel there today?\n",
117 |     "    \"weather.com\"\n",
118 |     "\"\"\""
119 |    ]
120 |   },
121 |   {
122 |    "cell_type": "markdown",
123 |    "id": "2128e40f-11ef-41ed-ac1f-7feeb5546224",
124 |    "metadata": {},
125 |    "source": [
126 |     "> Note: search was modified to return expected results in the event of an exception. High volumes of student traffic sometimes cause rate limit exceptions.\n"
127 |    ]
128 |   },
129 |   {
130 |    "cell_type": "code",
131 |    "execution_count": null,
132 |    "id": "10084a02-2928-4945-9f7c-ad3f5b33caf7",
133 |    "metadata": {
134 |     "height": 402
135 |    },
136 |    "outputs": [],
137 |    "source": [
138 |     "import requests\n",
139 |     "from bs4 import BeautifulSoup\n",
140 |     "from duckduckgo_search import DDGS\n",
141 |     "import re\n",
142 |     "\n",
143 |     "ddg = DDGS()\n",
144 |     "\n",
145 |     "\n",
146 |     "def search(query, max_results=6):\n",
147 |     "    try:\n",
148 |     "        results = ddg.text(query, max_results=max_results)\n",
149 |     "        return [i[\"href\"] for i in results]\n",
150 |     "    except Exception as e:\n",
151 |     "        print(f\"returning previous results due to exception reaching ddg.\")\n",
152 |     "        results = [  # cover case where DDG rate limits due to high traffic volume\n",
153 |     "            \"https://weather.com/weather/today/l/USCA0987:1:US\",\n",
154 |     "            \"https://weather.com/weather/hourbyhour/l/54f9d8baac32496f6b5497b4bf7a277c3e2e6cc5625de69680e6169e7e38e9a8\",\n",
155 |     "        ]\n",
156 |     "        return results\n",
157 |     "\n",
158 |     "\n",
159 |     "for i in search(query):\n",
160 |     "    print(i)"
161 |    ]
162 |   },
163 |   {
164 |    "cell_type": "code",
165 |    "execution_count": null,
166 |    "id": "d31c432d-8e22-412f-b302-961ace0b00bd",
167 |    "metadata": {
168 |     "height": 283
169 |    },
170 |    "outputs": [],
171 |    "source": [
172 |     "def scrape_weather_info(url):\n",
173 |     "    \"\"\"Scrape content from the given URL\"\"\"\n",
174 |     "    if not url:\n",
175 |     "        return \"Weather information could not be found.\"\n",
176 |     "\n",
177 |     "    # fetch data\n",
178 |     "    headers = {\"User-Agent\": \"Mozilla/5.0\"}\n",
179 |     "    response = requests.get(url, headers=headers)\n",
180 |     "    if response.status_code != 200:\n",
181 |     "        return \"Failed to retrieve the webpage.\"\n",
182 |     "\n",
183 |     "    # parse result\n",
184 |     "    soup = BeautifulSoup(response.text, \"html.parser\")\n",
185 |     "    return soup"
186 |    ]
187 |   },
188 |   {
189 |    "cell_type": "markdown",
190 |    "id": "0b67d8ad-a439-4c91-9dfe-7c84998ef644",
191 |    "metadata": {},
192 |    "source": [
193 |     "> Note: This produces a long output, you may want to right click and clear the cell output after you look at it briefly to avoid scrolling past it.\n"
194 |    ]
195 |   },
196 |   {
197 |    "cell_type": "code",
198 |    "execution_count": null,
199 |    "id": "714d1205-f8fc-4912-b148-2a45da99219c",
200 |    "metadata": {
201 |     "height": 164
202 |    },
203 |    "outputs": [],
204 |    "source": [
205 |     "# use DuckDuckGo to find websites and take the first result\n",
206 |     "url = search(query)[0]\n",
207 |     "\n",
208 |     "# scrape first wesbsite\n",
209 |     "soup = scrape_weather_info(url)\n",
210 |     "\n",
211 |     "print(f\"Website: {url}\\n\\n\")\n",
212 |     "print(str(soup.body)[:50000])  # limit long outputs"
213 |    ]
214 |   },
215 |   {
216 |    "cell_type": "code",
217 |    "execution_count": null,
218 |    "id": "6cb3ef4c-58b3-401b-b104-0d51e553d982",
219 |    "metadata": {
220 |     "height": 251
221 |    },
222 |    "outputs": [],
223 |    "source": [
224 |     "# extract text\n",
225 |     "weather_data = []\n",
226 |     "for tag in soup.find_all([\"h1\", \"h2\", \"h3\", \"p\"]):\n",
227 |     "    text = tag.get_text(\" \", strip=True)\n",
228 |     "    weather_data.append(text)\n",
229 |     "\n",
230 |     "# combine all elements into a single string\n",
231 |     "weather_data = \"\\n\".join(weather_data)\n",
232 |     "\n",
233 |     "# remove all spaces from the combined text\n",
234 |     "weather_data = re.sub(r\"\\s+\", \" \", weather_data)\n",
235 |     "\n",
236 |     "print(f\"Website: {url}\\n\\n\")\n",
237 |     "print(weather_data)"
238 |    ]
239 |   },
240 |   {
241 |    "cell_type": "markdown",
242 |    "id": "92db676d-d8d9-4558-8dcf-1b20fcb48e45",
243 |    "metadata": {},
244 |    "source": [
245 |     "## Agentic Search\n"
246 |    ]
247 |   },
248 |   {
249 |    "cell_type": "code",
250 |    "execution_count": null,
251 |    "id": "dc3293b7-a50c-43c8-a022-8975e1e444b8",
252 |    "metadata": {
253 |     "height": 132
254 |    },
255 |    "outputs": [],
256 |    "source": [
257 |     "# run search\n",
258 |     "result = client.search(query, max_results=1)\n",
259 |     "\n",
260 |     "# print first result\n",
261 |     "data = result[\"results\"][0][\"content\"]\n",
262 |     "\n",
263 |     "print(data)"
264 |    ]
265 |   },
266 |   {
267 |    "cell_type": "code",
268 |    "execution_count": null,
269 |    "id": "0722c3d4-4cbf-43bf-81b0-50f634c4ce61",
270 |    "metadata": {
271 |     "height": 266
272 |    },
273 |    "outputs": [],
274 |    "source": [
275 |     "import json\n",
276 |     "from pygments import highlight, lexers, formatters\n",
277 |     "\n",
278 |     "# parse JSON\n",
279 |     "parsed_json = json.loads(data.replace(\"'\", '\"'))\n",
280 |     "\n",
281 |     "# pretty print JSON with syntax highlighting\n",
282 |     "formatted_json = json.dumps(parsed_json, indent=4)\n",
283 |     "colorful_json = highlight(\n",
284 |     "    formatted_json, lexers.JsonLexer(), formatters.TerminalFormatter()\n",
285 |     ")\n",
286 |     "\n",
287 |     "print(colorful_json)"
288 |    ]
289 |   }
290 |  ],
291 |  "metadata": {
292 |   "kernelspec": {
293 |    "display_name": "agents-dev-env",
294 |    "language": "python",
295 |    "name": "agents-dev-env"
296 |   },
297 |   "language_info": {
298 |    "codemirror_mode": {
299 |     "name": "ipython",
300 |     "version": 3
301 |    },
302 |    "file_extension": ".py",
303 |    "mimetype": "text/x-python",
304 |    "name": "python",
305 |    "nbconvert_exporter": "python",
306 |    "pygments_lexer": "ipython3",
307 |    "version": "3.10.14"
308 |   }
309 |  },
310 |  "nbformat": 4,
311 |  "nbformat_minor": 5
312 | }
313 | 


--------------------------------------------------------------------------------
/Lab_2/Lab_2.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |  "cells": [
  3 |   {
  4 |    "cell_type": "markdown",
  5 |    "id": "21fa2e13-567d-4509-9023-c99fb230f31f",
  6 |    "metadata": {},
  7 |    "source": [
  8 |     "# Lab 2: LangGraph Components\n"
  9 |    ]
 10 |   },
 11 |   {
 12 |    "cell_type": "markdown",
 13 |    "id": "fcd7ea45",
 14 |    "metadata": {},
 15 |    "source": [
 16 |     "## Setting Up the Environment"
 17 |    ]
 18 |   },
 19 |   {
 20 |    "cell_type": "code",
 21 |    "execution_count": null,
 22 |    "id": "f5762271-8736-4e94-9444-8c92bd0e8074",
 23 |    "metadata": {
 24 |     "height": 47
 25 |    },
 26 |    "outputs": [],
 27 |    "source": [
 28 |     "from dotenv import load_dotenv\n",
 29 |     "import json\n",
 30 |     "import os\n",
 31 |     "import re\n",
 32 |     "import sys\n",
 33 |     "import warnings\n",
 34 |     "\n",
 35 |     "import boto3\n",
 36 |     "from botocore.config import Config\n",
 37 |     "\n",
 38 |     "warnings.filterwarnings(\"ignore\")\n",
 39 |     "import logging\n",
 40 |     "\n",
 41 |     "# import local modules\n",
 42 |     "dir_current = os.path.abspath(\"\")\n",
 43 |     "dir_parent = os.path.dirname(dir_current)\n",
 44 |     "if dir_parent not in sys.path:\n",
 45 |     "    sys.path.append(dir_parent)\n",
 46 |     "from utils import utils\n",
 47 |     "\n",
 48 |     "# Set basic configs\n",
 49 |     "logger = utils.set_logger()\n",
 50 |     "pp = utils.set_pretty_printer()\n",
 51 |     "\n",
 52 |     "# Load environment variables from .env file or Secret Manager\n",
 53 |     "_ = load_dotenv(\"../.env\")\n",
 54 |     "aws_region = os.getenv(\"AWS_REGION\")\n",
 55 |     "tavily_ai_api_key = utils.get_tavily_api(\"TAVILY_API_KEY\", aws_region)\n",
 56 |     "\n",
 57 |     "# Set bedrock configs\n",
 58 |     "bedrock_config = Config(\n",
 59 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
 60 |     ")\n",
 61 |     "\n",
 62 |     "# Create a bedrock runtime client\n",
 63 |     "bedrock_rt = boto3.client(\n",
 64 |     "    \"bedrock-runtime\", region_name=aws_region, config=bedrock_config\n",
 65 |     ")\n",
 66 |     "\n",
 67 |     "# Create a bedrock client to check available models\n",
 68 |     "bedrock = boto3.client(\"bedrock\", region_name=aws_region, config=bedrock_config)\n"
 69 |    ]
 70 |   },
 71 |   {
 72 |    "cell_type": "markdown",
 73 |    "id": "d6a89cc9",
 74 |    "metadata": {},
 75 |    "source": [
 76 |     "## LangGraph as a State Machine\n",
 77 |     "\n",
 78 |     "For solution architects familiar with system design, LangGraph can be thought of as a state machine for language models. Just as a state machine in software engineering defines a set of states and transitions between them, LangGraph allows us to define the states of our conversation (represented by nodes) and the transitions between them (represented by edges).\n",
 79 |     "\n",
 80 |     "**Analogy**: Think of LangGraph as a traffic control system for a smart city. Each intersection (node) represents a decision point, and the roads between them (edges) represent possible paths. The traffic lights (conditional edges) determine which path to take based on current conditions. In our case, the \"traffic\" is the flow of information and decisions in our AI agent.\n"
 81 |    ]
 82 |   },
 83 |   {
 84 |    "cell_type": "code",
 85 |    "execution_count": null,
 86 |    "id": "d0168aee-bce9-4d60-b827-f86a88187e31",
 87 |    "metadata": {
 88 |     "height": 115
 89 |    },
 90 |    "outputs": [],
 91 |    "source": [
 92 |     "import operator\n",
 93 |     "from typing import Annotated, TypedDict\n",
 94 |     "\n",
 95 |     "from langchain_aws import ChatBedrockConverse\n",
 96 |     "from langchain_community.tools.tavily_search import TavilySearchResults\n",
 97 |     "from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage, ToolMessage\n",
 98 |     "from langgraph.graph import END, StateGraph"
 99 |    ]
100 |   },
101 |   {
102 |    "cell_type": "code",
103 |    "execution_count": null,
104 |    "id": "2589c5b6-6cc2-4594-9a17-dccdcf676054",
105 |    "metadata": {
106 |     "height": 64
107 |    },
108 |    "outputs": [],
109 |    "source": [
110 |     "tool = TavilySearchResults(max_results=4)  # increased number of results\n",
111 |     "print(type(tool))\n",
112 |     "print(tool.name)"
113 |    ]
114 |   },
115 |   {
116 |    "cell_type": "markdown",
117 |    "id": "e196c186-af55-4f2d-b569-b7d63a859304",
118 |    "metadata": {},
119 |    "source": [
120 |     "> If you are not familiar with python typing annotation, you can refer to the [python documents](https://docs.python.org/3/library/typing.html).\n"
121 |    ]
122 |   },
123 |   {
124 |    "cell_type": "markdown",
125 |    "id": "5e40395e",
126 |    "metadata": {},
127 |    "source": [
128 |     "## The Agent State Concept\n",
129 |     "\n",
130 |     "The AgentState class is crucial for maintaining context throughout the conversation. For data scientists, this can be compared to maintaining state in a recurrent neural network.\n",
131 |     "\n",
132 |     "**Analogy**: Think of the AgentState as a sophisticated notepad. As you brainstorm ideas (process queries), you jot down key points (messages). This notepad doesn't just record; it has a special property where new notes (messages) are seamlessly integrated with existing ones, maintaining a coherent flow of thought.\n",
133 |     "At the same time, you can always go back in time and rewrite some parts of it - this is what we call \"time-travel\".\n"
134 |    ]
135 |   },
136 |   {
137 |    "cell_type": "code",
138 |    "execution_count": null,
139 |    "id": "a2ba84ec-c172-4de7-ac55-e3158a531b23",
140 |    "metadata": {
141 |     "height": 47
142 |    },
143 |    "outputs": [],
144 |    "source": [
145 |     "class AgentState(TypedDict):\n",
146 |     "    messages: Annotated[list[AnyMessage], operator.add]"
147 |    ]
148 |   },
149 |   {
150 |    "cell_type": "markdown",
151 |    "id": "74c7ba73-e603-453b-b06f-5db92c567b19",
152 |    "metadata": {},
153 |    "source": [
154 |     "> Note: in `take_action` below, some logic was added to cover the case that the LLM returned a non-existent tool name.\n",
155 |     "\n",
156 |     "```python\n",
157 |     "if not t[\"name\"] in self.tools:  # check for bad tool name from LLM\n",
158 |     "    print(\"\\n ....bad tool name....\")\n",
159 |     "    result = \"bad tool name, retry\"  # instruct LLM to retry if bad\n",
160 |     "\n",
161 |     "```\n"
162 |    ]
163 |   },
164 |   {
165 |    "cell_type": "code",
166 |    "execution_count": null,
167 |    "id": "876d5092-b8ef-4e38-b4d7-0e80c609bf7a",
168 |    "metadata": {
169 |     "height": 727
170 |    },
171 |    "outputs": [],
172 |    "source": [
173 |     "class Agent:\n",
174 |     "\n",
175 |     "    def __init__(self, model, tools, system=\"\"):\n",
176 |     "        self.system = system\n",
177 |     "        graph = StateGraph(AgentState)\n",
178 |     "        graph.add_node(\"llm\", self.call_bedrock)\n",
179 |     "        graph.add_node(\"action\", self.take_action)\n",
180 |     "        graph.add_conditional_edges(\n",
181 |     "            \"llm\", self.exists_action, {True: \"action\", False: END}\n",
182 |     "        )\n",
183 |     "        graph.add_edge(\"action\", \"llm\")\n",
184 |     "        graph.set_entry_point(\"llm\")\n",
185 |     "        self.graph = graph.compile()\n",
186 |     "        self.tools = {t.name: t for t in tools}\n",
187 |     "        self.model = model.bind_tools(tools)\n",
188 |     "\n",
189 |     "    def exists_action(self, state: AgentState):\n",
190 |     "        result = state[\"messages\"][-1]\n",
191 |     "        return len(result.tool_calls) > 0\n",
192 |     "\n",
193 |     "    def call_bedrock(self, state: AgentState):\n",
194 |     "        messages = state[\"messages\"]\n",
195 |     "        if self.system:\n",
196 |     "            messages = [SystemMessage(content=self.system)] + messages\n",
197 |     "        message = self.model.invoke(messages)\n",
198 |     "        return {\"messages\": [message]}\n",
199 |     "\n",
200 |     "    def take_action(self, state: AgentState):\n",
201 |     "        tool_calls = state[\"messages\"][-1].tool_calls\n",
202 |     "        results = []\n",
203 |     "        for t in tool_calls:\n",
204 |     "            print(f\"Calling: {t}\")\n",
205 |     "            if not t[\"name\"] in self.tools:  # check for bad tool name from LLM\n",
206 |     "                print(\"\\n ....bad tool name....\")\n",
207 |     "                result = \"bad tool name, retry\"  # instruct LLM to retry if bad\n",
208 |     "            else:\n",
209 |     "                result = self.tools[t[\"name\"]].invoke(t[\"args\"])\n",
210 |     "            results.append(\n",
211 |     "                ToolMessage(tool_call_id=t[\"id\"], name=t[\"name\"], content=str(result))\n",
212 |     "            )\n",
213 |     "        print(\"Back to the model!\")\n",
214 |     "        return {\"messages\": results}"
215 |    ]
216 |   },
217 |   {
218 |    "cell_type": "markdown",
219 |    "id": "2cd50074",
220 |    "metadata": {},
221 |    "source": [
222 |     "An often overlooked feature is the `??` to inspect the code of a function or object in python.\n",
223 |     "\n",
224 |     "Lets inspect the `bind_tools` method on the `ChatBedrockConverse` class.\n",
225 |     "Can you spot if our tavily tool will be supported, and if there are any restrictions?\n",
226 |     "\n",
227 |     "If you are unsure, how would you check?\n"
228 |    ]
229 |   },
230 |   {
231 |    "cell_type": "code",
232 |    "execution_count": null,
233 |    "id": "a61313ff",
234 |    "metadata": {},
235 |    "outputs": [],
236 |    "source": [
237 |     "??ChatBedrockConverse.bind_tools"
238 |    ]
239 |   },
240 |   {
241 |    "cell_type": "code",
242 |    "execution_count": null,
243 |    "id": "10084a02-2928-4945-9f7c-ad3f5b33caf7",
244 |    "metadata": {
245 |     "height": 149
246 |    },
247 |    "outputs": [],
248 |    "source": [
249 |     "prompt = \"\"\"You are a smart research assistant. Use the search engine to look up information. \\\n",
250 |     "You are allowed to make multiple calls (either together or in sequence).\\\n",
251 |     "Whenever you can, try to call multiple tools at once, to bring down inference time!\\\n",
252 |     "Only look up information when you are sure of what you want. \\\n",
253 |     "If you need to look up some information before asking a follow up question, you are allowed to do that!\n",
254 |     "\"\"\"\n",
255 |     "\n",
256 |     "model = ChatBedrockConverse(\n",
257 |     "    client=bedrock_rt,\n",
258 |     "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
259 |     "    temperature=0,\n",
260 |     "    max_tokens=None,\n",
261 |     ")\n",
262 |     "\n",
263 |     "abot = Agent(model, [tool], system=prompt)"
264 |    ]
265 |   },
266 |   {
267 |    "cell_type": "code",
268 |    "execution_count": null,
269 |    "id": "a3d6f5f4-2392-41b9-ab96-7919840baa3e",
270 |    "metadata": {
271 |     "height": 64
272 |    },
273 |    "outputs": [],
274 |    "source": [
275 |     "# make sure to install pygraphviz if you haven't done so already using 'conda install --channel conda-forge pygraphviz'\n",
276 |     "from IPython.display import Image\n",
277 |     "\n",
278 |     "Image(abot.graph.get_graph().draw_png())"
279 |    ]
280 |   },
281 |   {
282 |    "cell_type": "code",
283 |    "execution_count": null,
284 |    "id": "83588e70-254f-4f83-a510-c8ae81e729b0",
285 |    "metadata": {
286 |     "height": 47
287 |    },
288 |    "outputs": [],
289 |    "source": [
290 |     "messages = [HumanMessage(content=\"What is the weather in sf?\")]\n",
291 |     "result = abot.graph.invoke({\"messages\": messages})"
292 |    ]
293 |   },
294 |   {
295 |    "cell_type": "code",
296 |    "execution_count": null,
297 |    "id": "89a06a8c-fcd4-4ca6-98f0-36c5809813e6",
298 |    "metadata": {
299 |     "height": 30,
300 |     "scrolled": true
301 |    },
302 |    "outputs": [],
303 |    "source": [
304 |     "for message in result[\"messages\"]:\n",
305 |     "    print(f\"{message}\\n\")"
306 |    ]
307 |   },
308 |   {
309 |    "cell_type": "code",
310 |    "execution_count": null,
311 |    "id": "6cb3ef4c-58b3-401b-b104-0d51e553d982",
312 |    "metadata": {
313 |     "height": 30
314 |    },
315 |    "outputs": [],
316 |    "source": [
317 |     "result[\"messages\"][-1].content"
318 |    ]
319 |   },
320 |   {
321 |    "cell_type": "code",
322 |    "execution_count": null,
323 |    "id": "dc3293b7-a50c-43c8-a022-8975e1e444b8",
324 |    "metadata": {
325 |     "height": 47
326 |    },
327 |    "outputs": [],
328 |    "source": [
329 |     "messages = [HumanMessage(content=\"What is the weather in SF and LA?\")]\n",
330 |     "result = abot.graph.invoke({\"messages\": messages})"
331 |    ]
332 |   },
333 |   {
334 |    "cell_type": "code",
335 |    "execution_count": null,
336 |    "id": "0722c3d4-4cbf-43bf-81b0-50f634c4ce61",
337 |    "metadata": {
338 |     "height": 30
339 |    },
340 |    "outputs": [],
341 |    "source": [
342 |     "result[\"messages\"][-1].content"
343 |    ]
344 |   },
345 |   {
346 |    "cell_type": "markdown",
347 |    "id": "9f759e2d",
348 |    "metadata": {},
349 |    "source": [
350 |     "## 4. Parallel vs. Sequential Tool Calling\n",
351 |     "\n",
352 |     "The ability of the agent to make both parallel and sequential tool calls is a powerful feature that solution architects should pay attention to.\n",
353 |     "\n",
354 |     "**Deep Dive**:\n",
355 |     "\n",
356 |     "- Parallel tool calling is like a multithreaded application, where multiple independent tasks can be executed simultaneously. This is efficient for queries that require multiple pieces of independent information.\n",
357 |     "- Sequential tool calling is more like a pipeline, where the output of one operation becomes the input of the next. This is necessary for multi-step reasoning tasks.\n",
358 |     "\n",
359 |     "**Analogy**: Imagine a research team working on a complex project. Parallel tool calling is like assigning different team members to research different aspects simultaneously. Sequential tool calling is like a relay race, where each researcher builds on the findings of the previous one.\n",
360 |     "\n",
361 |     "Can you spot if we will have sequential or parallel tool calls and if we have parallel, would they really be executed in parallel?\n"
362 |    ]
363 |   },
364 |   {
365 |    "cell_type": "code",
366 |    "execution_count": null,
367 |    "id": "6b2f82fe-3ec4-4917-be51-9fb10d1317fa",
368 |    "metadata": {
369 |     "height": 183
370 |    },
371 |    "outputs": [],
372 |    "source": [
373 |     "# Note, the query was modified to produce more consistent results.\n",
374 |     "# Results may vary per run and over time as search information and models change.\n",
375 |     "\n",
376 |     "query = \"Who won the super bowl in 2024? In what state is the winning team headquarters located? \\\n",
377 |     "What is the GDP of that state? Answer each question.\"\n",
378 |     "messages = [HumanMessage(content=query)]\n",
379 |     "\n",
380 |     "model = ChatBedrockConverse(\n",
381 |     "    client=bedrock_rt,\n",
382 |     "    model=\"anthropic.claude-3-sonnet-20240229-v1:0\",\n",
383 |     "    temperature=0,\n",
384 |     "    max_tokens=None,\n",
385 |     ")\n",
386 |     "abot = Agent(model, [tool], system=prompt)\n",
387 |     "result = abot.graph.invoke({\"messages\": messages})"
388 |    ]
389 |   },
390 |   {
391 |    "cell_type": "code",
392 |    "execution_count": null,
393 |    "id": "ee0fe1c7-77e2-499c-a2f9-1f739bb6ddf0",
394 |    "metadata": {
395 |     "height": 30
396 |    },
397 |    "outputs": [],
398 |    "source": [
399 |     "print(result[\"messages\"][-1].content)"
400 |    ]
401 |   },
402 |   {
403 |    "cell_type": "markdown",
404 |    "id": "80ea5638",
405 |    "metadata": {},
406 |    "source": [
407 |     "# Exercise: How would you have to change the tool definition, to allow for parallel calling of the tool?\n",
408 |     "\n",
409 |     "> Note: you can omit the parallel execution with async\n"
410 |    ]
411 |   },
412 |   {
413 |    "cell_type": "code",
414 |    "execution_count": null,
415 |    "id": "0833ba0b",
416 |    "metadata": {},
417 |    "outputs": [],
418 |    "source": []
419 |   }
420 |  ],
421 |  "metadata": {
422 |   "kernelspec": {
423 |    "display_name": "agents-dev-env",
424 |    "language": "python",
425 |    "name": "agents-dev-env"
426 |   },
427 |   "language_info": {
428 |    "codemirror_mode": {
429 |     "name": "ipython",
430 |     "version": 3
431 |    },
432 |    "file_extension": ".py",
433 |    "mimetype": "text/x-python",
434 |    "name": "python",
435 |    "nbconvert_exporter": "python",
436 |    "pygments_lexer": "ipython3",
437 |    "version": "3.10.14"
438 |   }
439 |  },
440 |  "nbformat": 4,
441 |  "nbformat_minor": 5
442 | }
443 | 


--------------------------------------------------------------------------------
/Lab_4/Lab_4.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |  "cells": [
  3 |   {
  4 |    "cell_type": "markdown",
  5 |    "id": "b5789bc3-b1ae-42c7-94a8-2ef4f89946fc",
  6 |    "metadata": {},
  7 |    "source": [
  8 |     "# Lab 4: Persistence and Streaming"
  9 |    ]
 10 |   },
 11 |   {
 12 |    "cell_type": "markdown",
 13 |    "id": "35b724f8",
 14 |    "metadata": {},
 15 |    "source": [
 16 |     "## Environment Setup\n",
 17 |     "\n",
 18 |     "We begin by establishing our agent environment. This process involves loading the necessary environment variables, importing required modules, initializing our Tavily search tool, defining the agent state, and finally, constructing our agent.\n"
 19 |    ]
 20 |   },
 21 |   {
 22 |    "cell_type": "code",
 23 |    "execution_count": null,
 24 |    "id": "f5762271-8736-4e94-9444-8c92bd0e8074",
 25 |    "metadata": {
 26 |     "height": 64
 27 |    },
 28 |    "outputs": [],
 29 |    "source": [
 30 |     "from dotenv import load_dotenv\n",
 31 |     "import os\n",
 32 |     "import sys\n",
 33 |     "import json\n",
 34 |     "import re\n",
 35 |     "import pprint\n",
 36 |     "import boto3\n",
 37 |     "from botocore.client import Config\n",
 38 |     "import warnings\n",
 39 |     "\n",
 40 |     "warnings.filterwarnings(\"ignore\")\n",
 41 |     "import logging\n",
 42 |     "\n",
 43 |     "# import local modules\n",
 44 |     "dir_current = os.path.abspath(\"\")\n",
 45 |     "dir_parent = os.path.dirname(dir_current)\n",
 46 |     "if dir_parent not in sys.path:\n",
 47 |     "    sys.path.append(dir_parent)\n",
 48 |     "from utils import utils\n",
 49 |     "\n",
 50 |     "bedrock_config = Config(\n",
 51 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
 52 |     ")\n",
 53 |     "\n",
 54 |     "# Set basic configs\n",
 55 |     "logger = utils.set_logger()\n",
 56 |     "pp = utils.set_pretty_printer()\n",
 57 |     "\n",
 58 |     "# Load environment variables from .env file or Secret Manager\n",
 59 |     "_ = load_dotenv(\"../.env\")\n",
 60 |     "aws_region = os.getenv(\"AWS_REGION\")\n",
 61 |     "tavily_ai_api_key = utils.get_tavily_api(\"TAVILY_API_KEY\", aws_region)\n",
 62 |     "\n",
 63 |     "# Set bedrock configs\n",
 64 |     "bedrock_config = Config(\n",
 65 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
 66 |     ")\n",
 67 |     "\n",
 68 |     "# Create a bedrock runtime client\n",
 69 |     "bedrock_rt = boto3.client(\n",
 70 |     "    \"bedrock-runtime\", region_name=aws_region, config=bedrock_config\n",
 71 |     ")\n",
 72 |     "\n",
 73 |     "# Create a bedrock client to check available models\n",
 74 |     "bedrock = boto3.client(\"bedrock\", region_name=aws_region, config=bedrock_config)\n"
 75 |    ]
 76 |   },
 77 |   {
 78 |    "cell_type": "code",
 79 |    "execution_count": null,
 80 |    "id": "d0168aee-bce9-4d60-b827-f86a88187e31",
 81 |    "metadata": {
 82 |     "height": 115
 83 |    },
 84 |    "outputs": [],
 85 |    "source": [
 86 |     "from langgraph.graph import StateGraph, END\n",
 87 |     "from typing import TypedDict, Annotated\n",
 88 |     "import operator\n",
 89 |     "from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage\n",
 90 |     "from langchain_aws import ChatBedrockConverse\n",
 91 |     "from langchain_community.tools.tavily_search import TavilySearchResults\n",
 92 |     "from langgraph.checkpoint.memory import MemorySaver\n",
 93 |     "\n",
 94 |     "memory = MemorySaver()"
 95 |    ]
 96 |   },
 97 |   {
 98 |    "cell_type": "code",
 99 |    "execution_count": null,
100 |    "id": "da06a64f-a2d5-4a66-8090-9ada0930c684",
101 |    "metadata": {
102 |     "height": 30
103 |    },
104 |    "outputs": [],
105 |    "source": [
106 |     "tool = TavilySearchResults(max_results=2)"
107 |    ]
108 |   },
109 |   {
110 |    "cell_type": "markdown",
111 |    "id": "c989adc7",
112 |    "metadata": {},
113 |    "source": [
114 |     "## Implementing Persistence\n",
115 |     "\n",
116 |     "We now turn our attention to implementing persistence. To achieve this, we introduce the concept of a checkpointer in LangGraph. The checkpointer's function is to create state snapshots after and between each node in our agent's processing graph.\n",
117 |     "\n",
118 |     "#RESOURCE For a more comprehensive understanding of LangGraph's capabilities and usage, refer to the official LangGraph documentation.\n",
119 |     "\n",
120 |     "In this implementation, we utilize a SQLite saver as our checkpointer. This lightweight solution leverages SQLite, a built-in database engine. While we use an in-memory database for this demonstration, it's important to note that this can be easily adapted to connect to an external database for production environments. LangGraph also supports other persistence solutions, including Redis and Postgres, for scenarios requiring more robust database systems.\n",
121 |     "\n",
122 |     "After initializing the checkpointer, we pass it to the `graph.compile` method. We've enhanced our agent to accept a `checkpointer` parameter, which we set to our memory object.\n"
123 |    ]
124 |   },
125 |   {
126 |    "cell_type": "code",
127 |    "execution_count": null,
128 |    "id": "2589c5b6-6cc2-4594-9a17-dccdcf676054",
129 |    "metadata": {
130 |     "height": 47
131 |    },
132 |    "outputs": [],
133 |    "source": [
134 |     "class AgentState(TypedDict):\n",
135 |     "    messages: Annotated[list[AnyMessage], operator.add]"
136 |    ]
137 |   },
138 |   {
139 |    "cell_type": "markdown",
140 |    "id": "01441e5e",
141 |    "metadata": {},
142 |    "source": [
143 |     "## The Agent Class: A Detailed Examination\n",
144 |     "\n",
145 |     "The `Agent` class serves as the cornerstone of our implementation, orchestrating the interactions between the language model (Claude), tools (such as the Tavily search), and the overall conversation flow. Let's examine its key components:\n",
146 |     "\n",
147 |     "1. `__init__` method: This initializer sets up the agent with a model, tools, checkpointer, and an optional system message. It constructs the state graph that defines the agent's behavior.\n",
148 |     "\n",
149 |     "2. `call_bedrock` method: This method is responsible for invoking the Claude model via Amazon Bedrock. It processes the current state (messages) and returns the model's response.\n",
150 |     "\n",
151 |     "3. `exists_action` method: This method evaluates whether the latest message from the model includes any tool calls (actions to be executed).\n",
152 |     "\n",
153 |     "4. `take_action` method: This method executes any tool calls specified by the model and returns the results.\n",
154 |     "\n",
155 |     "The `Agent` class utilizes a `StateGraph` to manage the conversation flow, enabling complex interactions while maintaining a clear and manageable structure. This design choice facilitates the implementation of persistence and streaming capabilities.\n",
156 |     "\n",
157 |     "## Streaming Implementation\n",
158 |     "\n",
159 |     "With our agent now configured, we can implement streaming functionality. There are two primary aspects of streaming to consider:\n",
160 |     "\n",
161 |     "1. Message Streaming: This involves streaming individual messages, including the AI message that determines the next action and the observation message that represents the action's result.\n",
162 |     "\n",
163 |     "2. Token Streaming: This involves streaming each token of the language model's response as it's generated.\n",
164 |     "\n",
165 |     "We'll begin by implementing message streaming. We create a human message (e.g., \"What is the weather in SF?\") and introduce a thread config. This thread config is crucial for managing multiple conversations simultaneously within the persistent checkpointer, a necessity for production applications serving multiple users.\n",
166 |     "\n",
167 |     "We invoke the graph using the `stream` method instead of `invoke`, passing our messages dictionary and thread config. This returns a stream of events representing real-time updates to the state.\n",
168 |     "\n",
169 |     "Upon execution, we observe a stream of results: first, an AI message from Claude determining the action to take, followed by a tool message containing the Tavily search results, and finally, another AI message from Claude answering our initial query.\n"
170 |    ]
171 |   },
172 |   {
173 |    "cell_type": "code",
174 |    "execution_count": null,
175 |    "id": "a2ba84ec-c172-4de7-ac55-e3158a531b23",
176 |    "metadata": {
177 |     "height": 574
178 |    },
179 |    "outputs": [],
180 |    "source": [
181 |     "class Agent:\n",
182 |     "    def __init__(self, model, tools, checkpointer, system=\"\"):\n",
183 |     "        self.system = system\n",
184 |     "        graph = StateGraph(AgentState)\n",
185 |     "        graph.add_node(\"llm\", self.call_bedrock)\n",
186 |     "        graph.add_node(\"action\", self.take_action)\n",
187 |     "        graph.add_conditional_edges(\n",
188 |     "            \"llm\", self.exists_action, {True: \"action\", False: END}\n",
189 |     "        )\n",
190 |     "        graph.add_edge(\"action\", \"llm\")\n",
191 |     "        graph.set_entry_point(\"llm\")\n",
192 |     "        self.graph = graph.compile(checkpointer=checkpointer)\n",
193 |     "        self.tools = {t.name: t for t in tools}\n",
194 |     "        self.model = model.bind_tools(tools)\n",
195 |     "\n",
196 |     "    def call_bedrock(self, state: AgentState):\n",
197 |     "        messages = state[\"messages\"]\n",
198 |     "        if self.system:\n",
199 |     "            messages = [SystemMessage(content=self.system)] + messages\n",
200 |     "        message = self.model.invoke(messages)\n",
201 |     "        return {\"messages\": [message]}\n",
202 |     "\n",
203 |     "    def exists_action(self, state: AgentState):\n",
204 |     "        result = state[\"messages\"][-1]\n",
205 |     "        return len(result.tool_calls) > 0\n",
206 |     "\n",
207 |     "    def take_action(self, state: AgentState):\n",
208 |     "        tool_calls = state[\"messages\"][-1].tool_calls\n",
209 |     "        results = []\n",
210 |     "        for t in tool_calls:\n",
211 |     "            print(f\"Calling: {t}\")\n",
212 |     "            result = self.tools[t[\"name\"]].invoke(t[\"args\"])\n",
213 |     "            results.append(\n",
214 |     "                ToolMessage(tool_call_id=t[\"id\"], name=t[\"name\"], content=str(result))\n",
215 |     "            )\n",
216 |     "        print(\"Back to the model!\")\n",
217 |     "        return {\"messages\": results}"
218 |    ]
219 |   },
220 |   {
221 |    "cell_type": "code",
222 |    "execution_count": null,
223 |    "id": "876d5092-b8ef-4e38-b4d7-0e80c609bf7a",
224 |    "metadata": {
225 |     "height": 132
226 |    },
227 |    "outputs": [],
228 |    "source": [
229 |     "prompt = \"\"\"You are a smart research assistant. Use the search engine to look up information. \\\n",
230 |     "You are allowed to make multiple calls (either together or in sequence). \\\n",
231 |     "Only look up information when you are sure of what you want. \\\n",
232 |     "If you need to look up some information before asking a follow up question, you are allowed to do that!\n",
233 |     "\"\"\"\n",
234 |     "\n",
235 |     "\n",
236 |     "model = ChatBedrockConverse(\n",
237 |     "    client=bedrock_rt,\n",
238 |     "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
239 |     "    temperature=0,\n",
240 |     "    max_tokens=None,\n",
241 |     ")\n",
242 |     "abot = Agent(model, [tool], system=prompt, checkpointer=memory)"
243 |    ]
244 |   },
245 |   {
246 |    "cell_type": "code",
247 |    "execution_count": null,
248 |    "id": "10084a02-2928-4945-9f7c-ad3f5b33caf7",
249 |    "metadata": {
250 |     "height": 30
251 |    },
252 |    "outputs": [],
253 |    "source": [
254 |     "messages = [HumanMessage(content=\"What is the weather in sf?\")]"
255 |    ]
256 |   },
257 |   {
258 |    "cell_type": "code",
259 |    "execution_count": null,
260 |    "id": "714d1205-f8fc-4912-b148-2a45da99219c",
261 |    "metadata": {
262 |     "height": 30
263 |    },
264 |    "outputs": [],
265 |    "source": [
266 |     "thread = {\"configurable\": {\"thread_id\": \"1\"}}"
267 |    ]
268 |   },
269 |   {
270 |    "cell_type": "code",
271 |    "execution_count": null,
272 |    "id": "83588e70-254f-4f83-a510-c8ae81e729b0",
273 |    "metadata": {
274 |     "height": 64
275 |    },
276 |    "outputs": [],
277 |    "source": [
278 |     "for event in abot.graph.stream({\"messages\": messages}, thread):\n",
279 |     "    for v in event.values():\n",
280 |     "        print(v[\"messages\"])"
281 |    ]
282 |   },
283 |   {
284 |    "cell_type": "markdown",
285 |    "id": "070f625b",
286 |    "metadata": {},
287 |    "source": [
288 |     "## Demonstrating Persistence\n",
289 |     "\n",
290 |     "To illustrate the effectiveness of our persistence implementation, we continue the conversation with a follow-up question: \"What about in LA?\". By using the same thread ID, we ensure continuity from the previous interaction. Claude maintains context, understanding that we're still inquiring about weather conditions due to the persistence provided by our checkpoint system.\n",
291 |     "\n",
292 |     "We can further emphasize the importance of thread ID by altering it and posing the question, \"Which one is warmer?\". With the original thread ID, Claude can accurately compare temperatures. However, changing the thread ID results in Claude losing context, as it no longer has access to the conversation history.\n"
293 |    ]
294 |   },
295 |   {
296 |    "cell_type": "code",
297 |    "execution_count": null,
298 |    "id": "6cb3ef4c-58b3-401b-b104-0d51e553d982",
299 |    "metadata": {
300 |     "height": 98
301 |    },
302 |    "outputs": [],
303 |    "source": [
304 |     "messages = [HumanMessage(content=\"What about in la?\")]\n",
305 |     "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
306 |     "for event in abot.graph.stream({\"messages\": messages}, thread):\n",
307 |     "    for v in event.values():\n",
308 |     "        print(v)"
309 |    ]
310 |   },
311 |   {
312 |    "cell_type": "code",
313 |    "execution_count": null,
314 |    "id": "dc3293b7-a50c-43c8-a022-8975e1e444b8",
315 |    "metadata": {
316 |     "height": 98
317 |    },
318 |    "outputs": [],
319 |    "source": [
320 |     "messages = [HumanMessage(content=\"Which one is warmer?\")]\n",
321 |     "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
322 |     "for event in abot.graph.stream({\"messages\": messages}, thread):\n",
323 |     "    for v in event.values():\n",
324 |     "        print(v)"
325 |    ]
326 |   },
327 |   {
328 |    "cell_type": "code",
329 |    "execution_count": null,
330 |    "id": "0722c3d4-4cbf-43bf-81b0-50f634c4ce61",
331 |    "metadata": {
332 |     "height": 98
333 |    },
334 |    "outputs": [],
335 |    "source": [
336 |     "messages = [HumanMessage(content=\"Which one is warmer?\")]\n",
337 |     "thread = {\"configurable\": {\"thread_id\": \"2\"}}\n",
338 |     "for event in abot.graph.stream({\"messages\": messages}, thread):\n",
339 |     "    for v in event.values():\n",
340 |     "        print(v)"
341 |    ]
342 |   },
343 |   {
344 |    "cell_type": "markdown",
345 |    "id": "c607bb30",
346 |    "metadata": {},
347 |    "source": [
348 |     "## Token-Level Streaming\n",
349 |     "\n",
350 |     "For a more granular approach to streaming, we implement token-level updates using the `astream_events` method. This asynchronous method necessitates an async checkpointer, which we implement using `AsyncSqliteSaver`.\n",
351 |     "\n",
352 |     "Asynchronous programming allows our application to handle multiple operations concurrently without blocking the main execution thread. In the context of streaming tokens from an AI model, this translates to processing and displaying tokens as they're generated, resulting in a more responsive user experience. The `astream_events` method leverages this asynchronous approach to efficiently stream token-level updates from Claude.\n",
353 |     "\n",
354 |     "We initiate a new conversation with a fresh thread ID and iterate over the events, specifically looking for events of type \"on_chat_model_stream\". Upon encountering these events, we extract and display the content.\n",
355 |     "\n",
356 |     "When executed, we observe tokens streaming in real-time. We see Claude invoke the function (which doesn't generate streamable content), followed by the final response streaming token by token.\n"
357 |    ]
358 |   },
359 |   {
360 |    "cell_type": "code",
361 |    "execution_count": null,
362 |    "id": "6b2f82fe-3ec4-4917-be51-9fb10d1317fa",
363 |    "metadata": {
364 |     "height": 81
365 |    },
366 |    "outputs": [],
367 |    "source": [
368 |     "from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver\n",
369 |     "\n",
370 |     "# # If you are using a newer version of LangGraph, the package was separated:\n",
371 |     "# # !pip install langgraph-checkpoint-sqlite\n",
372 |     "\n",
373 |     "# from langgraph.checkpoint.memory import MemorySaver\n",
374 |     "# from langgraph.checkpoint.sqlite import SqliteSaver\n",
375 |     "# from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver\n",
376 |     "\n",
377 |     "async with AsyncSqliteSaver.from_conn_string(\"checkpoints.db\") as memory:\n",
378 |     "    abot = Agent(model, [tool], system=prompt, checkpointer=memory)\n",
379 |     "\n",
380 |     "    messages = [HumanMessage(content=\"What is the weather in SF?\")]\n",
381 |     "    thread = {\"configurable\": {\"thread_id\": \"4\"}}\n",
382 |     "    async for event in abot.graph.astream_events(\n",
383 |     "        {\"messages\": messages}, thread, version=\"v1\"\n",
384 |     "    ):\n",
385 |     "        kind = event[\"event\"]\n",
386 |     "        if kind == \"on_chat_model_stream\":\n",
387 |     "            content = event[\"data\"][\"chunk\"].content\n",
388 |     "            if content:\n",
389 |     "                # Empty content in the context of Amazon Bedrock means\n",
390 |     "                # that the model is asking for a tool to be invoked.\n",
391 |     "                # So we only print non-empty content\n",
392 |     "                print(content, end=\"|\")"
393 |    ]
394 |   },
395 |   {
396 |    "cell_type": "markdown",
397 |    "id": "4f8e0e9e",
398 |    "metadata": {},
399 |    "source": [
400 |     "## Conclusion\n",
401 |     "\n",
402 |     "This lab has provided a comprehensive exploration of persistence and streaming implementation using Anthropic's Claude model on Amazon Bedrock. While these concepts are straightforward to implement, they offer powerful capabilities for building production-grade AI applications.\n",
403 |     "\n",
404 |     "The ability to manage multiple simultaneous conversations, coupled with a robust memory system for conversation resumption, is crucial for scalable AI solutions. Moreover, the capacity to stream both final tokens and intermediate messages provides unparalleled visibility into the AI's decision-making process.\n",
405 |     "\n",
406 |     "Persistence also plays a vital role in enabling human-in-the-loop interactions, a topic we will explore in greater depth in our subsequent lab.\n",
407 |     "\n",
408 |     " To gain a deeper understanding of the practical implications of these concepts, we recommend exploring real-world case studies of persistence and streaming in production AI applications.\n"
409 |    ]
410 |   },
411 |   {
412 |    "cell_type": "markdown",
413 |    "id": "df424a98",
414 |    "metadata": {},
415 |    "source": []
416 |   }
417 |  ],
418 |  "metadata": {
419 |   "kernelspec": {
420 |    "display_name": "agents-dev-env",
421 |    "language": "python",
422 |    "name": "agents-dev-env"
423 |   },
424 |   "language_info": {
425 |    "codemirror_mode": {
426 |     "name": "ipython",
427 |     "version": 3
428 |    },
429 |    "file_extension": ".py",
430 |    "mimetype": "text/x-python",
431 |    "name": "python",
432 |    "nbconvert_exporter": "python",
433 |    "pygments_lexer": "ipython3",
434 |    "version": "3.10.14"
435 |   }
436 |  },
437 |  "nbformat": 4,
438 |  "nbformat_minor": 5
439 | }
440 | 


--------------------------------------------------------------------------------
/Lab_1/Lab_1.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |  "cells": [
  3 |   {
  4 |    "cell_type": "markdown",
  5 |    "id": "9690ac72-5d95-4cbf-875a-ae0e835593c9",
  6 |    "metadata": {},
  7 |    "source": [
  8 |     "# Lab 1: Building a ReAct Agent from Scratch\n",
  9 |     "\n",
 10 |     "## The ReAct Pattern\n",
 11 |     "\n",
 12 |     "In this section, we'll construct an AI agent using the ReAct (Reasoning and Acting) pattern. If you're unfamiliar with this concept, don't worry—we'll break it down step by step.\n",
 13 |     "\n",
 14 |     "The ReAct pattern is a framework for structuring an AI's problem-solving process, mirroring human cognitive patterns:\n",
 15 |     "\n",
 16 |     "1. **Reason** about the current situation\n",
 17 |     "2. **Decide** on an action to take\n",
 18 |     "3. **Observe** the results of that action\n",
 19 |     "4. **Repeat** until the task is complete\n",
 20 |     "\n",
 21 |     "To illustrate this concept, consider how an experienced software engineer might approach debugging a complex system:\n",
 22 |     "\n",
 23 |     "1. **Reason**: Analyze the error logs and system state (e.g., \"The database connection is timing out\")\n",
 24 |     "2. **Act**: Implement a diagnostic action (e.g., \"Run a database connection test\")\n",
 25 |     "3. **Observe**: Examine the results of the diagnostic (e.g., \"The test shows high latency\")\n",
 26 |     "4. **Repeat**: Continue this process, perhaps checking network configurations next, until the issue is resolved\n",
 27 |     "\n",
 28 |     "Our AI agent will employ a similar methodology to tackle problems. As we develop this agent, pay attention to the division of labor between the AI model (the \"brain\" that reasons and decides) and our Python code (the \"body\" that interacts with the environment and manages the process flow).\n",
 29 |     "\n",
 30 |     "This notebook is based on the following [notebook from Simon Willison](https://til.simonwillison.net/llms/python-react-pattern).\n"
 31 |    ]
 32 |   },
 33 |   {
 34 |    "cell_type": "markdown",
 35 |    "id": "0705a93c",
 36 |    "metadata": {},
 37 |    "source": [
 38 |     "## Setting Up the Environment\n",
 39 |     "\n",
 40 |     "Let's begin by importing the necessary libraries and configuring our environment.\n",
 41 |     "\n",
 42 |     "### Initializing the Bedrock Client\n",
 43 |     "\n",
 44 |     "To communicate with the Claude model via Amazon Bedrock, we need to establish a client connection. Think of this client as an API gateway that enables our code to send requests to the AI model and receive responses.\n",
 45 |     "\n",
 46 |     "We'll use the `boto3` library, which is the Amazon Web Services (AWS) SDK for Python. For those unfamiliar with AWS, `boto3` can be thought of as a comprehensive toolkit that facilitates Python's interaction with various AWS services, including Bedrock.\n",
 47 |     "\n",
 48 |     "For comprehensive instructions on configuring `boto3` with AWS credentials, refer to the [AWS documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html).\n",
 49 |     "\n",
 50 |     "In a production setting, you would implement secure AWS credential management. For the purposes of this section, we'll assume the credentials are pre-configured in your environment.\n"
 51 |    ]
 52 |   },
 53 |   {
 54 |    "cell_type": "code",
 55 |    "execution_count": null,
 56 |    "id": "e17ac2a6",
 57 |    "metadata": {},
 58 |    "outputs": [],
 59 |    "source": [
 60 |     "from dotenv import load_dotenv\n",
 61 |     "import os\n",
 62 |     "import sys\n",
 63 |     "import boto3\n",
 64 |     "import re\n",
 65 |     "from botocore.config import Config\n",
 66 |     "import warnings\n",
 67 |     "\n",
 68 |     "warnings.filterwarnings(\"ignore\")\n",
 69 |     "import logging\n",
 70 |     "\n",
 71 |     "# import local modules\n",
 72 |     "dir_current = os.path.abspath(\"\")\n",
 73 |     "dir_parent = os.path.dirname(dir_current)\n",
 74 |     "if dir_parent not in sys.path:\n",
 75 |     "    sys.path.append(dir_parent)\n",
 76 |     "from utils import utils\n",
 77 |     "\n",
 78 |     "# Set basic configs\n",
 79 |     "logger = utils.set_logger()\n",
 80 |     "pp = utils.set_pretty_printer()\n",
 81 |     "\n",
 82 |     "# Load environment variables from .env file\n",
 83 |     "_ = load_dotenv(\"../.env\")\n",
 84 |     "aws_region = os.getenv(\"AWS_REGION\")\n",
 85 |     "\n",
 86 |     "# Set bedrock configs\n",
 87 |     "bedrock_config = Config(\n",
 88 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
 89 |     ")\n",
 90 |     "\n",
 91 |     "# Create a bedrock runtime client in your aws region.\n",
 92 |     "# If you do not have the AWS CLI profile setup, you can authenticate with aws access key, secret and session token.\n",
 93 |     "# For more details check https://docs.aws.amazon.com/cli/v1/userguide/cli-authentication-short-term.html\n",
 94 |     "bedrock_rt = boto3.client(\n",
 95 |     "    \"bedrock-runtime\",\n",
 96 |     "    region_name=aws_region,\n",
 97 |     "    config=bedrock_config,\n",
 98 |     ")"
 99 |    ]
100 |   },
101 |   {
102 |    "cell_type": "markdown",
103 |    "id": "650d7c33",
104 |    "metadata": {},
105 |    "source": [
106 |     "First we are going to define a few inference parameters and test our connection via `boto3` to Amazon Bedrock.\n"
107 |    ]
108 |   },
109 |   {
110 |    "cell_type": "code",
111 |    "execution_count": null,
112 |    "id": "dc3293b7-a50c-43c8-a022-8975e1e444b8",
113 |    "metadata": {
114 |     "height": 30
115 |    },
116 |    "outputs": [],
117 |    "source": [
118 |     "# Set inference parameters\n",
119 |     "temperature = 0.0\n",
120 |     "top_k = 200\n",
121 |     "inference_config = {\"temperature\": temperature}\n",
122 |     "\n",
123 |     "additional_model_fields = {\"top_k\": top_k}\n",
124 |     "model_id = \"anthropic.claude-3-sonnet-20240229-v1:0\"\n",
125 |     "system_prompts = [{\"text\": \"You are a helpful agent.\"}]\n",
126 |     "message_1 = {\"role\": \"user\", \"content\": [{\"text\": \"Hello world\"}]}\n",
127 |     "\n",
128 |     "# Instantiate messages list\n",
129 |     "messages = []\n",
130 |     "messages.append(message_1)\n",
131 |     "\n",
132 |     "# Send the message.\n",
133 |     "response = bedrock_rt.converse(\n",
134 |     "    modelId=model_id,\n",
135 |     "    messages=messages,\n",
136 |     "    system=system_prompts,\n",
137 |     "    inferenceConfig=inference_config,\n",
138 |     "    additionalModelRequestFields=additional_model_fields,\n",
139 |     ")\n",
140 |     "\n",
141 |     "pp.pprint(response)\n",
142 |     "print(\"\\n\\n\")\n",
143 |     "pp.pprint(response[\"output\"][\"message\"][\"content\"][0][\"text\"])"
144 |    ]
145 |   },
146 |   {
147 |    "cell_type": "markdown",
148 |    "id": "d68601a7",
149 |    "metadata": {},
150 |    "source": [
151 |     "## Designing the Agent Class\n",
152 |     "\n",
153 |     "With our Bedrock client set up, we'll now create our Agent class. This class will serve as the core of our AI agent, encapsulating the logic for interacting with the Claude model and maintaining the conversation state.\n",
154 |     "\n",
155 |     "The ReAct pattern, which our agent will implement, consists of three primary steps:\n",
156 |     "\n",
157 |     "1. **Reasoning (Thought)**: The agent assesses the current situation and formulates a plan. For instance, \"To calculate the total weight of two dog breeds, I need to look up their individual weights and then sum them.\"\n",
158 |     "\n",
159 |     "2. **Acting (Action)**: Based on its reasoning, the agent selects an appropriate action. For example, \"I will query the average weight of a Border Collie.\"\n",
160 |     "\n",
161 |     "3. **Observing (Observation)**: The agent processes the feedback from its action. In our case, this might be \"The average weight of a Border Collie is 30-55 pounds.\"\n",
162 |     "\n",
163 |     "This pattern enables the agent to decompose complex tasks into manageable steps and adapt its strategy based on new information.\n",
164 |     "\n",
165 |     "Our Agent class will implement this pattern by maintaining a conversation history (`self.messages`) and providing methods to interact with the Claude model (`__call__` and `execute`).\n"
166 |    ]
167 |   },
168 |   {
169 |    "cell_type": "code",
170 |    "execution_count": null,
171 |    "id": "ee0fe1c7-77e2-499c-a2f9-1f739bb6ddf0",
172 |    "metadata": {
173 |     "height": 387
174 |    },
175 |    "outputs": [],
176 |    "source": [
177 |     "class Agent:\n",
178 |     "    def __init__(self, system=\"\"):\n",
179 |     "        self.system = system\n",
180 |     "        self.messages = []\n",
181 |     "        if self.system:\n",
182 |     "            self.system = [{\"text\": self.system}]\n",
183 |     "        self.bedrock_client = boto3.client(service_name=\"bedrock-runtime\")\n",
184 |     "\n",
185 |     "    def __call__(self, message):\n",
186 |     "        self.messages.append({\"role\": \"user\", \"content\": [{\"text\": message}]})\n",
187 |     "        result = self.execute()\n",
188 |     "        self.messages.append({\"role\": \"assistant\", \"content\": [{\"text\": result}]})\n",
189 |     "        return result\n",
190 |     "\n",
191 |     "    def execute(self):\n",
192 |     "        inference_config = {\n",
193 |     "            \"temperature\": 0.0,\n",
194 |     "            \"stopSequences\": [\n",
195 |     "                \"<PAUSE>\"\n",
196 |     "            ],  # we will explore later why this is important!\n",
197 |     "        }\n",
198 |     "\n",
199 |     "        additional_model_fields = {\"top_k\": 200}\n",
200 |     "\n",
201 |     "        response = self.bedrock_client.converse(\n",
202 |     "            modelId=\"anthropic.claude-3-sonnet-20240229-v1:0\",\n",
203 |     "            messages=self.messages,\n",
204 |     "            system=self.system,\n",
205 |     "            inferenceConfig=inference_config,\n",
206 |     "            additionalModelRequestFields=additional_model_fields,\n",
207 |     "        )\n",
208 |     "        return response[\"output\"][\"message\"][\"content\"][0][\"text\"]"
209 |    ]
210 |   },
211 |   {
212 |    "cell_type": "markdown",
213 |    "id": "10183105",
214 |    "metadata": {},
215 |    "source": [
216 |     "## Crafting the Prompt\n",
217 |     "\n",
218 |     "The prompt serves as a set of instructions for the AI model, crucial in defining its behavior and available actions.\n",
219 |     "\n",
220 |     "In our implementation, we're directing the model to:\n",
221 |     "\n",
222 |     "- Adhere to the ReAct pattern (Thought, Action, Observation cycle)\n",
223 |     "- Utilize specific formats for each step (e.g., prefixing thoughts with \"Thought:\")\n",
224 |     "- Restrict itself to the provided actions (in this case, a calculator and a dog weight lookup function)\n",
225 |     "\n",
226 |     "We also include a sample interaction to demonstrate the expected response format. This is analogous to providing a completed template before asking someone to fill out a complex form.\n",
227 |     "\n",
228 |     "As you can see from the prompt, the agent class, and the inference parameters, we are asking the model to stop generating after it predicts the `<PAUSE>` token. However, it is always better to be safe than sorry, so we add `<PAUSE>` to the `stopSequences`, ensuring the end of token generation after that!\n"
229 |    ]
230 |   },
231 |   {
232 |    "cell_type": "code",
233 |    "execution_count": null,
234 |    "id": "98f303b1-a4d0-408c-8cc0-515ff980717f",
235 |    "metadata": {
236 |     "height": 557
237 |    },
238 |    "outputs": [],
239 |    "source": [
240 |     "prompt = \"\"\"\n",
241 |     "You run in a loop of Thought, Action, <PAUSE>, Observation.\n",
242 |     "At the end of the loop you output an Answer\n",
243 |     "Use Thought to describe your thoughts about the question you have been asked.\n",
244 |     "Use Action to run one of the actions available to you - then return PAUSE.\n",
245 |     "Observation will be the result of running those actions.\n",
246 |     "\n",
247 |     "Your available actions are:\n",
248 |     "\n",
249 |     "calculate:\n",
250 |     "e.g. calculate: 4 * 7 / 3\n",
251 |     "Runs a calculation and returns the number - uses Python so be sure to use floating point syntax if necessary\n",
252 |     "\n",
253 |     "average_dog_weight:\n",
254 |     "e.g. average_dog_weight: Collie\n",
255 |     "returns average weight of a dog when given the breed\n",
256 |     "\n",
257 |     "If available, always call a tool to inform your decisions, never use your parametric knowledge when a tool can be called. \n",
258 |     "\n",
259 |     "When you have decided that you need to call a tool, output <PAUSE> and stop thereafter! \n",
260 |     "\n",
261 |     "Example session:\n",
262 |     "\n",
263 |     "Question: How much does a Bulldog weigh?\n",
264 |     "Thought: I should look the dogs weight using average_dog_weight\n",
265 |     "Action: average_dog_weight: Bulldog\n",
266 |     "<PAUSE>\n",
267 |     "----- execution stops here -----\n",
268 |     "You will be called again with this:\n",
269 |     "\n",
270 |     "Observation: A Bulldog weights 51 lbs\n",
271 |     "\n",
272 |     "You then output:\n",
273 |     "\n",
274 |     "Answer: A bulldog weights 51 lbs\n",
275 |     "\"\"\".strip()"
276 |    ]
277 |   },
278 |   {
279 |    "cell_type": "markdown",
280 |    "id": "0383697a",
281 |    "metadata": {},
282 |    "source": [
283 |     "## Implementing Helper Functions\n",
284 |     "\n",
285 |     "To empower our agent with practical capabilities, we'll define several helper functions. These functions will serve as the \"actions\" our agent can perform. In this example, we're providing:\n",
286 |     "\n",
287 |     "1. A basic calculator function\n",
288 |     "2. A function to retrieve average dog weights\n",
289 |     "\n",
290 |     "In a more sophisticated application, these functions could cover a diverse range of operations, from web scraping to database queries to API calls. They are the agent's versatile interface with external data sources and systems, offering a wide range of possibilities.\n"
291 |    ]
292 |   },
293 |   {
294 |    "cell_type": "code",
295 |    "execution_count": null,
296 |    "id": "bf4dcb93-6298-4cfd-b3ce-61dfac7fb35f",
297 |    "metadata": {
298 |     "height": 302
299 |    },
300 |    "outputs": [],
301 |    "source": [
302 |     "def calculate(what):\n",
303 |     "    return eval(what, {\"__builtins__\": {}}, {})\n",
304 |     "\n",
305 |     "\n",
306 |     "def average_dog_weight(name):\n",
307 |     "    if name in \"Scottish Terrier\":\n",
308 |     "        return \"Scottish Terriers average 20 lbs\"\n",
309 |     "    elif name in \"Border Collie\":\n",
310 |     "        return \"a Border Collies average weight is 37 lbs\"\n",
311 |     "    elif name in \"Toy Poodle\":\n",
312 |     "        return \"a toy poodles average weight is 7 lbs\"\n",
313 |     "    else:\n",
314 |     "        return \"An average dog weights 50 lbs\"\n",
315 |     "\n",
316 |     "\n",
317 |     "known_actions = {\"calculate\": calculate, \"average_dog_weight\": average_dog_weight}"
318 |    ]
319 |   },
320 |   {
321 |    "cell_type": "markdown",
322 |    "id": "ed6c3246",
323 |    "metadata": {},
324 |    "source": [
325 |     "## Testing the Agent\n",
326 |     "\n",
327 |     "With our agent and its action set defined, we'll conduct an initial test using a straightforward query about a toy poodle's weight.\n",
328 |     "\n",
329 |     "This test will illuminate the agent's information processing flow:\n",
330 |     "\n",
331 |     "1. It will reason about the necessary steps (identifying the need to look up the weight)\n",
332 |     "2. It will execute an action (invoking the `average_dog_weight` function)\n",
333 |     "3. It will process the observation (the returned weight of a toy poodle)\n",
334 |     "4. It will synthesize this information into a coherent response\n"
335 |    ]
336 |   },
337 |   {
338 |    "cell_type": "code",
339 |    "execution_count": null,
340 |    "id": "932883a4-c722-42bb-aec0-b4f41c5c81a4",
341 |    "metadata": {
342 |     "height": 30
343 |    },
344 |    "outputs": [],
345 |    "source": [
346 |     "abot = Agent(prompt)"
347 |    ]
348 |   },
349 |   {
350 |    "cell_type": "code",
351 |    "execution_count": null,
352 |    "id": "ff362f49-dcf1-4ea1-a86c-e516e9ab897d",
353 |    "metadata": {
354 |     "height": 47
355 |    },
356 |    "outputs": [],
357 |    "source": [
358 |     "result = abot(\"How much does a toy poodle weigh?\")\n",
359 |     "print(result)"
360 |    ]
361 |   },
362 |   {
363 |    "cell_type": "code",
364 |    "execution_count": null,
365 |    "id": "a7e15a20-83d7-434c-8551-bce8dcc32be0",
366 |    "metadata": {
367 |     "height": 30
368 |    },
369 |    "outputs": [],
370 |    "source": [
371 |     "result = average_dog_weight(\"Toy Poodle\")\n",
372 |     "result"
373 |    ]
374 |   },
375 |   {
376 |    "cell_type": "code",
377 |    "execution_count": null,
378 |    "id": "a833d3ce-bd31-4319-811d-decff226b970",
379 |    "metadata": {
380 |     "height": 30
381 |    },
382 |    "outputs": [],
383 |    "source": [
384 |     "next_prompt = \"Observation: {}\".format(result)"
385 |    ]
386 |   },
387 |   {
388 |    "cell_type": "code",
389 |    "execution_count": null,
390 |    "id": "76e93cce-6eab-4c7c-ac64-e9993fdb30d6",
391 |    "metadata": {
392 |     "height": 30
393 |    },
394 |    "outputs": [],
395 |    "source": [
396 |     "abot(next_prompt)"
397 |    ]
398 |   },
399 |   {
400 |    "cell_type": "code",
401 |    "execution_count": null,
402 |    "id": "fd2d0990-a932-423f-9ff3-5cada58c5f32",
403 |    "metadata": {
404 |     "height": 30
405 |    },
406 |    "outputs": [],
407 |    "source": [
408 |     "abot.messages"
409 |    ]
410 |   },
411 |   {
412 |    "cell_type": "code",
413 |    "execution_count": null,
414 |    "id": "27cde654-64e2-48bc-80a9-0ed668ccb7dc",
415 |    "metadata": {
416 |     "height": 30
417 |    },
418 |    "outputs": [],
419 |    "source": [
420 |     "abot = Agent(prompt)\n",
421 |     "abot.messages"
422 |    ]
423 |   },
424 |   {
425 |    "cell_type": "code",
426 |    "execution_count": null,
427 |    "id": "4871f644-b131-4065-b7ce-b82c20a41f11",
428 |    "metadata": {
429 |     "height": 64
430 |    },
431 |    "outputs": [],
432 |    "source": [
433 |     "question = \"\"\"I have 2 dogs, a border collie and a scottish terrier. \\\n",
434 |     "What is their combined weight\"\"\"\n",
435 |     "abot(question)"
436 |    ]
437 |   },
438 |   {
439 |    "cell_type": "code",
440 |    "execution_count": null,
441 |    "id": "8c3d8070-3f36-4cf0-a677-508e54359c8f",
442 |    "metadata": {
443 |     "height": 47
444 |    },
445 |    "outputs": [],
446 |    "source": [
447 |     "next_prompt = \"Observation: {}\".format(average_dog_weight(\"Border Collie\"))\n",
448 |     "print(next_prompt)"
449 |    ]
450 |   },
451 |   {
452 |    "cell_type": "code",
453 |    "execution_count": null,
454 |    "id": "98f3be1d-cc4c-41fa-9863-3e386e88e305",
455 |    "metadata": {
456 |     "height": 30
457 |    },
458 |    "outputs": [],
459 |    "source": [
460 |     "abot(next_prompt)"
461 |    ]
462 |   },
463 |   {
464 |    "cell_type": "code",
465 |    "execution_count": null,
466 |    "id": "0ad8a6cc-65d4-4ce7-87aa-4e67d7c23d7b",
467 |    "metadata": {
468 |     "height": 47
469 |    },
470 |    "outputs": [],
471 |    "source": [
472 |     "next_prompt = \"Observation: {}\".format(average_dog_weight(\"Scottish Terrier\"))\n",
473 |     "print(next_prompt)"
474 |    ]
475 |   },
476 |   {
477 |    "cell_type": "code",
478 |    "execution_count": null,
479 |    "id": "592b5e62-a203-433c-92a0-3783f490cde1",
480 |    "metadata": {
481 |     "height": 30
482 |    },
483 |    "outputs": [],
484 |    "source": [
485 |     "abot(next_prompt)"
486 |    ]
487 |   },
488 |   {
489 |    "cell_type": "code",
490 |    "execution_count": null,
491 |    "id": "14fa923c-7e4f-42d1-965f-0f8ccd50fbd7",
492 |    "metadata": {
493 |     "height": 47
494 |    },
495 |    "outputs": [],
496 |    "source": [
497 |     "next_prompt = \"Observation: {}\".format(eval(\"37 + 20\"), {\"__builtins__\": {}}, {})\n",
498 |     "print(next_prompt)"
499 |    ]
500 |   },
501 |   {
502 |    "cell_type": "code",
503 |    "execution_count": null,
504 |    "id": "570c6245-2837-4ac5-983b-95f61f3ac10d",
505 |    "metadata": {
506 |     "height": 30
507 |    },
508 |    "outputs": [],
509 |    "source": [
510 |     "abot(next_prompt)"
511 |    ]
512 |   },
513 |   {
514 |    "cell_type": "markdown",
515 |    "id": "8b970c97",
516 |    "metadata": {},
517 |    "source": [
518 |     "### A word on stopSequences\n",
519 |     "\n",
520 |     "Try to remove the `stopSequence` parameter from the above agent class.\n",
521 |     "\n",
522 |     "Now think about a few of the questions below before continuing:\n",
523 |     "\n",
524 |     "- How does your agent perform now?\n",
525 |     "- When should you use `stopSequences`, when could they become a burden for your application?\n",
526 |     "- Does the prompt comply with the [prompting standard of Anthropic](https://docs.anthropic.com/en/docs/prompt-engineering)?\n",
527 |     "\n",
528 |     "At the end of the notebook, try to complete the exercise so that the agent follows your instructions without having to use `stopSequences`.\n"
529 |    ]
530 |   },
531 |   {
532 |    "cell_type": "markdown",
533 |    "id": "6b46f2ac-f717-4ab9-b548-f34b74071d76",
534 |    "metadata": {},
535 |    "source": [
536 |     "## Implementing the Reasoning Loop\n",
537 |     "\n",
538 |     "To enhance our agent's autonomy, we'll implement an iterative loop that enables it to reason, act, and observe multiple times in pursuit of an answer. This loop will continue until the agent reaches a conclusion or hits a predefined maximum number of iterations.\n",
539 |     "\n",
540 |     "This approach mirrors how a human expert might tackle a complex problem, gathering information and working through multiple steps until arriving at a solution. The loop empowers our agent to handle more intricate queries that demand multiple steps or data points.\n"
541 |    ]
542 |   },
543 |   {
544 |    "cell_type": "code",
545 |    "execution_count": null,
546 |    "id": "6b910915-b087-4d35-afff-0ec30a5852f1",
547 |    "metadata": {
548 |     "height": 30
549 |    },
550 |    "outputs": [],
551 |    "source": [
552 |     "action_re = re.compile(\n",
553 |     "    \"^Action: (\\w+): (.*)$\"\n",
554 |     ")  # python regular expression to selection action"
555 |    ]
556 |   },
557 |   {
558 |    "cell_type": "code",
559 |    "execution_count": null,
560 |    "id": "c4feb6cc-5129-4a99-bb45-851bc07b5709",
561 |    "metadata": {
562 |     "height": 421
563 |    },
564 |    "outputs": [],
565 |    "source": [
566 |     "def query(question, max_turns=5):\n",
567 |     "    i = 0\n",
568 |     "    bot = Agent(prompt)\n",
569 |     "    next_prompt = question\n",
570 |     "    while i < max_turns:\n",
571 |     "        i += 1\n",
572 |     "        result = bot(next_prompt)\n",
573 |     "        print(result)\n",
574 |     "        actions = [action_re.match(a) for a in result.split(\"\\n\") if action_re.match(a)]\n",
575 |     "        if actions:\n",
576 |     "            # There is an action to run\n",
577 |     "            action, action_input = actions[0].groups()\n",
578 |     "            if action not in known_actions:\n",
579 |     "                raise Exception(\"Unknown action: {}: {}\".format(action, action_input))\n",
580 |     "            print(\" -- running {} {}\".format(action, action_input))\n",
581 |     "            observation = known_actions[action](action_input)\n",
582 |     "            print(\"Observation:\", observation)\n",
583 |     "            next_prompt = \"Observation: {}\".format(observation)\n",
584 |     "        else:\n",
585 |     "            return"
586 |    ]
587 |   },
588 |   {
589 |    "cell_type": "markdown",
590 |    "id": "24d50e0f",
591 |    "metadata": {},
592 |    "source": [
593 |     "## Final Evaluation\n",
594 |     "\n",
595 |     "To conclude, we'll test our fully-implemented agent with a more complex query requiring multiple steps of reasoning and action. We'll task it with calculating the combined weight of two distinct dog breeds.\n",
596 |     "\n",
597 |     "This comprehensive test will showcase the agent's ability to:\n",
598 |     "\n",
599 |     "1. Deconstruct a complex query into manageable sub-tasks\n",
600 |     "2. Retrieve information for multiple breeds\n",
601 |     "3. Perform calculations using the gathered data\n",
602 |     "4. Integrate all of this information into a coherent final response\n",
603 |     "\n",
604 |     "By working through this practical example, you'll gain valuable insights into the construction of AI agents capable of solving multi-step problems. Moreover, you'll see firsthand how Amazon Bedrock and model providers like Anthropic's Claude can be effectively utilized. This knowledge will empower you to develop more flexible and diverse AI applications in your future projects.\n"
605 |    ]
606 |   },
607 |   {
608 |    "cell_type": "code",
609 |    "execution_count": null,
610 |    "id": "e85a02b4-96cc-4b01-8792-397a774eb499",
611 |    "metadata": {
612 |     "height": 64
613 |    },
614 |    "outputs": [],
615 |    "source": [
616 |     "question = \"\"\"I have 2 dogs, a border collie and a scottish terrier. \\\n",
617 |     "What is their combined weight\"\"\"\n",
618 |     "query(question)"
619 |    ]
620 |   },
621 |   {
622 |    "cell_type": "markdown",
623 |    "id": "df1b1e6b",
624 |    "metadata": {},
625 |    "source": [
626 |     "# Exercise - Rewrite the Agent to use Anthropic style prompting!\n"
627 |    ]
628 |   },
629 |   {
630 |    "cell_type": "code",
631 |    "execution_count": null,
632 |    "id": "ae8b86a6-5e20-4252-b1d8-009b8318345a",
633 |    "metadata": {
634 |     "height": 30
635 |    },
636 |    "outputs": [],
637 |    "source": []
638 |   }
639 |  ],
640 |  "metadata": {
641 |   "kernelspec": {
642 |    "display_name": "agents-dev-env",
643 |    "language": "python",
644 |    "name": "agents-dev-env"
645 |   },
646 |   "language_info": {
647 |    "codemirror_mode": {
648 |     "name": "ipython",
649 |     "version": 3
650 |    },
651 |    "file_extension": ".py",
652 |    "mimetype": "text/x-python",
653 |    "name": "python",
654 |    "nbconvert_exporter": "python",
655 |    "pygments_lexer": "ipython3",
656 |    "version": "3.10.14"
657 |   }
658 |  },
659 |  "nbformat": 4,
660 |  "nbformat_minor": 5
661 | }
662 | 


--------------------------------------------------------------------------------
/Lab_6/helper.py:
--------------------------------------------------------------------------------
  1 | from dotenv import load_dotenv
  2 | import os
  3 | import sys
  4 | sys.path.append('../utils')
  5 | import utils
  6 | 
  7 | # Load environment variables from .env file or Secret Manager
  8 | _ = load_dotenv("../.env")
  9 | aws_region = os.getenv("AWS_REGION")
 10 | tavily_ai_api_key = utils.get_tavily_api("TAVILY_API_KEY", aws_region)
 11 | 
 12 | import warnings
 13 | warnings.filterwarnings("ignore", message=".*TqdmWarning.*")
 14 | 
 15 | from langgraph.graph import StateGraph, END
 16 | from typing import TypedDict, Annotated, List
 17 | import operator
 18 | from langgraph.checkpoint.sqlite import SqliteSaver
 19 | from langchain_core.messages import (
 20 |     AnyMessage,
 21 |     SystemMessage,
 22 |     HumanMessage,
 23 |     AIMessage,
 24 |     ChatMessage,
 25 | )
 26 | 
 27 | import boto3
 28 | from langchain_openai import ChatOpenAI
 29 | from langchain_aws import ChatBedrockConverse
 30 | from langchain_core.pydantic_v1 import BaseModel
 31 | from tavily import TavilyClient
 32 | import os
 33 | import sqlite3
 34 | 
 35 | 
 36 | # for the output parser
 37 | from typing import List
 38 | from langchain.output_parsers import PydanticOutputParser
 39 | from langchain_core.prompts import PromptTemplate
 40 | from langchain_core.pydantic_v1 import BaseModel, Field
 41 | import json
 42 | 
 43 | 
 44 | class AgentState(TypedDict):
 45 |     task: str
 46 |     lnode: str
 47 |     plan: str
 48 |     draft: str
 49 |     critique: str
 50 |     content: List[str]
 51 |     queries: List[str]
 52 |     revision_number: int
 53 |     max_revisions: int
 54 |     count: Annotated[int, operator.add]
 55 | 
 56 | 
 57 | class Queries(BaseModel):
 58 |     queries: List[str] = Field(description="List of research queries")
 59 | 
 60 | 
 61 | class ewriter:
 62 |     def __init__(self):
 63 | 
 64 |         self.bedrock_rt = boto3.client("bedrock-runtime", region_name=aws_region)
 65 |         self.tavily = TavilyClient(api_key=tavily_ai_api_key)
 66 |         self.model = ChatBedrockConverse(
 67 |             client=self.bedrock_rt,
 68 |             model_id="anthropic.claude-3-haiku-20240307-v1:0",
 69 |             temperature=0,
 70 |             max_tokens=None,
 71 |         )
 72 | 
 73 |         # self.model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
 74 |         self.PLAN_PROMPT = (
 75 |             "You are an expert writer tasked with writing a high level outline of a short 3 paragraph essay. "
 76 |             "Write such an outline for the user provided topic. Give the three main headers of an outline of "
 77 |             "the essay along with any relevant notes or instructions for the sections. "
 78 |         )
 79 |         self.WRITER_PROMPT = """You are an essay assistant tasked with writing excellent 5-paragraph essays.\
 80 |             Generate the best essay possible for the user's request and the initial outline. \
 81 |             If the user provides critique, respond with a revised version of your previous attempts. \
 82 |             Utilize all the information below as needed: 
 83 | 
 84 |             ------
 85 |             <content>
 86 |             {content}
 87 |             </content>"""
 88 |         self.RESEARCH_PLAN_PROMPT = (
 89 |             "You are a researcher charged with providing information that can "
 90 |             "be used when writing the following essay. Generate a list of search "
 91 |             "queries that will gather "
 92 |             "any relevant information. Only generate 3 queries max."
 93 |         )
 94 |         self.REFLECTION_PROMPT = (
 95 |             "You are a teacher grading an 3 paragraph essay submission. "
 96 |             "Generate critique and recommendations for the user's submission. "
 97 |             "Provide detailed recommendations, including requests for length, depth, style, etc."
 98 |         )
 99 |         self.RESEARCH_CRITIQUE_PROMPT = (
100 |             "You are a researcher charged with providing information that can "
101 |             "be used when making any requested revisions (as outlined below). "
102 |             "Generate a list of search queries that will gather any relevant information. "
103 |             "Only generate 2 queries max."
104 |         )
105 |         builder = StateGraph(AgentState)
106 |         builder.add_node("planner", self.plan_node)
107 |         builder.add_node("research_plan", self.research_plan_node)
108 |         builder.add_node("generate", self.generation_node)
109 |         builder.add_node("reflect", self.reflection_node)
110 |         builder.add_node("research_critique", self.research_critique_node)
111 |         builder.set_entry_point("planner")
112 |         builder.add_conditional_edges(
113 |             "generate", self.should_continue, {END: END, "reflect": "reflect"}
114 |         )
115 |         builder.add_edge("planner", "research_plan")
116 |         builder.add_edge("research_plan", "generate")
117 |         builder.add_edge("reflect", "research_critique")
118 |         builder.add_edge("research_critique", "generate")
119 |         memory = SqliteSaver(conn=sqlite3.connect(":memory:", check_same_thread=False))
120 |         self.graph = builder.compile(
121 |             checkpointer=memory,
122 |             interrupt_after=[
123 |                 "planner",
124 |                 "generate",
125 |                 "reflect",
126 |                 "research_plan",
127 |                 "research_critique",
128 |             ],
129 |         )
130 | 
131 |     def plan_node(self, state: AgentState):
132 |         messages = [
133 |             SystemMessage(content=self.PLAN_PROMPT),
134 |             HumanMessage(content=state["task"]),
135 |         ]
136 |         response = self.model.invoke(messages)
137 |         return {
138 |             "plan": response.content,
139 |             "lnode": "planner",
140 |             "count": 1,
141 |         }
142 | 
143 |     def research_plan_node(self, state: AgentState):
144 |         # Set up the Pydantic output parser
145 |         parser = PydanticOutputParser(pydantic_object=Queries)
146 | 
147 |         # Create a prompt template with format instructions
148 |         prompt = PromptTemplate(
149 |             template="Generate research queries based on the given task.\n{format_instructions}\nTask: {task}\n",
150 |             input_variables=["task"],
151 |             partial_variables={"format_instructions": parser.get_format_instructions()},
152 |         )
153 | 
154 |         # Use the model with the new prompt and parser
155 |         queries_output = self.model.invoke(prompt.format_prompt(task=state["task"]))
156 | 
157 |         # Extract the content from the AIMessage
158 |         queries_text = queries_output.content
159 | 
160 |         # Extract the JSON string from the content
161 |         json_start = queries_text.find("{")
162 |         json_end = queries_text.rfind("}") + 1
163 |         json_str = queries_text[json_start:json_end]
164 | 
165 |         # Parse the JSON string
166 |         queries_dict = json.loads(json_str)
167 | 
168 |         # Create a Queries object from the parsed JSON
169 |         parsed_queries = Queries(**queries_dict)
170 | 
171 |         content = state["content"] or []
172 |         for q in parsed_queries.queries:
173 |             response = self.tavily.search(query=q, max_results=2)
174 |             for r in response["results"]:
175 |                 content.append(r["content"])
176 |         return {
177 |             "content": content,
178 |             "queries": parsed_queries.queries,
179 |             "lnode": "research_plan",
180 |             "count": 1,
181 |         }
182 | 
183 |     def generation_node(self, state: AgentState):
184 |         content = "\n\n".join(state["content"] or [])
185 |         user_message = HumanMessage(
186 |             content=f"{state['task']}\n\nHere is my plan:\n\n{state['plan']}"
187 |         )
188 |         messages = [
189 |             SystemMessage(content=self.WRITER_PROMPT.format(content=content)),
190 |             user_message,
191 |         ]
192 |         response = self.model.invoke(messages)
193 |         return {
194 |             "draft": response.content,
195 |             "revision_number": state.get("revision_number", 1) + 1,
196 |             "lnode": "generate",
197 |             "count": 1,
198 |         }
199 | 
200 |     def reflection_node(self, state: AgentState):
201 |         messages = [
202 |             SystemMessage(content=self.REFLECTION_PROMPT),
203 |             HumanMessage(content=state["draft"]),
204 |         ]
205 |         response = self.model.invoke(messages)
206 |         return {
207 |             "critique": response.content,
208 |             "lnode": "reflect",
209 |             "count": 1,
210 |         }
211 | 
212 |     def research_critique_node(self, state: AgentState):
213 |         # Set up the Pydantic output parser
214 |         parser = PydanticOutputParser(pydantic_object=Queries)
215 | 
216 |         # Create a prompt template with format instructions
217 |         prompt = PromptTemplate(
218 |             template="Generate research queries based on the given critique.\n{format_instructions}\nCritique: {critique}\n",
219 |             input_variables=["critique"],
220 |             partial_variables={"format_instructions": parser.get_format_instructions()},
221 |         )
222 | 
223 |         # Use the model with the new prompt and parser
224 |         queries_output = self.model.invoke(
225 |             prompt.format_prompt(critique=state["critique"])
226 |         )
227 | 
228 |         # Extract the content from the AIMessage
229 |         queries_text = queries_output.content
230 | 
231 |         # Extract the JSON string from the content
232 |         json_start = queries_text.find("{")
233 |         json_end = queries_text.rfind("}") + 1
234 |         json_str = queries_text[json_start:json_end]
235 | 
236 |         # Parse the JSON string
237 |         queries_dict = json.loads(json_str)
238 | 
239 |         # Create a Queries object from the parsed JSON
240 |         parsed_queries = Queries(**queries_dict)
241 | 
242 |         content = state["content"] or []
243 |         for q in parsed_queries.queries:
244 |             response = self.tavily.search(query=q, max_results=2)
245 |             for r in response["results"]:
246 |                 content.append(r["content"])
247 |         return {
248 |             "content": content,
249 |             "lnode": "research_critique",
250 |             "count": 1,
251 |         }
252 | 
253 |     def should_continue(self, state):
254 |         if state["revision_number"] > state["max_revisions"]:
255 |             return END
256 |         return "reflect"
257 | 
258 | 
259 | import gradio as gr
260 | import time
261 | 
262 | 
263 | class writer_gui:
264 |     def __init__(self, graph):
265 |         self.graph = graph
266 |         self.partial_message = ""
267 |         self.response = {}
268 |         self.max_iterations = 10
269 |         self.iterations = []
270 |         self.threads = []
271 |         self.thread_id = -1
272 |         self.thread = {"configurable": {"thread_id": str(self.thread_id)}}
273 |         # self.sdisps = {} #global
274 |         self.demo = self.create_interface()
275 | 
276 |     def run_agent(self, start, topic, stop_after):
277 |         # global partial_message, thread_id,thread
278 |         # global response, max_iterations, iterations, threads
279 |         if start:
280 |             self.iterations.append(0)
281 |             config = {
282 |                 "task": topic,
283 |                 "max_revisions": 2,
284 |                 "revision_number": 0,
285 |                 "lnode": "",
286 |                 "planner": "no plan",
287 |                 "draft": "no draft",
288 |                 "critique": "no critique",
289 |                 "content": [
290 |                     "no content",
291 |                 ],
292 |                 "queries": "no queries",
293 |                 "count": 0,
294 |             }
295 |             self.thread_id += 1  # new agent, new thread
296 |             self.threads.append(self.thread_id)
297 |         else:
298 |             config = None
299 |         self.thread = {"configurable": {"thread_id": str(self.thread_id)}}
300 |         while self.iterations[self.thread_id] < self.max_iterations:
301 |             self.response = self.graph.invoke(config, self.thread)
302 |             self.iterations[self.thread_id] += 1
303 |             self.partial_message += str(self.response)
304 |             self.partial_message += f"\n------------------\n\n"
305 |             ## fix
306 |             lnode, nnode, _, rev, acount = self.get_disp_state()
307 |             yield self.partial_message, lnode, nnode, self.thread_id, rev, acount
308 |             config = None  # need
309 |             # print(f"run_agent:{lnode}")
310 |             if not nnode:
311 |                 # print("Hit the end")
312 |                 return
313 |             if lnode in stop_after:
314 |                 # print(f"stopping due to stop_after {lnode}")
315 |                 return
316 |             else:
317 |                 # print(f"Not stopping on lnode {lnode}")
318 |                 pass
319 |         return
320 | 
321 |     def get_disp_state(
322 |         self,
323 |     ):
324 |         current_state = self.graph.get_state(self.thread)
325 |         lnode = current_state.values["lnode"]
326 |         acount = current_state.values["count"]
327 |         rev = current_state.values["revision_number"]
328 |         nnode = current_state.next
329 |         # print  (lnode,nnode,self.thread_id,rev,acount)
330 |         return lnode, nnode, self.thread_id, rev, acount
331 | 
332 |     def get_state(self, key):
333 |         current_values = self.graph.get_state(self.thread)
334 |         if key in current_values.values:
335 |             lnode, nnode, self.thread_id, rev, astep = self.get_disp_state()
336 |             new_label = f"last_node: {lnode}, thread_id: {self.thread_id}, rev: {rev}, step: {astep}"
337 |             return gr.update(label=new_label, value=current_values.values[key])
338 |         else:
339 |             return ""
340 | 
341 |     def get_content(
342 |         self,
343 |     ):
344 |         current_values = self.graph.get_state(self.thread)
345 |         if "content" in current_values.values:
346 |             content = current_values.values["content"]
347 |             lnode, nnode, thread_id, rev, astep = self.get_disp_state()
348 |             new_label = f"last_node: {lnode}, thread_id: {self.thread_id}, rev: {rev}, step: {astep}"
349 |             return gr.update(
350 |                 label=new_label, value="\n\n".join(item for item in content) + "\n\n"
351 |             )
352 |         else:
353 |             return ""
354 | 
355 |     def update_hist_pd(
356 |         self,
357 |     ):
358 |         # print("update_hist_pd")
359 |         hist = []
360 |         # curiously, this generator returns the latest first
361 |         for state in self.graph.get_state_history(self.thread):
362 |             if state.metadata["step"] < 1:
363 |                 continue
364 |             thread_ts = state.config["configurable"]["thread_ts"]
365 |             tid = state.config["configurable"]["thread_id"]
366 |             count = state.values["count"]
367 |             lnode = state.values["lnode"]
368 |             rev = state.values["revision_number"]
369 |             nnode = state.next
370 |             st = f"{tid}:{count}:{lnode}:{nnode}:{rev}:{thread_ts}"
371 |             hist.append(st)
372 |         return gr.Dropdown(
373 |             label="update_state from: thread:count:last_node:next_node:rev:thread_ts",
374 |             choices=hist,
375 |             value=hist[0],
376 |             interactive=True,
377 |         )
378 | 
379 |     def find_config(self, thread_ts):
380 |         for state in self.graph.get_state_history(self.thread):
381 |             config = state.config
382 |             if config["configurable"]["thread_ts"] == thread_ts:
383 |                 return config
384 |         return None
385 | 
386 |     def copy_state(self, hist_str):
387 |         """result of selecting an old state from the step pulldown. Note does not change thread.
388 |         This copies an old state to a new current state.
389 |         """
390 |         thread_ts = hist_str.split(":")[-1]
391 |         # print(f"copy_state from {thread_ts}")
392 |         config = self.find_config(thread_ts)
393 |         # print(config)
394 |         state = self.graph.get_state(config)
395 |         self.graph.update_state(
396 |             self.thread, state.values, as_node=state.values["lnode"]
397 |         )
398 |         new_state = self.graph.get_state(self.thread)  # should now match
399 |         new_thread_ts = new_state.config["configurable"]["thread_ts"]
400 |         tid = new_state.config["configurable"]["thread_id"]
401 |         count = new_state.values["count"]
402 |         lnode = new_state.values["lnode"]
403 |         rev = new_state.values["revision_number"]
404 |         nnode = new_state.next
405 |         return lnode, nnode, new_thread_ts, rev, count
406 | 
407 |     def update_thread_pd(
408 |         self,
409 |     ):
410 |         # print("update_thread_pd")
411 |         return gr.Dropdown(
412 |             label="choose thread",
413 |             choices=threads,
414 |             value=self.thread_id,
415 |             interactive=True,
416 |         )
417 | 
418 |     def switch_thread(self, new_thread_id):
419 |         # print(f"switch_thread{new_thread_id}")
420 |         self.thread = {"configurable": {"thread_id": str(new_thread_id)}}
421 |         self.thread_id = new_thread_id
422 |         return
423 | 
424 |     def modify_state(self, key, asnode, new_state):
425 |         """gets the current state, modifes a single value in the state identified by key, and updates state with it.
426 |         note that this will create a new 'current state' node. If you do this multiple times with different keys, it will create
427 |         one for each update. Note also that it doesn't resume after the update
428 |         """
429 |         current_values = self.graph.get_state(self.thread)
430 |         current_values.values[key] = new_state
431 |         self.graph.update_state(self.thread, current_values.values, as_node=asnode)
432 |         return
433 | 
434 |     def create_interface(self):
435 |         with gr.Blocks(
436 |             theme=gr.themes.Default(spacing_size="sm", text_size="sm"),
437 |             analytics_enabled=False 
438 |         ) as demo:
439 | 
440 |             def updt_disp():
441 |                 """general update display on state change"""
442 |                 current_state = self.graph.get_state(self.thread)
443 |                 hist = []
444 |                 # curiously, this generator returns the latest first
445 |                 for state in self.graph.get_state_history(self.thread):
446 |                     if state.metadata["step"] < 1:  # ignore early states
447 |                         continue
448 |                     s_thread_ts = state.config["configurable"]["thread_ts"]
449 |                     s_tid = state.config["configurable"]["thread_id"]
450 |                     s_count = state.values["count"]
451 |                     s_lnode = state.values["lnode"]
452 |                     s_rev = state.values["revision_number"]
453 |                     s_nnode = state.next
454 |                     st = f"{s_tid}:{s_count}:{s_lnode}:{s_nnode}:{s_rev}:{s_thread_ts}"
455 |                     hist.append(st)
456 |                 if not current_state.metadata:  # handle init call
457 |                     return {}
458 |                 else:
459 |                     return {
460 |                         topic_bx: current_state.values["task"],
461 |                         lnode_bx: current_state.values["lnode"],
462 |                         count_bx: current_state.values["count"],
463 |                         revision_bx: current_state.values["revision_number"],
464 |                         nnode_bx: current_state.next,
465 |                         threadid_bx: self.thread_id,
466 |                         thread_pd: gr.Dropdown(
467 |                             label="choose thread",
468 |                             choices=self.threads,
469 |                             value=self.thread_id,
470 |                             interactive=True,
471 |                         ),
472 |                         step_pd: gr.Dropdown(
473 |                             label="update_state from: thread:count:last_node:next_node:rev:thread_ts",
474 |                             choices=hist,
475 |                             value=hist[0],
476 |                             interactive=True,
477 |                         ),
478 |                     }
479 | 
480 |             def get_snapshots():
481 |                 new_label = f"thread_id: {self.thread_id}, Summary of snapshots"
482 |                 sstate = ""
483 |                 for state in self.graph.get_state_history(self.thread):
484 |                     for key in ["plan", "draft", "critique"]:
485 |                         if key in state.values:
486 |                             state.values[key] = state.values[key][:80] + "..."
487 |                     if "content" in state.values:
488 |                         for i in range(len(state.values["content"])):
489 |                             state.values["content"][i] = (
490 |                                 state.values["content"][i][:20] + "..."
491 |                             )
492 |                     if "writes" in state.metadata:
493 |                         state.metadata["writes"] = "not shown"
494 |                     sstate += str(state) + "\n\n"
495 |                 return gr.update(label=new_label, value=sstate)
496 | 
497 |             def vary_btn(stat):
498 |                 # print(f"vary_btn{stat}")
499 |                 return gr.update(variant=stat)
500 | 
501 |             with gr.Tab("Agent"):
502 |                 with gr.Row():
503 |                     topic_bx = gr.Textbox(label="Essay Topic", value="Pizza Shop")
504 |                     gen_btn = gr.Button(
505 |                         "Generate Essay", scale=0, min_width=80, variant="primary"
506 |                     )
507 |                     cont_btn = gr.Button("Continue Essay", scale=0, min_width=80)
508 |                 with gr.Row():
509 |                     lnode_bx = gr.Textbox(label="last node", min_width=100)
510 |                     nnode_bx = gr.Textbox(label="next node", min_width=100)
511 |                     threadid_bx = gr.Textbox(label="Thread", scale=0, min_width=80)
512 |                     revision_bx = gr.Textbox(label="Draft Rev", scale=0, min_width=80)
513 |                     count_bx = gr.Textbox(label="count", scale=0, min_width=80)
514 |                 with gr.Accordion("Manage Agent", open=False):
515 |                     checks = list(self.graph.nodes.keys())
516 |                     checks.remove("__start__")
517 |                     stop_after = gr.CheckboxGroup(
518 |                         checks,
519 |                         label="Interrupt After State",
520 |                         value=checks,
521 |                         scale=0,
522 |                         min_width=400,
523 |                     )
524 |                     with gr.Row():
525 |                         thread_pd = gr.Dropdown(
526 |                             choices=self.threads,
527 |                             interactive=True,
528 |                             label="select thread",
529 |                             min_width=120,
530 |                             scale=0,
531 |                         )
532 |                         step_pd = gr.Dropdown(
533 |                             choices=["N/A"],
534 |                             interactive=True,
535 |                             label="select step",
536 |                             min_width=160,
537 |                             scale=1,
538 |                         )
539 |                 live = gr.Textbox(label="Live Agent Output", lines=5, max_lines=5)
540 | 
541 |                 # actions
542 |                 sdisps = [
543 |                     topic_bx,
544 |                     lnode_bx,
545 |                     nnode_bx,
546 |                     threadid_bx,
547 |                     revision_bx,
548 |                     count_bx,
549 |                     step_pd,
550 |                     thread_pd,
551 |                 ]
552 |                 thread_pd.input(self.switch_thread, [thread_pd], None).then(
553 |                     fn=updt_disp, inputs=None, outputs=sdisps
554 |                 )
555 |                 step_pd.input(self.copy_state, [step_pd], None).then(
556 |                     fn=updt_disp, inputs=None, outputs=sdisps
557 |                 )
558 |                 gen_btn.click(
559 |                     vary_btn, gr.Number("secondary", visible=False), gen_btn
560 |                 ).then(
561 |                     fn=self.run_agent,
562 |                     inputs=[gr.Number(True, visible=False), topic_bx, stop_after],
563 |                     outputs=[live],
564 |                     show_progress=True,
565 |                 ).then(
566 |                     fn=updt_disp, inputs=None, outputs=sdisps
567 |                 ).then(
568 |                     vary_btn, gr.Number("primary", visible=False), gen_btn
569 |                 ).then(
570 |                     vary_btn, gr.Number("primary", visible=False), cont_btn
571 |                 )
572 |                 cont_btn.click(
573 |                     vary_btn, gr.Number("secondary", visible=False), cont_btn
574 |                 ).then(
575 |                     fn=self.run_agent,
576 |                     inputs=[gr.Number(False, visible=False), topic_bx, stop_after],
577 |                     outputs=[live],
578 |                 ).then(
579 |                     fn=updt_disp, inputs=None, outputs=sdisps
580 |                 ).then(
581 |                     vary_btn, gr.Number("primary", visible=False), cont_btn
582 |                 )
583 | 
584 |             with gr.Tab("Plan"):
585 |                 with gr.Row():
586 |                     refresh_btn = gr.Button("Refresh")
587 |                     modify_btn = gr.Button("Modify")
588 |                 plan = gr.Textbox(label="Plan", lines=10, interactive=True)
589 |                 refresh_btn.click(
590 |                     fn=self.get_state,
591 |                     inputs=gr.Number("plan", visible=False),
592 |                     outputs=plan,
593 |                 )
594 |                 modify_btn.click(
595 |                     fn=self.modify_state,
596 |                     inputs=[
597 |                         gr.Number("plan", visible=False),
598 |                         gr.Number("planner", visible=False),
599 |                         plan,
600 |                     ],
601 |                     outputs=None,
602 |                 ).then(fn=updt_disp, inputs=None, outputs=sdisps)
603 |             with gr.Tab("Research Content"):
604 |                 refresh_btn = gr.Button("Refresh")
605 |                 content_bx = gr.Textbox(label="content", lines=10)
606 |                 refresh_btn.click(fn=self.get_content, inputs=None, outputs=content_bx)
607 |             with gr.Tab("Draft"):
608 |                 with gr.Row():
609 |                     refresh_btn = gr.Button("Refresh")
610 |                     modify_btn = gr.Button("Modify")
611 |                 draft_bx = gr.Textbox(label="draft", lines=10, interactive=True)
612 |                 refresh_btn.click(
613 |                     fn=self.get_state,
614 |                     inputs=gr.Number("draft", visible=False),
615 |                     outputs=draft_bx,
616 |                 )
617 |                 modify_btn.click(
618 |                     fn=self.modify_state,
619 |                     inputs=[
620 |                         gr.Number("draft", visible=False),
621 |                         gr.Number("generate", visible=False),
622 |                         draft_bx,
623 |                     ],
624 |                     outputs=None,
625 |                 ).then(fn=updt_disp, inputs=None, outputs=sdisps)
626 |             with gr.Tab("Critique"):
627 |                 with gr.Row():
628 |                     refresh_btn = gr.Button("Refresh")
629 |                     modify_btn = gr.Button("Modify")
630 |                 critique_bx = gr.Textbox(label="Critique", lines=10, interactive=True)
631 |                 refresh_btn.click(
632 |                     fn=self.get_state,
633 |                     inputs=gr.Number("critique", visible=False),
634 |                     outputs=critique_bx,
635 |                 )
636 |                 modify_btn.click(
637 |                     fn=self.modify_state,
638 |                     inputs=[
639 |                         gr.Number("critique", visible=False),
640 |                         gr.Number("reflect", visible=False),
641 |                         critique_bx,
642 |                     ],
643 |                     outputs=None,
644 |                 ).then(fn=updt_disp, inputs=None, outputs=sdisps)
645 |             with gr.Tab("StateSnapShots"):
646 |                 with gr.Row():
647 |                     refresh_btn = gr.Button("Refresh")
648 |                 snapshots = gr.Textbox(label="State Snapshots Summaries")
649 |                 refresh_btn.click(fn=get_snapshots, inputs=None, outputs=snapshots)
650 |         return demo
651 | 
652 |     def launch(self):
653 |             self.demo.launch(share=True)
654 | 
655 | 
656 | if __name__ == "__main__":
657 |     MultiAgent = ewriter()
658 |     app = writer_gui(MultiAgent.graph)
659 |     app.launch()
660 | 


--------------------------------------------------------------------------------
/Lab_6/Lab_6.ipynb:
--------------------------------------------------------------------------------
  1 | {
  2 |  "cells": [
  3 |   {
  4 |    "cell_type": "markdown",
  5 |    "id": "911b3b37-3b29-4833-94f2-bfe47af00c83",
  6 |    "metadata": {},
  7 |    "source": [
  8 |     "# Lab 6: Essay Writer\n",
  9 |     "\n",
 10 |     "## Setup and Imports\n",
 11 |     "\n",
 12 |     "In this section, we're building a more complex project: an AI Essay Writer. We'll start by setting up our environment and importing the necessary libraries. We're using Amazon Bedrock and Anthropic's Claude model, so our imports reflect that.\n",
 13 |     "We're setting up logging, configuring Bedrock, and retrieving our Tavily API key. Tavily is a research tool we'll use to gather information for our essays. Make sure you have your Tavily API key stored securely.\n",
 14 |     "\n",
 15 |     "A few requirements for the UI that we build at the end:\n",
 16 |     "\n",
 17 |     "1. Add your tavily ai key to the `.env` file.\n",
 18 |     "\n",
 19 |     "2. Pay attention that you run on LangGraph 0.0.53.\n",
 20 |     "\n",
 21 |     "3. Run the helper.py from the CLI and open the web browser to step through the graph or run it in the notebook directly!\n"
 22 |    ]
 23 |   },
 24 |   {
 25 |    "cell_type": "code",
 26 |    "execution_count": null,
 27 |    "id": "f5762271-8736-4e94-9444-8c92bd0e8074",
 28 |    "metadata": {
 29 |     "height": 64
 30 |    },
 31 |    "outputs": [],
 32 |    "source": [
 33 |     "from dotenv import load_dotenv\n",
 34 |     "import os\n",
 35 |     "import sys\n",
 36 |     "import json, re\n",
 37 |     "import pprint\n",
 38 |     "import boto3\n",
 39 |     "from botocore.client import Config\n",
 40 |     "import warnings\n",
 41 |     "\n",
 42 |     "warnings.filterwarnings(\"ignore\")\n",
 43 |     "import logging\n",
 44 |     "\n",
 45 |     "# import local modules\n",
 46 |     "dir_current = os.path.abspath(\"\")\n",
 47 |     "dir_parent = os.path.dirname(dir_current)\n",
 48 |     "if dir_parent not in sys.path:\n",
 49 |     "    sys.path.append(dir_parent)\n",
 50 |     "from utils import utils\n",
 51 |     "\n",
 52 |     "# Set basic configs\n",
 53 |     "logger = utils.set_logger()\n",
 54 |     "pp = utils.set_pretty_printer()\n",
 55 |     "\n",
 56 |     "# Load environment variables from .env file or Secret Manager\n",
 57 |     "_ = load_dotenv(\"../.env\")\n",
 58 |     "aws_region = os.getenv(\"AWS_REGION\")\n",
 59 |     "tavily_ai_api_key = utils.get_tavily_api(\"TAVILY_API_KEY\", aws_region)\n",
 60 |     "\n",
 61 |     "# Set bedrock configs\n",
 62 |     "bedrock_config = Config(\n",
 63 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
 64 |     ")\n",
 65 |     "\n",
 66 |     "# Create a bedrock runtime client\n",
 67 |     "bedrock_rt = boto3.client(\n",
 68 |     "    \"bedrock-runtime\", region_name=aws_region, config=bedrock_config\n",
 69 |     ")\n"
 70 |    ]
 71 |   },
 72 |   {
 73 |    "cell_type": "markdown",
 74 |    "id": "252ba3e5",
 75 |    "metadata": {},
 76 |    "source": [
 77 |     "## Defining the Agent State\n",
 78 |     "\n",
 79 |     "Now, let's define our agent's state. This is more complex than in previous lessons because our essay writer has multiple steps and needs to keep track of various pieces of information.\n",
 80 |     "\n",
 81 |     "We're creating a TypedDict called AgentState. It includes:\n",
 82 |     "\n",
 83 |     "- task: The essay topic or question\n",
 84 |     "- plan: The outline of the essay\n",
 85 |     "- draft: The current version of the essay\n",
 86 |     "- critique: Feedback on the current draft\n",
 87 |     "- content: Research information from Tavily\n",
 88 |     "- revision_number: How many revisions we've made\n",
 89 |     "- max_revisions: The maximum number of revisions we want to make\n",
 90 |     "\n",
 91 |     "These elements will help us manage the essay writing process and know when to stop revising.\n"
 92 |    ]
 93 |   },
 94 |   {
 95 |    "cell_type": "code",
 96 |    "execution_count": null,
 97 |    "id": "d0168aee-bce9-4d60-b827-f86a88187e31",
 98 |    "metadata": {
 99 |     "height": 132
100 |    },
101 |    "outputs": [],
102 |    "source": [
103 |     "from langgraph.graph import StateGraph, END\n",
104 |     "from typing import TypedDict, Annotated, List\n",
105 |     "import operator\n",
106 |     "from langgraph.checkpoint.memory import MemorySaver\n",
107 |     "\n",
108 |     "from langchain_core.messages import (\n",
109 |     "    AnyMessage,\n",
110 |     "    SystemMessage,\n",
111 |     "    HumanMessage,\n",
112 |     "    AIMessage,\n",
113 |     "    ChatMessage,\n",
114 |     ")\n",
115 |     "\n",
116 |     "memory = MemorySaver()"
117 |    ]
118 |   },
119 |   {
120 |    "cell_type": "code",
121 |    "execution_count": null,
122 |    "id": "2589c5b6-6cc2-4594-9a17-dccdcf676054",
123 |    "metadata": {
124 |     "height": 149
125 |    },
126 |    "outputs": [],
127 |    "source": [
128 |     "class AgentState(TypedDict):\n",
129 |     "    task: str\n",
130 |     "    plan: str\n",
131 |     "    draft: str\n",
132 |     "    critique: str\n",
133 |     "    content: List[str]\n",
134 |     "    revision_number: int\n",
135 |     "    max_revisions: int"
136 |    ]
137 |   },
138 |   {
139 |    "cell_type": "markdown",
140 |    "id": "d6aac3dc",
141 |    "metadata": {},
142 |    "source": [
143 |     "## Setting up the Model\n",
144 |     "\n",
145 |     "We're using Anthropic's Claude model via Amazon Bedrock. We're setting it up with a temperature of 0 for more consistent outputs. The model we're using is claude-3-haiku, which is well-suited for this task.\n"
146 |    ]
147 |   },
148 |   {
149 |    "cell_type": "code",
150 |    "execution_count": null,
151 |    "id": "a2ba84ec-c172-4de7-ac55-e3158a531b23",
152 |    "metadata": {
153 |     "height": 47
154 |    },
155 |    "outputs": [],
156 |    "source": [
157 |     "from langchain_aws import ChatBedrockConverse\n",
158 |     "\n",
159 |     "model = ChatBedrockConverse(\n",
160 |     "    client=bedrock_rt,\n",
161 |     "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
162 |     "    temperature=0,\n",
163 |     "    max_tokens=None,\n",
164 |     ")"
165 |    ]
166 |   },
167 |   {
168 |    "cell_type": "markdown",
169 |    "id": "6d61f1e5",
170 |    "metadata": {},
171 |    "source": [
172 |     "## Defining Prompts\n",
173 |     "\n",
174 |     "Our essay writer uses several prompts for different stages of the process:\n",
175 |     "\n",
176 |     "1. **PLAN_PROMPT**: This instructs the model to create an essay outline.\n",
177 |     "2. **WRITER_PROMPT**: This guides the model in writing the essay based on the plan and research.\n",
178 |     "3. **REFLECTION_PROMPT**: This tells the model how to critique the essay.\n",
179 |     "4. **RESEARCH_PLAN_PROMPT** and **RESEARCH_CRITIQUE_PROMPT**: These help generate search queries for our research step.\n",
180 |     "\n",
181 |     "Each prompt is carefully crafted to guide the model in performing its specific task within the essay writing process.\n"
182 |    ]
183 |   },
184 |   {
185 |    "cell_type": "code",
186 |    "execution_count": null,
187 |    "id": "876d5092-b8ef-4e38-b4d7-0e80c609bf7a",
188 |    "metadata": {
189 |     "height": 74
190 |    },
191 |    "outputs": [],
192 |    "source": [
193 |     "PLAN_PROMPT = \"\"\"You are an expert writer tasked with writing a high level outline of an essay. \\\n",
194 |     "Write such an outline for the user provided topic. Give an outline of the essay along with any relevant notes \\\n",
195 |     "or instructions for the sections.\"\"\""
196 |    ]
197 |   },
198 |   {
199 |    "cell_type": "code",
200 |    "execution_count": null,
201 |    "id": "10084a02-2928-4945-9f7c-ad3f5b33caf7",
202 |    "metadata": {
203 |     "height": 149
204 |    },
205 |    "outputs": [],
206 |    "source": [
207 |     "WRITER_PROMPT = \"\"\"You are an essay assistant tasked with writing excellent 5-paragraph essays.\\\n",
208 |     "Generate the best essay possible for the user's request and the initial outline. \\\n",
209 |     "If the user provides critique, respond with a revised version of your previous attempts. \\\n",
210 |     "Utilize all the information below as needed: \n",
211 |     "\n",
212 |     "------\n",
213 |     "<content>\n",
214 |     "{content}\n",
215 |     "</content>\"\"\""
216 |    ]
217 |   },
218 |   {
219 |    "cell_type": "code",
220 |    "execution_count": null,
221 |    "id": "714d1205-f8fc-4912-b148-2a45da99219c",
222 |    "metadata": {
223 |     "height": 64
224 |    },
225 |    "outputs": [],
226 |    "source": [
227 |     "REFLECTION_PROMPT = \"\"\"You are a teacher grading an essay submission. \\\n",
228 |     "Generate critique and recommendations for the user's submission. \\\n",
229 |     "Provide detailed recommendations, including requests for length, depth, style, etc.\"\"\""
230 |    ]
231 |   },
232 |   {
233 |    "cell_type": "code",
234 |    "execution_count": null,
235 |    "id": "83588e70-254f-4f83-a510-c8ae81e729b0",
236 |    "metadata": {
237 |     "height": 81
238 |    },
239 |    "outputs": [],
240 |    "source": [
241 |     "RESEARCH_PLAN_PROMPT = \"\"\"You are a researcher charged with providing information that can \\\n",
242 |     "be used when writing the following essay. Generate a list of search queries that will gather \\\n",
243 |     "any relevant information. Only generate 3 queries max.\"\"\""
244 |    ]
245 |   },
246 |   {
247 |    "cell_type": "code",
248 |    "execution_count": null,
249 |    "id": "6cb3ef4c-58b3-401b-b104-0d51e553d982",
250 |    "metadata": {
251 |     "height": 81
252 |    },
253 |    "outputs": [],
254 |    "source": [
255 |     "RESEARCH_CRITIQUE_PROMPT = \"\"\"You are a researcher charged with providing information that can \\\n",
256 |     "be used when making any requested revisions (as outlined below). \\\n",
257 |     "Generate a list of search queries that will gather any relevant information. Only generate 3 queries max.\"\"\""
258 |    ]
259 |   },
260 |   {
261 |    "cell_type": "markdown",
262 |    "id": "ddb9c562",
263 |    "metadata": {},
264 |    "source": [
265 |     "### Excursion on strucutred output generation with Anthropic Models:\n",
266 |     "\n",
267 |     "Take a look at the prompts above.\n",
268 |     "\n",
269 |     "- Do they follow the prompting guide of Anthropic?\n",
270 |     "- Do you think you could get more consistent outputs by requesting an answer structure that would look something like this, e.g. for the REFLECTION_PROMPT:\n",
271 |     "\n",
272 |     "```xml\n",
273 |     "<answer>\n",
274 |     "  <overall_assessment>\n",
275 |     "    <strengths>\n",
276 |     "      <strength_point></strength_point>\n",
277 |     "      <strength_point></strength_point>\n",
278 |     "    </strengths>\n",
279 |     "    <weaknesses>\n",
280 |     "      <weakness_point></weakness_point>\n",
281 |     "      <weakness_point></weakness_point>\n",
282 |     "    </weaknesses>\n",
283 |     "  </overall_assessment>\n",
284 |     "\n",
285 |     "...\n",
286 |     "\n",
287 |     "  <style_and_language>\n",
288 |     "    <clarity>\n",
289 |     "      <comment></comment>\n",
290 |     "      <recommendation></recommendation>\n",
291 |     "    </clarity>\n",
292 |     "    <tone>\n",
293 |     "      <comment></comment>\n",
294 |     "      <recommendation></recommendation>\n",
295 |     "    </tone>\n",
296 |     "    <grammar_and_mechanics>\n",
297 |     "      <comment></comment>\n",
298 |     "      <recommendation></recommendation>\n",
299 |     "    </grammar_and_mechanics>\n",
300 |     "  </style_and_language>\n",
301 |     "\n",
302 |     "  <length_assessment>\n",
303 |     "    <comment></comment>\n",
304 |     "    <recommendation></recommendation>\n",
305 |     "  </length_assessment>\n",
306 |     "\n",
307 |     "  <conclusion>\n",
308 |     "    <overall_recommendation></overall_recommendation>\n",
309 |     "    <priority_improvements>\n",
310 |     "      <improvement></improvement>\n",
311 |     "      <improvement></improvement>\n",
312 |     "    </priority_improvements>\n",
313 |     "  </conclusion>\n",
314 |     "</answer>\n",
315 |     "```\n",
316 |     "\n",
317 |     "- What would the advantage and disadvantage of such a structure be?\n",
318 |     "- Ask yourself, if the extra tokens are worth it? When should you invest in a detailed, token intensive prompt, vs. a more freeform one?\n",
319 |     "\n",
320 |     "- How would you parse this output?\n",
321 |     "\n",
322 |     "**Hint:**\n",
323 |     "You can combine XMLOutput-Parser from LangChain and PyDantic models.\n",
324 |     "\n",
325 |     "For reference, below you can find how to use the XMLOutput parser, without any dependencies on methods like `.with_structured_output(...)`, which is a very recent addition to langchain.\n"
326 |    ]
327 |   },
328 |   {
329 |    "cell_type": "code",
330 |    "execution_count": null,
331 |    "id": "35ff3ddc",
332 |    "metadata": {},
333 |    "outputs": [],
334 |    "source": [
335 |     "from langchain_core.output_parsers.xml import XMLOutputParser\n",
336 |     "\n",
337 |     "# Create the XMLOutputParser with our Pydantic model\n",
338 |     "essay_critique_parser = XMLOutputParser()\n",
339 |     "\n",
340 |     "# Example usage\n",
341 |     "xml_string = \"\"\"\n",
342 |     "<answer>\n",
343 |     "  <overall_assessment>\n",
344 |     "    <strengths>\n",
345 |     "      <strength_point>Clear thesis statement</strength_point>\n",
346 |     "      <strength_point>Well-structured paragraphs</strength_point>\n",
347 |     "    </strengths>\n",
348 |     "    <weaknesses>\n",
349 |     "      <weakness_point>Lack of detailed examples</weakness_point>\n",
350 |     "      <weakness_point>Some grammatical errors</weakness_point>\n",
351 |     "    </weaknesses>\n",
352 |     "  </overall_assessment>\n",
353 |     "  <content_evaluation>\n",
354 |     "    <depth_of_analysis>\n",
355 |     "      <comment>The analysis lacks depth in some areas.</comment>\n",
356 |     "      <recommendation>Expand on key points with more detailed explanations.</recommendation>\n",
357 |     "    </depth_of_analysis>\n",
358 |     "    <argument_quality>\n",
359 |     "      <comment>Arguments are logical but could be stronger.</comment>\n",
360 |     "      <recommendation>Provide more evidence to support your claims.</recommendation>\n",
361 |     "    </argument_quality>\n",
362 |     "    <evidence_use>\n",
363 |     "      <comment>Limited use of supporting evidence.</comment>\n",
364 |     "      <recommendation>Incorporate more relevant examples and data.</recommendation>\n",
365 |     "    </evidence_use>\n",
366 |     "  </content_evaluation>\n",
367 |     "  <structure_and_organization>\n",
368 |     "    <comment>The essay has a clear structure but transitions could be improved.</comment>\n",
369 |     "    <recommendation>Work on smoother transitions between paragraphs.</recommendation>\n",
370 |     "  </structure_and_organization>\n",
371 |     "  <style_and_language>\n",
372 |     "    <clarity>\n",
373 |     "      <comment>Writing is generally clear but some sentences are convoluted.</comment>\n",
374 |     "      <recommendation>Simplify complex sentences for better readability.</recommendation>\n",
375 |     "    </clarity>\n",
376 |     "    <tone>\n",
377 |     "      <comment>The tone is appropriate for an academic essay.</comment>\n",
378 |     "      <recommendation>Maintain this formal tone throughout.</recommendation>\n",
379 |     "    </tone>\n",
380 |     "    <grammar_and_mechanics>\n",
381 |     "      <comment>There are a few grammatical errors and typos.</comment>\n",
382 |     "      <recommendation>Proofread carefully to eliminate these errors.</recommendation>\n",
383 |     "    </grammar_and_mechanics>\n",
384 |     "  </style_and_language>\n",
385 |     "  <length_assessment>\n",
386 |     "    <comment>The essay meets the required length.</comment>\n",
387 |     "    <recommendation>No changes needed in terms of length.</recommendation>\n",
388 |     "  </length_assessment>\n",
389 |     "  <conclusion>\n",
390 |     "    <overall_recommendation>This is a solid essay that could be improved with more depth and better proofreading.</overall_recommendation>\n",
391 |     "    <priority_improvements>\n",
392 |     "      <improvement>Deepen analysis with more detailed explanations and examples.</improvement>\n",
393 |     "      <improvement>Carefully proofread to eliminate grammatical errors and typos.</improvement>\n",
394 |     "    </priority_improvements>\n",
395 |     "  </conclusion>\n",
396 |     "</answer>\n",
397 |     "\"\"\"\n",
398 |     "\n",
399 |     "# Parse the XML string\n",
400 |     "parsed_critique = essay_critique_parser.parse(xml_string)\n",
401 |     "parsed_critique"
402 |    ]
403 |   },
404 |   {
405 |    "cell_type": "markdown",
406 |    "id": "04d75275",
407 |    "metadata": {},
408 |    "source": [
409 |     "However, we can also use `.with_structured_output` to get the answer we need!\n"
410 |    ]
411 |   },
412 |   {
413 |    "cell_type": "code",
414 |    "execution_count": null,
415 |    "id": "0cefa30b",
416 |    "metadata": {},
417 |    "outputs": [],
418 |    "source": [
419 |     "from langchain_aws import ChatBedrock\n",
420 |     "from langchain_core.pydantic_v1 import BaseModel, Field\n",
421 |     "\n",
422 |     "\n",
423 |     "class StructuredOutput(BaseModel):\n",
424 |     "    title: str = Field(..., description=\"The title of the response\")\n",
425 |     "    content: str = Field(..., description=\"The main content of the response\")\n",
426 |     "    summary: str = Field(..., description=\"A brief summary of the content\")\n",
427 |     "\n",
428 |     "\n",
429 |     "llm = ChatBedrock(\n",
430 |     "    model_id=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
431 |     "    model_kwargs={\"temperature\": 0},\n",
432 |     ")\n",
433 |     "\n",
434 |     "structured_llm = llm.with_structured_output(StructuredOutput)\n",
435 |     "\n",
436 |     "response = structured_llm.invoke(\"Tell me about artificial intelligence\")\n",
437 |     "response"
438 |    ]
439 |   },
440 |   {
441 |    "cell_type": "code",
442 |    "execution_count": null,
443 |    "id": "ea7c615e",
444 |    "metadata": {},
445 |    "outputs": [],
446 |    "source": [
447 |     "print(response.title)"
448 |    ]
449 |   },
450 |   {
451 |    "cell_type": "code",
452 |    "execution_count": null,
453 |    "id": "0386ba8b",
454 |    "metadata": {},
455 |    "outputs": [],
456 |    "source": [
457 |     "structured_llm = llm.with_structured_output(\n",
458 |     "    StructuredOutput, method=\"xml_mode\"\n",
459 |     ")  # try xml_mode\n",
460 |     "response = structured_llm.invoke(\"Tell me about artificial intelligence\")\n",
461 |     "print(f\"The title:\\t{response.title}\\n\")\n",
462 |     "response"
463 |    ]
464 |   },
465 |   {
466 |    "cell_type": "code",
467 |    "execution_count": null,
468 |    "id": "dc3293b7-a50c-43c8-a022-8975e1e444b8",
469 |    "metadata": {
470 |     "height": 81
471 |    },
472 |    "outputs": [],
473 |    "source": [
474 |     "from langchain_core.pydantic_v1 import BaseModel\n",
475 |     "\n",
476 |     "\n",
477 |     "class Queries(BaseModel):\n",
478 |     "    queries: List[str]"
479 |    ]
480 |   },
481 |   {
482 |    "cell_type": "markdown",
483 |    "id": "7bfafd4a",
484 |    "metadata": {},
485 |    "source": [
486 |     "### Setting up Tavily Client\n",
487 |     "\n",
488 |     "We're using the Tavily API for research. We import the TavilyClient and initialize it with our API key. This will allow us to perform web searches to gather information for our essays.\n"
489 |    ]
490 |   },
491 |   {
492 |    "cell_type": "code",
493 |    "execution_count": null,
494 |    "id": "0722c3d4-4cbf-43bf-81b0-50f634c4ce61",
495 |    "metadata": {
496 |     "height": 64
497 |    },
498 |    "outputs": [],
499 |    "source": [
500 |     "from tavily import TavilyClient\n",
501 |     "import os\n",
502 |     "\n",
503 |     "tavily = TavilyClient(api_key=tavily_ai_api_key)"
504 |    ]
505 |   },
506 |   {
507 |    "cell_type": "markdown",
508 |    "id": "b4114e17",
509 |    "metadata": {},
510 |    "source": [
511 |     "## Defining Node Functions\n",
512 |     "\n",
513 |     "Now we're creating the individual components of our essay writing process. Each function represents a node in our graph:\n",
514 |     "\n",
515 |     "1. plan_node: Creates the essay outline.\n",
516 |     "2. research_plan_node: Generates search queries and fetches information based on the plan.\n",
517 |     "3. generation_node: Writes the essay draft.\n",
518 |     "4. reflection_node: Critiques the current draft.\n",
519 |     "5. research_critique_node: Performs additional research based on the critique.\n",
520 |     "6. should_continue: Decides whether to continue revising or stop.\n",
521 |     "\n",
522 |     "Each of these functions interacts with our Claude model and updates the agent's state accordingly.\n"
523 |    ]
524 |   },
525 |   {
526 |    "cell_type": "code",
527 |    "execution_count": null,
528 |    "id": "6b2f82fe-3ec4-4917-be51-9fb10d1317fa",
529 |    "metadata": {
530 |     "height": 132
531 |    },
532 |    "outputs": [],
533 |    "source": [
534 |     "def plan_node(state: AgentState):\n",
535 |     "    messages = [SystemMessage(content=PLAN_PROMPT), HumanMessage(content=state[\"task\"])]\n",
536 |     "    response = model.invoke(messages)\n",
537 |     "    return {\"plan\": response.content}"
538 |    ]
539 |   },
540 |   {
541 |    "cell_type": "code",
542 |    "execution_count": null,
543 |    "id": "65ee9b47",
544 |    "metadata": {},
545 |    "outputs": [],
546 |    "source": [
547 |     "from typing import List\n",
548 |     "from langchain.output_parsers import PydanticOutputParser\n",
549 |     "from langchain_core.prompts import PromptTemplate\n",
550 |     "from langchain_core.pydantic_v1 import BaseModel, Field\n",
551 |     "import json\n",
552 |     "\n",
553 |     "\n",
554 |     "class Queries(BaseModel):\n",
555 |     "    queries: List[str] = Field(description=\"List of research queries\")\n",
556 |     "\n",
557 |     "\n",
558 |     "def research_plan_node(state: AgentState):\n",
559 |     "    # Set up the Pydantic output parser\n",
560 |     "    parser = PydanticOutputParser(pydantic_object=Queries)\n",
561 |     "\n",
562 |     "    # Create a prompt template with format instructions\n",
563 |     "    prompt = PromptTemplate(\n",
564 |     "        template=\"Generate research queries based on the given task.\\n{format_instructions}\\nTask: {task}\\n\",\n",
565 |     "        input_variables=[\"task\"],\n",
566 |     "        partial_variables={\"format_instructions\": parser.get_format_instructions()},\n",
567 |     "    )\n",
568 |     "\n",
569 |     "    # Use the model with the new prompt and parser\n",
570 |     "    queries_output = model.invoke(prompt.format_prompt(task=state[\"task\"]))\n",
571 |     "\n",
572 |     "    # Extract the content from the AIMessage\n",
573 |     "    queries_text = queries_output.content\n",
574 |     "\n",
575 |     "    # Extract the JSON string from the content\n",
576 |     "    json_start = queries_text.find(\"{\")\n",
577 |     "    json_end = queries_text.rfind(\"}\") + 1\n",
578 |     "    json_str = queries_text[json_start:json_end]\n",
579 |     "\n",
580 |     "    # Parse the JSON string\n",
581 |     "    queries_dict = json.loads(json_str)\n",
582 |     "\n",
583 |     "    # Create a Queries object from the parsed JSON\n",
584 |     "    parsed_queries = Queries(**queries_dict)\n",
585 |     "\n",
586 |     "    content = state[\"content\"] or []\n",
587 |     "    for q in parsed_queries.queries:\n",
588 |     "        response = tavily.search(query=q, max_results=2)\n",
589 |     "        for r in response[\"results\"]:\n",
590 |     "            content.append(r[\"content\"])\n",
591 |     "    return {\"content\": content}"
592 |    ]
593 |   },
594 |   {
595 |    "cell_type": "code",
596 |    "execution_count": null,
597 |    "id": "98f303b1-a4d0-408c-8cc0-515ff980717f",
598 |    "metadata": {
599 |     "height": 285
600 |    },
601 |    "outputs": [],
602 |    "source": [
603 |     "def generation_node(state: AgentState):\n",
604 |     "    content = \"\\n\\n\".join(state[\"content\"] or [])\n",
605 |     "    user_message = HumanMessage(\n",
606 |     "        content=f\"{state['task']}\\n\\nHere is my plan:\\n\\n{state['plan']}\"\n",
607 |     "    )\n",
608 |     "    messages = [\n",
609 |     "        SystemMessage(content=WRITER_PROMPT.format(content=content)),\n",
610 |     "        user_message,\n",
611 |     "    ]\n",
612 |     "    response = model.invoke(messages)\n",
613 |     "    return {\n",
614 |     "        \"draft\": response.content,\n",
615 |     "        \"revision_number\": state.get(\"revision_number\", 1) + 1,\n",
616 |     "    }"
617 |    ]
618 |   },
619 |   {
620 |    "cell_type": "code",
621 |    "execution_count": null,
622 |    "id": "bf4dcb93-6298-4cfd-b3ce-61dfac7fb35f",
623 |    "metadata": {
624 |     "height": 132
625 |    },
626 |    "outputs": [],
627 |    "source": [
628 |     "def reflection_node(state: AgentState):\n",
629 |     "    messages = [\n",
630 |     "        SystemMessage(content=REFLECTION_PROMPT),\n",
631 |     "        HumanMessage(content=state[\"draft\"]),\n",
632 |     "    ]\n",
633 |     "    response = model.invoke(messages)\n",
634 |     "    return {\"critique\": response.content}"
635 |    ]
636 |   },
637 |   {
638 |    "cell_type": "code",
639 |    "execution_count": null,
640 |    "id": "27526845",
641 |    "metadata": {},
642 |    "outputs": [],
643 |    "source": [
644 |     "def research_critique_node(state: AgentState):\n",
645 |     "    # Set up the Pydantic output parser\n",
646 |     "    parser = PydanticOutputParser(pydantic_object=Queries)\n",
647 |     "\n",
648 |     "    # Create a prompt template with format instructions\n",
649 |     "    prompt = PromptTemplate(\n",
650 |     "        template=\"Generate research queries based on the given critique.\\n{format_instructions}\\nCritique: {critique}\\n\",\n",
651 |     "        input_variables=[\"critique\"],\n",
652 |     "        partial_variables={\"format_instructions\": parser.get_format_instructions()},\n",
653 |     "    )\n",
654 |     "\n",
655 |     "    # Use the model with the new prompt and parser\n",
656 |     "    queries_output = model.invoke(prompt.format_prompt(critique=state[\"critique\"]))\n",
657 |     "\n",
658 |     "    # Extract the content from the AIMessage\n",
659 |     "    queries_text = queries_output.content\n",
660 |     "\n",
661 |     "    # Extract the JSON string from the content\n",
662 |     "    json_start = queries_text.find(\"{\")\n",
663 |     "    json_end = queries_text.rfind(\"}\") + 1\n",
664 |     "    json_str = queries_text[json_start:json_end]\n",
665 |     "\n",
666 |     "    # Parse the JSON string\n",
667 |     "    queries_dict = json.loads(json_str)\n",
668 |     "\n",
669 |     "    # Create a Queries object from the parsed JSON\n",
670 |     "    parsed_queries = Queries(**queries_dict)\n",
671 |     "\n",
672 |     "    content = state[\"content\"] or []\n",
673 |     "    for q in parsed_queries.queries:\n",
674 |     "        response = tavily.search(query=q, max_results=2)\n",
675 |     "        for r in response[\"results\"]:\n",
676 |     "            content.append(r[\"content\"])\n",
677 |     "    return {\"content\": content}"
678 |    ]
679 |   },
680 |   {
681 |    "cell_type": "code",
682 |    "execution_count": null,
683 |    "id": "ff362f49-dcf1-4ea1-a86c-e516e9ab897d",
684 |    "metadata": {
685 |     "height": 81
686 |    },
687 |    "outputs": [],
688 |    "source": [
689 |     "def should_continue(state):\n",
690 |     "    if state[\"revision_number\"] > state[\"max_revisions\"]:\n",
691 |     "        return END\n",
692 |     "    return \"reflect\""
693 |    ]
694 |   },
695 |   {
696 |    "cell_type": "markdown",
697 |    "id": "7044a487",
698 |    "metadata": {},
699 |    "source": [
700 |     "## Building the Graph\n",
701 |     "\n",
702 |     "With our nodes defined, we can now build our graph. We're using LangGraph's StateGraph to create a flow for our essay writing process. We add each node to the graph, set the entry point to the planner, and define the edges between nodes.\n",
703 |     "\n",
704 |     "The key part here is the conditional edge after the generate node. It uses our should_continue function to decide whether to reflect and revise, or to end the process.\n"
705 |    ]
706 |   },
707 |   {
708 |    "cell_type": "code",
709 |    "execution_count": null,
710 |    "id": "a7e15a20-83d7-434c-8551-bce8dcc32be0",
711 |    "metadata": {
712 |     "height": 30
713 |    },
714 |    "outputs": [],
715 |    "source": [
716 |     "builder = StateGraph(AgentState)"
717 |    ]
718 |   },
719 |   {
720 |    "cell_type": "code",
721 |    "execution_count": null,
722 |    "id": "54ab2c74-f32e-490c-a85d-932d11444210",
723 |    "metadata": {
724 |     "height": 98
725 |    },
726 |    "outputs": [],
727 |    "source": [
728 |     "builder.add_node(\"planner\", plan_node)\n",
729 |     "builder.add_node(\"generate\", generation_node)\n",
730 |     "builder.add_node(\"reflect\", reflection_node)\n",
731 |     "builder.add_node(\"research_plan\", research_plan_node)\n",
732 |     "builder.add_node(\"research_critique\", research_critique_node)"
733 |    ]
734 |   },
735 |   {
736 |    "cell_type": "code",
737 |    "execution_count": null,
738 |    "id": "a833d3ce-bd31-4319-811d-decff226b970",
739 |    "metadata": {
740 |     "height": 30
741 |    },
742 |    "outputs": [],
743 |    "source": [
744 |     "builder.set_entry_point(\"planner\")"
745 |    ]
746 |   },
747 |   {
748 |    "cell_type": "code",
749 |    "execution_count": null,
750 |    "id": "76e93cce-6eab-4c7c-ac64-e9993fdb30d6",
751 |    "metadata": {
752 |     "height": 115
753 |    },
754 |    "outputs": [],
755 |    "source": [
756 |     "builder.add_conditional_edges(\n",
757 |     "    \"generate\", should_continue, {END: END, \"reflect\": \"reflect\"}\n",
758 |     ")"
759 |    ]
760 |   },
761 |   {
762 |    "cell_type": "code",
763 |    "execution_count": null,
764 |    "id": "fd2d0990-a932-423f-9ff3-5cada58c5f32",
765 |    "metadata": {
766 |     "height": 98
767 |    },
768 |    "outputs": [],
769 |    "source": [
770 |     "builder.add_edge(\"planner\", \"research_plan\")\n",
771 |     "builder.add_edge(\"research_plan\", \"generate\")\n",
772 |     "\n",
773 |     "builder.add_edge(\"reflect\", \"research_critique\")\n",
774 |     "builder.add_edge(\"research_critique\", \"generate\")"
775 |    ]
776 |   },
777 |   {
778 |    "cell_type": "code",
779 |    "execution_count": null,
780 |    "id": "27cde654-64e2-48bc-80a9-0ed668ccb7dc",
781 |    "metadata": {
782 |     "height": 30
783 |    },
784 |    "outputs": [],
785 |    "source": [
786 |     "graph = builder.compile(checkpointer=memory)"
787 |    ]
788 |   },
789 |   {
790 |    "cell_type": "code",
791 |    "execution_count": null,
792 |    "id": "4871f644-b131-4065-b7ce-b82c20a41f11",
793 |    "metadata": {
794 |     "height": 64
795 |    },
796 |    "outputs": [],
797 |    "source": [
798 |     "from IPython.display import Image\n",
799 |     "\n",
800 |     "Image(graph.get_graph().draw_png())"
801 |    ]
802 |   },
803 |   {
804 |    "cell_type": "markdown",
805 |    "id": "392cfcff",
806 |    "metadata": {},
807 |    "source": [
808 |     "## Running the Graph\n",
809 |     "\n",
810 |     "To test our essay writer, we're using the graph.stream method. This allows us to see each step of the process as it happens. We're asking it to write an essay on the difference between LangChain and LangSmith, with a maximum of two revisions.\n",
811 |     "\n",
812 |     "As it runs, you'll see outputs from each node, showing you how the essay evolves through the planning, research, writing, and revision stages.\n"
813 |    ]
814 |   },
815 |   {
816 |    "cell_type": "code",
817 |    "execution_count": null,
818 |    "id": "98f3be1d-cc4c-41fa-9863-3e386e88e305",
819 |    "metadata": {
820 |     "height": 132
821 |    },
822 |    "outputs": [],
823 |    "source": [
824 |     "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
825 |     "for s in graph.stream(\n",
826 |     "    {\n",
827 |     "        \"task\": \"what is the difference between langchain and langsmith\",\n",
828 |     "        \"max_revisions\": 2,\n",
829 |     "        \"revision_number\": 1,\n",
830 |     "    },\n",
831 |     "    thread,\n",
832 |     "):\n",
833 |     "    print(s)"
834 |    ]
835 |   },
836 |   {
837 |    "cell_type": "markdown",
838 |    "id": "4d1664b5-75e0-46b7-9c2b-4ac9171f4597",
839 |    "metadata": {},
840 |    "source": [
841 |     "## Essay Writer Interface\n",
842 |     "\n",
843 |     "Finally, lets use a simple GUI using Gradio to make it easy to interact with our essay writer. \n",
844 |     "\n",
845 |     "**IMPORTANT NOTE**: To use Gradio within the Amazon SageMaker Code Editor, the app needs to be launched in `shared=True` mode, which creates a public link. Review the [Security and File Access](https://www.gradio.app/guides/sharing-your-app#security-and-file-access) to make sure you understand the security implications.\n",
846 |     "\n",
847 |     "This GUI allows you to input an essay topic, generate the essay, and see the results of each step in the process. You can also interrupt the process after each step, view the current state of the essay, and even modify the topic or plan if you want to guide the essay in a different direction.\n",
848 |     "\n",
849 |     "This GUI makes it easy to experiment with the essay writer and see how changes to the input or the process affect the final output.\n",
850 |     "\n",
851 |     "And that concludes our AI Essay Writer project! You now have a complex, multi-step AI agent that can research, write, and refine essays on a wide range of topics. This project demonstrates how you can combine different AI and API services to create a powerful, practical application.\n"
852 |    ]
853 |   },
854 |   {
855 |    "cell_type": "code",
856 |    "execution_count": null,
857 |    "id": "0ad8a6cc-65d4-4ce7-87aa-4e67d7c23d7b",
858 |    "metadata": {
859 |     "height": 30
860 |    },
861 |    "outputs": [],
862 |    "source": [
863 |     "#set magic variables to allow for a reload when changing code without restarting the kernel\n",
864 |     "%load_ext autoreload\n",
865 |     "%autoreload 2\n",
866 |     "\n",
867 |     "import gradio as gr\n",
868 |     "from helper import ewriter, writer_gui\n",
869 |     "\n",
870 |     "MultiAgent = ewriter()\n",
871 |     "app = writer_gui(MultiAgent.graph)\n",
872 |     "app.launch()"
873 |    ]
874 |   },
875 |   {
876 |    "cell_type": "markdown",
877 |    "id": "37002eb0",
878 |    "metadata": {},
879 |    "source": [
880 |     "## Exercise - Rewrite the Essay writers prompts and parsers\n"
881 |    ]
882 |   },
883 |   {
884 |    "cell_type": "code",
885 |    "execution_count": null,
886 |    "id": "f79b56d9",
887 |    "metadata": {},
888 |    "outputs": [],
889 |    "source": []
890 |   },
891 |   {
892 |    "cell_type": "markdown",
893 |    "id": "e21704cb",
894 |    "metadata": {},
895 |    "source": []
896 |   }
897 |  ],
898 |  "metadata": {
899 |   "kernelspec": {
900 |    "display_name": "agents-dev-env",
901 |    "language": "python",
902 |    "name": "agents-dev-env"
903 |   },
904 |   "language_info": {
905 |    "codemirror_mode": {
906 |     "name": "ipython",
907 |     "version": 3
908 |    },
909 |    "file_extension": ".py",
910 |    "mimetype": "text/x-python",
911 |    "name": "python",
912 |    "nbconvert_exporter": "python",
913 |    "pygments_lexer": "ipython3",
914 |    "version": "3.10.14"
915 |   }
916 |  },
917 |  "nbformat": 4,
918 |  "nbformat_minor": 5
919 | }
920 | 


--------------------------------------------------------------------------------
/Lab_5/Lab_5.ipynb:
--------------------------------------------------------------------------------
   1 | {
   2 |  "cells": [
   3 |   {
   4 |    "cell_type": "markdown",
   5 |    "id": "8ad976db-fb87-4a92-b0d2-06defc098339",
   6 |    "metadata": {},
   7 |    "source": [
   8 |     "# Lab 5: Human in the Loop\n"
   9 |    ]
  10 |   },
  11 |   {
  12 |    "cell_type": "markdown",
  13 |    "id": "9644a30b",
  14 |    "metadata": {},
  15 |    "source": [
  16 |     "## Setup and Initialization\n",
  17 |     "\n",
  18 |     "We start by setting up our environment and importing necessary modules. A crucial component is the checkpointer:\n",
  19 |     "\n",
  20 |     "```python\n",
  21 |     "from langgraph.checkpoint.memory import MemorySaver\n",
  22 |     "memory = MemorySaver()\n",
  23 |     "```\n",
  24 |     "\n",
  25 |     "The checkpointer functions similarly to a distributed version control system like Git, but for AI states. It allows us to create \"commits\" of our AI's state at different points in its decision-making process, and \"branch\" or \"revert\" as needed.\n",
  26 |     "\n",
  27 |     "Be aware that excessive use of check-pointing can lead to significant memory usage, especially with large state objects. Implement a strategy to manage or expire old checkpoints in long-running applications. Think of usage of Time-To-Live for those checkpoints to be implemented in your database, if you no longer have a need for those checkpoints.\n",
  28 |     "We will be using SQLite as our database of choice for our checkpoints, but you might want to switch to either Redis or Postgres for your production applications.\n"
  29 |    ]
  30 |   },
  31 |   {
  32 |    "cell_type": "code",
  33 |    "execution_count": null,
  34 |    "id": "f5762271-8736-4e94-9444-8c92bd0e8074",
  35 |    "metadata": {
  36 |     "height": 64
  37 |    },
  38 |    "outputs": [],
  39 |    "source": [
  40 |     "from dotenv import load_dotenv\n",
  41 |     "import os\n",
  42 |     "import sys\n",
  43 |     "import json, re\n",
  44 |     "import pprint\n",
  45 |     "import boto3\n",
  46 |     "from botocore.client import Config\n",
  47 |     "import warnings\n",
  48 |     "\n",
  49 |     "warnings.filterwarnings(\"ignore\")\n",
  50 |     "import logging\n",
  51 |     "\n",
  52 |     "# import local modules\n",
  53 |     "dir_current = os.path.abspath(\"\")\n",
  54 |     "dir_parent = os.path.dirname(dir_current)\n",
  55 |     "if dir_parent not in sys.path:\n",
  56 |     "    sys.path.append(dir_parent)\n",
  57 |     "from utils import utils\n",
  58 |     "\n",
  59 |     "bedrock_config = Config(\n",
  60 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
  61 |     ")\n",
  62 |     "\n",
  63 |     "# Set basic configs\n",
  64 |     "logger = utils.set_logger()\n",
  65 |     "pp = utils.set_pretty_printer()\n",
  66 |     "\n",
  67 |     "# Load environment variables from .env file or Secret Manager\n",
  68 |     "_ = load_dotenv(\"../.env\")\n",
  69 |     "aws_region = os.getenv(\"AWS_REGION\")\n",
  70 |     "tavily_ai_api_key = utils.get_tavily_api(\"TAVILY_API_KEY\", aws_region)\n",
  71 |     "\n",
  72 |     "# Set bedrock configs\n",
  73 |     "bedrock_config = Config(\n",
  74 |     "    connect_timeout=120, read_timeout=120, retries={\"max_attempts\": 0}\n",
  75 |     ")\n",
  76 |     "\n",
  77 |     "# Create a bedrock runtime client\n",
  78 |     "bedrock_rt = boto3.client(\n",
  79 |     "    \"bedrock-runtime\", region_name=aws_region, config=bedrock_config\n",
  80 |     ")\n",
  81 |     "\n",
  82 |     "# Create a bedrock client to check available models\n",
  83 |     "bedrock = boto3.client(\"bedrock\", region_name=aws_region, config=bedrock_config)\n"
  84 |    ]
  85 |   },
  86 |   {
  87 |    "cell_type": "code",
  88 |    "execution_count": null,
  89 |    "id": "d0168aee-bce9-4d60-b827-f86a88187e31",
  90 |    "metadata": {
  91 |     "height": 166
  92 |    },
  93 |    "outputs": [],
  94 |    "source": [
  95 |     "from langgraph.graph import StateGraph, END\n",
  96 |     "from typing import TypedDict, Annotated\n",
  97 |     "import operator\n",
  98 |     "from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage\n",
  99 |     "from langchain_aws import ChatBedrockConverse\n",
 100 |     "from langchain_community.tools.tavily_search import TavilySearchResults\n",
 101 |     "from langgraph.checkpoint.memory import MemorySaver\n",
 102 |     "\n",
 103 |     "memory = MemorySaver()"
 104 |    ]
 105 |   },
 106 |   {
 107 |    "cell_type": "markdown",
 108 |    "id": "8d4c04fb",
 109 |    "metadata": {},
 110 |    "source": [
 111 |     "Next, we set up our agent state with a custom message handling function:\n"
 112 |    ]
 113 |   },
 114 |   {
 115 |    "cell_type": "code",
 116 |    "execution_count": null,
 117 |    "id": "2589c5b6-6cc2-4594-9a17-dccdcf676054",
 118 |    "metadata": {
 119 |     "height": 557
 120 |    },
 121 |    "outputs": [],
 122 |    "source": [
 123 |     "from uuid import uuid4\n",
 124 |     "from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage\n",
 125 |     "\n",
 126 |     "\"\"\"\n",
 127 |     "In previous examples we've annotated the `messages` state key\n",
 128 |     "with the default `operator.add` or `+` reducer, which always\n",
 129 |     "appends new messages to the end of the existing messages array.\n",
 130 |     "\n",
 131 |     "Now, to support replacing existing messages, we annotate the\n",
 132 |     "`messages` key with a customer reducer function, which replaces\n",
 133 |     "messages with the same `id`, and appends them otherwise.\n",
 134 |     "\"\"\"\n",
 135 |     "\n",
 136 |     "\n",
 137 |     "def reduce_messages(\n",
 138 |     "    left: list[AnyMessage], right: list[AnyMessage]\n",
 139 |     ") -> list[AnyMessage]:\n",
 140 |     "    # assign ids to messages that don't have them\n",
 141 |     "    for message in right:\n",
 142 |     "        if not message.id:\n",
 143 |     "            message.id = str(uuid4())\n",
 144 |     "    # merge the new messages with the existing messages\n",
 145 |     "    merged = left.copy()\n",
 146 |     "    for message in right:\n",
 147 |     "        for i, existing in enumerate(merged):\n",
 148 |     "            # replace any existing messages with the same id\n",
 149 |     "            if existing.id == message.id:\n",
 150 |     "                merged[i] = message\n",
 151 |     "                break\n",
 152 |     "        else:\n",
 153 |     "            # append any new messages to the end\n",
 154 |     "            merged.append(message)\n",
 155 |     "    return merged\n",
 156 |     "\n",
 157 |     "\n",
 158 |     "class AgentState(TypedDict):\n",
 159 |     "    messages: Annotated[list[AnyMessage], reduce_messages]"
 160 |    ]
 161 |   },
 162 |   {
 163 |    "cell_type": "markdown",
 164 |    "id": "f864ced5",
 165 |    "metadata": {},
 166 |    "source": [
 167 |     "This function acts like a smart message broker in a microservices architecture. It ensures that our conversation history remains consistent and up-to-date, crucial for maintaining context in long-running AI interactions.\n",
 168 |     "\n",
 169 |     "## Tool and Agent Setup\n",
 170 |     "\n",
 171 |     "We integrate the Tavily search tool:\n"
 172 |    ]
 173 |   },
 174 |   {
 175 |    "cell_type": "code",
 176 |    "execution_count": null,
 177 |    "id": "a2ba84ec-c172-4de7-ac55-e3158a531b23",
 178 |    "metadata": {
 179 |     "height": 30
 180 |    },
 181 |    "outputs": [],
 182 |    "source": [
 183 |     "tool = TavilySearchResults(max_results=2)"
 184 |    ]
 185 |   },
 186 |   {
 187 |    "cell_type": "markdown",
 188 |    "id": "d58320dc-c542-4a3c-bcf1-088909f9b075",
 189 |    "metadata": {},
 190 |    "source": [
 191 |     "## Human Intervention in the Agent Class\n",
 192 |     "\n",
 193 |     "The `Agent` class implements several key features that enable human intervention in the AI's decision-making process:\n",
 194 |     "\n",
 195 |     "1. **Interrupt Before Action**:\n",
 196 |     "\n",
 197 |     "   ```python\n",
 198 |     "   self.graph = graph.compile(\n",
 199 |     "       checkpointer=checkpointer, interrupt_before=[\"action\"]\n",
 200 |     "   )\n",
 201 |     "   ```\n",
 202 |     "\n",
 203 |     "   This line sets up an interrupt before the \"action\" node. It allows human oversight before any action is taken, providing an opportunity for review or modification.\n",
 204 |     "   You can add interrupts before any node you wish for as we will see in the next lab.\n",
 205 |     "\n",
 206 |     "2. **State Examination**:\n",
 207 |     "   The `exists_action` method prints the current state:\n",
 208 |     "\n",
 209 |     "   ```python\n",
 210 |     "   def exists_action(self, state: AgentState):\n",
 211 |     "       print(state)\n",
 212 |     "       # ...\n",
 213 |     "   ```\n",
 214 |     "\n",
 215 |     "   This allows humans to inspect the AI's current state, including its reasoning and intended actions.\n",
 216 |     "\n",
 217 |     "3. **Action Visibility**:\n",
 218 |     "   In the `take_action` method, each tool call is printed:\n",
 219 |     "\n",
 220 |     "   ```python\n",
 221 |     "   print(f\"Calling: {t}\")\n",
 222 |     "   ```\n",
 223 |     "\n",
 224 |     "   This provides visibility into what actions the AI is about to take, allowing for potential human intervention.\n",
 225 |     "\n",
 226 |     "4. **Checkpointing**:\n",
 227 |     "   The `checkpointer` parameter in the constructor allows for saving and loading states. This enables \"time travel\" functionality, where humans can revisit and potentially modify previous decision points.\n",
 228 |     "\n",
 229 |     "5. **Modifiable State**:\n",
 230 |     "   The `AgentState` is a mutable structure. While not shown in this class directly, it allows for state modification, enabling humans to alter the AI's context or decisions.\n",
 231 |     "\n",
 232 |     "These features collectively create a framework for human-in-the-loop AI, where human operators can monitor, intervene, and guide the AI's decision-making process at critical junctures.\n"
 233 |    ]
 234 |   },
 235 |   {
 236 |    "cell_type": "code",
 237 |    "execution_count": null,
 238 |    "id": "46a0e94e-d015-4106-b439-dbcd2fcb8bb0",
 239 |    "metadata": {
 240 |     "height": 642
 241 |    },
 242 |    "outputs": [],
 243 |    "source": [
 244 |     "class Agent:\n",
 245 |     "    def __init__(self, model, tools, system=\"\", checkpointer=None):\n",
 246 |     "        self.system = system\n",
 247 |     "        graph = StateGraph(AgentState)\n",
 248 |     "        graph.add_node(\"llm\", self.call_bedrock)\n",
 249 |     "        graph.add_node(\"action\", self.take_action)\n",
 250 |     "        graph.add_conditional_edges(\n",
 251 |     "            \"llm\", self.exists_action, {True: \"action\", False: END}\n",
 252 |     "        )\n",
 253 |     "        graph.add_edge(\"action\", \"llm\")\n",
 254 |     "        graph.set_entry_point(\"llm\")\n",
 255 |     "        self.graph = graph.compile(\n",
 256 |     "            checkpointer=checkpointer, interrupt_before=[\"action\"]\n",
 257 |     "        )\n",
 258 |     "        self.tools = {t.name: t for t in tools}\n",
 259 |     "        self.model = model.bind_tools(tools)\n",
 260 |     "\n",
 261 |     "    def call_bedrock(self, state: AgentState):\n",
 262 |     "        messages = state[\"messages\"]\n",
 263 |     "        if self.system:\n",
 264 |     "            messages = [SystemMessage(content=self.system)] + messages\n",
 265 |     "        message = self.model.invoke(messages)\n",
 266 |     "        return {\"messages\": [message]}\n",
 267 |     "\n",
 268 |     "    def exists_action(self, state: AgentState):\n",
 269 |     "        print(state)\n",
 270 |     "        result = state[\"messages\"][-1]\n",
 271 |     "        return len(result.tool_calls) > 0\n",
 272 |     "\n",
 273 |     "    def take_action(self, state: AgentState):\n",
 274 |     "        tool_calls = state[\"messages\"][-1].tool_calls\n",
 275 |     "        results = []\n",
 276 |     "        for t in tool_calls:\n",
 277 |     "            print(f\"Calling: {t}\")\n",
 278 |     "            result = self.tools[t[\"name\"]].invoke(t[\"args\"])\n",
 279 |     "            results.append(\n",
 280 |     "                ToolMessage(tool_call_id=t[\"id\"], name=t[\"name\"], content=str(result))\n",
 281 |     "            )\n",
 282 |     "        print(\"Back to the model!\")\n",
 283 |     "        return {\"messages\": results}"
 284 |    ]
 285 |   },
 286 |   {
 287 |    "cell_type": "markdown",
 288 |    "id": "8320d5d6",
 289 |    "metadata": {},
 290 |    "source": [
 291 |     "The `interrupt_before=[\"action\"]` parameter implements a critical control point in our AI pipeline.\n",
 292 |     "\n",
 293 |     "It's analogous to implementing approval gates in a CI/CD pipeline, ensuring that no critical action is taken without necessary checks.\n",
 294 |     "\n",
 295 |     "To deepen your understanding of AI safety and control mechanisms, explore the paper [\"Concrete Problems in AI Safety\"](https://arxiv.org/abs/1606.06565) from researchers at Google, Stanford, and Berkeley from 2016.\n",
 296 |     "\n",
 297 |     "## Initializing and Running the Agent\n",
 298 |     "\n",
 299 |     "We initialize our agent with Claude on Amazon Bedrock:\n"
 300 |    ]
 301 |   },
 302 |   {
 303 |    "cell_type": "code",
 304 |    "execution_count": null,
 305 |    "id": "10084a02-2928-4945-9f7c-ad3f5b33caf7",
 306 |    "metadata": {
 307 |     "height": 132
 308 |    },
 309 |    "outputs": [],
 310 |    "source": [
 311 |     "prompt = \"\"\"You are a smart research assistant. Use the search engine to look up information. \\\n",
 312 |     "You are allowed to make multiple calls (either together or in sequence). \\\n",
 313 |     "Only look up information when you are sure of what you want. \\\n",
 314 |     "If you need to look up some information before asking a follow up question, you are allowed to do that!\n",
 315 |     "\"\"\"\n",
 316 |     "model = ChatBedrockConverse(\n",
 317 |     "    client=bedrock_rt,\n",
 318 |     "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
 319 |     "    temperature=0,\n",
 320 |     "    max_tokens=None,\n",
 321 |     ")\n",
 322 |     "abot = Agent(model, [tool], system=prompt, checkpointer=memory)"
 323 |    ]
 324 |   },
 325 |   {
 326 |    "cell_type": "code",
 327 |    "execution_count": null,
 328 |    "id": "714d1205-f8fc-4912-b148-2a45da99219c",
 329 |    "metadata": {
 330 |     "height": 98
 331 |    },
 332 |    "outputs": [],
 333 |    "source": [
 334 |     "messages = [HumanMessage(content=\"Whats the weather in Berlin?\")]\n",
 335 |     "thread = {\"configurable\": {\"thread_id\": \"1\"}}\n",
 336 |     "for event in abot.graph.stream({\"messages\": messages}, thread):\n",
 337 |     "    for v in event.values():\n",
 338 |     "        print(v)"
 339 |    ]
 340 |   },
 341 |   {
 342 |    "cell_type": "code",
 343 |    "execution_count": null,
 344 |    "id": "4873049c",
 345 |    "metadata": {},
 346 |    "outputs": [],
 347 |    "source": [
 348 |     "thread"
 349 |    ]
 350 |   },
 351 |   {
 352 |    "cell_type": "markdown",
 353 |    "id": "7c64a9cf",
 354 |    "metadata": {},
 355 |    "source": [
 356 |     "## Examining and Modifying State\n",
 357 |     "\n",
 358 |     "The ability to examine and modify the agent's state is a powerful feature:\n"
 359 |    ]
 360 |   },
 361 |   {
 362 |    "cell_type": "code",
 363 |    "execution_count": null,
 364 |    "id": "83588e70-254f-4f83-a510-c8ae81e729b0",
 365 |    "metadata": {
 366 |     "height": 30
 367 |    },
 368 |    "outputs": [],
 369 |    "source": [
 370 |     "abot.graph.get_state(thread)"
 371 |    ]
 372 |   },
 373 |   {
 374 |    "cell_type": "markdown",
 375 |    "id": "a45e23bc",
 376 |    "metadata": {},
 377 |    "source": [
 378 |     "This capability is similar to having a live debugger in a production environment, allowing you to inspect and modify the state of a running process.\n"
 379 |    ]
 380 |   },
 381 |   {
 382 |    "cell_type": "code",
 383 |    "execution_count": null,
 384 |    "id": "6cb3ef4c-58b3-401b-b104-0d51e553d982",
 385 |    "metadata": {
 386 |     "height": 30
 387 |    },
 388 |    "outputs": [],
 389 |    "source": [
 390 |     "abot.graph.get_state(thread).next"
 391 |    ]
 392 |   },
 393 |   {
 394 |    "cell_type": "markdown",
 395 |    "id": "58f2735f",
 396 |    "metadata": {},
 397 |    "source": [
 398 |     "As you can see, the next node to be executed is the `('action',)` node.\n",
 399 |     "\n",
 400 |     "Remember, we compiled our graph with the interrupt set before the `action` node.\n",
 401 |     "\n",
 402 |     "```python\n",
 403 |     "        self.graph = graph.compile(\n",
 404 |     "            checkpointer=checkpointer, interrupt_before=[\"action\"]\n",
 405 |     "        )\n",
 406 |     "```\n"
 407 |    ]
 408 |   },
 409 |   {
 410 |    "cell_type": "markdown",
 411 |    "id": "f1f404d5-a3be-42c1-9990-b1e1ee011163",
 412 |    "metadata": {},
 413 |    "source": [
 414 |     "### ...continue after interrupt\n",
 415 |     "\n",
 416 |     "Now that stopped before the action (calling tavily-ai), we can continue. Lets see how that works.\n"
 417 |    ]
 418 |   },
 419 |   {
 420 |    "cell_type": "code",
 421 |    "execution_count": null,
 422 |    "id": "dc3293b7-a50c-43c8-a022-8975e1e444b8",
 423 |    "metadata": {
 424 |     "height": 64
 425 |    },
 426 |    "outputs": [],
 427 |    "source": [
 428 |     "for event in abot.graph.stream(None, thread):\n",
 429 |     "    for v in event.values():\n",
 430 |     "        print(v)"
 431 |    ]
 432 |   },
 433 |   {
 434 |    "cell_type": "code",
 435 |    "execution_count": null,
 436 |    "id": "0722c3d4-4cbf-43bf-81b0-50f634c4ce61",
 437 |    "metadata": {
 438 |     "height": 30
 439 |    },
 440 |    "outputs": [],
 441 |    "source": [
 442 |     "abot.graph.get_state(thread)"
 443 |    ]
 444 |   },
 445 |   {
 446 |    "cell_type": "code",
 447 |    "execution_count": null,
 448 |    "id": "6b2f82fe-3ec4-4917-be51-9fb10d1317fa",
 449 |    "metadata": {
 450 |     "height": 30
 451 |    },
 452 |    "outputs": [],
 453 |    "source": [
 454 |     "abot.graph.get_state(thread).next"
 455 |    ]
 456 |   },
 457 |   {
 458 |    "cell_type": "markdown",
 459 |    "id": "25b2da1c",
 460 |    "metadata": {},
 461 |    "source": [
 462 |     "### Adding your human input\n",
 463 |     "\n",
 464 |     "Next, we will see how to add your approval for executing the `action` to search the web.\n",
 465 |     "\n",
 466 |     "Depending if you are running in JupyterLab or in e.g. VS Code, the input box should pop up below the code cell, or on the top.\n",
 467 |     "\n",
 468 |     "![input box example, asking for user input](../assets/lab5_1.png)\n",
 469 |     "\n",
 470 |     "\n",
 471 |     "If you are fine with what will be searched, then please add a `y` for yes. Anything else, will abort the operation.\n"
 472 |    ]
 473 |   },
 474 |   {
 475 |    "cell_type": "code",
 476 |    "execution_count": null,
 477 |    "id": "ee0fe1c7-77e2-499c-a2f9-1f739bb6ddf0",
 478 |    "metadata": {
 479 |     "height": 251
 480 |    },
 481 |    "outputs": [],
 482 |    "source": [
 483 |     "messages = [HumanMessage(\"Whats the weather in LA?\")]\n",
 484 |     "thread = {\"configurable\": {\"thread_id\": \"2\"}}\n",
 485 |     "for event in abot.graph.stream({\"messages\": messages}, thread):\n",
 486 |     "    for v in event.values():\n",
 487 |     "        print(v)\n",
 488 |     "\n",
 489 |     "while abot.graph.get_state(thread).next:\n",
 490 |     "    print(\"\\n\", abot.graph.get_state(thread), \"\\n\")\n",
 491 |     "    _input = input(\"proceed?\")\n",
 492 |     "    if _input != \"y\":\n",
 493 |     "        print(\"aborting\")\n",
 494 |     "        break\n",
 495 |     "    for event in abot.graph.stream(None, thread):\n",
 496 |     "        for v in event.values():\n",
 497 |     "            print(v)"
 498 |    ]
 499 |   },
 500 |   {
 501 |    "cell_type": "markdown",
 502 |    "id": "1789a791",
 503 |    "metadata": {},
 504 |    "source": [
 505 |     "As you can see, the next node to be executed is the `('action',)` node.\n",
 506 |     "\n",
 507 |     "Remember, we compiled our graph with the interrupt set before the `action` node.\n",
 508 |     "\n",
 509 |     "```python\n",
 510 |     "        self.graph = graph.compile(\n",
 511 |     "            checkpointer=checkpointer, interrupt_before=[\"action\"]\n",
 512 |     "        )\n",
 513 |     "```\n",
 514 |     "\n",
 515 |     "---\n"
 516 |    ]
 517 |   },
 518 |   {
 519 |    "cell_type": "markdown",
 520 |    "id": "587c5a21",
 521 |    "metadata": {},
 522 |    "source": [
 523 |     "Now one thing that can be very interesting, is to only stop and ask for human input, whenever you are calling a certain set of tools, and not all of them.\n",
 524 |     "\n",
 525 |     "Think about how that could be achieved, before continue reading.\n"
 526 |    ]
 527 |   },
 528 |   {
 529 |    "cell_type": "markdown",
 530 |    "id": "11891ceb",
 531 |    "metadata": {},
 532 |    "source": [
 533 |     "### Stop at only specific tool calls\n",
 534 |     "\n",
 535 |     "1. Parsing the tool call and stopping execution only if the tool call parameter `name` (tool name) is the same as as the stopping parameter.\n",
 536 |     "2. Adding all tools where you need to stop to an extra node.\n",
 537 |     "\n",
 538 |     "In general, option 2 is a cleaner and easier to debug. If you would like to see how this is implemented in a bigger example, please head over to the [customer support agent](https://langchain-ai.github.io/langgraph/tutorials/customer-support/customer-support/#part-3-conditional-interrupt) from the LangGraph examples to see how sensitive tools are handled.\n",
 539 |     "\n",
 540 |     "Here is a sneak peak of the graph:\n",
 541 |     "\n",
 542 |     "![customer support agent graph with sensitive and safe tools](../assets/lab5_2.png)\n"
 543 |    ]
 544 |   },
 545 |   {
 546 |    "cell_type": "markdown",
 547 |    "id": "7bbe5689-54ab-49ca-9055-6e5216abd523",
 548 |    "metadata": {},
 549 |    "source": [
 550 |     "## Modify State\n",
 551 |     "\n",
 552 |     "Run until the interrupt and then modify the state.\n"
 553 |    ]
 554 |   },
 555 |   {
 556 |    "cell_type": "code",
 557 |    "execution_count": null,
 558 |    "id": "98f303b1-a4d0-408c-8cc0-515ff980717f",
 559 |    "metadata": {
 560 |     "height": 98
 561 |    },
 562 |    "outputs": [],
 563 |    "source": [
 564 |     "messages = [HumanMessage(\"Whats the weather in LA?\")]\n",
 565 |     "thread = {\"configurable\": {\"thread_id\": \"3\"}}\n",
 566 |     "for event in abot.graph.stream({\"messages\": messages}, thread):\n",
 567 |     "    for v in event.values():\n",
 568 |     "        print(v)"
 569 |    ]
 570 |   },
 571 |   {
 572 |    "cell_type": "code",
 573 |    "execution_count": null,
 574 |    "id": "bf4dcb93-6298-4cfd-b3ce-61dfac7fb35f",
 575 |    "metadata": {
 576 |     "height": 30
 577 |    },
 578 |    "outputs": [],
 579 |    "source": [
 580 |     "abot.graph.get_state(thread)"
 581 |    ]
 582 |   },
 583 |   {
 584 |    "cell_type": "code",
 585 |    "execution_count": null,
 586 |    "id": "932883a4-c722-42bb-aec0-b4f41c5c81a4",
 587 |    "metadata": {
 588 |     "height": 30
 589 |    },
 590 |    "outputs": [],
 591 |    "source": [
 592 |     "current_values = abot.graph.get_state(thread)"
 593 |    ]
 594 |   },
 595 |   {
 596 |    "cell_type": "code",
 597 |    "execution_count": null,
 598 |    "id": "ff362f49-dcf1-4ea1-a86c-e516e9ab897d",
 599 |    "metadata": {
 600 |     "height": 30
 601 |    },
 602 |    "outputs": [],
 603 |    "source": [
 604 |     "current_values.values[\"messages\"][-1]"
 605 |    ]
 606 |   },
 607 |   {
 608 |    "cell_type": "code",
 609 |    "execution_count": null,
 610 |    "id": "a7e15a20-83d7-434c-8551-bce8dcc32be0",
 611 |    "metadata": {
 612 |     "height": 30
 613 |    },
 614 |    "outputs": [],
 615 |    "source": [
 616 |     "current_values.values[\"messages\"][-1].tool_calls"
 617 |    ]
 618 |   },
 619 |   {
 620 |    "cell_type": "code",
 621 |    "execution_count": null,
 622 |    "id": "54ab2c74-f32e-490c-a85d-932d11444210",
 623 |    "metadata": {
 624 |     "height": 115
 625 |    },
 626 |    "outputs": [],
 627 |    "source": [
 628 |     "_id = current_values.values[\"messages\"][-1].tool_calls[0][\"id\"]\n",
 629 |     "current_values.values[\"messages\"][-1].tool_calls = [\n",
 630 |     "    {\n",
 631 |     "        \"name\": \"tavily_search_results_json\",\n",
 632 |     "        \"args\": {\"query\": \"current weather in Munich\"},\n",
 633 |     "        \"id\": _id,\n",
 634 |     "    }\n",
 635 |     "]"
 636 |    ]
 637 |   },
 638 |   {
 639 |    "cell_type": "code",
 640 |    "execution_count": null,
 641 |    "id": "a833d3ce-bd31-4319-811d-decff226b970",
 642 |    "metadata": {
 643 |     "height": 30
 644 |    },
 645 |    "outputs": [],
 646 |    "source": [
 647 |     "abot.graph.update_state(thread, current_values.values)"
 648 |    ]
 649 |   },
 650 |   {
 651 |    "cell_type": "code",
 652 |    "execution_count": null,
 653 |    "id": "76e93cce-6eab-4c7c-ac64-e9993fdb30d6",
 654 |    "metadata": {
 655 |     "height": 30
 656 |    },
 657 |    "outputs": [],
 658 |    "source": [
 659 |     "abot.graph.get_state(thread)"
 660 |    ]
 661 |   },
 662 |   {
 663 |    "cell_type": "code",
 664 |    "execution_count": null,
 665 |    "id": "fd2d0990-a932-423f-9ff3-5cada58c5f32",
 666 |    "metadata": {
 667 |     "height": 64
 668 |    },
 669 |    "outputs": [],
 670 |    "source": [
 671 |     "for event in abot.graph.stream(None, thread):\n",
 672 |     "    for v in event.values():\n",
 673 |     "        print(v)"
 674 |    ]
 675 |   },
 676 |   {
 677 |    "cell_type": "markdown",
 678 |    "id": "e75b870b-e0df-46f1-b29b-fb151ebcbcc3",
 679 |    "metadata": {},
 680 |    "source": [
 681 |     "## Time Travel\n"
 682 |    ]
 683 |   },
 684 |   {
 685 |    "cell_type": "code",
 686 |    "execution_count": null,
 687 |    "id": "27cde654-64e2-48bc-80a9-0ed668ccb7dc",
 688 |    "metadata": {
 689 |     "height": 98
 690 |    },
 691 |    "outputs": [],
 692 |    "source": [
 693 |     "states = []\n",
 694 |     "for state in abot.graph.get_state_history(thread):\n",
 695 |     "    print(state)\n",
 696 |     "    print(\"--\")\n",
 697 |     "    states.append(state)"
 698 |    ]
 699 |   },
 700 |   {
 701 |    "cell_type": "markdown",
 702 |    "id": "449896c8-6ec6-4166-b640-9cd1530336a0",
 703 |    "metadata": {},
 704 |    "source": [
 705 |     "To fetch the same state as was filmed, the offset below is changed to `-3` from `-1`. This accounts for the initial state `__start__` and the first state that are now stored to state memory with the latest version of software.\n"
 706 |    ]
 707 |   },
 708 |   {
 709 |    "cell_type": "code",
 710 |    "execution_count": null,
 711 |    "id": "4871f644-b131-4065-b7ce-b82c20a41f11",
 712 |    "metadata": {
 713 |     "height": 30
 714 |    },
 715 |    "outputs": [],
 716 |    "source": [
 717 |     "to_replay = states[-3]"
 718 |    ]
 719 |   },
 720 |   {
 721 |    "cell_type": "code",
 722 |    "execution_count": null,
 723 |    "id": "8c3d8070-3f36-4cf0-a677-508e54359c8f",
 724 |    "metadata": {
 725 |     "height": 30
 726 |    },
 727 |    "outputs": [],
 728 |    "source": [
 729 |     "to_replay"
 730 |    ]
 731 |   },
 732 |   {
 733 |    "cell_type": "code",
 734 |    "execution_count": null,
 735 |    "id": "98f3be1d-cc4c-41fa-9863-3e386e88e305",
 736 |    "metadata": {
 737 |     "height": 64
 738 |    },
 739 |    "outputs": [],
 740 |    "source": [
 741 |     "for event in abot.graph.stream(None, to_replay.config):\n",
 742 |     "    for k, v in event.items():\n",
 743 |     "        print(v)"
 744 |    ]
 745 |   },
 746 |   {
 747 |    "cell_type": "markdown",
 748 |    "id": "005353dc-630b-4b40-ab19-56c57ac06611",
 749 |    "metadata": {},
 750 |    "source": [
 751 |     "## Go back in time and edit\n"
 752 |    ]
 753 |   },
 754 |   {
 755 |    "cell_type": "code",
 756 |    "execution_count": null,
 757 |    "id": "0ad8a6cc-65d4-4ce7-87aa-4e67d7c23d7b",
 758 |    "metadata": {
 759 |     "height": 30
 760 |    },
 761 |    "outputs": [],
 762 |    "source": [
 763 |     "to_replay"
 764 |    ]
 765 |   },
 766 |   {
 767 |    "cell_type": "code",
 768 |    "execution_count": null,
 769 |    "id": "592b5e62-a203-433c-92a0-3783f490cde1",
 770 |    "metadata": {
 771 |     "height": 81
 772 |    },
 773 |    "outputs": [],
 774 |    "source": [
 775 |     "_id = to_replay.values[\"messages\"][-1].tool_calls[0][\"id\"]\n",
 776 |     "to_replay.values[\"messages\"][-1].tool_calls = [\n",
 777 |     "    {\n",
 778 |     "        \"name\": \"tavily_search_results_json\",\n",
 779 |     "        \"args\": {\"query\": \"current weather in LA, accuweather\"},\n",
 780 |     "        \"id\": _id,\n",
 781 |     "    }\n",
 782 |     "]"
 783 |    ]
 784 |   },
 785 |   {
 786 |    "cell_type": "code",
 787 |    "execution_count": null,
 788 |    "id": "14fa923c-7e4f-42d1-965f-0f8ccd50fbd7",
 789 |    "metadata": {
 790 |     "height": 30
 791 |    },
 792 |    "outputs": [],
 793 |    "source": [
 794 |     "branch_state = abot.graph.update_state(to_replay.config, to_replay.values)"
 795 |    ]
 796 |   },
 797 |   {
 798 |    "cell_type": "code",
 799 |    "execution_count": null,
 800 |    "id": "570c6245-2837-4ac5-983b-95f61f3ac10d",
 801 |    "metadata": {
 802 |     "height": 81
 803 |    },
 804 |    "outputs": [],
 805 |    "source": [
 806 |     "for event in abot.graph.stream(None, branch_state):\n",
 807 |     "    for k, v in event.items():\n",
 808 |     "        if k != \"__end__\":\n",
 809 |     "            print(v)"
 810 |    ]
 811 |   },
 812 |   {
 813 |    "cell_type": "markdown",
 814 |    "id": "f1016ce7-368d-4922-9044-fcdfd47c273b",
 815 |    "metadata": {},
 816 |    "source": [
 817 |     "## Add message to a state at a given time\n"
 818 |    ]
 819 |   },
 820 |   {
 821 |    "cell_type": "code",
 822 |    "execution_count": null,
 823 |    "id": "6b910915-b087-4d35-afff-0ec30a5852f1",
 824 |    "metadata": {
 825 |     "height": 30
 826 |    },
 827 |    "outputs": [],
 828 |    "source": [
 829 |     "to_replay"
 830 |    ]
 831 |   },
 832 |   {
 833 |    "cell_type": "code",
 834 |    "execution_count": null,
 835 |    "id": "c4feb6cc-5129-4a99-bb45-851bc07b5709",
 836 |    "metadata": {
 837 |     "height": 30
 838 |    },
 839 |    "outputs": [],
 840 |    "source": [
 841 |     "_id = to_replay.values[\"messages\"][-1].tool_calls[0][\"id\"]"
 842 |    ]
 843 |   },
 844 |   {
 845 |    "cell_type": "code",
 846 |    "execution_count": null,
 847 |    "id": "e85a02b4-96cc-4b01-8792-397a774eb499",
 848 |    "metadata": {
 849 |     "height": 98
 850 |    },
 851 |    "outputs": [],
 852 |    "source": [
 853 |     "# Lets update the humidity to something that is impossible\n",
 854 |     "\n",
 855 |     "state_update = {\n",
 856 |     "    \"messages\": [\n",
 857 |     "        ToolMessage(\n",
 858 |     "            tool_call_id=_id,\n",
 859 |     "            name=\"tavily_search_results_json\",\n",
 860 |     "            content=\"23 degree celcius, 110 percent humidity\",\n",
 861 |     "        )\n",
 862 |     "    ]\n",
 863 |     "}"
 864 |    ]
 865 |   },
 866 |   {
 867 |    "cell_type": "code",
 868 |    "execution_count": null,
 869 |    "id": "ae8b86a6-5e20-4252-b1d8-009b8318345a",
 870 |    "metadata": {
 871 |     "height": 81
 872 |    },
 873 |    "outputs": [],
 874 |    "source": [
 875 |     "branch_and_add = abot.graph.update_state(\n",
 876 |     "    to_replay.config, state_update, as_node=\"action\"\n",
 877 |     ")"
 878 |    ]
 879 |   },
 880 |   {
 881 |    "cell_type": "code",
 882 |    "execution_count": null,
 883 |    "id": "af925917-b746-48c9-ac74-62fefbe5246c",
 884 |    "metadata": {
 885 |     "height": 64
 886 |    },
 887 |    "outputs": [],
 888 |    "source": [
 889 |     "for event in abot.graph.stream(None, branch_and_add):\n",
 890 |     "    for k, v in event.items():\n",
 891 |     "        print(v)"
 892 |    ]
 893 |   },
 894 |   {
 895 |    "cell_type": "markdown",
 896 |    "id": "a6b1de51-f1dd-4719-89e6-01c21d2b304e",
 897 |    "metadata": {},
 898 |    "source": [
 899 |     "# Extra Practice\n"
 900 |    ]
 901 |   },
 902 |   {
 903 |    "cell_type": "markdown",
 904 |    "id": "31e06033-59fd-4d6b-8891-fb8bc2d7a037",
 905 |    "metadata": {},
 906 |    "source": [
 907 |     "## Build a small graph\n",
 908 |     "\n",
 909 |     "This is a small simple graph you can tinker with if you want more insight into controlling state memory.\n"
 910 |    ]
 911 |   },
 912 |   {
 913 |    "cell_type": "code",
 914 |    "execution_count": null,
 915 |    "id": "7614683e-1c13-4518-b464-263fb91be761",
 916 |    "metadata": {
 917 |     "height": 64
 918 |    },
 919 |    "outputs": [],
 920 |    "source": [
 921 |     "from dotenv import load_dotenv\n",
 922 |     "\n",
 923 |     "_ = load_dotenv()"
 924 |    ]
 925 |   },
 926 |   {
 927 |    "cell_type": "code",
 928 |    "execution_count": null,
 929 |    "id": "efbef86a-75fe-416f-99ab-70b181d934dc",
 930 |    "metadata": {
 931 |     "height": 81
 932 |    },
 933 |    "outputs": [],
 934 |    "source": [
 935 |     "from langgraph.graph import StateGraph, END\n",
 936 |     "from typing import TypedDict, Annotated\n",
 937 |     "import operator"
 938 |    ]
 939 |   },
 940 |   {
 941 |    "cell_type": "markdown",
 942 |    "id": "67c972f7-5a14-49da-9d82-37105a6637ed",
 943 |    "metadata": {},
 944 |    "source": [
 945 |     "Define a simple 2 node graph with the following state: -`lnode`: last node -`scratch`: a scratchpad location -`count` : a counter that is incremented each step\n"
 946 |    ]
 947 |   },
 948 |   {
 949 |    "cell_type": "code",
 950 |    "execution_count": null,
 951 |    "id": "b846f637-5e98-4a7d-9ca0-5144302b7cef",
 952 |    "metadata": {
 953 |     "height": 81
 954 |    },
 955 |    "outputs": [],
 956 |    "source": [
 957 |     "class AgentState(TypedDict):\n",
 958 |     "    lnode: str\n",
 959 |     "    scratch: str\n",
 960 |     "    count: Annotated[int, operator.add]"
 961 |    ]
 962 |   },
 963 |   {
 964 |    "cell_type": "code",
 965 |    "execution_count": null,
 966 |    "id": "b8bb1c6d-41b7-4f5b-807b-ecbdec82d8d6",
 967 |    "metadata": {
 968 |     "height": 183
 969 |    },
 970 |    "outputs": [],
 971 |    "source": [
 972 |     "def node1(state: AgentState):\n",
 973 |     "    print(f\"node1, count:{state['count']}\")\n",
 974 |     "    return {\n",
 975 |     "        \"lnode\": \"node_1\",\n",
 976 |     "        \"count\": 1,\n",
 977 |     "    }\n",
 978 |     "\n",
 979 |     "\n",
 980 |     "def node2(state: AgentState):\n",
 981 |     "    print(f\"node2, count:{state['count']}\")\n",
 982 |     "    return {\n",
 983 |     "        \"lnode\": \"node_2\",\n",
 984 |     "        \"count\": 1,\n",
 985 |     "    }"
 986 |    ]
 987 |   },
 988 |   {
 989 |    "cell_type": "markdown",
 990 |    "id": "6adc403b-10d2-4a6c-bd21-c84bb8e0c8fc",
 991 |    "metadata": {},
 992 |    "source": [
 993 |     "The graph goes N1->N2->N1... but breaks after count reaches 3.\n"
 994 |    ]
 995 |   },
 996 |   {
 997 |    "cell_type": "code",
 998 |    "execution_count": null,
 999 |    "id": "b2c6d249-4793-41c5-8e80-40f1eedc4baf",
1000 |    "metadata": {
1001 |     "height": 47
1002 |    },
1003 |    "outputs": [],
1004 |    "source": [
1005 |     "def should_continue(state):\n",
1006 |     "    return state[\"count\"] < 3"
1007 |    ]
1008 |   },
1009 |   {
1010 |    "cell_type": "code",
1011 |    "execution_count": null,
1012 |    "id": "4e3721d9-4508-48fa-9be0-93af20144072",
1013 |    "metadata": {
1014 |     "height": 166
1015 |    },
1016 |    "outputs": [],
1017 |    "source": [
1018 |     "builder = StateGraph(AgentState)\n",
1019 |     "builder.add_node(\"Node1\", node1)\n",
1020 |     "builder.add_node(\"Node2\", node2)\n",
1021 |     "\n",
1022 |     "builder.add_edge(\"Node1\", \"Node2\")\n",
1023 |     "builder.add_conditional_edges(\"Node2\", should_continue, {True: \"Node1\", False: END})\n",
1024 |     "builder.set_entry_point(\"Node1\")"
1025 |    ]
1026 |   },
1027 |   {
1028 |    "cell_type": "code",
1029 |    "execution_count": null,
1030 |    "id": "8d35d70c-daeb-49c1-9b7c-90cc2a7ca142",
1031 |    "metadata": {
1032 |     "height": 47
1033 |    },
1034 |    "outputs": [],
1035 |    "source": [
1036 |     "memory = MemorySaver()\n",
1037 |     "graph = builder.compile(checkpointer=memory)"
1038 |    ]
1039 |   },
1040 |   {
1041 |    "cell_type": "markdown",
1042 |    "id": "c57051f5-f7fb-4be5-a2f5-21cb1aa4fecb",
1043 |    "metadata": {},
1044 |    "source": [
1045 |     "### Run it!\n",
1046 |     "\n",
1047 |     "Now, set the thread and run!\n"
1048 |    ]
1049 |   },
1050 |   {
1051 |    "cell_type": "code",
1052 |    "execution_count": null,
1053 |    "id": "edba3a1a-84a7-45eb-a6ee-74a612b68d54",
1054 |    "metadata": {
1055 |     "height": 47
1056 |    },
1057 |    "outputs": [],
1058 |    "source": [
1059 |     "thread = {\"configurable\": {\"thread_id\": str(1)}}\n",
1060 |     "graph.invoke({\"count\": 0, \"scratch\": \"hi\"}, thread)"
1061 |    ]
1062 |   },
1063 |   {
1064 |    "cell_type": "markdown",
1065 |    "id": "2ce7b035-bfdb-4642-8750-47d188277423",
1066 |    "metadata": {},
1067 |    "source": [
1068 |     "### Look at current state\n"
1069 |    ]
1070 |   },
1071 |   {
1072 |    "cell_type": "markdown",
1073 |    "id": "43324904-a6f4-4299-86de-e0218e4cb225",
1074 |    "metadata": {},
1075 |    "source": [
1076 |     "Get the current state. Note the `values` which are the AgentState. Note the `config` and the `thread_ts`. You will be using those to refer to snapshots below.\n"
1077 |    ]
1078 |   },
1079 |   {
1080 |    "cell_type": "code",
1081 |    "execution_count": null,
1082 |    "id": "f2513b41-31b2-46e6-84f6-0c519f697973",
1083 |    "metadata": {
1084 |     "height": 30
1085 |    },
1086 |    "outputs": [],
1087 |    "source": [
1088 |     "graph.get_state(thread)"
1089 |    ]
1090 |   },
1091 |   {
1092 |    "cell_type": "markdown",
1093 |    "id": "7f7be4ae-24f0-4362-bd99-55ad38fe0112",
1094 |    "metadata": {},
1095 |    "source": [
1096 |     "View all the statesnapshots in memory. You can use the displayed `count` agentstate variable to help track what you see. Notice the most recent snapshots are returned by the iterator first. Also note that there is a handy `step` variable in the metadata that counts the number of steps in the graph execution. This is a bit detailed - but you can also notice that the _parent_config_ is the _config_ of the previous node. At initial startup, additional states are inserted into memory to create a parent. This is something to check when you branch or _time travel_ below.\n"
1097 |    ]
1098 |   },
1099 |   {
1100 |    "cell_type": "markdown",
1101 |    "id": "18647be0-c7c6-4ec5-9454-54a034fcd053",
1102 |    "metadata": {},
1103 |    "source": [
1104 |     "### Look at state history\n"
1105 |    ]
1106 |   },
1107 |   {
1108 |    "cell_type": "code",
1109 |    "execution_count": null,
1110 |    "id": "686fa197-c97f-4ae6-82d2-938bbe5542c1",
1111 |    "metadata": {
1112 |     "height": 47
1113 |    },
1114 |    "outputs": [],
1115 |    "source": [
1116 |     "for state in graph.get_state_history(thread):\n",
1117 |     "    print(state, \"\\n\")"
1118 |    ]
1119 |   },
1120 |   {
1121 |    "cell_type": "markdown",
1122 |    "id": "dd7850ff-0748-43d7-8956-074fa9fd819f",
1123 |    "metadata": {},
1124 |    "source": [
1125 |     "Store just the `config` into an list. Note the sequence of counts on the right. `get_state_history` returns the most recent snapshots first.\n"
1126 |    ]
1127 |   },
1128 |   {
1129 |    "cell_type": "code",
1130 |    "execution_count": null,
1131 |    "id": "9f68e604-4f53-46c0-8f06-e7726ec9dcf6",
1132 |    "metadata": {
1133 |     "height": 81
1134 |    },
1135 |    "outputs": [],
1136 |    "source": [
1137 |     "states = []\n",
1138 |     "for state in graph.get_state_history(thread):\n",
1139 |     "    states.append(state.config)\n",
1140 |     "    print(state.config, state.values[\"count\"])"
1141 |    ]
1142 |   },
1143 |   {
1144 |    "cell_type": "markdown",
1145 |    "id": "90790095-4080-4e76-b538-47caac7d9699",
1146 |    "metadata": {},
1147 |    "source": [
1148 |     "Grab an early state.\n"
1149 |    ]
1150 |   },
1151 |   {
1152 |    "cell_type": "code",
1153 |    "execution_count": null,
1154 |    "id": "e1db574e-f158-44cf-b921-f1f4466c314d",
1155 |    "metadata": {
1156 |     "height": 30
1157 |    },
1158 |    "outputs": [],
1159 |    "source": [
1160 |     "states[-3]"
1161 |    ]
1162 |   },
1163 |   {
1164 |    "cell_type": "markdown",
1165 |    "id": "2cc233eb-f388-4ecd-bfb5-dac3568d37ce",
1166 |    "metadata": {},
1167 |    "source": [
1168 |     "This is the state after Node1 completed for the first time. Note `next` is `Node2`and `count` is 1.\n"
1169 |    ]
1170 |   },
1171 |   {
1172 |    "cell_type": "code",
1173 |    "execution_count": null,
1174 |    "id": "240f5039-7916-4c2a-88d5-d363b2898e70",
1175 |    "metadata": {
1176 |     "height": 30
1177 |    },
1178 |    "outputs": [],
1179 |    "source": [
1180 |     "graph.get_state(states[-3])"
1181 |    ]
1182 |   },
1183 |   {
1184 |    "cell_type": "markdown",
1185 |    "id": "4872468e-4d23-4840-ae14-c06c1ab4f161",
1186 |    "metadata": {},
1187 |    "source": [
1188 |     "### Go Back in Time\n",
1189 |     "\n",
1190 |     "Use that state in `invoke` to go back in time. Notice it uses states[-3] as _current_state_ and continues to node2,\n"
1191 |    ]
1192 |   },
1193 |   {
1194 |    "cell_type": "code",
1195 |    "execution_count": null,
1196 |    "id": "e3049179-b901-4557-a9c4-78afb3d53d27",
1197 |    "metadata": {
1198 |     "height": 30
1199 |    },
1200 |    "outputs": [],
1201 |    "source": [
1202 |     "graph.invoke(None, states[-3])"
1203 |    ]
1204 |   },
1205 |   {
1206 |    "cell_type": "markdown",
1207 |    "id": "5bf7789a-df3d-4c7e-8899-d96a99d45717",
1208 |    "metadata": {},
1209 |    "source": [
1210 |     "Notice the new states are now in state history. Notice the counts on the far right.\n"
1211 |    ]
1212 |   },
1213 |   {
1214 |    "cell_type": "code",
1215 |    "execution_count": null,
1216 |    "id": "37fa44ec-9bd1-484b-a415-cc5f50b6e799",
1217 |    "metadata": {
1218 |     "height": 64
1219 |    },
1220 |    "outputs": [],
1221 |    "source": [
1222 |     "thread = {\"configurable\": {\"thread_id\": str(1)}}\n",
1223 |     "for state in graph.get_state_history(thread):\n",
1224 |     "    print(state.config, state.values[\"count\"])"
1225 |    ]
1226 |   },
1227 |   {
1228 |    "cell_type": "markdown",
1229 |    "id": "d3c8d305-6752-4cf6-a8cb-7babf3bbd643",
1230 |    "metadata": {},
1231 |    "source": [
1232 |     "You can see the details below. Lots of text, but try to find the node that start the new branch. Notice the parent _config_ is not the previous entry in the stack, but is the entry from state[-3].\n"
1233 |    ]
1234 |   },
1235 |   {
1236 |    "cell_type": "code",
1237 |    "execution_count": null,
1238 |    "id": "07d25697-fcf0-4a26-9485-3e195b0af225",
1239 |    "metadata": {
1240 |     "height": 64
1241 |    },
1242 |    "outputs": [],
1243 |    "source": [
1244 |     "thread = {\"configurable\": {\"thread_id\": str(1)}}\n",
1245 |     "for state in graph.get_state_history(thread):\n",
1246 |     "    print(state, \"\\n\")"
1247 |    ]
1248 |   },
1249 |   {
1250 |    "cell_type": "markdown",
1251 |    "id": "d4653f91-761c-4185-8ca2-83de758308c7",
1252 |    "metadata": {},
1253 |    "source": [
1254 |     "### Modify State\n",
1255 |     "\n",
1256 |     "Let's start by starting a fresh thread and running to clean out history.\n"
1257 |    ]
1258 |   },
1259 |   {
1260 |    "cell_type": "code",
1261 |    "execution_count": null,
1262 |    "id": "0eaf346c-ce22-4b6a-b18f-d200dfd991de",
1263 |    "metadata": {
1264 |     "height": 47
1265 |    },
1266 |    "outputs": [],
1267 |    "source": [
1268 |     "thread2 = {\"configurable\": {\"thread_id\": str(2)}}\n",
1269 |     "graph.invoke({\"count\": 0, \"scratch\": \"hi\"}, thread2)"
1270 |    ]
1271 |   },
1272 |   {
1273 |    "cell_type": "code",
1274 |    "execution_count": null,
1275 |    "id": "90fc1797-25e6-4931-b173-2e520b71c372",
1276 |    "metadata": {
1277 |     "height": 64
1278 |    },
1279 |    "outputs": [],
1280 |    "source": [
1281 |     "from IPython.display import Image\n",
1282 |     "\n",
1283 |     "Image(graph.get_graph().draw_png())"
1284 |    ]
1285 |   },
1286 |   {
1287 |    "cell_type": "code",
1288 |    "execution_count": null,
1289 |    "id": "4ede9215-3133-4ad9-8dd5-6a8288ebe055",
1290 |    "metadata": {
1291 |     "height": 81
1292 |    },
1293 |    "outputs": [],
1294 |    "source": [
1295 |     "states2 = []\n",
1296 |     "for state in graph.get_state_history(thread2):\n",
1297 |     "    states2.append(state.config)\n",
1298 |     "    print(state.config, state.values[\"count\"])"
1299 |    ]
1300 |   },
1301 |   {
1302 |    "cell_type": "markdown",
1303 |    "id": "23ab2c5f-cb1d-4851-8925-ce0af1f12a40",
1304 |    "metadata": {},
1305 |    "source": [
1306 |     "Start by grabbing a state.\n"
1307 |    ]
1308 |   },
1309 |   {
1310 |    "cell_type": "code",
1311 |    "execution_count": null,
1312 |    "id": "65313f6a-c7a9-49e3-aab2-0eb905e74a91",
1313 |    "metadata": {
1314 |     "height": 47
1315 |    },
1316 |    "outputs": [],
1317 |    "source": [
1318 |     "save_state = graph.get_state(states2[-3])\n",
1319 |     "save_state"
1320 |    ]
1321 |   },
1322 |   {
1323 |    "cell_type": "markdown",
1324 |    "id": "c81b354b-8c2d-4150-9150-ac9d8486af02",
1325 |    "metadata": {},
1326 |    "source": [
1327 |     "Now modify the values. One subtle item to note: Recall when agent state was defined, `count` used `operator.add` to indicate that values are _added_ to the current value. Here, `-3` will be added to the current count value rather than replace it.\n"
1328 |    ]
1329 |   },
1330 |   {
1331 |    "cell_type": "code",
1332 |    "execution_count": null,
1333 |    "id": "dae6527b-912f-4f86-a07b-1fae684aaa77",
1334 |    "metadata": {
1335 |     "height": 64
1336 |    },
1337 |    "outputs": [],
1338 |    "source": [
1339 |     "save_state.values[\"count\"] = -3\n",
1340 |     "save_state.values[\"scratch\"] = \"hello\"\n",
1341 |     "save_state"
1342 |    ]
1343 |   },
1344 |   {
1345 |    "cell_type": "markdown",
1346 |    "id": "904bb1eb-2e39-472f-ae56-d8e95e2f1ab6",
1347 |    "metadata": {},
1348 |    "source": [
1349 |     "Now update the state. This creates a new entry at the _top_, or _latest_ entry in memory. This will become the current state.\n"
1350 |    ]
1351 |   },
1352 |   {
1353 |    "cell_type": "code",
1354 |    "execution_count": null,
1355 |    "id": "bcd78670-b119-45fa-934b-2ce2d477b4c2",
1356 |    "metadata": {
1357 |     "height": 30
1358 |    },
1359 |    "outputs": [],
1360 |    "source": [
1361 |     "graph.update_state(thread2, save_state.values)"
1362 |    ]
1363 |   },
1364 |   {
1365 |    "cell_type": "markdown",
1366 |    "id": "1bb3beed-1c44-4364-bde9-b29ea5e8ca30",
1367 |    "metadata": {},
1368 |    "source": [
1369 |     "Current state is at the top. You can match the `thread_ts`.\n",
1370 |     "Notice the `parent_config`, `thread_ts` of the new node - it is the previous node.\n"
1371 |    ]
1372 |   },
1373 |   {
1374 |    "cell_type": "code",
1375 |    "execution_count": null,
1376 |    "id": "231c7011-afe4-44bf-9e19-16220004912f",
1377 |    "metadata": {
1378 |     "height": 81
1379 |    },
1380 |    "outputs": [],
1381 |    "source": [
1382 |     "for i, state in enumerate(graph.get_state_history(thread2)):\n",
1383 |     "    if i >= 3:  # print latest 3\n",
1384 |     "        break\n",
1385 |     "    print(state, \"\\n\")"
1386 |    ]
1387 |   },
1388 |   {
1389 |    "cell_type": "markdown",
1390 |    "id": "62d34cd8-99b6-4224-b345-fe8475f2a602",
1391 |    "metadata": {},
1392 |    "source": [
1393 |     "### Try again with `as_node`\n",
1394 |     "\n",
1395 |     "When writing using `update_state()`, you want to define to the graph logic which node should be assumed as the writer. What this does is allow th graph logic to find the node on the graph. After writing the values, the `next()` value is computed by travesing the graph using the new state. In this case, the state we have was written by `Node1`. The graph can then compute the next state as being `Node2`. Note that in some graphs, this may involve going through conditional edges! Let's try this out.\n"
1396 |    ]
1397 |   },
1398 |   {
1399 |    "cell_type": "code",
1400 |    "execution_count": null,
1401 |    "id": "7bfda2f6-5887-40fa-a733-a4357ad857d8",
1402 |    "metadata": {
1403 |     "height": 30
1404 |    },
1405 |    "outputs": [],
1406 |    "source": [
1407 |     "graph.update_state(thread2, save_state.values, as_node=\"Node1\")"
1408 |    ]
1409 |   },
1410 |   {
1411 |    "cell_type": "code",
1412 |    "execution_count": null,
1413 |    "id": "0788ac0f-5f64-49b2-b335-3094bbb19143",
1414 |    "metadata": {
1415 |     "height": 81
1416 |    },
1417 |    "outputs": [],
1418 |    "source": [
1419 |     "for i, state in enumerate(graph.get_state_history(thread2)):\n",
1420 |     "    if i >= 3:  # print latest 3\n",
1421 |     "        break\n",
1422 |     "    print(state, \"\\n\")"
1423 |    ]
1424 |   },
1425 |   {
1426 |    "cell_type": "markdown",
1427 |    "id": "18a8f52e-d1a6-46a0-a519-6fda0cb91624",
1428 |    "metadata": {},
1429 |    "source": [
1430 |     "`invoke` will run from the current state if not given a particular `thread_ts`. This is now the entry that was just added.\n"
1431 |    ]
1432 |   },
1433 |   {
1434 |    "cell_type": "code",
1435 |    "execution_count": null,
1436 |    "id": "9650d355-07c1-4a7d-a21d-4231673a7ee9",
1437 |    "metadata": {
1438 |     "height": 30
1439 |    },
1440 |    "outputs": [],
1441 |    "source": [
1442 |     "graph.invoke(None, thread2)"
1443 |    ]
1444 |   },
1445 |   {
1446 |    "cell_type": "markdown",
1447 |    "id": "c59d3069-8fe7-4a05-9bfd-8e3a2e4590ad",
1448 |    "metadata": {},
1449 |    "source": [
1450 |     "Print out the state history, notice the `scratch` value change on the latest entries.\n"
1451 |    ]
1452 |   },
1453 |   {
1454 |    "cell_type": "code",
1455 |    "execution_count": null,
1456 |    "id": "35088704-2b18-4d44-ba0e-dccbfde6ed9e",
1457 |    "metadata": {
1458 |     "height": 47
1459 |    },
1460 |    "outputs": [],
1461 |    "source": [
1462 |     "for state in graph.get_state_history(thread2):\n",
1463 |     "    print(state, \"\\n\")"
1464 |    ]
1465 |   },
1466 |   {
1467 |    "cell_type": "markdown",
1468 |    "id": "a5ea1bba-cf12-4493-948b-08b300054742",
1469 |    "metadata": {},
1470 |    "source": [
1471 |     "Continue to experiment!\n"
1472 |    ]
1473 |   },
1474 |   {
1475 |    "cell_type": "markdown",
1476 |    "id": "2c6782f1",
1477 |    "metadata": {},
1478 |    "source": []
1479 |   },
1480 |   {
1481 |    "cell_type": "markdown",
1482 |    "id": "199a41bb",
1483 |    "metadata": {},
1484 |    "source": []
1485 |   }
1486 |  ],
1487 |  "metadata": {
1488 |   "kernelspec": {
1489 |    "display_name": "agents-dev-env",
1490 |    "language": "python",
1491 |    "name": "agents-dev-env"
1492 |   },
1493 |   "language_info": {
1494 |    "codemirror_mode": {
1495 |     "name": "ipython",
1496 |     "version": 3
1497 |    },
1498 |    "file_extension": ".py",
1499 |    "mimetype": "text/x-python",
1500 |    "name": "python",
1501 |    "nbconvert_exporter": "python",
1502 |    "pygments_lexer": "ipython3",
1503 |    "version": "3.10.14"
1504 |   }
1505 |  },
1506 |  "nbformat": 4,
1507 |  "nbformat_minor": 5
1508 | }
1509 | 


--------------------------------------------------------------------------------