langgraph-api 0.2.130__py3-none-any.whl → 0.2.134__py3-none-any.whl

langgraph_api/__init__.py

@@ -1 +1 @@
-__version__ = "0.2.130"
+__version__ = "0.2.134"

langgraph_api/api/assistants.py

@@ -16,9 +16,16 @@ from langgraph_api.feature_flags import USE_RUNTIME_CONTEXT_API
 from langgraph_api.graph import get_assistant_id, get_graph
 from langgraph_api.js.base import BaseRemotePregel
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
+from langgraph_api.schema import ASSISTANT_FIELDS
 from langgraph_api.serde import ajson_loads
-from langgraph_api.utils import fetchone, get_pagination_headers, validate_uuid
+from langgraph_api.utils import (
+    fetchone,
+    get_pagination_headers,
+    validate_select_columns,
+    validate_uuid,
+)
 from langgraph_api.validation import (
+    AssistantCountRequest,
     AssistantCreate,
     AssistantPatch,
     AssistantSearchRequest,
@@ -61,7 +68,8 @@ def _get_configurable_jsonschema(graph: Pregel) -> dict:
     in favor of graph.get_context_jsonschema().
     """
     # Otherwise, use the config_schema method.
-    config_schema = graph.config_schema()
+    # TODO: Remove this when we no longer support langgraph < 0.6
+    config_schema = graph.config_schema()  # type: ignore[deprecated]
     model_fields = getattr(config_schema, "model_fields", None) or getattr(
         config_schema, "__fields__", None
     )
@@ -87,11 +95,11 @@ def _state_jsonschema(graph: Pregel) -> dict | None:
     for k in graph.stream_channels_list:
         v = graph.channels[k]
         try:
-            create_model(k, __root__=(v.UpdateType, None)).schema()
+            create_model(k, __root__=(v.UpdateType, None)).model_json_schema()
             fields[k] = (v.UpdateType, None)
         except Exception:
             fields[k] = (Any, None)
-    return create_model(graph.get_name("State"), **fields).schema()
+    return create_model(graph.get_name("State"), **fields).model_json_schema()
 
 
 def _graph_schemas(graph: Pregel) -> dict:
@@ -132,7 +140,7 @@ def _graph_schemas(graph: Pregel) -> dict:
             logger.warning(
                 f"Failed to get context schema for graph {graph.name} with error: `{str(e)}`"
             )
-            context_schema = graph.config_schema()
+            context_schema = graph.config_schema()  # type: ignore[deprecated]
     else:
         context_schema = None
 
@@ -172,6 +180,7 @@ async def search_assistants(
 ) -> ApiResponse:
     """List assistants."""
     payload = await request.json(AssistantSearchRequest)
+    select = validate_select_columns(payload.get("select") or None, ASSISTANT_FIELDS)
     offset = int(payload.get("offset") or 0)
     async with connect() as conn:
         assistants_iter, next_offset = await Assistants.search(
@@ -182,6 +191,7 @@ async def search_assistants(
             offset=offset,
             sort_by=payload.get("sort_by"),
             sort_order=payload.get("sort_order"),
+            select=select,
         )
     assistants, response_headers = await get_pagination_headers(
         assistants_iter, next_offset, offset
@@ -189,6 +199,21 @@ async def search_assistants(
     return ApiResponse(assistants, headers=response_headers)
 
 
+@retry_db
+async def count_assistants(
+    request: ApiRequest,
+) -> ApiResponse:
+    """Count assistants."""
+    payload = await request.json(AssistantCountRequest)
+    async with connect() as conn:
+        count = await Assistants.count(
+            conn,
+            graph_id=payload.get("graph_id"),
+            metadata=payload.get("metadata"),
+        )
+    return ApiResponse(count)
+
+
 @retry_db
 async def get_assistant(
     request: ApiRequest,
@@ -366,7 +391,7 @@ async def patch_assistant(
 
 
 @retry_db
-async def delete_assistant(request: ApiRequest) -> ApiResponse:
+async def delete_assistant(request: ApiRequest) -> Response:
     """Delete an assistant by ID."""
     assistant_id = request.path_params["assistant_id"]
     validate_uuid(assistant_id, "Invalid assistant ID: must be a UUID")
@@ -414,6 +439,7 @@ async def set_latest_assistant_version(request: ApiRequest) -> ApiResponse:
 assistants_routes: list[BaseRoute] = [
     ApiRoute("/assistants", create_assistant, methods=["POST"]),
     ApiRoute("/assistants/search", search_assistants, methods=["POST"]),
+    ApiRoute("/assistants/count", count_assistants, methods=["POST"]),
     ApiRoute(
         "/assistants/{assistant_id}/latest",
         set_latest_assistant_version,

langgraph_api/api/meta.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 import langgraph.version
 from starlette.responses import JSONResponse, PlainTextResponse
 
@@ -43,7 +45,7 @@ async def meta_metrics(request: ApiRequest):
 
     # collect stats
     metrics = get_metrics()
-    worker_metrics = metrics["workers"]
+    worker_metrics = cast(dict[str, int], metrics["workers"])
     workers_max = worker_metrics["max"]
     workers_active = worker_metrics["active"]
     workers_available = worker_metrics["available"]

langgraph_api/api/openapi.py

@@ -25,7 +25,7 @@ def set_custom_spec(spec: dict):
 
 
 @lru_cache(maxsize=1)
-def get_openapi_spec() -> str:
+def get_openapi_spec() -> bytes:
     # patch the graph_id enums
     graph_ids = list(GRAPHS.keys())
     for schema in (

langgraph_api/api/runs.py

@@ -1,9 +1,8 @@
 import asyncio
 from collections.abc import AsyncIterator
-from typing import Literal
+from typing import Literal, cast
 
 import orjson
-from langgraph.checkpoint.base.id import uuid6
 from starlette.exceptions import HTTPException
 from starlette.responses import Response, StreamingResponse
 
@@ -11,9 +10,17 @@ from langgraph_api import config
 from langgraph_api.asyncio import ValueEvent, aclosing
 from langgraph_api.models.run import create_valid_run
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
+from langgraph_api.schema import CRON_FIELDS, RUN_FIELDS
 from langgraph_api.sse import EventSourceResponse
-from langgraph_api.utils import fetchone, get_pagination_headers, validate_uuid
+from langgraph_api.utils import (
+    fetchone,
+    get_pagination_headers,
+    uuid7,
+    validate_select_columns,
+    validate_uuid,
+)
 from langgraph_api.validation import (
+    CronCountRequest,
     CronCreate,
     CronSearch,
     RunBatchCreate,
@@ -92,7 +99,7 @@ async def stream_run(
     thread_id = request.path_params["thread_id"]
     payload = await request.json(RunCreateStateful)
     on_disconnect = payload.get("on_disconnect", "continue")
-    run_id = uuid6()
+    run_id = uuid7()
     sub = asyncio.create_task(Runs.Stream.subscribe(run_id))
 
     try:
@@ -132,7 +139,7 @@ async def stream_run_stateless(
     """Create a stateless run."""
     payload = await request.json(RunCreateStateless)
     on_disconnect = payload.get("on_disconnect", "continue")
-    run_id = uuid6()
+    run_id = uuid7()
     sub = asyncio.create_task(Runs.Stream.subscribe(run_id))
 
     try:
@@ -173,7 +180,7 @@ async def wait_run(request: ApiRequest):
     thread_id = request.path_params["thread_id"]
     payload = await request.json(RunCreateStateful)
     on_disconnect = payload.get("on_disconnect", "continue")
-    run_id = uuid6()
+    run_id = uuid7()
     sub = asyncio.create_task(Runs.Stream.subscribe(run_id))
 
     try:
@@ -255,7 +262,7 @@ async def wait_run_stateless(request: ApiRequest):
     """Create a stateless run, wait for the output."""
     payload = await request.json(RunCreateStateless)
     on_disconnect = payload.get("on_disconnect", "continue")
-    run_id = uuid6()
+    run_id = uuid7()
     sub = asyncio.create_task(Runs.Stream.subscribe(run_id))
 
     try:
@@ -338,6 +345,9 @@ async def list_runs(
     limit = int(request.query_params.get("limit", 10))
     offset = int(request.query_params.get("offset", 0))
     status = request.query_params.get("status")
+    select = validate_select_columns(
+        request.query_params.getlist("select") or None, RUN_FIELDS
+    )
 
     async with connect() as conn, conn.pipeline():
         thread, runs = await asyncio.gather(
@@ -348,6 +358,7 @@ async def list_runs(
                 limit=limit,
                 offset=offset,
                 status=status,
+                select=select,
             ),
         )
     await fetchone(thread)
@@ -425,7 +436,10 @@ async def cancel_run(
     wait_str = request.query_params.get("wait", "false")
     wait = wait_str.lower() in {"true", "yes", "1"}
     action_str = request.query_params.get("action", "interrupt")
-    action = action_str if action_str in {"interrupt", "rollback"} else "interrupt"
+    action = cast(
+        Literal["interrupt", "rollback"],
+        action_str if action_str in {"interrupt", "rollback"} else "interrupt",
+    )
 
     async with connect() as conn:
         await Runs.cancel(
@@ -471,8 +485,9 @@ async def cancel_runs(
         for rid in run_ids:
             validate_uuid(rid, "Invalid run ID: must be a UUID")
     action_str = request.query_params.get("action", "interrupt")
-    action: Literal["interrupt", "rollback"] = (
-        action_str if action_str in ("interrupt", "rollback") else "interrupt"
+    action = cast(
+        Literal["interrupt", "rollback"],
+        action_str if action_str in ("interrupt", "rollback") else "interrupt",
     )
 
     async with connect() as conn:
@@ -557,6 +572,7 @@ async def delete_cron(request: ApiRequest):
 async def search_crons(request: ApiRequest):
     """List all cron jobs for an assistant"""
     payload = await request.json(CronSearch)
+    select = validate_select_columns(payload.get("select") or None, CRON_FIELDS)
     if assistant_id := payload.get("assistant_id"):
         validate_uuid(assistant_id, "Invalid assistant ID: must be a UUID")
     if thread_id := payload.get("thread_id"):
@@ -572,6 +588,7 @@ async def search_crons(request: ApiRequest):
             offset=offset,
             sort_by=payload.get("sort_by"),
             sort_order=payload.get("sort_order"),
+            select=select,
         )
     crons, response_headers = await get_pagination_headers(
         crons_iter, next_offset, offset
@@ -579,6 +596,24 @@ async def search_crons(request: ApiRequest):
     return ApiResponse(crons, headers=response_headers)
 
 
+@retry_db
+async def count_crons(request: ApiRequest):
+    """Count cron jobs."""
+    payload = await request.json(CronCountRequest)
+    if assistant_id := payload.get("assistant_id"):
+        validate_uuid(assistant_id, "Invalid assistant ID: must be a UUID")
+    if thread_id := payload.get("thread_id"):
+        validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+
+    async with connect() as conn:
+        count = await Crons.count(
+            conn,
+            assistant_id=assistant_id,
+            thread_id=thread_id,
+        )
+    return ApiResponse(count)
+
+
 runs_routes = [
     ApiRoute("/runs/stream", stream_run_stateless, methods=["POST"]),
     ApiRoute("/runs/wait", wait_run_stateless, methods=["POST"]),
@@ -595,6 +630,11 @@ runs_routes = [
         if config.FF_CRONS_ENABLED and plus_features_enabled()
         else None
     ),
+    (
+        ApiRoute("/runs/crons/count", count_crons, methods=["POST"])
+        if config.FF_CRONS_ENABLED and plus_features_enabled()
+        else None
+    ),
     ApiRoute("/threads/{thread_id}/runs/{run_id}/join", join_run, methods=["GET"]),
     ApiRoute(
         "/threads/{thread_id}/runs/{run_id}/stream",

langgraph_api/api/threads.py

@@ -5,9 +5,16 @@ from starlette.responses import Response
 from starlette.routing import BaseRoute
 
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
+from langgraph_api.schema import THREAD_FIELDS
 from langgraph_api.state import state_snapshot_to_thread_state
-from langgraph_api.utils import fetchone, get_pagination_headers, validate_uuid
+from langgraph_api.utils import (
+    fetchone,
+    get_pagination_headers,
+    validate_select_columns,
+    validate_uuid,
+)
 from langgraph_api.validation import (
+    ThreadCountRequest,
     ThreadCreate,
     ThreadPatch,
     ThreadSearchRequest,
@@ -58,6 +65,7 @@ async def search_threads(
 ):
     """List threads."""
     payload = await request.json(ThreadSearchRequest)
+    select = validate_select_columns(payload.get("select") or None, THREAD_FIELDS)
     limit = int(payload.get("limit") or 10)
     offset = int(payload.get("offset") or 0)
     async with connect() as conn:
@@ -70,6 +78,7 @@ async def search_threads(
             offset=offset,
             sort_by=payload.get("sort_by"),
             sort_order=payload.get("sort_order"),
+            select=select,
         )
     threads, response_headers = await get_pagination_headers(
         threads_iter, next_offset, offset
@@ -77,6 +86,22 @@ async def search_threads(
     return ApiResponse(threads, headers=response_headers)
 
 
+@retry_db
+async def count_threads(
+    request: ApiRequest,
+):
+    """Count threads."""
+    payload = await request.json(ThreadCountRequest)
+    async with connect() as conn:
+        count = await Threads.count(
+            conn,
+            status=payload.get("status"),
+            values=payload.get("values"),
+            metadata=payload.get("metadata"),
+        )
+    return ApiResponse(count)
+
+
 @retry_db
 async def get_thread_state(
     request: ApiRequest,
@@ -260,6 +285,7 @@ async def copy_thread(request: ApiRequest):
 threads_routes: list[BaseRoute] = [
     ApiRoute("/threads", endpoint=create_thread, methods=["POST"]),
     ApiRoute("/threads/search", endpoint=search_threads, methods=["POST"]),
+    ApiRoute("/threads/count", endpoint=count_threads, methods=["POST"]),
     ApiRoute("/threads/{thread_id}", endpoint=get_thread, methods=["GET"]),
     ApiRoute("/threads/{thread_id}", endpoint=patch_thread, methods=["PATCH"]),
     ApiRoute("/threads/{thread_id}", endpoint=delete_thread, methods=["DELETE"]),

langgraph_api/api/ui.py

@@ -56,6 +56,8 @@ async def handle_ui(request: ApiRequest) -> Response:
 
         # Use http:// protocol if accessing a localhost service
         def is_host(needle: str) -> bool:
+            if not isinstance(host, str):
+                return False
             return host.startswith(needle + ":") or host == needle
 
         protocol = "http:" if is_host("localhost") or is_host("127.0.0.1") else ""

langgraph_api/asgi_transport.py

@@ -13,7 +13,7 @@ from httpx import AsyncByteStream, Request, Response
 if typing.TYPE_CHECKING:  # pragma: no cover
     import asyncio
 
-    import trio
+    import trio  # type: ignore[unresolved-import]
 
     Event = asyncio.Event | trio.Event
 
@@ -37,7 +37,7 @@ def is_running_trio() -> bool:
 
 def create_event() -> Event:
     if is_running_trio():
-        import trio
+        import trio  # type: ignore[unresolved-import]
 
         return trio.Event()
 

langgraph_api/asyncio.py

@@ -119,7 +119,7 @@ def create_task(
 
 def run_coroutine_threadsafe(
     coro: Coroutine[Any, Any, T], ignore_exceptions: tuple[type[Exception], ...] = ()
-) -> concurrent.futures.Future[T | None]:
+) -> concurrent.futures.Future[T] | concurrent.futures.Future[None]:
     if _MAIN_LOOP is None:
         raise RuntimeError("No event loop set")
     future = asyncio.run_coroutine_threadsafe(coro, _MAIN_LOOP)
@@ -226,7 +226,7 @@ def to_aiter(*args: T) -> AsyncIterator[T]:
 V = TypeVar("V")
 
 
-class aclosing(Generic[V], AbstractAsyncContextManager):
+class aclosing(Generic[V], AbstractAsyncContextManager[V]):
     """Async context manager for safely finalizing an asynchronously cleaned-up
     resource such as an async generator, calling its ``aclose()`` method.
 
@@ -255,14 +255,16 @@ class aclosing(Generic[V], AbstractAsyncContextManager):
         await self.thing.aclose()
 
 
-async def aclosing_aiter(aiter: AsyncIterator[T]) -> AsyncIterator[T]:
-    if hasattr(aiter, "__aenter__"):
-        async with aiter:
-            async for item in aiter:
+async def aclosing_aiter(
+    aiterator: AsyncIterator[T],
+) -> AsyncIterator[T]:
+    if hasattr(aiterator, "__aenter__"):
+        async with aiterator:  # type: ignore[invalid-context-manager]
+            async for item in aiterator:
                 yield item
     else:
-        async with aclosing(aiter):
-            async for item in aiter:
+        async with aclosing(aiterator):
+            async for item in aiterator:
                 yield item
 
 

langgraph_api/auth/custom.py

@@ -251,14 +251,15 @@ def _get_auth_instance(path: str | None = None) -> Auth | Literal["js"] | None:
         deps := _get_dependencies(auth_instance._authenticate_handler)
     ):
         auth_instance._authenticate_handler = _solve_fastapi_dependencies(
-            auth_instance._authenticate_handler, deps
+            auth_instance._authenticate_handler,  # type: ignore[invalid-argument-type]
+            deps,
         )
     logger.info(f"Loaded auth instance from path {path}: {auth_instance}")
     return auth_instance
 
 
 def _extract_arguments_from_scope(
-    scope: dict[str, Any],
+    scope: Mapping[str, Any],
     param_names: set[str],
     request: Request | None = None,
     response: Response | None = None,
@@ -283,7 +284,11 @@ def _extract_arguments_from_scope(
     if "path" in param_names:
         args["path"] = scope["path"]
     if "query_params" in param_names:
-        args["query_params"] = QueryParams(scope.get("query_string"))
+        query_params = scope.get("query_string")
+        if query_params:
+            args["query_params"] = QueryParams(query_params)
+        else:
+            args["query_params"] = QueryParams()
     if "headers" in param_names:
         args["headers"] = dict(scope.get("headers", {}))
     if "authorization" in param_names:
@@ -595,7 +600,7 @@ def _load_auth_obj(path: str) -> Auth | Literal["js"]:
                 raise ValueError(f"Could not load file: {module_name}")
             module = importlib.util.module_from_spec(modspec)
             sys.modules[modname] = module
-            modspec.loader.exec_module(module)
+            modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
         else:
             # Load from Python module
             module = importlib.import_module(module_name)

langgraph_api/auth/langsmith/client.py

@@ -14,7 +14,7 @@ from langgraph_api.config import LANGSMITH_AUTH_ENDPOINT
 _client: "JsonHttpClient"
 
 
-def is_retriable_error(exception: Exception) -> bool:
+def is_retriable_error(exception: BaseException) -> bool:
     if isinstance(exception, httpx.TransportError):
         return True
     if isinstance(exception, httpx.HTTPStatusError):

langgraph_api/cli.py

@@ -204,7 +204,7 @@ def run_server(
         mount_prefix = os.environ.get("LANGGRAPH_MOUNT_PREFIX")
     if isinstance(env, str | pathlib.Path):
         try:
-            from dotenv.main import DotEnv
+            from dotenv.main import DotEnv  # type: ignore[unresolved-import]
 
             env_vars = DotEnv(dotenv_path=env).dict() or {}
             logger.debug(f"Loaded environment variables from {env}: {sorted(env_vars)}")
@@ -216,7 +216,7 @@ def run_server(
 
     if debug_port is not None:
         try:
-            import debugpy
+            import debugpy  # type: ignore[unresolved-import]
         except ImportError:
             logger.warning("debugpy is not installed. Debugging will not be available.")
             logger.info("To enable debugging, install debugpy: pip install debugpy")
@@ -301,6 +301,7 @@ def run_server(
         def _open_browser():
             nonlocal studio_origin, full_studio_url
             import time
+            import urllib.error
             import urllib.request
             import webbrowser
             from concurrent.futures import ThreadPoolExecutor
@@ -377,8 +378,8 @@ For production use, please use LangGraph Platform.
             reload=reload,
             env_file=env_file,
             access_log=False,
-            reload_includes=reload_includes,
-            reload_excludes=reload_excludes,
+            reload_includes=list(reload_includes) if reload_includes else None,
+            reload_excludes=list(reload_excludes) if reload_excludes else None,
             log_config={
                 "version": 1,
                 "incremental": False,

langgraph_api/config.py

@@ -287,7 +287,7 @@ if THREAD_TTL is None and CHECKPOINTER_CONFIG is not None:
 N_JOBS_PER_WORKER = env("N_JOBS_PER_WORKER", cast=int, default=10)
 BG_JOB_TIMEOUT_SECS = env("BG_JOB_TIMEOUT_SECS", cast=float, default=3600)
 FF_CRONS_ENABLED = env("FF_CRONS_ENABLED", cast=bool, default=True)
-FF_RICH_THREADS = env("FF_RICH_THREADS", cast=bool, default=False)
+FF_RICH_THREADS = env("FF_RICH_THREADS", cast=bool, default=True)
 
 # auth
 

langgraph_api/executor_entrypoint.py

@@ -0,0 +1,23 @@
+import argparse
+import asyncio
+import json
+import logging.config
+import pathlib
+
+from langgraph_api.queue_entrypoint import main
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--grpc-port", type=int, default=50051)
+    args = parser.parse_args()
+    with open(pathlib.Path(__file__).parent.parent / "logging.json") as file:
+        loaded_config = json.load(file)
+        logging.config.dictConfig(loaded_config)
+    try:
+        import uvloop  # type: ignore[unresolved-import]
+
+        uvloop.install()
+    except ImportError:
+        pass
+    asyncio.run(main(grpc_port=args.grpc_port, entrypoint_name="python-executor"))

langgraph_api/graph.py

@@ -9,7 +9,7 @@ import warnings
 from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from itertools import filterfalse
-from typing import TYPE_CHECKING, Any, NamedTuple, cast
+from typing import TYPE_CHECKING, Any, NamedTuple, TypeGuard, cast
 from uuid import UUID, uuid5
 
 import orjson
@@ -35,10 +35,10 @@ logger = structlog.stdlib.get_logger(__name__)
 
 GraphFactoryFromConfig = Callable[[Config], Pregel | StateGraph]
 GraphFactory = Callable[[], Pregel | StateGraph]
-GraphValue = Pregel | GraphFactory
+GraphValue = Pregel | GraphFactory | GraphFactoryFromConfig
 
 
-GRAPHS: dict[str, Pregel | GraphFactoryFromConfig | GraphFactory] = {}
+GRAPHS: dict[str, GraphValue] = {}
 NAMESPACE_GRAPH = UUID("6ba7b821-9dad-11d1-80b4-00c04fd430c8")
 FACTORY_ACCEPTS_CONFIG: dict[str, bool] = {}
 
@@ -110,11 +110,23 @@ async def _generate_graph(value: Any) -> AsyncIterator[Any]:
         yield value
 
 
-def is_js_graph(graph_id: str) -> bool:
+def is_js_graph(graph_id: str) -> TypeGuard[BaseRemotePregel]:
     """Return whether a graph is a JS graph."""
     return graph_id in GRAPHS and isinstance(GRAPHS[graph_id], BaseRemotePregel)
 
 
+def is_factory(
+    value: GraphValue, graph_id: str
+) -> TypeGuard[GraphFactoryFromConfig | GraphFactory]:
+    return graph_id in FACTORY_ACCEPTS_CONFIG
+
+
+def factory_accepts_config(
+    value: GraphValue, graph_id: str
+) -> TypeGuard[GraphFactoryFromConfig]:
+    return FACTORY_ACCEPTS_CONFIG.get(graph_id, False)
+
+
 @asynccontextmanager
 async def get_graph(
     graph_id: str,
@@ -128,7 +140,7 @@ async def get_graph(
 
     assert_graph_exists(graph_id)
     value = GRAPHS[graph_id]
-    if graph_id in FACTORY_ACCEPTS_CONFIG:
+    if is_factory(value, graph_id):
         config = lg_config.ensure_config(config)
 
         if store is not None:
@@ -139,6 +151,8 @@ async def get_graph(
                 runtime = config["configurable"].get(CONFIG_KEY_RUNTIME)
                 if runtime is None:
                     patched_runtime = Runtime(store=store)
+                elif isinstance(runtime, dict):
+                    patched_runtime = Runtime(**(runtime | {"store": store}))
                 elif runtime.store is None:
                     patched_runtime = cast(Runtime, runtime).override(store=store)
                 else:
@@ -156,7 +170,7 @@ async def get_graph(
         ):
             config["configurable"][CONFIG_KEY_CHECKPOINTER] = checkpointer
         var_child_runnable_config.set(config)
-        value = value(config) if FACTORY_ACCEPTS_CONFIG[graph_id] else value()
+        value = value(config) if factory_accepts_config(value, graph_id) else value()
     try:
         async with _generate_graph(value) as graph_obj:
             if isinstance(graph_obj, StateGraph):
@@ -451,7 +465,7 @@ def _graph_from_spec(spec: GraphSpec) -> GraphValue:
                 raise ValueError(f"Could not find python file for graph: {spec}")
             module = importlib.util.module_from_spec(modspec)
             sys.modules[modname] = module
-            modspec.loader.exec_module(module)
+            modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
         except ImportError as e:
             e.add_note(f"Could not import python module for graph:\n{spec}")
             if config.API_VARIANT == "local_dev":
@@ -565,7 +579,9 @@ def _graph_from_spec(spec: GraphSpec) -> GraphValue:
 @functools.lru_cache(maxsize=1)
 def _get_init_embeddings() -> Callable[[str, ...], "Embeddings"] | None:
     try:
-        from langchain.embeddings import init_embeddings
+        from langchain.embeddings import (  # type: ignore[unresolved-import]
+            init_embeddings,
+        )
 
         return init_embeddings
     except ImportError:
@@ -606,7 +622,7 @@ def resolve_embeddings(index_config: dict) -> "Embeddings":
                     raise ValueError(f"Could not find embeddings file: {module_name}")
                 module = importlib.util.module_from_spec(modspec)
                 sys.modules[modname] = module
-                modspec.loader.exec_module(module)
+                modspec.loader.exec_module(module)  # type: ignore[possibly-unbound-attribute]
             else:
                 # Load from Python module
                 module = importlib.import_module(module_name)