diff --git a/.codespellignore b/.codespellignore new file mode 100644 index 0000000..e69de29 diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..342da42 --- /dev/null +++ b/.env.example @@ -0,0 +1,9 @@ +# To separate your traces from other application +LANGSMITH_PROJECT=memory-graph + +# The following depend on your selected configuration + +## LLM choice: +# ANTHROPIC_API_KEY=.... +# FIREWORKS_API_KEY=... +# OPENAI_API_KEY=... diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml new file mode 100644 index 0000000..4191038 --- /dev/null +++ b/.github/workflows/integration-tests.yml @@ -0,0 +1,43 @@ +# This workflow will run integration tests for the current project once per day + +name: Integration Tests + +on: + schedule: + - cron: "37 14 * * *" # Run at 7:37 AM Pacific Time (14:37 UTC) every day + workflow_dispatch: # Allows triggering the workflow manually in GitHub UI + +# If another scheduled run starts while this workflow is still running, +# cancel the earlier run in favor of the next run. +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + integration-tests: + name: Integration Tests + strategy: + matrix: + os: [ubuntu-latest] + python-version: ["3.11", "3.12"] + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + curl -LsSf https://astral.sh/uv/install.sh | sh + uv venv + uv pip install -r pyproject.toml --extra dev + uv pip install -U pytest-asyncio vcrpy + - name: Run integration tests + env: + ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} + TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }} + LANGSMITH_API_KEY: ${{ secrets.LANGSMITH_API_KEY }} + LANGSMITH_TRACING: true + run: | + uv run pytest tests/integration_tests diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml new file mode 100644 index 0000000..055407c --- /dev/null +++ b/.github/workflows/unit-tests.yml @@ -0,0 +1,57 @@ +# This workflow will run unit tests for the current project + +name: CI + +on: + push: + branches: ["main"] + pull_request: + workflow_dispatch: # Allows triggering the workflow manually in GitHub UI + +# If another push to the same PR or branch happens while this workflow is still running, +# cancel the earlier run in favor of the next run. +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + unit-tests: + name: Unit Tests + strategy: + matrix: + os: [ubuntu-latest] + python-version: ["3.11", "3.12"] + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + curl -LsSf https://astral.sh/uv/install.sh | sh + uv venv + uv pip install -r pyproject.toml + - name: Lint with ruff + run: | + uv pip install ruff + uv run ruff check . + - name: Lint with mypy + run: | + uv pip install mypy + uv run mypy --strict src/ + - name: Check README spelling + uses: codespell-project/actions-codespell@v2 + with: + ignore_words_file: .codespellignore + path: README.md + - name: Check code spelling + uses: codespell-project/actions-codespell@v2 + with: + ignore_words_file: .codespellignore + path: src/ + - name: Run tests with pytest + run: | + uv pip install pytest + uv run pytest tests/unit_tests diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b8af3ff --- /dev/null +++ b/.gitignore @@ -0,0 +1,164 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ +.DS_Store +uv.lock diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..57d0481 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 LangChain + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..2f6583a --- /dev/null +++ b/Makefile @@ -0,0 +1,65 @@ +.PHONY: all format lint test tests test_watch integration_tests docker_tests help extended_tests + +# Default target executed when no arguments are given to make. +all: help + +# Define a variable for the test file path. +TEST_FILE ?= tests/unit_tests/ + +test: + python -m pytest $(TEST_FILE) + +test_watch: + python -m ptw --snapshot-update --now . -- -vv tests/unit_tests + +test_profile: + python -m pytest -vv tests/unit_tests/ --profile-svg + +extended_tests: + python -m pytest --only-extended $(TEST_FILE) + + +###################### +# LINTING AND FORMATTING +###################### + +# Define a variable for Python and notebook files. +PYTHON_FILES=src/ +MYPY_CACHE=.mypy_cache +lint: PYTHON_FILES=src +format: PYTHON_FILES=. +lint_diff format_diff: PYTHON_FILES=$(shell git diff --name-only --diff-filter=d main | grep -E '\.py$$|\.ipynb$$') +lint_package: PYTHON_FILES=src +lint_tests: PYTHON_FILES=tests +lint_tests: MYPY_CACHE=.mypy_cache_test + +lint lint_diff lint_package lint_tests: + python -m ruff check . + [ "$(PYTHON_FILES)" = "" ] || python -m ruff format $(PYTHON_FILES) --diff + [ "$(PYTHON_FILES)" = "" ] || python -m ruff check --select I $(PYTHON_FILES) + [ "$(PYTHON_FILES)" = "" ] || python -m mypy --strict $(PYTHON_FILES) + [ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) && python -m mypy --strict $(PYTHON_FILES) --cache-dir $(MYPY_CACHE) + +format format_diff: + ruff format $(PYTHON_FILES) + ruff check --select I --fix $(PYTHON_FILES) + +spell_check: + codespell --toml pyproject.toml + +spell_fix: + codespell --toml pyproject.toml -w + +###################### +# HELP +###################### + +help: + @echo '----' + @echo 'format - run code formatters' + @echo 'lint - run linters' + @echo 'test - run unit tests' + @echo 'tests - run unit tests' + @echo 'test TEST_FILE= - run all tests in file' + @echo 'test_watch - run unit tests in watch mode' + diff --git a/README.md b/README.md new file mode 100644 index 0000000..ca76785 --- /dev/null +++ b/README.md @@ -0,0 +1,224 @@ +# LangGraph ReAct Memory Agent + +[![CI](https://github.com/langchain-ai/memory-agent/actions/workflows/unit-tests.yml/badge.svg)](https://github.com/langchain-ai/memory-agent/actions/workflows/unit-tests.yml) +[![Integration Tests](https://github.com/langchain-ai/memory-agent/actions/workflows/integration-tests.yml/badge.svg)](https://github.com/langchain-ai/memory-agent/actions/workflows/integration-tests.yml) +[![Open in - LangGraph Studio](https://img.shields.io/badge/Open_in-LangGraph_Studio-00324d.svg?logo=data:image/svg%2bxml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSI4NS4zMzMiIGhlaWdodD0iODUuMzMzIiB2ZXJzaW9uPSIxLjAiIHZpZXdCb3g9IjAgMCA2NCA2NCI+PHBhdGggZD0iTTEzIDcuOGMtNi4zIDMuMS03LjEgNi4zLTYuOCAyNS43LjQgMjQuNi4zIDI0LjUgMjUuOSAyNC41QzU3LjUgNTggNTggNTcuNSA1OCAzMi4zIDU4IDcuMyA1Ni43IDYgMzIgNmMtMTIuOCAwLTE2LjEuMy0xOSAxLjhtMzcuNiAxNi42YzIuOCAyLjggMy40IDQuMiAzLjQgNy42cy0uNiA0LjgtMy40IDcuNkw0Ny4yIDQzSDE2LjhsLTMuNC0zLjRjLTQuOC00LjgtNC44LTEwLjQgMC0xNS4ybDMuNC0zLjRoMzAuNHoiLz48cGF0aCBkPSJNMTguOSAyNS42Yy0xLjEgMS4zLTEgMS43LjQgMi41LjkuNiAxLjcgMS44IDEuNyAyLjcgMCAxIC43IDIuOCAxLjYgNC4xIDEuNCAxLjkgMS40IDIuNS4zIDMuMi0xIC42LS42LjkgMS40LjkgMS41IDAgMi43LS41IDIuNy0xIDAtLjYgMS4xLS44IDIuNi0uNGwyLjYuNy0xLjgtMi45Yy01LjktOS4zLTkuNC0xMi4zLTExLjUtOS44TTM5IDI2YzAgMS4xLS45IDIuNS0yIDMuMi0yLjQgMS41LTIuNiAzLjQtLjUgNC4yLjguMyAyIDEuNyAyLjUgMy4xLjYgMS41IDEuNCAyLjMgMiAyIDEuNS0uOSAxLjItMy41LS40LTMuNS0yLjEgMC0yLjgtMi44LS44LTMuMyAxLjYtLjQgMS42LS41IDAtLjYtMS4xLS4xLTEuNS0uNi0xLjItMS42LjctMS43IDMuMy0yLjEgMy41LS41LjEuNS4yIDEuNi4zIDIuMiAwIC43LjkgMS40IDEuOSAxLjYgMi4xLjQgMi4zLTIuMy4yLTMuMi0uOC0uMy0yLTEuNy0yLjUtMy4xLTEuMS0zLTMtMy4zLTMtLjUiLz48L3N2Zz4=)](https://langgraph-studio.vercel.app/templates/open?githubUrl=https://github.com/langchain-ai/memory-agent) + +This repo provides a simple example of a ReAct-style agent with a tool to save memories. This is a simple way to let an agent persist important information to reuse later. In this case, we save all memories scoped to a configurable `user_id`, which lets the bot learn a user's preferences across conversational threads. + +![Memory Diagram](./static/memory_graph.png) + +## Getting Started + +This quickstart will get your memory service deployed on [LangGraph Cloud](https://langchain-ai.github.io/langgraph/cloud/). Once created, you can interact with it from any API. + +Assuming you have already [installed LangGraph Studio](https://github.com/langchain-ai/langgraph-studio?tab=readme-ov-file#download), to set up: + +1. Create a `.env` file. + +```bash +cp .env.example .env +``` + +2. Define required API keys in your `.env` file. + + + +### Setup Model + +The defaults values for `model` are shown below: + +```yaml +model: anthropic/claude-3-5-sonnet-20240620 +``` + +Follow the instructions below to get set up, or pick one of the additional options. + +#### Anthropic + +To use Anthropic's chat models: + +1. Sign up for an [Anthropic API key](https://console.anthropic.com/) if you haven't already. +2. Once you have your API key, add it to your `.env` file: + +``` +ANTHROPIC_API_KEY=your-api-key +``` + +#### OpenAI + +To use OpenAI's chat models: + +1. Sign up for an [OpenAI API key](https://platform.openai.com/signup). +2. Once you have your API key, add it to your `.env` file: + +``` +OPENAI_API_KEY=your-api-key +``` + + + +3. Open in LangGraph studio. Navigate to the `memory_agent` graph and have a conversation with it! Try sending some messages saying your name and other things the bot should remember. + +Assuming the bot saved some memories, create a _new_ thread using the `+` icon. Then chat with the bot again - if you've completed your setup correctly, the bot should now have access to the memories you've saved! + +You can review the saved memories by clicking the "memory" button. + +![Memories Explorer](./static/memories.png) + +## How it works + +This chat bot reads from your memory graph's `Store` to easily list extracted memories. If it calls a tool, LangGraph will route to the `store_memory` node to save the information to the store. + +## How to evaluate + +Memory management can be challenging to get right, especially if you add additional tools for the bot to choose between. +To tune the frequency and quality of memories your bot is saving, we recommend starting from an evaluation set, adding to it over time as you find and address common errors in your service. + +We have provided a few example evaluation cases in [the test file here](./tests/integration_tests/test_graph.py). As you can see, the metrics themselves don't have to be terribly complicated, especially not at the outset. + +We use [LangSmith's @unit decorator](https://docs.smith.langchain.com/how_to_guides/evaluation/unit_testing#write-a-test) to sync all the evaluations to LangSmith so you can better optimize your system and identify the root cause of any issues that may arise. + +## How to customize + +1. Customize memory content: we've defined a simple memory structure `content: str, context: str` for each memory, but you could structure them in other ways. +2. Provide additional tools: the bot will be more useful if you connect it to other functions. +3. Select a different model: We default to anthropic/claude-3-5-sonnet-20240620. You can select a compatible chat model using provider/model-name via configuration. Example: openai/gpt-4. +4. Customize the prompts: We provide a default prompt in the [prompts.py](src/memory_agent/prompts.py) file. You can easily update this via configuration. + + + diff --git a/langgraph.json b/langgraph.json new file mode 100644 index 0000000..baefb19 --- /dev/null +++ b/langgraph.json @@ -0,0 +1,9 @@ +{ + "dockerfile_lines": [], + "graphs": { + "agent": "./src/memory_agent/graph.py:graph" + }, + "env": ".env", + "python_version": "3.11", + "dependencies": ["."] +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..5a4f9a0 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,58 @@ +[project] +name = "memory-agent" +version = "0.0.1" +description = "An agent with a tool to save long-term memory." +authors = [ + { name = "William Fu-Hinthorn", email = "13333726+hinthornw@users.noreply.github.com" }, +] +readme = "README.md" +license = { text = "MIT" } +requires-python = ">=3.9" +dependencies = [ + "langgraph>=0.2.34,<0.3.0", + # Optional (for selecting different models) + "langchain-openai>=0.2.1", + "langchain-anthropic>=0.2.1", + "langchain>=0.3.1", + "langchain-core>=0.3.8", + "python-dotenv>=1.0.1", + "langgraph-sdk>=0.1.32", +] + +[project.optional-dependencies] +dev = ["mypy>=1.11.1", "ruff>=0.6.1", "pytest-asyncio"] + +[build-system] +requires = ["setuptools>=73.0.0", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +packages = ["memory_agent"] +[tool.setuptools.package-dir] +"memory_agent" = "src/memory_agent" +"langgraph.templates.memory_agent" = "src/memory_agent" + + +[tool.setuptools.package-data] +"*" = ["py.typed"] + +[tool.ruff] +lint.select = [ + "E", # pycodestyle + "F", # pyflakes + "I", # isort + "D", # pydocstyle + "D401", # First line should be in imperative mood + "T201", + "UP", +] +lint.ignore = ["UP006", "UP007", "UP035", "D417", "E501"] +include = ["*.py", "*.pyi", "*.ipynb"] +[tool.ruff.lint.per-file-ignores] +"tests/*" = ["D", "UP"] +"ntbk/*" = ["D", "UP", "T201"] +[tool.ruff.lint.pydocstyle] +convention = "google" + +[tool.mypy] +ignore_errors = true diff --git a/src/memory_agent/__init__.py b/src/memory_agent/__init__.py new file mode 100644 index 0000000..17b501a --- /dev/null +++ b/src/memory_agent/__init__.py @@ -0,0 +1,5 @@ +"""Enrichment for a pre-defined schema.""" + +from memory_agent.graph import graph + +__all__ = ["graph"] diff --git a/src/memory_agent/configuration.py b/src/memory_agent/configuration.py new file mode 100644 index 0000000..6288ba8 --- /dev/null +++ b/src/memory_agent/configuration.py @@ -0,0 +1,42 @@ +"""Define the configurable parameters for the agent.""" + +import os +from dataclasses import dataclass, field, fields +from typing import Any, Optional + +from langchain_core.runnables import RunnableConfig +from typing_extensions import Annotated + +from memory_agent import prompts + + +@dataclass(kw_only=True) +class Configuration: + """Main configuration class for the memory graph system.""" + + user_id: str = "default" + """The ID of the user to remember in the conversation.""" + model: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = field( + default="anthropic/claude-3-5-sonnet-20240620", + metadata={ + "description": "The name of the language model to use for the agent. " + "Should be in the form: provider/model-name." + }, + ) + system_prompt: str = prompts.SYSTEM_PROMPT + + @classmethod + def from_runnable_config( + cls, config: Optional[RunnableConfig] = None + ) -> "Configuration": + """Create a Configuration instance from a RunnableConfig.""" + configurable = ( + config["configurable"] if config and "configurable" in config else {} + ) + values: dict[str, Any] = { + f.name: os.environ.get(f.name.upper(), configurable.get(f.name)) + for f in fields(cls) + if f.init + } + + return cls(**{k: v for k, v in values.items() if v}) diff --git a/src/memory_agent/graph.py b/src/memory_agent/graph.py new file mode 100644 index 0000000..09c9ff4 --- /dev/null +++ b/src/memory_agent/graph.py @@ -0,0 +1,105 @@ +"""Graphs that extract memories on a schedule.""" + +import asyncio +import logging +from datetime import datetime + +from langchain.chat_models import init_chat_model +from langchain_core.runnables import RunnableConfig +from langgraph.graph import END, StateGraph +from langgraph.store.base import BaseStore + +from memory_agent import configuration, tools, utils +from memory_agent.state import State + +logger = logging.getLogger(__name__) + +# Initialize the language model to be used for memory extraction +llm = init_chat_model() + + +async def call_model(state: State, config: RunnableConfig, *, store: BaseStore) -> dict: + """Extract the user's state from the conversation and update the memory.""" + configurable = configuration.Configuration.from_runnable_config(config) + + # Retrieve the most recent memories for context + memories = await store.asearch( + ("memories", config["configurable"]["user_id"]), limit=10 + ) + + # Format memories for inclusion in the prompt + formatted = "\n".join(f"[{mem.key}]: {mem.value}" for mem in memories) + if formatted: + formatted = f""" + +{formatted} +""" + + # Prepare the system prompt with user memories and current time + # This helps the model understand the context and temporal relevance + sys = configurable.system_prompt.format( + user_info=formatted, time=datetime.now().isoformat() + ) + + # Invoke the language model with the prepared prompt and tools + # "bind_tools" gives the LLM the JSON schema for all tools in the list so it knows how + # to use them. + msg = await llm.bind_tools([tools.upsert_memory]).ainvoke( + [{"role": "system", "content": sys}, *state.messages], + {"configurable": utils.split_model_and_provider(configurable.model)}, + ) + return {"messages": [msg]} + + +async def store_memory(state: State, config: RunnableConfig, *, store: BaseStore): + # Extract tool calls from the last message + tool_calls = state.messages[-1].tool_calls + + # Concurrently execute all upsert_memory calls + saved_memories = await asyncio.gather( + *( + tools.upsert_memory(**tc["args"], config=config, store=store) + for tc in tool_calls + ) + ) + + # Format the results of memory storage operations + # This provides confirmation to the model that the actions it took were completed + results = [ + { + "role": "tool", + "content": mem, + "tool_call_id": tc["id"], + } + for tc, mem in zip(tool_calls, saved_memories) + ] + return {"messages": results} + + +def route_message(state: State): + """Determine the next step based on the presence of tool calls.""" + msg = state.messages[-1] + if msg.tool_calls: + # If there are tool calls, we need to store memories + return "store_memory" + # Otherwise, finish; user can send the next message + return END + + +# Create the graph + all nodes +builder = StateGraph(State, config_schema=configuration.Configuration) + +# Define the flow of the memory extraction process +builder.add_node(call_model) +builder.add_edge("__start__", "call_model") +builder.add_node(store_memory) +builder.add_conditional_edges("call_model", route_message, ["store_memory", END]) +# Right now, we're returning control to the user after storing a memory +# Depending on the model, you may want to route back to the model +# to let it first store memories, then generate a response +builder.add_edge("store_memory", "call_model") +graph = builder.compile() +graph.name = "MemoryAgent" + + +__all__ = ["graph"] diff --git a/src/memory_agent/prompts.py b/src/memory_agent/prompts.py new file mode 100644 index 0000000..66811a9 --- /dev/null +++ b/src/memory_agent/prompts.py @@ -0,0 +1,7 @@ +"""Define default prompts.""" + +SYSTEM_PROMPT = """You are a helpful and friendly chatbot. Get to know the user! \ +Ask questions! Be spontaneous! +{user_info} + +System Time: {time}""" diff --git a/src/memory_agent/state.py b/src/memory_agent/state.py new file mode 100644 index 0000000..15a543b --- /dev/null +++ b/src/memory_agent/state.py @@ -0,0 +1,22 @@ +"""Define the shared values.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from langchain_core.messages import AnyMessage +from langgraph.graph import add_messages +from typing_extensions import Annotated + + +@dataclass(kw_only=True) +class State: + """Main graph state.""" + + messages: Annotated[list[AnyMessage], add_messages] + """The messages in the conversation.""" + + +__all__ = [ + "State", +] diff --git a/src/memory_agent/tools.py b/src/memory_agent/tools.py new file mode 100644 index 0000000..a8667f6 --- /dev/null +++ b/src/memory_agent/tools.py @@ -0,0 +1,43 @@ +"""Define he agent's tools.""" + +import uuid +from typing import Annotated, Optional + +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import InjectedToolArg +from langgraph.store.base import BaseStore + +from memory_agent.configuration import Configuration + + +async def upsert_memory( + content: str, + context: str, + *, + memory_id: Optional[uuid.UUID] = None, + # Hide these arguments from the model. + config: Annotated[RunnableConfig, InjectedToolArg], + store: Annotated[BaseStore, InjectedToolArg], +): + """Upsert a memory in the database. + + If a memory conflicts with an existing one, then just UPDATE the + existing one by passing in memory_id - don't create two memories + that are the same. If the user corrects a memory, UPDATE it. + + Args: + content: The main content of the memory. For example: + "User expressed interest in learning about French." + context: Additional context for the memory. For example: + "This was mentioned while discussing career options in Europe." + memory_id: ONLY PROVIDE IF UPDATING AN EXISTING MEMORY. + The memory to overwrite. + """ + mem_id = memory_id or uuid.uuid4() + user_id = Configuration.from_runnable_config(config).user_id + await store.aput( + ("memories", user_id), + key=str(mem_id), + value={"content": content, "context": context}, + ) + return f"Stored memory {memory_id}" diff --git a/src/memory_agent/utils.py b/src/memory_agent/utils.py new file mode 100644 index 0000000..5fd28e2 --- /dev/null +++ b/src/memory_agent/utils.py @@ -0,0 +1,11 @@ +"""Utility functions used in our graph.""" + + +def split_model_and_provider(fully_specified_name: str) -> dict: + """Initialize the configured chat model.""" + if "/" in fully_specified_name: + provider, model = fully_specified_name.split("/", maxsplit=1) + else: + provider = None + model = fully_specified_name + return {"model": model, "provider": provider} diff --git a/static/memories.png b/static/memories.png new file mode 100644 index 0000000..a7f2610 Binary files /dev/null and b/static/memories.png differ diff --git a/static/memory_graph.png b/static/memory_graph.png new file mode 100644 index 0000000..e64c23f Binary files /dev/null and b/static/memory_graph.png differ diff --git a/tests/integration_tests/__init__.py b/tests/integration_tests/__init__.py new file mode 100644 index 0000000..d02981b --- /dev/null +++ b/tests/integration_tests/__init__.py @@ -0,0 +1 @@ +"""Define any integration tests you want in this directory.""" diff --git a/tests/integration_tests/test_graph.py b/tests/integration_tests/test_graph.py new file mode 100644 index 0000000..5a859b1 --- /dev/null +++ b/tests/integration_tests/test_graph.py @@ -0,0 +1,55 @@ +from typing import List + +import langsmith as ls +import pytest +from langgraph.checkpoint.memory import MemorySaver +from langgraph.store.memory import InMemoryStore + +from memory_agent.graph import builder + + +@pytest.mark.asyncio +@ls.unit +@pytest.mark.parametrize( + "conversation", + [ + ["My name is Alice and I love pizza. Remember this."], + [ + "Hi, I'm Bob and I enjoy playing tennis. Remember this.", + "Yes, I also have a pet dog named Max.", + "Max is a golden retriever and he's 5 years old. Please remember this too.", + ], + [ + "Hello, I'm Charlie. I work as a software engineer and I'm passionate about AI. Remember this.", + "I specialize in machine learning algorithms and I'm currently working on a project involving natural language processing.", + "My main goal is to improve sentiment analysis accuracy in multi-lingual texts. It's challenging but exciting.", + "We've made some progress using transformer models, but we're still working on handling context and idioms across languages.", + "Chinese and English have been the most challenging pair so far due to their vast differences in structure and cultural contexts.", + ], + ], + ids=["short", "medium", "long"], +) +async def test_memory_storage(conversation: List[str]): + mem_store = InMemoryStore() + + graph = builder.compile(store=mem_store, checkpointer=MemorySaver()) + user_id = "test-user" + config = { + "configurable": {}, + "user_id": user_id, + } + + for content in conversation: + await graph.ainvoke( + {"messages": [("user", content)]}, + {**config, "thread_id": "thread"}, + ) + + namespace = ("memories", user_id) + memories = mem_store.search(namespace) + + ls.expect(len(memories)).to_be_greater_than(0) + + bad_namespace = ("memories", "wrong-user") + bad_memories = mem_store.search(bad_namespace) + ls.expect(len(bad_memories)).to_equal(0) diff --git a/tests/unit_tests/__init__.py b/tests/unit_tests/__init__.py new file mode 100644 index 0000000..f2900f2 --- /dev/null +++ b/tests/unit_tests/__init__.py @@ -0,0 +1 @@ +"""Define any unit tests you may want in this directory.""" diff --git a/tests/unit_tests/test_configuration.py b/tests/unit_tests/test_configuration.py new file mode 100644 index 0000000..43db0bd --- /dev/null +++ b/tests/unit_tests/test_configuration.py @@ -0,0 +1,5 @@ +from memory_agent.configuration import Configuration + + +def test_configuration_from_none() -> None: + Configuration.from_runnable_config()