Update KnowledgeBase, InMemoryKnowledgeBase, ActionQueryKnowledgeBase
MODIFY
rasa_sdk/knowledge_base/actions.py
rasa_sdk/knowledge_base/actions.py
@@ -65,7 +65,7 @@ class ActionQueryKnowledgeBase(Action):
text=f"Did not find a valid value for attribute '{attribute_name}' for object '{object_name}'."
)
- def utter_objects(
+ async def utter_objects(
self,
dispatcher: CollectingDispatcher,
object_type: Text,
@@ -84,9 +84,17 @@ class ActionQueryKnowledgeBase(Action):
text=f"Found the following objects of type '{object_type}':"
)
+ if utils.is_coroutine_action(
+ self.knowledge_base.get_representation_function_of_object
+ ):
+ repr_function = await self.knowledge_base.get_representation_function_of_object(
+ object_type
+ )
+ else:
repr_function = self.knowledge_base.get_representation_function_of_object(
object_type
)
+
for i, obj in enumerate(objects, 1):
dispatcher.utter_message(text=f"{i}: {repr_function(obj)}")
else:
@@ -152,7 +160,7 @@ class ActionQueryKnowledgeBase(Action):
object_type = tracker.get_slot(SLOT_OBJECT_TYPE)
if utils.is_coroutine_action(self.knowledge_base.get_attributes_of_object):
object_attributes = await self.knowledge_base.get_attributes_of_object(
- object_type # type: ignore
+ object_type
)
else:
object_attributes = self.knowledge_base.get_attributes_of_object(
@@ -164,7 +172,7 @@ class ActionQueryKnowledgeBase(Action):
attributes = get_attribute_slots(tracker, object_attributes)
# query the knowledge base
if utils.is_coroutine_action(self.knowledge_base.get_objects):
- objects = await self.knowledge_base.get_objects(object_type, attributes) # type: ignore
+ objects = await self.knowledge_base.get_objects(object_type, attributes)
else:
objects = self.knowledge_base.get_objects(object_type, attributes)
@@ -178,7 +186,7 @@ class ActionQueryKnowledgeBase(Action):
if utils.is_coroutine_action(self.knowledge_base.get_key_attribute_of_object):
key_attribute = await self.knowledge_base.get_key_attribute_of_object(
- object_type # type: ignore
+ object_type
)
else:
key_attribute = self.knowledge_base.get_key_attribute_of_object(object_type)
@@ -251,7 +259,7 @@ class ActionQueryKnowledgeBase(Action):
object_representation = repr_function(object_of_interest)
if utils.is_coroutine_action(self.knowledge_base.get_key_attribute_of_object):
key_attribute = await self.knowledge_base.get_key_attribute_of_object(
- object_type # type: ignore
+ object_type
)
else:
key_attribute = self.knowledge_base.get_key_attribute_of_object(object_type)

MODIFY
rasa_sdk/knowledge_base/storage.py
rasa_sdk/knowledge_base/storage.py
@@ -5,6 +5,8 @@ import random
from typing import Text, Callable, Dict, List, Any, Optional
from collections import defaultdict
+from rasa_sdk import utils
+
logger = logging.getLogger(__name__)
@@ -29,7 +31,7 @@ class KnowledgeBase:
self.key_attribute = defaultdict(lambda: "id")
self.representation_function = defaultdict(lambda: lambda obj: obj["name"])
- def get_attributes_of_object(self, object_type: Text) -> List[Text]:
+ async def get_attributes_of_object(self, object_type: Text) -> List[Text]:
"""
Returns a list of all attributes that belong to the provided object type.
@@ -40,7 +42,7 @@ class KnowledgeBase:
"""
raise NotImplementedError("Method is not implemented.")
- def get_key_attribute_of_object(self, object_type: Text) -> Text:
+ async def get_key_attribute_of_object(self, object_type: Text) -> Text:
"""
Returns the key attribute for the given object type.
@@ -51,7 +53,9 @@ class KnowledgeBase:
"""
return self.key_attribute[object_type]
- def get_representation_function_of_object(self, object_type: Text) -> Callable:
+ async def get_representation_function_of_object(
+ self, object_type: Text
+ ) -> Callable:
"""
Returns a lamdba function that takes the object and returns a string
representation of it.
@@ -73,7 +77,7 @@ class KnowledgeBase:
"""
self.ordinal_mention_mapping = mapping
- def get_objects(
+ async def get_objects(
self, object_type: Text, attributes: List[Dict[Text, Text]], limit: int = 5
) -> List[Dict[Text, Any]]:
"""
@@ -89,7 +93,9 @@ class KnowledgeBase:
"""
raise NotImplementedError("Method is not implemented.")
- def get_object(self, object_type: Text, object_identifier: Text) -> Dict[Text, Any]:
+ async def get_object(
+ self, object_type: Text, object_identifier: Text
+ ) -> Dict[Text, Any]:
"""
Returns the object of the given type that matches the given object identifier.
@@ -158,7 +164,7 @@ class InMemoryKnowledgeBase(KnowledgeBase):
"""
self.key_attribute[object_type] = key_attribute
- def get_attributes_of_object(self, object_type: Text) -> List[Text]:
+ async def get_attributes_of_object(self, object_type: Text) -> List[Text]:
if object_type not in self.data or not self.data[object_type]:
return []
@@ -166,7 +172,7 @@ class InMemoryKnowledgeBase(KnowledgeBase):
return list(first_object.keys())
- def get_objects(
+ async def get_objects(
self, object_type: Text, attributes: List[Dict[Text, Text]], limit: int = 5
) -> List[Dict[Text, Any]]:
if object_type not in self.data:
@@ -190,13 +196,17 @@ class InMemoryKnowledgeBase(KnowledgeBase):
return objects[:limit]
- def get_object(
+ async def get_object(
self, object_type: Text, object_identifier: Text
) -> Optional[Dict[Text, Any]]:
if object_type not in self.data:
return None
objects = self.data[object_type]
+
+ if utils.is_coroutine_action(self.get_key_attribute_of_object):
+ key_attribute = await self.get_key_attribute_of_object(object_type)
+ else:
key_attribute = self.get_key_attribute_of_object(object_type)
# filter the objects by its key attribute, for example, 'id'
@@ -211,7 +221,13 @@ class InMemoryKnowledgeBase(KnowledgeBase):
# if the object was referred to directly, we need to compare the representation
# of each object with the given object identifier
if not objects_of_interest:
+ if utils.is_coroutine_action(self.get_representation_function_of_object):
+ repr_function = await self.get_representation_function_of_object(
+ object_type
+ )
+ else:
repr_function = self.get_representation_function_of_object(object_type)
+
objects_of_interest = list(
filter(
lambda obj: str(object_identifier).lower()

MODIFY
tests/knowledge_base/test_actions.py
tests/knowledge_base/test_actions.py
import pytest
-import asyncio
from rasa_sdk import Tracker
from rasa_sdk.events import SlotSet
@@ -117,16 +116,14 @@ def compare_slots(slot_list_1, slot_list_2):
),
],
)
-def test_action_run(data_file, slots, expected_slots):
+async def test_action_run(data_file, slots, expected_slots):
knowledge_base = InMemoryKnowledgeBase(data_file)
action = ActionQueryKnowledgeBase(knowledge_base)
dispatcher = CollectingDispatcher()
tracker = Tracker("default", slots, {}, [], False, None, {}, "action_listen")
- loop = asyncio.new_event_loop()
- actual_slots = loop.run_until_complete(action.run(dispatcher, tracker, {}))
- loop.close()
+ actual_slots = await action.run(dispatcher, tracker, {})
compare_slots(expected_slots, actual_slots)
compare_slots(actual_slots, expected_slots)

MODIFY
tests/knowledge_base/test_storage.py
tests/knowledge_base/test_storage.py
@@ -12,10 +12,10 @@ from rasa_sdk.knowledge_base.storage import InMemoryKnowledgeBase
("restaurant", [{"name": "cuisine", "value": "Italian"}], 2),
],
)
-def test_query_entities(data_file, object_type, attributes, expected_length):
+async def test_query_entities(data_file, object_type, attributes, expected_length):
knowledge_base = InMemoryKnowledgeBase(data_file)
- entities = knowledge_base.get_objects(
+ entities = await knowledge_base.get_objects(
object_type=object_type, attributes=attributes
)
assert expected_length == len(entities)
@@ -43,10 +43,10 @@ def test_query_entities(data_file, object_type, attributes, expected_length):
("hotel", None, None),
],
)
-def test_query_object(data_file, object_type, object_identifier, expected_value):
+async def test_query_object(data_file, object_type, object_identifier, expected_value):
knowledge_base = InMemoryKnowledgeBase(data_file)
- actual_value = knowledge_base.get_object(
+ actual_value = await knowledge_base.get_object(
object_type=object_type, object_identifier=object_identifier
)
@@ -57,10 +57,12 @@ def test_query_object(data_file, object_type, object_identifier, expected_value)
"object_type,expected_attributes",
[("restaurant", ["id", "name", "cuisine", "wifi"])],
)
-def test_get_attributes_for(data_file, object_type, expected_attributes):
+async def test_get_attributes_for(data_file, object_type, expected_attributes):
knowledge_base = InMemoryKnowledgeBase(data_file)
- actual_attributes = knowledge_base.get_attributes_of_object(object_type=object_type)
+ actual_attributes = await knowledge_base.get_attributes_of_object(
+ object_type=object_type
+ )
assert set(expected_attributes) == set(actual_attributes)
@@ -69,14 +71,14 @@ def test_get_attributes_for(data_file, object_type, expected_attributes):
"object_type,set_key_attribute,expected_key_attribute",
[("restaurant", None, "id"), ("restaurant", "name", "name")],
)
-def test_key_attribute_of_object(
+async def test_key_attribute_of_object(
data_file, object_type, set_key_attribute, expected_key_attribute
):
knowledge_base = InMemoryKnowledgeBase(data_file)
if set_key_attribute:
knowledge_base.set_key_attribute_of_object(object_type, set_key_attribute)
- actual_key_attribute = knowledge_base.get_key_attribute_of_object(
+ actual_key_attribute = await knowledge_base.get_key_attribute_of_object(
object_type=object_type
)
@@ -94,7 +96,7 @@ def test_key_attribute_of_object(
),
],
)
-def test_get_representation_function_of_object(
+async def test_get_representation_function_of_object(
data_file, object_type, set_repr_function, expected_repr_function
):
knowledge_base = InMemoryKnowledgeBase(data_file)
@@ -103,10 +105,10 @@ def test_get_representation_function_of_object(
object_type, set_repr_function
)
- actual_repr_function = knowledge_base.get_representation_function_of_object(
+ actual_repr_function = await knowledge_base.get_representation_function_of_object(
object_type=object_type
)
- dummy_object = knowledge_base.get_object(object_type, 1)
+ dummy_object = await knowledge_base.get_object(object_type, 1)
assert expected_repr_function(dummy_object) == actual_repr_function(dummy_object)

Update KnowledgeBase, InMemoryKnowledgeBase, ActionQueryKnowledgeBase