mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 14:42:28 +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)
```
44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
from collections.abc import AsyncIterator
|
|
|
|
import pytest
|
|
from psycopg import AsyncConnection
|
|
from psycopg.errors import UndefinedTable
|
|
from psycopg.rows import DictRow, dict_row
|
|
|
|
from tests.embed_test_utils import CharacterEmbeddings
|
|
|
|
DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
async def conn() -> AsyncIterator[AsyncConnection[DictRow]]:
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_URI, autocommit=True, prepare_threshold=0, row_factory=dict_row
|
|
) as conn:
|
|
yield conn
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=True)
|
|
async def clear_test_db(conn: AsyncConnection[DictRow]) -> None:
|
|
"""Delete all tables before each test."""
|
|
try:
|
|
await conn.execute("DELETE FROM checkpoints")
|
|
await conn.execute("DELETE FROM checkpoint_blobs")
|
|
await conn.execute("DELETE FROM checkpoint_writes")
|
|
await conn.execute("DELETE FROM checkpoint_migrations")
|
|
except UndefinedTable:
|
|
pass
|
|
try:
|
|
await conn.execute("DELETE FROM store_migrations")
|
|
await conn.execute("DELETE FROM store")
|
|
except UndefinedTable:
|
|
pass
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_embeddings() -> CharacterEmbeddings:
|
|
return CharacterEmbeddings(dims=500)
|
|
|
|
|
|
VECTOR_TYPES = ["vector", "halfvec"]
|