Slot Improvements

chatsky/__rebuild_pydantic_models__.py

@@ -1,10 +1,11 @@
 # flake8: noqa: F401
 
 from chatsky.core.service.types import ExtraHandlerRuntimeInfo, ComponentExecutionState
-from chatsky.core import Context, Script
+from chatsky.core import Context, Message, Script
 from chatsky.core.script import Node
 from chatsky.core.pipeline import Pipeline
-from chatsky.slots.slots import SlotManager
+from chatsky.slots.standard_slots import FunctionSlot
+from chatsky.slots.slot_manager import SlotManager
 from chatsky.core.context import FrameworkData, ServiceState
 from chatsky.core.service import PipelineComponent
 
@@ -15,3 +16,4 @@
 ExtraHandlerRuntimeInfo.model_rebuild()
 FrameworkData.model_rebuild()
 ServiceState.model_rebuild()
+FunctionSlot.model_rebuild()

chatsky/conditions/__init__.py

@@ -9,5 +9,5 @@
     Not,
     HasCallbackQuery,
 )
-from chatsky.conditions.slots import SlotsExtracted
+from chatsky.conditions.slots import SlotsExtracted, SlotValueEquals
 from chatsky.conditions.service import ServiceFinished

chatsky/conditions/slots.py

@@ -5,10 +5,10 @@
 """
 
 from __future__ import annotations
-from typing import Literal, List
+from typing import Literal, List, Any
 
 from chatsky.core import Context, BaseCondition
-from chatsky.slots.slots import SlotName
+from chatsky.slots.slot_manager import SlotName
 
 
 class SlotsExtracted(BaseCondition):
@@ -36,3 +36,27 @@ async def call(self, ctx: Context) -> bool:
             return all(manager.is_slot_extracted(slot) for slot in self.slots)
         elif self.mode == "any":
             return any(manager.is_slot_extracted(slot) for slot in self.slots)
+
+
+class SlotValueEquals(BaseCondition):
+    """
+    Check if :py:attr:`.slot_name`'s extracted value is equal to a given value.
+
+    :raises KeyError: If the slot with the specified name does not exist.
+    """
+
+    slot_name: SlotName
+    """
+    Name of the slot that needs to be checked.
+    """
+    value: Any
+    """
+    The value which the slot's extracted value is supposed to be checked against.
+    """
+
+    def __init__(self, slot_name: SlotName, value: Any):
+        super().__init__(slot_name=slot_name, value=value)
+
+    async def call(self, ctx: Context) -> bool:
+        manager = ctx.framework_data.slot_manager
+        return manager.get_extracted_slot(self.slot_name).value == self.value

chatsky/core/context.py

@@ -25,7 +25,7 @@
 from pydantic import BaseModel, Field
 
 from chatsky.core.message import Message, MessageInitTypes
-from chatsky.slots.slots import SlotManager
+from chatsky.slots.slot_manager import SlotManager
 from chatsky.core.node_label import AbsoluteNodeLabel, AbsoluteNodeLabelInitTypes
 
 if TYPE_CHECKING:

chatsky/core/pipeline.py

@@ -21,7 +21,7 @@
 
 from chatsky.messengers.console import CLIMessengerInterface
 from chatsky.messengers.common import MessengerInterface
-from chatsky.slots.slots import GroupSlot
+from chatsky.slots import GroupSlot
 from chatsky.core.service.group import ServiceGroup, ServiceGroupInitTypes
 from chatsky.core.service.extra import ComponentExtraHandlerInitTypes, BeforeHandler, AfterHandler
 from .service import Service

chatsky/processing/slots.py

@@ -9,7 +9,7 @@
 import logging
 from typing import List
 
-from chatsky.slots.slots import SlotName
+from chatsky.slots.slot_manager import SlotName
 from chatsky.core import Context, BaseProcessing
 from chatsky.responses.slots import FilledTemplate
 
@@ -24,16 +24,16 @@ class Extract(BaseProcessing):
 
     slots: List[SlotName]
     """A list of slot names to extract."""
-    success_only: bool = True
+    save_on_failure: bool = True
     """If set, only successfully extracted values will be stored in the slot storage."""
 
-    def __init__(self, *slots: SlotName, success_only: bool = True):
-        super().__init__(slots=slots, success_only=success_only)
+    def __init__(self, *slots: SlotName, save_on_failure: bool = True):
+        super().__init__(slots=slots, save_on_failure=save_on_failure)
 
     async def call(self, ctx: Context):
         manager = ctx.framework_data.slot_manager
         results = await asyncio.gather(
-            *(manager.extract_slot(slot, ctx, self.success_only) for slot in self.slots), return_exceptions=True
+            *(manager.extract_slot(slot, ctx, self.save_on_failure) for slot in self.slots), return_exceptions=True
         )
 
         for result in results:

chatsky/slots/__init__.py

@@ -1 +1,2 @@
-from chatsky.slots.slots import GroupSlot, ValueSlot, RegexpSlot, FunctionSlot
+from chatsky.slots.standard_slots import RegexpSlot, RegexpGroupSlot, FunctionSlot
+from chatsky.slots.base_slots import GroupSlot, ValueSlot