diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index f5fe83362..fb00f80c3 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -19,6 +19,7 @@ from langgraph.store.base import ( Result, SearchOp, ensure_embeddings, + tokenize_path, ) from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.postgres.base import ( @@ -70,6 +71,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con self.supports_pipeline = Capabilities().has_pipeline() self.embedding_config = embedding if self.embedding_config: + self.embedding_config = self.embedding_config.copy() + self.embedding_config["__tokenized_fields"] = [ + (p, tokenize_path(p)) if p != "__root__" else (p, p) + for p in (self.embedding_config.get("text_fields") or ["__root__"]) + ] self.embeddings: Optional[Embeddings] = ensure_embeddings( self.embedding_config.get("embed"), aembed=self.embedding_config.get("aembed"), @@ -259,7 +265,6 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con else: async with ( self.lock, - conn.transaction(), conn.cursor(binary=True) as cur, ): yield cur diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index f9151f6ff..93f04bf75 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -42,6 +42,8 @@ from langgraph.store.base import ( SearchItem, SearchOp, ensure_embeddings, + get_text_at_path, + tokenize_path, ) if TYPE_CHECKING: @@ -345,11 +347,7 @@ class BasePostgresStore(Generic[C]): # Then handle embeddings if configured if self.embedding_config: - text_fields = self.embedding_config.get("text_fields", ["__root__"]) - if isinstance(text_fields, str): - text_fields = [text_fields] - elif text_fields is None: - text_fields = ["__root__"] + paths = self.embedding_config["__tokenized_fields"] for op in inserts: if op.index is False: continue @@ -357,12 +355,12 @@ class BasePostgresStore(Generic[C]): ns = _namespace_to_text(op.namespace) k = op.key - for field in text_fields: - for text in _extract_text_by_path(value, field): + for path, tokenized_path in paths: + for text in get_text_at_path(value, tokenized_path): vector_values.append( "(%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" ) - embedding_request_params.append((ns, k, field, text)) + embedding_request_params.append((ns, k, path, text)) values_str = ",".join(values) query = f""" @@ -556,6 +554,11 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): self.lock = threading.Lock() self.embedding_config = embedding if self.embedding_config: + self.embedding_config = self.embedding_config.copy() + self.embedding_config["__tokenized_fields"] = [ + (p, tokenize_path(p)) if p != "__root__" else (p, p) + for p in (self.embedding_config.get("text_fields") or ["__root__"]) + ] self.embeddings: Optional[Embeddings] = ensure_embeddings( self.embedding_config.get("embed"), aembed=self.embedding_config.get("aembed"), @@ -931,160 +934,6 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]: return tuple(namespace.split(".")) -def _tokenize_path(path: str) -> list[str]: - """Tokenize a path into components. - - Handles: - - Simple paths: "field1.field2" - - Array indexing: "[0]", "[*]", "[-1]" - - Wildcards: "*" - - Multi-field selection: "{field1,field2}" - """ - if not path: - return [] - - tokens = [] - current: list[str] = [] - i = 0 - while i < len(path): - char = path[i] - - if char == "[": # Handle array index - if current: - tokens.append("".join(current)) - current = [] - bracket_count = 1 - index_chars = ["["] - i += 1 - while i < len(path) and bracket_count > 0: - if path[i] == "[": - bracket_count += 1 - elif path[i] == "]": - bracket_count -= 1 - index_chars.append(path[i]) - i += 1 - tokens.append("".join(index_chars)) - continue - - elif char == "{": # Handle multi-field selection - if current: - tokens.append("".join(current)) - current = [] - brace_count = 1 - field_chars = ["{"] - i += 1 - while i < len(path) and brace_count > 0: - if path[i] == "{": - brace_count += 1 - elif path[i] == "}": - brace_count -= 1 - field_chars.append(path[i]) - i += 1 - tokens.append("".join(field_chars)) - continue - - elif char == ".": # Handle regular field - if current: - tokens.append("".join(current)) - current = [] - else: - current.append(char) - i += 1 - - if current: - tokens.append("".join(current)) - - return tokens - - -def _extract_text_by_path(obj: Any, path: str) -> list[str]: - """Extract text from an object using a path expression. - - Supports: - - Simple paths: "field1.field2" - - Array indexing: "[0]", "[*]", "[-1]" - - Wildcards: "*" - - Multi-field selection: "{field1,field2}" - - Nested paths in multi-field: "{field1,nested.field2}" - """ - if not path or path == "__root__": - return [json.dumps(obj, sort_keys=True)] - - def _extract_from_obj(obj: Any, tokens: list[str], pos: int) -> list[str]: - if pos >= len(tokens): - if isinstance(obj, (str, int, float, bool)): - return [str(obj)] - elif obj is None: - return [] - elif isinstance(obj, (list, dict)): - return [json.dumps(obj, sort_keys=True)] - return [] - - token = tokens[pos] - results = [] - - if token.startswith("[") and token.endswith("]"): - if not isinstance(obj, list): - return [] - - index = token[1:-1] - if index == "*": - for item in obj: - results.extend(_extract_from_obj(item, tokens, pos + 1)) - else: - try: - idx = int(index) - if idx < 0: - idx = len(obj) + idx - if 0 <= idx < len(obj): - results.extend(_extract_from_obj(obj[idx], tokens, pos + 1)) - except (ValueError, IndexError): - return [] - - elif token.startswith("{") and token.endswith("}"): - if not isinstance(obj, dict): - return [] - - fields = [f.strip() for f in token[1:-1].split(",")] - for field in fields: - nested_tokens = _tokenize_path(field) - if nested_tokens: - current_obj: Optional[dict] = obj - for nested_token in nested_tokens: - if ( - isinstance(current_obj, dict) - and nested_token in current_obj - ): - current_obj = current_obj[nested_token] - else: - current_obj = None - break - if current_obj is not None: - if isinstance(current_obj, (str, int, float, bool)): - results.append(str(current_obj)) - elif isinstance(current_obj, (list, dict)): - results.append(json.dumps(current_obj, sort_keys=True)) - - # Handle wildcard - elif token == "*": - if isinstance(obj, dict): - for value in obj.values(): - results.extend(_extract_from_obj(value, tokens, pos + 1)) - elif isinstance(obj, list): - for item in obj: - results.extend(_extract_from_obj(item, tokens, pos + 1)) - - # Handle regular field - else: - if isinstance(obj, dict) and token in obj: - results.extend(_extract_from_obj(obj[token], tokens, pos + 1)) - - return results - - tokens = _tokenize_path(path) - return _extract_from_obj(obj, tokens, 0) - - def _get_distance_operator(store: Any) -> tuple[str, str]: """Get the distance operator and score expression based on config.""" if not store.embedding_config: diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index 39f9455a7..b4b4debc0 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -1,6 +1,5 @@ # type: ignore -import json from uuid import uuid4 import pytest @@ -22,7 +21,6 @@ from langgraph.store.base import ( SearchOp, ) from langgraph.store.postgres import PostgresStore -from langgraph.store.postgres.base import _extract_text_by_path @pytest.fixture(scope="function", params=["default", "pipe", "pool"]) @@ -535,71 +533,3 @@ def test_vector_search_edge_cases(vector_store: PostgresStore) -> None: special_query = "test!@#$%^&*()" results = vector_store.search(("test",), query=special_query) assert len(results) == 1 - - -def test_extract_text_by_path(): - nested_data = { - "name": "test", - "info": { - "age": 25, - "tags": ["a", "b", "c"], - "metadata": {"created": "2024-01-01", "updated": "2024-01-02"}, - }, - "items": [ - {"id": 1, "value": "first", "tags": ["x", "y"]}, - {"id": 2, "value": "second", "tags": ["y", "z"]}, - {"id": 3, "value": "third", "tags": ["z", "w"]}, - ], - "empty": None, - "zeros": [0, 0.0, "0"], - "empty_list": [], - "empty_dict": {}, - } - - assert _extract_text_by_path(nested_data, "__root__") == [ - json.dumps(nested_data, sort_keys=True) - ] - - assert _extract_text_by_path(nested_data, "name") == ["test"] - assert _extract_text_by_path(nested_data, "info.age") == ["25"] - - assert _extract_text_by_path(nested_data, "info.metadata.created") == ["2024-01-01"] - - assert _extract_text_by_path(nested_data, "items[0].value") == ["first"] - assert _extract_text_by_path(nested_data, "items[-1].value") == ["third"] - assert _extract_text_by_path(nested_data, "items[1].tags[0]") == ["y"] - - values = _extract_text_by_path(nested_data, "items[*].value") - assert set(values) == {"first", "second", "third"} - - metadata_dates = _extract_text_by_path(nested_data, "info.metadata.*") - assert set(metadata_dates) == {"2024-01-01", "2024-01-02"} - name_and_age = _extract_text_by_path(nested_data, "{name,info.age}") - assert set(name_and_age) == {"test", "25"} - - item_fields = _extract_text_by_path(nested_data, "items[*].{id,value}") - assert set(item_fields) == {"1", "2", "3", "first", "second", "third"} - - all_tags = _extract_text_by_path(nested_data, "items[*].tags[*]") - assert set(all_tags) == {"x", "y", "z", "w"} - - assert _extract_text_by_path(None, "any.path") == [] - assert _extract_text_by_path({}, "any.path") == [] - assert _extract_text_by_path(nested_data, "") == [ - json.dumps(nested_data, sort_keys=True) - ] - assert _extract_text_by_path(nested_data, "nonexistent") == [] - assert _extract_text_by_path(nested_data, "items[99].value") == [] - assert _extract_text_by_path(nested_data, "items[*].nonexistent") == [] - - assert _extract_text_by_path(nested_data, "empty") == [] - assert _extract_text_by_path(nested_data, "empty_list") == ["[]"] - assert _extract_text_by_path(nested_data, "empty_dict") == ["{}"] - - zeros = _extract_text_by_path(nested_data, "zeros[*]") - assert set(zeros) == {"0", "0.0"} - - assert _extract_text_by_path(nested_data, "items[].value") == [] - assert _extract_text_by_path(nested_data, "items[abc].value") == [] - assert _extract_text_by_path(nested_data, "{unclosed") == [] - assert _extract_text_by_path(nested_data, "nested[{invalid}]") == []