Use new thing

src/memory_service/graph.py

@@ -12,12 +12,11 @@ from langchain.chat_models import init_chat_model
 from langchain_core.runnables import RunnableConfig
 from langgraph.constants import Send
 from langgraph.graph import StateGraph
-from langgraph.store.store import Store
-from trustcall import create_extractor
-
+from langgraph.store.base import BaseStore
 from memory_service import _configuration as configuration
 from memory_service import _utils as utils
 from memory_service import state as schemas
+from trustcall import create_extractor
 
 logger = logging.getLogger("memory")
 
@@ -36,7 +35,7 @@ async def _extract_memory(
 
 
 async def handle_patch_memory(
-    state: schemas.ProcessorState, config: RunnableConfig, *, store: Store
+    state: schemas.ProcessorState, config: RunnableConfig, *, store: BaseStore
 ) -> dict:
     """Extract the user's state from the conversation and update the memory."""
     configurable = configuration.Configuration.from_runnable_config(config)
@@ -59,7 +58,7 @@ async def handle_patch_memory(
 
 
 async def handle_insertion_memory(
-    state: schemas.ProcessorState, config: RunnableConfig, *, store: Store
+    state: schemas.ProcessorState, config: RunnableConfig, *, store: BaseStore
 ) -> dict:
     """Upsert memory events."""
     configurable = configuration.Configuration.from_runnable_config(config)

tests/evals/test_store.py

@@ -4,14 +4,14 @@ from typing import Literal, Optional
 
 import langsmith as ls
 from langgraph.store.memory import MemoryStore
-from langgraph.store.store import Store
+from langgraph.store.base import BaseStore
 from pydantic import BaseModel, Field
 
 from memory_service.graph import builder
 
 
 class User(BaseModel):
-    """Store all important information about a user here."""
+    """BaseStore all important information about a user here."""
 
     preferred_name: Optional[str] = None
     current_age: Optional[str] = None
@@ -28,7 +28,7 @@ class User(BaseModel):
         description="Other preferences the user has expressed that informs how you should interact with them."
     )
     relationships: list[str] = Field(
-        description="Store information about friends, family members, coworkers, and other important relationships the user has here. Include relevant information about htem."
+        description="BaseStore information about friends, family members, coworkers, and other important relationships the user has here. Include relevant information about htem."
     )
 
 
@@ -69,8 +69,7 @@ async def test_patch_memory_stored():
         {"messages": [("user", "My name is Bob. I like fun things")]}, config
     )
     namespace = ("user_states", user_id, "User")
-    store = Store(mem_store)
-    memories = store.search(namespace, query=None, filter=None, weights=None)
+    memories = mem_store.search(namespace, query=None, filter=None, weights=None)
     ls.expect(len(memories)).to_equal(1)
     mem = memories[0]
     ls.expect(mem.value.get("preferred_name")).to_equal("Bob")
@@ -83,14 +82,14 @@ async def test_patch_memory_stored():
         },
         config,
     )
-    memories = store.search(namespace, query=None, filter=None, weights=None)
+    memories = mem_store.search(namespace, query=None, filter=None, weights=None)
     ls.expect(len(memories)).to_equal(1)
     mem = memories[0]
     ls.expect(mem.value.get("preferred_name")).to_equal("Robert")
 
     # Check that searching by a different namespace returns no memories
     bad_namespace = ("user_states", "my-bad-test-user", "User")
-    memories = store.search(bad_namespace, query=None, filter=None, weights=None)
+    memories = mem_store.search(bad_namespace, query=None, filter=None, weights=None)
     ls.expect(memories).against(lambda x: not x)
 
 
@@ -167,8 +166,7 @@ async def test_insertion_memory_stored():
         {**config, "thread_id": thread_id},
     )
     namespace = ("events", user_id, "Relationship")
-    store = Store(mem_store)
-    memories = store.search(namespace, query=None, filter=None, weights=None)
+    memories = mem_store.search(namespace, query=None, filter=None, weights=None)
     ls.expect(len(memories)).to_be_greater_than(1)
     # Check for Joanne's relationship
     joanne_relationship = next(
@@ -210,7 +208,9 @@ async def test_insertion_memory_stored():
     )
 
     # Check the memories again
-    updated_memories = store.search(namespace, query=None, filter=None, weights=None)
+    updated_memories = mem_store.search(
+        namespace, query=None, filter=None, weights=None
+    )
     ls.expect(len(updated_memories)).to_equal(
         3
     )  # Now there should be 3 objects: Nick, Joanne, and Anthony