Partial context updates

.gitignore

@@ -3,6 +3,7 @@
 dist/
 venv/
 build/
+dbs/
 docs/source/apiref
 docs/source/_misc
 docs/source/release_notes.rst
@@ -33,3 +34,4 @@ dbs
 benchmarks
 benchmark_results_files.json
 uploaded_benchmarks
+chatsky/utils/logging/*.log

chatsky/__rebuild_pydantic_models__.py

@@ -5,13 +5,18 @@
 from chatsky.core.script import Node
 from chatsky.core.pipeline import Pipeline
 from chatsky.slots.slots import SlotManager
+from chatsky.context_storages import DBContextStorage, MemoryContextStorage
+from chatsky.core.ctx_dict import ContextDict
+from chatsky.context_storages.file import SerializableStorage
 from chatsky.core.context import FrameworkData, ServiceState
 from chatsky.core.service import PipelineComponent
 
+ContextDict.model_rebuild()
 PipelineComponent.model_rebuild()
 Pipeline.model_rebuild()
 Script.model_rebuild()
 Context.model_rebuild()
 ExtraHandlerRuntimeInfo.model_rebuild()
 FrameworkData.model_rebuild()
 ServiceState.model_rebuild()
+SerializableStorage.model_rebuild()

chatsky/conditions/standard.py

@@ -202,7 +202,7 @@ def __init__(
         super().__init__(flow_labels=flow_labels, labels=labels, last_n_indices=last_n_indices)
 
     async def call(self, ctx: Context) -> bool:
-        labels = list(ctx.labels.values())[-self.last_n_indices :]  # noqa: E203
+        labels = await ctx.labels[-self.last_n_indices :]  # noqa: E203
         for label in labels:
             if label.flow_name in self.flow_labels or label in self.labels:
                 return True

chatsky/context_storages/__init__.py

@@ -1,11 +1,10 @@
 # -*- coding: utf-8 -*-
 
-from .database import DBContextStorage, threadsafe_method, context_storage_factory
-from .json import JSONContextStorage, json_available
-from .pickle import PickleContextStorage, pickle_available
+from .database import DBContextStorage, context_storage_factory
+from .file import JSONContextStorage, PickleContextStorage, ShelveContextStorage, json_available, pickle_available
 from .sql import SQLContextStorage, postgres_available, mysql_available, sqlite_available, sqlalchemy_available
 from .ydb import YDBContextStorage, ydb_available
 from .redis import RedisContextStorage, redis_available
+from .memory import MemoryContextStorage
 from .mongo import MongoContextStorage, mongo_available
-from .shelve import ShelveContextStorage
 from .protocol import PROTOCOLS, get_protocol_install_suggestion

chatsky/context_storages/database.py

@@ -8,181 +8,153 @@
 This class implements the basic functionality and can be extended to add additional features as needed.
 """
 
-import asyncio
-import importlib
-import threading
-from functools import wraps
 from abc import ABC, abstractmethod
-from typing import Callable, Hashable, Optional
+from asyncio import Lock
+from importlib import import_module
+from pathlib import Path
+from typing import Any, Callable, Coroutine, Dict, List, Literal, Optional, Set, Tuple, Union
 
 from .protocol import PROTOCOLS
-from chatsky.core import Context
 
+_SUBSCRIPT_TYPE = Union[Literal["__all__"], int, Set[str]]
+_SUBSCRIPT_DICT = Dict[str, Union[_SUBSCRIPT_TYPE, Literal["__none__"]]]
 
-class DBContextStorage(ABC):
-    r"""
-    An abstract interface for `chatsky` DB context storages.
-    It includes the most essential methods of the python `dict` class.
-    Can not be instantiated.
-
-    :param path: Parameter `path` should be set with the URI of the database.
-        It includes a prefix and the required connection credentials.
-        Example: postgresql+asyncpg://user:password@host:port/database
-        In the case of classes that save data to hard drive instead of external databases
-        you need to specify the location of the file, like you do in sqlite.
-        Keep in mind that in Windows you will have to use double backslashes '\\'
-        instead of forward slashes '/' when defining the file path.
-
-    """
 
-    def __init__(self, path: str):
+class DBContextStorage(ABC):
+    _main_table_name: Literal["main"] = "main"
+    _turns_table_name: Literal["turns"] = "turns"
+    _key_column_name: Literal["key"] = "key"
+    _id_column_name: Literal["id"] = "id"
+    _current_turn_id_column_name: Literal["current_turn_id"] = "current_turn_id"
+    _created_at_column_name: Literal["created_at"] = "created_at"
+    _updated_at_column_name: Literal["updated_at"] = "updated_at"
+    _misc_column_name: Literal["misc"] = "misc"
+    _framework_data_column_name: Literal["framework_data"] = "framework_data"
+    _labels_field_name: Literal["labels"] = "labels"
+    _requests_field_name: Literal["requests"] = "requests"
+    _responses_field_name: Literal["responses"] = "responses"
+    _default_subscript_value: int = 3
+
+    def __init__(
+        self,
+        path: str,
+        rewrite_existing: bool = False,
+        configuration: Optional[_SUBSCRIPT_DICT] = None,
+    ):
         _, _, file_path = path.partition("://")
+        configuration = configuration if configuration is not None else dict()
         self.full_path = path
         """Full path to access the context storage, as it was provided by user."""
-        self.path = file_path
-        """`full_path` without a prefix defining db used"""
-        self._lock = threading.Lock()
-        """Threading for methods that require single thread access."""
-
-    def __getitem__(self, key: Hashable) -> Context:
-        """
-        Synchronous method for accessing stored Context.
-
-        :param key: Hashable key used to store Context instance.
-        :return: The stored context, associated with the given key.
-        """
-        return asyncio.run(self.get_item_async(key))
+        self.path = Path(file_path)
+        """`full_path` without a prefix defining db used."""
+        self.rewrite_existing = rewrite_existing
+        """Whether to rewrite existing data in the storage."""
+        self._subscripts = dict()
+        self._sync_lock = Lock()
+        for field in (self._labels_field_name, self._requests_field_name, self._responses_field_name):
+            value = configuration.get(field, self._default_subscript_value)
+            self._subscripts[field] = 0 if value == "__none__" else value
+
+    @staticmethod
+    def _synchronously_lock(method: Coroutine):
+        def setup_lock(condition: Callable[["DBContextStorage"], bool] = lambda _: True):
+            async def lock(self: "DBContextStorage", *args, **kwargs):
+                if condition(self):
+                    async with self._sync_lock:
+                        return await method(self, *args, **kwargs)
+                else:
+                    return await method(self, *args, **kwargs)
+
+            return lock
+
+        return setup_lock
+
+    @staticmethod
+    def _verify_field_name(method: Callable):
+        def verifier(self: "DBContextStorage", *args, **kwargs):
+            field_name = args[1] if len(args) >= 1 else kwargs.get("field_name", None)
+            if field_name is None:
+                raise ValueError(f"For method {method.__name__} argument 'field_name' is not found!")
+            elif field_name not in (self._labels_field_name, self._requests_field_name, self._responses_field_name):
+                raise ValueError(f"Invalid value '{field_name}' for method '{method.__name__}' argument 'field_name'!")
+            else:
+                return method(self, *args, **kwargs)
+
+        return verifier
 
     @abstractmethod
-    async def get_item_async(self, key: Hashable) -> Context:
+    async def load_main_info(self, ctx_id: str) -> Optional[Tuple[int, int, int, bytes, bytes]]:
         """
-        Asynchronous method for accessing stored Context.
-
-        :param key: Hashable key used to store Context instance.
-        :return: The stored context, associated with the given key.
+        Load main information about the context storage.
         """
         raise NotImplementedError
 
-    def __setitem__(self, key: Hashable, value: Context):
-        """
-        Synchronous method for storing Context.
-
-        :param key: Hashable key used to store Context instance.
-        :param value: Context to store.
-        """
-        return asyncio.run(self.set_item_async(key, value))
-
     @abstractmethod
-    async def set_item_async(self, key: Hashable, value: Context):
+    async def update_main_info(
+        self, ctx_id: str, turn_id: int, crt_at: int, upd_at: int, misc: bytes, fw_data: bytes
+    ) -> None:
         """
-        Asynchronous method for storing Context.
-
-        :param key: Hashable key used to store Context instance.
-        :param value: Context to store.
+        Update main information about the context storage.
         """
         raise NotImplementedError
 
-    def __delitem__(self, key: Hashable):
-        """
-        Synchronous method for removing stored Context.
-
-        :param key: Hashable key used to identify Context instance for deletion.
-        """
-        return asyncio.run(self.del_item_async(key))
-
     @abstractmethod
-    async def del_item_async(self, key: Hashable):
+    async def delete_context(self, ctx_id: str) -> None:
         """
-        Asynchronous method for removing stored Context.
-
-        :param key: Hashable key used to identify Context instance for deletion.
+        Delete context from context storage.
         """
         raise NotImplementedError
 
-    def __contains__(self, key: Hashable) -> bool:
-        """
-        Synchronous method for finding whether any Context is stored with given key.
-
-        :param key: Hashable key used to check if Context instance is stored.
-        :return: True if there is Context accessible by given key, False otherwise.
-        """
-        return asyncio.run(self.contains_async(key))
-
     @abstractmethod
-    async def contains_async(self, key: Hashable) -> bool:
+    async def load_field_latest(self, ctx_id: str, field_name: str) -> List[Tuple[int, bytes]]:
         """
-        Asynchronous method for finding whether any Context is stored with given key.
-
-        :param key: Hashable key used to check if Context instance is stored.
-        :return: True if there is Context accessible by given key, False otherwise.
+        Load the latest field data.
         """
         raise NotImplementedError
 
-    def __len__(self) -> int:
+    @abstractmethod
+    async def load_field_keys(self, ctx_id: str, field_name: str) -> List[int]:
         """
-        Synchronous method for retrieving number of stored Contexts.
-
-        :return: The number of stored Contexts.
+        Load all field keys.
         """
-        return asyncio.run(self.len_async())
+        raise NotImplementedError
 
     @abstractmethod
-    async def len_async(self) -> int:
+    async def load_field_items(self, ctx_id: str, field_name: str, keys: List[int]) -> List[Tuple[int, bytes]]:
         """
-        Asynchronous method for retrieving number of stored Contexts.
-
-        :return: The number of stored Contexts.
+        Load field items.
         """
         raise NotImplementedError
 
-    def get(self, key: Hashable, default: Optional[Context] = None) -> Context:
+    @abstractmethod
+    async def update_field_items(self, ctx_id: str, field_name: str, items: List[Tuple[int, bytes]]) -> None:
         """
-        Synchronous method for accessing stored Context, returning default if no Context is stored with the given key.
-
-        :param key: Hashable key used to store Context instance.
-        :param default: Optional default value to be returned if no Context is found.
-        :return: The stored context, associated with the given key or default value.
+        Update field items.
         """
-        return asyncio.run(self.get_async(key, default))
-
-    async def get_async(self, key: Hashable, default: Optional[Context] = None) -> Context:
-        """
-        Asynchronous method for accessing stored Context, returning default if no Context is stored with the given key.
-
-        :param key: Hashable key used to store Context instance.
-        :param default: Optional default value to be returned if no Context is found.
-        :return: The stored context, associated with the given key or default value.
-        """
-        try:
-            return await self.get_item_async(str(key))
-        except KeyError:
-            return default
+        raise NotImplementedError
 
-    def clear(self):
+    @_verify_field_name
+    async def delete_field_keys(self, ctx_id: str, field_name: str, keys: List[int]) -> None:
         """
-        Synchronous method for clearing context storage, removing all the stored Contexts.
+        Delete field keys.
         """
-        return asyncio.run(self.clear_async())
+        await self.update_field_items(ctx_id, field_name, [(k, None) for k in keys])
 
     @abstractmethod
-    async def clear_async(self):
+    async def clear_all(self) -> None:
         """
-        Asynchronous method for clearing context storage, removing all the stored Contexts.
+        Clear all the chatsky tables and records.
         """
         raise NotImplementedError
 
-
-def threadsafe_method(func: Callable):
-    """
-    A decorator that makes sure methods of an object instance are threadsafe.
-    """
-
-    @wraps(func)
-    def _synchronized(self, *args, **kwargs):
-        with self._lock:
-            return func(self, *args, **kwargs)
-
-    return _synchronized
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, DBContextStorage):
+            return False
+        return (
+            self.full_path == other.full_path
+            and self.path == other.path
+            and self.rewrite_existing == other.rewrite_existing
+        )
 
 
 def context_storage_factory(path: str, **kwargs) -> DBContextStorage:
@@ -209,20 +181,28 @@ def context_storage_factory(path: str, **kwargs) -> DBContextStorage:
     json://file.json
     When using sqlite backend your prefix should contain three slashes if you use Windows, or four in other cases:
     sqlite:////file.db
+
+    For MemoryContextStorage pass an empty string as ``path``.
+
     If you want to use additional parameters in class constructors, you can pass them to this function as kwargs.
 
     :param path: Path to the file.
     """
-    prefix, _, _ = path.partition("://")
-    if "sql" in prefix:
-        prefix = prefix.split("+")[0]  # this takes care of alternative sql drivers
-    assert (
-        prefix in PROTOCOLS
-    ), f"""
-    URI path should be prefixed with one of the following:\n
-    {", ".join(PROTOCOLS.keys())}.\n
-    For more information, see the function doc:\n{context_storage_factory.__doc__}
-    """
-    _class, module = PROTOCOLS[prefix]["class"], PROTOCOLS[prefix]["module"]
-    target_class = getattr(importlib.import_module(f".{module}", package="chatsky.context_storages"), _class)
+    if path == "":
+        module = "memory"
+        _class = "MemoryContextStorage"
+    else:
+        prefix, _, _ = path.partition("://")
+        if any(prefix.startswith(sql_prefix) for sql_prefix in ("sqlite", "mysql", "postgresql")):
+            prefix = prefix.split("+")[0]  # this takes care of alternative sql drivers
+        if prefix not in PROTOCOLS:
+            raise ValueError(
+                f"""
+        URI path should be prefixed with one of the following:\n
+        {", ".join(PROTOCOLS.keys())}.\n
+        For more information, see the function doc:\n{context_storage_factory.__doc__}
+        """
+            )
+        _class, module = PROTOCOLS[prefix]["class"], PROTOCOLS[prefix]["module"]
+    target_class = getattr(import_module(f".{module}", package="chatsky.context_storages"), _class)
     return target_class(path, **kwargs)