chg: goobredis fixes

.prompts.d/annotations.txt

@@ -9,23 +9,23 @@
 
     using @src/goob_ai/backend/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/backend/cache/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/cache/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/cache/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/backend/cache/goobredis.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/cache/goobredis.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/cache/goobredis.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @goobredis.py as context, I want you to make sure to add typing annotations to each function or class in side of @goobredis.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @goobredis.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/backend/cache/goobredis_session.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/cache/goobredis_session.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/cache/goobredis_session.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @goobredis_session.py as context, I want you to make sure to add typing annotations to each function or class in side of @goobredis_session.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @goobredis_session.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/backend/discord/discord_client.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/discord/discord_client.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/discord/discord_client.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @discord_client.py as context, I want you to make sure to add typing annotations to each function or class in side of @discord_client.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @discord_client.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
     using @src/goob_ai/backend/ml/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/ml/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/ml/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/backend/ml/devices.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/ml/devices.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/ml/devices.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @devices.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/backend/ml/devices.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/backend/ml/devices.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/base.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/base.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/base.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @base.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/base.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/base.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
     using @src/goob_ai/bot_logger/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/bot_logger/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/bot_logger/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/chromadb_manager.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/chromadb_manager.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/chromadb_manager.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @chromadb_manager.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/chromadb_manager.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/chromadb_manager.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
     using @src/goob_ai/cli.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/cli.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/cli.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
@@ -149,13 +149,13 @@
 
     using @src/goob_ai/services/__init__.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/__init__.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/__init__.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/services/chroma_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/chroma_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/chroma_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @chroma_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/chroma_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/chroma_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
     using @src/goob_ai/services/cromadb_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/cromadb_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/cromadb_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
     using @src/goob_ai/services/image_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/image_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/image_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
-    using @src/goob_ai/services/pinecone_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/pinecone_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/pinecone_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
+    using @pinecone_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/pinecone_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/pinecone_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 
     using @src/goob_ai/services/screencrop_service.py as context, I want you to make sure to add typing annotations to each function or class in side of @src/goob_ai/services/screencrop_service.py. Be sure to include return types when necessary. add descriptive docstrings to all functions and classes inside of @src/goob_ai/services/screencrop_service.py. Please use pep257 convention, in the style of 'google'. Update existing docstrings if they do not conform to is. Don't forget to include any necessary typing modules at the top of the file.
 

src/goob_ai/backend/cache/goobredis.py

@@ -106,6 +106,7 @@ async def _connect(self) -> None:
         )
         self._pool = redis.asyncio.Redis(connection_pool=self._conn_pool)
         self._pubsub_channels: dict[str, PubSub] = {}
+        self.auto_close_connection_pool: bool = self._pool.auto_close_connection_pool
 
     async def finalize(self) -> None:
         """Finalize the Redis client."""
@@ -293,28 +294,6 @@ async def unsubscribe(self, channel_name: str) -> None:
             finally:
                 await p.__aexit__(None, None, None)
 
-    # async def subscribe(self, channel_name: str) -> AsyncGenerator[Any, None]:
-    #     """
-    #     Subscribe to a Redis channel.
-
-    #     Args:
-    #         channel_name: The name of the channel to subscribe to.
-
-    #     Yields:
-    #         Messages received from the subscribed channel.
-
-    #     Raises:
-    #         NoRedisConfigured: If no Redis is configured.
-    #     """
-    #     if self._pool is None:
-    #         raise NoRedisConfigured()
-
-    #     p: PubSub = await self._pool.pubsub().__aenter__()
-    #     self._pubsub_channels[channel_name] = p
-    #     await p.subscribe(channel_name)
-    #     async for message in self._listener(p):
-    #         yield message
-
     async def subscribe(self, channel_name: str) -> AsyncGenerator[Any, None]:
         """
         Subscribe to a Redis channel.
@@ -351,6 +330,22 @@ async def _listener(self, p: PubSub) -> AsyncGenerator[Any, None]:
             if message is not None:
                 yield message["data"]
 
+    async def aclose(self, close_connection_pool: Optional[bool] = None) -> None:
+        """
+        Closes Redis client connection
+
+        :param close_connection_pool: decides whether to close the connection pool used
+        by this Redis client, overriding Redis.auto_close_connection_pool. By default,
+        let Redis.auto_close_connection_pool decide whether to close the connection
+        pool.
+        """
+        conn = self._pool
+        if conn:
+            self._pool = None
+            await self._pool.release(conn)
+        if close_connection_pool or (close_connection_pool is None and self.auto_close_connection_pool):
+            await self._pool.disconnect()
+
 
 #     def set(self, key, value, expiration=3600):
 #         try:

tests/backend/cache/test_goobredis.py

@@ -8,6 +8,8 @@
 
 from typing import TYPE_CHECKING
 
+import redis.asyncio as redis
+
 from goob_ai.backend.cache.goobredis import get_driver
 from loguru import logger as LOGGER
 
@@ -22,10 +24,54 @@
 
     from pytest_mock.plugin import MockerFixture
 
+##########################################################################
+
+
+@pytest.fixture(name="create_redis")
+async def create_redis(request):
+    """Wrapper around redis.create_redis."""
+    # (single_connection,) = request.param
+
+    teardown_clients = []
+
+    async def client_factory(
+        **kwargs,
+    ) -> GoobRedisClient:
+        client = await get_driver()
+
+        async def teardown():
+            try:
+                await client.aclose()
+                # await client.connection_pool.disconnect()
+            except redis.ConnectionError:
+                pass
+
+        teardown_clients.append(teardown)
+        return client
+
+    yield client_factory
+
+    for teardown in teardown_clients:
+        await teardown()
+
+
+@pytest.fixture(name="r")
+async def r(create_redis):
+    return await create_redis()
+
+
+@pytest.fixture(name="decoded_r")
+async def decoded_r(create_redis):
+    return await create_redis(decode_responses=True)
+
+
+##########################################################################
+
 
 # @pytest.mark.app_settings({"applications": ["guillotina", "guillotina.contrib.redis"]})
-async def test_redis_ops(caplog: LogCaptureFixture):
-    driver = await get_driver()
+async def test_redis_ops(caplog: LogCaptureFixture, create_redis, **kwargs):
+    driver: GoobRedisClient = await get_driver()
+    r = await create_redis(**kwargs)
     assert driver.initialized
     assert driver.pool is not None