25 lines
803 B
Python
25 lines
803 B
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import os
|
||
|
|
from collections.abc import Iterator
|
||
|
|
from contextlib import contextmanager
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def open_story_checkpointer() -> Iterator[object]:
|
||
|
|
database_url = os.environ.get("DATABASE_URL", "")
|
||
|
|
if database_url.startswith(("postgres://", "postgresql://")):
|
||
|
|
try:
|
||
|
|
from langgraph.checkpoint.postgres import PostgresSaver
|
||
|
|
except ImportError as exc:
|
||
|
|
raise RuntimeError(
|
||
|
|
"Spark story workflows require langgraph-checkpoint-postgres; install project dependencies"
|
||
|
|
) from exc
|
||
|
|
with PostgresSaver.from_conn_string(database_url) as saver:
|
||
|
|
saver.setup()
|
||
|
|
yield saver
|
||
|
|
return
|
||
|
|
from langgraph.checkpoint.memory import MemorySaver
|
||
|
|
|
||
|
|
yield MemorySaver()
|