Class VectorDBClient

AskAgent/python/core/retriever.py:351–1051

Unified client for vector database operations. This class routes operations to the appropriate client implementation based on the database type specified in configuration.

class VectorDBClient:
    """
    Unified client for vector database operations. This class routes operations to the appropriate
    client implementation based on the database type specified in configuration.
    """

    def __init__(self, endpoint_name: str | None = None, query_params: dict[str, Any] | None = None):
        """
        Initialize the database client.

        Args:
            endpoint_name: Optional name of the endpoint to use (for backward compatibility)
            query_params: Optional query parameters for overriding endpoint
        """
        self.query_params = query_params or {}
        self.endpoint_name = endpoint_name  # Store the endpoint name
        self.db_type = None  # Will be set based on the primary endpoint

        # Check if query_params specifies a database endpoint override
        if self.query_params:
            # Check for 'db' or 'retrieval_backend' parameter
            param_endpoint = self.query_params.get('db') or self.query_params.get('retrieval_backend')
            if CONFIG.is_development_mode():
                print(f"[RETRIEVER] Development mode - param_endpoint from query_params: {param_endpoint}")
            if param_endpoint:
                # Handle case where param_endpoint might be a list
                if isinstance(param_endpoint, list):
                    if len(param_endpoint) > 0:
                        param_endpoint = param_endpoint[0]
                        logger.warning(f"'db' parameter was a list, using first element: {param_endpoint}")
                    else:
                        logger.error("'db' parameter is an empty list")
                        param_endpoint = None

                if param_endpoint:
                    logger.info(f"Using database endpoint from params: {param_endpoint}")
                    endpoint_name = param_endpoint

        # If specific endpoint requested, validate and use it
        if endpoint_name:
            try:
                if endpoint_name not in CONFIG.retrieval_endpoints:
                    available_endpoints = list(CONFIG.retrieval_endpoints.keys())
                    error_msg = f"Invalid endpoint: '{endpoint_name}'. Available endpoints: {', '.join(available_endpoints)}"
                    logger.error(error_msg)
                    raise ValueError(error_msg)
            except TypeError as e:
                # This can happen if endpoint_name is unhashable (e.g., a list)
                error_msg = f"Invalid endpoint name type: {type(endpoint_name).__name__}. Expected string, got: {endpoint_name}"
                logger.error(error_msg)
                raise ValueError(error_msg) from e

            # For backward compatibility, use only the specified endpoint
            endpoint_config = CONFIG.retrieval_endpoints[endpoint_name]
            self.enabled_endpoints = {endpoint_name: endpoint_config}
            self.db_type = endpoint_config.db_type  # Set db_type from the endpoint
            logger.info(f"VectorDBClient initialized with specific endpoint: {endpoint_name}")
        else: