#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

import contextlib

import logging

import os

import socket

import tempfile

import unittest

from dataclasses import dataclass

from typing import Callable

from typing import List

from typing import Optional

from typing import cast

from apache_beam.ml.rag.types import Chunk

from apache_beam.ml.rag.utils import retry_with_backoff

# pylint: disable=ungrouped-imports

try:

import yaml

from pymilvus import CollectionSchema

from pymilvus import FieldSchema

from pymilvus import MilvusClient

from pymilvus.exceptions import MilvusException

from pymilvus.milvus_client import IndexParams

from testcontainers.core.config import testcontainers_config

from testcontainers.core.generic import DbContainer

from testcontainers.milvus import MilvusContainer

from apache_beam.ml.rag.enrichment.milvus_search import MilvusConnectionParameters

except ImportError as e:

raise unittest.SkipTest(f'RAG test util dependencies not installed: {str(e)}')

_LOGGER = logging.getLogger(__name__)

@dataclass

class VectorDBContainerInfo:

"""Container information for vector database test instances.

Holds connection details and container reference for testing with

vector databases like Milvus in containerized environments.

"""

container: DbContainer

host: str

port: int

user: str = ""

password: str = ""

token: str = ""

id: str = "default"

@property

def uri(self) -> str:

return f"http://{self.host}:{self.port}"

class TestHelpers:

@staticmethod

def find_free_port():

"""Find a free port on the local machine."""

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:

# Bind to port 0, which asks OS to assign a free port.

s.bind(('', 0))

s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

# Return the port number assigned by OS.

return s.getsockname()[1]

class CustomMilvusContainer(MilvusContainer):

"""Custom Milvus container with configurable ports and environment setup.

Extends MilvusContainer to provide custom port binding and environment

configuration for testing with standalone Milvus instances.

"""

def __init__( # pylint: disable=bad-super-call

self,

image: str,

service_container_port,

healthcheck_container_port,

**kwargs,

) -> None:

# Skip the parent class's constructor and go straight to

# GenericContainer.

super(

MilvusContainer,

self,

).__init__(

image=image, **kwargs)

self.port = service_container_port

self.healthcheck_port = healthcheck_container_port

self.with_exposed_ports(service_container_port, healthcheck_container_port)

# Get free host ports.

service_host_port = TestHelpers.find_free_port()

healthcheck_host_port = TestHelpers.find_free_port()

# Bind container and host ports.

self.with_bind_ports(service_container_port, service_host_port)

self.with_bind_ports(healthcheck_container_port, healthcheck_host_port)

self.cmd = "milvus run standalone"

# Set environment variables needed for Milvus.

envs = {

"ETCD_USE_EMBED": "true",

"ETCD_DATA_DIR": "/var/lib/milvus/etcd",

"COMMON_STORAGETYPE": "local",

"METRICS_PORT": str(healthcheck_container_port)

}

for env, value in envs.items():

self.with_env(env, value)

class MilvusTestHelpers:

"""Helper utilities for testing Milvus vector database operations.

Provides static methods for managing test containers, configuration files,

and chunk comparison utilities for Milvus-based integration tests.

"""

# IMPORTANT: When upgrading the Milvus server version, ensure the pymilvus

# Python SDK client in setup.py is updated to match. Referring to the Milvus

# release notes compatibility matrix at

# https://milvus.io/docs/release_notes.md or PyPI at

# https://pypi.org/project/pymilvus/ for version compatibility.

# Example: Milvus v2.6.0 requires pymilvus==2.6.0 (exact match required).

@staticmethod

def _wait_for_milvus_grpc(uri: str) -> None:

"""Wait until Milvus accepts RPCs.

Docker may report started before gRPC is ready.

"""

def list_collections_probe():

client = MilvusClient(uri=uri)

try:

client.list_collections()

finally:

client.close()

retry_with_backoff(

list_collections_probe,

max_retries=25,

retry_delay=2.0,

retry_backoff_factor=1.2,

operation_name="Milvus client connection after container start",

exception_types=(MilvusException, ))

@staticmethod

def start_db_container(

image="milvusdb/milvus:v2.5.10",

max_vec_fields=5,

vector_client_max_retries=3,

tc_max_retries=None) -> Optional[VectorDBContainerInfo]:

service_container_port = TestHelpers.find_free_port()

healthcheck_container_port = TestHelpers.find_free_port()

user_yaml_creator = MilvusTestHelpers.create_user_yaml

with user_yaml_creator(service_container_port, max_vec_fields) as cfg:

info = None

original_tc_max_tries = testcontainers_config.max_tries

if tc_max_retries is not None:

testcontainers_config.max_tries = tc_max_retries

for i in range(vector_client_max_retries):

vector_db_container: Optional[CustomMilvusContainer] = None

try:

vector_db_container = CustomMilvusContainer(

image=image,

service_container_port=service_container_port,

healthcheck_container_port=healthcheck_container_port)

mapped_container = vector_db_container.with_volume_mapping(

cfg, "/milvus/configs/user.yaml")

assert mapped_container is not None

running_container: CustomMilvusContainer = mapped_container

vector_db_container = running_container

running_container.start()

host = running_container.get_container_host_ip()

port = running_container.get_exposed_port(service_container_port)

info = VectorDBContainerInfo(running_container, host, port)

MilvusTestHelpers._wait_for_milvus_grpc(info.uri)

_LOGGER.info(

"milvus db container started successfully on %s.", info.uri)

break

except Exception as e:

stdout_logs = stderr_logs = ""

if vector_db_container is not None:

raw_out, raw_err = vector_db_container.get_logs()

stdout_logs = raw_out.decode("utf-8")

stderr_logs = raw_err.decode("utf-8")

_LOGGER.warning(

"Retry %d/%d: Failed to start Milvus DB container. Reason: %s. "

"STDOUT logs:\n%s\nSTDERR logs:\n%s",

i + 1,

vector_client_max_retries,

e,

stdout_logs,

stderr_logs)

if i == vector_client_max_retries - 1:

_LOGGER.error(

"Unable to start milvus db container for I/O tests after %d "

"retries. Tests cannot proceed. STDOUT logs:\n%s\n"

"STDERR logs:\n%s",

vector_client_max_retries,

stdout_logs,

stderr_logs)

raise e

finally:

testcontainers_config.max_tries = original_tc_max_tries

return info

@staticmethod

def stop_db_container(db_info: VectorDBContainerInfo):

if db_info is None:

_LOGGER.warning("Milvus db info is None. Skipping stop operation.")

return

_LOGGER.debug("Stopping milvus db container.")

db_info.container.stop()

_LOGGER.info("milvus db container stopped successfully.")

@staticmethod

def initialize_db_with_data(

connc_params: MilvusConnectionParameters, config: dict):

# Open the connection to the milvus db with retry.

def create_client():

return MilvusClient(**connc_params.__dict__)

client = retry_with_backoff(

create_client,

max_retries=5,

retry_delay=2.0,

operation_name="Test Milvus client connection",

exception_types=(MilvusException, ))

# Configure schema.

field_schemas: List[FieldSchema] = cast(List[FieldSchema], config["fields"])

schema = CollectionSchema(

fields=field_schemas, functions=config["functions"])

# Create collection with the schema.

collection_name = config["collection_name"]

index_function: Callable[[], IndexParams] = cast(

Callable[[], IndexParams], config["index"])

client.create_collection(

collection_name=collection_name,

schema=schema,

index_params=index_function())

# Assert that collection was created.

collection_error = f"Expected collection '{collection_name}' to be created."

assert client.has_collection(collection_name), collection_error

# Gather all fields we have excluding 'sparse_embedding_bm25' special field.

fields = list(map(lambda field: field.name, field_schemas))

# Prep data for indexing. Currently we can't insert sparse vectors for BM25

# sparse embedding field as it would be automatically generated by Milvus

# through the registered BM25 function.

data_ready_to_index = []

for doc in config["corpus"]:

item = {}

for field in fields:

if field.startswith("dense_embedding"):

item[field] = doc["dense_embedding"]

elif field == "sparse_embedding_inner_product":

item[field] = doc["sparse_embedding"]

elif field == "sparse_embedding_bm25":

# It is automatically generated by Milvus from the content field.

continue

else:

item[field] = doc[field]

data_ready_to_index.append(item)

# Index data.

result = client.insert(

collection_name=collection_name, data=data_ready_to_index)

# Assert that the intended data has been properly indexed.

insertion_err = f'failed to insert the {result["insert_count"]} data points'

assert result["insert_count"] == len(data_ready_to_index), insertion_err

# Release the collection from memory. It will be loaded lazily when the

# enrichment handler is invoked.

client.release_collection(collection_name)

# Close the connection to the Milvus database, as no further preparation

# operations are needed before executing the enrichment handler.

client.close()

return collection_name

@staticmethod

@contextlib.contextmanager

def create_user_yaml(service_port: int, max_vector_field_num=5):

"""Creates a temporary user.yaml file for Milvus configuration.

This user yaml file overrides Milvus default configurations. It sets

the Milvus service port to the specified container service port. The

default for maxVectorFieldNum is 4, but we need 5

(one unique field for each metric).

Args:

service_port: Port number for the Milvus service.

max_vector_field_num: Max number of vec fields allowed per collection.

Yields:

str: Path to the created temporary yaml file.

"""

with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml',

delete=False) as temp_file:

# Define the content for user.yaml.

user_config = {

'proxy': {

'maxVectorFieldNum': max_vector_field_num, 'port': service_port

},

'etcd': {

'use': {

'embed': True

}, 'data': {

'dir': '/var/lib/milvus/etcd'

}

}

}

# Write the content to the file.

yaml.dump(user_config, temp_file, default_flow_style=False)

path = temp_file.name

try:

yield path

finally:

if os.path.exists(path):

os.remove(path)

@staticmethod

def assert_chunks_equivalent(

actual_chunks: List[Chunk], expected_chunks: List[Chunk]):

"""assert_chunks_equivalent checks for presence rather than exact match"""

# Sort both lists by ID to ensure consistent ordering.

actual_sorted = sorted(actual_chunks, key=lambda c: c.id)

expected_sorted = sorted(expected_chunks, key=lambda c: c.id)

actual_len = len(actual_sorted)

expected_len = len(expected_sorted)

err_msg = (

f"Different number of chunks, actual: {actual_len}, "

f"expected: {expected_len}")

assert actual_len == expected_len, err_msg

for actual, expected in zip(actual_sorted, expected_sorted):

# Assert that IDs match.

assert actual.id == expected.id

# Assert that dense embeddings match.

err_msg = f"Dense embedding mismatch for chunk {actual.id}"

assert actual.dense_embedding == expected.dense_embedding, err_msg

# Assert that sparse embeddings match.

err_msg = f"Sparse embedding mismatch for chunk {actual.id}"

assert actual.sparse_embedding == expected.sparse_embedding, err_msg

# Assert that text content match.

err_msg = f"Text Content mismatch for chunk {actual.id}"

assert actual.content.text == expected.content.text, err_msg

# For enrichment_data, be more flexible.

# If "expected" has values for enrichment_data but actual doesn't, that's

# acceptable since vector search results can vary based on many factors

# including implementation details, vector database state, and slight

# variations in similarity calculations.

# First ensure the enrichment data key exists.

err_msg = f"Missing enrichment_data key in chunk {actual.id}"

assert 'enrichment_data' in actual.metadata, err_msg

# For enrichment_data, ensure consistent ordering of results.

actual_data = actual.metadata['enrichment_data']

expected_data = expected.metadata['enrichment_data']

# If actual has enrichment data, then perform detailed validation.

if actual_data:

# Ensure the id key exist.

err_msg = f"Missing id key in metadata {actual.id}"

assert 'id' in actual_data, err_msg

# Validate IDs have consistent ordering.

actual_ids = sorted(actual_data['id'])

expected_ids = sorted(expected_data['id'])

err_msg = f"IDs in enrichment_data don't match for chunk {actual.id}"

assert actual_ids == expected_ids, err_msg

# Ensure the distance key exist.

err_msg = f"Missing distance key in metadata {actual.id}"

assert 'distance' in actual_data, err_msg

# Validate distances exist and have same length as IDs.

actual_distances = actual_data['distance']

expected_distances = expected_data['distance']

err_msg = (

"Number of distances doesn't match number of IDs for "

f"chunk {actual.id}")

assert len(actual_distances) == len(expected_distances), err_msg

# Ensure the fields key exist.

err_msg = f"Missing fields key in metadata {actual.id}"

assert 'fields' in actual_data, err_msg

# Validate fields have consistent content.

# Sort fields by 'id' to ensure consistent ordering.

actual_fields_sorted = sorted(

actual_data['fields'], key=lambda f: f.get('id', 0))

expected_fields_sorted = sorted(

expected_data['fields'], key=lambda f: f.get('id', 0))

# Compare field IDs.

actual_field_ids = [f.get('id') for f in actual_fields_sorted]

expected_field_ids = [f.get('id') for f in expected_fields_sorted]

err_msg = f"Field IDs don't match for chunk {actual.id}"

assert actual_field_ids == expected_field_ids, err_msg

# Compare field content.

for a_f, e_f in zip(actual_fields_sorted, expected_fields_sorted):

# Ensure the id key exist.

err_msg = f"Missing id key in metadata.fields {actual.id}"

assert 'id' in a_f, err_msg

err_msg = f"Field ID mismatch chunk {actual.id}"

assert a_f['id'] == e_f['id'], err_msg

# Validate field metadata.

err_msg = f"Field Metadata doesn't match for chunk {actual.id}"

assert a_f['metadata'] == e_f['metadata'], err_msg

if __name__ == '__main__':

unittest.main()