mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-11 10:45:18 +02:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
38fac0099c | ||
|
|
3ac3fb414e |
@@ -12,6 +12,10 @@ Generative user interfaces (Generative UI) allows agents to go beyond text and g
|
||||
|
||||
LangGraph Platform supports colocating your React components with your graph code. This allows you to focus on building specific UI components for your graph while easily plugging into existing chat interfaces such as [Agent Chat](https://agentchat.vercel.app) and loading the code only when actually needed.
|
||||
|
||||
!!! warning "LangGraph.js only"
|
||||
|
||||
Currently only LangGraph.js supports Generative UI. Support for Python is coming soon.
|
||||
|
||||
## Tutorial
|
||||
|
||||
### 1. Define and configure UI components
|
||||
@@ -70,105 +74,58 @@ CSS and Tailwind 4.x is also supported out of the box, so you can freely use Tai
|
||||
|
||||
### 2. Send the UI components in your graph
|
||||
|
||||
=== "Python"
|
||||
Use the `typedUi` utility to emit UI elements from your agent nodes:
|
||||
|
||||
```python title="src/agent.py"
|
||||
import uuid
|
||||
from typing import Annotated, Sequence, TypedDict
|
||||
```typescript title="src/agent/index.ts"
|
||||
import {
|
||||
typedUi,
|
||||
uiMessageReducer,
|
||||
} from "@langchain/langgraph-sdk/react-ui/server";
|
||||
|
||||
from langchain_core.messages import AIMessage, BaseMessage
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.graph.ui import AnyUIMessage, ui_message_reducer, push_ui_message
|
||||
import { ChatOpenAI } from "@langchain/openai";
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
import { z } from "zod";
|
||||
|
||||
import type ComponentMap from "./ui.js";
|
||||
|
||||
class AgentState(TypedDict): # noqa: D101
|
||||
messages: Annotated[Sequence[BaseMessage], add_messages]
|
||||
ui: Annotated[Sequence[AnyUIMessage], ui_message_reducer]
|
||||
import {
|
||||
Annotation,
|
||||
MessagesAnnotation,
|
||||
StateGraph,
|
||||
type LangGraphRunnableConfig,
|
||||
} from "@langchain/langgraph";
|
||||
|
||||
const AgentState = Annotation.Root({
|
||||
...MessagesAnnotation.spec,
|
||||
ui: Annotation({ reducer: uiMessageReducer, default: () => [] }),
|
||||
});
|
||||
|
||||
async def weather(state: AgentState):
|
||||
class WeatherOutput(TypedDict):
|
||||
city: str
|
||||
export const graph = new StateGraph(AgentState)
|
||||
.addNode("weather", async (state, config) => {
|
||||
// Provide the type of the component map to ensure
|
||||
// type safety of `ui.push()` calls as well as
|
||||
// pushing the messages to the `ui` and sending a custom event as well.
|
||||
const ui = typedUi<typeof ComponentMap>(config);
|
||||
|
||||
weather: WeatherOutput = (
|
||||
await ChatOpenAI(model="gpt-4o-mini")
|
||||
.with_structured_output(WeatherOutput)
|
||||
.with_config({"tags": ["nostream"]})
|
||||
.ainvoke(state["messages"])
|
||||
)
|
||||
const weather = await new ChatOpenAI({ model: "gpt-4o-mini" })
|
||||
.withStructuredOutput(z.object({ city: z.string() }))
|
||||
.withConfig({ tags: ["langsmith:nostream"] })
|
||||
.invoke(state.messages);
|
||||
|
||||
message = AIMessage(
|
||||
id=str(uuid.uuid4()),
|
||||
content=f"Here's the weather for {weather['city']}",
|
||||
)
|
||||
const response = {
|
||||
id: uuidv4(),
|
||||
type: "ai",
|
||||
content: `Here's the weather for ${weather.city}`,
|
||||
};
|
||||
|
||||
# Emit UI elements associated with the message
|
||||
push_ui_message("weather", weather, message=message)
|
||||
return {"messages": [message]}
|
||||
// Emit UI elements with associated AI message
|
||||
ui.push({ name: "weather", props: weather }, { message: response });
|
||||
|
||||
|
||||
workflow = StateGraph(AgentState)
|
||||
workflow.add_node(weather)
|
||||
workflow.add_edge("__start__", "weather")
|
||||
graph = workflow.compile()
|
||||
```
|
||||
|
||||
=== "JS"
|
||||
|
||||
Use the `typedUi` utility to emit UI elements from your agent nodes:
|
||||
|
||||
```typescript title="src/agent/index.ts"
|
||||
import {
|
||||
typedUi,
|
||||
uiMessageReducer,
|
||||
} from "@langchain/langgraph-sdk/react-ui/server";
|
||||
|
||||
import { ChatOpenAI } from "@langchain/openai";
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
import { z } from "zod";
|
||||
|
||||
import type ComponentMap from "./ui.js";
|
||||
|
||||
import {
|
||||
Annotation,
|
||||
MessagesAnnotation,
|
||||
StateGraph,
|
||||
type LangGraphRunnableConfig,
|
||||
} from "@langchain/langgraph";
|
||||
|
||||
const AgentState = Annotation.Root({
|
||||
...MessagesAnnotation.spec,
|
||||
ui: Annotation({ reducer: uiMessageReducer, default: () => [] }),
|
||||
});
|
||||
|
||||
export const graph = new StateGraph(AgentState)
|
||||
.addNode("weather", async (state, config) => {
|
||||
// Provide the type of the component map to ensure
|
||||
// type safety of `ui.push()` calls as well as
|
||||
// pushing the messages to the `ui` and sending a custom event as well.
|
||||
const ui = typedUi<typeof ComponentMap>(config);
|
||||
|
||||
const weather = await new ChatOpenAI({ model: "gpt-4o-mini" })
|
||||
.withStructuredOutput(z.object({ city: z.string() }))
|
||||
.withConfig({ tags: ["nostream"] })
|
||||
.invoke(state.messages);
|
||||
|
||||
const response = {
|
||||
id: uuidv4(),
|
||||
type: "ai",
|
||||
content: `Here's the weather for ${weather.city}`,
|
||||
};
|
||||
|
||||
// Emit UI elements associated with the AI message
|
||||
ui.push({ name: "weather", props: weather }, { message: response });
|
||||
|
||||
return { messages: [response] };
|
||||
})
|
||||
.addEdge("__start__", "weather")
|
||||
.compile();
|
||||
```
|
||||
return { messages: [response] };
|
||||
})
|
||||
.addEdge("__start__", "weather")
|
||||
.compile();
|
||||
```
|
||||
|
||||
### 3. Handle UI elements in your React application
|
||||
|
||||
@@ -337,29 +294,18 @@ const { thread, submit } = useStream({
|
||||
|
||||
### Remove UI messages from state
|
||||
|
||||
Similar to how messages can be removed from the state by appending a RemoveMessage you can remove an UI message from the state by calling `remove_ui_message` / `ui.delete` with the ID of the UI message.
|
||||
Similar to how messages can be removed from the state by appending a RemoveMessage you can remove an UI message from the state by calling `ui.delete` with the ID of the UI message.
|
||||
|
||||
=== "Python"
|
||||
```tsx
|
||||
// pushed message
|
||||
const message = ui.push({ name: "weather", props: { city: "London" } });
|
||||
|
||||
```python
|
||||
from langgraph.graph.ui import push_ui_message, delete_ui_message
|
||||
// remove said message
|
||||
ui.delete(message.id);
|
||||
|
||||
# push message
|
||||
message = push_ui_message("weather", {"city": "London"})
|
||||
|
||||
# remove said message
|
||||
delete_ui_message(message["id"])
|
||||
```
|
||||
|
||||
=== "JS"
|
||||
|
||||
```tsx
|
||||
// push message
|
||||
const message = ui.push({ name: "weather", props: { city: "London" } });
|
||||
|
||||
// remove said message
|
||||
ui.delete(message.id);
|
||||
```
|
||||
// return new state to persist changes
|
||||
return { ui: ui.items };
|
||||
```
|
||||
|
||||
## Learn more
|
||||
|
||||
|
||||
@@ -4,10 +4,6 @@ LangGraph has a built-in persistence layer, implemented through checkpointers. W
|
||||
|
||||

|
||||
|
||||
!!! info "LangGraph API handles checkpointing automatically"
|
||||
|
||||
When using the LangGraph API, you don't need to implement or configure checkpointers manually. The API handles all persistence infrastructure for you behind the scenes.
|
||||
|
||||
## Threads
|
||||
|
||||
A thread is a unique ID or [thread identifier](#threads) assigned to each checkpoint saved by a checkpointer. When invoking graph with a checkpointer, you **must** specify a `thread_id` as part of the `configurable` portion of the config:
|
||||
@@ -30,7 +26,7 @@ Let's see what checkpoints are saved when a simple graph is invoked as follows:
|
||||
|
||||
```python
|
||||
from langgraph.graph import StateGraph, START, END
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from typing import Annotated
|
||||
from typing_extensions import TypedDict
|
||||
from operator import add
|
||||
@@ -53,7 +49,7 @@ workflow.add_edge(START, "node_a")
|
||||
workflow.add_edge("node_a", "node_b")
|
||||
workflow.add_edge("node_b", END)
|
||||
|
||||
checkpointer = InMemorySaver()
|
||||
checkpointer = MemorySaver()
|
||||
graph = workflow.compile(checkpointer=checkpointer)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
@@ -227,10 +223,6 @@ But, what if we want to retain some information *across threads*? Consider the c
|
||||
|
||||
With checkpointers alone, we cannot share information across threads. This motivates the need for the [`Store`](../reference/store.md#langgraph.store.base.BaseStore) interface. As an illustration, we can define an `InMemoryStore` to store information about a user across threads. We simply compile our graph with a checkpointer, as before, and with our new `in_memory_store` variable.
|
||||
|
||||
!!! info "LangGraph API handles stores automatically"
|
||||
|
||||
When using the LangGraph API, you don't need to implement or configure stores manually. The API handles all storage infrastructure for you behind the scenes.
|
||||
|
||||
### Basic Usage
|
||||
|
||||
First, let's showcase this in isolation without using LangGraph.
|
||||
@@ -332,10 +324,10 @@ store.put(
|
||||
With this all in place, we use the `in_memory_store` in LangGraph. The `in_memory_store` works hand-in-hand with the checkpointer: the checkpointer saves state to threads, as discussed above, and the `in_memory_store` allows us to store arbitrary information for access *across* threads. We compile the graph with both the checkpointer and the `in_memory_store` as follows.
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
|
||||
# We need this because we want to enable threads (conversations)
|
||||
checkpointer = InMemorySaver()
|
||||
checkpointer = MemorySaver()
|
||||
|
||||
# ... Define the graph ...
|
||||
|
||||
@@ -448,7 +440,6 @@ Under the hood, checkpointing is powered by checkpointer objects that conform to
|
||||
* `langgraph-checkpoint-sqlite`: An implementation of LangGraph checkpointer that uses SQLite database ([SqliteSaver][langgraph.checkpoint.sqlite.SqliteSaver] / [AsyncSqliteSaver][langgraph.checkpoint.sqlite.aio.AsyncSqliteSaver]). Ideal for experimentation and local workflows. Needs to be installed separately.
|
||||
* `langgraph-checkpoint-postgres`: An advanced checkpointer that uses Postgres database ([PostgresSaver][langgraph.checkpoint.postgres.PostgresSaver] / [AsyncPostgresSaver][langgraph.checkpoint.postgres.aio.AsyncPostgresSaver]), used in LangGraph Cloud. Ideal for using in production. Needs to be installed separately.
|
||||
|
||||
|
||||
### Checkpointer interface
|
||||
|
||||
Each checkpointer conforms to [BaseCheckpointSaver][langgraph.checkpoint.base.BaseCheckpointSaver] interface and implements the following methods:
|
||||
@@ -461,7 +452,7 @@ Each checkpointer conforms to [BaseCheckpointSaver][langgraph.checkpoint.base.Ba
|
||||
If the checkpointer is used with asynchronous graph execution (i.e. executing the graph via `.ainvoke`, `.astream`, `.abatch`), asynchronous versions of the above methods will be used (`.aput`, `.aput_writes`, `.aget_tuple`, `.alist`).
|
||||
|
||||
!!! note Note
|
||||
For running your graph asynchronously, you can use `InMemorySaver`, or async versions of Sqlite/Postgres checkpointers -- `AsyncSqliteSaver` / `AsyncPostgresSaver` checkpointers.
|
||||
For running your graph asynchronously, you can use `MemorySaver`, or async versions of Sqlite/Postgres checkpointers -- `AsyncSqliteSaver` / `AsyncPostgresSaver` checkpointers.
|
||||
|
||||
### Serializer
|
||||
|
||||
|
||||
@@ -16,10 +16,6 @@
|
||||
" - [Memory](../../concepts/memory/)\n",
|
||||
" - [Chat Models](https://python.langchain.com/docs/concepts/chat_models/)\n",
|
||||
"\n",
|
||||
"!!! info \"Not needed for LangGraph API users\"\n",
|
||||
"\n",
|
||||
" If you're using the LangGraph API, you needn't manually implement a checkpointer. The API automatically handles checkpointing for you. This guide is relevant when implementing LangGraph in your own custom server.\n",
|
||||
"\n",
|
||||
"Many AI applications need memory to share context across multiple interactions on the same [thread](../../concepts/persistence#threads) (e.g., multiple turns of a conversation). In LangGraph functional API, this kind of memory can be added to any [entrypoint()][langgraph.func.entrypoint] workflow using [thread-level persistence](https://langchain-ai.github.io/langgraph/concepts/persistence).\n",
|
||||
"\n",
|
||||
"When creating a LangGraph workflow, you can set it up to persist its results by using a [checkpointer](https://langchain-ai.github.io/langgraph/reference/checkpoints/#basecheckpointsaver):\n",
|
||||
|
||||
@@ -31,10 +31,6 @@
|
||||
" </p>\n",
|
||||
"</div> \n",
|
||||
"\n",
|
||||
"!!! info \"Not needed for LangGraph API users\"\n",
|
||||
"\n",
|
||||
" If you're using the LangGraph API, you needn't manually implement a checkpointer. The API automatically handles checkpointing for you. This guide is relevant when implementing LangGraph in your own custom server.\n",
|
||||
"\n",
|
||||
"Many AI applications need memory to share context across multiple interactions. In LangGraph, this kind of memory can be added to any [StateGraph](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.StateGraph) using [thread-level persistence](https://langchain-ai.github.io/langgraph/concepts/persistence) .\n",
|
||||
"\n",
|
||||
"When creating any LangGraph graph, you can set it up to persist its state by adding a [checkpointer](https://langchain-ai.github.io/langgraph/reference/checkpoints/#basecheckpointsaver) when compiling the graph:\n",
|
||||
|
||||
@@ -26,10 +26,6 @@
|
||||
" </p>\n",
|
||||
"</div> \n",
|
||||
"\n",
|
||||
"!!! info \"Not needed for LangGraph API users\"\n",
|
||||
"\n",
|
||||
" If you're using the LangGraph API, you needn't manually implement a checkpointer. The API automatically handles checkpointing for you. This guide is relevant when implementing LangGraph in your own custom server.\n",
|
||||
"\n",
|
||||
"When creating LangGraph agents, you can also set them up so that they persist their state. This allows you to do things like interact with an agent multiple times and have it remember previous interactions.\n",
|
||||
"\n",
|
||||
"This how-to guide shows how to use `Postgres` as the backend for persisting checkpoint state using the [`langgraph-checkpoint-postgres`](https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres) library.\n",
|
||||
@@ -150,20 +146,20 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e9342c62-dbb4-40f6-9271-7393f1ca48c4",
|
||||
"id": "f54f9371",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Use sync connection\n",
|
||||
"## Use async connection\n",
|
||||
"\n",
|
||||
"This sets up a synchronous connection to the database. \n",
|
||||
"For most production server use cases, we recommend using the async connection to the database.\n",
|
||||
"\n",
|
||||
"Synchronous connections execute operations in a blocking manner, meaning each operation waits for completion before moving to the next one. The `DB_URI` is the database connection URI, with the protocol used for connecting to a PostgreSQL database, authentication, and host where database is running. The connection_kwargs dictionary defines additional parameters for the database connection."
|
||||
"Async connections allow non-blocking database operations. This means other parts of your application can continue running while waiting for database operations to complete. It's particularly useful in high-concurrency scenarios or when dealing with I/O-bound operations."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"id": "2b9d13b1-9d72-48a0-b63a-adc062c06c29",
|
||||
"execution_count": null,
|
||||
"id": "2d8d6c9a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -172,8 +168,8 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"id": "3fe36f67-073a-4fd7-a8f8-da196dd46a0d",
|
||||
"execution_count": null,
|
||||
"id": "351b4251",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -183,6 +179,238 @@
|
||||
"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "8e1dd27d",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### With a connection pool"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "a89a237a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from psycopg_pool import AsyncConnectionPool\n",
|
||||
"\n",
|
||||
"async with AsyncConnectionPool(\n",
|
||||
" # Example configuration\n",
|
||||
" conninfo=DB_URI,\n",
|
||||
" max_size=20,\n",
|
||||
" kwargs=connection_kwargs,\n",
|
||||
") as pool:\n",
|
||||
" checkpointer = AsyncPostgresSaver(pool)\n",
|
||||
"\n",
|
||||
" # NOTE: you need to call .setup() the first time you're using your checkpointer\n",
|
||||
" await checkpointer.setup()\n",
|
||||
"\n",
|
||||
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
||||
" config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
|
||||
" res = await graph.ainvoke(\n",
|
||||
" {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" checkpoint = await checkpointer.aget(config)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "2bb1b8fd",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'v': 1,\n",
|
||||
" 'id': '1ef559b7-5cc9-6460-8003-8655824c0944',\n",
|
||||
" 'ts': '2024-08-08T15:32:45.640793+00:00',\n",
|
||||
" 'current_tasks': {},\n",
|
||||
" 'pending_sends': [],\n",
|
||||
" 'versions_seen': {'agent': {'tools': '00000000000000000000000000000004.022986cd20ae85c77ea298a383f69ba8',\n",
|
||||
" 'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'},\n",
|
||||
" 'tools': {'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'},\n",
|
||||
" '__input__': {},\n",
|
||||
" '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}},\n",
|
||||
" 'channel_versions': {'agent': '00000000000000000000000000000005.065d90dd7f7cd091f0233855210bb2af',\n",
|
||||
" 'tools': '00000000000000000000000000000005.',\n",
|
||||
" 'messages': '00000000000000000000000000000005.d869fc7231619df0db74feed624efe41',\n",
|
||||
" '__start__': '00000000000000000000000000000002.',\n",
|
||||
" 'start:agent': '00000000000000000000000000000003.',\n",
|
||||
" 'branch:agent:should_continue:tools': '00000000000000000000000000000004.'},\n",
|
||||
" 'channel_values': {'agent': 'agent',\n",
|
||||
" 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='d883b8a0-99de-486d-91a2-bcfa7f25dc05'),\n",
|
||||
" AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_H6TAYfyd6AnaCrkQGs6Q2fVp', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-6f542f84-ad73-444c-8ef7-b5ea75a2e09b-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_H6TAYfyd6AnaCrkQGs6Q2fVp', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}),\n",
|
||||
" ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='c0e52254-77a4-4ea9-a2b7-61dd2d65ec68', tool_call_id='call_H6TAYfyd6AnaCrkQGs6Q2fVp'),\n",
|
||||
" AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-977140d4-7582-40c3-b2b6-31b542c430a3-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"checkpoint"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "a68094dc",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Note: If you are using this in an ASGI web framework or Starlette (or FastAPI), we'd recommend creating the connection pool within a [**lifespan event.**](https://www.starlette.io/lifespan/), similar to the pseudocode below:\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"import contextlib\n",
|
||||
"\n",
|
||||
"from starlette.applications import Starlette\n",
|
||||
"from starlette.requests import Request\n",
|
||||
"from starlette.responses import Response\n",
|
||||
"from starlette.routing import Route\n",
|
||||
"from psycopg_pool import AsyncConnectionPool\n",
|
||||
"from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver\n",
|
||||
"\n",
|
||||
"@contextlib.asynccontextmanager\n",
|
||||
"async def lifespan(app):\n",
|
||||
" async with AsyncConnectionPool(\n",
|
||||
" # Example configuration\n",
|
||||
" conninfo=DB_URI,\n",
|
||||
" max_size=20,\n",
|
||||
" kwargs=connection_kwargs,\n",
|
||||
" ) as pool:\n",
|
||||
" checkpointer = AsyncPostgresSaver(pool)\n",
|
||||
"\n",
|
||||
" # NOTE: you need to call .setup() the first time you're using your checkpointer\n",
|
||||
" await checkpointer.setup()\n",
|
||||
" yield {\"checkpointer\": checkpointer}\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
||||
"\n",
|
||||
"async def my_route(request: Request):\n",
|
||||
" checkpointer = request.state.checkpointer\n",
|
||||
" agent = graph.copy({\"checkpointer\": checkpointer})\n",
|
||||
" await agent.ainvoke(request)\n",
|
||||
" return Response(...)\n",
|
||||
"\n",
|
||||
"routes = [\n",
|
||||
" Route(\"/\", my_route),\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"app = Starlette(routes=routes, lifespan=lifespan)\n",
|
||||
"```"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6e53287e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### With a connection"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "7f0b0f76",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from psycopg import AsyncConnection\n",
|
||||
"\n",
|
||||
"async with await AsyncConnection.connect(DB_URI, **connection_kwargs) as conn:\n",
|
||||
" checkpointer = AsyncPostgresSaver(conn)\n",
|
||||
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
||||
" config = {\"configurable\": {\"thread_id\": \"5\"}}\n",
|
||||
" res = await graph.ainvoke(\n",
|
||||
" {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
|
||||
" )\n",
|
||||
" checkpoint_tuple = await checkpointer.aget_tuple(config)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "6a81c382",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"CheckpointTuple(config={'configurable': {'thread_id': '5', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-65b4-60ca-8003-1ef4b620559a'}}, checkpoint={'v': 1, 'id': '1ef559b7-65b4-60ca-8003-1ef4b620559a', 'ts': '2024-08-08T15:32:46.575814+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'agent': {'tools': '00000000000000000000000000000004.022986cd20ae85c77ea298a383f69ba8', 'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, 'tools': {'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'agent': '00000000000000000000000000000005.065d90dd7f7cd091f0233855210bb2af', 'tools': '00000000000000000000000000000005.', 'messages': '00000000000000000000000000000005.1557a6006d58f736d5cb2dd5c5f10111', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000003.', 'branch:agent:should_continue:tools': '00000000000000000000000000000004.'}, 'channel_values': {'agent': 'agent', 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='935e7732-b288-49bd-9ec2-1f7610cc38cb'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_94KtjtPmsiaj7T8yXvL7Ef31', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-790c929a-7982-49e7-af67-2cbe4a86373b-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_94KtjtPmsiaj7T8yXvL7Ef31', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='b2dc1073-abc4-4492-8982-434a7e32e445', tool_call_id='call_94KtjtPmsiaj7T8yXvL7Ef31'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-7e8a7f16-d8e1-457a-89f3-192102396449-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}, metadata={'step': 3, 'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'logprobs': None, 'model_name': 'gpt-4o-mini-2024-07-18', 'token_usage': {'total_tokens': 97, 'prompt_tokens': 88, 'completion_tokens': 9}, 'finish_reason': 'stop', 'system_fingerprint': 'fp_48196bc67a'}, id='run-7e8a7f16-d8e1-457a-89f3-192102396449-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '5', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-62ae-6128-8002-c04af82bcd41'}}, pending_writes=[])"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"checkpoint_tuple"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "26bd7fbb",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### With a connection string"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d93f68f7",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:\n",
|
||||
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
||||
" config = {\"configurable\": {\"thread_id\": \"6\"}}\n",
|
||||
" res = await graph.ainvoke(\n",
|
||||
" {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
|
||||
" )\n",
|
||||
" checkpoint_tuples = [c async for c in checkpointer.alist(config)]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "717a28ea",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-723c-67de-8003-63bd4eab35af'}}, checkpoint={'v': 1, 'id': '1ef559b7-723c-67de-8003-63bd4eab35af', 'ts': '2024-08-08T15:32:47.890003+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'agent': {'tools': '00000000000000000000000000000004.022986cd20ae85c77ea298a383f69ba8', 'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, 'tools': {'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'agent': '00000000000000000000000000000005.065d90dd7f7cd091f0233855210bb2af', 'tools': '00000000000000000000000000000005.', 'messages': '00000000000000000000000000000005.b6fe2a26011590cfe8fd6a39151a9e92', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000003.', 'branch:agent:should_continue:tools': '00000000000000000000000000000004.'}, 'channel_values': {'agent': 'agent', 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='977ddb90-9991-44cb-9f73-361c6dd21396'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-47b10c48-4db3-46d8-b4fa-e021818e01c5-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='798c520f-4f9a-4f6d-a389-da721eb4d4ce', tool_call_id='call_QIFCuh4zfP9owpjToycJiZf7'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-4a34e05d-8bcf-41ad-adc3-715919fde64c-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}, metadata={'step': 3, 'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'logprobs': None, 'model_name': 'gpt-4o-mini-2024-07-18', 'token_usage': {'total_tokens': 97, 'prompt_tokens': 88, 'completion_tokens': 9}, 'finish_reason': 'stop', 'system_fingerprint': 'fp_48196bc67a'}, id='run-4a34e05d-8bcf-41ad-adc3-715919fde64c-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6bf5-63c6-8002-ed990dbbc96e'}}, pending_writes=None),\n",
|
||||
" CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6bf5-63c6-8002-ed990dbbc96e'}}, checkpoint={'v': 1, 'id': '1ef559b7-6bf5-63c6-8002-ed990dbbc96e', 'ts': '2024-08-08T15:32:47.231667+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'agent': {'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, 'tools': {'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'agent': '00000000000000000000000000000004.', 'tools': '00000000000000000000000000000004.022986cd20ae85c77ea298a383f69ba8', 'messages': '00000000000000000000000000000004.c9074f2a41f05486b5efb86353dc75c0', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000003.', 'branch:agent:should_continue:tools': '00000000000000000000000000000004.'}, 'channel_values': {'tools': 'tools', 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='977ddb90-9991-44cb-9f73-361c6dd21396'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-47b10c48-4db3-46d8-b4fa-e021818e01c5-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='798c520f-4f9a-4f6d-a389-da721eb4d4ce', tool_call_id='call_QIFCuh4zfP9owpjToycJiZf7')]}}, metadata={'step': 2, 'source': 'loop', 'writes': {'tools': {'messages': [ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='798c520f-4f9a-4f6d-a389-da721eb4d4ce', tool_call_id='call_QIFCuh4zfP9owpjToycJiZf7')]}}}, parent_config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6be0-6926-8001-1a8ce73baf9e'}}, pending_writes=None),\n",
|
||||
" CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6be0-6926-8001-1a8ce73baf9e'}}, checkpoint={'v': 1, 'id': '1ef559b7-6be0-6926-8001-1a8ce73baf9e', 'ts': '2024-08-08T15:32:47.223198+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'agent': {'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'agent': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af', 'messages': '00000000000000000000000000000003.097b5407d709b297591f1ef5d50c8368', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000003.', 'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'}, 'channel_values': {'agent': 'agent', 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='977ddb90-9991-44cb-9f73-361c6dd21396'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-47b10c48-4db3-46d8-b4fa-e021818e01c5-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})], 'branch:agent:should_continue:tools': 'agent'}}, metadata={'step': 1, 'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'function', 'function': {'name': 'get_weather', 'arguments': '{\"city\":\"nyc\"}'}}]}, response_metadata={'logprobs': None, 'model_name': 'gpt-4o-mini-2024-07-18', 'token_usage': {'total_tokens': 73, 'prompt_tokens': 58, 'completion_tokens': 15}, 'finish_reason': 'tool_calls', 'system_fingerprint': 'fp_48196bc67a'}, id='run-47b10c48-4db3-46d8-b4fa-e021818e01c5-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})]}}}, parent_config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-663d-60b4-8000-10a8922bffbf'}}, pending_writes=None),\n",
|
||||
" CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-663d-60b4-8000-10a8922bffbf'}}, checkpoint={'v': 1, 'id': '1ef559b7-663d-60b4-8000-10a8922bffbf', 'ts': '2024-08-08T15:32:46.631935+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'messages': '00000000000000000000000000000002.2a79db8da664e437bdb25ea804457ca7', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='977ddb90-9991-44cb-9f73-361c6dd21396')], 'start:agent': '__start__'}}, metadata={'step': 0, 'source': 'loop', 'writes': None}, parent_config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6637-6d4e-bfff-6cecf690c3cb'}}, pending_writes=None),\n",
|
||||
" CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6637-6d4e-bfff-6cecf690c3cb'}}, checkpoint={'v': 1, 'id': '1ef559b7-6637-6d4e-bfff-6cecf690c3cb', 'ts': '2024-08-08T15:32:46.629806+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'__input__': {}}, 'channel_versions': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}, 'channel_values': {'__start__': {'messages': [['human', \"what's the weather in nyc\"]]}}}, metadata={'step': -1, 'source': 'input', 'writes': {'messages': [['human', \"what's the weather in nyc\"]]}}, parent_config=None, pending_writes=None)]"
|
||||
]
|
||||
},
|
||||
"metadata": {},
|
||||
"output_type": "display_data"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"checkpoint_tuples"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e9342c62-dbb4-40f6-9271-7393f1ca48c4",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Use sync connection\n",
|
||||
"\n",
|
||||
"This sets up a synchronous connection to the database. \n",
|
||||
"\n",
|
||||
"Synchronous connections execute operations in a blocking manner, meaning each operation waits for completion before moving to the next one. The `DB_URI` is the database connection URI, with the protocol used for connecting to a PostgreSQL database, authentication, and host where database is running. The connection_kwargs dictionary defines additional parameters for the database connection."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e39fc712-9e1c-4831-9077-dd07b0c13594",
|
||||
@@ -391,193 +619,6 @@
|
||||
"source": [
|
||||
"checkpoint_tuples"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c0a47d3e-e588-48fc-a5d4-2145dff17e77",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Use async connection\n",
|
||||
"\n",
|
||||
"This sets up an asynchronous connection to the database. \n",
|
||||
"\n",
|
||||
"Async connections allow non-blocking database operations. This means other parts of your application can continue running while waiting for database operations to complete. It's particularly useful in high-concurrency scenarios or when dealing with I/O-bound operations."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ee6b6cf7-d8f7-4777-a48d-93b5855fe681",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### With a connection pool"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"id": "4faf6087-73cc-4957-9a4f-f3509a32a740",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from psycopg_pool import AsyncConnectionPool\n",
|
||||
"\n",
|
||||
"async with AsyncConnectionPool(\n",
|
||||
" # Example configuration\n",
|
||||
" conninfo=DB_URI,\n",
|
||||
" max_size=20,\n",
|
||||
" kwargs=connection_kwargs,\n",
|
||||
") as pool:\n",
|
||||
" checkpointer = AsyncPostgresSaver(pool)\n",
|
||||
"\n",
|
||||
" # NOTE: you need to call .setup() the first time you're using your checkpointer\n",
|
||||
" await checkpointer.setup()\n",
|
||||
"\n",
|
||||
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
||||
" config = {\"configurable\": {\"thread_id\": \"4\"}}\n",
|
||||
" res = await graph.ainvoke(\n",
|
||||
" {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" checkpoint = await checkpointer.aget(config)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 14,
|
||||
"id": "e0c42044-4de6-4742-8e00-fe295d50c95a",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"{'v': 1,\n",
|
||||
" 'id': '1ef559b7-5cc9-6460-8003-8655824c0944',\n",
|
||||
" 'ts': '2024-08-08T15:32:45.640793+00:00',\n",
|
||||
" 'current_tasks': {},\n",
|
||||
" 'pending_sends': [],\n",
|
||||
" 'versions_seen': {'agent': {'tools': '00000000000000000000000000000004.022986cd20ae85c77ea298a383f69ba8',\n",
|
||||
" 'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'},\n",
|
||||
" 'tools': {'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'},\n",
|
||||
" '__input__': {},\n",
|
||||
" '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}},\n",
|
||||
" 'channel_versions': {'agent': '00000000000000000000000000000005.065d90dd7f7cd091f0233855210bb2af',\n",
|
||||
" 'tools': '00000000000000000000000000000005.',\n",
|
||||
" 'messages': '00000000000000000000000000000005.d869fc7231619df0db74feed624efe41',\n",
|
||||
" '__start__': '00000000000000000000000000000002.',\n",
|
||||
" 'start:agent': '00000000000000000000000000000003.',\n",
|
||||
" 'branch:agent:should_continue:tools': '00000000000000000000000000000004.'},\n",
|
||||
" 'channel_values': {'agent': 'agent',\n",
|
||||
" 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='d883b8a0-99de-486d-91a2-bcfa7f25dc05'),\n",
|
||||
" AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_H6TAYfyd6AnaCrkQGs6Q2fVp', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-6f542f84-ad73-444c-8ef7-b5ea75a2e09b-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_H6TAYfyd6AnaCrkQGs6Q2fVp', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}),\n",
|
||||
" ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='c0e52254-77a4-4ea9-a2b7-61dd2d65ec68', tool_call_id='call_H6TAYfyd6AnaCrkQGs6Q2fVp'),\n",
|
||||
" AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-977140d4-7582-40c3-b2b6-31b542c430a3-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}"
|
||||
]
|
||||
},
|
||||
"execution_count": 14,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"checkpoint"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "56552584-9eb8-40df-a6a0-44151018b509",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### With a connection"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 15,
|
||||
"id": "386b78bc-2f73-49ba-a2a4-47bce6fc49b7",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from psycopg import AsyncConnection\n",
|
||||
"\n",
|
||||
"async with await AsyncConnection.connect(DB_URI, **connection_kwargs) as conn:\n",
|
||||
" checkpointer = AsyncPostgresSaver(conn)\n",
|
||||
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
||||
" config = {\"configurable\": {\"thread_id\": \"5\"}}\n",
|
||||
" res = await graph.ainvoke(\n",
|
||||
" {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
|
||||
" )\n",
|
||||
" checkpoint_tuple = await checkpointer.aget_tuple(config)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"id": "d1ed1344-c923-4a46-b04e-cc3646737d48",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"CheckpointTuple(config={'configurable': {'thread_id': '5', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-65b4-60ca-8003-1ef4b620559a'}}, checkpoint={'v': 1, 'id': '1ef559b7-65b4-60ca-8003-1ef4b620559a', 'ts': '2024-08-08T15:32:46.575814+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'agent': {'tools': '00000000000000000000000000000004.022986cd20ae85c77ea298a383f69ba8', 'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, 'tools': {'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'agent': '00000000000000000000000000000005.065d90dd7f7cd091f0233855210bb2af', 'tools': '00000000000000000000000000000005.', 'messages': '00000000000000000000000000000005.1557a6006d58f736d5cb2dd5c5f10111', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000003.', 'branch:agent:should_continue:tools': '00000000000000000000000000000004.'}, 'channel_values': {'agent': 'agent', 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='935e7732-b288-49bd-9ec2-1f7610cc38cb'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_94KtjtPmsiaj7T8yXvL7Ef31', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-790c929a-7982-49e7-af67-2cbe4a86373b-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_94KtjtPmsiaj7T8yXvL7Ef31', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='b2dc1073-abc4-4492-8982-434a7e32e445', tool_call_id='call_94KtjtPmsiaj7T8yXvL7Ef31'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-7e8a7f16-d8e1-457a-89f3-192102396449-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}, metadata={'step': 3, 'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'logprobs': None, 'model_name': 'gpt-4o-mini-2024-07-18', 'token_usage': {'total_tokens': 97, 'prompt_tokens': 88, 'completion_tokens': 9}, 'finish_reason': 'stop', 'system_fingerprint': 'fp_48196bc67a'}, id='run-7e8a7f16-d8e1-457a-89f3-192102396449-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '5', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-62ae-6128-8002-c04af82bcd41'}}, pending_writes=[])"
|
||||
]
|
||||
},
|
||||
"execution_count": 16,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"checkpoint_tuple"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "2f7e486a-3e63-41d7-b84b-6743f0a5764c",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### With a connection string"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 17,
|
||||
"id": "6a39d1ff-ca37-4457-8b52-07d33b59c36e",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:\n",
|
||||
" graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
|
||||
" config = {\"configurable\": {\"thread_id\": \"6\"}}\n",
|
||||
" res = await graph.ainvoke(\n",
|
||||
" {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n",
|
||||
" )\n",
|
||||
" checkpoint_tuples = [c async for c in checkpointer.alist(config)]"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 18,
|
||||
"id": "2b6d73ca-519e-45f7-90c2-1b8596624505",
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-723c-67de-8003-63bd4eab35af'}}, checkpoint={'v': 1, 'id': '1ef559b7-723c-67de-8003-63bd4eab35af', 'ts': '2024-08-08T15:32:47.890003+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'agent': {'tools': '00000000000000000000000000000004.022986cd20ae85c77ea298a383f69ba8', 'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, 'tools': {'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'agent': '00000000000000000000000000000005.065d90dd7f7cd091f0233855210bb2af', 'tools': '00000000000000000000000000000005.', 'messages': '00000000000000000000000000000005.b6fe2a26011590cfe8fd6a39151a9e92', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000003.', 'branch:agent:should_continue:tools': '00000000000000000000000000000004.'}, 'channel_values': {'agent': 'agent', 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='977ddb90-9991-44cb-9f73-361c6dd21396'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-47b10c48-4db3-46d8-b4fa-e021818e01c5-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='798c520f-4f9a-4f6d-a389-da721eb4d4ce', tool_call_id='call_QIFCuh4zfP9owpjToycJiZf7'), AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 88, 'total_tokens': 97}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'stop', 'logprobs': None}, id='run-4a34e05d-8bcf-41ad-adc3-715919fde64c-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}, metadata={'step': 3, 'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='The weather in NYC might be cloudy.', response_metadata={'logprobs': None, 'model_name': 'gpt-4o-mini-2024-07-18', 'token_usage': {'total_tokens': 97, 'prompt_tokens': 88, 'completion_tokens': 9}, 'finish_reason': 'stop', 'system_fingerprint': 'fp_48196bc67a'}, id='run-4a34e05d-8bcf-41ad-adc3-715919fde64c-0', usage_metadata={'input_tokens': 88, 'output_tokens': 9, 'total_tokens': 97})]}}}, parent_config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6bf5-63c6-8002-ed990dbbc96e'}}, pending_writes=None),\n",
|
||||
" CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6bf5-63c6-8002-ed990dbbc96e'}}, checkpoint={'v': 1, 'id': '1ef559b7-6bf5-63c6-8002-ed990dbbc96e', 'ts': '2024-08-08T15:32:47.231667+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'agent': {'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, 'tools': {'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'agent': '00000000000000000000000000000004.', 'tools': '00000000000000000000000000000004.022986cd20ae85c77ea298a383f69ba8', 'messages': '00000000000000000000000000000004.c9074f2a41f05486b5efb86353dc75c0', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000003.', 'branch:agent:should_continue:tools': '00000000000000000000000000000004.'}, 'channel_values': {'tools': 'tools', 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='977ddb90-9991-44cb-9f73-361c6dd21396'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-47b10c48-4db3-46d8-b4fa-e021818e01c5-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73}), ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='798c520f-4f9a-4f6d-a389-da721eb4d4ce', tool_call_id='call_QIFCuh4zfP9owpjToycJiZf7')]}}, metadata={'step': 2, 'source': 'loop', 'writes': {'tools': {'messages': [ToolMessage(content='It might be cloudy in nyc', name='get_weather', id='798c520f-4f9a-4f6d-a389-da721eb4d4ce', tool_call_id='call_QIFCuh4zfP9owpjToycJiZf7')]}}}, parent_config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6be0-6926-8001-1a8ce73baf9e'}}, pending_writes=None),\n",
|
||||
" CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6be0-6926-8001-1a8ce73baf9e'}}, checkpoint={'v': 1, 'id': '1ef559b7-6be0-6926-8001-1a8ce73baf9e', 'ts': '2024-08-08T15:32:47.223198+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'agent': {'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'agent': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af', 'messages': '00000000000000000000000000000003.097b5407d709b297591f1ef5d50c8368', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000003.', 'branch:agent:should_continue:tools': '00000000000000000000000000000003.065d90dd7f7cd091f0233855210bb2af'}, 'channel_values': {'agent': 'agent', 'messages': [HumanMessage(content=\"what's the weather in nyc\", id='977ddb90-9991-44cb-9f73-361c6dd21396'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'function': {'arguments': '{\"city\":\"nyc\"}', 'name': 'get_weather'}, 'type': 'function'}]}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 58, 'total_tokens': 73}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_48196bc67a', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-47b10c48-4db3-46d8-b4fa-e021818e01c5-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})], 'branch:agent:should_continue:tools': 'agent'}}, metadata={'step': 1, 'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'function', 'function': {'name': 'get_weather', 'arguments': '{\"city\":\"nyc\"}'}}]}, response_metadata={'logprobs': None, 'model_name': 'gpt-4o-mini-2024-07-18', 'token_usage': {'total_tokens': 73, 'prompt_tokens': 58, 'completion_tokens': 15}, 'finish_reason': 'tool_calls', 'system_fingerprint': 'fp_48196bc67a'}, id='run-47b10c48-4db3-46d8-b4fa-e021818e01c5-0', tool_calls=[{'name': 'get_weather', 'args': {'city': 'nyc'}, 'id': 'call_QIFCuh4zfP9owpjToycJiZf7', 'type': 'tool_call'}], usage_metadata={'input_tokens': 58, 'output_tokens': 15, 'total_tokens': 73})]}}}, parent_config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-663d-60b4-8000-10a8922bffbf'}}, pending_writes=None),\n",
|
||||
" CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-663d-60b4-8000-10a8922bffbf'}}, checkpoint={'v': 1, 'id': '1ef559b7-663d-60b4-8000-10a8922bffbf', 'ts': '2024-08-08T15:32:46.631935+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}}, 'channel_versions': {'messages': '00000000000000000000000000000002.2a79db8da664e437bdb25ea804457ca7', '__start__': '00000000000000000000000000000002.', 'start:agent': '00000000000000000000000000000002.d6f25946c3108fc12f27abbcf9b4cedc'}, 'channel_values': {'messages': [HumanMessage(content=\"what's the weather in nyc\", id='977ddb90-9991-44cb-9f73-361c6dd21396')], 'start:agent': '__start__'}}, metadata={'step': 0, 'source': 'loop', 'writes': None}, parent_config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6637-6d4e-bfff-6cecf690c3cb'}}, pending_writes=None),\n",
|
||||
" CheckpointTuple(config={'configurable': {'thread_id': '6', 'checkpoint_ns': '', 'checkpoint_id': '1ef559b7-6637-6d4e-bfff-6cecf690c3cb'}}, checkpoint={'v': 1, 'id': '1ef559b7-6637-6d4e-bfff-6cecf690c3cb', 'ts': '2024-08-08T15:32:46.629806+00:00', 'current_tasks': {}, 'pending_sends': [], 'versions_seen': {'__input__': {}}, 'channel_versions': {'__start__': '00000000000000000000000000000001.0e148ae3debe753278387e84f786e863'}, 'channel_values': {'__start__': {'messages': [['human', \"what's the weather in nyc\"]]}}}, metadata={'step': -1, 'source': 'input', 'writes': {'messages': [['human', \"what's the weather in nyc\"]]}}, parent_config=None, pending_writes=None)]"
|
||||
]
|
||||
},
|
||||
"execution_count": 18,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"checkpoint_tuples"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
{% extends "base.html" %}
|
||||
|
||||
{% block extrahead %}
|
||||
<meta name="algolia-site-verification" content="165B7E7C89E49946" />
|
||||
<style>
|
||||
@import url("https://fonts.googleapis.com/css2?family=Public+Sans&display=swap");
|
||||
:root {
|
||||
|
||||
@@ -38,8 +38,6 @@ class InMemorySaver(
|
||||
Only use `InMemorySaver` for debugging or testing purposes.
|
||||
For production use cases we recommend installing [langgraph-checkpoint-postgres](https://pypi.org/project/langgraph-checkpoint-postgres/) and using `PostgresSaver` / `AsyncPostgresSaver`.
|
||||
|
||||
If you are using the LangGraph Platform, no checkpointer needs to be specified. The correct managed checkpointer will be used automatically.
|
||||
|
||||
Args:
|
||||
serde (Optional[SerializerProtocol]): The serializer to use for serializing and deserializing checkpoints. Defaults to None.
|
||||
|
||||
|
||||
@@ -778,22 +778,7 @@ def _update_graph_paths(
|
||||
FileNotFoundError: If the local file (module) does not actually exist on disk.
|
||||
IsADirectoryError: If `module_str` points to a directory instead of a file.
|
||||
"""
|
||||
for graph_id, data in config["graphs"].items():
|
||||
if isinstance(data, dict):
|
||||
# Then we're looking for a 'path' key
|
||||
if "path" not in data:
|
||||
raise ValueError(
|
||||
f"Graph '{graph_id}' must contain a 'path' key if "
|
||||
f" it is a dictionary."
|
||||
)
|
||||
import_str = data["path"]
|
||||
elif isinstance(data, str):
|
||||
import_str = data
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Graph '{graph_id}' must be a string or a dictionary with a 'path' key."
|
||||
)
|
||||
|
||||
for graph_id, import_str in config["graphs"].items():
|
||||
module_str, _, attr_str = import_str.partition(":")
|
||||
if not module_str or not attr_str:
|
||||
message = (
|
||||
@@ -833,10 +818,7 @@ def _update_graph_paths(
|
||||
"Add its containing package to 'dependencies' list."
|
||||
)
|
||||
# update the config
|
||||
if isinstance(data, dict):
|
||||
config["graphs"][graph_id]["path"] = f"{module_str}:{attr_str}"
|
||||
else:
|
||||
config["graphs"][graph_id] = f"{module_str}:{attr_str}"
|
||||
config["graphs"][graph_id] = f"{module_str}:{attr_str}"
|
||||
|
||||
|
||||
def _update_auth_path(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-cli"
|
||||
version = "0.1.89"
|
||||
version = "0.1.84"
|
||||
description = "CLI for interacting with LangGraph API"
|
||||
authors = []
|
||||
license = "MIT"
|
||||
|
||||
@@ -196,50 +196,6 @@ def test_dockerfile_command_basic() -> None:
|
||||
assert save_path.exists()
|
||||
|
||||
|
||||
def test_dockerfile_command_new_style_config() -> None:
|
||||
"""Test `dockerfile` command with a new style config.
|
||||
|
||||
This config format allows specifying agent data as a dictionary.
|
||||
{
|
||||
"graphs": {
|
||||
"agent1": {
|
||||
"path": ... # path to graph definition,
|
||||
... # other fields
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
runner = CliRunner()
|
||||
config_content = {
|
||||
"dependencies": ["./my_agent"],
|
||||
"graphs": {
|
||||
"agent": {
|
||||
"path": "./my_agent/agent.py:graph",
|
||||
"description": "This is a test agent",
|
||||
}
|
||||
},
|
||||
"env": ".env",
|
||||
}
|
||||
with temporary_config_folder(config_content) as temp_dir:
|
||||
save_path = temp_dir / "Dockerfile"
|
||||
# Add agent.py file
|
||||
agent_path = temp_dir / "my_agent" / "agent.py"
|
||||
agent_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
agent_path.touch()
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
["dockerfile", str(save_path), "--config", str(temp_dir / "config.json")],
|
||||
)
|
||||
|
||||
# Assert command was successful
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "✅ Created: Dockerfile" in result.output
|
||||
|
||||
# Check if Dockerfile was created
|
||||
assert save_path.exists()
|
||||
|
||||
|
||||
def test_dockerfile_command_with_docker_compose() -> None:
|
||||
"""Test the 'dockerfile' command with Docker Compose configuration."""
|
||||
runner = CliRunner()
|
||||
|
||||
@@ -9,7 +9,6 @@ from bench.fanout_to_subgraph import fanout_to_subgraph, fanout_to_subgraph_sync
|
||||
from bench.pydantic_state import pydantic_state
|
||||
from bench.react_agent import react_agent
|
||||
from bench.sequential import create_sequential
|
||||
from bench.wide_dict import wide_dict
|
||||
from bench.wide_state import wide_state
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from langgraph.graph import StateGraph
|
||||
@@ -26,7 +25,6 @@ async def arun(graph: Pregel, input: dict):
|
||||
"configurable": {"thread_id": str(uuid4())},
|
||||
"recursion_limit": 1000000000,
|
||||
},
|
||||
checkpoint_during=False,
|
||||
)
|
||||
]
|
||||
)
|
||||
@@ -43,7 +41,6 @@ async def arun_first_event_latency(graph: Pregel, input: dict) -> None:
|
||||
"configurable": {"thread_id": str(uuid4())},
|
||||
"recursion_limit": 1000000000,
|
||||
},
|
||||
checkpoint_during=False,
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -63,7 +60,6 @@ def run(graph: Pregel, input: dict):
|
||||
"configurable": {"thread_id": str(uuid4())},
|
||||
"recursion_limit": 1000000000,
|
||||
},
|
||||
checkpoint_during=False,
|
||||
)
|
||||
]
|
||||
)
|
||||
@@ -80,7 +76,6 @@ def run_first_event_latency(graph: Pregel, input: dict) -> None:
|
||||
"configurable": {"thread_id": str(uuid4())},
|
||||
"recursion_limit": 1000000000,
|
||||
},
|
||||
checkpoint_during=False,
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -256,102 +251,6 @@ benchmarks = (
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_dict_25x300",
|
||||
wide_dict(300).compile(checkpointer=None),
|
||||
wide_dict(300).compile(checkpointer=None),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(5)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_dict_25x300_checkpoint",
|
||||
wide_dict(300).compile(checkpointer=MemorySaver()),
|
||||
wide_dict(300).compile(checkpointer=MemorySaver()),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(5)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_dict_15x600",
|
||||
wide_dict(600).compile(checkpointer=None),
|
||||
wide_dict(600).compile(checkpointer=None),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(3)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_dict_15x600_checkpoint",
|
||||
wide_dict(600).compile(checkpointer=MemorySaver()),
|
||||
wide_dict(600).compile(checkpointer=MemorySaver()),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(5)
|
||||
}
|
||||
for i in range(3)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_dict_9x1200",
|
||||
wide_dict(1200).compile(checkpointer=None),
|
||||
wide_dict(1200).compile(checkpointer=None),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(3)
|
||||
}
|
||||
for i in range(3)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"wide_dict_9x1200_checkpoint",
|
||||
wide_dict(1200).compile(checkpointer=MemorySaver()),
|
||||
wide_dict(1200).compile(checkpointer=MemorySaver()),
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(3)
|
||||
}
|
||||
for i in range(3)
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
"sequential_10",
|
||||
create_sequential(10).compile(),
|
||||
|
||||
@@ -1,153 +0,0 @@
|
||||
import operator
|
||||
from functools import partial
|
||||
from random import choice
|
||||
from typing import Annotated, Optional, Sequence
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph.state import StateGraph
|
||||
|
||||
|
||||
def wide_dict(n: int) -> StateGraph:
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, operator.add]
|
||||
trigger_events: Annotated[list, operator.add]
|
||||
"""The external events that are converted by the graph."""
|
||||
primary_issue_medium: Annotated[str, lambda x, y: y or x]
|
||||
autoresponse: Annotated[Optional[dict], lambda _, y: y] # Always overwrite
|
||||
issue: Annotated[dict | None, lambda x, y: y if y else x]
|
||||
relevant_rules: Optional[list[dict]]
|
||||
"""SOPs fetched from the rulebook that are relevant to the current conversation."""
|
||||
memory_docs: Optional[list[dict]]
|
||||
"""Memory docs fetched from the memory service that are relevant to the current conversation."""
|
||||
categorizations: Annotated[list[dict], operator.add]
|
||||
"""The issue categorizations auto-generated by the AI."""
|
||||
responses: Annotated[list[dict], operator.add]
|
||||
"""The draft responses recommended by the AI."""
|
||||
|
||||
user_info: Annotated[Optional[dict], lambda x, y: y if y is not None else x]
|
||||
"""The current user state (by email)."""
|
||||
crm_info: Annotated[Optional[dict], lambda x, y: y if y is not None else x]
|
||||
"""The CRM information for organization the current user is from."""
|
||||
email_thread_id: Annotated[
|
||||
Optional[str], lambda x, y: y if y is not None else x
|
||||
]
|
||||
"""The current email thread ID."""
|
||||
slack_participants: Annotated[dict, operator.or_]
|
||||
"""The growing list of current slack participants."""
|
||||
bot_id: Optional[str]
|
||||
"""The ID of the bot user in the slack channel."""
|
||||
notified_assignees: Annotated[dict, operator.or_]
|
||||
|
||||
list_fields = {
|
||||
"messages",
|
||||
"trigger_events",
|
||||
"categorizations",
|
||||
"responses",
|
||||
"memory_docs",
|
||||
"relevant_rules",
|
||||
}
|
||||
dict_fields = {
|
||||
"user_info",
|
||||
"crm_info",
|
||||
"slack_participants",
|
||||
"notified_assignees",
|
||||
"autoresponse",
|
||||
"issue",
|
||||
}
|
||||
|
||||
def read_write(read: str, write: Sequence[str], input: State) -> dict:
|
||||
val = input.get(read)
|
||||
val = {val: val} if isinstance(val, str) else val
|
||||
val_single = val[-1] if isinstance(val, list) else val
|
||||
val_list = val if isinstance(val, list) else [val]
|
||||
return {
|
||||
k: val_list
|
||||
if k in list_fields
|
||||
else val_single
|
||||
if k in dict_fields
|
||||
else "".join(choice("abcdefghijklmnopqrstuvwxyz") for _ in range(n))
|
||||
for k in write
|
||||
}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_edge(START, "one")
|
||||
builder.add_node(
|
||||
"one",
|
||||
partial(read_write, "messages", ["trigger_events", "primary_issue_medium"]),
|
||||
)
|
||||
builder.add_edge("one", "two")
|
||||
builder.add_node(
|
||||
"two",
|
||||
partial(read_write, "trigger_events", ["autoresponse", "issue"]),
|
||||
)
|
||||
builder.add_edge("two", "three")
|
||||
builder.add_edge("two", "four")
|
||||
builder.add_node(
|
||||
"three",
|
||||
partial(read_write, "autoresponse", ["relevant_rules"]),
|
||||
)
|
||||
builder.add_node(
|
||||
"four",
|
||||
partial(
|
||||
read_write,
|
||||
"trigger_events",
|
||||
["categorizations", "responses", "memory_docs"],
|
||||
),
|
||||
)
|
||||
builder.add_node(
|
||||
"five",
|
||||
partial(
|
||||
read_write,
|
||||
"categorizations",
|
||||
[
|
||||
"user_info",
|
||||
"crm_info",
|
||||
"email_thread_id",
|
||||
"slack_participants",
|
||||
"bot_id",
|
||||
"notified_assignees",
|
||||
],
|
||||
),
|
||||
)
|
||||
builder.add_edge(["three", "four"], "five")
|
||||
builder.add_edge("five", "six")
|
||||
builder.add_node(
|
||||
"six",
|
||||
partial(read_write, "responses", ["messages"]),
|
||||
)
|
||||
builder.add_conditional_edges(
|
||||
"six", lambda state: END if len(state["messages"]) > n else "one"
|
||||
)
|
||||
|
||||
return builder
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
|
||||
import uvloop
|
||||
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
|
||||
graph = wide_dict(1000).compile(checkpointer=MemorySaver())
|
||||
input = {
|
||||
"messages": [
|
||||
{
|
||||
str(i) * 10: {
|
||||
str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
|
||||
for j in range(50)
|
||||
}
|
||||
for i in range(50)
|
||||
}
|
||||
]
|
||||
}
|
||||
config = {"configurable": {"thread_id": "1"}, "recursion_limit": 20000000000}
|
||||
|
||||
async def run():
|
||||
async for c in graph.astream(input, config=config):
|
||||
print(c.keys())
|
||||
|
||||
uvloop.install()
|
||||
asyncio.run(run())
|
||||
@@ -1,7 +1,6 @@
|
||||
import operator
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
from random import choice
|
||||
from typing import Annotated, Optional, Sequence
|
||||
|
||||
from langgraph.constants import END, START
|
||||
@@ -50,34 +49,12 @@ def wide_state(n: int) -> StateGraph:
|
||||
"""The ID of the bot user in the slack channel."""
|
||||
notified_assignees: Annotated[dict, operator.or_] = field(default_factory=dict)
|
||||
|
||||
list_fields = {
|
||||
"messages",
|
||||
"trigger_events",
|
||||
"categorizations",
|
||||
"responses",
|
||||
"memory_docs",
|
||||
"relevant_rules",
|
||||
}
|
||||
dict_fields = {
|
||||
"user_info",
|
||||
"crm_info",
|
||||
"slack_participants",
|
||||
"notified_assignees",
|
||||
"autoresponse",
|
||||
"issue",
|
||||
}
|
||||
|
||||
def read_write(read: str, write: Sequence[str], input: State) -> dict:
|
||||
val = getattr(input, read)
|
||||
val = {val: val} if isinstance(val, str) else val
|
||||
val_single = val[-1] if isinstance(val, list) else val
|
||||
val_list = val if isinstance(val, list) else [val]
|
||||
return {
|
||||
k: val_list
|
||||
if k in list_fields
|
||||
else val_single
|
||||
if k in dict_fields
|
||||
else "".join(choice("abcdefghijklmnopqrstuvwxyz") for _ in range(n))
|
||||
k: val_list if isinstance(getattr(input, k), list) else val_single
|
||||
for k in write
|
||||
}
|
||||
|
||||
|
||||
@@ -83,8 +83,6 @@ CONFIG_KEY_PREVIOUS = sys.intern("__pregel_previous")
|
||||
# holds the previous return value from a stateful Pregel graph.
|
||||
CONFIG_KEY_RUNNER_SUBMIT = sys.intern("__pregel_runner_submit")
|
||||
# holds a function that receives tasks from runner, executes them and returns results
|
||||
CONFIG_KEY_CHECKPOINT_DURING = sys.intern("__pregel_checkpoint_during")
|
||||
# holds a boolean indicating whether to checkpoint during the run (or only at the end)
|
||||
|
||||
# --- Other constants ---
|
||||
PUSH = sys.intern("__pregel_push")
|
||||
|
||||
@@ -1,206 +0,0 @@
|
||||
from typing import Any, Literal, Optional, Union
|
||||
from uuid import uuid4
|
||||
|
||||
from langchain_core.messages import AnyMessage
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND
|
||||
from langgraph.utils.config import get_config, get_stream_writer
|
||||
|
||||
|
||||
class UIMessage(TypedDict):
|
||||
"""A message type for UI updates in LangGraph.
|
||||
|
||||
This TypedDict represents a UI message that can be sent to update the UI state.
|
||||
It contains information about the UI component to render and its properties.
|
||||
|
||||
Attributes:
|
||||
type: Literal type indicating this is a UI message.
|
||||
id: Unique identifier for the UI message.
|
||||
name: Name of the UI component to render.
|
||||
props: Properties to pass to the UI component.
|
||||
metadata: Additional metadata about the UI message.
|
||||
"""
|
||||
|
||||
type: Literal["ui"]
|
||||
id: str
|
||||
name: str
|
||||
props: dict[str, Any]
|
||||
metadata: dict[str, Any]
|
||||
|
||||
|
||||
class RemoveUIMessage(TypedDict):
|
||||
"""A message type for removing UI components in LangGraph.
|
||||
|
||||
This TypedDict represents a message that can be sent to remove a UI component
|
||||
from the current state.
|
||||
|
||||
Attributes:
|
||||
type: Literal type indicating this is a remove-ui message.
|
||||
id: Unique identifier of the UI message to remove.
|
||||
"""
|
||||
|
||||
type: Literal["remove-ui"]
|
||||
id: str
|
||||
|
||||
|
||||
AnyUIMessage = Union[UIMessage, RemoveUIMessage]
|
||||
|
||||
|
||||
def push_ui_message(
|
||||
name: str,
|
||||
props: dict[str, Any],
|
||||
*,
|
||||
id: Optional[str] = None,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
message: Optional[AnyMessage] = None,
|
||||
state_key: str = "ui",
|
||||
) -> UIMessage:
|
||||
"""Push a new UI message to update the UI state.
|
||||
|
||||
This function creates and sends a UI message that will be rendered in the UI.
|
||||
It also updates the graph state with the new UI message.
|
||||
|
||||
Args:
|
||||
name: Name of the UI component to render.
|
||||
props: Properties to pass to the UI component.
|
||||
id: Optional unique identifier for the UI message.
|
||||
If not provided, a random UUID will be generated.
|
||||
metadata: Optional additional metadata about the UI message.
|
||||
message: Optional message object to associate with the UI message.
|
||||
state_key: Key in the graph state where the UI messages are stored.
|
||||
Defaults to "ui".
|
||||
|
||||
Returns:
|
||||
The created UI message.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
push_ui_message(
|
||||
name="component-name",
|
||||
props={"content": "Hello world"},
|
||||
)
|
||||
|
||||
"""
|
||||
writer = get_stream_writer()
|
||||
config = get_config()
|
||||
|
||||
message_id = None
|
||||
if message:
|
||||
if isinstance(message, dict) and "id" in message:
|
||||
message_id = message.get("id")
|
||||
elif hasattr(message, "id"):
|
||||
message_id = message.id
|
||||
|
||||
evt: UIMessage = {
|
||||
"type": "ui",
|
||||
"id": id or str(uuid4()),
|
||||
"name": name,
|
||||
"props": props,
|
||||
"metadata": {
|
||||
**(config.get("metadata") or {}),
|
||||
"tags": config.get("tags", None),
|
||||
"name": config.get("run_name", None),
|
||||
"run_id": config.get("run_id", None),
|
||||
**(metadata or {}),
|
||||
**({"message_id": message_id} if message_id else {}),
|
||||
},
|
||||
}
|
||||
|
||||
writer(evt)
|
||||
config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])
|
||||
|
||||
return evt
|
||||
|
||||
|
||||
def delete_ui_message(id: str, *, state_key: str = "ui") -> RemoveUIMessage:
|
||||
"""Delete a UI message by ID from the UI state.
|
||||
|
||||
This function creates and sends a message to remove a UI component from the current state.
|
||||
It also updates the graph state to remove the UI message.
|
||||
|
||||
Args:
|
||||
id: Unique identifier of the UI component to remove.
|
||||
state_key: Key in the graph state where the UI messages are stored. Defaults to "ui".
|
||||
|
||||
Returns:
|
||||
The remove UI message.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
delete_ui_message("message-123")
|
||||
|
||||
"""
|
||||
writer = get_stream_writer()
|
||||
config = get_config()
|
||||
|
||||
evt: RemoveUIMessage = {"type": "remove-ui", "id": id}
|
||||
|
||||
writer(evt)
|
||||
config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])
|
||||
|
||||
return evt
|
||||
|
||||
|
||||
def ui_message_reducer(
|
||||
left: Union[list[AnyUIMessage], AnyUIMessage],
|
||||
right: Union[list[AnyUIMessage], AnyUIMessage],
|
||||
) -> list[AnyUIMessage]:
|
||||
"""Merge two lists of UI messages, supporting removing UI messages.
|
||||
|
||||
This function combines two lists of UI messages, handling both regular UI messages
|
||||
and `remove-ui` messages. When a `remove-ui` message is encountered, it removes any
|
||||
UI message with the matching ID from the current state.
|
||||
|
||||
Args:
|
||||
left: First list of UI messages or single UI message.
|
||||
right: Second list of UI messages or single UI message.
|
||||
|
||||
Returns:
|
||||
Combined list of UI messages with removals applied.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
messages = ui_message_reducer(
|
||||
[{"type": "ui", "id": "1", "name": "Chat", "props": {}}],
|
||||
{"type": "remove-ui", "id": "1"}
|
||||
)
|
||||
|
||||
"""
|
||||
if not isinstance(left, list):
|
||||
left = [left]
|
||||
|
||||
if not isinstance(right, list):
|
||||
right = [right]
|
||||
|
||||
# merge messages
|
||||
merged = left.copy()
|
||||
merged_by_id = {m.get("id"): i for i, m in enumerate(merged)}
|
||||
ids_to_remove = set()
|
||||
|
||||
for msg in right:
|
||||
msg_id = msg.get("id")
|
||||
|
||||
if (existing_idx := merged_by_id.get(msg_id)) is not None:
|
||||
if msg.get("type") == "remove-ui":
|
||||
ids_to_remove.add(msg_id)
|
||||
else:
|
||||
ids_to_remove.discard(msg_id)
|
||||
merged[existing_idx] = msg
|
||||
else:
|
||||
if msg.get("type") == "remove-ui":
|
||||
raise ValueError(
|
||||
f"Attempting to delete an UI message with an ID that doesn't exist ('{msg_id}')"
|
||||
)
|
||||
|
||||
merged_by_id[msg_id] = len(merged)
|
||||
merged.append(msg)
|
||||
|
||||
merged = [m for m in merged if m.get("id") not in ids_to_remove]
|
||||
return merged
|
||||
@@ -39,6 +39,7 @@ from langchain_core.runnables.utils import (
|
||||
ConfigurableFieldSpec,
|
||||
get_unique_config_specs,
|
||||
)
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import Self
|
||||
|
||||
@@ -53,7 +54,6 @@ from langgraph.checkpoint.base import (
|
||||
)
|
||||
from langgraph.constants import (
|
||||
CONF,
|
||||
CONFIG_KEY_CHECKPOINT_DURING,
|
||||
CONFIG_KEY_CHECKPOINT_ID,
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
CONFIG_KEY_CHECKPOINTER,
|
||||
@@ -125,11 +125,6 @@ from langgraph.utils.fields import get_enhanced_type_hints
|
||||
from langgraph.utils.pydantic import create_model, is_supported_by_pydantic
|
||||
from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined]
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = None # type: ignore
|
||||
|
||||
WriteValue = Union[Callable[[Input], Output], Any]
|
||||
|
||||
|
||||
@@ -2099,7 +2094,6 @@ class Pregel(PregelProtocol):
|
||||
output_keys: Optional[Union[str, Sequence[str]]] = None,
|
||||
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
|
||||
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
|
||||
checkpoint_during: Optional[bool] = None,
|
||||
debug: Optional[bool] = None,
|
||||
subgraphs: bool = False,
|
||||
) -> Iterator[Union[dict[str, Any], Any]]:
|
||||
@@ -2121,7 +2115,6 @@ class Pregel(PregelProtocol):
|
||||
output_keys: The keys to stream, defaults to all non-context channels.
|
||||
interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph.
|
||||
interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph.
|
||||
checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved.
|
||||
debug: Whether to print debug information during execution, defaults to False.
|
||||
subgraphs: Whether to stream subgraphs, defaults to False.
|
||||
|
||||
@@ -2283,9 +2276,6 @@ class Pregel(PregelProtocol):
|
||||
config[CONF][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put(
|
||||
((), "custom", c)
|
||||
)
|
||||
# set checkpointing mode for subgraphs
|
||||
if checkpoint_during is not None:
|
||||
config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during
|
||||
with SyncPregelLoop(
|
||||
input,
|
||||
input_model=self.input_model,
|
||||
@@ -2301,9 +2291,6 @@ class Pregel(PregelProtocol):
|
||||
interrupt_after=interrupt_after_,
|
||||
manager=run_manager,
|
||||
debug=debug,
|
||||
checkpoint_during=checkpoint_during
|
||||
if checkpoint_during is not None
|
||||
else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True),
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
migrate_checkpoint=self._migrate_checkpoint,
|
||||
) as loop:
|
||||
@@ -2386,7 +2373,6 @@ class Pregel(PregelProtocol):
|
||||
output_keys: Optional[Union[str, Sequence[str]]] = None,
|
||||
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
|
||||
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
|
||||
checkpoint_during: Optional[bool] = None,
|
||||
debug: Optional[bool] = None,
|
||||
subgraphs: bool = False,
|
||||
) -> AsyncIterator[Union[dict[str, Any], Any]]:
|
||||
@@ -2408,7 +2394,6 @@ class Pregel(PregelProtocol):
|
||||
output_keys: The keys to stream, defaults to all non-context channels.
|
||||
interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph.
|
||||
interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph.
|
||||
checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved.
|
||||
debug: Whether to print debug information during execution, defaults to False.
|
||||
subgraphs: Whether to stream subgraphs, defaults to False.
|
||||
|
||||
@@ -2544,17 +2529,13 @@ class Pregel(PregelProtocol):
|
||||
run_id=config.get("run_id"),
|
||||
)
|
||||
# if running from astream_log() run each proc with streaming
|
||||
do_stream = (
|
||||
next(
|
||||
(
|
||||
cast(_StreamingCallbackHandler, h)
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if _StreamingCallbackHandler is not None
|
||||
else False
|
||||
do_stream = next(
|
||||
(
|
||||
cast(_StreamingCallbackHandler, h)
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
),
|
||||
None,
|
||||
)
|
||||
try:
|
||||
# assign defaults
|
||||
@@ -2590,9 +2571,6 @@ class Pregel(PregelProtocol):
|
||||
stream.put_nowait, ((), "custom", c)
|
||||
)
|
||||
)
|
||||
# set checkpointing mode for subgraphs
|
||||
if checkpoint_during is not None:
|
||||
config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during
|
||||
async with AsyncPregelLoop(
|
||||
input,
|
||||
input_model=self.input_model,
|
||||
@@ -2608,9 +2586,6 @@ class Pregel(PregelProtocol):
|
||||
interrupt_after=interrupt_after_,
|
||||
manager=run_manager,
|
||||
debug=debug,
|
||||
checkpoint_during=checkpoint_during
|
||||
if checkpoint_during is not None
|
||||
else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True),
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
migrate_checkpoint=self._migrate_checkpoint,
|
||||
) as loop:
|
||||
@@ -2686,7 +2661,6 @@ class Pregel(PregelProtocol):
|
||||
output_keys: Optional[Union[str, Sequence[str]]] = None,
|
||||
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
|
||||
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
|
||||
checkpoint_during: Optional[bool] = None,
|
||||
debug: Optional[bool] = None,
|
||||
**kwargs: Any,
|
||||
) -> Union[dict[str, Any], Any]:
|
||||
@@ -2718,7 +2692,6 @@ class Pregel(PregelProtocol):
|
||||
output_keys=output_keys,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
checkpoint_during=checkpoint_during,
|
||||
debug=debug,
|
||||
**kwargs,
|
||||
):
|
||||
@@ -2740,7 +2713,6 @@ class Pregel(PregelProtocol):
|
||||
output_keys: Optional[Union[str, Sequence[str]]] = None,
|
||||
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
|
||||
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
|
||||
checkpoint_during: Optional[bool] = None,
|
||||
debug: Optional[bool] = None,
|
||||
**kwargs: Any,
|
||||
) -> Union[dict[str, Any], Any]:
|
||||
@@ -2773,7 +2745,6 @@ class Pregel(PregelProtocol):
|
||||
output_keys=output_keys,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
checkpoint_during=checkpoint_during,
|
||||
debug=debug,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
@@ -63,7 +63,6 @@ from langgraph.constants import (
|
||||
RESUME,
|
||||
SCHEDULED,
|
||||
TAG_HIDDEN,
|
||||
TASKS,
|
||||
)
|
||||
from langgraph.errors import (
|
||||
CheckpointNotLatest,
|
||||
@@ -156,7 +155,7 @@ class PregelLoop(LoopProtocol):
|
||||
manager: Union[None, AsyncParentRunManager, ParentRunManager]
|
||||
interrupt_after: Union[All, Sequence[str]]
|
||||
interrupt_before: Union[All, Sequence[str]]
|
||||
checkpoint_during: bool
|
||||
checkpoint_every_step: bool
|
||||
debug: bool
|
||||
|
||||
checkpointer_get_next_version: GetNextVersion
|
||||
@@ -181,7 +180,6 @@ class PregelLoop(LoopProtocol):
|
||||
channels: Mapping[str, BaseChannel]
|
||||
managed: ManagedValueMapping
|
||||
checkpoint: Checkpoint
|
||||
checkpoint_id_saved: str
|
||||
checkpoint_ns: tuple[str, ...]
|
||||
checkpoint_config: RunnableConfig
|
||||
checkpoint_metadata: CheckpointMetadata
|
||||
@@ -217,7 +215,7 @@ class PregelLoop(LoopProtocol):
|
||||
debug: bool = False,
|
||||
migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
|
||||
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
|
||||
checkpoint_during: bool = True,
|
||||
checkpoint_every_step: bool = True,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
step=0,
|
||||
@@ -243,7 +241,7 @@ class PregelLoop(LoopProtocol):
|
||||
)
|
||||
self._migrate_checkpoint = migrate_checkpoint
|
||||
self.trigger_to_nodes = trigger_to_nodes
|
||||
self.checkpoint_during = checkpoint_during
|
||||
self.checkpoint_every_step = checkpoint_every_step
|
||||
self.debug = debug
|
||||
if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
|
||||
self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
|
||||
@@ -296,19 +294,29 @@ class PregelLoop(LoopProtocol):
|
||||
"""Put writes for a task, to be read by the next tick."""
|
||||
if not writes:
|
||||
return
|
||||
# always checkpoint writes containing Send, as they are fetched from the
|
||||
# parent checkpoint, not the current one
|
||||
checkpoint_during = self.checkpoint_during or any(w[0] == TASKS for w in writes)
|
||||
# deduplicate writes to special channels, last write wins
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes):
|
||||
writes = list({w[0]: w for w in writes}.values())
|
||||
# remove existing writes for this task
|
||||
self.checkpoint_pending_writes = [
|
||||
w for w in self.checkpoint_pending_writes if w[0] != task_id
|
||||
]
|
||||
# save writes
|
||||
self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
|
||||
if checkpoint_during and self.checkpointer_put_writes is not None:
|
||||
for c, v in writes:
|
||||
if (
|
||||
c in WRITES_IDX_MAP
|
||||
and (
|
||||
idx := next(
|
||||
(
|
||||
i
|
||||
for i, w in enumerate(self.checkpoint_pending_writes)
|
||||
if w[0] == task_id and w[1] == c
|
||||
),
|
||||
None,
|
||||
)
|
||||
)
|
||||
is not None
|
||||
):
|
||||
self.checkpoint_pending_writes[idx] = (task_id, c, v)
|
||||
else:
|
||||
self.checkpoint_pending_writes.append((task_id, c, v))
|
||||
if self.checkpointer_put_writes is not None:
|
||||
config = patch_configurable(
|
||||
self.checkpoint_config,
|
||||
{
|
||||
@@ -341,46 +349,6 @@ class PregelLoop(LoopProtocol):
|
||||
if hasattr(self, "tasks"):
|
||||
self._output_writes(task_id, writes)
|
||||
|
||||
def _put_pending_writes(self) -> None:
|
||||
if self.checkpointer_put_writes is None:
|
||||
return
|
||||
if not self.checkpoint_pending_writes:
|
||||
return
|
||||
# patch config
|
||||
config = patch_configurable(
|
||||
self.checkpoint_config,
|
||||
{
|
||||
CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
|
||||
CONFIG_KEY_CHECKPOINT_NS, ""
|
||||
),
|
||||
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
|
||||
},
|
||||
)
|
||||
# group by task id
|
||||
by_task = defaultdict(list)
|
||||
for task_id, channel, value in self.checkpoint_pending_writes:
|
||||
by_task[task_id].append((channel, value))
|
||||
# submit writes to checkpointer
|
||||
for task_id, writes in by_task.items():
|
||||
if self.checkpointer_put_writes_accepts_task_path and hasattr(
|
||||
self, "tasks"
|
||||
):
|
||||
task = self.tasks.get(task_id)
|
||||
self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
config,
|
||||
writes,
|
||||
task_id,
|
||||
task_path_str(task.path) if task else "",
|
||||
)
|
||||
else:
|
||||
self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
config,
|
||||
writes,
|
||||
task_id,
|
||||
)
|
||||
|
||||
def accept_push(
|
||||
self, task: PregelExecutableTask, write_idx: int, call: Optional[Call] = None
|
||||
) -> Optional[PregelExecutableTask]:
|
||||
@@ -743,44 +711,32 @@ class PregelLoop(LoopProtocol):
|
||||
|
||||
def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
|
||||
# assign step and parents
|
||||
exiting = metadata is self.checkpoint_metadata
|
||||
if exiting and self.checkpoint["id"] == self.checkpoint_id_saved:
|
||||
# checkpoint already saved
|
||||
return
|
||||
if not exiting:
|
||||
metadata["step"] = self.step
|
||||
metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {})
|
||||
self.checkpoint_metadata = metadata
|
||||
# debug flag
|
||||
if self.debug:
|
||||
print_step_checkpoint(
|
||||
metadata,
|
||||
self.channels,
|
||||
(
|
||||
[self.stream_keys]
|
||||
if isinstance(self.stream_keys, str)
|
||||
else self.stream_keys
|
||||
),
|
||||
)
|
||||
self.checkpoint_id_prev = self.checkpoint["id"] if self.step > -1 else None
|
||||
# do checkpoint?
|
||||
do_checkpoint = self._checkpointer_put_after_previous is not None and (
|
||||
exiting or self.checkpoint_during
|
||||
)
|
||||
# create new checkpoint
|
||||
self.checkpoint = create_checkpoint(
|
||||
self.checkpoint,
|
||||
self.channels if do_checkpoint else None,
|
||||
self.step,
|
||||
id=self.checkpoint["id"] if exiting else None,
|
||||
)
|
||||
metadata["step"] = self.step
|
||||
metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {})
|
||||
# debug flag
|
||||
if self.debug:
|
||||
print_step_checkpoint(
|
||||
metadata,
|
||||
self.channels,
|
||||
(
|
||||
[self.stream_keys]
|
||||
if isinstance(self.stream_keys, str)
|
||||
else self.stream_keys
|
||||
),
|
||||
)
|
||||
# bail if no checkpointer
|
||||
if do_checkpoint and self._checkpointer_put_after_previous is not None:
|
||||
if self._checkpointer_put_after_previous is not None:
|
||||
for k, v in self.config["metadata"].items():
|
||||
if k in EXCLUDED_METADATA_KEYS:
|
||||
continue
|
||||
metadata.setdefault(k, v) # type: ignore
|
||||
|
||||
# create new checkpoint
|
||||
self.checkpoint = create_checkpoint(
|
||||
self.checkpoint, self.channels, self.step
|
||||
)
|
||||
self.checkpoint_metadata = metadata
|
||||
|
||||
self.prev_checkpoint_config = (
|
||||
self.checkpoint_config
|
||||
if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF]
|
||||
@@ -791,8 +747,6 @@ class PregelLoop(LoopProtocol):
|
||||
**self.checkpoint_config,
|
||||
CONF: {
|
||||
**self.checkpoint_config[CONF],
|
||||
# this is guaranteed to be set by code above
|
||||
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint_id_prev,
|
||||
CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
|
||||
CONFIG_KEY_CHECKPOINT_NS, ""
|
||||
),
|
||||
@@ -823,9 +777,8 @@ class PregelLoop(LoopProtocol):
|
||||
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
|
||||
},
|
||||
}
|
||||
if not exiting:
|
||||
# increment step
|
||||
self.step += 1
|
||||
# increment step
|
||||
self.step += 1
|
||||
|
||||
def _update_mv(self, key: str, values: Sequence[Any]) -> None:
|
||||
raise NotImplementedError
|
||||
@@ -836,10 +789,6 @@ class PregelLoop(LoopProtocol):
|
||||
exc_value: Optional[BaseException],
|
||||
traceback: Optional[TracebackType],
|
||||
) -> Optional[bool]:
|
||||
# persist current checkpoint and writes
|
||||
if not self.checkpoint_during:
|
||||
self._put_checkpoint(self.checkpoint_metadata)
|
||||
self._put_pending_writes()
|
||||
# suppress interrupt
|
||||
suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
|
||||
if suppress:
|
||||
@@ -958,7 +907,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
|
||||
debug: bool = False,
|
||||
migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
|
||||
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
|
||||
checkpoint_during: bool = True,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
input,
|
||||
@@ -977,7 +925,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
|
||||
debug=debug,
|
||||
migrate_checkpoint=migrate_checkpoint,
|
||||
trigger_to_nodes=trigger_to_nodes,
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
self.stack = ExitStack()
|
||||
if checkpointer:
|
||||
@@ -1057,7 +1004,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
|
||||
},
|
||||
}
|
||||
self.prev_checkpoint_config = saved.parent_config
|
||||
self.checkpoint_id_saved = saved.checkpoint["id"]
|
||||
self.checkpoint = saved.checkpoint
|
||||
self.checkpoint_metadata = saved.metadata
|
||||
self.checkpoint_pending_writes = (
|
||||
@@ -1108,7 +1054,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
|
||||
debug: bool = False,
|
||||
migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
|
||||
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
|
||||
checkpoint_during: bool = True,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
input,
|
||||
@@ -1127,7 +1072,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
|
||||
debug=debug,
|
||||
migrate_checkpoint=migrate_checkpoint,
|
||||
trigger_to_nodes=trigger_to_nodes,
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
self.stack = AsyncExitStack()
|
||||
if checkpointer:
|
||||
@@ -1207,7 +1151,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
|
||||
},
|
||||
}
|
||||
self.prev_checkpoint_config = saved.parent_config
|
||||
self.checkpoint_id_saved = saved.checkpoint["id"]
|
||||
self.checkpoint = saved.checkpoint
|
||||
self.checkpoint_metadata = saved.metadata
|
||||
self.checkpoint_pending_writes = (
|
||||
|
||||
@@ -7,7 +7,6 @@ from typing import (
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
@@ -16,16 +15,11 @@ from uuid import UUID, uuid4
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
from langchain_core.messages import BaseMessage
|
||||
from langchain_core.outputs import ChatGenerationChunk, LLMResult
|
||||
from langchain_core.tracers._streaming import T, _StreamingCallbackHandler
|
||||
|
||||
from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM
|
||||
from langgraph.types import StreamChunk
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = object # type: ignore
|
||||
|
||||
T = TypeVar("T")
|
||||
Meta = tuple[tuple[str, ...], dict[str, Any]]
|
||||
|
||||
|
||||
|
||||
@@ -10,6 +10,7 @@ from typing import (
|
||||
cast,
|
||||
)
|
||||
|
||||
import orjson
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langchain_core.runnables.graph import (
|
||||
Edge as DrawableEdge,
|
||||
@@ -34,8 +35,6 @@ from typing_extensions import Self
|
||||
from langgraph.checkpoint.base import CheckpointMetadata
|
||||
from langgraph.constants import (
|
||||
CONF,
|
||||
CONFIG_KEY_CHECKPOINT_ID,
|
||||
CONFIG_KEY_CHECKPOINT_MAP,
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
CONFIG_KEY_STREAM,
|
||||
INTERRUPT,
|
||||
@@ -47,14 +46,6 @@ from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode
|
||||
from langgraph.types import Command, Interrupt, StreamProtocol
|
||||
from langgraph.utils.config import merge_configs
|
||||
|
||||
CONF_DROPLIST = frozenset(
|
||||
(
|
||||
CONFIG_KEY_CHECKPOINT_MAP,
|
||||
CONFIG_KEY_CHECKPOINT_ID,
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class RemoteException(Exception):
|
||||
"""Exception raised when an error occurs in the remote graph."""
|
||||
@@ -299,26 +290,47 @@ class RemoteGraph(PregelProtocol):
|
||||
}
|
||||
|
||||
def _sanitize_config(self, config: RunnableConfig) -> RunnableConfig:
|
||||
"""Sanitize the config to remove non-serializable fields."""
|
||||
sanitized: RunnableConfig = {}
|
||||
reserved_configurable_keys = frozenset(
|
||||
[
|
||||
"callbacks",
|
||||
"checkpoint_map",
|
||||
"checkpoint_id",
|
||||
"checkpoint_ns",
|
||||
]
|
||||
)
|
||||
|
||||
def _sanitize_obj(obj: Any) -> Any:
|
||||
"""Remove non-JSON serializable fields from the given object."""
|
||||
if isinstance(obj, dict):
|
||||
return {k: _sanitize_obj(v) for k, v in obj.items()}
|
||||
elif isinstance(obj, list):
|
||||
return [_sanitize_obj(v) for v in obj]
|
||||
else:
|
||||
try:
|
||||
orjson.dumps(obj)
|
||||
return obj
|
||||
except orjson.JSONEncodeError:
|
||||
return None
|
||||
|
||||
# Remove non-JSON serializable fields from the config.
|
||||
config = _sanitize_obj(config)
|
||||
|
||||
# Only include configurable keys that are not reserved and
|
||||
# not starting with "__pregel_" prefix.
|
||||
new_configurable = {
|
||||
k: v
|
||||
for k, v in config["configurable"].items()
|
||||
if k not in reserved_configurable_keys and not k.startswith("__pregel_")
|
||||
}
|
||||
|
||||
sanitized: RunnableConfig = {
|
||||
"tags": config.get("tags") or [],
|
||||
"metadata": config.get("metadata") or {},
|
||||
"configurable": new_configurable,
|
||||
}
|
||||
if "recursion_limit" in config:
|
||||
sanitized["recursion_limit"] = config["recursion_limit"]
|
||||
if "tags" in config:
|
||||
sanitized["tags"] = [tag for tag in config["tags"] if isinstance(tag, str)]
|
||||
if "metadata" in config:
|
||||
sanitized["metadata"] = {}
|
||||
for k, v in config["metadata"].items():
|
||||
if isinstance(k, str) and isinstance(v, (str, int, float, bool)):
|
||||
sanitized["metadata"][k] = v
|
||||
if "configurable" in config:
|
||||
sanitized["configurable"] = {}
|
||||
for k, v in config["configurable"].items():
|
||||
if (
|
||||
isinstance(k, str)
|
||||
and k not in CONF_DROPLIST
|
||||
and isinstance(v, (str, int, float, bool))
|
||||
):
|
||||
sanitized["configurable"][k] = v
|
||||
|
||||
return sanitized
|
||||
|
||||
def get_state(
|
||||
@@ -642,10 +654,9 @@ class RemoteGraph(PregelProtocol):
|
||||
# raise interrupt or errors
|
||||
if chunk.event.startswith("updates"):
|
||||
if isinstance(chunk.data, dict) and INTERRUPT in chunk.data:
|
||||
if caller_ns:
|
||||
raise GraphInterrupt(
|
||||
[Interrupt(**i) for i in chunk.data[INTERRUPT]]
|
||||
)
|
||||
raise GraphInterrupt(
|
||||
[Interrupt(**i) for i in chunk.data[INTERRUPT]]
|
||||
)
|
||||
elif chunk.event.startswith("error"):
|
||||
raise RemoteException(chunk.data)
|
||||
# filter for what was actually requested
|
||||
@@ -737,10 +748,9 @@ class RemoteGraph(PregelProtocol):
|
||||
# raise interrupt or errors
|
||||
if chunk.event.startswith("updates"):
|
||||
if isinstance(chunk.data, dict) and INTERRUPT in chunk.data:
|
||||
if caller_ns:
|
||||
raise GraphInterrupt(
|
||||
[Interrupt(**i) for i in chunk.data[INTERRUPT]]
|
||||
)
|
||||
raise GraphInterrupt(
|
||||
[Interrupt(**i) for i in chunk.data[INTERRUPT]]
|
||||
)
|
||||
elif chunk.event.startswith("error"):
|
||||
raise RemoteException(chunk.data)
|
||||
# filter for what was actually requested
|
||||
|
||||
@@ -36,6 +36,7 @@ from langchain_core.runnables.config import (
|
||||
var_child_runnable_config,
|
||||
)
|
||||
from langchain_core.runnables.utils import Input, Output
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
from typing_extensions import TypeGuard
|
||||
|
||||
from langgraph.constants import (
|
||||
@@ -53,11 +54,6 @@ from langgraph.utils.config import (
|
||||
patch_config,
|
||||
)
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = None # type: ignore
|
||||
|
||||
|
||||
def _set_config_context(
|
||||
config: RunnableConfig,
|
||||
@@ -687,15 +683,13 @@ class RunnableSeq(Runnable):
|
||||
iterator = step.stream(input, config, **kwargs)
|
||||
else:
|
||||
iterator = step.transform(iterator, config)
|
||||
if _StreamingCallbackHandler is not None and (
|
||||
stream_handler := next(
|
||||
(
|
||||
cast(_StreamingCallbackHandler, h)
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if stream_handler := next(
|
||||
(
|
||||
cast(_StreamingCallbackHandler, h)
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
),
|
||||
None,
|
||||
):
|
||||
# populates streamed_output in astream_log() output if needed
|
||||
iterator = stream_handler.tap_output_iter(run_manager.run_id, iterator)
|
||||
@@ -755,15 +749,13 @@ class RunnableSeq(Runnable):
|
||||
aiterator = step.atransform(aiterator, config)
|
||||
if hasattr(aiterator, "aclose"):
|
||||
stack.push_async_callback(aiterator.aclose)
|
||||
if _StreamingCallbackHandler is not None and (
|
||||
stream_handler := next(
|
||||
(
|
||||
cast(_StreamingCallbackHandler, h)
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if stream_handler := next(
|
||||
(
|
||||
cast(_StreamingCallbackHandler, h)
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
),
|
||||
None,
|
||||
):
|
||||
# populates streamed_output in astream_log() output if needed
|
||||
aiterator = stream_handler.tap_output_aiter(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph"
|
||||
version = "0.3.27"
|
||||
version = "0.3.24"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
license = "MIT"
|
||||
|
||||
@@ -1,20 +1,20 @@
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from tests.conftest import (
|
||||
REGULAR_CHECKPOINTERS_ASYNC,
|
||||
REGULAR_CHECKPOINTERS_SYNC,
|
||||
ALL_CHECKPOINTERS_ASYNC,
|
||||
ALL_CHECKPOINTERS_SYNC,
|
||||
awith_checkpointer,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_interruption_without_state_updates(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
|
||||
) -> None:
|
||||
"""Test interruption without state updates. This test confirms that
|
||||
interrupting doesn't require a state key having been updated in the prev step"""
|
||||
@@ -40,27 +40,20 @@ def test_interruption_without_state_updates(
|
||||
initial_input = {"input": "hello world"}
|
||||
thread = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
graph.invoke(initial_input, thread, checkpoint_during=checkpoint_during)
|
||||
graph.invoke(initial_input, thread, debug=True)
|
||||
assert graph.get_state(thread).next == ("step_2",)
|
||||
n_checkpoints = len([c for c in graph.get_state_history(thread)])
|
||||
assert n_checkpoints == (3 if checkpoint_during else 1)
|
||||
|
||||
graph.invoke(None, thread, checkpoint_during=checkpoint_during)
|
||||
graph.invoke(None, thread, debug=True)
|
||||
assert graph.get_state(thread).next == ("step_3",)
|
||||
n_checkpoints = len([c for c in graph.get_state_history(thread)])
|
||||
assert n_checkpoints == (4 if checkpoint_during else 2)
|
||||
|
||||
graph.invoke(None, thread, checkpoint_during=checkpoint_during)
|
||||
graph.invoke(None, thread, debug=True)
|
||||
assert graph.get_state(thread).next == ()
|
||||
n_checkpoints = len([c for c in graph.get_state_history(thread)])
|
||||
assert n_checkpoints == (5 if checkpoint_during else 3)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_interruption_without_state_updates_async(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
checkpointer_name: str, mocker: MockerFixture
|
||||
):
|
||||
"""Test interruption without state updates. This test confirms that
|
||||
interrupting doesn't require a state key having been updated in the prev step"""
|
||||
|
||||
@@ -85,17 +78,11 @@ async def test_interruption_without_state_updates_async(
|
||||
initial_input = {"input": "hello world"}
|
||||
thread = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during)
|
||||
await graph.ainvoke(initial_input, thread, debug=True)
|
||||
assert (await graph.aget_state(thread)).next == ("step_2",)
|
||||
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
|
||||
assert n_checkpoints == (3 if checkpoint_during else 1)
|
||||
|
||||
await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during)
|
||||
await graph.ainvoke(None, thread, debug=True)
|
||||
assert (await graph.aget_state(thread)).next == ("step_3",)
|
||||
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
|
||||
assert n_checkpoints == (4 if checkpoint_during else 2)
|
||||
|
||||
await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during)
|
||||
await graph.ainvoke(None, thread, debug=True)
|
||||
assert (await graph.aget_state(thread)).next == ()
|
||||
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
|
||||
assert n_checkpoints == (5 if checkpoint_during else 3)
|
||||
|
||||
@@ -7258,10 +7258,9 @@ def test_branch_then(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_send_dedupe_on_resume(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
) -> None:
|
||||
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
|
||||
|
||||
@@ -7317,7 +7316,7 @@ def test_send_dedupe_on_resume(
|
||||
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert graph.invoke(["0"], thread1, checkpoint_during=checkpoint_during) == [
|
||||
assert graph.invoke(["0"], thread1, debug=1) == [
|
||||
"0",
|
||||
"1",
|
||||
"3.1",
|
||||
@@ -7334,11 +7333,12 @@ def test_send_dedupe_on_resume(
|
||||
pytest.xfail("TODO: shallow checkpointer reports wrong next set")
|
||||
assert state.next == ("flaky",)
|
||||
# check history
|
||||
history = [c for c in graph.get_state_history(thread1)]
|
||||
assert len(history) == (4 if checkpoint_during else 1)
|
||||
if "shallow" not in checkpointer_name:
|
||||
history = [c for c in graph.get_state_history(thread1)]
|
||||
assert len(history) == 4
|
||||
|
||||
# resume execution
|
||||
assert graph.invoke(None, thread1, checkpoint_during=checkpoint_during) == [
|
||||
assert graph.invoke(None, thread1, debug=1) == [
|
||||
"0",
|
||||
"1",
|
||||
"3.1",
|
||||
@@ -7358,7 +7358,6 @@ def test_send_dedupe_on_resume(
|
||||
assert state.next == ()
|
||||
# check history
|
||||
history = [c for c in graph.get_state_history(thread1)]
|
||||
assert len(history) == (6 if checkpoint_during else 2)
|
||||
expected_history = [
|
||||
StateSnapshot(
|
||||
values=[
|
||||
@@ -7495,9 +7494,13 @@ def test_send_dedupe_on_resume(
|
||||
name="flaky",
|
||||
path=("__pregel_push", 1),
|
||||
error=None,
|
||||
interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),),
|
||||
interrupts=(
|
||||
Interrupt(
|
||||
value="Bahh", resumable=False, ns=None, when="during"
|
||||
),
|
||||
),
|
||||
state=None,
|
||||
result=["flaky|4"] if checkpoint_during else None,
|
||||
result=["flaky|4"],
|
||||
),
|
||||
PregelTask(
|
||||
id=AnyStr(),
|
||||
@@ -7634,11 +7637,10 @@ def test_send_dedupe_on_resume(
|
||||
),
|
||||
),
|
||||
]
|
||||
if checkpoint_during:
|
||||
assert history == expected_history
|
||||
else:
|
||||
assert history[0] == expected_history[0]
|
||||
assert history[1] == expected_history[2]
|
||||
if "shallow" in checkpointer_name:
|
||||
expected_history = expected_history[:1]
|
||||
|
||||
assert history == expected_history
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
|
||||
@@ -1115,14 +1115,10 @@ def test_invoke_checkpoint_two(
|
||||
assert checkpoint["channel_values"].get("total") == 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_pending_writes_resume(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
checkpointer: BaseCheckpointSaver = request.getfixturevalue(
|
||||
f"checkpointer_{checkpointer_name}"
|
||||
)
|
||||
@@ -1148,19 +1144,17 @@ def test_pending_writes_resume(
|
||||
self.calls = 0
|
||||
|
||||
one = AwhileMaker(0.1, {"value": 2})
|
||||
two = AwhileMaker(0.2, ConnectionError("I'm not good"))
|
||||
two = AwhileMaker(0.3, ConnectionError("I'm not good"))
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("one", one)
|
||||
builder.add_node(
|
||||
"two", two, retry=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False)
|
||||
)
|
||||
builder.add_node("two", two, retry=RetryPolicy(max_attempts=2))
|
||||
builder.add_edge(START, "one")
|
||||
builder.add_edge(START, "two")
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
|
||||
thread1: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
with pytest.raises(ConnectionError, match="I'm not good"):
|
||||
graph.invoke({"value": 1}, thread1, checkpoint_during=checkpoint_during)
|
||||
graph.invoke({"value": 1}, thread1)
|
||||
|
||||
# both nodes should have been called once
|
||||
assert one.calls == 1
|
||||
@@ -1206,7 +1200,7 @@ def test_pending_writes_resume(
|
||||
|
||||
# resume execution
|
||||
with pytest.raises(ConnectionError, match="I'm not good"):
|
||||
graph.invoke(None, thread1, checkpoint_during=checkpoint_during)
|
||||
graph.invoke(None, thread1)
|
||||
|
||||
# node "one" succeeded previously, so shouldn't be called again
|
||||
assert one.calls == 1
|
||||
@@ -1220,9 +1214,7 @@ def test_pending_writes_resume(
|
||||
# resume execution, without exception
|
||||
two.rtn = {"value": 3}
|
||||
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
|
||||
assert graph.invoke(None, thread1, checkpoint_during=checkpoint_during) == {
|
||||
"value": 6
|
||||
}
|
||||
assert graph.invoke(None, thread1) == {"value": 6}
|
||||
|
||||
if "shallow" in checkpointer_name:
|
||||
assert len(list(checkpointer.list(thread1))) == 1
|
||||
@@ -1231,7 +1223,7 @@ def test_pending_writes_resume(
|
||||
# check all final checkpoints
|
||||
checkpoints = [c for c in checkpointer.list(thread1)]
|
||||
# we should have 3
|
||||
assert len(checkpoints) == (3 if checkpoint_during else 2)
|
||||
assert len(checkpoints) == 3
|
||||
# the last one not too interesting for this test
|
||||
assert checkpoints[0] == CheckpointTuple(
|
||||
config={
|
||||
@@ -1333,26 +1325,15 @@ def test_pending_writes_resume(
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"]
|
||||
if checkpoint_during
|
||||
else AnyStr(),
|
||||
"checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
pending_writes=UnsortedSequence(
|
||||
(AnyStr(), "value", 2),
|
||||
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
|
||||
(AnyStr(), "value", 3),
|
||||
)
|
||||
if checkpoint_during
|
||||
else UnsortedSequence(
|
||||
(AnyStr(), "value", 2),
|
||||
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
|
||||
# the write against the previous checkpoint is not saved, as it is
|
||||
# produced in a run where only the next checkpoint (the last) is saved
|
||||
),
|
||||
)
|
||||
if not checkpoint_during:
|
||||
return
|
||||
assert checkpoints[2] == CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
@@ -1510,14 +1491,8 @@ def test_send_sequences() -> None:
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_imp_task(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> None:
|
||||
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
|
||||
mapper_calls = 0
|
||||
|
||||
@@ -1583,7 +1558,7 @@ def test_imp_task(
|
||||
}
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert [*graph.stream([0, 1], thread1, checkpoint_during=checkpoint_during)] == [
|
||||
assert [*graph.stream([0, 1], thread1)] == [
|
||||
{"mapper": "00"},
|
||||
{"mapper": "11"},
|
||||
{
|
||||
@@ -1599,23 +1574,17 @@ def test_imp_task(
|
||||
]
|
||||
assert mapper_calls == 2
|
||||
|
||||
assert graph.invoke(
|
||||
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
|
||||
) == [
|
||||
assert graph.invoke(Command(resume="answer"), thread1) == [
|
||||
"00answer",
|
||||
"11answer",
|
||||
]
|
||||
assert mapper_calls == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_imp_nested(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
|
||||
|
||||
def mynode(input: list[str]) -> list[str]:
|
||||
@@ -1657,7 +1626,7 @@ def test_imp_nested(
|
||||
}
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert [*graph.stream([0, 1], thread1, checkpoint_during=checkpoint_during)] == [
|
||||
assert [*graph.stream([0, 1], thread1)] == [
|
||||
{"submapper": "0"},
|
||||
{"mapper": "00"},
|
||||
{"submapper": "1"},
|
||||
@@ -1674,22 +1643,16 @@ def test_imp_nested(
|
||||
},
|
||||
]
|
||||
|
||||
assert graph.invoke(
|
||||
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
|
||||
) == [
|
||||
assert graph.invoke(Command(resume="answer"), thread1) == [
|
||||
"00answera",
|
||||
"11answera",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_imp_stream_order(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
|
||||
|
||||
@task()
|
||||
@@ -1712,10 +1675,7 @@ def test_imp_stream_order(
|
||||
return fut_baz.result()
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert [
|
||||
c
|
||||
for c in graph.stream({"a": "0"}, thread1, checkpoint_during=checkpoint_during)
|
||||
] == [
|
||||
assert [c for c in graph.stream({"a": "0"}, thread1)] == [
|
||||
{
|
||||
"foo": (
|
||||
"0foo",
|
||||
@@ -3683,14 +3643,10 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None:
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_subgraph_checkpoint_true(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Unsupported combo")
|
||||
|
||||
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
|
||||
|
||||
class InnerState(TypedDict):
|
||||
@@ -3722,12 +3678,7 @@ def test_subgraph_checkpoint_true(
|
||||
app = graph.compile(checkpointer=checkpointer)
|
||||
|
||||
config = {"configurable": {"thread_id": "2"}}
|
||||
assert [
|
||||
c
|
||||
for c in app.stream(
|
||||
{"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c for c in app.stream({"my_key": ""}, config, subgraphs=True)] == [
|
||||
(("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}),
|
||||
(("inner",), {"inner_2": {"my_key": " and there"}}),
|
||||
((), {"inner": {"my_key": " got here and there"}}),
|
||||
@@ -3752,14 +3703,10 @@ def test_subgraph_checkpoint_true(
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_subgraph_checkpoint_true_interrupt(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Unsupported combo")
|
||||
|
||||
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
|
||||
|
||||
# Define subgraph
|
||||
@@ -3798,18 +3745,15 @@ def test_subgraph_checkpoint_true_interrupt(
|
||||
builder.add_edge(START, "node_1")
|
||||
builder.add_edge("node_1", "node_2")
|
||||
|
||||
checkpointer = MemorySaver()
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
assert graph.invoke(
|
||||
{"foo": "foo"}, config, checkpoint_during=checkpoint_during
|
||||
) == {"foo": "hi! foo"}
|
||||
assert graph.invoke({"foo": "foo"}, config) == {"foo": "hi! foo"}
|
||||
assert graph.get_state(config, subgraphs=True).tasks[0].state.values == {
|
||||
"bar": "hi! foo"
|
||||
}
|
||||
assert graph.invoke(
|
||||
Command(resume="baz"), config, checkpoint_during=checkpoint_during
|
||||
) == {"foo": "hi! foobaz"}
|
||||
assert graph.invoke(Command(resume="baz"), config) == {"foo": "hi! foobaz"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
@@ -3925,14 +3869,10 @@ def test_stream_buffering_single_node(
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_nested_graph_interrupts_parallel(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Unsupported combo")
|
||||
|
||||
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
|
||||
|
||||
class InnerState(TypedDict):
|
||||
@@ -3979,11 +3919,11 @@ def test_nested_graph_interrupts_parallel(
|
||||
|
||||
# test invoke w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
assert app.invoke({"my_key": ""}, config, checkpoint_during=checkpoint_during) == {
|
||||
assert app.invoke({"my_key": ""}, config, debug=True) == {
|
||||
"my_key": " and parallel",
|
||||
}
|
||||
|
||||
assert app.invoke(None, config, checkpoint_during=checkpoint_during) == {
|
||||
assert app.invoke(None, config, debug=True) == {
|
||||
"my_key": "got here and there and parallel and back again",
|
||||
}
|
||||
|
||||
@@ -3992,17 +3932,13 @@ def test_nested_graph_interrupts_parallel(
|
||||
# - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream)
|
||||
# test stream updates w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "2"}}
|
||||
assert [
|
||||
*app.stream(
|
||||
{"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream({"my_key": ""}, config, subgraphs=True)] == [
|
||||
# we got to parallel node first
|
||||
((), {"outer_1": {"my_key": " and parallel"}}),
|
||||
((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}),
|
||||
((), {"__interrupt__": ()}),
|
||||
]
|
||||
assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [
|
||||
assert [*app.stream(None, config)] == [
|
||||
{"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}},
|
||||
{"inner": {"my_key": "got here and there"}},
|
||||
{"outer_2": {"my_key": " and back again"}},
|
||||
@@ -4010,22 +3946,11 @@ def test_nested_graph_interrupts_parallel(
|
||||
|
||||
# test stream values w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "3"}}
|
||||
assert [
|
||||
*app.stream(
|
||||
{"my_key": ""},
|
||||
config,
|
||||
stream_mode="values",
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": " and parallel"},
|
||||
]
|
||||
assert [
|
||||
*app.stream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
@@ -4034,28 +3959,15 @@ def test_nested_graph_interrupts_parallel(
|
||||
# test interrupts BEFORE the parallel node
|
||||
app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"])
|
||||
config = {"configurable": {"thread_id": "4"}}
|
||||
assert [
|
||||
*app.stream(
|
||||
{"my_key": ""},
|
||||
config,
|
||||
stream_mode="values",
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
] == [{"my_key": ""}]
|
||||
assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [
|
||||
{"my_key": ""}
|
||||
]
|
||||
# while we're waiting for the node w/ interrupt inside to finish
|
||||
assert [
|
||||
*app.stream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": " and parallel"},
|
||||
]
|
||||
assert [
|
||||
*app.stream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
@@ -4064,43 +3976,24 @@ def test_nested_graph_interrupts_parallel(
|
||||
# test interrupts AFTER the parallel node
|
||||
app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"])
|
||||
config = {"configurable": {"thread_id": "5"}}
|
||||
assert [
|
||||
*app.stream(
|
||||
{"my_key": ""},
|
||||
config,
|
||||
stream_mode="values",
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": " and parallel"},
|
||||
]
|
||||
assert [
|
||||
*app.stream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
]
|
||||
assert [
|
||||
*app.stream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_doubly_nested_graph_interrupts(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Unsupported combo")
|
||||
|
||||
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
|
||||
|
||||
class State(TypedDict):
|
||||
@@ -4154,13 +4047,11 @@ def test_doubly_nested_graph_interrupts(
|
||||
|
||||
# test invoke w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
assert app.invoke(
|
||||
{"my_key": "my value"}, config, checkpoint_during=checkpoint_during
|
||||
) == {
|
||||
assert app.invoke({"my_key": "my value"}, config, debug=True) == {
|
||||
"my_key": "hi my value",
|
||||
}
|
||||
|
||||
assert app.invoke(None, config, checkpoint_during=checkpoint_during) == {
|
||||
assert app.invoke(None, config, debug=True) == {
|
||||
"my_key": "hi my value here and there and back again",
|
||||
}
|
||||
|
||||
@@ -4169,14 +4060,12 @@ def test_doubly_nested_graph_interrupts(
|
||||
config = {
|
||||
"configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append}
|
||||
}
|
||||
assert [
|
||||
*app.stream({"my_key": "my value"}, config, checkpoint_during=checkpoint_during)
|
||||
] == [
|
||||
assert [*app.stream({"my_key": "my value"}, config)] == [
|
||||
{"parent_1": {"my_key": "hi my value"}},
|
||||
{"__interrupt__": ()},
|
||||
]
|
||||
assert nodes == ["parent_1", "grandchild_1"]
|
||||
assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [
|
||||
assert [*app.stream(None, config)] == [
|
||||
{"child": {"my_key": "hi my value here and there"}},
|
||||
{"parent_2": {"my_key": "hi my value here and there and back again"}},
|
||||
]
|
||||
@@ -4191,22 +4080,11 @@ def test_doubly_nested_graph_interrupts(
|
||||
|
||||
# test stream values w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "3"}}
|
||||
assert [
|
||||
*app.stream(
|
||||
{"my_key": "my value"},
|
||||
config,
|
||||
stream_mode="values",
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [
|
||||
{"my_key": "my value"},
|
||||
{"my_key": "hi my value"},
|
||||
]
|
||||
assert [
|
||||
*app.stream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{"my_key": "hi my value"},
|
||||
{"my_key": "hi my value here and there"},
|
||||
{"my_key": "hi my value here and there and back again"},
|
||||
|
||||
@@ -1947,14 +1947,10 @@ async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str)
|
||||
assert checkpoint["channel_values"].get("total") == 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_pending_writes_resume(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
class State(TypedDict):
|
||||
value: Annotated[int, operator.add]
|
||||
|
||||
@@ -1976,12 +1972,10 @@ async def test_pending_writes_resume(
|
||||
self.calls = 0
|
||||
|
||||
one = AwhileMaker(0.1, {"value": 2})
|
||||
two = AwhileMaker(0.2, ConnectionError("I'm not good"))
|
||||
two = AwhileMaker(0.3, ConnectionError("I'm not good"))
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("one", one)
|
||||
builder.add_node(
|
||||
"two", two, retry=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False)
|
||||
)
|
||||
builder.add_node("two", two, retry=RetryPolicy(max_attempts=2))
|
||||
builder.add_edge(START, "one")
|
||||
builder.add_edge(START, "two")
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
@@ -1989,9 +1983,7 @@ async def test_pending_writes_resume(
|
||||
|
||||
thread1: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
with pytest.raises(ConnectionError, match="I'm not good"):
|
||||
await graph.ainvoke(
|
||||
{"value": 1}, thread1, checkpoint_during=checkpoint_during
|
||||
)
|
||||
await graph.ainvoke({"value": 1}, thread1)
|
||||
|
||||
# both nodes should have been called once
|
||||
assert one.calls == 1
|
||||
@@ -2042,7 +2034,7 @@ async def test_pending_writes_resume(
|
||||
|
||||
# resume execution
|
||||
with pytest.raises(ConnectionError, match="I'm not good"):
|
||||
await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during)
|
||||
await graph.ainvoke(None, thread1)
|
||||
|
||||
# node "one" succeeded previously, so shouldn't be called again
|
||||
assert one.calls == 1
|
||||
@@ -2056,9 +2048,7 @@ async def test_pending_writes_resume(
|
||||
# resume execution, without exception
|
||||
two.rtn = {"value": 3}
|
||||
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
|
||||
assert await graph.ainvoke(
|
||||
None, thread1, checkpoint_during=checkpoint_during
|
||||
) == {"value": 6}
|
||||
assert await graph.ainvoke(None, thread1) == {"value": 6}
|
||||
|
||||
if "shallow" in checkpointer_name:
|
||||
assert len([c async for c in checkpointer.alist(thread1)]) == 1
|
||||
@@ -2067,7 +2057,7 @@ async def test_pending_writes_resume(
|
||||
# check all final checkpoints
|
||||
checkpoints = [c async for c in checkpointer.alist(thread1)]
|
||||
# we should have 3
|
||||
assert len(checkpoints) == (3 if checkpoint_during else 2)
|
||||
assert len(checkpoints) == 3
|
||||
# the last one not too interesting for this test
|
||||
assert checkpoints[0] == CheckpointTuple(
|
||||
config={
|
||||
@@ -2173,26 +2163,15 @@ async def test_pending_writes_resume(
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": checkpoints[2].config["configurable"][
|
||||
"checkpoint_id"
|
||||
]
|
||||
if checkpoint_during
|
||||
else AnyStr(),
|
||||
],
|
||||
}
|
||||
},
|
||||
pending_writes=UnsortedSequence(
|
||||
(AnyStr(), "value", 2),
|
||||
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
|
||||
(AnyStr(), "value", 3),
|
||||
)
|
||||
if checkpoint_during
|
||||
else UnsortedSequence(
|
||||
(AnyStr(), "value", 2),
|
||||
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
|
||||
# the write against the previous checkpoint is not saved, as it is
|
||||
# produced in a run where only the next checkpoint (the last) is saved
|
||||
),
|
||||
)
|
||||
if not checkpoint_during:
|
||||
return
|
||||
assert checkpoints[2] == CheckpointTuple(
|
||||
config={
|
||||
"configurable": {
|
||||
@@ -2230,7 +2209,7 @@ async def test_pending_writes_resume(
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
|
||||
async def test_run_from_checkpoint_id_retains_previous_writes(
|
||||
checkpointer_name: str,
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
|
||||
) -> None:
|
||||
class MyState(TypedDict):
|
||||
myval: Annotated[int, operator.add]
|
||||
@@ -2275,8 +2254,8 @@ async def test_run_from_checkpoint_id_retains_previous_writes(
|
||||
history = [c async for c in graph.aget_state_history(thread1)]
|
||||
|
||||
assert len(history) == 4
|
||||
assert history[0].values == {"myval": 4, "otherval": False}
|
||||
assert history[-1].values == {"myval": 0}
|
||||
assert history[0].values == {"myval": 4, "otherval": False}
|
||||
|
||||
second_run_config = {
|
||||
**thread1,
|
||||
@@ -2453,12 +2432,8 @@ async def test_send_sequences(checkpointer_name: str) -> None:
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
async def test_imp_task(checkpointer_name: str) -> None:
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
mapper_calls = 0
|
||||
|
||||
@@ -2478,12 +2453,7 @@ async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None
|
||||
|
||||
tracer = FakeTracer()
|
||||
thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]}
|
||||
assert [
|
||||
c
|
||||
async for c in graph.astream(
|
||||
[0, 1], thread1, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in graph.astream([0, 1], thread1)] == [
|
||||
{"mapper": "00"},
|
||||
{"mapper": "11"},
|
||||
{
|
||||
@@ -2507,9 +2477,7 @@ async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None
|
||||
assert any(r.inputs == {"input": 0} for r in mapper_runs)
|
||||
assert any(r.inputs == {"input": 1} for r in mapper_runs)
|
||||
|
||||
assert await graph.ainvoke(
|
||||
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
|
||||
) == [
|
||||
assert await graph.ainvoke(Command(resume="answer"), thread1) == [
|
||||
"00answer",
|
||||
"11answer",
|
||||
]
|
||||
@@ -2517,12 +2485,8 @@ async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
async def test_imp_nested(checkpointer_name: str) -> None:
|
||||
async def mynode(input: list[str]) -> list[str]:
|
||||
return [it + "a" for it in input]
|
||||
|
||||
@@ -2562,12 +2526,7 @@ async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> No
|
||||
}
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert [
|
||||
c
|
||||
async for c in graph.astream(
|
||||
[0, 1], thread1, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in graph.astream([0, 1], thread1)] == [
|
||||
{"submapper": "0"},
|
||||
{"mapper": "00"},
|
||||
{"submapper": "1"},
|
||||
@@ -2584,21 +2543,15 @@ async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> No
|
||||
},
|
||||
]
|
||||
|
||||
assert await graph.ainvoke(
|
||||
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
|
||||
) == [
|
||||
assert await graph.ainvoke(Command(resume="answer"), thread1) == [
|
||||
"00answera",
|
||||
"11answera",
|
||||
]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
async def test_imp_task_cancel(checkpointer_name: str) -> None:
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
mapper_calls = 0
|
||||
mapper_cancels = 0
|
||||
@@ -2624,12 +2577,7 @@ async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool)
|
||||
return [m + answer for m in mapped]
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert [
|
||||
c
|
||||
async for c in graph.astream(
|
||||
[0, 1], thread1, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in graph.astream([0, 1], thread1)] == [
|
||||
{"mapper": "00"},
|
||||
{
|
||||
"__interrupt__": (
|
||||
@@ -2645,9 +2593,7 @@ async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool)
|
||||
assert mapper_calls == 2
|
||||
assert mapper_cancels == 1
|
||||
|
||||
assert await graph.ainvoke(
|
||||
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
|
||||
) == [
|
||||
assert await graph.ainvoke(Command(resume="answer"), thread1) == [
|
||||
"00answer",
|
||||
]
|
||||
assert mapper_calls == 3
|
||||
@@ -2655,14 +2601,8 @@ async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool)
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_imp_sync_from_async(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
async def test_imp_sync_from_async(checkpointer_name: str) -> None:
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
|
||||
@task()
|
||||
@@ -2685,12 +2625,7 @@ async def test_imp_sync_from_async(
|
||||
return fut_baz.result()
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert [
|
||||
c
|
||||
async for c in graph.astream(
|
||||
{"a": "0"}, thread1, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in graph.astream({"a": "0"}, thread1)] == [
|
||||
{"foo": {"a": "0foo", "b": "bar"}},
|
||||
{"bar": {"a": "0foobar", "c": "bark"}},
|
||||
{"baz": {"a": "0foobarbaz", "c": "something else"}},
|
||||
@@ -2699,14 +2634,8 @@ async def test_imp_sync_from_async(
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_imp_stream_order(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
if not checkpoint_during and "shallow" in checkpointer_name:
|
||||
pytest.skip("Checkpointing during execution not supported")
|
||||
|
||||
async def test_imp_stream_order(checkpointer_name: str) -> None:
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
|
||||
@task()
|
||||
@@ -2730,12 +2659,7 @@ async def test_imp_stream_order(
|
||||
return await fut_baz
|
||||
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert [
|
||||
c
|
||||
async for c in graph.astream(
|
||||
{"a": "0"}, thread1, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in graph.astream({"a": "0"}, thread1)] == [
|
||||
{"foo": {"a": "0foo", "b": "bar"}},
|
||||
{"bar": {"a": "0foobar", "c": "bark"}},
|
||||
{"baz": {"a": "0foobarbaz", "c": "something else"}},
|
||||
@@ -2743,11 +2667,8 @@ async def test_imp_stream_order(
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
|
||||
async def test_send_dedupe_on_resume(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
|
||||
class InterruptOnce:
|
||||
ticks: int = 0
|
||||
|
||||
@@ -2798,9 +2719,7 @@ async def test_send_dedupe_on_resume(
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
assert await graph.ainvoke(
|
||||
["0"], thread1, checkpoint_during=checkpoint_during
|
||||
) == [
|
||||
assert await graph.ainvoke(["0"], thread1, debug=1) == [
|
||||
"0",
|
||||
"1",
|
||||
"3.1",
|
||||
@@ -2812,9 +2731,7 @@ async def test_send_dedupe_on_resume(
|
||||
assert builder.nodes["2"].runnable.func.ticks == 3
|
||||
assert builder.nodes["flaky"].runnable.func.ticks == 1
|
||||
# resume execution
|
||||
assert await graph.ainvoke(
|
||||
None, thread1, checkpoint_during=checkpoint_during
|
||||
) == [
|
||||
assert await graph.ainvoke(None, thread1, debug=1) == [
|
||||
"0",
|
||||
"1",
|
||||
"3.1",
|
||||
@@ -2831,8 +2748,7 @@ async def test_send_dedupe_on_resume(
|
||||
assert builder.nodes["flaky"].runnable.func.ticks == 2
|
||||
# check history
|
||||
history = [c async for c in graph.aget_state_history(thread1)]
|
||||
assert len(history) == (6 if checkpoint_during else 2)
|
||||
expected_history = [
|
||||
assert history == [
|
||||
StateSnapshot(
|
||||
values=[
|
||||
"0",
|
||||
@@ -2968,9 +2884,13 @@ async def test_send_dedupe_on_resume(
|
||||
name="flaky",
|
||||
path=("__pregel_push", 1),
|
||||
error=None,
|
||||
interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),),
|
||||
interrupts=(
|
||||
Interrupt(
|
||||
value="Bahh", resumable=False, ns=None, when="during"
|
||||
),
|
||||
),
|
||||
state=None,
|
||||
result=["flaky|4"] if checkpoint_during else None,
|
||||
result=["flaky|4"],
|
||||
),
|
||||
PregelTask(
|
||||
id=AnyStr(),
|
||||
@@ -3107,11 +3027,6 @@ async def test_send_dedupe_on_resume(
|
||||
),
|
||||
),
|
||||
]
|
||||
if checkpoint_during:
|
||||
assert history == expected_history
|
||||
else:
|
||||
assert history[0] == expected_history[0]
|
||||
assert history[1] == expected_history[2]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
@@ -5433,132 +5348,6 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
|
||||
assert times_called == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
|
||||
async def test_subgraph_checkpoint_true(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
class InnerState(TypedDict):
|
||||
my_key: Annotated[str, operator.add]
|
||||
my_other_key: str
|
||||
|
||||
def inner_1(state: InnerState):
|
||||
return {"my_key": " got here", "my_other_key": state["my_key"]}
|
||||
|
||||
def inner_2(state: InnerState):
|
||||
return {"my_key": " and there"}
|
||||
|
||||
inner = StateGraph(InnerState)
|
||||
inner.add_node("inner_1", inner_1)
|
||||
inner.add_node("inner_2", inner_2)
|
||||
inner.add_edge("inner_1", "inner_2")
|
||||
inner.set_entry_point("inner_1")
|
||||
inner.set_finish_point("inner_2")
|
||||
|
||||
class State(TypedDict):
|
||||
my_key: str
|
||||
|
||||
graph = StateGraph(State)
|
||||
graph.add_node("inner", inner.compile(checkpointer=True))
|
||||
graph.add_edge(START, "inner")
|
||||
graph.add_conditional_edges(
|
||||
"inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
|
||||
)
|
||||
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
app = graph.compile(checkpointer=checkpointer)
|
||||
|
||||
config = {"configurable": {"thread_id": "2"}}
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
{"my_key": ""},
|
||||
config,
|
||||
subgraphs=True,
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
] == [
|
||||
(("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}),
|
||||
(("inner",), {"inner_2": {"my_key": " and there"}}),
|
||||
((), {"inner": {"my_key": " got here and there"}}),
|
||||
(
|
||||
("inner",),
|
||||
{
|
||||
"inner_1": {
|
||||
"my_key": " got here",
|
||||
"my_other_key": " got here and there got here and there",
|
||||
}
|
||||
},
|
||||
),
|
||||
(("inner",), {"inner_2": {"my_key": " and there"}}),
|
||||
(
|
||||
(),
|
||||
{
|
||||
"inner": {
|
||||
"my_key": " got here and there got here and there got here and there"
|
||||
}
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
|
||||
async def test_subgraph_checkpoint_true_interrupt(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
# Define subgraph
|
||||
class SubgraphState(TypedDict):
|
||||
# note that none of these keys are shared with the parent graph state
|
||||
bar: str
|
||||
baz: str
|
||||
|
||||
def subgraph_node_1(state: SubgraphState):
|
||||
baz_value = interrupt("Provide baz value")
|
||||
return {"baz": baz_value}
|
||||
|
||||
def subgraph_node_2(state: SubgraphState):
|
||||
return {"bar": state["bar"] + state["baz"]}
|
||||
|
||||
subgraph_builder = StateGraph(SubgraphState)
|
||||
subgraph_builder.add_node(subgraph_node_1)
|
||||
subgraph_builder.add_node(subgraph_node_2)
|
||||
subgraph_builder.add_edge(START, "subgraph_node_1")
|
||||
subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2")
|
||||
subgraph = subgraph_builder.compile(checkpointer=True)
|
||||
|
||||
class ParentState(TypedDict):
|
||||
foo: str
|
||||
|
||||
def node_1(state: ParentState):
|
||||
return {"foo": "hi! " + state["foo"]}
|
||||
|
||||
async def node_2(state: ParentState, config: RunnableConfig):
|
||||
response = await subgraph.ainvoke({"bar": state["foo"]})
|
||||
return {"foo": response["bar"]}
|
||||
|
||||
builder = StateGraph(ParentState)
|
||||
builder.add_node("node_1", node_1)
|
||||
builder.add_node("node_2", node_2)
|
||||
builder.add_edge(START, "node_1")
|
||||
builder.add_edge("node_1", "node_2")
|
||||
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
graph = builder.compile(checkpointer=checkpointer)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
assert await graph.ainvoke(
|
||||
{"foo": "foo"}, config, checkpoint_during=checkpoint_during
|
||||
) == {"foo": "hi! foo"}
|
||||
assert (await graph.aget_state(config, subgraphs=True)).tasks[
|
||||
0
|
||||
].state.values == {"bar": "hi! foo"}
|
||||
assert await graph.ainvoke(
|
||||
Command(resume="baz"), config, checkpoint_during=checkpoint_during
|
||||
) == {"foo": "hi! foobaz"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None:
|
||||
class InnerState(TypedDict):
|
||||
@@ -5667,11 +5456,8 @@ async def test_stream_buffering_single_node(checkpointer_name: str) -> None:
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_nested_graph_interrupts_parallel(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None:
|
||||
class InnerState(TypedDict):
|
||||
my_key: Annotated[str, operator.add]
|
||||
my_other_key: str
|
||||
@@ -5720,13 +5506,11 @@ async def test_nested_graph_interrupts_parallel(
|
||||
|
||||
# test invoke w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
assert await app.ainvoke(
|
||||
{"my_key": ""}, config, checkpoint_during=checkpoint_during
|
||||
) == {
|
||||
assert await app.ainvoke({"my_key": ""}, config, debug=True) == {
|
||||
"my_key": " and parallel",
|
||||
}
|
||||
|
||||
assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == {
|
||||
assert await app.ainvoke(None, config, debug=True) == {
|
||||
"my_key": "got here and there and parallel and back again",
|
||||
}
|
||||
|
||||
@@ -5736,13 +5520,7 @@ async def test_nested_graph_interrupts_parallel(
|
||||
# test stream updates w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "2"}}
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
{"my_key": ""},
|
||||
config,
|
||||
subgraphs=True,
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
c async for c in app.astream({"my_key": ""}, config, subgraphs=True)
|
||||
] == [
|
||||
# we got to parallel node first
|
||||
((), {"outer_1": {"my_key": " and parallel"}}),
|
||||
@@ -5752,12 +5530,7 @@ async def test_nested_graph_interrupts_parallel(
|
||||
),
|
||||
((), {"__interrupt__": ()}),
|
||||
]
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
None, config, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream(None, config)] == [
|
||||
{"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}},
|
||||
{"inner": {"my_key": "got here and there"}},
|
||||
{"outer_2": {"my_key": " and back again"}},
|
||||
@@ -5766,23 +5539,12 @@ async def test_nested_graph_interrupts_parallel(
|
||||
# test stream values w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "3"}}
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
{"my_key": ""},
|
||||
config,
|
||||
stream_mode="values",
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
c async for c in app.astream({"my_key": ""}, config, stream_mode="values")
|
||||
] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": " and parallel"},
|
||||
]
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
@@ -5792,32 +5554,16 @@ async def test_nested_graph_interrupts_parallel(
|
||||
app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"])
|
||||
config = {"configurable": {"thread_id": "4"}}
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
{"my_key": ""},
|
||||
config,
|
||||
stream_mode="values",
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
c async for c in app.astream({"my_key": ""}, config, stream_mode="values")
|
||||
] == [
|
||||
{"my_key": ""},
|
||||
]
|
||||
# while we're waiting for the node w/ interrupt inside to finish
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": " and parallel"},
|
||||
]
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
@@ -5827,42 +5573,23 @@ async def test_nested_graph_interrupts_parallel(
|
||||
app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"])
|
||||
config = {"configurable": {"thread_id": "5"}}
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
{"my_key": ""},
|
||||
config,
|
||||
stream_mode="values",
|
||||
checkpoint_during=checkpoint_during,
|
||||
)
|
||||
c async for c in app.astream({"my_key": ""}, config, stream_mode="values")
|
||||
] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": " and parallel"},
|
||||
]
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
]
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_during", [True, False])
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
async def test_doubly_nested_graph_interrupts(
|
||||
checkpointer_name: str, checkpoint_during: bool
|
||||
) -> None:
|
||||
async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None:
|
||||
class State(TypedDict):
|
||||
my_key: str
|
||||
|
||||
@@ -5915,13 +5642,11 @@ async def test_doubly_nested_graph_interrupts(
|
||||
|
||||
# test invoke w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
assert await app.ainvoke(
|
||||
{"my_key": "my value"}, config, checkpoint_during=checkpoint_during
|
||||
) == {
|
||||
assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == {
|
||||
"my_key": "hi my value",
|
||||
}
|
||||
|
||||
assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == {
|
||||
assert await app.ainvoke(None, config, debug=True) == {
|
||||
"my_key": "hi my value here and there and back again",
|
||||
}
|
||||
|
||||
@@ -5930,22 +5655,12 @@ async def test_doubly_nested_graph_interrupts(
|
||||
config = {
|
||||
"configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append}
|
||||
}
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
{"my_key": "my value"}, config, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream({"my_key": "my value"}, config)] == [
|
||||
{"parent_1": {"my_key": "hi my value"}},
|
||||
{"__interrupt__": ()},
|
||||
]
|
||||
assert nodes == ["parent_1", "grandchild_1"]
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
None, config, checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream(None, config)] == [
|
||||
{"child": {"my_key": "hi my value here and there"}},
|
||||
{"parent_2": {"my_key": "hi my value here and there and back again"}},
|
||||
]
|
||||
@@ -5963,21 +5678,13 @@ async def test_doubly_nested_graph_interrupts(
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
{"my_key": "my value"},
|
||||
config,
|
||||
stream_mode="values",
|
||||
checkpoint_during=checkpoint_during,
|
||||
{"my_key": "my value"}, config, stream_mode="values"
|
||||
)
|
||||
] == [
|
||||
{"my_key": "my value"},
|
||||
{"my_key": "hi my value"},
|
||||
]
|
||||
assert [
|
||||
c
|
||||
async for c in app.astream(
|
||||
None, config, stream_mode="values", checkpoint_during=checkpoint_during
|
||||
)
|
||||
] == [
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": "hi my value"},
|
||||
{"my_key": "hi my value here and there"},
|
||||
{"my_key": "hi my value here and there and back again"},
|
||||
|
||||
@@ -437,17 +437,15 @@ def test_stream():
|
||||
sync_client=mock_sync_client,
|
||||
)
|
||||
|
||||
# test raising graph interrupt if invoked as a subgraph
|
||||
# stream modes doesn't include 'updates'
|
||||
stream_parts = []
|
||||
with pytest.raises(GraphInterrupt) as exc:
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
# pretend we invoked this as a subgraph
|
||||
config={
|
||||
"configurable": {"thread_id": "thread_1", "checkpoint_ns": "some_ns"}
|
||||
},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode="values",
|
||||
):
|
||||
pass
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert exc.value.args[0] == [
|
||||
Interrupt(
|
||||
@@ -458,15 +456,6 @@ def test_stream():
|
||||
)
|
||||
]
|
||||
|
||||
# stream modes doesn't include 'updates'
|
||||
stream_parts = []
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode="values",
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
{"chunk": "data1"},
|
||||
{"chunk": "data2"},
|
||||
@@ -481,62 +470,62 @@ def test_stream():
|
||||
|
||||
# default stream_mode is updates
|
||||
stream_parts = []
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
{"chunk": "data3"},
|
||||
{"chunk": "data4"},
|
||||
{"__interrupt__": ()},
|
||||
]
|
||||
|
||||
# list stream_mode includes mode names
|
||||
stream_parts = []
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
("updates", {"chunk": "data3"}),
|
||||
("updates", {"chunk": "data4"}),
|
||||
("updates", {"__interrupt__": ()}),
|
||||
]
|
||||
|
||||
# subgraphs + list modes
|
||||
stream_parts = []
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
((), "updates", {"chunk": "data3"}),
|
||||
((), "updates", {"chunk": "data4"}),
|
||||
((), "updates", {"__interrupt__": ()}),
|
||||
]
|
||||
|
||||
# subgraphs + single mode
|
||||
stream_parts = []
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
for stream_part in remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
((), {"chunk": "data3"}),
|
||||
((), {"chunk": "data4"}),
|
||||
((), {"__interrupt__": ()}),
|
||||
]
|
||||
|
||||
|
||||
@@ -572,17 +561,15 @@ async def test_astream():
|
||||
client=mock_async_client,
|
||||
)
|
||||
|
||||
# test raising graph interrupt if invoked as a subgraph
|
||||
# stream modes doesn't include 'updates'
|
||||
stream_parts = []
|
||||
with pytest.raises(GraphInterrupt) as exc:
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
# pretend we invoked this as a subgraph
|
||||
config={
|
||||
"configurable": {"thread_id": "thread_1", "checkpoint_ns": "some_ns"}
|
||||
},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode="values",
|
||||
):
|
||||
pass
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert exc.value.args[0] == [
|
||||
Interrupt(
|
||||
@@ -593,15 +580,6 @@ async def test_astream():
|
||||
)
|
||||
]
|
||||
|
||||
# stream modes doesn't include 'updates'
|
||||
stream_parts = []
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode="values",
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
{"chunk": "data1"},
|
||||
{"chunk": "data2"},
|
||||
@@ -618,62 +596,62 @@ async def test_astream():
|
||||
|
||||
# default stream_mode is updates
|
||||
stream_parts = []
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
{"chunk": "data3"},
|
||||
{"chunk": "data4"},
|
||||
{"__interrupt__": ()},
|
||||
]
|
||||
|
||||
# list stream_mode includes mode names
|
||||
stream_parts = []
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
("updates", {"chunk": "data3"}),
|
||||
("updates", {"chunk": "data4"}),
|
||||
("updates", {"__interrupt__": ()}),
|
||||
]
|
||||
|
||||
# subgraphs + list modes
|
||||
stream_parts = []
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
((), "updates", {"chunk": "data3"}),
|
||||
((), "updates", {"chunk": "data4"}),
|
||||
((), "updates", {"__interrupt__": ()}),
|
||||
]
|
||||
|
||||
# subgraphs + single mode
|
||||
stream_parts = []
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
((), {"chunk": "data3"}),
|
||||
((), {"chunk": "data4"}),
|
||||
((), {"__interrupt__": ()}),
|
||||
]
|
||||
|
||||
async_iter = MagicMock()
|
||||
@@ -686,33 +664,33 @@ async def test_astream():
|
||||
|
||||
# subgraphs + list modes
|
||||
stream_parts = []
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
stream_mode=["updates"],
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
(("my", "subgraph"), "updates", {"chunk": "data3"}),
|
||||
(("hello", "subgraph"), "updates", {"chunk": "data4"}),
|
||||
(("bye", "subgraph"), "updates", {"__interrupt__": ()}),
|
||||
]
|
||||
|
||||
# subgraphs + single mode
|
||||
stream_parts = []
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
with pytest.raises(GraphInterrupt):
|
||||
async for stream_part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
config={"configurable": {"thread_id": "thread_1"}},
|
||||
subgraphs=True,
|
||||
):
|
||||
stream_parts.append(stream_part)
|
||||
|
||||
assert stream_parts == [
|
||||
(("my", "subgraph"), {"chunk": "data3"}),
|
||||
(("hello", "subgraph"), {"chunk": "data4"}),
|
||||
(("bye", "subgraph"), {"__interrupt__": ()}),
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -6,10 +6,6 @@ client.cjs
|
||||
client.js
|
||||
client.d.ts
|
||||
client.d.cts
|
||||
auth.cjs
|
||||
auth.js
|
||||
auth.d.ts
|
||||
auth.d.cts
|
||||
react.cjs
|
||||
react.js
|
||||
react.d.ts
|
||||
|
||||
@@ -14,7 +14,6 @@ export const config = {
|
||||
entrypoints: {
|
||||
index: "index",
|
||||
client: "client",
|
||||
auth: "auth/index",
|
||||
react: "react/index",
|
||||
"react-ui": "react-ui/index",
|
||||
"react-ui/server": "react-ui/server/index",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@langchain/langgraph-sdk",
|
||||
"version": "0.0.64",
|
||||
"version": "0.0.62",
|
||||
"description": "Client library for interacting with the LangGraph API",
|
||||
"type": "module",
|
||||
"packageManager": "yarn@1.22.19",
|
||||
@@ -72,15 +72,6 @@
|
||||
"import": "./client.js",
|
||||
"require": "./client.cjs"
|
||||
},
|
||||
"./auth": {
|
||||
"types": {
|
||||
"import": "./auth.d.ts",
|
||||
"require": "./auth.d.cts",
|
||||
"default": "./auth.d.ts"
|
||||
},
|
||||
"import": "./auth.js",
|
||||
"require": "./auth.cjs"
|
||||
},
|
||||
"./react": {
|
||||
"types": {
|
||||
"import": "./react.d.ts",
|
||||
@@ -120,10 +111,6 @@
|
||||
"client.js",
|
||||
"client.d.ts",
|
||||
"client.d.cts",
|
||||
"auth.cjs",
|
||||
"auth.js",
|
||||
"auth.d.ts",
|
||||
"auth.d.cts",
|
||||
"react.cjs",
|
||||
"react.js",
|
||||
"react.d.ts",
|
||||
|
||||
@@ -1,80 +0,0 @@
|
||||
const HTTP_STATUS_MAPPING: { [key: number]: string } = {
|
||||
100: "Continue",
|
||||
101: "Switching Protocols",
|
||||
102: "Processing",
|
||||
103: "Early Hints",
|
||||
200: "OK",
|
||||
201: "Created",
|
||||
202: "Accepted",
|
||||
203: "Non-Authoritative Information",
|
||||
204: "No Content",
|
||||
205: "Reset Content",
|
||||
206: "Partial Content",
|
||||
207: "Multi-Status",
|
||||
208: "Already Reported",
|
||||
226: "IM Used",
|
||||
300: "Multiple Choices",
|
||||
301: "Moved Permanently",
|
||||
302: "Found",
|
||||
303: "See Other",
|
||||
304: "Not Modified",
|
||||
305: "Use Proxy",
|
||||
307: "Temporary Redirect",
|
||||
308: "Permanent Redirect",
|
||||
400: "Bad Request",
|
||||
401: "Unauthorized",
|
||||
402: "Payment Required",
|
||||
403: "Forbidden",
|
||||
404: "Not Found",
|
||||
405: "Method Not Allowed",
|
||||
406: "Not Acceptable",
|
||||
407: "Proxy Authentication Required",
|
||||
408: "Request Timeout",
|
||||
409: "Conflict",
|
||||
410: "Gone",
|
||||
411: "Length Required",
|
||||
412: "Precondition Failed",
|
||||
413: "Request Entity Too Large",
|
||||
414: "Request-URI Too Long",
|
||||
415: "Unsupported Media Type",
|
||||
416: "Requested Range Not Satisfiable",
|
||||
417: "Expectation Failed",
|
||||
418: "I'm a Teapot",
|
||||
421: "Misdirected Request",
|
||||
422: "Unprocessable Entity",
|
||||
423: "Locked",
|
||||
424: "Failed Dependency",
|
||||
425: "Too Early",
|
||||
426: "Upgrade Required",
|
||||
428: "Precondition Required",
|
||||
429: "Too Many Requests",
|
||||
431: "Request Header Fields Too Large",
|
||||
451: "Unavailable For Legal Reasons",
|
||||
500: "Internal Server Error",
|
||||
501: "Not Implemented",
|
||||
502: "Bad Gateway",
|
||||
503: "Service Unavailable",
|
||||
504: "Gateway Timeout",
|
||||
505: "HTTP Version Not Supported",
|
||||
506: "Variant Also Negotiates",
|
||||
507: "Insufficient Storage",
|
||||
508: "Loop Detected",
|
||||
510: "Not Extended",
|
||||
511: "Network Authentication Required",
|
||||
};
|
||||
|
||||
export class HTTPException extends Error {
|
||||
status: number;
|
||||
headers: HeadersInit;
|
||||
|
||||
constructor(
|
||||
status: number,
|
||||
options?: { message?: string; headers?: HeadersInit; cause?: unknown },
|
||||
) {
|
||||
super(options?.message ?? HTTP_STATUS_MAPPING[status] ?? "Unknown error", {
|
||||
cause: options?.cause,
|
||||
});
|
||||
this.status = status;
|
||||
this.headers = options?.headers ?? {};
|
||||
}
|
||||
}
|
||||
@@ -1,36 +0,0 @@
|
||||
import type {
|
||||
AuthenticateCallback,
|
||||
AnyCallback,
|
||||
CallbackEvent,
|
||||
OnCallback,
|
||||
BaseAuthReturn,
|
||||
ToUserLike,
|
||||
BaseUser,
|
||||
} from "./types.js";
|
||||
|
||||
export class Auth<
|
||||
TExtra = {},
|
||||
TAuthReturn extends BaseAuthReturn = BaseAuthReturn,
|
||||
TUser extends BaseUser = ToUserLike<TAuthReturn>,
|
||||
> {
|
||||
"~handlerCache": {
|
||||
authenticate?: AuthenticateCallback<BaseAuthReturn>;
|
||||
callbacks?: Record<string, AnyCallback>;
|
||||
} = {};
|
||||
|
||||
authenticate<T extends BaseAuthReturn>(
|
||||
cb: AuthenticateCallback<T>,
|
||||
): Auth<TExtra, T> {
|
||||
this["~handlerCache"].authenticate = cb;
|
||||
return this as unknown as Auth<TExtra, T>;
|
||||
}
|
||||
|
||||
on<T extends CallbackEvent>(event: T, callback: OnCallback<T, TUser>): this {
|
||||
this["~handlerCache"].callbacks ??= {};
|
||||
this["~handlerCache"].callbacks[event as string] = callback as AnyCallback;
|
||||
return this;
|
||||
}
|
||||
}
|
||||
|
||||
export type { Filters, ResourceActionType } from "./types.js";
|
||||
export { HTTPException } from "./error.js";
|
||||
@@ -1,345 +0,0 @@
|
||||
type Maybe<T> = T | null | undefined;
|
||||
type PromiseMaybe<T> = Promise<T> | T;
|
||||
|
||||
interface AssistantConfig {
|
||||
tags?: Maybe<string[]>;
|
||||
recursion_limit?: Maybe<number>;
|
||||
configurable?: Maybe<{
|
||||
thread_id?: Maybe<string>;
|
||||
thread_ts?: Maybe<string>;
|
||||
[key: string]: unknown;
|
||||
}>;
|
||||
}
|
||||
|
||||
interface AssistantCreate {
|
||||
assistant_id?: Maybe<string>;
|
||||
metadata?: Maybe<Record<string, unknown>>;
|
||||
config?: Maybe<AssistantConfig>;
|
||||
if_exists?: Maybe<"raise" | "do_nothing">;
|
||||
name?: Maybe<string>;
|
||||
graph_id: string;
|
||||
}
|
||||
|
||||
interface AssistantRead {
|
||||
assistant_id: string;
|
||||
metadata?: Maybe<Record<string, unknown>>;
|
||||
}
|
||||
|
||||
interface AssistantUpdate {
|
||||
assistant_id: string;
|
||||
metadata?: Maybe<Record<string, unknown>>;
|
||||
config?: Maybe<AssistantConfig>;
|
||||
graph_id?: Maybe<string>;
|
||||
name?: Maybe<string>;
|
||||
version?: Maybe<number>;
|
||||
}
|
||||
|
||||
interface AssistantDelete {
|
||||
assistant_id: string;
|
||||
}
|
||||
|
||||
interface AssistantSearch {
|
||||
graph_id?: Maybe<string>;
|
||||
metadata?: Maybe<Record<string, unknown>>;
|
||||
limit?: Maybe<number>;
|
||||
offset?: Maybe<number>;
|
||||
}
|
||||
|
||||
interface ThreadCreate {
|
||||
thread_id?: Maybe<string>;
|
||||
metadata?: Maybe<Record<string, unknown>>;
|
||||
if_exists?: Maybe<"raise" | "do_nothing">;
|
||||
}
|
||||
|
||||
interface ThreadRead {
|
||||
thread_id?: Maybe<string>;
|
||||
}
|
||||
|
||||
interface ThreadUpdate {
|
||||
thread_id?: Maybe<string>;
|
||||
metadata?: Maybe<Record<string, unknown>>;
|
||||
action?: Maybe<"interrupt" | "rollback">;
|
||||
}
|
||||
|
||||
interface ThreadDelete {
|
||||
thread_id?: Maybe<string>;
|
||||
run_id?: Maybe<string>;
|
||||
}
|
||||
|
||||
interface ThreadSearch {
|
||||
thread_id?: Maybe<string>;
|
||||
status?: Maybe<"idle" | "busy" | "interrupted" | "error" | (string & {})>;
|
||||
metadata?: Maybe<Record<string, unknown>>;
|
||||
values?: Maybe<Record<string, unknown>>;
|
||||
limit?: Maybe<number>;
|
||||
offset?: Maybe<number>;
|
||||
}
|
||||
|
||||
interface CronCreate {
|
||||
payload?: Maybe<Record<string, unknown>>;
|
||||
schedule: string;
|
||||
cron_id?: Maybe<string>;
|
||||
thread_id?: Maybe<string>;
|
||||
user_id?: Maybe<string>;
|
||||
end_time?: Maybe<string>;
|
||||
}
|
||||
|
||||
interface CronRead {
|
||||
cron_id: string;
|
||||
}
|
||||
|
||||
interface CronUpdate {
|
||||
cron_id: string;
|
||||
payload?: Maybe<Record<string, unknown>>;
|
||||
schedule?: Maybe<string>;
|
||||
}
|
||||
|
||||
interface CronDelete {
|
||||
cron_id: string;
|
||||
}
|
||||
|
||||
interface CronSearch {
|
||||
assistant_id?: Maybe<string>;
|
||||
thread_id?: Maybe<string>;
|
||||
limit?: Maybe<number>;
|
||||
offset?: Maybe<number>;
|
||||
}
|
||||
|
||||
interface StorePut {
|
||||
namespace: string[];
|
||||
key: string;
|
||||
value: Record<string, unknown>;
|
||||
}
|
||||
|
||||
interface StoreGet {
|
||||
namespace: Maybe<string[]>;
|
||||
key: string;
|
||||
}
|
||||
|
||||
interface StoreSearch {
|
||||
namespace?: Maybe<string[]>;
|
||||
filter?: Maybe<Record<string, unknown>>;
|
||||
limit?: Maybe<number>;
|
||||
offset?: Maybe<number>;
|
||||
query?: Maybe<string>;
|
||||
}
|
||||
|
||||
interface StoreListNamespaces {
|
||||
namespace?: Maybe<string[]>;
|
||||
suffix?: Maybe<string[]>;
|
||||
max_depth?: Maybe<number>;
|
||||
limit?: Maybe<number>;
|
||||
offset?: Maybe<number>;
|
||||
}
|
||||
|
||||
interface StoreDelete {
|
||||
namespace?: Maybe<string[]>;
|
||||
key: string;
|
||||
}
|
||||
|
||||
interface RunsCreate {
|
||||
thread_id?: Maybe<string>;
|
||||
assistant_id: string;
|
||||
run_id: string;
|
||||
status: Maybe<
|
||||
"pending" | "running" | "error" | "success" | "timeout" | "interrupted"
|
||||
>;
|
||||
metadata?: Maybe<Record<string, unknown>>;
|
||||
prevent_insert_if_inflight?: Maybe<boolean>;
|
||||
multitask_strategy?: Maybe<"interrupt" | "rollback" | "reject" | "enqueue">;
|
||||
if_not_exists?: Maybe<"reject" | "create">;
|
||||
after_seconds?: Maybe<number>;
|
||||
kwargs: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export interface ResourceActionType {
|
||||
["threads:create"]: ThreadCreate;
|
||||
["threads:read"]: ThreadRead;
|
||||
["threads:update"]: ThreadUpdate;
|
||||
["threads:delete"]: ThreadDelete;
|
||||
["threads:search"]: ThreadSearch;
|
||||
["threads:create_run"]: RunsCreate;
|
||||
|
||||
["assistants:create"]: AssistantCreate;
|
||||
["assistants:read"]: AssistantRead;
|
||||
["assistants:update"]: AssistantUpdate;
|
||||
["assistants:delete"]: AssistantDelete;
|
||||
["assistants:search"]: AssistantSearch;
|
||||
|
||||
["crons:create"]: CronCreate;
|
||||
["crons:read"]: CronRead;
|
||||
["crons:update"]: CronUpdate;
|
||||
["crons:delete"]: CronDelete;
|
||||
["crons:search"]: CronSearch;
|
||||
|
||||
["store:put"]: StorePut;
|
||||
["store:get"]: StoreGet;
|
||||
["store:search"]: StoreSearch;
|
||||
["store:list_namespaces"]: StoreListNamespaces;
|
||||
["store:delete"]: StoreDelete;
|
||||
}
|
||||
interface ResourceType {
|
||||
threads:
|
||||
| "threads:create"
|
||||
| "threads:read"
|
||||
| "threads:update"
|
||||
| "threads:delete"
|
||||
| "threads:search"
|
||||
| "threads:create_run";
|
||||
|
||||
assistants:
|
||||
| "assistants:create"
|
||||
| "assistants:read"
|
||||
| "assistants:update"
|
||||
| "assistants:delete"
|
||||
| "assistants:search";
|
||||
crons:
|
||||
| "crons:create"
|
||||
| "crons:read"
|
||||
| "crons:update"
|
||||
| "crons:delete"
|
||||
| "crons:search";
|
||||
|
||||
store:
|
||||
| "store:put"
|
||||
| "store:get"
|
||||
| "store:search"
|
||||
| "store:list_namespaces"
|
||||
| "store:delete";
|
||||
}
|
||||
interface ActionType {
|
||||
"*:create": "threads:create" | "assistants:create" | "crons:create";
|
||||
|
||||
"*:read": "threads:read" | "assistants:read" | "crons:read";
|
||||
|
||||
"*:update": "threads:update" | "assistants:update" | "crons:update";
|
||||
|
||||
"*:delete":
|
||||
| "threads:delete"
|
||||
| "assistants:delete"
|
||||
| "crons:delete"
|
||||
| "store:delete";
|
||||
|
||||
"*:search":
|
||||
| "threads:search"
|
||||
| "assistants:search"
|
||||
| "crons:search"
|
||||
| "store:search";
|
||||
|
||||
"*:create_run": "threads:create_run";
|
||||
|
||||
"*:put": "store:put";
|
||||
|
||||
"*:get": "store:get";
|
||||
|
||||
"*:list_namespaces": "store:list_namespaces";
|
||||
}
|
||||
|
||||
export type BaseAuthReturn =
|
||||
| {
|
||||
is_authenticated?: boolean;
|
||||
display_name?: string;
|
||||
identity: string;
|
||||
permissions: string[];
|
||||
}
|
||||
| string;
|
||||
|
||||
export interface BaseUser {
|
||||
is_authenticated: boolean;
|
||||
display_name: string;
|
||||
identity: string;
|
||||
permissions: string[];
|
||||
}
|
||||
|
||||
export type ToUserLike<T extends BaseAuthReturn> = T extends string
|
||||
? {
|
||||
is_authenticated: boolean;
|
||||
display_name: string;
|
||||
identity: string;
|
||||
permissions: string[];
|
||||
}
|
||||
: Omit<T, "is_authenticated" | "display_name"> & {
|
||||
is_authenticated: boolean;
|
||||
display_name: string;
|
||||
};
|
||||
|
||||
type CallbackParameter<
|
||||
Resource extends string = string,
|
||||
Action extends string = string,
|
||||
Value extends unknown = unknown,
|
||||
TUser extends BaseUser = BaseUser,
|
||||
> = {
|
||||
resource: Resource;
|
||||
action: Action;
|
||||
value: Value;
|
||||
user: TUser;
|
||||
permissions: string[];
|
||||
};
|
||||
|
||||
type ContextMap = {
|
||||
[ActionType in keyof ResourceActionType]: CallbackParameter<
|
||||
ActionType extends `${infer Resource}:${string}` ? Resource : never,
|
||||
ActionType,
|
||||
ResourceActionType[ActionType],
|
||||
BaseUser
|
||||
>;
|
||||
};
|
||||
|
||||
type ActionCallbackParameter<
|
||||
T extends keyof ActionType,
|
||||
TUser extends BaseUser = BaseUser,
|
||||
> = ContextMap[ActionType[T]] & { user: TUser };
|
||||
type AuthCallbackParameter<
|
||||
T extends keyof ResourceActionType,
|
||||
TUser extends BaseUser = BaseUser,
|
||||
> = ContextMap[T] & { user: TUser };
|
||||
type ResourceCallbackParameter<
|
||||
T extends keyof ResourceType,
|
||||
TUser extends BaseUser = BaseUser,
|
||||
> = ContextMap[ResourceType[T]] & { user: TUser };
|
||||
|
||||
export type Filters<TKey extends string | number | symbol> = {
|
||||
[key in TKey]: string | { [op in "$contains" | "$eq"]?: string };
|
||||
};
|
||||
|
||||
export interface AuthenticateCallback<T extends BaseAuthReturn> {
|
||||
(request: Request): PromiseMaybe<T>;
|
||||
}
|
||||
|
||||
type OnKey = keyof ResourceType | keyof ActionType | keyof ResourceActionType;
|
||||
|
||||
type OnSingleParameter<
|
||||
T extends OnKey,
|
||||
TUser extends BaseUser = BaseUser,
|
||||
> = T extends keyof ResourceType
|
||||
? ResourceCallbackParameter<T, TUser>
|
||||
: T extends keyof ActionType
|
||||
? ActionCallbackParameter<T, TUser>
|
||||
: T extends keyof ResourceActionType
|
||||
? AuthCallbackParameter<T, TUser>
|
||||
: never;
|
||||
|
||||
type OnParameter<
|
||||
T extends "*" | OnKey | OnKey[],
|
||||
TUser extends BaseUser = BaseUser,
|
||||
> = T extends OnKey[]
|
||||
? OnSingleParameter<T[number], TUser>
|
||||
: T extends "*"
|
||||
? AuthCallbackParameter<keyof ResourceActionType, TUser>
|
||||
: T extends OnKey
|
||||
? OnSingleParameter<T, TUser>
|
||||
: never;
|
||||
|
||||
export type AnyCallback = (
|
||||
request: CallbackParameter,
|
||||
) => void | boolean | Filters<string>;
|
||||
|
||||
export type CallbackEvent = "*" | OnKey | OnKey[];
|
||||
|
||||
export type OnCallback<
|
||||
T extends CallbackEvent,
|
||||
TUser extends BaseUser = BaseUser,
|
||||
TMetadata extends Record<string, unknown> = Record<string, unknown>,
|
||||
> = (
|
||||
request: OnParameter<T, TUser>,
|
||||
) => void | boolean | Filters<keyof TMetadata>;
|
||||
@@ -340,7 +340,6 @@ export class AssistantsClient extends BaseClient {
|
||||
assistantId?: string;
|
||||
ifExists?: OnConflictBehavior;
|
||||
name?: string;
|
||||
description?: string;
|
||||
}): Promise<Assistant> {
|
||||
return this.fetch<Assistant>("/assistants", {
|
||||
method: "POST",
|
||||
@@ -351,7 +350,6 @@ export class AssistantsClient extends BaseClient {
|
||||
assistant_id: payload.assistantId,
|
||||
if_exists: payload.ifExists,
|
||||
name: payload.name,
|
||||
description: payload.description,
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -369,7 +367,6 @@ export class AssistantsClient extends BaseClient {
|
||||
config?: Config;
|
||||
metadata?: Metadata;
|
||||
name?: string;
|
||||
description?: string;
|
||||
},
|
||||
): Promise<Assistant> {
|
||||
return this.fetch<Assistant>(`/assistants/${assistantId}`, {
|
||||
@@ -379,7 +376,6 @@ export class AssistantsClient extends BaseClient {
|
||||
config: payload.config,
|
||||
metadata: payload.metadata,
|
||||
name: payload.name,
|
||||
description: payload.description,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
@@ -113,9 +113,6 @@ export interface AssistantBase {
|
||||
|
||||
/** The name of the assistant */
|
||||
name: string;
|
||||
|
||||
/** The description of the assistant */
|
||||
description?: string;
|
||||
}
|
||||
|
||||
export interface AssistantVersion extends AssistantBase {}
|
||||
|
||||
@@ -2,7 +2,11 @@
|
||||
"extends": "@tsconfig/recommended",
|
||||
"compilerOptions": {
|
||||
"target": "ES2021",
|
||||
"lib": ["ES2021", "ES2022.Object", "ES2022.Error", "DOM"],
|
||||
"lib": [
|
||||
"ES2021",
|
||||
"ES2022.Object",
|
||||
"DOM"
|
||||
],
|
||||
"module": "NodeNext",
|
||||
"moduleResolution": "nodenext",
|
||||
"esModuleInterop": true,
|
||||
@@ -18,14 +22,24 @@
|
||||
"jsx": "react-jsx",
|
||||
"outDir": "dist"
|
||||
},
|
||||
"include": ["src/**/*"],
|
||||
"exclude": ["node_modules", "dist", "coverage"],
|
||||
"include": [
|
||||
"src/**/*"
|
||||
],
|
||||
"exclude": [
|
||||
"node_modules",
|
||||
"dist",
|
||||
"coverage"
|
||||
],
|
||||
"includeVersion": true,
|
||||
"typedocOptions": {
|
||||
"entryPoints": ["src/client.ts"],
|
||||
"entryPoints": [
|
||||
"src/client.ts"
|
||||
],
|
||||
"readme": "none",
|
||||
"out": "docs",
|
||||
"plugin": ["typedoc-plugin-markdown"],
|
||||
"plugin": [
|
||||
"typedoc-plugin-markdown"
|
||||
],
|
||||
"excludePrivate": true,
|
||||
"excludeProtected": true,
|
||||
"excludeExternals": false
|
||||
|
||||
Reference in New Issue
Block a user