langgraph-api 0.0.1__py3-none-any.whl

langgraph_api/lifespan.py

@@ -0,0 +1,41 @@
+from contextlib import asynccontextmanager
+
+from starlette.applications import Starlette
+
+import langgraph_api.config as config
+from langgraph_api.asyncio import SimpleTaskGroup
+from langgraph_api.cron_scheduler import cron_scheduler
+from langgraph_api.graph import collect_graphs_from_env, stop_remote_graphs
+from langgraph_api.http import start_http_client, stop_http_client
+from langgraph_api.metadata import metadata_loop
+from langgraph_api.queue import queue
+from langgraph_license.validation import get_license_status, plus_features_enabled
+from langgraph_storage.database import start_pool, stop_pool
+
+
+@asynccontextmanager
+async def lifespan(app: Starlette):
+    if not await get_license_status():
+        raise ValueError(
+            "License verification failed. Please ensure proper configuration:\n"
+            "- For local development, set a valid LANGSMITH_API_KEY for an account with LangGraph Cloud access "
+            "in the environment defined in your langgraph.json file.\n"
+            "- For production, configure the LANGGRAPH_CLOUD_LICENSE_KEY environment variable "
+            "with your LangGraph Cloud license key.\n"
+            "Review your configuration settings and try again. If issues persist, "
+            "contact support for assistance."
+        )
+    await start_http_client()
+    await start_pool()
+    await collect_graphs_from_env(True)
+    try:
+        async with SimpleTaskGroup(cancel=True) as tg:
+            tg.create_task(metadata_loop())
+            tg.create_task(queue(config.N_JOBS_PER_WORKER, config.BG_JOB_TIMEOUT_SECS))
+            if config.FF_CRONS_ENABLED and plus_features_enabled():
+                tg.create_task(cron_scheduler())
+            yield
+    finally:
+        await stop_remote_graphs()
+        await stop_http_client()
+        await stop_pool()

langgraph_api/logging.py

@@ -0,0 +1,121 @@
+import logging
+import os
+
+import structlog
+from starlette.config import Config
+from structlog.typing import EventDict
+
+from langgraph_api.metadata import append_log
+from langgraph_api.serde import json_dumpb
+
+# env
+
+log_env = Config()
+
+LOG_JSON = log_env("LOG_JSON", cast=bool, default=False)
+LOG_COLOR = log_env("LOG_COLOR", cast=bool, default=True)
+LOG_LEVEL = log_env("LOG_LEVEL", cast=str, default="INFO")
+
+logging.getLogger().setLevel(LOG_LEVEL.upper())
+logging.getLogger("psycopg").setLevel(logging.WARNING)
+
+# custom processors
+
+
+class AddPrefixedEnvVars:
+    def __init__(self, prefix: str) -> None:
+        self.kv = {
+            key.removeprefix(prefix).lower(): value
+            for key, value in os.environ.items()
+            if key.startswith(prefix)
+        }
+
+    def __call__(
+        self, logger: logging.Logger, method_name: str, event_dict: EventDict
+    ) -> EventDict:
+        event_dict.update(self.kv)
+        return event_dict
+
+
+class JSONRenderer:
+    def __call__(
+        self, logger: logging.Logger, method_name: str, event_dict: EventDict
+    ) -> str:
+        """
+        The return type of this depends on the return type of self._dumps.
+        """
+        return json_dumpb(event_dict).decode()
+
+
+LEVELS = logging.getLevelNamesMapping()
+
+
+class TapForMetadata:
+    def __call__(
+        self, logger: logging.Logger, method_name: str, event_dict: EventDict
+    ) -> str:
+        """
+        Tap WARN and above logs for metadata. Exclude user loggers.
+        """
+        if (
+            event_dict["logger"].startswith("langgraph")
+            and LEVELS[event_dict["level"].upper()] > LEVELS["INFO"]
+        ):
+            append_log(event_dict.copy())
+        return event_dict
+
+
+# shared config, for both logging and structlog
+
+shared_processors = [
+    structlog.stdlib.add_logger_name,
+    structlog.stdlib.add_log_level,
+    structlog.stdlib.PositionalArgumentsFormatter(),
+    structlog.stdlib.ExtraAdder(),
+    AddPrefixedEnvVars("LANGSMITH_LANGGRAPH_"),  # injected by docker build
+    structlog.processors.TimeStamper(fmt="iso", utc=True),
+    structlog.processors.StackInfoRenderer(),
+    structlog.processors.format_exc_info,
+    structlog.processors.UnicodeDecoder(),
+]
+
+
+# configure logging, used by logging.json, applied by uvicorn
+
+renderer = (
+    JSONRenderer() if LOG_JSON else structlog.dev.ConsoleRenderer(colors=LOG_COLOR)
+)
+
+
+class Formatter(structlog.stdlib.ProcessorFormatter):
+    def __init__(self, *args, **kwargs) -> None:
+        if len(args) == 3:
+            fmt, datefmt, style = args
+            kwargs["fmt"] = fmt
+            kwargs["datefmt"] = datefmt
+            kwargs["style"] = style
+        else:
+            raise RuntimeError("Invalid number of arguments")
+        super().__init__(
+            processors=[
+                structlog.stdlib.ProcessorFormatter.remove_processors_meta,
+                TapForMetadata(),
+                renderer,
+            ],
+            foreign_pre_chain=shared_processors,
+            **kwargs,
+        )
+
+
+# configure structlog
+
+structlog.configure(
+    processors=[
+        structlog.stdlib.filter_by_level,
+        *shared_processors,
+        structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
+    ],
+    logger_factory=structlog.stdlib.LoggerFactory(),
+    wrapper_class=structlog.stdlib.BoundLogger,
+    cache_logger_on_first_use=True,
+)

langgraph_api/metadata.py

@@ -0,0 +1,101 @@
+import asyncio
+import os
+from datetime import UTC, datetime
+
+import langgraph.version
+import orjson
+import structlog
+
+from langgraph_api.config import LANGGRAPH_CLOUD_LICENSE_KEY, LANGSMITH_API_KEY
+from langgraph_api.http import http_request
+from langgraph_license.validation import plus_features_enabled
+
+logger = structlog.stdlib.get_logger(__name__)
+
+INTERVAL = 300
+REVISION = os.getenv("LANGSMITH_LANGGRAPH_API_REVISION")
+VARIANT = os.getenv("LANGSMITH_LANGGRAPH_API_VARIANT")
+if VARIANT == "cloud":
+    HOST = "saas"
+elif os.getenv("LANGSMITH_HOST_PROJECT_ID"):
+    HOST = "byoc"
+else:
+    HOST = "self-hosted"
+PLAN = "enterprise" if plus_features_enabled() else "developer"
+
+LOGS: list[dict] = []
+RUN_COUNTER = 0
+NODE_COUNTER = 0
+FROM_TIMESTAMP = datetime.now(UTC).isoformat()
+
+
+def incr_runs(*, incr: int = 1) -> None:
+    global RUN_COUNTER
+    RUN_COUNTER += incr
+
+
+def incr_nodes(_, *, incr: int = 1) -> None:
+    global NODE_COUNTER
+    NODE_COUNTER += incr
+
+
+def append_log(log: dict) -> None:
+    if not LANGGRAPH_CLOUD_LICENSE_KEY and not LANGSMITH_API_KEY:
+        return
+
+    global LOGS
+    LOGS.append(log)
+
+
+async def metadata_loop() -> None:
+    if not LANGGRAPH_CLOUD_LICENSE_KEY and not LANGSMITH_API_KEY:
+        return
+
+    logger.info("Starting metadata loop")
+
+    global RUN_COUNTER, NODE_COUNTER, FROM_TIMESTAMP
+    while True:
+        # because we always read and write from coroutines in main thread
+        # we don't need a lock as long as there's no awaits in this block
+        from_timestamp = FROM_TIMESTAMP
+        to_timestamp = datetime.now(UTC).isoformat()
+        nodes = NODE_COUNTER
+        runs = RUN_COUNTER
+        logs = LOGS.copy()
+        LOGS.clear()
+        RUN_COUNTER = 0
+        NODE_COUNTER = 0
+        FROM_TIMESTAMP = to_timestamp
+
+        payload = {
+            "license_key": LANGGRAPH_CLOUD_LICENSE_KEY,
+            "api_key": LANGSMITH_API_KEY,
+            "from_timestamp": from_timestamp,
+            "to_timestamp": to_timestamp,
+            "tags": {
+                "langgraph.python.version": langgraph.version.__version__,
+                "langgraph.platform.revision": REVISION,
+                "langgraph.platform.variant": VARIANT,
+                "langgraph.platform.host": HOST,
+                "langgraph.platform.plan": PLAN,
+            },
+            "measures": {
+                "langgraph.platform.runs": runs,
+                "langgraph.platform.nodes": nodes,
+            },
+            "logs": logs,
+        }
+        try:
+            await http_request(
+                "POST",
+                "https://api.smith.langchain.com/v1/metadata/submit",
+                body=orjson.dumps(payload),
+                headers={"Content-Type": "application/json"},
+            )
+        except Exception as e:
+            # retry on next iteration
+            incr_runs(incr=runs)
+            incr_nodes("", incr=nodes)
+            FROM_TIMESTAMP = from_timestamp
+            logger.warning("Failed to submit metadata", exc_info=e)
+        await asyncio.sleep(INTERVAL)

langgraph_api/models/run.py

@@ -0,0 +1,229 @@
+import asyncio
+import uuid
+from collections.abc import Mapping, Sequence
+from typing import Any, TypedDict
+from uuid import UUID
+
+from langgraph.checkpoint.base.id import uuid6
+from starlette.exceptions import HTTPException
+
+from langgraph_api.graph import get_assistant_id
+from langgraph_api.schema import (
+    All,
+    Config,
+    IfNotExists,
+    MetadataInput,
+    MultitaskStrategy,
+    OnCompletion,
+    Run,
+    RunCommand,
+    StreamMode,
+)
+from langgraph_api.utils import AsyncConnectionProto
+from langgraph_storage.ops import Runs, logger
+
+
+class RunCreateDict(TypedDict):
+    """Payload for creating a run."""
+
+    assistant_id: str
+    """Assistant ID to use for this run."""
+    checkpoint_id: str | None
+    """Checkpoint ID to start from. Defaults to the latest checkpoint."""
+    input: Sequence[dict] | dict[str, Any] | None
+    """Input to the run. Pass null to resume from the current state of the thread."""
+    command: RunCommand | None
+    """One or more commands to update the graph's state and send messages to nodes."""
+    metadata: MetadataInput
+    """Metadata for the run."""
+    config: Config | None
+    """Additional configuration for the run."""
+    webhook: str | None
+    """Webhook to call when the run is complete."""
+
+    interrupt_before: All | list[str] | None
+    """Interrupt execution before entering these nodes."""
+    interrupt_after: All | list[str] | None
+    """Interrupt execution after leaving these nodes."""
+
+    multitask_strategy: MultitaskStrategy
+    """Strategy to handle concurrent runs on the same thread. Only relevant if
+    there is a pending/inflight run on the same thread. One of:
+    - "reject": Reject the new run.
+    - "interrupt": Interrupt the current run, keeping steps completed until now,
+       and start a new one.
+    - "rollback": Cancel and delete the existing run, rolling back the thread to
+      the state before it had started, then start the new run.
+    - "enqueue": Queue up the new run to start after the current run finishes.
+    """
+    on_completion: OnCompletion
+    """What to do when the run completes. One of:
+    - "keep": Keep the thread in the database.
+    - "delete": Delete the thread from the database.
+    """
+    stream_mode: list[StreamMode] | StreamMode
+    """One or more of "values", "messages", "updates" or "events".
+    - "values": Stream the thread state any time it changes.
+    - "messages": Stream chat messages from thread state and calls to chat models, 
+      token-by-token where possible.
+    - "updates": Stream the state updates returned by each node.
+    - "events": Stream all events produced by sub-runs (eg. nodes, LLMs, etc.).
+    - "custom": Stream custom events produced by your nodes.
+    """
+    stream_subgraphs: bool | None
+    """Stream output from subgraphs. By default, streams only the top graph."""
+    feedback_keys: list[str] | None
+    """Pass one or more feedback_keys if you want to request short-lived signed URLs
+    for submitting feedback to LangSmith with this key for this run."""
+    after_seconds: int | None
+    """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."""
+
+
+def ensure_ids(
+    assistant_id: str,
+    thread_id: str | None,
+    payload: RunCreateDict,
+) -> tuple[uuid.UUID, uuid.UUID | None, uuid.UUID | None]:
+    try:
+        results = [UUID(assistant_id)]
+    except ValueError:
+        raise HTTPException(status_code=422, detail="Invalid assistant ID") from None
+    if thread_id:
+        try:
+            results.append(UUID(thread_id))
+        except ValueError:
+            raise HTTPException(status_code=422, detail="Invalid thread ID") from None
+    else:
+        results.append(None)
+    if payload.get("checkpoint_id"):
+        try:
+            results.append(UUID(payload.get("checkpoint_id")))
+        except ValueError:
+            raise HTTPException(
+                status_code=422, detail="Invalid checkpoint ID"
+            ) from None
+    else:
+        results.append(None)
+    return tuple(results)
+
+
+def assign_defaults(
+    payload: RunCreateDict,
+):
+    if payload.get("stream_mode"):
+        stream_mode = (
+            payload["stream_mode"]
+            if isinstance(payload["stream_mode"], list)
+            else [payload["stream_mode"]]
+        )
+    else:
+        stream_mode = ["values"]
+    multitask_strategy = payload.get("multitask_strategy") or "reject"
+    prevent_insert_if_inflight = multitask_strategy == "reject"
+    return stream_mode, multitask_strategy, prevent_insert_if_inflight
+
+
+async def create_valid_run(
+    conn: AsyncConnectionProto,
+    thread_id: str | None,
+    payload: RunCreateDict,
+    user_id: str | None,
+    headers: Mapping[str, str],
+    barrier: asyncio.Barrier | None = None,
+    run_id: UUID | None = None,
+) -> Run:
+    # get assistant_id
+    assistant_id = get_assistant_id(payload["assistant_id"])
+
+    # ensure UUID validity defaults
+    assistant_id, thread_id, checkpoint_id = ensure_ids(
+        assistant_id, thread_id, payload
+    )
+    stream_mode, multitask_strategy, prevent_insert_if_inflight = assign_defaults(
+        payload
+    )
+    temporary = thread_id is None and payload.get("on_completion", "delete") == "delete"
+    run_id = run_id or uuid6()
+
+    # assign custom headers and checkpoint to config
+    config = payload.get("config") or {}
+    config.setdefault("configurable", {})
+    if checkpoint_id:
+        config["configurable"]["checkpoint_id"] = str(checkpoint_id)
+    if checkpoint := payload.get("checkpoint"):
+        config["configurable"].update(checkpoint)
+    for key, value in headers.items():
+        if key.startswith("x-"):
+            if key in (
+                "x-api-key",
+                "x-tenant-id",
+                "x-service-key",
+            ):
+                continue
+            config["configurable"][key] = value
+
+    # try to insert
+    run_coro = Runs.put(
+        conn,
+        assistant_id,
+        {
+            "input": payload.get("input"),
+            "command": payload.get("command"),
+            "config": config,
+            "stream_mode": stream_mode,
+            "interrupt_before": payload.get("interrupt_before"),
+            "interrupt_after": payload.get("interrupt_after"),
+            "webhook": payload.get("webhook"),
+            "feedback_keys": payload.get("feedback_keys"),
+            "temporary": temporary,
+            "subgraphs": payload.get("stream_subgraphs", False),
+        },
+        metadata=payload.get("metadata"),
+        status="pending",
+        user_id=user_id,
+        thread_id=thread_id,
+        run_id=run_id,
+        multitask_strategy=multitask_strategy,
+        prevent_insert_if_inflight=prevent_insert_if_inflight,
+        after_seconds=payload.get("after_seconds", 0),
+        if_not_exists=payload.get("if_not_exists", "reject"),
+    )
+    run_ = await run_coro
+
+    if barrier:
+        await barrier.wait()
+
+    # abort if thread, assistant, etc not found
+    try:
+        first = await anext(run_)
+    except StopAsyncIteration:
+        raise HTTPException(
+            status_code=404, detail="Thread or assistant not found."
+        ) from None
+
+    # handle multitask strategy
+    inflight_runs = [run async for run in run_]
+    if first["run_id"] == run_id:
+        logger.info("Created run", run_id=str(run_id), thread_id=str(thread_id))
+        # inserted, proceed
+        if multitask_strategy in ("interrupt", "rollback") and inflight_runs:
+            try:
+                await Runs.cancel(
+                    conn,
+                    [run["run_id"] for run in inflight_runs],
+                    thread_id=thread_id,
+                    action=multitask_strategy,
+                )
+            except HTTPException:
+                # if we can't find the inflight runs again, we can proceeed
+                pass
+        return first
+    elif multitask_strategy == "reject":
+        raise HTTPException(
+            status_code=409,
+            detail="Thread is already running a task. Wait for it to finish or choose a different multitask strategy.",
+        )
+    else:
+        raise NotImplementedError

langgraph_api/patch.py

@@ -0,0 +1,42 @@
+from typing import Any
+
+from starlette.responses import Response, StreamingResponse
+from starlette.types import Send
+
+"""
+Patch Response.render and StreamingResponse.stream_response
+to recognize bytearrays and memoryviews as bytes-like objects.
+"""
+
+
+def Response_render(self, content: Any) -> bytes:
+    if content is None:
+        return b""
+    if isinstance(content, (bytes, bytearray, memoryview)):  # noqa: UP038
+        return content
+    return content.encode(self.charset)  # type: ignore
+
+
+async def StreamingResponse_stream_response(self, send: Send) -> None:
+    await send(
+        {
+            "type": "http.response.start",
+            "status": self.status_code,
+            "headers": self.raw_headers,
+        }
+    )
+    async for chunk in self.body_iterator:
+        if not isinstance(chunk, (bytes, bytearray, memoryview)):  # noqa: UP038
+            chunk = chunk.encode(self.charset)
+        await send({"type": "http.response.body", "body": chunk, "more_body": True})
+
+    await send({"type": "http.response.body", "body": b"", "more_body": False})
+
+
+# patch StreamingResponse.stream_response
+
+StreamingResponse.stream_response = StreamingResponse_stream_response
+
+# patch Response.render
+
+Response.render = Response_render