langgraph-api 0.0.47__py3-none-any.whl → 0.1.0__py3-none-any.whl

langgraph_storage/checkpoint.py

@@ -1,123 +0,0 @@
-import logging
-import os
-import uuid
-
-from langchain_core.runnables import RunnableConfig
-from langgraph.checkpoint.base import (
-    Checkpoint,
-    CheckpointMetadata,
-    CheckpointTuple,
-    SerializerProtocol,
-)
-from langgraph.checkpoint.memory import MemorySaver, PersistentDict
-
-from langgraph_api.serde import Serializer
-
-logger = logging.getLogger(__name__)
-
-_EXCLUDED_KEYS = {"checkpoint_ns", "checkpoint_id", "run_id", "thread_id"}
-DISABLE_FILE_PERSISTENCE = (
-    os.getenv("LANGGRAPH_DISABLE_FILE_PERSISTENCE", "false").lower() == "true"
-)
-
-
-class InMemorySaver(MemorySaver):
-    def __init__(
-        self,
-        *,
-        serde: SerializerProtocol | None = None,
-    ) -> None:
-        self.filename = os.path.join(".langgraph_api", ".langgraph_checkpoint.")
-        i = 0
-
-        def factory(*args):
-            nonlocal i
-            i += 1
-
-            thisfname = self.filename + str(i) + ".pckl"
-            d = PersistentDict(*args, filename=thisfname)
-            if not os.path.exists(".langgraph_api"):
-                os.mkdir(".langgraph_api")
-            try:
-                d.load()
-            except FileNotFoundError:
-                pass
-            except ModuleNotFoundError:
-                logger.error(
-                    "Unable to load cached data - your code has changed in a way that's incompatible with the cache."
-                    "\nThis usually happens when you've:"
-                    "\n  - Renamed or moved classes"
-                    "\n  - Changed class structures"
-                    "\n  - Pulled updates that modified class definitions in a way that's incompatible with the cache"
-                    "\n\nRemoving invalid cache data stored at path: .langgraph_api"
-                )
-                os.remove(self.filename)
-            except Exception as e:
-                logger.error("Failed to load cached data: %s", str(e))
-                os.remove(self.filename)
-            return d
-
-        super().__init__(
-            serde=serde if serde is not None else Serializer(),
-            factory=factory if not DISABLE_FILE_PERSISTENCE else None,
-        )
-
-    def put(
-        self,
-        config: RunnableConfig,
-        checkpoint: Checkpoint,
-        metadata: CheckpointMetadata,
-        new_versions: dict[str, str | int | float],
-    ) -> RunnableConfig:
-        # TODO: Should this be done in OSS as well?
-        metadata = {
-            **{
-                k: v
-                for k, v in config["configurable"].items()
-                if not k.startswith("__") and k not in _EXCLUDED_KEYS
-            },
-            **config.get("metadata", {}),
-            **metadata,
-        }
-        if not isinstance(checkpoint["id"], uuid.UUID):
-            # Avoid type inconsistencies
-            checkpoint = checkpoint.copy()
-            checkpoint["id"] = str(checkpoint["id"])
-        return super().put(config, checkpoint, metadata, new_versions)
-
-    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
-        if isinstance(config["configurable"].get("checkpoint_id"), uuid.UUID):
-            # Avoid type inconsistencies....
-            config = config.copy()
-
-            config["configurable"] = {
-                **config["configurable"],
-                "checkpoint_id": str(config["configurable"]["checkpoint_id"]),
-            }
-        return super().get_tuple(config)
-
-    def clear(self):
-        self.storage.clear()
-        self.writes.clear()
-        for suffix in ["1", "2"]:
-            file_path = f"{self.filename}{suffix}.pckl"
-            if os.path.exists(file_path):
-                os.remove(file_path)
-
-
-MEMORY = InMemorySaver()
-
-
-def Checkpointer(*args, unpack_hook=None, **kwargs):
-    if unpack_hook is not None:
-        saver = InMemorySaver(
-            serde=Serializer(__unpack_ext_hook__=unpack_hook), **kwargs
-        )
-        saver.writes = MEMORY.writes
-        saver.blobs = MEMORY.blobs
-        saver.storage = MEMORY.storage
-        return saver
-    return MEMORY
-
-
-__all__ = ["Checkpointer"]

langgraph_storage/database.py

@@ -1,200 +0,0 @@
-import asyncio
-import os
-import uuid
-from collections import defaultdict
-from collections.abc import AsyncIterator
-from contextlib import asynccontextmanager
-from datetime import datetime
-from typing import Any, NotRequired, TypedDict
-from uuid import UUID
-
-import structlog
-from langgraph.checkpoint.memory import PersistentDict
-
-from langgraph_api import config as langgraph_config
-from langgraph_api.utils import AsyncConnectionProto
-from langgraph_storage import store
-from langgraph_storage.inmem_stream import start_stream, stop_stream
-
-logger = structlog.stdlib.get_logger(__name__)
-
-
-class Assistant(TypedDict):
-    assistant_id: UUID
-    graph_id: str
-    name: str
-    description: str | None
-    created_at: NotRequired[datetime]
-    updated_at: NotRequired[datetime]
-    config: dict[str, Any]
-    metadata: dict[str, Any]
-
-
-class Thread(TypedDict):
-    thread_id: UUID
-    created_at: NotRequired[datetime]
-    updated_at: NotRequired[datetime]
-    metadata: dict[str, Any]
-    status: str
-
-
-class Run(TypedDict):
-    run_id: UUID
-    thread_id: UUID
-    assistant_id: UUID
-    created_at: NotRequired[datetime]
-    updated_at: NotRequired[datetime]
-    metadata: dict[str, Any]
-    status: str
-
-
-class RunEvent(TypedDict):
-    event_id: UUID
-    run_id: UUID
-    received_at: NotRequired[datetime]
-    span_id: UUID
-    event: str
-    name: str
-    tags: list[Any]
-    data: dict[str, Any]
-    metadata: dict[str, Any]
-
-
-class AssistantVersion(TypedDict):
-    assistant_id: UUID
-    version: int
-    graph_id: str
-    config: dict[str, Any]
-    metadata: dict[str, Any]
-    created_at: NotRequired[datetime]
-    name: str
-
-
-class GlobalStore(PersistentDict):
-    def __init__(self, *args: Any, filename: str, **kwargs: Any) -> None:
-        super().__init__(*args, filename=filename, **kwargs)
-        self.clear()
-
-    def clear(self):
-        assistants = self.get("assistants", [])
-        super().clear()
-        self["runs"] = []
-        self["threads"] = []
-        self["assistants"] = [
-            a for a in assistants if a["metadata"].get("created_by") == "system"
-        ]
-        self["assistant_versions"] = []
-
-
-OPS_FILENAME = os.path.join(".langgraph_api", ".langgraph_ops.pckl")
-RETRY_COUNTER_FILENAME = os.path.join(".langgraph_api", ".langgraph_retry_counter.pckl")
-
-
-class InMemoryRetryCounter:
-    def __init__(self):
-        self._counters: dict[uuid.UUID, int] = PersistentDict(
-            int, filename=RETRY_COUNTER_FILENAME
-        )
-        self._locks: dict[uuid.UUID, asyncio.Lock] = defaultdict(asyncio.Lock)
-
-    async def increment(self, run_id: uuid.UUID) -> int:
-        async with self._locks[run_id]:
-            self._counters[run_id] += 1
-            return self._counters[run_id]
-
-    def close(self):
-        self._counters.close()
-
-
-# Global retry counter for in-memory implementation
-GLOBAL_RETRY_COUNTER = InMemoryRetryCounter()
-GLOBAL_STORE = GlobalStore(filename=OPS_FILENAME)
-
-
-class InMemConnectionProto:
-    def __init__(self):
-        self.filename = OPS_FILENAME
-        self.store = GLOBAL_STORE
-        self.retry_counter = GLOBAL_RETRY_COUNTER
-        self.can_execute = False
-
-    @asynccontextmanager
-    async def pipeline(self):
-        yield None
-
-    async def execute(self, query: str, *args, **kwargs):
-        return None
-
-    def clear(self):
-        self.store.clear()
-        keys = list(self.retry_counter._counters)
-        for key in keys:
-            del self.retry_counter._counters[key]
-        keys = list(self.retry_counter._locks)
-        for key in keys:
-            del self.retry_counter._locks[key]
-        if os.path.exists(self.filename):
-            os.remove(self.filename)
-
-
-@asynccontextmanager
-async def connect(*, __test__: bool = False) -> AsyncIterator[AsyncConnectionProto]:
-    yield InMemConnectionProto()
-
-
-async def start_pool() -> None:
-    if store._STORE_CONFIG is None:
-        if langgraph_config.STORE_CONFIG:
-            config_ = langgraph_config.STORE_CONFIG
-            store.set_store_config(config_)
-
-    if not os.path.exists(".langgraph_api"):
-        os.mkdir(".langgraph_api")
-    if os.path.exists(OPS_FILENAME):
-        try:
-            GLOBAL_STORE.load()
-        except ModuleNotFoundError:
-            logger.error(
-                "Unable to load cached data - your code has changed in a way that's incompatible with the cache."
-                "\nThis usually happens when you've:"
-                "\n  - Renamed or moved classes"
-                "\n  - Changed class structures"
-                "\n  - Pulled updates that modified class definitions in a way that's incompatible with the cache"
-                "\n\nRemoving invalid cache data stored at path: .langgraph_api"
-            )
-            await asyncio.to_thread(os.remove, OPS_FILENAME)
-            await asyncio.to_thread(os.remove, RETRY_COUNTER_FILENAME)
-        except Exception as e:
-            logger.error("Failed to load cached data: %s", str(e))
-            await asyncio.to_thread(os.remove, OPS_FILENAME)
-            await asyncio.to_thread(os.remove, RETRY_COUNTER_FILENAME)
-    for k in ["runs", "threads", "assistants", "assistant_versions"]:
-        if not GLOBAL_STORE.get(k):
-            GLOBAL_STORE[k] = []
-    for k in ["crons"]:
-        if not GLOBAL_STORE.get(k):
-            GLOBAL_STORE[k] = {}
-    await start_stream()
-
-
-async def stop_pool() -> None:
-    await asyncio.to_thread(GLOBAL_STORE.close)
-    await asyncio.to_thread(GLOBAL_RETRY_COUNTER.close)
-    from langgraph_storage.checkpoint import Checkpointer
-    from langgraph_storage.store import STORE
-
-    await asyncio.to_thread(STORE.close)
-
-    async with Checkpointer():
-        pass
-    await stop_stream()
-
-
-async def healthcheck() -> None:
-    # What could possibly go wrong?
-    pass
-
-
-def pool_stats() -> dict[str, dict[str, int]]:
-    # TODO??
-    return {}

langgraph_storage/inmem_stream.py

@@ -1,109 +0,0 @@
-import asyncio
-import logging
-from collections import defaultdict
-from dataclasses import dataclass
-from uuid import UUID
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class Message:
-    topic: bytes
-    data: bytes
-
-
-class ContextQueue(asyncio.Queue):
-    """Queue that supports async context manager protocol"""
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: object | None,
-    ) -> None:
-        # Clear the queue
-        while not self.empty():
-            try:
-                self.get_nowait()
-            except asyncio.QueueEmpty:
-                break
-
-
-class StreamManager:
-    def __init__(self):
-        self.queues = defaultdict(list)  # Dict[UUID, List[asyncio.Queue]]
-        self.control_queues = defaultdict(list)
-
-    def get_queues(self, run_id: UUID) -> list[asyncio.Queue]:
-        return self.queues[run_id]
-
-    async def put(self, run_id: UUID, message: Message) -> None:
-        topic = message.topic.decode()
-        if "control" in topic:
-            self.control_queues[run_id].append(message)
-        queues = self.queues.get(run_id, [])
-        coros = [queue.put(message) for queue in queues]
-        results = await asyncio.gather(*coros, return_exceptions=True)
-        for result in results:
-            if isinstance(result, Exception):
-                logger.exception(f"Failed to put message in queue: {result}")
-
-    async def add_queue(self, run_id: UUID) -> asyncio.Queue:
-        queue = ContextQueue()
-        self.queues[run_id].append(queue)
-        for control_msg in self.control_queues[run_id]:
-            try:
-                await queue.put(control_msg)
-            except Exception:
-                logger.exception(
-                    f"Failed to put control message in queue: {control_msg}"
-                )
-
-        return queue
-
-    async def remove_queue(self, run_id: UUID, queue: asyncio.Queue):
-        if run_id in self.queues:
-            self.queues[run_id].remove(queue)
-            if not self.queues[run_id]:
-                del self.queues[run_id]
-
-
-# Global instance
-stream_manager = StreamManager()
-
-
-async def start_stream() -> None:
-    """Initialize the queue system.
-    In this in-memory implementation, we just need to ensure we have a clean StreamManager instance.
-    """
-    global stream_manager
-    stream_manager = StreamManager()
-
-
-async def stop_stream() -> None:
-    """Clean up the queue system.
-    Clear all queues and stored control messages."""
-    global stream_manager
-
-    # Send 'done' message to all active queues before clearing
-    for run_id in list(stream_manager.queues.keys()):
-        control_message = Message(topic=f"run:{run_id}:control".encode(), data=b"done")
-
-        for queue in stream_manager.queues[run_id]:
-            try:
-                await queue.put(control_message)
-            except (Exception, RuntimeError):
-                pass  # Ignore errors during shutdown
-
-    # Clear all stored data
-    stream_manager.queues.clear()
-    stream_manager.control_queues.clear()
-
-
-def get_stream_manager() -> StreamManager:
-    """Get the global stream manager instance."""
-    return stream_manager