Merge pull request #84 from langchain-ai/brace/new-memory

feat: New memory system

.gitignore

@@ -24,7 +24,12 @@ yarn-error.log
 src/tests/data/screenshots/
 src/agents/verify-reddit-post/nodes/tests/data/
 src/agents/verify-reddit-post/nodes/tests/data/openai_o1_vs_recent_leetcode_questions.json
+src/agents/curate-data/nodes/tweets/tests/data/*
 src/clients/reddit/.secrets/
+
 # LangGraph API
 .langgraph_api
-src/agents/curate-data/nodes/tweets/tests/data/*
+
+__pycache__/
+.mypy_cache/
+.ruff_cache/

.vscode/settings.json

@@ -6,5 +6,6 @@
     "subreddits",
     "Supabase",
     "Userless"
-  ]
+  ],
+  "python.languageServer": "None"
 }

langgraph.json

@@ -1,5 +1,6 @@
 {
   "node_version": "20",
+  "python_version": "3.12",
   "graphs": {
     "ingest_data": "./src/agents/ingest-data/ingest-data-graph.ts:graph",
     "generate_post": "./src/agents/generate-post/generate-post-graph.ts:generatePostGraph",
@@ -10,9 +11,10 @@
     "verify_reddit_post": "./src/agents/verify-reddit-post/verify-reddit-post-graph.ts:verifyRedditPostGraph",
     "verify_tweet": "./src/agents/verify-tweet/verify-tweet-graph.ts:verifyTweetGraph",
     "supervisor": "./src/agents/supervisor/supervisor-graph.ts:supervisorGraph",
-    "generate_report": "./src/agents/generate-report/index.ts:generateReportGraph"
+    "generate_report": "./src/agents/generate-report/index.ts:generateReportGraph",
+    "memory": "./memory-v2/memory_v2/graph.py:graph"
   },
   "env": ".env",
-  "dependencies": ["."],
+  "dependencies": [".", "./memory-v2"],
   "dockerfile_lines": ["RUN npx -y [email protected] install --with-deps"]
 }

memory-v2/Makefile

@@ -0,0 +1,66 @@
+.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)
+
+integration_tests:
+	python -m pytest tests/integration_tests 
+
+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 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=<test_file>   - run all tests in file'
+	@echo 'test_watch                   - run unit tests in watch mode'

memory-v2/langgraph.json

@@ -0,0 +1,7 @@
+{
+  "graphs": {
+    "reflection_v2": "./memory_v2/graph.py:graph"
+  },
+  "env": "../.env",
+  "dependencies": ["."]
+}

memory-v2/memory_v2/graph.py

@@ -0,0 +1,146 @@
+"""The reflection graph."""
+
+from typing import Any, Dict
+
+from langchain_anthropic import ChatAnthropic
+from langgraph.graph import StateGraph
+from langgraph.store.base import BaseStore
+from langmem.prompts.looping import (
+    Prompt,
+    create_prompt_optimizer,
+)
+
+from memory_v2.state import State
+
+REFLECTION_PROMPT = """You are helping an AI assistant learn by optimizing its prompt.
+
+<context>
+You are provided with two social media posts: the original, and the revised version.
+The post was revised based on user feedback, which is also provided to you.
+</context>
+
+You are an AI assistant tasked with analyzing social media post revisions and user feedback to determine if a new rule should be created for future post modifications.
+Your goal is to identify patterns in the changes requested by the user and decide if these changes should be applied automatically in the future.
+
+You will be given three pieces of information:
+
+1. The original social media post:
+<original_post>
+{ORIGINAL_POST}
+</original_post>
+
+2. The revised post:
+<new_post>
+{NEW_POST}
+</new_post>
+
+3. The user's response to the revision:
+<user_response>
+{USER_RESPONSE}
+</user_response>
+
+Carefully analyze these three elements, paying attention to the following:
+1. What specific changes were made between the original and new post?
+2. How did the user respond to these changes?
+3. Is there a clear pattern or preference expressed by the user?
+4. Could this preference be generalized into a rule for future posts?
+
+Based on your analysis, decide if a new rule should be created. Consider the following:
+1. Is the change specific enough to be applied consistently?
+2. Would applying this change automatically improve future posts?
+3. Is there any potential downside to always making this change?
+
+If you determine that a new rule should be created, formulate it clearly and concisely. The rule should be specific enough to be applied consistently but general enough to cover similar situations in the future.
+You should not be generating a rule which is specific to this post, like business logic. The rule, if created, should be applicable to any future post.
+
+Provide your analysis and decision in the following format:
+
+<analysis>
+[Your detailed analysis of the changes and user response]
+</analysis>
+
+<decision>
+[Your decision on whether a new rule should be created, along with your reasoning]
+</decision>
+
+If applicable, call the 'new_rule' tool to create the new rule. If no new rule is needed, simply write "No new rule required."
+
+Remember to be thorough in your analysis, clear in your decision-making, and precise in your rule formulation if one is needed."""
+
+
+REFLECTIONS_NAMESPACE = ("reflection_rules",)
+REFLECTIONS_KEY = "rules"
+PROMPT_KEY = "prompt"
+
+
+async def aget_reflections(store: BaseStore) -> str:
+    """Get reflections from the store."""
+    reflections = await store.aget(REFLECTIONS_NAMESPACE, REFLECTIONS_KEY)
+
+    if not reflections:
+        return "No prompt rules have been created yet."
+
+    ruleset = reflections.value.get(
+        PROMPT_KEY, "No prompt rules have been created yet."
+    )
+
+    return ruleset
+
+
+async def aput_reflections(store: BaseStore, reflections: str) -> None:
+    """Put reflections in the store."""
+    await store.aput(REFLECTIONS_NAMESPACE, REFLECTIONS_KEY, {PROMPT_KEY: reflections})
+
+
+async def reflection(state: State, store: BaseStore) -> Dict[str, Any]:
+    """Process reflection and update rules based on user interaction."""
+    model = ChatAnthropic(model="claude-3-5-sonnet-latest", temperature=0)
+
+    current_reflections_prompt = await aget_reflections(store)
+
+    update_instructions = """Analyze the following to determine if rules prompt updates are needed:
+1. Current rules prompt (current_prompt)
+2. Generated social media post (session)
+3. User feedback on the post (feedback)
+
+If the user's feedback explicitly requests changes:
+1. Create or update rules that directly address the feedback
+2. Keep each rule clear, specific, and concise
+3. If a new rule conflicts with an existing one, use the new rule
+4. Only add rules that are explicitly mentioned in the user's feedback
+
+Guidelines for updates:
+- Do not infer or assume rules beyond what's explicitly stated
+- Do not add rules based on implicit feedback
+- Do not overgeneralize the feedback
+- Combine existing rules if it improves clarity without losing specificity
+
+Output only the updated rules prompt, with no additional context or instructions."""
+
+    feedback = state.user_response
+
+    prompt = Prompt(
+        name="Update Prompt",
+        prompt=current_reflections_prompt,
+        update_instructions=update_instructions,
+        feedback=feedback,
+    )
+
+    sessions = state.original_post
+
+    optimizer = create_prompt_optimizer(model, kind="metaprompt")
+
+    result = await optimizer(sessions, prompt)
+
+    await aput_reflections(store, result)
+
+    return {}
+
+
+# Define a new graph
+workflow = StateGraph(State)
+workflow.add_node("reflection", reflection)
+workflow.add_edge("__start__", "reflection")
+
+graph = workflow.compile()
+graph.name = "Reflection Graph"

memory-v2/memory_v2/state.py

@@ -0,0 +1,15 @@
+"""Define the state structures for the agent."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class State:
+    """The state of the memory graph."""
+
+    original_post: str = ""
+    """The original post that the user submitted feedback on"""
+    user_response: str = ""
+    """The user's feedback on the new post"""