Making DocumentDatabase and ChromaDatabase with async functions

Author: mc-dorzo

src/parlant/adapters/db/chroma/database.py

@@ -20,7 +20,8 @@
 import operator
 from pathlib import Path
 from typing import Generic, Optional, Sequence, TypeVar, TypedDict, cast
-from typing_extensions import override
+import chromadb.api
+from typing_extensions import override, Self
 import chromadb
 
 from parlant.core.common import Version
@@ -67,18 +68,15 @@ def __init__(
         dir_path: Path,
         embedder_factory: EmbedderFactory,
     ) -> None:
+        self._dir_path = dir_path
         self._logger = logger
         self._embedder_factory = embedder_factory
 
-        self._chroma_client = chromadb.PersistentClient(str(dir_path))
-        self._collections: dict[str, ChromaCollection[BaseDocument]] = (
-            self._load_chromadb_collections()
-        )
+        self._chroma_client: chromadb.api.ClientAPI
+        self._collections: dict[str, ChromaCollection[BaseDocument]] = {}
 
-    def _load_chromadb_collections(
-        self,
-    ) -> dict[str, ChromaCollection[BaseDocument]]:
-        collections: dict[str, ChromaCollection[BaseDocument]] = {}
+    async def __aenter__(self) -> Self:
+        self._chroma_client = chromadb.PersistentClient(str(self._dir_path))
         for chromadb_collection in self._chroma_client.list_collections():
             embedder_module = importlib.import_module(
                 chromadb_collection.metadata["embedder_module_path"]
@@ -94,7 +92,7 @@ def _load_chromadb_collections(
                 embedding_function=None,
             )
 
-            collections[chromadb_collection.name] = ChromaCollection(
+            self._collections[chromadb_collection.name] = ChromaCollection(
                 logger=self._logger,
                 chromadb_collection=chroma_collection,
                 name=chromadb_collection.name,
@@ -103,9 +101,17 @@ def _load_chromadb_collections(
                 ),
                 embedder=embedder,
             )
-        return collections
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[object],
+    ) -> None:
+        pass
 
-    def create_collection(
+    async def create_collection(
         self,
         name: str,
         schema: type[TDocument],
@@ -133,7 +139,7 @@ def create_collection(
 
         return cast(ChromaCollection[TDocument], self._collections[name])
 
-    def get_collection(
+    async def get_collection(
         self,
         name: str,
     ) -> ChromaCollection[TDocument]:
@@ -142,7 +148,7 @@ def get_collection(
 
         raise ValueError(f'ChromaDB collection "{name}" not found.')
 
-    def get_or_create_collection(
+    async def get_or_create_collection(
         self,
         name: str,
         schema: type[TDocument],
@@ -171,7 +177,7 @@ def get_or_create_collection(
 
         return cast(ChromaCollection[TDocument], self._collections[name])
 
-    def delete_collection(
+    async def delete_collection(
         self,
         name: str,
     ) -> None:

src/parlant/adapters/db/chroma/glossary.py

@@ -16,9 +16,9 @@
 from datetime import datetime, timezone
 from itertools import chain
 from typing import Optional, Sequence
-from typing_extensions import override, TypedDict
+from typing_extensions import override, TypedDict, Self
 
-from parlant.adapters.db.chroma.database import ChromaDatabase
+from parlant.adapters.db.chroma.database import ChromaCollection, ChromaDatabase
 from parlant.core.common import (
     ItemNotFoundError,
     UniqueId,
@@ -49,12 +49,28 @@ def __init__(
         chroma_db: ChromaDatabase,
         embedder_type: type[Embedder],
     ):
-        self._collection = chroma_db.get_or_create_collection(
+        self._embedder_type = embedder_type
+        self._embedder = embedder_type()
+
+        self._chroma_db = chroma_db
+        self._collection: ChromaCollection[_TermDocument]
+
+    async def __aenter__(self) -> Self:
+        self._collection = await self._chroma_db.get_or_create_collection(
             name="glossary",
             schema=_TermDocument,
-            embedder_type=embedder_type,
+            embedder_type=self._embedder_type,
         )
-        self._embedder = embedder_type()
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[object],
+    ) -> None:
+        pass
 
     def _serialize(self, term: Term, term_set: str, content: str) -> _TermDocument:
         return _TermDocument(

src/parlant/adapters/db/json_file.py

@@ -19,7 +19,7 @@
 import operator
 from pathlib import Path
 from typing import Any, Mapping, Optional, Sequence, cast
-from typing_extensions import override
+from typing_extensions import override, Self
 import aiofiles
 
 from parlant.core.persistence.document_database import (
@@ -57,7 +57,7 @@ async def flush(self) -> None:
         async with self._lock:
             await self._flush_unlocked()
 
-    async def __aenter__(self) -> JSONFileDocumentDatabase:
+    async def __aenter__(self) -> Self:
         async with self._lock:
             raw_data = await self._load_data()
 
@@ -122,7 +122,7 @@ async def _save_data(
             await file.write(json_string)
 
     @override
-    def create_collection(
+    async def create_collection(
         self,
         name: str,
         schema: type[TDocument],
@@ -138,7 +138,7 @@ def create_collection(
         return cast(JSONFileDocumentCollection[TDocument], self._collections[name])
 
     @override
-    def get_collection(
+    async def get_collection(
         self,
         name: str,
     ) -> JSONFileDocumentCollection[TDocument]:
@@ -147,7 +147,7 @@ def get_collection(
         raise ValueError(f'Collection "{name}" does not exists')
 
     @override
-    def get_or_create_collection(
+    async def get_or_create_collection(
         self,
         name: str,
         schema: type[TDocument],
@@ -164,7 +164,7 @@ def get_or_create_collection(
         return cast(JSONFileDocumentCollection[TDocument], self._collections[name])
 
     @override
-    def delete_collection(
+    async def delete_collection(
         self,
         name: str,
     ) -> None:

src/parlant/adapters/db/transient.py

@@ -36,7 +36,7 @@ def __init__(self) -> None:
         self._collections: dict[str, _TransientDocumentCollection[BaseDocument]] = {}
 
     @override
-    def create_collection(
+    async def create_collection(
         self,
         name: str,
         schema: type[TDocument],
@@ -52,7 +52,7 @@ def create_collection(
         return cast(_TransientDocumentCollection[TDocument], self._collections[name])
 
     @override
-    def get_collection(
+    async def get_collection(
         self,
         name: str,
     ) -> _TransientDocumentCollection[TDocument]:
@@ -61,7 +61,7 @@ def get_collection(
         raise ValueError(f'Collection "{name}" does not exist')
 
     @override
-    def get_or_create_collection(
+    async def get_or_create_collection(
         self,
         name: str,
         schema: type[TDocument],
@@ -72,13 +72,13 @@ def get_or_create_collection(
         annotations = get_type_hints(schema)
         assert "id" in annotations and annotations["id"] == ObjectId
 
-        return self.create_collection(
+        return await self.create_collection(
             name=name,
             schema=schema,
         )
 
     @override
-    def delete_collection(
+    async def delete_collection(
         self,
         name: str,
     ) -> None:

src/parlant/bin/server.py

@@ -224,19 +224,25 @@ async def setup_container(nlp_service_name: str) -> AsyncIterator[Container]:
         JSONFileDocumentDatabase(LOGGER, PARLANT_HOME_DIR / "services.json")
     )
 
-    c[AgentStore] = AgentDocumentStore(agents_db)
-    c[ContextVariableStore] = ContextVariableDocumentStore(context_variables_db)
-    c[TagStore] = TagDocumentStore(tags_db)
-    c[CustomerStore] = CustomerDocumentStore(customers_db)
-    c[GuidelineStore] = GuidelineDocumentStore(guidelines_db)
-    c[GuidelineToolAssociationStore] = GuidelineToolAssociationDocumentStore(
-        guideline_tool_associations_db
-    )
-    c[GuidelineConnectionStore] = GuidelineConnectionDocumentStore(guideline_connections_db)
-    c[SessionStore] = SessionDocumentStore(sessions_db)
+    c[AgentStore] = await EXIT_STACK.enter_async_context(AgentDocumentStore(agents_db))
+    c[ContextVariableStore] = await EXIT_STACK.enter_async_context(
+        ContextVariableDocumentStore(context_variables_db)
+    )
+    c[TagStore] = await EXIT_STACK.enter_async_context(TagDocumentStore(tags_db))
+    c[CustomerStore] = await EXIT_STACK.enter_async_context(CustomerDocumentStore(customers_db))
+    c[GuidelineStore] = await EXIT_STACK.enter_async_context(GuidelineDocumentStore(guidelines_db))
+    c[GuidelineToolAssociationStore] = await EXIT_STACK.enter_async_context(
+        GuidelineToolAssociationDocumentStore(guideline_tool_associations_db)
+    )
+    c[GuidelineConnectionStore] = await EXIT_STACK.enter_async_context(
+        GuidelineConnectionDocumentStore(guideline_connections_db)
+    )
+    c[SessionStore] = await EXIT_STACK.enter_async_context(SessionDocumentStore(sessions_db))
     c[SessionListener] = PollingSessionListener
 
-    c[EvaluationStore] = EvaluationDocumentStore(evaluations_db)
+    c[EvaluationStore] = await EXIT_STACK.enter_async_context(
+        EvaluationDocumentStore(evaluations_db)
+    )
     c[EvaluationListener] = PollingEvaluationListener
 
     c[EventEmitterFactory] = Singleton(EventPublisherFactory)

src/parlant/core/agents.py

@@ -16,10 +16,11 @@
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import NewType, Optional, Sequence, cast
-from typing_extensions import override, TypedDict
+from typing_extensions import override, TypedDict, Self
 
 from parlant.core.common import ItemNotFoundError, UniqueId, Version, generate_id
 from parlant.core.persistence.document_database import (
+    DocumentCollection,
     DocumentDatabase,
     ObjectId,
 )
@@ -93,10 +94,23 @@ def __init__(
         self,
         database: DocumentDatabase,
     ):
-        self._collection = database.get_or_create_collection(
+        self._database = database
+        self._collection: DocumentCollection[_AgentDocument]
+
+    async def __aenter__(self) -> Self:
+        self._collection = await self._database.get_or_create_collection(
             name="agents",
             schema=_AgentDocument,
         )
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[object],
+    ) -> None:
+        pass
 
     def _serialize(self, agent: Agent) -> _AgentDocument:
         return _AgentDocument(

src/parlant/core/context_variables.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import NewType, Optional, Sequence, cast
-from typing_extensions import TypedDict, override
+from typing_extensions import TypedDict, override, Self
 from datetime import datetime, timezone
 from dataclasses import dataclass
 
@@ -27,6 +27,7 @@
     generate_id,
 )
 from parlant.core.persistence.document_database import (
+    DocumentCollection,
     DocumentDatabase,
     ObjectId,
 )
@@ -158,15 +159,29 @@ class ContextVariableDocumentStore(ContextVariableStore):
     VERSION = Version.from_string("0.1.0")
 
     def __init__(self, database: DocumentDatabase):
-        self._variable_collection = database.get_or_create_collection(
+        self._database = database
+        self._variable_collection: DocumentCollection[_ContextVariableDocument]
+        self._value_collection: DocumentCollection[_ContextVariableValueDocument]
+
+    async def __aenter__(self) -> Self:
+        self._variable_collection = await self._database.get_or_create_collection(
             name="variables",
             schema=_ContextVariableDocument,
         )
 
-        self._value_collection = database.get_or_create_collection(
+        self._value_collection = await self._database.get_or_create_collection(
             name="values",
             schema=_ContextVariableValueDocument,
         )
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[object],
+    ) -> None:
+        pass
 
     def _serialize_context_variable(
         self,