apache_kafka

Apache Kafka Message Queue.

KafkaMessageQueueConfig

Bases: BaseSettings

Kafka message queue configuration.

Parameters:

Name	Type	Description	Default
type	Literal[str]		'kafka'
url	str		'localhost:9092'
host	str | None		None
port	int | None		None

Source code in llama_deploy/message_queues/apache_kafka.py

class KafkaMessageQueueConfig(BaseSettings):
    """Kafka message queue configuration."""

    model_config = SettingsConfigDict(env_prefix="KAFKA_")

    type: Literal["kafka"] = Field(default="kafka", exclude=True)
    url: str = DEFAULT_URL
    host: str | None = None
    port: int | None = None

    @model_validator(mode="after")
    def update_url(self) -> "KafkaMessageQueueConfig":
        if self.host and self.port:
            self.url = f"{self.host}:{self.port}"
        return self

KafkaMessageQueue

Bases: AbstractMessageQueue

Apache Kafka integration with aiokafka.

This class implements a traditional message broker using Apache Kafka. - Topics are created with N partitions - Consumers are registered to a single group to implement a competing consumer scheme where only one consumer subscribed to a topic gets the message - Default round-robin assignment is used

Attributes:

Name	Type	Description
url	str	The broker url string to connect to the Kafka server

Examples:

from llama_deploy.message_queues.apache_kafka import KafkaMessageQueue

message_queue = KafkaMessageQueue()  # uses the default url

Source code in llama_deploy/message_queues/apache_kafka.py

class KafkaMessageQueue(AbstractMessageQueue):
    """Apache Kafka integration with aiokafka.

    This class implements a traditional message broker using Apache Kafka.
        - Topics are created with N partitions
        - Consumers are registered to a single group to implement a competing
        consumer scheme where only one consumer subscribed to a topic gets the
        message
            - Default round-robin assignment is used

    Attributes:
        url (str): The broker url string to connect to the Kafka server

    Examples:
        ```python
        from llama_deploy.message_queues.apache_kafka import KafkaMessageQueue

        message_queue = KafkaMessageQueue()  # uses the default url
        ```
    """

    def __init__(self, config: KafkaMessageQueueConfig | None = None) -> None:
        self._config = config or KafkaMessageQueueConfig()
        self._kafka_consumers: dict[str, Any] = {}

    @classmethod
    def from_url_params(
        cls,
        host: str,
        port: int | None = None,
    ) -> "KafkaMessageQueue":
        """Convenience constructor from url params.

        Args:
            host (str): host for rabbitmq server
            port (Optional[int], optional): port for rabbitmq server. Defaults to None.

        Returns:
            KafkaMessageQueue: An Apache Kafka MessageQueue integration.
        """
        url = f"{host}:{port}" if port else f"{host}"
        return cls(KafkaMessageQueueConfig(url=url))

    def _create_new_topic(
        self,
        topic_name: str,
        num_partitions: int | None = None,
        replication_factor: int | None = None,
        **kwargs: Dict[str, Any],
    ) -> None:
        """Create a new topic.

        Use kafka-python-ng instead of aio-kafka as latter has issues with
        resolving api_version with broker.

        TODO: convert to aiokafka once this it is resolved there.
        """
        try:
            from kafka.admin import KafkaAdminClient, NewTopic
            from kafka.errors import TopicAlreadyExistsError
        except ImportError:
            raise ImportError(
                "kafka-python-ng is not installed. "
                "Please install it using `pip install kafka-python-ng`."
            )

        admin_client = KafkaAdminClient(bootstrap_servers=self._config.url)
        try:
            topic = NewTopic(
                name=topic_name,
                num_partitions=num_partitions or DEFAULT_TOPIC_PARTITIONS,
                replication_factor=replication_factor
                or DEFAULT_TOPIC_REPLICATION_FACTOR,
                **kwargs,
            )
            admin_client.create_topics(new_topics=[topic])
            logger.info(f"New topic {topic_name} created.")
        except TopicAlreadyExistsError:
            logger.info(f"Topic {topic_name} already exists.")
            pass

    async def _publish(self, message: QueueMessage, topic: str) -> Any:
        """Publish message to the queue."""
        try:
            from aiokafka import AIOKafkaProducer
        except ImportError:
            raise ImportError(
                "aiokafka is not installed. "
                "Please install it using `pip install aiokafka`."
            )

        producer = AIOKafkaProducer(bootstrap_servers=self._config.url)
        await producer.start()
        try:
            message_body = json.dumps(message.model_dump()).encode("utf-8")
            await producer.send_and_wait(topic, message_body)
            logger.info(f"published message {message.id_}")
        finally:
            await producer.stop()

    async def cleanup_local(
        self, message_types: List[str], *args: Any, **kwargs: Dict[str, Any]
    ) -> None:
        """Cleanup for local runs.

        Use kafka-python-ng instead of aio-kafka as latter has issues with
        resolving api_version with broker when using admin client.

        TODO: convert to aiokafka once this it is resolved there.
        """
        try:
            from kafka.admin import KafkaAdminClient
        except ImportError:
            raise ImportError(
                "aiokafka is not installed. "
                "Please install it using `pip install aiokafka`."
            )

        admin_client = KafkaAdminClient(bootstrap_servers=self._config.url)
        active_topics = admin_client.list_topics()
        topics_to_delete = [el for el in message_types if el in active_topics]
        admin_client.delete_consumer_groups(DEFAULT_GROUP_ID)
        if topics_to_delete:
            admin_client.delete_topics(topics_to_delete)

    async def deregister_consumer(self, consumer: BaseMessageQueueConsumer) -> Any:
        """Deregister a consumer."""
        if consumer.id_ in self._kafka_consumers:
            await self._kafka_consumers[consumer.id_].stop()

    async def launch_local(self) -> asyncio.Task:
        """Launch the message queue locally, in-process.

        Launches a dummy task.
        """
        return asyncio.create_task(self.processing_loop())

    async def launch_server(self) -> None:
        """Launch server."""
        pass

    async def processing_loop(self) -> None:
        pass

    async def register_consumer(
        self, consumer: BaseMessageQueueConsumer, topic: str | None = None
    ) -> Callable[..., Coroutine[Any, Any, None]]:
        """Register a new consumer."""
        try:
            from aiokafka import AIOKafkaConsumer
        except ImportError:
            raise ImportError(
                "aiokafka is not installed. "
                "Please install it using `pip install aiokafka`."
            )

        # register topic
        if topic is None:
            raise ValueError("Topic must be a valid string")

        if consumer.id_ in self._kafka_consumers:
            msg = f"Consumer {consumer.id_} already registered for topic {topic}"
            raise ValueError(msg)

        self._create_new_topic(topic)
        kafka_consumer = AIOKafkaConsumer(
            topic,
            bootstrap_servers=self._config.url,
            group_id=DEFAULT_GROUP_ID,
            auto_offset_reset="earliest",
        )
        self._kafka_consumers[consumer.id_] = kafka_consumer

        await kafka_consumer.start()

        logger.info(
            f"Registered consumer {consumer.id_} for message type '{consumer.message_type}' on topic '{topic}'",
        )

        async def start_consuming_callable() -> None:
            """StartConsumingCallable."""
            try:
                async for msg in kafka_consumer:
                    if msg.value is None:
                        raise RuntimeError("msg.value is None")
                    decoded_message = json.loads(msg.value.decode("utf-8"))
                    queue_message = QueueMessage.model_validate(decoded_message)
                    await consumer.process_message(queue_message)
            finally:
                stop_task = asyncio.create_task(kafka_consumer.stop())
                stop_task.add_done_callback(
                    lambda _: logger.info(
                        f"stopped kafka consumer {consumer.id_}: {consumer.message_type} on topic {topic}"
                    )
                )
                await asyncio.shield(stop_task)
                del self._kafka_consumers[consumer.id_]

        return start_consuming_callable

    def as_config(self) -> BaseModel:
        return self._config

from_url_params classmethod

from_url_params(host: str, port: int | None = None) -> KafkaMessageQueue

Convenience constructor from url params.

Parameters:

Name	Type	Description	Default
host	str	host for rabbitmq server	required
port	Optional[int]	port for rabbitmq server. Defaults to None.	None

Returns:

Name	Type	Description
KafkaMessageQueue	KafkaMessageQueue	An Apache Kafka MessageQueue integration.

Source code in llama_deploy/message_queues/apache_kafka.py

@classmethod
def from_url_params(
    cls,
    host: str,
    port: int | None = None,
) -> "KafkaMessageQueue":
    """Convenience constructor from url params.

    Args:
        host (str): host for rabbitmq server
        port (Optional[int], optional): port for rabbitmq server. Defaults to None.

    Returns:
        KafkaMessageQueue: An Apache Kafka MessageQueue integration.
    """
    url = f"{host}:{port}" if port else f"{host}"
    return cls(KafkaMessageQueueConfig(url=url))

cleanup_local async

cleanup_local(message_types: List[str], *args: Any, **kwargs: Dict[str, Any]) -> None

Cleanup for local runs.

Use kafka-python-ng instead of aio-kafka as latter has issues with resolving api_version with broker when using admin client.

TODO: convert to aiokafka once this it is resolved there.

Source code in llama_deploy/message_queues/apache_kafka.py

async def cleanup_local(
    self, message_types: List[str], *args: Any, **kwargs: Dict[str, Any]
) -> None:
    """Cleanup for local runs.

    Use kafka-python-ng instead of aio-kafka as latter has issues with
    resolving api_version with broker when using admin client.

    TODO: convert to aiokafka once this it is resolved there.
    """
    try:
        from kafka.admin import KafkaAdminClient
    except ImportError:
        raise ImportError(
            "aiokafka is not installed. "
            "Please install it using `pip install aiokafka`."
        )

    admin_client = KafkaAdminClient(bootstrap_servers=self._config.url)
    active_topics = admin_client.list_topics()
    topics_to_delete = [el for el in message_types if el in active_topics]
    admin_client.delete_consumer_groups(DEFAULT_GROUP_ID)
    if topics_to_delete:
        admin_client.delete_topics(topics_to_delete)

deregister_consumer async

deregister_consumer(consumer: BaseMessageQueueConsumer) -> Any

Deregister a consumer.

Source code in llama_deploy/message_queues/apache_kafka.py

async def deregister_consumer(self, consumer: BaseMessageQueueConsumer) -> Any:
    """Deregister a consumer."""
    if consumer.id_ in self._kafka_consumers:
        await self._kafka_consumers[consumer.id_].stop()

launch_local async

launch_local() -> Task

Launch the message queue locally, in-process.

Launches a dummy task.

Source code in llama_deploy/message_queues/apache_kafka.py

async def launch_local(self) -> asyncio.Task:
    """Launch the message queue locally, in-process.

    Launches a dummy task.
    """
    return asyncio.create_task(self.processing_loop())

launch_server async

launch_server() -> None

Launch server.

Source code in llama_deploy/message_queues/apache_kafka.py

async def launch_server(self) -> None:
    """Launch server."""
    pass

register_consumer async

register_consumer(consumer: BaseMessageQueueConsumer, topic: str | None = None) -> Callable[..., Coroutine[Any, Any, None]]

Register a new consumer.

Source code in llama_deploy/message_queues/apache_kafka.py

async def register_consumer(
    self, consumer: BaseMessageQueueConsumer, topic: str | None = None
) -> Callable[..., Coroutine[Any, Any, None]]:
    """Register a new consumer."""
    try:
        from aiokafka import AIOKafkaConsumer
    except ImportError:
        raise ImportError(
            "aiokafka is not installed. "
            "Please install it using `pip install aiokafka`."
        )

    # register topic
    if topic is None:
        raise ValueError("Topic must be a valid string")

    if consumer.id_ in self._kafka_consumers:
        msg = f"Consumer {consumer.id_} already registered for topic {topic}"
        raise ValueError(msg)

    self._create_new_topic(topic)
    kafka_consumer = AIOKafkaConsumer(
        topic,
        bootstrap_servers=self._config.url,
        group_id=DEFAULT_GROUP_ID,
        auto_offset_reset="earliest",
    )
    self._kafka_consumers[consumer.id_] = kafka_consumer

    await kafka_consumer.start()

    logger.info(
        f"Registered consumer {consumer.id_} for message type '{consumer.message_type}' on topic '{topic}'",
    )

    async def start_consuming_callable() -> None:
        """StartConsumingCallable."""
        try:
            async for msg in kafka_consumer:
                if msg.value is None:
                    raise RuntimeError("msg.value is None")
                decoded_message = json.loads(msg.value.decode("utf-8"))
                queue_message = QueueMessage.model_validate(decoded_message)
                await consumer.process_message(queue_message)
        finally:
            stop_task = asyncio.create_task(kafka_consumer.stop())
            stop_task.add_done_callback(
                lambda _: logger.info(
                    f"stopped kafka consumer {consumer.id_}: {consumer.message_type} on topic {topic}"
                )
            )
            await asyncio.shield(stop_task)
            del self._kafka_consumers[consumer.id_]

    return start_consuming_callable