langgraph-api 0.2.44__py3-none-any.whl → 0.2.46__py3-none-any.whl

langgraph_api/__init__.py

@@ -1 +1 @@
-__version__ = "0.2.44"
+__version__ = "0.2.46"

langgraph_api/api/assistants.py

@@ -189,40 +189,40 @@ async def get_assistant_graph(
     async with connect() as conn:
         assistant_ = await Assistants.get(conn, assistant_id)
         assistant = await fetchone(assistant_)
-        config = await ajson_loads(assistant["config"])
-        async with get_graph(
-            assistant["graph_id"],
-            config,
-            checkpointer=Checkpointer(conn),
-            store=(await api_store.get_store()),
-        ) as graph:
-            xray: bool | int = False
-            xray_query = request.query_params.get("xray")
-            if xray_query:
-                if xray_query in ("true", "True"):
-                    xray = True
-                elif xray_query in ("false", "False"):
-                    xray = False
-                else:
-                    try:
-                        xray = int(xray_query)
-                    except ValueError:
-                        raise HTTPException(422, detail="Invalid xray value") from None
-
-                    if xray <= 0:
-                        raise HTTPException(422, detail="Invalid xray value") from None
-
-            if isinstance(graph, BaseRemotePregel):
-                drawable_graph = await graph.fetch_graph(xray=xray)
-                return ApiResponse(drawable_graph.to_json())
+    config = await ajson_loads(assistant["config"])
+    async with get_graph(
+        assistant["graph_id"],
+        config,
+        checkpointer=Checkpointer(),
+        store=(await api_store.get_store()),
+    ) as graph:
+        xray: bool | int = False
+        xray_query = request.query_params.get("xray")
+        if xray_query:
+            if xray_query in ("true", "True"):
+                xray = True
+            elif xray_query in ("false", "False"):
+                xray = False
+            else:
+                try:
+                    xray = int(xray_query)
+                except ValueError:
+                    raise HTTPException(422, detail="Invalid xray value") from None
+
+                if xray <= 0:
+                    raise HTTPException(422, detail="Invalid xray value") from None
+
+        if isinstance(graph, BaseRemotePregel):
+            drawable_graph = await graph.fetch_graph(xray=xray)
+            return ApiResponse(drawable_graph.to_json())
 
-            try:
-                drawable_graph = await graph.aget_graph(xray=xray)
-                return ApiResponse(drawable_graph.to_json())
-            except NotImplementedError:
-                raise HTTPException(
-                    422, detail="The graph does not support visualization"
-                ) from None
+        try:
+            drawable_graph = await graph.aget_graph(xray=xray)
+            return ApiResponse(drawable_graph.to_json())
+        except NotImplementedError:
+            raise HTTPException(
+                422, detail="The graph does not support visualization"
+            ) from None
 
 
 @retry_db
@@ -239,7 +239,7 @@ async def get_assistant_subgraphs(
         async with get_graph(
             assistant["graph_id"],
             config,
-            checkpointer=Checkpointer(conn),
+            checkpointer=Checkpointer(),
             store=(await api_store.get_store()),
         ) as graph:
             namespace = request.path_params.get("namespace")
@@ -285,7 +285,7 @@ async def get_assistant_schemas(
         async with get_graph(
             assistant["graph_id"],
             config,
-            checkpointer=Checkpointer(conn),
+            checkpointer=Checkpointer(),
             store=(await api_store.get_store()),
         ) as graph:
             if isinstance(graph, BaseRemotePregel):

langgraph_api/api/runs.py

@@ -222,6 +222,9 @@ async def wait_run(request: ApiRequest):
         stream = asyncio.create_task(consume())
         while True:
             try:
+                if stream.done():
+                    # raise stream exception if any
+                    stream.result()
                 yield await asyncio.wait_for(last_chunk.wait(), timeout=5)
                 break
             except TimeoutError:
@@ -270,7 +273,10 @@ async def wait_run_stateless(request: ApiRequest):
         vchunk: bytes | None = None
         async with aclosing(
             Runs.Stream.join(
-                run["run_id"], thread_id=run["thread_id"], stream_mode=await sub
+                run["run_id"],
+                thread_id=run["thread_id"],
+                stream_mode=await sub,
+                ignore_404=True,
             )
         ) as stream:
             async for mode, chunk, _ in stream:
@@ -290,6 +296,9 @@ async def wait_run_stateless(request: ApiRequest):
         stream = asyncio.create_task(consume())
         while True:
             try:
+                if stream.done():
+                    # raise stream exception if any
+                    stream.result()
                 yield await asyncio.wait_for(last_chunk.wait(), timeout=5)
                 break
             except TimeoutError:

langgraph_api/js/remote.py

@@ -48,7 +48,6 @@ from langgraph_api.js.sse import SSEDecoder, aiter_lines_raw
 from langgraph_api.route import ApiResponse
 from langgraph_api.schema import Config
 from langgraph_api.serde import json_dumpb
-from langgraph_api.utils import AsyncConnectionProto
 
 logger = structlog.stdlib.get_logger(__name__)
 
@@ -452,27 +451,27 @@ async def run_js_http_process(paths_str: str, http_config: dict, watch: bool = F
                 attempt += 1
 
 
-def _get_passthrough_checkpointer(conn: AsyncConnectionProto):
-    from langgraph_runtime.checkpoint import Checkpointer
+class PassthroughSerialiser(SerializerProtocol):
+    def dumps(self, obj: Any) -> bytes:
+        return json_dumpb(obj)
 
-    class PassthroughSerialiser(SerializerProtocol):
-        def dumps(self, obj: Any) -> bytes:
-            return json_dumpb(obj)
+    def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
+        return "json", json_dumpb(obj)
 
-        def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
-            return "json", json_dumpb(obj)
+    def loads(self, data: bytes) -> Any:
+        return orjson.loads(data)
 
-        def loads(self, data: bytes) -> Any:
-            return orjson.loads(data)
+    def loads_typed(self, data: tuple[str, bytes]) -> Any:
+        type, payload = data
+        if type != "json":
+            raise ValueError(f"Unsupported type {type}")
+        return orjson.loads(payload)
 
-        def loads_typed(self, data: tuple[str, bytes]) -> Any:
-            type, payload = data
-            if type != "json":
-                raise ValueError(f"Unsupported type {type}")
-            return orjson.loads(payload)
 
-    checkpointer = Checkpointer(conn)
+def _get_passthrough_checkpointer():
+    from langgraph_runtime.checkpoint import Checkpointer
 
+    checkpointer = Checkpointer()
     # This checkpointer does not attempt to revive LC-objects.
     # Instead, it will pass through the JSON values as-is.
     checkpointer.serde = PassthroughSerialiser()
@@ -487,53 +486,46 @@ async def _get_passthrough_store():
 # Setup a HTTP server on top of CHECKPOINTER_SOCKET unix socket
 # used by `client.mts` to communicate with the Python checkpointer
 async def run_remote_checkpointer():
-    from langgraph_runtime.database import connect
-
     async def checkpointer_list(payload: dict):
         """Search checkpoints"""
 
         result = []
-        async with connect() as conn:
-            checkpointer = _get_passthrough_checkpointer(conn)
-            async for item in checkpointer.alist(
-                config=payload.get("config"),
-                limit=int(payload.get("limit") or 10),
-                before=payload.get("before"),
-                filter=payload.get("filter"),
-            ):
-                result.append(item)
+        checkpointer = _get_passthrough_checkpointer()
+        async for item in checkpointer.alist(
+            config=payload.get("config"),
+            limit=int(payload.get("limit") or 10),
+            before=payload.get("before"),
+            filter=payload.get("filter"),
+        ):
+            result.append(item)
 
         return result
 
     async def checkpointer_put(payload: dict):
         """Put the new checkpoint metadata"""
 
-        async with connect() as conn:
-            checkpointer = _get_passthrough_checkpointer(conn)
-            return await checkpointer.aput(
-                payload["config"],
-                payload["checkpoint"],
-                payload["metadata"],
-                payload.get("new_versions", {}),
-            )
+        checkpointer = _get_passthrough_checkpointer()
+        return await checkpointer.aput(
+            payload["config"],
+            payload["checkpoint"],
+            payload["metadata"],
+            payload.get("new_versions", {}),
+        )
 
     async def checkpointer_get_tuple(payload: dict):
         """Get actual checkpoint values (reads)"""
-
-        async with connect() as conn:
-            checkpointer = _get_passthrough_checkpointer(conn)
-            return await checkpointer.aget_tuple(config=payload["config"])
+        checkpointer = _get_passthrough_checkpointer()
+        return await checkpointer.aget_tuple(config=payload["config"])
 
     async def checkpointer_put_writes(payload: dict):
         """Put actual checkpoint values (writes)"""
 
-        async with connect() as conn:
-            checkpointer = _get_passthrough_checkpointer(conn)
-            return await checkpointer.aput_writes(
-                payload["config"],
-                payload["writes"],
-                payload["taskId"],
-            )
+        checkpointer = _get_passthrough_checkpointer()
+        return await checkpointer.aput_writes(
+            payload["config"],
+            payload["writes"],
+            payload["taskId"],
+        )
 
     async def store_batch(payload: dict):
         """Batch operations on the store"""
@@ -687,7 +679,11 @@ async def run_remote_checkpointer():
                 payload = orjson.loads(await request.body())
                 return ApiResponse(await cb(payload))
             except ValueError as exc:
+                await logger.error(exc)
                 return ApiResponse({"error": str(exc)}, status_code=400)
+            except Exception as exc:
+                await logger.error(exc)
+                return ApiResponse({"error": str(exc)}, status_code=500)
 
         return wrapped
 
@@ -800,15 +796,18 @@ async def js_healthcheck():
                 transport=httpx.AsyncHTTPTransport(verify=SSL),
             ) as checkpointer_client,
         ):
+            graph_passed = False
             try:
                 res = await graph_client.get("/ok")
                 res.raise_for_status()
+                graph_passed = True
                 res = await checkpointer_client.get("/ok")
                 res.raise_for_status()
                 return True
             except httpx.HTTPError as exc:
                 logger.warning(
                     "JS healthcheck failed. Either the JS server is not running or the event loop is blocked by a CPU-intensive task.",
+                    graph_passed=graph_passed,
                     error=exc,
                 )
                 raise HTTPException(

langgraph_api/logging.py

@@ -57,6 +57,16 @@ class AddPrefixedEnvVars:
         return event_dict
 
 
+class AddApiVersion:
+    def __call__(
+        self, logger: logging.Logger, method_name: str, event_dict: EventDict
+    ) -> EventDict:
+        from langgraph_api import __version__
+
+        event_dict["langgraph_api_version"] = __version__
+        return event_dict
+
+
 class AddLoggingContext:
     def __call__(
         self, logger: logging.Logger, method_name: str, event_dict: EventDict
@@ -90,11 +100,14 @@ shared_processors = [
     structlog.stdlib.PositionalArgumentsFormatter(),
     structlog.stdlib.ExtraAdder(),
     AddPrefixedEnvVars("LANGSMITH_LANGGRAPH_"),  # injected by docker build
+    AddApiVersion(),
     structlog.processors.TimeStamper(fmt="iso", utc=True),
     structlog.processors.StackInfoRenderer(),
-    structlog.processors.dict_tracebacks
-    if LOG_JSON
-    else structlog.processors.format_exc_info,
+    (
+        structlog.processors.dict_tracebacks
+        if LOG_JSON
+        else structlog.processors.format_exc_info
+    ),
     structlog.processors.UnicodeDecoder(),
     AddLoggingContext(),
 ]

langgraph_api/models/run.py

@@ -29,6 +29,13 @@ from langgraph_api.utils import AsyncConnectionProto, get_auth_ctx
 from langgraph_runtime.ops import Runs, logger
 
 
+class LangSmithTracer(TypedDict, total=False):
+    """Configuration for LangSmith tracing."""
+
+    example_id: str | None
+    project_name: str | None
+
+
 class RunCreateDict(TypedDict):
     """Payload for creating a run."""
 
@@ -87,6 +94,8 @@ class RunCreateDict(TypedDict):
     """Start the run after this many seconds. Defaults to 0."""
     if_not_exists: IfNotExists
     """Create the thread if it doesn't exist. If False, reply with 404."""
+    langsmith_tracer: LangSmithTracer | None
+    """Configuration for additional tracing with LangSmith."""
 
 
 def ensure_ids(
@@ -295,6 +304,9 @@ async def create_valid_run(
         user_id = None
     if not configurable.get("langgraph_request_id"):
         configurable["langgraph_request_id"] = request_id
+    if ls_tracing := payload.get("langsmith_tracer"):
+        configurable["__langsmith_project__"] = ls_tracing.get("project_name")
+        configurable["__langsmith_example_id__"] = ls_tracing.get("example_id")
     if request_start_time:
         configurable["__request_start_time_ms__"] = request_start_time
     after_seconds = payload.get("after_seconds", 0)

langgraph_api/stream.py

@@ -1,5 +1,5 @@
 from collections.abc import AsyncIterator, Callable
-from contextlib import AsyncExitStack, aclosing
+from contextlib import AsyncExitStack, aclosing, asynccontextmanager
 from functools import lru_cache
 from typing import Any, cast
 
@@ -20,6 +20,7 @@ from langgraph.errors import (
     InvalidUpdateError,
 )
 from langgraph.pregel.debug import CheckpointPayload, TaskResultPayload
+from langsmith.utils import get_tracer_project
 from pydantic import ValidationError
 from pydantic.v1 import ValidationError as ValidationErrorLegacy
 
@@ -31,7 +32,6 @@ from langgraph_api.js.base import BaseRemotePregel
 from langgraph_api.metadata import HOST, PLAN, USER_API_URL, incr_nodes
 from langgraph_api.schema import Run, StreamMode
 from langgraph_api.serde import json_dumpb
-from langgraph_api.utils import AsyncConnectionProto
 from langgraph_runtime.checkpoint import Checkpointer
 from langgraph_runtime.ops import Runs
 
@@ -71,9 +71,13 @@ def _preprocess_debug_checkpoint(payload: CheckpointPayload | None) -> dict[str,
     return payload
 
 
+@asynccontextmanager
+async def async_tracing_context(*args, **kwargs):
+    with langsmith.tracing_context(*args, **kwargs):
+        yield
+
+
 async def astream_state(
-    stack: AsyncExitStack,
-    conn: AsyncConnectionProto,
     run: Run,
     attempt: int,
     done: ValueEvent,
@@ -83,7 +87,6 @@ async def astream_state(
 ) -> AnyStream:
     """Stream messages from the runnable."""
     run_id = str(run["run_id"])
-    await stack.enter_async_context(conn.pipeline())
     # extract args from run
     kwargs = run["kwargs"].copy()
     kwargs.pop("webhook", None)
@@ -91,12 +94,13 @@ async def astream_state(
     subgraphs = kwargs.get("subgraphs", False)
     temporary = kwargs.pop("temporary", False)
     config = kwargs.pop("config")
+    stack = AsyncExitStack()
     graph = await stack.enter_async_context(
         get_graph(
             config["configurable"]["graph_id"],
             config,
             store=(await api_store.get_store()),
-            checkpointer=None if temporary else Checkpointer(conn),
+            checkpointer=None if temporary else Checkpointer(),
         )
     )
     input = kwargs.pop("input")
@@ -131,6 +135,25 @@ async def astream_state(
     use_astream_events = "events" in stream_mode or isinstance(graph, BaseRemotePregel)
     # yield metadata chunk
     yield "metadata", {"run_id": run_id, "attempt": attempt}
+
+    #  is a langsmith tracing project is specified, additionally pass that in to tracing context
+    if ls_project := config["configurable"].get("__langsmith_project__"):
+        updates = None
+        if example_id := config["configurable"].get("__langsmith_example_id__"):
+            updates = {"reference_example_id": example_id}
+
+        await stack.enter_async_context(
+            async_tracing_context(
+                replicas=[
+                    (
+                        ls_project,
+                        updates,
+                    ),
+                    (get_tracer_project(), None),
+                ]
+            )
+        )
+
     # stream run
     if use_astream_events:
         async with (
@@ -269,7 +292,7 @@ async def astream_state(
                         yield mode, chunk
                 # --- end shared logic with astream_events ---
     if is_remote_pregel:
-        # increament the remote runs
+        # increment the remote runs
         try:
             nodes_executed = await graph.fetch_nodes_executed()
             incr_nodes(None, incr=nodes_executed)