Redis Doesn't Support Streaming

python/packages/autogen-ext/src/autogen_ext/cache_store/redis.py

@@ -90,6 +90,7 @@ def set(self, key: str, value: T) -> None:
 
         This method handles both primitive values and complex objects:
         - Pydantic models are automatically serialized to JSON
+        - Lists containing Pydantic models are serialized to JSON
         - Primitive values (strings, numbers, etc.) are stored as-is
 
         Args:
@@ -101,10 +102,20 @@ def set(self, key: str, value: T) -> None:
                 # Serialize Pydantic models to JSON
                 serialized_value = value.model_dump_json().encode("utf-8")
                 self.cache.set(key, serialized_value)
+            elif isinstance(value, list):
+                # Serialize lists (which may contain Pydantic models) to JSON
+                serializable_list: list[Any] = []
+                for item in value:
+                    if isinstance(item, BaseModel):
+                        serializable_list.append(item.model_dump())
+                    else:
+                        serializable_list.append(item)
+                serialized_value = json.dumps(serializable_list).encode("utf-8")
+                self.cache.set(key, serialized_value)
             else:
                 # Backward compatibility for primitives
                 self.cache.set(key, cast(Any, value))
-        except (redis.RedisError, ConnectionError, UnicodeEncodeError):
+        except (redis.RedisError, ConnectionError, UnicodeEncodeError, TypeError):
             # Log the error but don't re-raise to maintain robustness
             pass
 

python/packages/autogen-ext/src/autogen_ext/models/cache/_chat_completion_cache.py

@@ -77,6 +77,81 @@ async def main():
 
         asyncio.run(main())
 
+    For Redis caching:
+
+    .. code-block:: python
+
+        import asyncio
+
+        from autogen_core.models import UserMessage
+        from autogen_ext.models.openai import OpenAIChatCompletionClient
+        from autogen_ext.models.cache import ChatCompletionCache, CHAT_CACHE_VALUE_TYPE
+        from autogen_ext.cache_store.redis import RedisStore
+        import redis
+
+
+        async def main():
+            # Initialize the original client
+            openai_model_client = OpenAIChatCompletionClient(model="gpt-4o")
+
+            # Initialize Redis cache store
+            redis_instance = redis.Redis()
+            cache_store = RedisStore[CHAT_CACHE_VALUE_TYPE](redis_instance)
+            cache_client = ChatCompletionCache(openai_model_client, cache_store)
+
+            response = await cache_client.create([UserMessage(content="Hello, how are you?", source="user")])
+            print(response)  # Should print response from OpenAI
+            response = await cache_client.create([UserMessage(content="Hello, how are you?", source="user")])
+            print(response)  # Should print cached response
+
+
+        asyncio.run(main())
+
+    For streaming with Redis caching:
+
+    .. code-block:: python
+
+        import asyncio
+
+        from autogen_core.models import UserMessage, CreateResult
+        from autogen_ext.models.openai import OpenAIChatCompletionClient
+        from autogen_ext.models.cache import ChatCompletionCache, CHAT_CACHE_VALUE_TYPE
+        from autogen_ext.cache_store.redis import RedisStore
+        import redis
+
+
+        async def main():
+            # Initialize the original client
+            openai_model_client = OpenAIChatCompletionClient(model="gpt-4o")
+
+            # Initialize Redis cache store
+            redis_instance = redis.Redis()
+            cache_store = RedisStore[CHAT_CACHE_VALUE_TYPE](redis_instance)
+            cache_client = ChatCompletionCache(openai_model_client, cache_store)
+
+            # First streaming call
+            async for chunk in cache_client.create_stream(
+                [UserMessage(content="List all countries in Africa", source="user")]
+            ):
+                if isinstance(chunk, CreateResult):
+                    print("\\n")
+                    print("Cached: ", chunk.cached)  # Should print False
+                else:
+                    print(chunk, end="")
+
+            # Second streaming call (cached)
+            async for chunk in cache_client.create_stream(
+                [UserMessage(content="List all countries in Africa", source="user")]
+            ):
+                if isinstance(chunk, CreateResult):
+                    print("\\n")
+                    print("Cached: ", chunk.cached)  # Should print True
+                else:
+                    print(chunk, end="")
+
+
+        asyncio.run(main())
+
     You can now use the `cached_client` as you would the original client, but with caching enabled.
 
     Args:
@@ -171,10 +246,18 @@ async def create(
         NOTE: cancellation_token is ignored for cached results.
         """
         cached_result, cache_key = self._check_cache(messages, tools, json_output, extra_create_args)
-        if cached_result:
-            assert isinstance(cached_result, CreateResult)
-            cached_result.cached = True
-            return cached_result
+        if cached_result is not None:
+            if isinstance(cached_result, CreateResult):
+                # Cache hit from previous non-streaming call
+                cached_result.cached = True
+                return cached_result
+            elif isinstance(cached_result, list):
+                # Cache hit from previous streaming call - extract the final CreateResult
+                for item in reversed(cached_result):
+                    if isinstance(item, CreateResult):
+                        item.cached = True
+                        return item
+                # If no CreateResult found in list, fall through to make actual call
 
         result = await self.client.create(
             messages,
@@ -212,13 +295,19 @@ async def _generator() -> AsyncGenerator[Union[str, CreateResult], None]:
                 json_output,
                 extra_create_args,
             )
-            if cached_result:
-                assert isinstance(cached_result, list)
-                for result in cached_result:
-                    if isinstance(result, CreateResult):
-                        result.cached = True
-                    yield result
-                return
+            if cached_result is not None:
+                if isinstance(cached_result, list):
+                    # Cache hit from previous streaming call
+                    for result in cached_result:
+                        if isinstance(result, CreateResult):
+                            result.cached = True
+                        yield result
+                    return
+                elif isinstance(cached_result, CreateResult):
+                    # Cache hit from previous non-streaming call - convert to streaming format
+                    cached_result.cached = True
+                    yield cached_result
+                    return
 
             result_stream = self.client.create_stream(
                 messages,
@@ -230,12 +319,14 @@ async def _generator() -> AsyncGenerator[Union[str, CreateResult], None]:
             )
 
             output_results: List[Union[str, CreateResult]] = []
-            self.store.set(cache_key, output_results)
 
             async for result in result_stream:
                 output_results.append(result)
                 yield result
 
+            # Store the complete results only after streaming is finished
+            self.store.set(cache_key, output_results)
+
         return _generator()
 
     async def close(self) -> None: