diff --git a/griptape/drivers/memory/conversation/redis_conversation_memory_driver.py b/griptape/drivers/memory/conversation/redis_conversation_memory_driver.py index a6c03884e..65ef03208 100644 --- a/griptape/drivers/memory/conversation/redis_conversation_memory_driver.py +++ b/griptape/drivers/memory/conversation/redis_conversation_memory_driver.py @@ -24,8 +24,8 @@ class RedisConversationMemoryDriver(BaseConversationMemoryDriver): Proper setup of the Redis instance and RediSearch is necessary for the driver to function correctly. Attributes: - host: The host of the Redis instance. - port: The port of the Redis instance. + host: The host of the Redis instance. Required unless client is provided. + port: The port of the Redis instance. Required unless client is provided. db: The database of the Redis instance. username: The username of the Redis instance. password: The password of the Redis instance. @@ -33,9 +33,9 @@ class RedisConversationMemoryDriver(BaseConversationMemoryDriver): conversation_id: The id of the conversation. """ - host: str = field(kw_only=True, metadata={"serializable": True}) + host: str | None = field(kw_only=True, default=None, metadata={"serializable": True}) username: str = field(kw_only=True, default="default", metadata={"serializable": False}) - port: int = field(kw_only=True, metadata={"serializable": True}) + port: int | None = field(kw_only=True, default=None, metadata={"serializable": True}) db: int = field(kw_only=True, default=0, metadata={"serializable": True}) password: str | None = field(default=None, kw_only=True, metadata={"serializable": False}) index: str = field(kw_only=True, metadata={"serializable": True}) @@ -43,18 +43,24 @@ class RedisConversationMemoryDriver(BaseConversationMemoryDriver): client: Redis = field( default=Factory( - lambda self: import_optional_dependency("redis").Redis( - host=self.host, - port=self.port, - db=self.db, - username=self.username, - password=self.password, - decode_responses=False, - ), + lambda self: self._create_client(), takes_self=True, ), ) + def _create_client(self) -> Redis: + if self.host is None or self.port is None: + raise ValueError("host and port are required when client is not provided") + + return import_optional_dependency("redis").Redis( + host=self.host, + port=self.port, + db=self.db, + username=self.username, + password=self.password, + decode_responses=False, + ) + def store(self, runs: list[Run], metadata: dict[str, Any]) -> None: self.client.hset(self.index, self.conversation_id, json.dumps(self._to_params_dict(runs, metadata))) diff --git a/tests/unit/drivers/memory/conversation/test_redis_conversation_memory_driver.py b/tests/unit/drivers/memory/conversation/test_redis_conversation_memory_driver.py index 9531228c8..c2e8ef1ab 100644 --- a/tests/unit/drivers/memory/conversation/test_redis_conversation_memory_driver.py +++ b/tests/unit/drivers/memory/conversation/test_redis_conversation_memory_driver.py @@ -1,3 +1,5 @@ +import json + import pytest import redis @@ -38,6 +40,50 @@ def test_store(self, driver): memory = BaseConversationMemory.from_json(TEST_MEMORY) assert driver.store(memory.runs, memory.meta) is None + @pytest.mark.parametrize("connection_params", [{}, {"host": HOST}, {"port": PORT}, {"host": HOST, "port": PORT}]) + def test_injected_client(self, mocker, connection_params): + client = mocker.MagicMock(spec=redis.Redis) + client.hget.return_value = TEST_DATA + constructor = mocker.patch.object(redis, "Redis") + driver = RedisConversationMemoryDriver( + client=client, index=INDEX, conversation_id=CONVERSATION_ID, **connection_params + ) + + assert driver.client is client + constructor.assert_not_called() + runs, metadata = driver.load() + client.hget.assert_called_once_with(INDEX, CONVERSATION_ID) + assert len(runs) == 1 + assert metadata == {"foo": "bar"} + driver.store(runs, metadata) + client.hset.assert_called_once() + index, conversation_id, stored_json = client.hset.call_args.args + assert (index, conversation_id) == (INDEX, CONVERSATION_ID) + stored = json.loads(stored_json) + assert stored["metadata"] == metadata + assert stored["runs"][0]["input"]["value"] == "Hi There, Hello" + assert stored["runs"][0]["output"]["value"] == "Hello! How can I assist you today?" + + def test_create_client(self, mocker): + constructor = mocker.patch.object(redis, "Redis") + driver = RedisConversationMemoryDriver( + host=HOST, port=PORT, username=USERNAME, password=PASSWORD, db=1, index=INDEX + ) + + assert driver.client is constructor.return_value + constructor.assert_called_once_with( + host=HOST, port=PORT, username=USERNAME, password=PASSWORD, db=1, decode_responses=False + ) + + @pytest.mark.parametrize("connection_params", [{}, {"host": HOST}, {"port": PORT}]) + def test_missing_connection_parameters(self, mocker, connection_params): + constructor = mocker.patch.object(redis, "Redis") + + with pytest.raises(ValueError, match="host and port are required when client is not provided"): + RedisConversationMemoryDriver(index=INDEX, **connection_params) + + constructor.assert_not_called() + def test_load(self, driver): runs, metadata = driver.load() assert len(runs) == 1