"""Kafka-based chat message history by using confluent-kafka-python. confluent-kafka-python is under Apache 2.0 license. https://github.com/confluentinc/confluent-kafka-python """ from __future__ import annotations import json import logging import time from enum import Enum from typing import TYPE_CHECKING, List, Optional, Sequence from langchain_core.chat_history import BaseChatMessageHistory from langchain_core.messages import BaseMessage, message_to_dict, messages_from_dict if TYPE_CHECKING: from confluent_kafka import TopicPartition from confluent_kafka.admin import AdminClient logger = logging.getLogger(__name__) BOOTSTRAP_SERVERS_CONFIG = "bootstrap.servers" DEFAULT_TTL_MS = 604800000 # 7 days DEFAULT_REPLICATION_FACTOR = 1 DEFAULT_PARTITION = 3 class ConsumeStartPosition(Enum): """Consume start position for Kafka consumer to get chat history messages. LAST_CONSUMED: Continue from the last consumed offset. EARLIEST: Start consuming from the beginning. LATEST: Start consuming from the latest offset. """ LAST_CONSUMED = 1 EARLIEST = 2 LATEST = 3 def ensure_topic_exists( admin_client: AdminClient, topic_name: str, replication_factor: int, partition: int, ttl_ms: int, ) -> int: """Create topic if it doesn't exist, and return the number of partitions. If the topic already exists, we don't change the topic configuration. """ from confluent_kafka.admin import NewTopic try: topic_metadata = admin_client.list_topics().topics if topic_name in topic_metadata: num_partitions = len(topic_metadata[topic_name].partitions) logger.info( f"Topic {topic_name} already exists with {num_partitions} partitions" ) return num_partitions except Exception as e: logger.error(f"Failed to list topics: {e}") raise e topics = [ NewTopic( topic_name, num_partitions=partition, replication_factor=replication_factor, config={"retention.ms": str(ttl_ms)}, ) ] try: futures = admin_client.create_topics(topics) for _, f in futures.items(): f.result() # result is None logger.info(f"Topic {topic_name} created") except Exception as e: logger.error(f"Failed to create topic {topic_name}: {e}") raise e return partition class KafkaChatMessageHistory(BaseChatMessageHistory): """Chat message history stored in Kafka. Setup: Install ``confluent-kafka-python``. .. code-block:: bash pip install confluent_kafka Instantiate: .. code-block:: python from langchain_community.chat_message_histories import KafkaChatMessageHistory history = KafkaChatMessageHistory( session_id="your_session_id", bootstrap_servers="host:port", ) Add and retrieve messages: .. code-block:: python # Add messages history.add_messages([message1, message2, message3, ...]) # Retrieve messages message_batch_0 = history.messages # retrieve messages after message_batch_0 message_batch_1 = history.messages # Reset to beginning and retrieve messages messages_from_beginning = history.messages_from_beginning() Retrieving messages is stateful. Internally, it uses Kafka consumer to read. The consumed offset is maintained persistently. To retrieve messages, you can use the following methods: - `messages`: continue consuming chat messages from last one. - `messages_from_beginning`: reset the consumer to the beginning of the chat history and return messages. Optional parameters: 1. `max_message_count`: maximum number of messages to return. 2. `max_time_sec`: maximum time in seconds to wait for messages. - `messages_from_latest`: reset to end of the chat history and try consuming messages. Optional parameters same as above. - `messages_from_last_consumed`: continuing from the last consumed message, similar to `messages`. Optional parameters same as above. `max_message_count` and `max_time_sec` are used to avoid blocking indefinitely when retrieving messages. As a result, the method to retrieve messages may not return all messages. Change `max_message_count` and `max_time_sec` to retrieve all history messages. """ # noqa: E501 def __init__( self, session_id: str, bootstrap_servers: str, ttl_ms: int = DEFAULT_TTL_MS, replication_factor: int = DEFAULT_REPLICATION_FACTOR, partition: int = DEFAULT_PARTITION, ): """ Args: session_id: The ID for single chat session. It is used as Kafka topic name. bootstrap_servers: Comma-separated host/port pairs to establish connection to Kafka cluster https://kafka.apache.org/documentation.html#adminclientconfigs_bootstrap.servers ttl_ms: Time-to-live (milliseconds) for automatic expiration of entries. Default 7 days. -1 for no expiration. It translates to https://kafka.apache.org/documentation.html#topicconfigs_retention.ms replication_factor: The replication factor for the topic. Default 1. partition: The number of partitions for the topic. Default 3. """ try: from confluent_kafka import Producer from confluent_kafka.admin import AdminClient except (ImportError, ModuleNotFoundError): raise ImportError( "Could not import confluent_kafka package. " "Please install it with `pip install confluent_kafka`." ) self.session_id = session_id self.bootstrap_servers = bootstrap_servers self.admin_client = AdminClient({BOOTSTRAP_SERVERS_CONFIG: bootstrap_servers}) self.num_partitions = ensure_topic_exists( self.admin_client, session_id, replication_factor, partition, ttl_ms ) self.producer = Producer({BOOTSTRAP_SERVERS_CONFIG: bootstrap_servers}) def add_messages( self, messages: Sequence[BaseMessage], flush_timeout_seconds: float = 5.0, ) -> None: """Add messages to the chat history by producing to the Kafka topic.""" try: for message in messages: self.producer.produce( topic=self.session_id, value=json.dumps(message_to_dict(message)), ) message_remaining = self.producer.flush(flush_timeout_seconds) if message_remaining > 0: logger.warning(f"{message_remaining} messages are still in-flight.") except Exception as e: logger.error(f"Failed to add messages to Kafka: {e}") raise e def __read_messages( self, consume_start_pos: ConsumeStartPosition, max_message_count: Optional[int], max_time_sec: Optional[float], ) -> List[BaseMessage]: """Retrieve messages from Kafka topic for the session. Please note this method is stateful. Internally, it uses Kafka consumer to consume messages, and maintains the consumed offset. Args: consume_start_pos: Start position for Kafka consumer. max_message_count: Maximum number of messages to consume. max_time_sec: Time limit in seconds to consume messages. Returns: List of messages. """ from confluent_kafka import OFFSET_BEGINNING, OFFSET_END, Consumer consumer_config = { BOOTSTRAP_SERVERS_CONFIG: self.bootstrap_servers, "group.id": self.session_id, "auto.offset.reset": "latest" if consume_start_pos == ConsumeStartPosition.LATEST else "earliest", } def assign_beginning( assigned_consumer: Consumer, assigned_partitions: list[TopicPartition] ) -> None: for p in assigned_partitions: p.offset = OFFSET_BEGINNING assigned_consumer.assign(assigned_partitions) def assign_latest( assigned_consumer: Consumer, assigned_partitions: list[TopicPartition] ) -> None: for p in assigned_partitions: p.offset = OFFSET_END assigned_consumer.assign(assigned_partitions) messages: List[dict] = [] consumer = Consumer(consumer_config) try: if consume_start_pos == ConsumeStartPosition.EARLIEST: consumer.subscribe([self.session_id], on_assign=assign_beginning) elif consume_start_pos == ConsumeStartPosition.LATEST: consumer.subscribe([self.session_id], on_assign=assign_latest) else: consumer.subscribe([self.session_id]) start_time_sec = time.time() while True: if ( max_time_sec is not None and time.time() - start_time_sec > max_time_sec ): break if max_message_count is not None and len(messages) >= max_message_count: break message = consumer.poll(timeout=1.0) if message is None: # poll timeout continue if message.error() is not None: # error logger.error(f"Consumer error: {message.error()}") continue if message.value() is None: # empty value logger.warning("Empty message value") continue messages.append(json.loads(message.value())) except Exception as e: logger.error(f"Failed to consume messages from Kafka: {e}") raise e finally: consumer.close() return messages_from_dict(messages) def messages_from_beginning( self, max_message_count: Optional[int] = 5, max_time_sec: Optional[float] = 5.0 ) -> List[BaseMessage]: """Retrieve messages from Kafka topic from the beginning. This method resets the consumer to the beginning and consumes messages. Args: max_message_count: Maximum number of messages to consume. max_time_sec: Time limit in seconds to consume messages. Returns: List of messages. """ return self.__read_messages( consume_start_pos=ConsumeStartPosition.EARLIEST, max_message_count=max_message_count, max_time_sec=max_time_sec, ) def messages_from_latest( self, max_message_count: Optional[int] = 5, max_time_sec: Optional[float] = 5.0 ) -> List[BaseMessage]: """Reset to the end offset. Try to consume messages if available. Args: max_message_count: Maximum number of messages to consume. max_time_sec: Time limit in seconds to consume messages. Returns: List of messages. """ return self.__read_messages( consume_start_pos=ConsumeStartPosition.LATEST, max_message_count=max_message_count, max_time_sec=max_time_sec, ) def messages_from_last_consumed( self, max_message_count: Optional[int] = 5, max_time_sec: Optional[float] = 5.0 ) -> List[BaseMessage]: """Retrieve messages from Kafka topic from the last consumed message. Please note this method is stateful. Internally, it uses Kafka consumer to consume messages, and maintains the commit offset. Args: max_message_count: Maximum number of messages to consume. max_time_sec: Time limit in seconds to consume messages. Returns: List of messages. """ return self.__read_messages( consume_start_pos=ConsumeStartPosition.LAST_CONSUMED, max_message_count=max_message_count, max_time_sec=max_time_sec, ) @property def messages(self) -> List[BaseMessage]: # type: ignore[override] """ Retrieve the messages for the session, from Kafka topic continuously from last consumed message. This method is stateful and maintains consumed(committed) offset based on consumer group. Alternatively, use messages_from_last_consumed() with specified parameters. Use messages_from_beginning() to read from the earliest message. Use messages_from_latest() to read from the latest message. """ return self.messages_from_last_consumed() def clear(self) -> None: """Clear the chat history by deleting the Kafka topic.""" try: futures = self.admin_client.delete_topics([self.session_id]) for _, f in futures.items(): f.result() # result is None logger.info(f"Topic {self.session_id} deleted") except Exception as e: logger.error(f"Failed to delete topic {self.session_id}: {e}") raise e def close(self) -> None: """Release the resources. Nothing to be released at this moment. """ pass