from __future__ import annotations import contextlib import json import logging import queue import re import threading import time import uuid from typing import TYPE_CHECKING, Any import clickzetta # type: ignore from pydantic import BaseModel, model_validator if TYPE_CHECKING: from clickzetta.connector.v0.connection import Connection # type: ignore from configs import dify_config from core.rag.datasource.vdb.field import Field, parse_metadata_json from core.rag.datasource.vdb.vector_base import BaseVector from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory from core.rag.embedding.embedding_base import Embeddings from core.rag.models.document import Document from models.dataset import Dataset logger = logging.getLogger(__name__) # ClickZetta Lakehouse Vector Database Configuration class ClickzettaConfig(BaseModel): """ Configuration class for Clickzetta connection. """ username: str password: str instance: str service: str = "api.clickzetta.com" workspace: str = "quick_start" vcluster: str = "default_ap" schema_name: str = "dify" # Renamed to avoid shadowing BaseModel.schema # Advanced settings batch_size: int = 20 # Reduced batch size to avoid large SQL statements enable_inverted_index: bool = True # Enable inverted index for full-text search analyzer_type: str = "chinese" # Analyzer type for full-text search: keyword, english, chinese, unicode analyzer_mode: str = "smart" # Analyzer mode: max_word, smart vector_distance_function: str = "cosine_distance" # l2_distance or cosine_distance @model_validator(mode="before") @classmethod def validate_config(cls, values: dict[str, Any]): """ Validate the configuration values. """ if not values.get("username"): raise ValueError("config CLICKZETTA_USERNAME is required") if not values.get("password"): raise ValueError("config CLICKZETTA_PASSWORD is required") if not values.get("instance"): raise ValueError("config CLICKZETTA_INSTANCE is required") if not values.get("service"): raise ValueError("config CLICKZETTA_SERVICE is required") if not values.get("workspace"): raise ValueError("config CLICKZETTA_WORKSPACE is required") if not values.get("vcluster"): raise ValueError("config CLICKZETTA_VCLUSTER is required") if not values.get("schema_name"): raise ValueError("config CLICKZETTA_SCHEMA is required") return values class ClickzettaConnectionPool: """ Global connection pool for ClickZetta connections. Manages connection reuse across ClickzettaVector instances. """ _instance: ClickzettaConnectionPool | None = None _lock = threading.Lock() def __init__(self): self._pools: dict[str, list[tuple[Connection, float]]] = {} # config_key -> [(connection, last_used_time)] self._pool_locks: dict[str, threading.Lock] = {} self._max_pool_size = 5 # Maximum connections per configuration self._connection_timeout = 300 # 5 minutes timeout self._cleanup_thread: threading.Thread | None = None self._shutdown = False self._start_cleanup_thread() @classmethod def get_instance(cls) -> ClickzettaConnectionPool: """Get singleton instance of connection pool.""" if cls._instance is None: with cls._lock: if cls._instance is None: cls._instance = cls() return cls._instance def _get_config_key(self, config: ClickzettaConfig) -> str: """Generate unique key for connection configuration.""" return ( f"{config.username}:{config.instance}:{config.service}:" f"{config.workspace}:{config.vcluster}:{config.schema_name}" ) def _create_connection(self, config: ClickzettaConfig) -> Connection: """Create a new ClickZetta connection.""" max_retries = 3 retry_delay = 1.0 for attempt in range(max_retries): try: connection = clickzetta.connect( username=config.username, password=config.password, instance=config.instance, service=config.service, workspace=config.workspace, vcluster=config.vcluster, schema=config.schema_name, ) # Configure connection session settings self._configure_connection(connection) logger.debug("Created new ClickZetta connection (attempt %d/%d)", attempt + 1, max_retries) return connection except Exception: logger.exception("ClickZetta connection attempt %d/%d failed", attempt + 1, max_retries) if attempt < max_retries - 1: time.sleep(retry_delay * (2**attempt)) else: raise raise RuntimeError(f"Failed to create ClickZetta connection after {max_retries} attempts") def _configure_connection(self, connection: Connection): """Configure connection session settings.""" try: with connection.cursor() as cursor: # Temporarily suppress ClickZetta client logging to reduce noise clickzetta_logger = logging.getLogger("clickzetta") original_level = clickzetta_logger.level clickzetta_logger.setLevel(logging.WARNING) try: # Use quote mode for string literal escaping cursor.execute("SET cz.sql.string.literal.escape.mode = 'quote'") # Apply performance optimization hints performance_hints = [ # Vector index optimization "SET cz.storage.parquet.vector.index.read.memory.cache = true", "SET cz.storage.parquet.vector.index.read.local.cache = false", # Query optimization "SET cz.sql.table.scan.push.down.filter = true", "SET cz.sql.table.scan.enable.ensure.filter = true", "SET cz.storage.always.prefetch.internal = true", "SET cz.optimizer.generate.columns.always.valid = true", "SET cz.sql.index.prewhere.enabled = true", # Storage optimization "SET cz.storage.parquet.enable.io.prefetch = false", "SET cz.optimizer.enable.mv.rewrite = false", "SET cz.sql.dump.as.lz4 = true", "SET cz.optimizer.limited.optimization.naive.query = true", "SET cz.sql.table.scan.enable.push.down.log = false", "SET cz.storage.use.file.format.local.stats = false", "SET cz.storage.local.file.object.cache.level = all", # Job execution optimization "SET cz.sql.job.fast.mode = true", "SET cz.storage.parquet.non.contiguous.read = true", "SET cz.sql.compaction.after.commit = true", ] for hint in performance_hints: cursor.execute(hint) finally: # Restore original logging level clickzetta_logger.setLevel(original_level) except Exception: logger.exception("Failed to configure connection, continuing with defaults") def _is_connection_valid(self, connection: Connection) -> bool: """Check if connection is still valid.""" try: with connection.cursor() as cursor: cursor.execute("SELECT 1") return True except Exception: return False def get_connection(self, config: ClickzettaConfig) -> Connection: """Get a connection from the pool or create a new one.""" config_key = self._get_config_key(config) # Ensure pool lock exists if config_key not in self._pool_locks: with self._lock: if config_key not in self._pool_locks: self._pool_locks[config_key] = threading.Lock() self._pools[config_key] = [] with self._pool_locks[config_key]: pool = self._pools[config_key] current_time = time.time() # Try to reuse existing connection while pool: connection, last_used = pool.pop(0) # Check if connection is not expired and still valid if current_time - last_used < self._connection_timeout and self._is_connection_valid(connection): logger.debug("Reusing ClickZetta connection from pool") return connection else: # Connection expired or invalid, close it with contextlib.suppress(Exception): connection.close() # No valid connection found, create new one return self._create_connection(config) def return_connection(self, config: ClickzettaConfig, connection: Connection): """Return a connection to the pool.""" config_key = self._get_config_key(config) if config_key not in self._pool_locks: # Pool was cleaned up, just close the connection with contextlib.suppress(Exception): connection.close() return with self._pool_locks[config_key]: pool = self._pools[config_key] # Only return to pool if not at capacity and connection is valid if len(pool) < self._max_pool_size and self._is_connection_valid(connection): pool.append((connection, time.time())) logger.debug("Returned ClickZetta connection to pool") else: # Pool full or connection invalid, close it with contextlib.suppress(Exception): connection.close() def _cleanup_expired_connections(self): """Clean up expired connections from all pools.""" current_time = time.time() with self._lock: for config_key in list(self._pools.keys()): if config_key not in self._pool_locks: continue with self._pool_locks[config_key]: pool = self._pools[config_key] valid_connections = [] for connection, last_used in pool: if current_time - last_used < self._connection_timeout: valid_connections.append((connection, last_used)) else: with contextlib.suppress(Exception): connection.close() self._pools[config_key] = valid_connections def _start_cleanup_thread(self): """Start background thread for connection cleanup.""" def cleanup_worker(): while not self._shutdown: try: time.sleep(60) # Cleanup every minute if not self._shutdown: self._cleanup_expired_connections() except Exception: logger.exception("Error in connection pool cleanup") self._cleanup_thread = threading.Thread(target=cleanup_worker, daemon=True) self._cleanup_thread.start() def shutdown(self): """Shutdown connection pool and close all connections.""" self._shutdown = True with self._lock: for config_key in list(self._pools.keys()): if config_key not in self._pool_locks: continue with self._pool_locks[config_key]: pool = self._pools[config_key] for connection, _ in pool: with contextlib.suppress(Exception): connection.close() pool.clear() class ClickzettaVector(BaseVector): """ Clickzetta vector storage implementation. """ # Class-level write queue and lock for serializing writes _write_queue: queue.Queue | None = None _write_thread: threading.Thread | None = None _write_lock = threading.Lock() _shutdown = False def __init__(self, collection_name: str, config: ClickzettaConfig): super().__init__(collection_name) self._config = config self._table_name = collection_name.replace("-", "_").lower() # Ensure valid table name self._connection_pool = ClickzettaConnectionPool.get_instance() self._init_write_queue() def _get_connection(self) -> Connection: """Get a connection from the pool.""" return self._connection_pool.get_connection(self._config) def _return_connection(self, connection: Connection): """Return a connection to the pool.""" self._connection_pool.return_connection(self._config, connection) class ConnectionContext: """Context manager for borrowing and returning connections.""" def __init__(self, vector_instance: ClickzettaVector): self.vector = vector_instance self.connection: Connection | None = None def __enter__(self) -> Connection: self.connection = self.vector._get_connection() return self.connection def __exit__(self, exc_type, exc_val, exc_tb): if self.connection: self.vector._return_connection(self.connection) def get_connection_context(self) -> ClickzettaVector.ConnectionContext: """Get a connection context manager.""" return self.ConnectionContext(self) def _parse_metadata(self, raw_metadata: str, row_id: str): """ Parse metadata from JSON string with proper error handling and fallback. Args: raw_metadata: Raw JSON string from database row_id: Row ID for fallback document_id Returns: Parsed metadata dict with guaranteed required fields """ try: if raw_metadata: # First parse may yield a string (double-encoded JSON) so use json.loads first_pass = json.loads(raw_metadata) # Handle double-encoded JSON if isinstance(first_pass, str): metadata = parse_metadata_json(first_pass) elif isinstance(first_pass, dict): metadata = first_pass else: metadata = {} else: metadata = {} except (json.JSONDecodeError, ValueError, TypeError): logger.exception("JSON parsing failed for metadata") # Fallback: extract document_id with regex doc_id_match = re.search(r'"document_id":\s*"([^"]+)"', raw_metadata or "") metadata = {"document_id": doc_id_match.group(1)} if doc_id_match else {} # Ensure required fields are set metadata["doc_id"] = row_id # segment id # Ensure document_id exists (critical for Dify's format_retrieval_documents) if "document_id" not in metadata: metadata["document_id"] = row_id # fallback to segment id return metadata @classmethod def _init_write_queue(cls): """Initialize the write queue and worker thread.""" with cls._write_lock: if cls._write_queue is None: cls._write_queue = queue.Queue() cls._write_thread = threading.Thread(target=cls._write_worker, daemon=True) cls._write_thread.start() logger.info("Started Clickzetta write worker thread") @classmethod def _write_worker(cls): """Worker thread that processes write tasks sequentially.""" while not cls._shutdown: try: # Get task from queue with timeout if cls._write_queue is not None: task = cls._write_queue.get(timeout=1) if task is None: # Shutdown signal break # Execute the write task func, args, kwargs, result_queue = task try: result = func(*args, **kwargs) result_queue.put((True, result)) except (RuntimeError, ValueError, TypeError, ConnectionError) as e: logger.exception("Write task failed") result_queue.put((False, e)) finally: cls._write_queue.task_done() else: break except queue.Empty: continue except (RuntimeError, ValueError, TypeError, ConnectionError) as e: logger.exception("Write worker error") def _execute_write(self, func, *args, **kwargs): """Execute a write operation through the queue.""" if ClickzettaVector._write_queue is None: raise RuntimeError("Write queue not initialized") result_queue: queue.Queue[tuple[bool, Any]] = queue.Queue() ClickzettaVector._write_queue.put((func, args, kwargs, result_queue)) # Wait for result success, result = result_queue.get() if not success: raise result return result def get_type(self) -> str: """Return the vector database type.""" return "clickzetta" def _ensure_connection(self) -> Connection: """Get a connection from the pool.""" return self._get_connection() def _table_exists(self) -> bool: """Check if the table exists.""" try: with self.get_connection_context() as connection: with connection.cursor() as cursor: cursor.execute(f"DESC {self._config.schema_name}.{self._table_name}") return True except Exception as e: error_message = str(e).lower() # Handle ClickZetta specific "table or view not found" errors if any( phrase in error_message for phrase in ["table or view not found", "czlh-42000", "semantic analysis exception"] ): logger.debug("Table %s.%s does not exist", self._config.schema_name, self._table_name) return False else: # For other connection/permission errors, log warning but return False to avoid blocking cleanup logger.exception( "Table existence check failed for %s.%s, assuming it doesn't exist", self._config.schema_name, self._table_name, ) return False def create(self, texts: list[Document], embeddings: list[list[float]], **kwargs): """Create the collection and add initial documents.""" # Execute table creation through write queue to avoid concurrent conflicts self._execute_write(self._create_table_and_indexes, embeddings) # Add initial texts if texts: self.add_texts(texts, embeddings, **kwargs) def _create_table_and_indexes(self, embeddings: list[list[float]]): """Create table and indexes (executed in write worker thread).""" # Check if table already exists to avoid unnecessary index creation if self._table_exists(): logger.info("Table %s.%s already exists, skipping creation", self._config.schema_name, self._table_name) return # Create table with vector and metadata columns dimension = len(embeddings[0]) if embeddings else 768 create_table_sql = f""" CREATE TABLE IF NOT EXISTS {self._config.schema_name}.{self._table_name} ( id STRING NOT NULL COMMENT 'Unique document identifier', {Field.CONTENT_KEY} STRING NOT NULL COMMENT 'Document text content for search and retrieval', {Field.METADATA_KEY} JSON COMMENT 'Document metadata including source, type, and other attributes', {Field.VECTOR} VECTOR(FLOAT, {dimension}) NOT NULL COMMENT 'High-dimensional embedding vector for semantic similarity search', PRIMARY KEY (id) ) COMMENT 'Dify RAG knowledge base vector storage table for document embeddings and content' """ with self.get_connection_context() as connection: with connection.cursor() as cursor: cursor.execute(create_table_sql) logger.info("Created table %s.%s", self._config.schema_name, self._table_name) # Create vector index self._create_vector_index(cursor) # Create inverted index for full-text search if enabled if self._config.enable_inverted_index: self._create_inverted_index(cursor) def _create_vector_index(self, cursor): """Create HNSW vector index for similarity search.""" # Use a fixed index name based on table and column name index_name = f"idx_{self._table_name}_vector" # First check if an index already exists on this column try: cursor.execute(f"SHOW INDEX FROM {self._config.schema_name}.{self._table_name}") existing_indexes = cursor.fetchall() for idx in existing_indexes: # Check if vector index already exists on the embedding column if Field.VECTOR in str(idx).lower(): logger.info("Vector index already exists on column %s", Field.VECTOR) return except (RuntimeError, ValueError) as e: logger.warning("Failed to check existing indexes: %s", e) index_sql = f""" CREATE VECTOR INDEX IF NOT EXISTS {index_name} ON TABLE {self._config.schema_name}.{self._table_name}({Field.VECTOR}) PROPERTIES ( "distance.function" = "{self._config.vector_distance_function}", "scalar.type" = "f32", "m" = "16", "ef.construction" = "128" ) """ try: cursor.execute(index_sql) logger.info("Created vector index: %s", index_name) except (RuntimeError, ValueError) as e: error_msg = str(e).lower() if "already exists" in error_msg or "already has index" in error_msg or "with the same type" in error_msg: logger.info("Vector index already exists: %s", e) else: logger.exception("Failed to create vector index") raise def _create_inverted_index(self, cursor): """Create inverted index for full-text search.""" # Use a fixed index name based on table name to avoid duplicates index_name = f"idx_{self._table_name}_text" # Check if an inverted index already exists on this column try: cursor.execute(f"SHOW INDEX FROM {self._config.schema_name}.{self._table_name}") existing_indexes = cursor.fetchall() for idx in existing_indexes: idx_str = str(idx).lower() # More precise check: look for inverted index specifically on the content column if ( "inverted" in idx_str and Field.CONTENT_KEY.lower() in idx_str and (index_name.lower() in idx_str or f"idx_{self._table_name}_text" in idx_str) ): logger.info("Inverted index already exists on column %s: %s", Field.CONTENT_KEY, idx) return except (RuntimeError, ValueError) as e: logger.warning("Failed to check existing indexes: %s", e) index_sql = f""" CREATE INVERTED INDEX IF NOT EXISTS {index_name} ON TABLE {self._config.schema_name}.{self._table_name} ({Field.CONTENT_KEY}) PROPERTIES ( "analyzer" = "{self._config.analyzer_type}", "mode" = "{self._config.analyzer_mode}" ) """ try: cursor.execute(index_sql) logger.info("Created inverted index: %s", index_name) except (RuntimeError, ValueError) as e: error_msg = str(e).lower() # Handle ClickZetta specific error messages if ( "already exists" in error_msg or "already has index" in error_msg or "with the same type" in error_msg or "cannot create inverted index" in error_msg ) and "already has index" in error_msg: logger.info("Inverted index already exists on column %s", Field.CONTENT_KEY) # Try to get the existing index name for logging try: cursor.execute(f"SHOW INDEX FROM {self._config.schema_name}.{self._table_name}") existing_indexes = cursor.fetchall() for idx in existing_indexes: if "inverted" in str(idx).lower() and Field.CONTENT_KEY.lower() in str(idx).lower(): logger.info("Found existing inverted index: %s", idx) break except (RuntimeError, ValueError): pass else: logger.warning("Failed to create inverted index: %s", e) # Continue without inverted index - full-text search will fall back to LIKE def add_texts(self, documents: list[Document], embeddings: list[list[float]], **kwargs) -> list[str]: """Add documents with embeddings to the collection.""" if not documents: return [] batch_size = self._config.batch_size total_batches = (len(documents) + batch_size - 1) // batch_size added_ids = [] for i in range(0, len(documents), batch_size): batch_docs = documents[i : i + batch_size] batch_embeddings = embeddings[i : i + batch_size] batch_doc_ids = [] for doc in batch_docs: metadata = doc.metadata if isinstance(doc.metadata, dict) else {} batch_doc_ids.append(self._safe_doc_id(metadata.get("doc_id", str(uuid.uuid4())))) added_ids.extend(batch_doc_ids) # Execute batch insert through write queue self._execute_write( self._insert_batch, batch_docs, batch_embeddings, batch_doc_ids, i, batch_size, total_batches ) return added_ids def _insert_batch( self, batch_docs: list[Document], batch_embeddings: list[list[float]], batch_doc_ids: list[str], batch_index: int, batch_size: int, total_batches: int, ): """Insert a batch of documents using parameterized queries (executed in write worker thread).""" if not batch_docs or not batch_embeddings: logger.warning("Empty batch provided, skipping insertion") return if len(batch_docs) != len(batch_embeddings): logger.error("Mismatch between docs (%d) and embeddings (%d)", len(batch_docs), len(batch_embeddings)) return # Prepare data for parameterized insertion data_rows = [] vector_dimension = len(batch_embeddings[0]) if batch_embeddings and batch_embeddings[0] else 768 for doc, embedding, doc_id in zip(batch_docs, batch_embeddings, batch_doc_ids): # Optimized: minimal checks for common case, fallback for edge cases metadata = doc.metadata if isinstance(doc.metadata, dict) else {} # Fast path for JSON serialization try: metadata_json = json.dumps(metadata, ensure_ascii=True) except (TypeError, ValueError): logger.warning("JSON serialization failed, using empty dict") metadata_json = "{}" content = doc.page_content or "" # According to ClickZetta docs, vector should be formatted as array string # for external systems: '[1.0, 2.0, 3.0]' vector_str = "[" + ",".join(map(str, embedding)) + "]" data_rows.append([doc_id, content, metadata_json, vector_str]) # Check if we have any valid data to insert if not data_rows: logger.warning("No valid documents to insert in batch %d/%d", batch_index // batch_size + 1, total_batches) return # Use parameterized INSERT with executemany for better performance and security # Cast JSON and VECTOR in SQL, pass raw data as parameters columns = f"id, {Field.CONTENT_KEY}, {Field.METADATA_KEY}, {Field.VECTOR}" insert_sql = ( f"INSERT INTO {self._config.schema_name}.{self._table_name} ({columns}) " f"VALUES (?, ?, CAST(? AS JSON), CAST(? AS VECTOR({vector_dimension})))" ) with self.get_connection_context() as connection: with connection.cursor() as cursor: try: # Set session-level hints for batch insert operations # Note: executemany doesn't support hints parameter, so we set them as session variables # Temporarily suppress ClickZetta client logging to reduce noise clickzetta_logger = logging.getLogger("clickzetta") original_level = clickzetta_logger.level clickzetta_logger.setLevel(logging.WARNING) try: cursor.execute("SET cz.sql.job.fast.mode = true") cursor.execute("SET cz.sql.compaction.after.commit = true") cursor.execute("SET cz.storage.always.prefetch.internal = true") finally: # Restore original logging level clickzetta_logger.setLevel(original_level) cursor.executemany(insert_sql, data_rows) logger.info( "Inserted batch %d/%d (%d valid docs using parameterized query with VECTOR(%d) cast)", batch_index // batch_size + 1, total_batches, len(data_rows), vector_dimension, ) except (RuntimeError, ValueError, TypeError, ConnectionError): logger.exception("Parameterized SQL execution failed for %d documents", len(data_rows)) logger.exception("SQL template: %s", insert_sql) logger.exception("Sample data row: %s", data_rows[0] if data_rows else "None") raise def text_exists(self, id: str) -> bool: """Check if a document exists by ID.""" # Check if table exists first if not self._table_exists(): return False safe_id = self._safe_doc_id(id) with self.get_connection_context() as connection: with connection.cursor() as cursor: cursor.execute( f"SELECT COUNT(*) FROM {self._config.schema_name}.{self._table_name} WHERE id = ?", binding_params=[safe_id], ) result = cursor.fetchone() return result[0] > 0 if result else False def delete_by_ids(self, ids: list[str]): """Delete documents by IDs.""" if not ids: return # Check if table exists before attempting delete if not self._table_exists(): logger.warning("Table %s.%s does not exist, skipping delete", self._config.schema_name, self._table_name) return # Execute delete through write queue self._execute_write(self._delete_by_ids_impl, ids) def _delete_by_ids_impl(self, ids: list[str]): """Implementation of delete by IDs (executed in write worker thread).""" safe_ids = [self._safe_doc_id(id) for id in ids] # Use parameterized query to prevent SQL injection placeholders = ",".join("?" for _ in safe_ids) sql = f"DELETE FROM {self._config.schema_name}.{self._table_name} WHERE id IN ({placeholders})" with self.get_connection_context() as connection: with connection.cursor() as cursor: cursor.execute(sql, binding_params=safe_ids) def delete_by_metadata_field(self, key: str, value: str): """Delete documents by metadata field.""" # Check if table exists before attempting delete if not self._table_exists(): logger.warning("Table %s.%s does not exist, skipping delete", self._config.schema_name, self._table_name) return # Execute delete through write queue self._execute_write(self._delete_by_metadata_field_impl, key, value) def _delete_by_metadata_field_impl(self, key: str, value: str): """Implementation of delete by metadata field (executed in write worker thread).""" with self.get_connection_context() as connection: with connection.cursor() as cursor: # Using JSON path to filter with parameterized query # Note: JSON path requires literal key name, cannot be parameterized # Use json_extract_string function for ClickZetta compatibility sql = ( f"DELETE FROM {self._config.schema_name}.{self._table_name} " f"WHERE json_extract_string({Field.METADATA_KEY}, '$.{key}') = ?" ) cursor.execute(sql, binding_params=[value]) def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]: """Search for documents by vector similarity.""" # Check if table exists first if not self._table_exists(): logger.warning( "Table %s.%s does not exist, returning empty results", self._config.schema_name, self._table_name, ) return [] top_k = kwargs.get("top_k", 10) score_threshold = kwargs.get("score_threshold", 0.0) document_ids_filter = kwargs.get("document_ids_filter") # Handle filter parameter from canvas (workflow) _ = kwargs.get("filter", {}) # Build filter clause filter_clauses = [] if document_ids_filter: safe_doc_ids = [str(id).replace("'", "''") for id in document_ids_filter] doc_ids_str = ",".join(f"'{id}'" for id in safe_doc_ids) # Use json_extract_string function for ClickZetta compatibility filter_clauses.append(f"json_extract_string({Field.METADATA_KEY}, '$.document_id') IN ({doc_ids_str})") # No need for dataset_id filter since each dataset has its own table # Add distance threshold based on distance function vector_dimension = len(query_vector) if self._config.vector_distance_function == "cosine_distance": # For cosine distance, smaller is better (0 = identical, 2 = opposite) distance_func = "COSINE_DISTANCE" if score_threshold > 0: query_vector_str = f"CAST('[{self._format_vector_simple(query_vector)}]' AS VECTOR({vector_dimension}))" filter_clauses.append(f"{distance_func}({Field.VECTOR}, {query_vector_str}) < {2 - score_threshold}") else: # For L2 distance, smaller is better distance_func = "L2_DISTANCE" if score_threshold > 0: query_vector_str = f"CAST('[{self._format_vector_simple(query_vector)}]' AS VECTOR({vector_dimension}))" filter_clauses.append(f"{distance_func}({Field.VECTOR}, {query_vector_str}) < {score_threshold}") where_clause = " AND ".join(filter_clauses) if filter_clauses else "1=1" # Execute vector search query query_vector_str = f"CAST('[{self._format_vector_simple(query_vector)}]' AS VECTOR({vector_dimension}))" search_sql = f""" SELECT id, {Field.CONTENT_KEY}, {Field.METADATA_KEY}, {distance_func}({Field.VECTOR}, {query_vector_str}) AS distance FROM {self._config.schema_name}.{self._table_name} WHERE {where_clause} ORDER BY distance LIMIT {top_k} """ documents = [] with self.get_connection_context() as connection: with connection.cursor() as cursor: # Use hints parameter for vector search optimization search_hints = { "hints": { "sdk.job.timeout": 60, # Increase timeout for vector search "cz.sql.job.fast.mode": True, "cz.storage.parquet.vector.index.read.memory.cache": True, } } cursor.execute(search_sql, search_hints) results = cursor.fetchall() for row in results: # Parse metadata using centralized method metadata = self._parse_metadata(row[2], row[0]) # Add score based on distance if self._config.vector_distance_function == "cosine_distance": metadata["score"] = 1 - (row[3] / 2) else: metadata["score"] = 1 / (1 + row[3]) doc = Document(page_content=row[1], metadata=metadata) documents.append(doc) return documents def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: """Search for documents using full-text search with inverted index.""" if not self._config.enable_inverted_index: logger.warning("Full-text search is not enabled. Enable inverted index in config.") return [] # Check if table exists first if not self._table_exists(): logger.warning( "Table %s.%s does not exist, returning empty results", self._config.schema_name, self._table_name, ) return [] top_k = kwargs.get("top_k", 10) document_ids_filter = kwargs.get("document_ids_filter") # Handle filter parameter from canvas (workflow) _ = kwargs.get("filter", {}) # Build filter clause filter_clauses = [] if document_ids_filter: safe_doc_ids = [str(id).replace("'", "''") for id in document_ids_filter] doc_ids_str = ",".join(f"'{id}'" for id in safe_doc_ids) # Use json_extract_string function for ClickZetta compatibility filter_clauses.append(f"json_extract_string({Field.METADATA_KEY}, '$.document_id') IN ({doc_ids_str})") # No need for dataset_id filter since each dataset has its own table # Use match_all function for full-text search # match_all requires all terms to be present # Use simple quote escaping for MATCH_ALL since it needs to be in the WHERE clause escaped_query = query.replace("'", "''") filter_clauses.append(f"MATCH_ALL({Field.CONTENT_KEY}, '{escaped_query}')") where_clause = " AND ".join(filter_clauses) # Execute full-text search query search_sql = f""" SELECT id, {Field.CONTENT_KEY}, {Field.METADATA_KEY} FROM {self._config.schema_name}.{self._table_name} WHERE {where_clause} LIMIT {top_k} """ documents = [] with self.get_connection_context() as connection: with connection.cursor() as cursor: try: # Use hints parameter for full-text search optimization fulltext_hints = { "hints": { "sdk.job.timeout": 30, # Timeout for full-text search "cz.sql.job.fast.mode": True, "cz.sql.index.prewhere.enabled": True, } } cursor.execute(search_sql, fulltext_hints) results = cursor.fetchall() for row in results: # Parse metadata from JSON string (may be double-encoded) try: if row[2]: # First parse may yield a string (double-encoded JSON) first_pass = json.loads(row[2]) if isinstance(first_pass, str): metadata = parse_metadata_json(first_pass) elif isinstance(first_pass, dict): metadata = first_pass else: metadata = {} else: metadata = {} except (json.JSONDecodeError, ValueError, TypeError): logger.exception("JSON parsing failed") # Fallback: extract document_id with regex doc_id_match = re.search(r'"document_id":\s*"([^"]+)"', str(row[2] or "")) metadata = {"document_id": doc_id_match.group(1)} if doc_id_match else {} # Ensure required fields are set metadata["doc_id"] = row[0] # segment id # Ensure document_id exists (critical for Dify's format_retrieval_documents) if "document_id" not in metadata: metadata["document_id"] = row[0] # fallback to segment id # Add a relevance score for full-text search metadata["score"] = 1.0 # Clickzetta doesn't provide relevance scores doc = Document(page_content=row[1], metadata=metadata) documents.append(doc) except (RuntimeError, ValueError, TypeError, ConnectionError): logger.exception("Full-text search failed") # Fallback to LIKE search if full-text search fails return self._search_by_like(query, **kwargs) return documents def _search_by_like(self, query: str, **kwargs: Any) -> list[Document]: """Fallback search using LIKE operator.""" # Check if table exists first if not self._table_exists(): logger.warning( "Table %s.%s does not exist, returning empty results", self._config.schema_name, self._table_name, ) return [] top_k = kwargs.get("top_k", 10) document_ids_filter = kwargs.get("document_ids_filter") # Handle filter parameter from canvas (workflow) _ = kwargs.get("filter", {}) # Build filter clause filter_clauses = [] if document_ids_filter: safe_doc_ids = [str(id).replace("'", "''") for id in document_ids_filter] doc_ids_str = ",".join(f"'{id}'" for id in safe_doc_ids) # Use json_extract_string function for ClickZetta compatibility filter_clauses.append(f"json_extract_string({Field.METADATA_KEY}, '$.document_id') IN ({doc_ids_str})") # No need for dataset_id filter since each dataset has its own table # Escape special characters for LIKE clause to prevent SQL injection from libs.helper import escape_like_pattern escaped_query = escape_like_pattern(query).replace("'", "''") filter_clauses.append(f"{Field.CONTENT_KEY} LIKE '%{escaped_query}%' ESCAPE '\\\\'") where_clause = " AND ".join(filter_clauses) search_sql = f""" SELECT id, {Field.CONTENT_KEY}, {Field.METADATA_KEY} FROM {self._config.schema_name}.{self._table_name} WHERE {where_clause} LIMIT {top_k} """ documents = [] with self.get_connection_context() as connection: with connection.cursor() as cursor: # Use hints parameter for LIKE search optimization like_hints = { "hints": { "sdk.job.timeout": 20, # Timeout for LIKE search "cz.sql.job.fast.mode": True, } } cursor.execute(search_sql, like_hints) results = cursor.fetchall() for row in results: # Parse metadata using centralized method metadata = self._parse_metadata(row[2], row[0]) metadata["score"] = 0.5 # Lower score for LIKE search doc = Document(page_content=row[1], metadata=metadata) documents.append(doc) return documents def delete(self): """Delete the entire collection.""" with self.get_connection_context() as connection: with connection.cursor() as cursor: cursor.execute(f"DROP TABLE IF EXISTS {self._config.schema_name}.{self._table_name}") def _format_vector_simple(self, vector: list[float]) -> str: """Simple vector formatting for SQL queries.""" return ",".join(map(str, vector)) def _safe_doc_id(self, doc_id: str) -> str: """Ensure doc_id is safe for SQL and doesn't contain special characters.""" if not doc_id: return str(uuid.uuid4()) # Remove or replace potentially problematic characters safe_id = str(doc_id) # Only allow alphanumeric, hyphens, underscores safe_id = "".join(c for c in safe_id if c.isalnum() or c in "-_") if not safe_id: # If all characters were removed return str(uuid.uuid4()) return safe_id[:255] # Limit length class ClickzettaVectorFactory(AbstractVectorFactory): """Factory for creating Clickzetta vector instances.""" def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> BaseVector: """Initialize a Clickzetta vector instance.""" # Get configuration from environment variables or dataset config config = ClickzettaConfig( username=dify_config.CLICKZETTA_USERNAME or "", password=dify_config.CLICKZETTA_PASSWORD or "", instance=dify_config.CLICKZETTA_INSTANCE or "", service=dify_config.CLICKZETTA_SERVICE or "api.clickzetta.com", workspace=dify_config.CLICKZETTA_WORKSPACE or "quick_start", vcluster=dify_config.CLICKZETTA_VCLUSTER or "default_ap", schema_name=dify_config.CLICKZETTA_SCHEMA or "dify", batch_size=dify_config.CLICKZETTA_BATCH_SIZE or 100, enable_inverted_index=dify_config.CLICKZETTA_ENABLE_INVERTED_INDEX or True, analyzer_type=dify_config.CLICKZETTA_ANALYZER_TYPE or "chinese", analyzer_mode=dify_config.CLICKZETTA_ANALYZER_MODE or "smart", vector_distance_function=dify_config.CLICKZETTA_VECTOR_DISTANCE_FUNCTION or "cosine_distance", ) # Use dataset collection name as table name collection_name = Dataset.gen_collection_name_by_id(dataset.id).lower() return ClickzettaVector(collection_name=collection_name, config=config)