langgraph-runtime-pg 0.11.1__py3-none-any.whl

langgraph_runtime_pg/lifespan.py

@@ -0,0 +1,198 @@
+"""Starlette lifespan: PG pool, Redis, graphs, and queue."""
+
+from __future__ import annotations
+
+import asyncio
+import signal
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any, cast
+
+import structlog
+from langchain_core.runnables.config import RunnableConfig, var_child_runnable_config
+from langgraph.constants import CONF
+from starlette.applications import Starlette
+
+from langgraph_runtime_pg import queue
+from langgraph_runtime_pg.database import start_pool, stop_pool
+
+logger = structlog.stdlib.get_logger(__name__)
+
+_LAST_LIFESPAN_ERROR: BaseException | None = None
+
+
+def get_last_error() -> BaseException | None:
+    return _LAST_LIFESPAN_ERROR
+
+
+@asynccontextmanager
+async def lifespan(
+    app: Starlette | None = None,
+    cancel_event: asyncio.Event | None = None,
+    taskset: set[asyncio.Task] | None = None,
+    **kwargs: Any,
+) -> AsyncIterator[None]:
+    import langgraph_api.config as config
+    from langgraph_api import (
+        __version__,
+        _checkpointer as api_checkpointer,
+        feature_flags,
+        graph,
+        store as api_store,
+    )
+    from langgraph_api.asyncio import SimpleTaskGroup, set_event_loop
+    from langgraph_api.http import (
+        start_http_client,
+        stop_http_client,
+        stop_webhook_http_client,
+    )
+    from langgraph_api.js.ui import start_ui_bundler, stop_ui_bundler
+    from langgraph_api.metadata import metadata_loop
+
+    from langgraph_runtime_pg import __version__ as runtime_version
+
+    await logger.ainfo(
+        f"Starting PG runtime with langgraph-api={__version__} "
+        f"and langgraph-runtime-pg={runtime_version}",
+        version=__version__,
+        runtime_version=runtime_version,
+    )
+    try:
+        current_loop = asyncio.get_running_loop()
+        set_event_loop(current_loop)
+    except RuntimeError:
+        await logger.aerror("Failed to set loop")
+
+    global _LAST_LIFESPAN_ERROR
+    _LAST_LIFESPAN_ERROR = None
+
+    async def _log_graph_load_failure(err: graph.GraphLoadError) -> None:
+        cause = err.__cause__ or err.cause
+        log_fields = err.log_fields()
+        log_fields["action"] = "fix_user_graph"
+        await logger.aerror(
+            f"Graph '{err.spec.id}' failed to load: {err.cause_message}",
+            **log_fields,
+        )
+        await logger.adebug(
+            "Full graph load failure traceback (internal)",
+            **{k: v for k, v in log_fields.items() if k != "user_traceback"},
+            exc_info=cause,
+        )
+
+    started_http = False
+    started_pool = False
+    started_checkpointer = False
+    started_ui = False
+    try:
+        await start_http_client()
+        started_http = True
+        await start_pool()
+        started_pool = True
+        await api_checkpointer.start_checkpointer()
+        started_checkpointer = True
+        await start_ui_bundler()
+        started_ui = True
+
+        async with SimpleTaskGroup(
+            cancel=True,
+            cancel_event=cancel_event,
+            taskgroup_name="Lifespan",
+        ) as tg:
+            tg.create_task(metadata_loop())
+            await api_store.collect_store_from_env()
+            store_instance = await api_store.get_store()
+            if not api_store.CUSTOM_STORE:
+                tg.create_task(store_instance.start_ttl_sweeper())
+            else:
+                await logger.ainfo("Using custom store. Skipping store TTL sweeper.")
+
+            if feature_flags.USE_RUNTIME_CONTEXT_API:
+                from langgraph._internal._constants import (
+                    CONFIG_KEY_RUNTIME,
+                )
+                from langgraph.runtime import Runtime
+
+                langgraph_config = cast(
+                    RunnableConfig,
+                    {CONF: {CONFIG_KEY_RUNTIME: Runtime(store=store_instance)}},
+                )
+            else:
+                from langgraph.constants import CONFIG_KEY_STORE
+
+                langgraph_config = cast(
+                    RunnableConfig,
+                    {CONF: {CONFIG_KEY_STORE: store_instance}},
+                )
+
+            var_child_runnable_config.set(langgraph_config)
+
+            graph.patch_packages_distributions()
+            try:
+                await graph.collect_graphs_from_env(True)
+            except graph.GraphLoadError as exc:
+                _LAST_LIFESPAN_ERROR = exc
+                await _log_graph_load_failure(exc)
+                raise
+            if config.N_JOBS_PER_WORKER > 0:
+                tg.create_task(queue_with_signal())
+
+            from langgraph_api import cron_scheduler
+
+            tg.create_task(cron_scheduler.cron_scheduler())
+
+            yield
+    except graph.GraphLoadError as exc:
+        _LAST_LIFESPAN_ERROR = exc
+        raise
+    except asyncio.CancelledError:
+        pass
+    finally:
+        try:
+            await api_store.exit_store()
+        except Exception:
+            logger.debug("exit_store failed during lifespan cleanup", exc_info=True)
+        if started_checkpointer:
+            try:
+                await api_checkpointer.exit_checkpointer()
+            except Exception:
+                logger.debug("exit_checkpointer failed during lifespan cleanup", exc_info=True)
+        if started_ui:
+            try:
+                await stop_ui_bundler()
+            except Exception:
+                logger.debug("stop_ui_bundler failed during lifespan cleanup", exc_info=True)
+        try:
+            await graph.stop_remote_graphs()
+        except Exception:
+            logger.debug("stop_remote_graphs failed during lifespan cleanup", exc_info=True)
+        if started_http:
+            try:
+                await stop_http_client()
+            except Exception:
+                logger.debug("stop_http_client failed during lifespan cleanup", exc_info=True)
+            try:
+                await stop_webhook_http_client()
+            except Exception:
+                logger.debug(
+                    "stop_webhook_http_client failed during lifespan cleanup",
+                    exc_info=True,
+                )
+        if started_pool:
+            try:
+                await stop_pool()
+            except Exception:
+                logger.debug("stop_pool failed during lifespan cleanup", exc_info=True)
+
+
+async def queue_with_signal() -> None:
+    try:
+        await queue.queue()
+    except asyncio.CancelledError:
+        pass
+    except Exception as exc:
+        logger.exception("Queue failed. Signaling shutdown", exc_info=exc)
+        signal.raise_signal(signal.SIGINT)
+
+
+lifespan.get_last_error = get_last_error  # type: ignore[attr-defined]

langgraph_runtime_pg/metrics.py

@@ -0,0 +1,7 @@
+"""Runtime metrics stub."""
+
+from __future__ import annotations
+
+
+def get_metrics() -> dict[str, dict[str, int]]:
+    return {}

langgraph_runtime_pg/migrate.py

@@ -0,0 +1,101 @@
+"""Alembic migration runner (``langgraph-runtime-pg-migrate``)."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from alembic import command
+from alembic.config import Config
+from alembic.script import ScriptDirectory
+
+from langgraph_runtime_pg.database import get_database_uri, to_psycopg_uri
+
+MIGRATIONS_DIR = Path(__file__).resolve().parent / "migrations"
+
+
+def to_sync_url(uri: str) -> str:
+    """Convert DATABASE_URI to a sync psycopg URL for Alembic."""
+    return "postgresql+psycopg://" + to_psycopg_uri(uri).removeprefix("postgresql://")
+
+
+def alembic_config(database_uri: str | None = None) -> Config:
+    uri = to_sync_url(database_uri or get_database_uri())
+    cfg = Config()
+    cfg.set_main_option("script_location", str(MIGRATIONS_DIR))
+    cfg.set_main_option("sqlalchemy.url", uri)
+    cfg.attributes["configure_logger"] = False
+    return cfg
+
+
+def _head_revision() -> str:
+    heads = ScriptDirectory(str(MIGRATIONS_DIR)).get_heads()
+    return heads[0] if len(heads) == 1 else "head"
+
+
+def upgrade_head(database_uri: str | None = None) -> str:
+    """Apply all pending migrations; return the head revision id."""
+    cfg = alembic_config(database_uri)
+    command.upgrade(cfg, "head")
+    return _head_revision()
+
+
+def stamp_head(database_uri: str | None = None) -> None:
+    """Mark DB at head without running DDL."""
+    command.stamp(alembic_config(database_uri), "head")
+
+
+def current(database_uri: str | None = None) -> None:
+    command.current(alembic_config(database_uri))
+
+
+def history(database_uri: str | None = None) -> None:
+    command.history(alembic_config(database_uri))
+
+
+def downgrade(revision: str, database_uri: str | None = None) -> None:
+    command.downgrade(alembic_config(database_uri), revision)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="langgraph-runtime-pg-migrate",
+        description="Apply langgraph_runtime_pg Postgres schema migrations",
+    )
+    parser.add_argument(
+        "command",
+        choices=("upgrade", "stamp", "current", "history", "downgrade"),
+        help="Migration action",
+    )
+    parser.add_argument(
+        "revision",
+        nargs="?",
+        default="head",
+        help="For downgrade: target revision (default unused for upgrade)",
+    )
+    args = parser.parse_args(argv)
+
+    try:
+        if args.command == "upgrade":
+            rev = upgrade_head()
+            print(f"upgraded to {rev}")
+        elif args.command == "stamp":
+            stamp_head()
+            print("stamped head")
+        elif args.command == "current":
+            current()
+        elif args.command == "history":
+            history()
+        elif args.command == "downgrade":
+            target = args.revision if args.revision != "head" else "-1"
+            downgrade(target)
+            print(f"downgraded to {target}")
+    except Exception as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

langgraph_runtime_pg/migrations/__init__.py

@@ -0,0 +1 @@
+"""Alembic script location for langgraph_runtime_pg."""

langgraph_runtime_pg/migrations/env.py

@@ -0,0 +1,51 @@
+"""Alembic environment for langgraph_runtime_pg schema migrations."""
+
+from __future__ import annotations
+
+from logging.config import fileConfig
+
+from alembic import context
+from sqlalchemy import engine_from_config, pool
+
+from langgraph_runtime_pg.models import Base
+
+config = context.config
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+target_metadata = Base.metadata
+
+
+def run_migrations_offline() -> None:
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+        compare_type=True,
+    )
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section, {}),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+    with connectable.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata,
+            compare_type=True,
+        )
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

langgraph_runtime_pg/migrations/script.py.mako

@@ -0,0 +1,27 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+"""
+
+from __future__ import annotations
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}