Upstash

UpstashChatStore

基础：BaseChatStore

使用 Redis 存储和检索聊天消息的 Upstash 聊天存储。

此类实现了 BaseChatStore 接口，并提供了在 Upstash Redis 数据库中管理聊天消息的方法。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

class UpstashChatStore(BaseChatStore):
    """
    Upstash chat store for storing and retrieving chat messages using Redis.

    This class implements the BaseChatStore interface and provides methods
    for managing chat messages in an Upstash Redis database.
    """

    _sync_redis_client: SyncRedis = PrivateAttr()
    _async_redis_client: AsyncRedis = PrivateAttr()

    ttl: Optional[int] = Field(default=None, description="Time to live in seconds.")

    class Config:
        arbitrary_types_allowed = True

    def __init__(
        self,
        redis_url: str = "",
        redis_token: str = "",
        ttl: Optional[int] = None,
    ):
        """
        Initialize the UpstashChatStore.

        Args:
            redis_url (str): The URL of the Upstash Redis instance.
            redis_token (str): The authentication token for the Upstash Redis instance.
            ttl (Optional[int]): Time to live in seconds for stored messages.

        Raises:
            ValueError: If redis_url or redis_token is empty.

        """
        if redis_url == "" or redis_token == "":
            raise ValueError("Please provide a valid URL and token")
        super().__init__(ttl=ttl)
        try:
            self._sync_redis_client = SyncRedis(url=redis_url, token=redis_token)
            self._async_redis_client = AsyncRedis(url=redis_url, token=redis_token)
        except Exception as error:
            logger.error(f"Upstash Redis client could not be initiated: {error}")

        # self.ttl = ttl

    @classmethod
    def class_name(cls) -> str:
        """
        Get the class name.

        Returns:
            str: The name of the class.

        """
        return "UpstashChatStore"

    def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
        """
        Set messages for a key.

        Args:
            key (str): The key to store the messages under.
            messages (List[ChatMessage]): The list of messages to store.

        """
        self._sync_redis_client.delete(key)
        for message in messages:
            self.add_message(key, message)

        if self.ttl:
            self._sync_redis_client.expire(key, self.ttl)

    async def async_set_messages(self, key: str, messages: List[ChatMessage]) -> None:
        """
        Set messages for a key.

        Args:
            key (str): The key to store the messages under.
            messages (List[ChatMessage]): The list of messages to store.

        """
        await self._async_redis_client.delete(key)
        for message in messages:
            await self.async_add_message(key, message)

        if self.ttl:
            await self._async_redis_client.expire(key, self.ttl)

    def get_messages(self, key: str) -> List[ChatMessage]:
        """
        Get messages for a key.

        Args:
            key (str): The key to retrieve messages from.

        Returns:
            List[ChatMessage]: The list of retrieved messages.

        """
        items = self._sync_redis_client.lrange(key, 0, -1)
        if len(items) == 0:
            return []

        return [ChatMessage.parse_raw(item) for item in items]

    async def async_get_messages(self, key: str) -> List[ChatMessage]:
        """
        Get messages for a key.

        Args:
            key (str): The key to retrieve messages from.

        Returns:
            List[ChatMessage]: The list of retrieved messages.

        """
        items = await self._async_redis_client.lrange(key, 0, -1)
        if len(items) == 0:
            return []

        return [ChatMessage.parse_raw(item) for item in items]

    def add_message(
        self, key: str, message: ChatMessage, idx: Optional[int] = None
    ) -> None:
        """
        Add a message to a key.

        Args:
            key (str): The key to add the message to.
            message (ChatMessage): The message to add.
            idx (Optional[int]): The index at which to insert the message.

        """
        if idx is None:
            message_json = json.dumps(_message_to_dict(message))
            self._sync_redis_client.rpush(key, message_json)
        else:
            self._insert_element_at_index(key, message, idx)

        if self.ttl:
            self._sync_redis_client.expire(key, self.ttl)

    async def async_add_message(
        self, key: str, message: ChatMessage, idx: Optional[int] = None
    ) -> None:
        """
        Add a message to a key.

        Args:
            key (str): The key to add the message to.
            message (ChatMessage): The message to add.
            idx (Optional[int]): The index at which to insert the message.

        """
        if idx is None:
            message_json = json.dumps(_message_to_dict(message))
            await self._async_redis_client.rpush(key, message_json)
        else:
            await self._async_insert_element_at_index(key, message, idx)

        if self.ttl:
            await self._async_redis_client.expire(key, self.ttl)

    def delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
        """
        Delete messages for a key.

        Args:
            key (str): The key to delete messages from.

        Returns:
            Optional[List[ChatMessage]]: Always returns None in this implementation.

        """
        self._sync_redis_client.delete(key)
        return None

    async def async_delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
        """
        Delete messages for a key.

        Args:
            key (str): The key to delete messages from.

        Returns:
            Optional[List[ChatMessage]]: Always returns None in this implementation.

        """
        await self._async_redis_client.delete(key)
        return None

    def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
        """
        Delete a message from a key.

        Args:
            key (str): The key to delete the message from.
            idx (int): The index of the message to delete.

        Returns:
            Optional[ChatMessage]: The deleted message, or None if not found or an error occurred.

        """
        try:
            deleted_message = self._sync_redis_client.lindex(key, idx)
            if deleted_message is None:
                return None

            placeholder = f"{key}:{idx}:deleted"

            self._sync_redis_client.lset(key, idx, placeholder)
            self._sync_redis_client.lrem(key, 1, placeholder)
            if self.ttl:
                self._sync_redis_client.expire(key, self.ttl)

            return deleted_message

        except Exception as e:
            logger.error(f"Error deleting message at index {idx} from {key}: {e}")
            return None

    async def async_delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
        """
        Delete a message from a key.

        Args:
            key (str): The key to delete the message from.
            idx (int): The index of the message to delete.

        Returns:
            Optional[ChatMessage]: The deleted message, or None if not found or an error occurred.

        """
        try:
            deleted_message = await self._async_redis_client.lindex(key, idx)
            if deleted_message is None:
                return None

            placeholder = f"{key}:{idx}:deleted"

            await self._async_redis_client.lset(key, idx, placeholder)
            await self._async_redis_client.lrem(key, 1, placeholder)
            if self.ttl:
                await self._async_redis_client.expire(key, self.ttl)

            return deleted_message

        except Exception as e:
            logger.error(f"Error deleting message at index {idx} from {key}: {e}")
            return None

    def delete_last_message(self, key: str) -> Optional[ChatMessage]:
        """
        Delete the last message from a key.

        Args:
            key (str): The key to delete the last message from.

        Returns:
            Optional[ChatMessage]: The deleted message, or None if the list is empty.

        """
        deleted_message = self._sync_redis_client.rpop(key)
        return ChatMessage.parse_raw(deleted_message) if deleted_message else None

    async def async_delete_last_message(self, key: str) -> Optional[ChatMessage]:
        """
        Delete the last message from a key.

        Args:
            key (str): The key to delete the last message from.

        Returns:
            Optional[ChatMessage]: The deleted message, or None if the list is empty.

        """
        deleted_message = await self._async_redis_client.rpop(key)
        if deleted_message:
            return ChatMessage.parse_raw(deleted_message)
        return None

    def get_keys(self) -> List[str]:
        """
        Get all keys.

        Returns:
            List[str]: A list of all keys in the Redis store.

        """
        keys = self._sync_redis_client.keys("*")
        return keys if isinstance(keys, list) else [keys]

    async def async_get_keys(self) -> List[str]:
        """
        Get all keys.

        Returns:
            List[str]: A list of all keys in the Redis store.

        """
        keys = await self._async_redis_client.keys("*")
        return keys if isinstance(keys, list) else [keys]

    def _insert_element_at_index(
        self, key: str, message: ChatMessage, idx: int
    ) -> List[ChatMessage]:
        """
        Insert a message at a specific index.

        Args:
            key (str): The key of the list to insert into.
            message (ChatMessage): The message to insert.
            idx (int): The index at which to insert the message.

        Returns:
            List[ChatMessage]: The updated list of messages.

        """
        current_list = self.get_messages(key)
        current_list.insert(idx, message)

        self._sync_redis_client.delete(key)

        self.set_messages(key, current_list)

        return current_list

    async def _async_insert_element_at_index(
        self, key: str, message: ChatMessage, idx: int
    ) -> List[ChatMessage]:
        """
        Insert a message at a specific index.

        Args:
            key (str): The key of the list to insert into.
            message (ChatMessage): The message to insert.
            idx (int): The index at which to insert the message.

        Returns:
            List[ChatMessage]: The updated list of messages.

        """
        current_list = await self.async_get_messages(key)
        current_list.insert(idx, message)

        await self.async_delete_messages(key)

        await self.async_set_messages(key, current_list)

        return current_list

class_name classmethod

class_name() -> str

获取类名。

返回值

名称	类型	描述
str	str	类的名称。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

@classmethod
def class_name(cls) -> str:
    """
    Get the class name.

    Returns:
        str: The name of the class.

    """
    return "UpstashChatStore"

set_messages

set_messages(key: str, messages: List[ChatMessage]) -> None

为键设置消息。

参数

名称	类型	描述	默认值
key	str	存储消息的键。	必需
messages	List[ChatMessage]	要存储的消息列表。	必需

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
    """
    Set messages for a key.

    Args:
        key (str): The key to store the messages under.
        messages (List[ChatMessage]): The list of messages to store.

    """
    self._sync_redis_client.delete(key)
    for message in messages:
        self.add_message(key, message)

    if self.ttl:
        self._sync_redis_client.expire(key, self.ttl)

async_set_messages async

async_set_messages(key: str, messages: List[ChatMessage]) -> None

为键设置消息。

参数

名称	类型	描述	默认值
key	str	存储消息的键。	必需
messages	List[ChatMessage]	要存储的消息列表。	必需

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

async def async_set_messages(self, key: str, messages: List[ChatMessage]) -> None:
    """
    Set messages for a key.

    Args:
        key (str): The key to store the messages under.
        messages (List[ChatMessage]): The list of messages to store.

    """
    await self._async_redis_client.delete(key)
    for message in messages:
        await self.async_add_message(key, message)

    if self.ttl:
        await self._async_redis_client.expire(key, self.ttl)

get_messages

get_messages(key: str) -> List[ChatMessage]

获取某个键的消息。

参数

名称	类型	描述	默认值
key	str	用于检索消息的键。	必需

返回值

类型	描述
List[ChatMessage]	List[ChatMessage]：检索到的消息列表。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

def get_messages(self, key: str) -> List[ChatMessage]:
    """
    Get messages for a key.

    Args:
        key (str): The key to retrieve messages from.

    Returns:
        List[ChatMessage]: The list of retrieved messages.

    """
    items = self._sync_redis_client.lrange(key, 0, -1)
    if len(items) == 0:
        return []

    return [ChatMessage.parse_raw(item) for item in items]

async_get_messages async

async_get_messages(key: str) -> List[ChatMessage]

获取某个键的消息。

参数

名称	类型	描述	默认值
key	str	用于检索消息的键。	必需

返回值

类型	描述
List[ChatMessage]	List[ChatMessage]：检索到的消息列表。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

async def async_get_messages(self, key: str) -> List[ChatMessage]:
    """
    Get messages for a key.

    Args:
        key (str): The key to retrieve messages from.

    Returns:
        List[ChatMessage]: The list of retrieved messages.

    """
    items = await self._async_redis_client.lrange(key, 0, -1)
    if len(items) == 0:
        return []

    return [ChatMessage.parse_raw(item) for item in items]

add_message

add_message(key: str, message: ChatMessage, idx: Optional[int] = None) -> None

向某个键添加消息。

参数

名称	类型	描述	默认值
key	str	用于添加消息的键。	必需
message	ChatMessage	要添加的消息。	必需
idx	可选[int]	插入消息的索引位置。	无

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

def add_message(
    self, key: str, message: ChatMessage, idx: Optional[int] = None
) -> None:
    """
    Add a message to a key.

    Args:
        key (str): The key to add the message to.
        message (ChatMessage): The message to add.
        idx (Optional[int]): The index at which to insert the message.

    """
    if idx is None:
        message_json = json.dumps(_message_to_dict(message))
        self._sync_redis_client.rpush(key, message_json)
    else:
        self._insert_element_at_index(key, message, idx)

    if self.ttl:
        self._sync_redis_client.expire(key, self.ttl)

async_add_message async

async_add_message(key: str, message: ChatMessage, idx: Optional[int] = None) -> None

向某个键添加消息。

参数

名称	类型	描述	默认值
key	str	用于添加消息的键。	必需
message	ChatMessage	要添加的消息。	必需
idx	可选[int]	插入消息的索引位置。	无

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

async def async_add_message(
    self, key: str, message: ChatMessage, idx: Optional[int] = None
) -> None:
    """
    Add a message to a key.

    Args:
        key (str): The key to add the message to.
        message (ChatMessage): The message to add.
        idx (Optional[int]): The index at which to insert the message.

    """
    if idx is None:
        message_json = json.dumps(_message_to_dict(message))
        await self._async_redis_client.rpush(key, message_json)
    else:
        await self._async_insert_element_at_index(key, message, idx)

    if self.ttl:
        await self._async_redis_client.expire(key, self.ttl)

delete_messages

delete_messages(key: str) -> Optional[List[ChatMessage]]

删除某个键的消息。

参数

名称	类型	描述	默认值
key	str	用于删除消息的键。	必需

返回值

类型	描述
可选[列表[ChatMessage]]	可选[列表[ChatMessage]]：在此实现中总是返回 None。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

def delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
    """
    Delete messages for a key.

    Args:
        key (str): The key to delete messages from.

    Returns:
        Optional[List[ChatMessage]]: Always returns None in this implementation.

    """
    self._sync_redis_client.delete(key)
    return None

async_delete_messages async

async_delete_messages(key: str) -> Optional[List[ChatMessage]]

删除某个键的消息。

参数

名称	类型	描述	默认值
key	str	用于删除消息的键。	必需

返回值

类型	描述
可选[列表[ChatMessage]]	可选[列表[ChatMessage]]：在此实现中总是返回 None。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

async def async_delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
    """
    Delete messages for a key.

    Args:
        key (str): The key to delete messages from.

    Returns:
        Optional[List[ChatMessage]]: Always returns None in this implementation.

    """
    await self._async_redis_client.delete(key)
    return None

delete_message

delete_message(key: str, idx: int) -> Optional[ChatMessage]

从某个键删除一条消息。

参数

名称	类型	描述	默认值
key	str	用于删除消息的键。	必需
idx	int	要删除消息的索引。	必需

返回值

类型	描述
可选[ChatMessage]	可选[ChatMessage]：已删除的消息，如果未找到或发生错误则为 None。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
    """
    Delete a message from a key.

    Args:
        key (str): The key to delete the message from.
        idx (int): The index of the message to delete.

    Returns:
        Optional[ChatMessage]: The deleted message, or None if not found or an error occurred.

    """
    try:
        deleted_message = self._sync_redis_client.lindex(key, idx)
        if deleted_message is None:
            return None

        placeholder = f"{key}:{idx}:deleted"

        self._sync_redis_client.lset(key, idx, placeholder)
        self._sync_redis_client.lrem(key, 1, placeholder)
        if self.ttl:
            self._sync_redis_client.expire(key, self.ttl)

        return deleted_message

    except Exception as e:
        logger.error(f"Error deleting message at index {idx} from {key}: {e}")
        return None

async_delete_message async

async_delete_message(key: str, idx: int) -> Optional[ChatMessage]

从某个键删除一条消息。

参数

名称	类型	描述	默认值
key	str	用于删除消息的键。	必需
idx	int	要删除消息的索引。	必需

返回值

类型	描述
可选[ChatMessage]	可选[ChatMessage]：已删除的消息，如果未找到或发生错误则为 None。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

async def async_delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
    """
    Delete a message from a key.

    Args:
        key (str): The key to delete the message from.
        idx (int): The index of the message to delete.

    Returns:
        Optional[ChatMessage]: The deleted message, or None if not found or an error occurred.

    """
    try:
        deleted_message = await self._async_redis_client.lindex(key, idx)
        if deleted_message is None:
            return None

        placeholder = f"{key}:{idx}:deleted"

        await self._async_redis_client.lset(key, idx, placeholder)
        await self._async_redis_client.lrem(key, 1, placeholder)
        if self.ttl:
            await self._async_redis_client.expire(key, self.ttl)

        return deleted_message

    except Exception as e:
        logger.error(f"Error deleting message at index {idx} from {key}: {e}")
        return None

delete_last_message

delete_last_message(key: str) -> Optional[ChatMessage]

从某个键删除最后一条消息。

参数

名称	类型	描述	默认值
key	str	用于删除最后一条消息的键。	必需

返回值

类型	描述
可选[ChatMessage]	可选[ChatMessage]：已删除的消息，如果列表为空则为 None。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

def delete_last_message(self, key: str) -> Optional[ChatMessage]:
    """
    Delete the last message from a key.

    Args:
        key (str): The key to delete the last message from.

    Returns:
        Optional[ChatMessage]: The deleted message, or None if the list is empty.

    """
    deleted_message = self._sync_redis_client.rpop(key)
    return ChatMessage.parse_raw(deleted_message) if deleted_message else None

async_delete_last_message async

async_delete_last_message(key: str) -> Optional[ChatMessage]

从某个键删除最后一条消息。

参数

名称	类型	描述	默认值
key	str	用于删除最后一条消息的键。	必需

返回值

类型	描述
可选[ChatMessage]	可选[ChatMessage]：已删除的消息，如果列表为空则为 None。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

async def async_delete_last_message(self, key: str) -> Optional[ChatMessage]:
    """
    Delete the last message from a key.

    Args:
        key (str): The key to delete the last message from.

    Returns:
        Optional[ChatMessage]: The deleted message, or None if the list is empty.

    """
    deleted_message = await self._async_redis_client.rpop(key)
    if deleted_message:
        return ChatMessage.parse_raw(deleted_message)
    return None

get_keys

get_keys() -> List[str]

获取所有键。

返回值

类型	描述
列表[str]	列表[str]：Redis 存储中的所有键的列表。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

def get_keys(self) -> List[str]:
    """
    Get all keys.

    Returns:
        List[str]: A list of all keys in the Redis store.

    """
    keys = self._sync_redis_client.keys("*")
    return keys if isinstance(keys, list) else [keys]

async_get_keys async

async_get_keys() -> List[str]

获取所有键。

返回值

类型	描述
列表[str]	列表[str]：Redis 存储中的所有键的列表。

源代码位于 llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-upstash/llama_index/storage/chat_store/upstash/base.py

async def async_get_keys(self) -> List[str]:
    """
    Get all keys.

    Returns:
        List[str]: A list of all keys in the Redis store.

    """
    keys = await self._async_redis_client.keys("*")
    return keys if isinstance(keys, list) else [keys]