mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 22:52:29 +02:00
William FHandGitHub
d767af421b
- Initializing the store with an 'embedding config' -> this contains the
'dims' (used to create the table) and the encoder object (rn langchain
embeddings object, though that is ......)
- Call setup() -> creates the vector table.
Each document has 1 or more vectors associated with it for each json
path in the embedding config.
Would welcome critique and requests!
Leaving the params as the defaults for pgvector but open to feedback if
you think it's important to be able to more transparently configure that
in setup()
```python
from typing import TypedDict, List, Dict, Any, Optional
from langchain_openai import OpenAIEmbeddings
from langgraph.graph import StateGraph
from langgraph.store.postgres import PostgresStore
emb_config = {
"dims": 1536, # OpenAI embedding dimensions
"embed": OpenAIEmbeddings(model="text-embedding-3-small"),
"distance_type": "cosine",
}
with PostgresStore.from_conn_string(
"postgres://postgres:postgres@localhost:5441",
embedding=emb_config,
) as store:
store.setup()
# Define the state type for our graph
class State(TypedDict):
query: str
results: Optional[List[Dict[str, Any]]]
def put_stuff(state: State) -> State:
docs = [
("doc1", {"text": "red apple in kitchen"}),
("doc2", {"text": "blue car in garage"}),
("doc3", {"text": "green apple on table"}),
]
for key, value in docs:
store.put(("docs",), key, value)
def search_stuff(state: State) -> State:
"""Search for documents using vector similarity."""
results = store.search(("docs",), query=state["query"])
return {"results": results}
builder = StateGraph(State)
builder.add_node(put_stuff)
builder.add_node(search_stuff)
builder.add_edge("__start__", "put_stuff")
builder.add_edge("put_stuff", "search_stuff")
# Compile
with PostgresStore.from_conn_string(
"postgres://postgres:postgres@localhost:5441",
embedding=emb_config,
) as store:
chain = builder.compile(store=store)
result = chain.invoke({"query": "sour apple"})
# Print results
for doc in result["results"]:
print(doc.key)
print(doc.value)
print(doc.response_metadata)
```
113 lines
4.0 KiB
Python
113 lines
4.0 KiB
Python
from typing import Any
|
|
|
|
import pytest
|
|
from langchain_core.runnables import RunnableConfig
|
|
|
|
from langgraph.checkpoint.base import (
|
|
Checkpoint,
|
|
CheckpointMetadata,
|
|
create_checkpoint,
|
|
empty_checkpoint,
|
|
)
|
|
from langgraph.checkpoint.postgres import PostgresSaver
|
|
from tests.conftest import DEFAULT_URI
|
|
|
|
|
|
class TestPostgresSaver:
|
|
@pytest.fixture(autouse=True)
|
|
def setup(self) -> None:
|
|
# objects for test setup
|
|
self.config_1: RunnableConfig = {
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
# for backwards compatibility testing
|
|
"thread_ts": "1",
|
|
"checkpoint_ns": "",
|
|
}
|
|
}
|
|
self.config_2: RunnableConfig = {
|
|
"configurable": {
|
|
"thread_id": "thread-2",
|
|
"checkpoint_id": "2",
|
|
"checkpoint_ns": "",
|
|
}
|
|
}
|
|
self.config_3: RunnableConfig = {
|
|
"configurable": {
|
|
"thread_id": "thread-2",
|
|
"checkpoint_id": "2-inner",
|
|
"checkpoint_ns": "inner",
|
|
}
|
|
}
|
|
|
|
self.chkpnt_1: Checkpoint = empty_checkpoint()
|
|
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
|
|
self.chkpnt_3: Checkpoint = empty_checkpoint()
|
|
|
|
self.metadata_1: CheckpointMetadata = {
|
|
"source": "input",
|
|
"step": 2,
|
|
"writes": {},
|
|
"score": 1,
|
|
}
|
|
self.metadata_2: CheckpointMetadata = {
|
|
"source": "loop",
|
|
"step": 1,
|
|
"writes": {"foo": "bar"},
|
|
"score": None,
|
|
}
|
|
self.metadata_3: CheckpointMetadata = {}
|
|
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
|
saver.setup()
|
|
|
|
def test_search(self) -> None:
|
|
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
|
# save checkpoints
|
|
saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
|
|
saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
|
|
saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
|
|
|
|
# call method / assertions
|
|
query_1 = {"source": "input"} # search by 1 key
|
|
query_2 = {
|
|
"step": 1,
|
|
"writes": {"foo": "bar"},
|
|
} # search by multiple keys
|
|
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
|
|
query_4 = {"source": "update", "step": 1} # no match
|
|
|
|
search_results_1 = list(saver.list(None, filter=query_1))
|
|
assert len(search_results_1) == 1
|
|
assert search_results_1[0].metadata == self.metadata_1
|
|
|
|
search_results_2 = list(saver.list(None, filter=query_2))
|
|
assert len(search_results_2) == 1
|
|
assert search_results_2[0].metadata == self.metadata_2
|
|
|
|
search_results_3 = list(saver.list(None, filter=query_3))
|
|
assert len(search_results_3) == 3
|
|
|
|
search_results_4 = list(saver.list(None, filter=query_4))
|
|
assert len(search_results_4) == 0
|
|
|
|
# search by config (defaults to checkpoints across all namespaces)
|
|
search_results_5 = list(
|
|
saver.list({"configurable": {"thread_id": "thread-2"}})
|
|
)
|
|
assert len(search_results_5) == 2
|
|
assert {
|
|
search_results_5[0].config["configurable"]["checkpoint_ns"],
|
|
search_results_5[1].config["configurable"]["checkpoint_ns"],
|
|
} == {"", "inner"}
|
|
|
|
# TODO: test before and limit params
|
|
|
|
def test_null_chars(self) -> None:
|
|
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
|
config = saver.put(self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {})
|
|
assert saver.get_tuple(config).metadata["my_key"] == "abc" # type: ignore
|
|
assert (
|
|
list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"] # type: ignore
|
|
== "abc"
|
|
)
|