langgraph-api 0.0.1__py3-none-any.whl

langgraph_storage/queue.py

@@ -0,0 +1,108 @@
+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_queue() -> 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_queue() -> 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

langgraph_storage/retry.py

@@ -0,0 +1,27 @@
+import asyncio
+import functools
+from collections.abc import Callable
+from typing import ParamSpec, TypeVar
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+RETRIABLE_EXCEPTIONS: tuple[type[BaseException], ...] = ()
+OVERLOADED_EXCEPTIONS: tuple[type[BaseException], ...] = ()
+
+
+def retry_db(func: Callable[P, T]) -> Callable[P, T]:
+    attempts = 3
+
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        for i in range(attempts):
+            if i == attempts - 1:
+                return await func(*args, **kwargs)
+            try:
+                return await func(*args, **kwargs)
+            except RETRIABLE_EXCEPTIONS:
+                await asyncio.sleep(0.01)
+
+    return wrapper

langgraph_storage/store.py

@@ -0,0 +1,28 @@
+import os
+
+from langgraph.checkpoint.memory import PersistentDict
+from langgraph.store.memory import InMemoryStore
+
+
+class DiskBackedInMemStore(InMemoryStore):
+    def __init__(self, *args, filename=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.filename = filename
+        self._data = PersistentDict(dict, filename=self.filename)
+        try:
+            self._data.load()
+        except FileNotFoundError:
+            pass
+
+    def close(self):
+        self._data.close()
+
+
+_STORE_FILE = os.path.join(".langgraph_api", "store.pckl")
+if not os.path.exists(".langgraph_api"):
+    os.mkdir(".langgraph_api")
+STORE = DiskBackedInMemStore(filename=_STORE_FILE)
+
+
+def Store(*args, **kwargs):
+    return STORE

langgraph_storage/ttl_dict.py

@@ -0,0 +1,54 @@
+import asyncio
+
+
+# TODO: Make this class thread-safe
+class TTLDict:
+    def __init__(self, timeout=10):
+        self._store = {}
+        self._timeouts = {}
+        self._tasks = {}  # Track tasks for key expiry
+        self._timeout = timeout
+
+    async def _remove_key_after_timeout(self, key):
+        await asyncio.sleep(self._timeouts[key])
+        if key in self._store:
+            del self._store[key]
+            del self._timeouts[key]
+
+    def set(self, key, value, timeout=None):
+        self._store[key] = value
+        key_timeout = timeout if timeout is not None else self._timeout
+        self._timeouts[key] = key_timeout
+
+        # Cancel the previous task if it exists
+        if key in self._tasks:
+            self._tasks[key].cancel()
+
+        # Schedule the removal of the key
+        self._tasks[key] = asyncio.create_task(self._remove_key_after_timeout(key))
+
+    def get(self, key, default=None):
+        return self._store.get(key, default)
+
+    def __contains__(self, key):
+        return key in self._store
+
+    def remove(self, key):
+        if key in self._store:
+            del self._store[key]
+            del self._timeouts[key]
+            if key in self._tasks:
+                self._tasks[key].cancel()
+                del self._tasks[key]
+
+    def keys(self):
+        return self._store.keys()
+
+    def values(self):
+        return self._store.values()
+
+    def items(self):
+        return self._store.items()
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self._store})"

logging.json

@@ -0,0 +1,22 @@
+{
+  "version": 1,
+  "incremental": false,
+  "disable_existing_loggers": false,
+  "formatters": {
+    "simple": {
+      "class": "langgraph_api.logging.Formatter"
+    }
+  },
+  "handlers": {
+    "console": {
+      "class": "logging.StreamHandler",
+      "formatter": "simple",
+      "stream": "ext://sys.stdout"
+    }
+  },
+  "root": {
+    "handlers": [
+      "console"
+    ]
+  }
+}