File size: 1,438 Bytes
22dcdfd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
from contextlib import AbstractAsyncContextManager
from langgraph.checkpoint.mongodb.aio import AsyncMongoDBSaver
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from core.settings import DatabaseType, settings
from memory.mongodb import get_mongo_saver
from memory.postgres import get_postgres_saver, get_postgres_store
from memory.sqlite import get_sqlite_saver, get_sqlite_store
def initialize_database() -> AbstractAsyncContextManager[
AsyncSqliteSaver | AsyncPostgresSaver | AsyncMongoDBSaver
]:
"""
Initialize the appropriate database checkpointer based on configuration.
Returns an initialized AsyncCheckpointer instance.
"""
if settings.DATABASE_TYPE == DatabaseType.POSTGRES:
return get_postgres_saver()
if settings.DATABASE_TYPE == DatabaseType.MONGO:
return get_mongo_saver()
else: # Default to SQLite
return get_sqlite_saver()
def initialize_store():
"""
Initialize the appropriate store based on configuration.
Returns an async context manager for the initialized store.
"""
if settings.DATABASE_TYPE == DatabaseType.POSTGRES:
return get_postgres_store()
# TODO: Add Mongo store - https://pypi.org/project/langgraph-store-mongodb/
else: # Default to SQLite
return get_sqlite_store()
__all__ = ["initialize_database", "initialize_store"]
|