dao-ai 0.0.25__py3-none-any.whl → 0.1.2__py3-none-any.whl

dao_ai/memory/core.py

@@ -1,8 +1,6 @@
 from typing import Any
 
-from databricks_langchain import (
-    DatabricksEmbeddings,
-)
+from databricks_langchain import DatabricksEmbeddings
 from langchain_core.embeddings.embeddings import Embeddings
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.memory import InMemorySaver
@@ -27,11 +25,13 @@ class InMemoryStoreManager(StoreManagerBase):
         self.store_model = store_model
 
     def store(self) -> BaseStore:
-        logger.debug("Creating InMemory store")
+        embedding_model: LLMModel = self.store_model.embedding_model
 
-        index: dict[str, Any] = None
+        logger.debug(
+            "Creating in-memory store", embeddings_enabled=embedding_model is not None
+        )
 
-        embedding_model: LLMModel = self.store_model.embedding_model
+        index: dict[str, Any] = None
 
         if embedding_model:
             embeddings: Embeddings = DatabricksEmbeddings(endpoint=embedding_model.name)
@@ -41,6 +41,11 @@ class InMemoryStoreManager(StoreManagerBase):
 
             dims: int = self.store_model.dims
             index = {"dims": dims, "embed": embed_texts}
+            logger.debug(
+                "Store embeddings configured",
+                endpoint=embedding_model.name,
+                dimensions=dims,
+            )
 
         store: BaseStore = InMemoryStore(index=index)
 
@@ -60,26 +65,39 @@ class StoreManager:
 
     @classmethod
     def instance(cls, store_model: StoreModel) -> StoreManagerBase:
-        store_manager: StoreManagerBase = None
-        match store_model.type:
+        store_manager: StoreManagerBase | None = None
+        match store_model.storage_type:
             case StorageType.MEMORY:
                 store_manager = cls.store_managers.get(store_model.name)
                 if store_manager is None:
                     store_manager = InMemoryStoreManager(store_model)
                     cls.store_managers[store_model.name] = store_manager
             case StorageType.POSTGRES:
-                from dao_ai.memory.postgres import PostgresStoreManager
+                # Route based on database configuration: instance_name -> Databricks, host -> Postgres
+                if store_model.database.is_lakebase:
+                    # Databricks Lakebase connection
+                    from dao_ai.memory.databricks import DatabricksStoreManager
 
-                store_manager = cls.store_managers.get(
-                    store_model.database.instance_name
-                )
-                if store_manager is None:
-                    store_manager = PostgresStoreManager(store_model)
-                    cls.store_managers[store_model.database.instance_name] = (
-                        store_manager
+                    store_manager = cls.store_managers.get(
+                        store_model.database.instance_name
                     )
+                    if store_manager is None:
+                        store_manager = DatabricksStoreManager(store_model)
+                        cls.store_managers[store_model.database.instance_name] = (
+                            store_manager
+                        )
+                else:
+                    # Standard PostgreSQL connection
+                    from dao_ai.memory.postgres import PostgresStoreManager
+
+                    # Use database name as key for standard PostgreSQL
+                    cache_key = f"{store_model.database.name}"
+                    store_manager = cls.store_managers.get(cache_key)
+                    if store_manager is None:
+                        store_manager = PostgresStoreManager(store_model)
+                        cls.store_managers[cache_key] = store_manager
             case _:
-                raise ValueError(f"Unknown store type: {store_model.type}")
+                raise ValueError(f"Unknown storage type: {store_model.storage_type}")
 
         return store_manager
 
@@ -89,8 +107,8 @@ class CheckpointManager:
 
     @classmethod
     def instance(cls, checkpointer_model: CheckpointerModel) -> CheckpointManagerBase:
-        checkpointer_manager: CheckpointManagerBase = None
-        match checkpointer_model.type:
+        checkpointer_manager: CheckpointManagerBase | None = None
+        match checkpointer_model.storage_type:
             case StorageType.MEMORY:
                 checkpointer_manager = cls.checkpoint_managers.get(
                     checkpointer_model.name
@@ -103,19 +121,36 @@ class CheckpointManager:
                         checkpointer_manager
                     )
             case StorageType.POSTGRES:
-                from dao_ai.memory.postgres import AsyncPostgresCheckpointerManager
+                # Route based on database configuration: instance_name -> Databricks, host -> Postgres
+                if checkpointer_model.database.is_lakebase:
+                    # Databricks Lakebase connection
+                    from dao_ai.memory.databricks import DatabricksCheckpointerManager
 
-                checkpointer_manager = cls.checkpoint_managers.get(
-                    checkpointer_model.database.instance_name
-                )
-                if checkpointer_manager is None:
-                    checkpointer_manager = AsyncPostgresCheckpointerManager(
-                        checkpointer_model
-                    )
-                    cls.checkpoint_managers[
+                    checkpointer_manager = cls.checkpoint_managers.get(
                         checkpointer_model.database.instance_name
-                    ] = checkpointer_manager
+                    )
+                    if checkpointer_manager is None:
+                        checkpointer_manager = DatabricksCheckpointerManager(
+                            checkpointer_model
+                        )
+                        cls.checkpoint_managers[
+                            checkpointer_model.database.instance_name
+                        ] = checkpointer_manager
+                else:
+                    # Standard PostgreSQL connection
+                    from dao_ai.memory.postgres import AsyncPostgresCheckpointerManager
+
+                    # Use database name as key for standard PostgreSQL
+                    cache_key = f"{checkpointer_model.database.name}"
+                    checkpointer_manager = cls.checkpoint_managers.get(cache_key)
+                    if checkpointer_manager is None:
+                        checkpointer_manager = AsyncPostgresCheckpointerManager(
+                            checkpointer_model
+                        )
+                        cls.checkpoint_managers[cache_key] = checkpointer_manager
             case _:
-                raise ValueError(f"Unknown store type: {checkpointer_model.type}")
+                raise ValueError(
+                    f"Unknown storage type: {checkpointer_model.storage_type}"
+                )
 
         return checkpointer_manager

dao_ai/memory/databricks.py

@@ -0,0 +1,402 @@
+"""
+Databricks-native memory storage implementations.
+
+Provides CheckpointSaver and DatabricksStore implementations using
+Databricks Lakebase for persistent storage, with async support.
+
+See:
+- https://api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html#databricks_langchain.CheckpointSaver
+- https://api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html#databricks_langchain.DatabricksStore
+"""
+
+import asyncio
+from collections.abc import AsyncIterator, Iterable, Sequence
+from functools import partial
+from typing import Any, Literal
+
+from databricks_langchain import (
+    CheckpointSaver as DatabricksCheckpointSaver,
+)
+from databricks_langchain import (
+    DatabricksEmbeddings,
+    DatabricksStore,
+)
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+)
+from langgraph.store.base import BaseStore, Item, Op, Result, SearchItem
+from loguru import logger
+
+from dao_ai.config import (
+    CheckpointerModel,
+    StoreModel,
+)
+from dao_ai.memory.base import (
+    CheckpointManagerBase,
+    StoreManagerBase,
+)
+
+# Type alias for namespace path
+NamespacePath = tuple[str, ...]
+
+# Sentinel for not-provided values
+NOT_PROVIDED = object()
+
+
+class AsyncDatabricksCheckpointSaver(DatabricksCheckpointSaver):
+    """
+    Async wrapper for DatabricksCheckpointSaver.
+
+    Provides async implementations of checkpoint methods by delegating
+    to the sync methods using asyncio.to_thread().
+    """
+
+    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        """Async version of get_tuple."""
+        thread_id = config.get("configurable", {}).get("thread_id", "unknown")
+        logger.trace("Fetching checkpoint", thread_id=thread_id, method="aget_tuple")
+        result = await asyncio.to_thread(self.get_tuple, config)
+        if result:
+            logger.trace("Checkpoint found", thread_id=thread_id)
+        else:
+            logger.trace("No checkpoint found", thread_id=thread_id)
+        return result
+
+    async def aget(self, config: RunnableConfig) -> Checkpoint | None:
+        """Async version of get."""
+        thread_id = config.get("configurable", {}).get("thread_id", "unknown")
+        logger.trace("Fetching checkpoint", thread_id=thread_id, method="aget")
+        result = await asyncio.to_thread(self.get, config)
+        if result:
+            logger.trace("Checkpoint found", thread_id=thread_id)
+        else:
+            logger.trace("No checkpoint found", thread_id=thread_id)
+        return result
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Async version of put."""
+        thread_id = config.get("configurable", {}).get("thread_id", "unknown")
+        checkpoint_id = checkpoint.get("id", "unknown")
+        logger.trace(
+            "Saving checkpoint", checkpoint_id=checkpoint_id, thread_id=thread_id
+        )
+        result = await asyncio.to_thread(
+            self.put, config, checkpoint, metadata, new_versions
+        )
+        logger.trace(
+            "Checkpoint saved", thread_id=thread_id, checkpoint_id=checkpoint_id
+        )
+        return result
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        """Async version of put_writes."""
+        thread_id = config.get("configurable", {}).get("thread_id", "unknown")
+        logger.trace(
+            "Saving checkpoint writes",
+            writes_count=len(writes),
+            thread_id=thread_id,
+            task_id=task_id,
+        )
+        await asyncio.to_thread(self.put_writes, config, writes, task_id, task_path)
+        logger.trace("Checkpoint writes saved", thread_id=thread_id, task_id=task_id)
+
+    async def alist(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        """Async version of list."""
+        thread_id = (
+            config.get("configurable", {}).get("thread_id", "unknown")
+            if config
+            else "all"
+        )
+        logger.trace("Listing checkpoints", thread_id=thread_id, limit=limit)
+        # Get all items from sync iterator in a thread
+        items = await asyncio.to_thread(
+            lambda: list(self.list(config, filter=filter, before=before, limit=limit))
+        )
+        logger.debug("Checkpoints listed", thread_id=thread_id, count=len(items))
+        for item in items:
+            yield item
+
+    async def adelete_thread(self, thread_id: str) -> None:
+        """Async version of delete_thread."""
+        logger.trace("Deleting thread", thread_id=thread_id)
+        await asyncio.to_thread(self.delete_thread, thread_id)
+        logger.debug("Thread deleted", thread_id=thread_id)
+
+
+class AsyncDatabricksStore(DatabricksStore):
+    """
+    Async wrapper for DatabricksStore.
+
+    Provides async implementations of store methods by delegating
+    to the sync methods using asyncio.to_thread().
+    """
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        """Async version of batch."""
+        ops_list = list(ops)
+        logger.trace("Executing batch operations", operations_count=len(ops_list))
+        result = await asyncio.to_thread(self.batch, ops_list)
+        logger.debug("Batch operations completed", operations_count=len(result))
+        return result
+
+    async def aget(
+        self,
+        namespace: tuple[str, ...],
+        key: str,
+        *,
+        refresh_ttl: bool | None = None,
+    ) -> Item | None:
+        """Async version of get."""
+        ns_str = "/".join(namespace)
+        logger.trace("Fetching store item", key=key, namespace=ns_str)
+        result = await asyncio.to_thread(
+            partial(self.get, namespace, key, refresh_ttl=refresh_ttl)
+        )
+        if result:
+            logger.trace("Store item found", key=key, namespace=ns_str)
+        else:
+            logger.trace("Store item not found", key=key, namespace=ns_str)
+        return result
+
+    async def aput(
+        self,
+        namespace: tuple[str, ...],
+        key: str,
+        value: dict[str, Any],
+        index: Literal[False] | list[str] | None = None,
+        *,
+        ttl: float | None = None,
+    ) -> None:
+        """Async version of put."""
+        ns_str = "/".join(namespace)
+        logger.trace("Storing item", key=key, namespace=ns_str, has_ttl=ttl is not None)
+        # Handle the ttl parameter - only pass if explicitly provided
+        if ttl is not None:
+            await asyncio.to_thread(
+                partial(self.put, namespace, key, value, index, ttl=ttl)
+            )
+        else:
+            await asyncio.to_thread(partial(self.put, namespace, key, value, index))
+        logger.trace("Item stored", key=key, namespace=ns_str)
+
+    async def adelete(self, namespace: tuple[str, ...], key: str) -> None:
+        """Async version of delete."""
+        ns_str = "/".join(namespace)
+        logger.trace("Deleting item", key=key, namespace=ns_str)
+        await asyncio.to_thread(self.delete, namespace, key)
+        logger.trace("Item deleted", key=key, namespace=ns_str)
+
+    async def asearch(
+        self,
+        namespace_prefix: tuple[str, ...],
+        /,
+        *,
+        query: str | None = None,
+        filter: dict[str, Any] | None = None,
+        limit: int = 10,
+        offset: int = 0,
+        refresh_ttl: bool | None = None,
+    ) -> list[SearchItem]:
+        """Async version of search."""
+        ns_str = "/".join(namespace_prefix)
+        logger.trace(
+            "Searching store", namespace_prefix=ns_str, query=query, limit=limit
+        )
+        result = await asyncio.to_thread(
+            partial(
+                self.search,
+                namespace_prefix,
+                query=query,
+                filter=filter,
+                limit=limit,
+                offset=offset,
+                refresh_ttl=refresh_ttl,
+            )
+        )
+        logger.debug(
+            "Store search completed", namespace_prefix=ns_str, results_count=len(result)
+        )
+        return result
+
+    async def alist_namespaces(
+        self,
+        *,
+        prefix: NamespacePath | None = None,
+        suffix: NamespacePath | None = None,
+        max_depth: int | None = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[tuple[str, ...]]:
+        """Async version of list_namespaces."""
+        prefix_str = "/".join(prefix) if prefix else "all"
+        logger.trace("Listing namespaces", prefix=prefix_str, limit=limit)
+        result = await asyncio.to_thread(
+            partial(
+                self.list_namespaces,
+                prefix=prefix,
+                suffix=suffix,
+                max_depth=max_depth,
+                limit=limit,
+                offset=offset,
+            )
+        )
+        logger.debug("Namespaces listed", count=len(result))
+        return result
+
+
+class DatabricksCheckpointerManager(CheckpointManagerBase):
+    """
+    Checkpointer manager using Databricks CheckpointSaver with async support.
+
+    Uses AsyncDatabricksCheckpointSaver which wraps databricks_langchain.CheckpointSaver
+    with async method implementations for LangGraph async streaming compatibility.
+
+    Required configuration via CheckpointerModel.database:
+    - instance_name: The Databricks Lakebase instance name
+    - workspace_client: WorkspaceClient (supports OBO, service principal, or default auth)
+
+    See: https://api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html#databricks_langchain.CheckpointSaver
+    """
+
+    def __init__(self, checkpointer_model: CheckpointerModel):
+        self.checkpointer_model = checkpointer_model
+        self._checkpointer: BaseCheckpointSaver | None = None
+
+    def checkpointer(self) -> BaseCheckpointSaver:
+        if self._checkpointer is None:
+            database = self.checkpointer_model.database
+            if database is None:
+                raise ValueError(
+                    "Database configuration is required for Databricks checkpointer. "
+                    "Please provide a 'database' field in the checkpointer configuration."
+                )
+
+            instance_name = database.instance_name
+            workspace_client = database.workspace_client
+
+            logger.debug(
+                "Creating Databricks checkpointer", instance_name=instance_name
+            )
+
+            checkpointer = AsyncDatabricksCheckpointSaver(
+                instance_name=instance_name,
+                workspace_client=workspace_client,
+            )
+
+            # Setup the checkpointer (creates necessary tables if needed)
+            logger.debug("Setting up checkpoint tables", instance_name=instance_name)
+            checkpointer.setup()
+            logger.success(
+                "Databricks checkpointer initialized", instance_name=instance_name
+            )
+
+            self._checkpointer = checkpointer
+
+        return self._checkpointer
+
+
+class DatabricksStoreManager(StoreManagerBase):
+    """
+    Store manager using Databricks DatabricksStore with async support.
+
+    Uses AsyncDatabricksStore which wraps databricks_langchain.DatabricksStore
+    with async method implementations for LangGraph async streaming compatibility.
+
+    Required configuration via StoreModel.database:
+    - instance_name: The Databricks Lakebase instance name
+    - workspace_client: WorkspaceClient (supports OBO, service principal, or default auth)
+
+    Optional configuration via StoreModel:
+    - embedding_model: LLMModel for embeddings (will be converted to DatabricksEmbeddings)
+    - dims: Embedding dimensions
+
+    See: https://api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html#databricks_langchain.DatabricksStore
+    """
+
+    def __init__(self, store_model: StoreModel):
+        self.store_model = store_model
+        self._store: BaseStore | None = None
+
+    def store(self) -> BaseStore:
+        if self._store is None:
+            database = self.store_model.database
+            if database is None:
+                raise ValueError(
+                    "Database configuration is required for Databricks store. "
+                    "Please provide a 'database' field in the store configuration."
+                )
+
+            instance_name = database.instance_name
+            workspace_client = database.workspace_client
+
+            # Build embeddings configuration if embedding_model is provided
+            embeddings: DatabricksEmbeddings | None = None
+            embedding_dims: int | None = None
+
+            if self.store_model.embedding_model is not None:
+                embedding_endpoint = self.store_model.embedding_model.name
+                embedding_dims = self.store_model.dims
+
+                logger.debug(
+                    "Configuring store embeddings",
+                    endpoint=embedding_endpoint,
+                    dimensions=embedding_dims,
+                )
+
+                embeddings = DatabricksEmbeddings(endpoint=embedding_endpoint)
+
+            logger.debug(
+                "Creating Databricks store",
+                instance_name=instance_name,
+                embeddings_enabled=embeddings is not None,
+            )
+
+            store = AsyncDatabricksStore(
+                instance_name=instance_name,
+                workspace_client=workspace_client,
+                embeddings=embeddings,
+                embedding_dims=embedding_dims,
+            )
+
+            # Setup the store (creates necessary tables if needed)
+            store.setup()
+            logger.success(
+                "Databricks store initialized",
+                instance_name=instance_name,
+                embeddings_enabled=embeddings is not None,
+            )
+            self._store = store
+
+        return self._store
+
+
+__all__ = [
+    "AsyncDatabricksCheckpointSaver",
+    "AsyncDatabricksStore",
+    "DatabricksCheckpointerManager",
+    "DatabricksStoreManager",
+]