From 46dd424a7eec665df12c9d0a35741f6b3bf57076 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Mon, 2 Dec 2024 19:53:58 +0100 Subject: [PATCH 01/13] feat(sdk): add ability to search runs via status --- libs/sdk-js/src/client.ts | 7 +++++++ libs/sdk-js/src/schema.ts | 2 +- libs/sdk-py/langgraph_sdk/client.py | 19 +++++++++++++++---- 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/libs/sdk-js/src/client.ts b/libs/sdk-js/src/client.ts index 010864419..ecf7147ca 100644 --- a/libs/sdk-js/src/client.ts +++ b/libs/sdk-js/src/client.ts @@ -7,6 +7,7 @@ import { GraphSchema, Metadata, Run, + RunStatus, Thread, ThreadState, Cron, @@ -944,12 +945,18 @@ export class RunsClient extends BaseClient { * Defaults to 0. */ offset?: number; + + /** + * Status of the run to filter by. + */ + status?: RunStatus; }, ): Promise { return this.fetch(`/threads/${threadId}/runs`, { params: { limit: options?.limit ?? 10, offset: options?.offset ?? 0, + status: options?.status ?? undefined, }, }); } diff --git a/libs/sdk-js/src/schema.ts b/libs/sdk-js/src/schema.ts index dd79e07ba..f68e3d42f 100644 --- a/libs/sdk-js/src/schema.ts +++ b/libs/sdk-js/src/schema.ts @@ -2,7 +2,7 @@ import type { JSONSchema7 } from "json-schema"; type Optional = T | null | undefined; -type RunStatus = +export type RunStatus = | "pending" | "running" | "error" diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index 81e8a8506..b7cc1a78c 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -51,6 +51,7 @@ from langgraph_sdk.schema import ( OnConflictBehavior, Run, RunCreate, + RunStatus, SearchItemsResponse, StreamMode, StreamPart, @@ -1684,7 +1685,12 @@ class RunsClient: return response async def list( - self, thread_id: str, *, limit: int = 10, offset: int = 0 + self, + thread_id: str, + *, + limit: int = 10, + offset: int = 0, + status: Optional[RunStatus] = None, ) -> List[Run]: """List runs. @@ -1692,6 +1698,7 @@ class RunsClient: thread_id: The thread ID to list runs for. limit: The maximum number of results to return. offset: The number of results to skip. + status: The status of the run to filter by. Returns: List[Run]: The runs for the thread. @@ -1705,9 +1712,13 @@ class RunsClient: ) """ # noqa: E501 - return await self.http.get( - f"/threads/{thread_id}/runs?limit={limit}&offset={offset}" - ) + params = { + "limit": limit, + "offset": offset, + } + if status is not None: + params["status"] = status + return await self.http.get(f"/threads/{thread_id}/runs", params=params) async def get(self, thread_id: str, run_id: str) -> Run: """Get a run. From 2b65308508fa829198cef3ca547b61b46ca8142e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 15 Nov 2024 14:01:16 -0800 Subject: [PATCH 02/13] WIP: Handle commands for subgraphs --- libs/langgraph/tests/test_pregel.py | 32 +++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index c2ed63d28..d45c3ae22 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -14471,3 +14471,35 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) }, tasks=(), ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_interrupt_subgraph(request: pytest.FixtureRequest, checkpointer_name: str): + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + baz: str + + def foo(state): + return {"baz": "foo"} + + def bar(state): + value = interrupt("Please provide baz value:") + return {"baz": value} + + child_builder = StateGraph(State) + child_builder.add_node(bar) + child_builder.add_edge(START, "bar") + + builder = StateGraph(State) + builder.add_node(foo) + builder.add_node("bar", child_builder.compile()) + builder.add_edge(START, "foo") + builder.add_edge("foo", "bar") + graph = builder.compile(checkpointer=checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert graph.invoke({"baz": ""}, thread1) + # Resume with answer + assert graph.invoke(Command(resume="bar"), thread1) From 6fc1c602ab151525904fcec033f327f75b0db34e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Dec 2024 13:53:11 -0800 Subject: [PATCH 03/13] Add one more --- libs/langgraph/tests/test_pregel.py | 153 ++++++++++++++++++++++++++++ 1 file changed, 153 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d45c3ae22..113a6c3c8 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8728,6 +8728,159 @@ def test_copy_checkpoint( ) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_dynamic_interrupt_subgraph( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class SubgraphState(TypedDict): + my_key: str + market: str + + tool_two_node_count = 0 + + def tool_two_node(s: SubgraphState) -> SubgraphState: + print("tool_two_node", s) + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + print("after interrupt") + return {"my_key": answer} + + subgraph = StateGraph(SubgraphState) + subgraph.add_node("do", tool_two_node, retry=RetryPolicy()) + subgraph.add_edge(START, "do") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", subgraph.compile()) + tool_two_graph.add_edge(START, "tool_two") + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert tool_two.invoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value", + "market": "DE", + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value"} + + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ) + }, + ] + # resume with answer + assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear tasks + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), + ), + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # clear the interrupt and next tasks + tool_two.update_state(thread1, None, as_node=END) + # interrupt and next tasks are cleared + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_start_branch_then( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str From 988dd237d26be56133718f06f91da107836cd982 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Tue, 3 Dec 2024 01:17:15 +0100 Subject: [PATCH 04/13] feat(sdk): pass cancel on disconnect when joining stream --- libs/sdk-js/src/client.ts | 17 +++++++++++++---- libs/sdk-py/langgraph_sdk/client.py | 11 +++++++++-- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/libs/sdk-js/src/client.ts b/libs/sdk-js/src/client.ts index ecf7147ca..14284bcff 100644 --- a/libs/sdk-js/src/client.ts +++ b/libs/sdk-js/src/client.ts @@ -1021,19 +1021,28 @@ export class RunsClient extends BaseClient { * * @param threadId The ID of the thread. * @param runId The ID of the run. - * @param signal An optional abort signal. * @returns An async generator yielding stream parts. */ async *joinStream( threadId: string, runId: string, - signal?: AbortSignal, + options?: + | { signal?: AbortSignal; cancelOnDisconnect?: boolean } + | AbortSignal, ): AsyncGenerator<{ event: StreamEvent; data: any }> { + const opts = + typeof options === "object" && + options != null && + options instanceof AbortSignal + ? { signal: options } + : options; + const response = await this.asyncCaller.fetch( ...this.prepareFetchOptions(`/threads/${threadId}/runs/${runId}/stream`, { method: "GET", timeoutMs: null, - signal, + signal: opts?.signal, + params: { cancel_on_disconnect: opts?.cancelOnDisconnect ? "1" : "0" }, }), ); @@ -1048,7 +1057,7 @@ export class RunsClient extends BaseClient { async start(ctrl) { parser = createParser((event) => { if ( - (signal && signal.aborted) || + (opts?.signal && opts.signal.aborted) || (event.type === "event" && event.data === "[DONE]") ) { ctrl.terminate(); diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index b7cc1a78c..7a959a477 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -1796,7 +1796,9 @@ class RunsClient: """ # noqa: E501 return await self.http.get(f"/threads/{thread_id}/runs/{run_id}/join") - def join_stream(self, thread_id: str, run_id: str) -> AsyncIterator[StreamPart]: + def join_stream( + self, thread_id: str, run_id: str, *, cancel_on_disconnect: bool = False + ) -> AsyncIterator[StreamPart]: """Stream output from a run in real-time, until the run is done. Output is not buffered, so any output produced before this call will not be received here. @@ -1804,6 +1806,7 @@ class RunsClient: Args: thread_id: The thread ID to join. run_id: The run ID to join. + cancel_on_disconnect: Whether to cancel the run when the stream is disconnected. Returns: None @@ -1816,7 +1819,11 @@ class RunsClient: ) """ # noqa: E501 - return self.http.stream(f"/threads/{thread_id}/runs/{run_id}/stream", "GET") + return self.http.stream( + f"/threads/{thread_id}/runs/{run_id}/stream", + "GET", + params={"cancel_on_disconnect": cancel_on_disconnect}, + ) async def delete(self, thread_id: str, run_id: str) -> None: """Delete a run. From 6a6c3ed84cfaaea6000ddcb1c6c9bc8267d15eea Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Tue, 3 Dec 2024 01:31:11 +0100 Subject: [PATCH 05/13] Bump to 0.0.30 --- libs/sdk-js/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/sdk-js/package.json b/libs/sdk-js/package.json index ded99a00b..3b373dabe 100644 --- a/libs/sdk-js/package.json +++ b/libs/sdk-js/package.json @@ -1,6 +1,6 @@ { "name": "@langchain/langgraph-sdk", - "version": "0.0.29", + "version": "0.0.30", "description": "Client library for interacting with the LangGraph API", "type": "module", "packageManager": "yarn@1.22.19", From a91bf116cbfae4aa37a6c69b3b4a66795485f043 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Tue, 3 Dec 2024 01:31:55 +0100 Subject: [PATCH 06/13] Bump to 0.1.41 --- libs/sdk-py/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/sdk-py/pyproject.toml b/libs/sdk-py/pyproject.toml index c7776a6cc..46464c6a6 100644 --- a/libs/sdk-py/pyproject.toml +++ b/libs/sdk-py/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-sdk" -version = "0.1.40" +version = "0.1.41" description = "SDK for interacting with LangGraph API" authors = [] license = "MIT" From efbd02a27da9abdb15d2c31aeb1c76ec2a430ccd Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Dec 2024 16:43:35 -0800 Subject: [PATCH 07/13] Implement support for interrupt/resume in subgraphs --- libs/langgraph/langgraph/pregel/algo.py | 4 +- libs/langgraph/tests/test_pregel.py | 31 +++- libs/langgraph/tests/test_pregel_async.py | 215 ++++++++++++++++++++++ 3 files changed, 241 insertions(+), 9 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 3c3008ce4..5d104a85f 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -595,7 +595,7 @@ def prepare_single_task( for tid, c, v in pending_writes if tid in (NULL_TASK_ID, task_id) and c == RESUME ), - MISSING, + configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING), ), }, ), @@ -720,7 +720,7 @@ def prepare_single_task( if tid in (NULL_TASK_ID, task_id) and c == RESUME ), - MISSING, + configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING), ), }, ), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 113a6c3c8..6ff70276f 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8741,14 +8741,12 @@ def test_dynamic_interrupt_subgraph( tool_two_node_count = 0 def tool_two_node(s: SubgraphState) -> SubgraphState: - print("tool_two_node", s) nonlocal tool_two_node_count tool_two_node_count += 1 if s["market"] == "DE": answer = interrupt("Just because...") else: answer = " all good" - print("after interrupt") return {"my_key": answer} subgraph = StateGraph(SubgraphState) @@ -8807,7 +8805,7 @@ def test_dynamic_interrupt_subgraph( ] # resume with answer assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ - {"tool_two": {"my_key": " my answer"}}, + {"tool_two": {"my_key": " my answer", "market": "DE"}}, ] # flow: interrupt -> clear tasks @@ -8817,7 +8815,12 @@ def test_dynamic_interrupt_subgraph( "my_key": "value ⛰️", "market": "DE", } - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + assert [ + c.metadata + for c in tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} + ) + ] == [ { "parents": {}, "source": "loop", @@ -8845,9 +8848,15 @@ def test_dynamic_interrupt_subgraph( Interrupt( value="Just because...", resumable=True, - ns=[AnyStr("tool_two:")], + ns=[AnyStr("tool_two:"), AnyStr("do:")], ), ), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("tool_two:"), + } + }, ), ), config=tool_two.checkpointer.get_tuple(thread1).config, @@ -8859,7 +8868,11 @@ def test_dynamic_interrupt_subgraph( "writes": None, "thread_id": "1", }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + parent_config=[ + *tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 + ) + ][-1].config, ) # clear the interrupt and next tasks tool_two.update_state(thread1, None, as_node=END) @@ -8877,7 +8890,11 @@ def test_dynamic_interrupt_subgraph( "writes": {}, "thread_id": "1", }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + parent_config=[ + *tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 + ) + ][-1].config, ) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 813fa8240..6c85e57ed 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -429,6 +429,189 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: ) +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: + class SubgraphState(TypedDict): + my_key: str + market: str + + tool_two_node_count = 0 + + def tool_two_node(s: SubgraphState) -> SubgraphState: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + subgraph = StateGraph(SubgraphState) + subgraph.add_node("do", tool_two_node, retry=RetryPolicy()) + subgraph.add_edge(START, "do") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", subgraph.compile()) + tool_two_graph.add_edge(START, "tool_two") + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert await tool_two.ainvoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value", + "market": "DE", + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value"} + + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good", + "market": "US", + } + + async with awith_checkpointer(checkpointer_name) as checkpointer: + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value ⛰️", "market": "DE"}, thread2 + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ) + }, + ] + # resume with answer + assert [ + c async for c in tool_two.astream(Command(resume=" my answer"), thread2) + ] == [ + {"tool_two": {"my_key": " my answer", "market": "DE"}}, + ] + + # flow: interrupt -> clear + thread1 = {"configurable": {"thread_id": "1"}} + thread1root = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value ⛰️", "market": "DE"}, thread1 + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ) + }, + ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1root)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("tool_two:"), + } + }, + ), + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1root, limit=2) + ][-1].config, + ) + + # clear the interrupt and next tasks + await tool_two.aupdate_state(thread1, None, as_node=END) + # interrupt is cleared, as well as the next tasks + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1root, limit=2) + ][-1].config, + ) + + @pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled") @pytest.mark.skipif( sys.version_info < (3, 11), @@ -12677,3 +12860,35 @@ async def test_parent_command(checkpointer_name: str) -> None: }, tasks=(), ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_interrupt_subgraph(checkpointer_name: str): + class State(TypedDict): + baz: str + + def foo(state): + return {"baz": "foo"} + + def bar(state): + value = interrupt("Please provide baz value:") + return {"baz": value} + + child_builder = StateGraph(State) + child_builder.add_node(bar) + child_builder.add_edge(START, "bar") + + builder = StateGraph(State) + builder.add_node(foo) + builder.add_node("bar", child_builder.compile()) + builder.add_edge(START, "foo") + builder.add_edge("foo", "bar") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert await graph.ainvoke({"baz": ""}, thread1) + # Resume with answer + assert await graph.ainvoke(Command(resume="bar"), thread1) From a3feaef2eb0a25ccdc0abbeedd8107c2b1b8e3b8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Dec 2024 16:45:07 -0800 Subject: [PATCH 08/13] lib: Handle Command in RemoteGraph --- libs/langgraph/langgraph/pregel/remote.py | 15 ++++++++++++++- libs/sdk-py/langgraph_sdk/client.py | 15 +++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index abe27eb28..8ee7ec884 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -1,3 +1,4 @@ +from dataclasses import asdict from typing import ( Any, AsyncIterator, @@ -41,7 +42,7 @@ from langgraph.constants import ( from langgraph.errors import GraphInterrupt from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode -from langgraph.types import Interrupt, StreamProtocol +from langgraph.types import Command, Interrupt, StreamProtocol from langgraph.utils.config import merge_configs @@ -597,11 +598,17 @@ class RemoteGraph(PregelProtocol): stream_modes, requested, req_single, stream = self._get_stream_modes( stream_mode, config ) + if isinstance(input, Command): + command: dict[str, Any] = asdict(input) + input = None + else: + command = None for chunk in sync_client.runs.stream( thread_id=sanitized_config["configurable"].get("thread_id"), assistant_id=self.name, input=input, + command=command, config=sanitized_config, stream_mode=stream_modes, interrupt_before=interrupt_before, @@ -680,11 +687,17 @@ class RemoteGraph(PregelProtocol): stream_modes, requested, req_single, stream = self._get_stream_modes( stream_mode, config ) + if isinstance(input, Command): + command: dict[str, Any] = asdict(input) + input = None + else: + command = None async for chunk in client.runs.stream( thread_id=sanitized_config["configurable"].get("thread_id"), assistant_id=self.name, input=input, + command=command, config=sanitized_config, stream_mode=stream_modes, interrupt_before=interrupt_before, diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index 81e8a8506..be909d336 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -3303,6 +3303,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3326,6 +3327,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3346,6 +3348,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3370,6 +3373,7 @@ class SyncRunsClient: assistant_id: The assistant ID or graph name to stream from. If using graph name, will default to first assistant created from that graph. input: The input to the graph. + command: The command to execute. stream_mode: The stream mode(s) to use. stream_subgraphs: Whether to stream output from subgraphs. metadata: Metadata to assign to the run. @@ -3420,6 +3424,7 @@ class SyncRunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "config": config, "metadata": metadata, "stream_mode": stream_mode, @@ -3453,6 +3458,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3472,6 +3478,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3492,6 +3499,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3514,6 +3522,7 @@ class SyncRunsClient: assistant_id: The assistant ID or graph name to stream from. If using graph name, will default to first assistant created from that graph. input: The input to the graph. + command: The command to execute. stream_mode: The stream mode(s) to use. stream_subgraphs: Whether to stream output from subgraphs. metadata: Metadata to assign to the run. @@ -3600,6 +3609,7 @@ class SyncRunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "stream_mode": stream_mode, "stream_subgraphs": stream_subgraphs, "config": config, @@ -3637,6 +3647,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, @@ -3657,6 +3668,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, @@ -3674,6 +3686,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, @@ -3695,6 +3708,7 @@ class SyncRunsClient: assistant_id: The assistant ID or graph name to run. If using graph name, will default to first assistant created from that graph. input: The input to the graph. + command: The command to execute. metadata: Metadata to assign to the run. config: The configuration for the assistant. checkpoint: The checkpoint to resume from. @@ -3761,6 +3775,7 @@ class SyncRunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "config": config, "metadata": metadata, "assistant_id": assistant_id, From d36e6ceaaf3c9c9d57ae751bf1ab9e965fec370e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Dec 2024 16:59:38 -0800 Subject: [PATCH 09/13] Fix --- libs/langgraph/tests/test_pregel_async.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 6c85e57ed..378ca1ac1 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -12862,6 +12862,10 @@ async def test_parent_command(checkpointer_name: str) -> None: ) +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_subgraph(checkpointer_name: str): class State(TypedDict): From 0071bd1e1cf91f84a9c364ee3eae8eacd9ad1001 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Dec 2024 17:01:12 -0800 Subject: [PATCH 10/13] Lint --- libs/langgraph/langgraph/pregel/remote.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 8ee7ec884..d45cdb310 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -28,6 +28,7 @@ from langgraph_sdk.client import ( get_sync_client, ) from langgraph_sdk.schema import Checkpoint, ThreadState +from langgraph_sdk.schema import Command as CommandSDK from langgraph_sdk.schema import StreamMode as StreamModeSDK from typing_extensions import Self @@ -599,7 +600,7 @@ class RemoteGraph(PregelProtocol): stream_mode, config ) if isinstance(input, Command): - command: dict[str, Any] = asdict(input) + command: Optional[CommandSDK] = cast(CommandSDK, asdict(input)) input = None else: command = None @@ -688,7 +689,7 @@ class RemoteGraph(PregelProtocol): stream_mode, config ) if isinstance(input, Command): - command: dict[str, Any] = asdict(input) + command: Optional[CommandSDK] = cast(CommandSDK, asdict(input)) input = None else: command = None From 20f091a27710f9c0b0a799786f73bc73b13daaa9 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Mon, 2 Dec 2024 17:08:24 -0800 Subject: [PATCH 11/13] [postgres] Sort Ascending (#2594) Adds a few of preliminaries: 1. Makes the returned "score" actually the result of the requested operation (cosine, inner_product, l2) 2. Sorts asc, etc. so that if you were to add an HNSW index (and not have any WHERE filters), it would be used 3. Drop the inner WHERE statement if no namespace or other filters are provided. See (2) for why. I don't yet add an index to the migrations since I think we need to agree on the right balance to ensure it's actually used in common query patterns. --- .../langgraph/store/postgres/base.py | 43 ++++++--- libs/checkpoint-postgres/tests/test_store.py | 87 +++++++++++++++++++ 2 files changed, 118 insertions(+), 12 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 2a908c90e..94c7d3e4e 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -321,7 +321,7 @@ class BasePostgresStore(Generic[C]): if op.query and self.index_config: embedding_requests.append((idx, op.query)) - score_operator = _get_distance_operator(self) + score_operator, post_operator = _get_distance_operator(self) vector_type = ( cast(PostgresIndexConfig, self.index_config) .get("ann_index_config", {}) @@ -351,18 +351,28 @@ class BasePostgresStore(Generic[C]): if not filter_conditions else " AND " + " AND ".join(filter_conditions) ) + if op.namespace_prefix: + prefix_filter_str = f"WHERE s.prefix LIKE %s {filter_str} " + ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",) + else: + ns_args = () + if filter_str: + prefix_filter_str = f"WHERE {filter_str} " + else: + prefix_filter_str = "" + base_query = f""" WITH scored AS ( - SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS score + SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS neg_score FROM store s JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key - WHERE s.prefix LIKE %s {filter_str} - ORDER BY {score_operator} DESC + {prefix_filter_str} + ORDER BY {score_operator} ASC LIMIT %s ) SELECT * FROM ( SELECT DISTINCT ON (prefix, key) - prefix, key, value, created_at, updated_at, score + prefix, key, value, created_at, updated_at, {post_operator} as score FROM scored ORDER BY prefix, key, score DESC ) AS unique_docs @@ -372,7 +382,7 @@ class BasePostgresStore(Generic[C]): """ params = [ _PLACEHOLDER, # Vector placeholder - f"{_namespace_to_text(op.namespace_prefix)}%", + *ns_args, *filter_params, _PLACEHOLDER, expanded_limit, @@ -702,7 +712,6 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): _paramslist[i] = embedding for (idx, _), (query, params) in zip(search_ops, queries): - # Execute the actual query cur.execute(query, params) rows = cast(list[Row], cur.fetchall()) results[idx] = [ @@ -915,7 +924,7 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]: return tuple(namespace.split(".")) -def _get_distance_operator(store: Any) -> str: +def _get_distance_operator(store: Any) -> tuple[str, str]: """Get the distance operator and score expression based on config.""" # Note: Today, we are not using ANN indices due to restrictions # on PGVector's support for mixing vector and non-vector filters @@ -936,12 +945,22 @@ def _get_distance_operator(store: Any) -> str: config = cast(PostgresIndexConfig, store.index_config) distance_type = config.get("distance_type", "cosine") + # Return the operator and the score expression + # The operator is used in the CTE and will be compatible with an ASCENDING ORDER + # sort clause. + # The score expression is used in the final query and will be compatible with + # a DESCENDING ORDER sort clause and the user's expectations of what the similarity score + # should be. if distance_type == "l2": - return "1 - (sv.embedding <-> %s::%s)" + # Final: "-(sv.embedding <-> %s::%s)" + # We return the "l2 similarity" so that the sorting order is the same + return "sv.embedding <-> %s::%s", "-scored.neg_score" elif distance_type == "inner_product": - return "-(sv.embedding <#> %s::%s)" - else: # cosine - return "1 - (sv.embedding <=> %s::%s)" + # Final: "-(sv.embedding <#> %s::%s)" + return "sv.embedding <#> %s::%s", "-(scored.neg_score)" + else: # cosine similarity + # Final: "1 - (sv.embedding <=> %s::%s)" + return "sv.embedding <=> %s::%s", "1 - scored.neg_score" def _ensure_index_config( diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index c9d220fe0..35dfa2150 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -634,6 +634,7 @@ def test_embed_with_path_operation_config( distance_type: str, ) -> None: """Test operation-level field configuration for vector search.""" + with _create_vector_store( vector_type, distance_type, @@ -695,3 +696,89 @@ def test_embed_with_path_operation_config( # assert len(results) == 3 # doc5_result = next(r for r in results if r.key == "doc5") # assert doc5_result.score is None + + +def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]: + """ + Compute cosine similarity between a vector X and a matrix Y. + Lazy import numpy for efficiency. + """ + + similarities = [] + for y in Y: + dot_product = sum(a * b for a, b in zip(X, y)) + norm1 = sum(a * a for a in X) ** 0.5 + norm2 = sum(a * a for a in y) ** 0.5 + similarity = dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0 + similarities.append(similarity) + + return similarities + + +def _inner_product(X: list[float], Y: list[list[float]]) -> list[float]: + """ + Compute inner product between a vector X and a matrix Y. + Lazy import numpy for efficiency. + """ + + similarities = [] + for y in Y: + similarity = sum(a * b for a, b in zip(X, y)) + similarities.append(similarity) + + return similarities + + +def _neg_l2_distance(X: list[float], Y: list[list[float]]) -> list[float]: + """ + Compute l2 distance between a vector X and a matrix Y. + Lazy import numpy for efficiency. + """ + + similarities = [] + for y in Y: + similarity = sum((a - b) ** 2 for a, b in zip(X, y)) ** 0.5 + similarities.append(-similarity) + + return similarities + + +@pytest.mark.parametrize( + "vector_type,distance_type", + [ + ("vector", "cosine"), + ("vector", "inner_product"), + ("halfvec", "l2"), + ], +) +@pytest.mark.parametrize("query", ["aaa", "bbb", "ccc", "abcd", "poisson"]) +def test_scores( + fake_embeddings: CharacterEmbeddings, + vector_type: str, + distance_type: str, + query: str, +) -> None: + """Test operation-level field configuration for vector search.""" + with _create_vector_store( + vector_type, + distance_type, + fake_embeddings, + text_fields=["key0"], + ) as store: + doc = { + "key0": "aaa", + } + store.put(("test",), "doc", doc, index=["key0", "key1"]) + + results = store.search((), query=query) + vec0 = fake_embeddings.embed_query(doc["key0"]) + vec1 = fake_embeddings.embed_query(query) + if distance_type == "cosine": + similarities = _cosine_similarity(vec1, [vec0]) + elif distance_type == "inner_product": + similarities = _inner_product(vec1, [vec0]) + elif distance_type == "l2": + similarities = _neg_l2_distance(vec1, [vec0]) + + assert len(results) == 1 + assert results[0].score == pytest.approx(similarities[0], abs=1e-3) From fe538d4bcbecf2eecb8a0cc2ace548908d24e608 Mon Sep 17 00:00:00 2001 From: Mingqi Hu Date: Tue, 3 Dec 2024 09:36:13 +0800 Subject: [PATCH 12/13] docs: Use edit mode as default to install template (#2590) Small change to install the dependencies with `edit` mode so that users or freshman can see the effect immediately when they change the template code. As below, `pip install -e .` It's very good to evaluate how agent works and easy to test & re-develop! --------- Signed-off-by: Mingqi Hu Co-authored-by: William FH <13333726+hinthornw@users.noreply.github.com> --- docs/docs/tutorials/langgraph-platform/local-server.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/docs/tutorials/langgraph-platform/local-server.md b/docs/docs/tutorials/langgraph-platform/local-server.md index 9c8e3ea9f..db7f53fea 100644 --- a/docs/docs/tutorials/langgraph-platform/local-server.md +++ b/docs/docs/tutorials/langgraph-platform/local-server.md @@ -35,10 +35,10 @@ Create a new app from the `react-agent` template. This template is a simple agen ## Install Dependencies -In the root of your new LangGraph app, install the dependencies: +In the root of your new LangGraph app, install the dependencies in `edit` mode so your local changes are used by the server: ```shell -pip install . +pip install -e . ``` ## Create a `.env` file From 15f0765d60f18a9c3a2015af3c6506232750643c Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Mon, 2 Dec 2024 17:42:29 -0800 Subject: [PATCH 13/13] Add IVFFlat and HNSW support (#2598) It seems that actually once i moved the operators & other things out, the query planner does do reasonable things and do sequential scanning if filtered N < some size but the index otherwise, even with namespace filtering. --- .../langgraph/store/postgres/base.py | 70 ++++++++++++++++++- 1 file changed, 67 insertions(+), 3 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 94c7d3e4e..28edd8998 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -56,6 +56,7 @@ class Migration(NamedTuple): sql: str params: Optional[dict[str, Any]] = None + condition: Optional[Callable[["BasePostgresStore"], bool]] = None MIGRATIONS: Sequence[str] = [ @@ -104,11 +105,29 @@ CREATE TABLE IF NOT EXISTS store_vectors ( ), }, ), - # TODO: Add an HNSW or IVFFlat index depending on config - # First must improve the search query when filtering by - # namespace + Migration( + """ +CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors + USING %(index_type)s (embedding %(ops)s)%(index_params)s; +""", + condition=lambda store: bool( + store.index_config and _get_index_params(store)[0] != "flat" + ), + params={ + "index_type": lambda store: _get_index_params(store)[0], + "ops": lambda store: _get_vector_type_ops(store), + "index_params": lambda store: ( + " WITH (" + + ", ".join(f"{k}={v}" for k, v in _get_index_params(store)[1].items()) + + ")" + if _get_index_params(store)[1] + else "" + ), + }, + ), ] + C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn]) @@ -140,6 +159,8 @@ class PoolConfig(TypedDict, total=False): class ANNIndexConfig(TypedDict, total=False): """Configuration for vector index in PostgreSQL store.""" + kind: Literal["hnsw", "ivfflat", "flat"] + """Type of index to use: 'hnsw' for Hierarchical Navigable Small World, or 'ivfflat' for Inverted File Flat.""" vector_type: Literal["vector", "halfvec"] """Type of vector storage to use. Options: @@ -148,6 +169,35 @@ class ANNIndexConfig(TypedDict, total=False): """ +class HNSWConfig(ANNIndexConfig, total=False): + """Configuration for HNSW (Hierarchical Navigable Small World) index.""" + + kind: Literal["hnsw"] # type: ignore[misc] + m: int + """Maximum number of connections per layer. Default is 16.""" + ef_construction: int + """Size of dynamic candidate list for index construction. Default is 64.""" + + +class IVFFlatConfig(ANNIndexConfig, total=False): + """IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff). + + Three keys to achieving good recall are: + 1. Create the index after the table has some data + 2. Choose an appropriate number of lists - a good place to start is rows / 1000 for up to 1M rows and sqrt(rows) for over 1M rows + 3. When querying, specify an appropriate number of probes (higher is better for recall, lower is better for speed) - a good place to start is sqrt(lists) + """ + + kind: Literal["ivfflat"] # type: ignore[misc] + nlist: int + """Number of inverted lists (clusters) for IVF index. + + Determines the number of clusters used in the index structure. + Higher values can improve search speed but increase index size and build time. + Typically set to the square root of the number of vectors in the index. + """ + + class PostgresIndexConfig(IndexConfig, total=False): """Configuration for vector embeddings in PostgreSQL store with pgvector-specific options. @@ -774,6 +824,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): for v, migration in enumerate( self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1 ): + if migration.condition and not migration.condition(self): + continue sql = migration.sql if migration.params: params = { @@ -832,6 +884,18 @@ def _get_vector_type_ops(store: BasePostgresStore) -> str: return f"{type_prefix}_{distance_suffix}" +def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]: + """Get the index type and configuration based on config.""" + if not store.index_config: + return "hnsw", {} + + config = cast(PostgresIndexConfig, store.index_config) + index_config = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy() + kind = index_config.pop("kind", "hnsw") + index_config.pop("vector_type", None) + return kind, index_config + + def _namespace_to_text( namespace: tuple[str, ...], handle_wildcards: bool = False ) -> str: