diff --git a/CHANGELOG.md b/CHANGELOG.md index 8acd7f357..a5535768c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,17 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). ### Changed +- **Knowledge Base embedding model binding (breaking / migration)** + The Knowledge Base now treats the **Model Hub ID** as the single source of truth for embedding model identity: + - `collection_metadata.embedding_model_id` stores the Hub ID (trimmed; no other normalization). + - Embeddings tables are named by Hub ID: `embeddings_{to_model_tag(hub_id)}`. + - The `model` field stored alongside each embedding vector is the Hub ID. + + **Migration / backward compatibility:** Older deployments may have created embeddings tables using the provider `model_name` + (e.g. `embeddings_text-embedding-v4`). During search and embedding reads, the system will **try the new Hub-ID table first** + and automatically **fall back to the legacy table name** derived from the resolved `model_name` when the new table is missing. + Rebuild/inference helpers were updated to prefer Hub IDs when they can be resolved from Model Hub metadata. + - **Knowledge Base upload: default parse method (breaking)** The default parse method on the KB detail upload form is now `"default"` instead of `"pypdf"`. The backend chooses the parser by file type (e.g. .docx, .pdf). If you rely on the previous default (always use PyPDF), select `"pypdf"` explicitly in the parse method dropdown when uploading. diff --git a/src/xagent/config.py b/src/xagent/config.py index 91b80d1f9..9139df310 100644 --- a/src/xagent/config.py +++ b/src/xagent/config.py @@ -189,12 +189,7 @@ def get_lancedb_path() -> Path: Priority: 1. LANCEDB_PATH environment variable - 2. Default to ./data/lancedb (relative to cwd) - - .. warning:: - Default to ``./data/lancedb``, which is **relative** to cwd, **NOT** - relative to ``storage_root``. This behavior is kept for backward - compatibility but may change in the future (see proposal #246). + 2. Default to STORAGE_ROOT/data/lancedb Returns: Path object for LanceDB directory @@ -203,8 +198,8 @@ def get_lancedb_path() -> Path: if env_path: return Path(env_path) - # Default: ./data/lancedb - return Path("data/lancedb") + # Default: storage_root/data/lancedb + return get_storage_root() / "data" / "lancedb" def get_default_sqlite_db_path() -> str: diff --git a/src/xagent/core/tools/core/RAG_tools/chunk/chunk_document.py b/src/xagent/core/tools/core/RAG_tools/chunk/chunk_document.py index 10b8db81b..22870d47d 100644 --- a/src/xagent/core/tools/core/RAG_tools/chunk/chunk_document.py +++ b/src/xagent/core/tools/core/RAG_tools/chunk/chunk_document.py @@ -11,7 +11,6 @@ import pandas as pd -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.config import ( DEFAULT_IMAGE_CONTEXT_SIZE, DEFAULT_TABLE_CONTEXT_SIZE, @@ -23,12 +22,9 @@ DocumentValidationError, ) from ..core.schemas import ChunkStrategy -from ..LanceDB.schema_manager import ensure_chunks_table +from ..storage.factory import get_vector_index_store from ..utils.hash_utils import compute_chunk_hash -from ..utils.lancedb_query_utils import query_to_list from ..utils.metadata_utils import deserialize_metadata, serialize_metadata -from ..utils.string_utils import build_lancedb_filter_expression -from ..utils.user_permissions import UserPermissions from .chunk_strategies import ( apply_fixed_size_strategy, apply_markdown_strategy, @@ -109,14 +105,6 @@ def chunk_document( f"Starting document chunking: doc_id={doc_id}, strategy={chunk_strategy}" ) - # Get database connection - try: - conn = get_connection_from_env() - ensure_chunks_table(conn) - except Exception as e: - logger.error(f"Database connection failed: {e}") - raise DatabaseOperationError(f"Failed to connect to database: {e}") from e - # Validate chunk parameters _validate_chunk_params(chunk_strategy, params) @@ -251,8 +239,7 @@ def _chunks_exist( ) -> bool: """Check if chunk records already exist.""" try: - conn = get_connection_from_env() - table = conn.open_table("chunks") + vector_store = get_vector_index_store() # Build safe filter expression using utility function query_filters = { @@ -261,8 +248,7 @@ def _chunks_exist( "parse_hash": parse_hash, "config_hash": config_hash, } - filter_expr = build_lancedb_filter_expression(query_filters) - return bool(table.count_rows(filter_expr) > 0) + return vector_store.count_rows_or_zero("chunks", filters=query_filters) > 0 except Exception as e: logger.error(f"Failed to check chunk existence: {e}") raise DatabaseOperationError(f"Database query failed: {e}") from e @@ -293,8 +279,7 @@ def _get_existing_chunks( List of existing chunks accessible to the user """ try: - conn = get_connection_from_env() - table = conn.open_table("chunks") + vector_store = get_vector_index_store() # Build safe filter expression using utility function query_filters = { @@ -303,25 +288,28 @@ def _get_existing_chunks( "parse_hash": parse_hash, "config_hash": config_hash, } - base_filter_expr = build_lancedb_filter_expression(query_filters) - # Add user permission filter for multi-tenancy - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - - # Combine filters - if user_filter_expr and base_filter_expr: - filter_expr = f"({base_filter_expr}) and ({user_filter_expr})" - elif user_filter_expr: - filter_expr = user_filter_expr - else: - filter_expr = base_filter_expr - - # OPTIMIZATION: Use count_rows() for memory-efficient existence check - if table.count_rows(filter_expr) == 0: + # OPTIMIZATION: Use count_rows_or_zero() for memory-efficient existence check + if ( + vector_store.count_rows_or_zero( + "chunks", filters=query_filters, user_id=user_id, is_admin=is_admin + ) + == 0 + ): return [] - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - chunks_data = query_to_list(table.search().where(filter_expr)) + # Use iter_batches to load chunks + chunks_data = [] + for batch in vector_store.iter_batches( + table_name="chunks", + filters=query_filters, + user_id=user_id, + is_admin=is_admin, + ): + # Convert batch to pandas for easier row-by-row processing + batch_df = batch.to_pandas() + for _, row in batch_df.iterrows(): + chunks_data.append(row.to_dict()) # Convert to expected format with metadata deserialization # Arrow/to_list() returns None instead of NaN, so direct None check is sufficient @@ -372,8 +360,7 @@ def _load_paragraphs( ) -> List[Dict[str, Any]]: """Load parsed content from parses table.""" try: - conn = get_connection_from_env() - table = conn.open_table("parses") + vector_store = get_vector_index_store() # Build safe filter expression using utility function query_filters = { @@ -381,26 +368,29 @@ def _load_paragraphs( "doc_id": doc_id, "parse_hash": parse_hash, } - base_filter_expr = build_lancedb_filter_expression(query_filters) - # Add user permission filter for multi-tenancy - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - - # Combine filters - if user_filter_expr and base_filter_expr: - filter_expr = f"({base_filter_expr}) and ({user_filter_expr})" - elif user_filter_expr: - filter_expr = user_filter_expr - else: - filter_expr = base_filter_expr - - # First check if any parse exists using efficient count_rows - if table.count_rows(filter_expr) == 0: + # First check if any parse exists using efficient count_rows_or_zero + if ( + vector_store.count_rows_or_zero( + "parses", filters=query_filters, user_id=user_id, is_admin=is_admin + ) + == 0 + ): return [] - # Only load data if parse exists - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - records = query_to_list(table.search().where(filter_expr)) + # Load data using iter_batches + records = [] + for batch in vector_store.iter_batches( + table_name="parses", + filters=query_filters, + user_id=user_id, + is_admin=is_admin, + ): + # Convert batch to pandas for easier row-by-row processing + batch_df = batch.to_pandas() + for _, row in batch_df.iterrows(): + records.append(row.to_dict()) + if not records: return [] record = records[0] @@ -445,11 +435,8 @@ def _write_chunks_to_db( user_id: Optional[int] = None, is_admin: bool = False, ) -> bool: - """Write chunk records to database.""" + """Write chunk records to database using abstraction layer.""" try: - conn = get_connection_from_env() - table = conn.open_table("chunks") - rows = [] for chunk in chunks: text = chunk["text"] @@ -477,11 +464,10 @@ def _write_chunks_to_db( if not rows: return False - # Use merge_insert for efficient upsert operation - # This handles cases where chunks might already exist (idempotent operation) - table.merge_insert( - ["collection", "doc_id", "parse_hash", "chunk_id"] - ).when_matched_update_all().when_not_matched_insert_all().execute(rows) + # Use abstraction layer for upsert + vector_store = get_vector_index_store() + vector_store.upsert_chunks(rows) + logger.info( f"Chunk records written to database: doc_id={doc_id}, parse_hash={parse_hash}, config_hash={config_hash}" ) diff --git a/src/xagent/core/tools/core/RAG_tools/core/config.py b/src/xagent/core/tools/core/RAG_tools/core/config.py index e74d359ba..9d623121b 100644 --- a/src/xagent/core/tools/core/RAG_tools/core/config.py +++ b/src/xagent/core/tools/core/RAG_tools/core/config.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Final, Mapping, Sequence @@ -55,13 +56,42 @@ Set to 0 to disable any artificial throttling. """ +DEFAULT_LANCEDB_BATCH_SIZE: Final[int] = 1000 +"""Default batch size for embedding writes to LanceDB (env: LANCEDB_BATCH_SIZE).""" + +DEFAULT_VECTOR_STORE_SCAN_LIMIT: Final[int] = 10_000 +"""Default max rows scanned in vector-store document listing operations.""" + +DEFAULT_VECTOR_STORE_EXTENDED_SCAN_LIMIT: Final[int] = 1_000_000 +"""Higher limit for operations like listing all documents in a collection or deleting a collection.""" + # Reserved int64 lower bound for internal system sentinel values. MIN_INT64: Final[int] = -(2**63) +"""Minimum 64-bit integer, used as internal sentinel value.""" # Stable expression that always matches no rows for unauthenticated reads. UNAUTHENTICATED_NO_ACCESS_FILTER: Final[str] = ( "(user_id IS NULL and user_id IS NOT NULL)" ) +"""A stable LanceDB filter expression that always matches no rows.""" + +ENABLE_AUTO_EMBEDDINGS_MIGRATION: Final[bool] = ( + os.getenv("ENABLE_AUTO_EMBEDDINGS_MIGRATION", "false").lower() == "true" +) +""" +Enable automatic forward migration of legacy embeddings tables. + +When disabled (default), the system will not automatically migrate data from +legacy table names (embeddings_{model_name}) to new Hub ID-based names +(embeddings_{hub_id}). This prevents unexpected data movement and performance +impact during normal operations. + +To enable automatic migration, set the environment variable: + ENABLE_AUTO_EMBEDDINGS_MIGRATION=true + +Automatic migration should only be enabled during controlled maintenance windows +or when explicitly executing migration tools. +""" # Parameters that affect parse hash PARSE_PARAM_WHITELIST: Final[Sequence[str]] = ( diff --git a/src/xagent/core/tools/core/RAG_tools/core/schemas.py b/src/xagent/core/tools/core/RAG_tools/core/schemas.py index 83dc083a9..7c0b6cf4c 100644 --- a/src/xagent/core/tools/core/RAG_tools/core/schemas.py +++ b/src/xagent/core/tools/core/RAG_tools/core/schemas.py @@ -799,6 +799,29 @@ class HybridSearchResponse(BaseModel): ) +class IndexResult(BaseModel): + """Structured result from index creation operations. + + This model replaces the previous string-based return format for create_index, + providing type-safe access to index status, advice, and FTS enabled state. + + Attributes: + status: Index creation status (e.g., "index_ready", "readonly", "failed") + advice: Optional advice message for further actions + fts_enabled: Whether FTS index is actually enabled (separate from vector index) + """ + + model_config = ConfigDict(frozen=True) + + status: str = Field(..., description="Index creation status") + advice: Optional[str] = Field( + default=None, description="Human-readable index advice" + ) + fts_enabled: bool = Field( + default=False, description="Whether FTS index is enabled on text column" + ) + + class SearchConfig(BaseModel): """Configuration for the unified document search pipeline.""" @@ -1319,7 +1342,15 @@ def is_initialized(self) -> bool: @classmethod def from_storage(cls, data: dict) -> "CollectionInfo": - """Factory method to load from LanceDB, handling migration automatically.""" + """Load from storage dict with in-memory schema normalization. + + Legacy rows (e.g. ``schema_version`` missing / ``0.0.0``) are upgraded + **in memory only** via :func:`~.migration_utils.migrate_collection_metadata` + with ``infer_embedding=False`` so this path does **not** open LanceDB or + scan embedding tables (read-side-effect-free). For full migration with + embedding inference, call ``migrate_collection_metadata(data)`` explicitly + (e.g. admin repair or write pipeline). + """ import json import math @@ -1342,14 +1373,12 @@ def from_storage(cls, data: dict) -> "CollectionInfo": if isinstance(value, float) and math.isnan(value): data[key] = None - # 3. Check version and migrate if needed + # 3. Check version and migrate if needed (no DB access on read path) current_version = "1.0.0" data_version = data.get("schema_version", "0.0.0") if data_version < current_version: - data = migrate_collection_metadata(data) - # Note: In LanceDB, we don't auto-save migrated data here - # It will be saved when the collection is next updated + data = migrate_collection_metadata(data, infer_embedding=False) return cls(**data) diff --git a/src/xagent/core/tools/core/RAG_tools/file/register_document.py b/src/xagent/core/tools/core/RAG_tools/file/register_document.py index b15485f00..175bcad03 100644 --- a/src/xagent/core/tools/core/RAG_tools/file/register_document.py +++ b/src/xagent/core/tools/core/RAG_tools/file/register_document.py @@ -16,7 +16,6 @@ import pandas as pd -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.exceptions import ( ConfigurationError, DatabaseOperationError, @@ -24,15 +23,15 @@ HashComputationError, ) from ..core.schemas import RegisterDocumentRequest, RegisterDocumentResponse -from ..LanceDB.schema_manager import ensure_documents_table +from ..storage.factory import get_vector_index_store from ..utils import check_file_type, compute_file_hash from ..utils.string_utils import ( - build_lancedb_filter_expression, generate_deterministic_doc_id, ) logger = logging.getLogger(__name__) + # Public entry with explicit arguments (for LG/CLI/FastAPI). Returns plain dict. # Internally constructs Pydantic request and delegates to _register_document. @@ -156,25 +155,26 @@ def _register_document(request: RegisterDocumentRequest) -> RegisterDocumentResp except Exception as e: raise HashComputationError(f"Failed to compute content hash: {e}") from e - # LanceDB operations + # LanceDB operations using abstraction layer try: - # Get LanceDB connection - db = get_connection_from_env() - - # Ensure documents table exists - ensure_documents_table(db) - - # Open the documents table - table = db.open_table("documents") + vector_store = get_vector_index_store() - # Check if document already exists (for idempotency) + # Check if document already exists (for idempotency) using count_rows query_filters = { "collection": collection, "doc_id": doc_id, } - filter_expr = build_lancedb_filter_expression(query_filters) - - exists = table.count_rows(filter_expr) > 0 + # For existence check, use admin mode to see all records including legacy data + # count_rows_or_zero returns 0 if table doesn't exist + exists = ( + vector_store.count_rows_or_zero( + "documents", + filters=query_filters, + user_id=request.user_id, + is_admin=True, + ) + > 0 + ) # Prepare document record doc_record = { @@ -191,10 +191,8 @@ def _register_document(request: RegisterDocumentRequest) -> RegisterDocumentResp "user_id": request.user_id, # Add user_id for multi-tenancy } - # Use merge_insert for efficient upsert operation - table.merge_insert( - ["collection", "doc_id"] - ).when_matched_update_all().when_not_matched_insert_all().execute([doc_record]) + # Use abstraction layer for upsert + vector_store.upsert_documents([doc_record]) created = not exists @@ -213,11 +211,11 @@ def _register_document(request: RegisterDocumentRequest) -> RegisterDocumentResp def get_document(db_dir: str, collection: str, doc_id: str) -> Optional[Any]: - """Retrieve a document record from LanceDB. + """Retrieve a document record from LanceDB using abstraction layer. Args: - db_dir: LanceDB directory path + db_dir: LanceDB directory path (unused, kept for compatibility) collection: Collection name to filter by (only returns documents from this collection) doc_id: Document ID to retrieve @@ -228,19 +226,23 @@ def get_document(db_dir: str, collection: str, doc_id: str) -> Optional[Any]: DatabaseOperationError: If database operation fails """ try: - db = get_connection_from_env() - ensure_documents_table(db) - table = db.open_table("documents") + vector_store = get_vector_index_store() - filter_expr = build_lancedb_filter_expression( - {"collection": collection, "doc_id": doc_id} - ) - if table.count_rows(filter_expr) == 0: + # Check if document exists + query_filters = {"collection": collection, "doc_id": doc_id} + if vector_store.count_rows_or_zero("documents", filters=query_filters) == 0: return None - # Convert to dict and handle datetime - record = table.search().where(filter_expr).to_pandas().iloc[0].to_dict() - return record + # Use iter_batches to load the document + for batch in vector_store.iter_batches( + table_name="documents", + filters=query_filters, + ): + batch_df = batch.to_pandas() + for _, row in batch_df.iterrows(): + return row.to_dict() + + return None except Exception as e: raise DatabaseOperationError(f"Failed to retrieve document: {e}") from e @@ -249,10 +251,10 @@ def get_document(db_dir: str, collection: str, doc_id: str) -> Optional[Any]: def list_documents( db_dir: str, collection: str, limit: int = 100 ) -> list[Dict[str, Any]]: - """List documents in the collection. + """List documents in the collection using abstraction layer. Args: - db_dir: LanceDB directory path + db_dir: LanceDB directory path (unused, kept for compatibility) collection: Collection name to filter by (only documents in this KB are returned) limit: Maximum number of documents to return @@ -263,13 +265,25 @@ def list_documents( DatabaseOperationError: If database operation fails """ try: - db = get_connection_from_env() - ensure_documents_table(db) - table = db.open_table("documents") - - filter_expr = build_lancedb_filter_expression({"collection": collection}) - results = table.search().where(filter_expr).limit(limit).to_pandas() - return list(results.to_dict("records")) + vector_store = get_vector_index_store() + query_filters = {"collection": collection} + + results = [] + for batch in vector_store.iter_batches( + table_name="documents", + filters=query_filters, + user_id=None, + is_admin=True, # Use admin mode to see all documents including legacy data + ): + batch_df = batch.to_pandas() + for _, row in batch_df.iterrows(): + results.append(row.to_dict()) + if len(results) >= limit: + break + if len(results) >= limit: + break + + return results except Exception as e: raise DatabaseOperationError(f"Failed to list documents: {e}") from e diff --git a/src/xagent/core/tools/core/RAG_tools/generate/format_generation_prompt.py b/src/xagent/core/tools/core/RAG_tools/generate/format_generation_prompt.py index 7c429f6d6..2da4031fb 100644 --- a/src/xagent/core/tools/core/RAG_tools/generate/format_generation_prompt.py +++ b/src/xagent/core/tools/core/RAG_tools/generate/format_generation_prompt.py @@ -12,43 +12,39 @@ def format_generation_prompt( """Formats a prompt template and contexts into a single string for LLM input. This function takes a base prompt template and a string of formatted contexts, - and combines them into a single, cohesive prompt string suitable for - sending to a Large Language Model (LLM). It ensures that both the - prompt template and contexts are provided. + and combines them into a cohesive prompt string. If the template contains + a "{context}" placeholder, it will be replaced with the formatted contexts. + Otherwise, the contexts will be appended after the template. Args: - prompt_template: The base template for the prompt, which may include placeholders. - formatted_contexts: A string containing the relevant contexts, already - formatted for LLM input (e.g., from search results). + prompt_template: The base template for the prompt. + formatted_contexts: A string containing the relevant contexts. Returns: A single string representing the full prompt ready for LLM consumption. Raises: ConfigurationError: If `prompt_template` is empty. - - Examples: - >>> template = "Answer the question based on the following context: {context}" - >>> contexts = "Context: The capital of France is Paris." - >>> full_prompt = format_generation_prompt(template, contexts) - >>> print(full_prompt) - Answer the question based on the following context: {context} - - Context: - The capital of France is Paris. - - Answer: """ if not prompt_template: raise ConfigurationError("Prompt template cannot be empty.") + if not formatted_contexts: - # NOTE: Depending on the use case, empty contexts might be valid. - # For RAG, we generally expect contexts. logger.warning( "Formatted contexts are empty, which might lead to non-grounded generation." ) - full_prompt = f"{prompt_template}\n\nContext:\n{formatted_contexts}\n\nAnswer:" - logger.debug(f"Formatted prompt length: {len(full_prompt)} chars.") + # Check if the template has a placeholder for context + if "{context}" in prompt_template: + try: + full_prompt = prompt_template.format(context=formatted_contexts) + except (KeyError, ValueError) as e: + logger.error(f"Failed to format prompt template: {e}") + # Fallback to appending if formatting fails + full_prompt = f"{prompt_template}\n\nContext:\n{formatted_contexts}" + else: + # Default behavior: append context and answer marker + full_prompt = f"{prompt_template}\n\nContext:\n{formatted_contexts}\n\nAnswer:" + logger.debug(f"Formatted prompt length: {len(full_prompt)} chars.") return full_prompt diff --git a/src/xagent/core/tools/core/RAG_tools/management/collection_manager.py b/src/xagent/core/tools/core/RAG_tools/management/collection_manager.py index 6c308f0b5..2f77bfedb 100644 --- a/src/xagent/core/tools/core/RAG_tools/management/collection_manager.py +++ b/src/xagent/core/tools/core/RAG_tools/management/collection_manager.py @@ -12,12 +12,11 @@ from functools import wraps from typing import Any, Awaitable, Callable, Optional, TypeVar -from ......providers.vector_store.lancedb import DBConnection, get_connection_from_env from ..core.parser_registry import get_supported_parsers, validate_parser_compatibility from ..core.schemas import CollectionInfo -from ..LanceDB.schema_manager import ensure_collection_metadata_table +from ..storage.factory import get_metadata_store, get_vector_index_store from ..utils.model_resolver import resolve_embedding_adapter -from ..utils.string_utils import escape_lancedb_string +from ..utils.tag_mapping import register_tag_mapping T = TypeVar("T") @@ -134,17 +133,7 @@ class CollectionManager: """ def __init__(self) -> None: - self._conn: Optional[DBConnection] = None - - async def _get_connection(self) -> DBConnection: - """Lazy initialization of LanceDB connection. - - Returns: - The LanceDB connection instance - """ - if self._conn is None: - self._conn = get_connection_from_env() - return self._conn + self._metadata_store = get_metadata_store() async def get_collection(self, collection_name: str) -> CollectionInfo: """Get collection metadata from storage. @@ -158,27 +147,11 @@ async def get_collection(self, collection_name: str) -> CollectionInfo: Raises: ValueError: If collection not found """ - conn = await self._get_connection() - - # Ensure table exists before accessing - ensure_collection_metadata_table(conn) - try: - # Try to read from collection_metadata table - table = conn.open_table("collection_metadata") - # Use safe parameterized query to prevent SQL injection - safe_name = escape_lancedb_string(collection_name) - result = table.search().where(f"name = '{safe_name}'").to_pandas() - - if result.empty: - raise ValueError(f"Collection '{collection_name}' not found") - - # Convert to dict and deserialize - data = result.iloc[0].to_dict() - return CollectionInfo.from_storage(data) + return await self._metadata_store.get_collection(collection_name) except Exception as e: - # Table might not exist yet, or other LanceDB errors + # Table might not exist yet, or other backend errors logger.debug(f"Error reading collection {collection_name}: {e}") raise ValueError(f"Collection '{collection_name}' not found") @@ -205,62 +178,9 @@ async def _save_collection_with_retry( Raises: Exception: If all retry attempts fail """ - conn = await self._get_connection() - - # Ensure table exists before accessing - ensure_collection_metadata_table(conn) - for attempt in range(max_retries): try: - # Prepare data for storage - data = collection.to_storage() - data["updated_at"] = datetime.now(timezone.utc).replace( - tzinfo=None - ) # Fresh timestamp - - # Upsert to LanceDB: delete existing then add new - table = conn.open_table("collection_metadata") - safe_name = escape_lancedb_string(collection.name) - - # Check if collection already exists - existing = table.search().where(f"name = '{safe_name}'").to_pandas() - if not existing.empty: - # Delete existing record - table.delete(f"name = '{safe_name}'") - - # Add new record - # Ensure data strictly matches table schema to prevent LanceDB schema errors - # (e.g. "missing=[owners]" or "contains null values") - import pyarrow as pa # type: ignore[import-not-found] - - clean_data: dict[str, Any] = {} - for field in table.schema: - val = data.get(field.name) - if val is None: - # Provide default for missing or None values if not nullable - if not field.nullable: - if pa.types.is_string( - field.type - ) or pa.types.is_large_string(field.type): - clean_data[field.name] = "" - elif pa.types.is_integer(field.type): - clean_data[field.name] = 0 - elif pa.types.is_floating(field.type): - clean_data[field.name] = 0.0 - elif pa.types.is_boolean(field.type): - clean_data[field.name] = False - elif pa.types.is_timestamp(field.type): - clean_data[field.name] = datetime.now( - timezone.utc - ).replace(tzinfo=None) - else: - clean_data[field.name] = "" - else: - clean_data[field.name] = None - else: - clean_data[field.name] = val - - table.add([clean_data]) + await self._metadata_store.save_collection(collection) return except Exception as e: @@ -637,22 +557,18 @@ def resolve_effective_embedding_model_sync( raise -def rebuild_collection_metadata() -> None: +async def rebuild_collection_metadata() -> None: """Rebuild collection_metadata table from existing data. This function reads all collections from documents/parses/chunks/embeddings tables and creates corresponding entries in the collection_metadata table. Use this to migrate existing data when collection_metadata table is missing or outdated. - - This is a synchronous blocking operation. """ - from xagent.providers.vector_store.lancedb import get_connection_from_env - from . import collections # Get all existing collections (use is_admin=True to bypass user filtering) - result = collections.list_collections(is_admin=True) + result = await collections.list_collections(is_admin=True) if result.status != "success": logger.error(f"Failed to list collections: {result.message}") @@ -662,10 +578,49 @@ def rebuild_collection_metadata() -> None: return # Get connection and find embeddings tables - conn = get_connection_from_env() - table_names = conn.table_names() # type: ignore[attr-defined] + vector_store = get_vector_index_store() + table_names = vector_store.list_table_names() embeddings_tables = [t for t in table_names if t.startswith("embeddings_")] + # Build lookup from legacy/new table tags to Hub model IDs. + hub_tag_to_id: dict[str, tuple[str, Optional[int]]] = {} + try: + from xagent.core.model.model import EmbeddingModelConfig + + from ..LanceDB.model_tag_utils import to_model_tag + from ..utils.model_resolver import _get_or_init_model_hub + + hub = _get_or_init_model_hub() + if hub is not None: + for cfg in hub.list().values(): + if not isinstance(cfg, EmbeddingModelConfig): + continue + register_tag_mapping( + hub_tag_to_id, + to_model_tag(cfg.id), + (cfg.id, cfg.dimension), + get_identity=lambda item: item[0], + logger=logger, + ) + register_tag_mapping( + hub_tag_to_id, + to_model_tag(cfg.model_name), + (cfg.id, cfg.dimension), + get_identity=lambda item: item[0], + logger=logger, + ) + except Exception as e: + logger.warning( + "Model hub initialization failed during collection metadata rebuild: " + "error_type=%s, error_message=%s, fallback_behavior=%s, impact=%s", + type(e).__name__, + str(e), + "legacy_model_resolution", + "May use suboptimal model selection or missing embeddings", + exc_info=True, + ) + hub_tag_to_id = {} + # Save each collection to metadata table for collection in result.collections: try: @@ -676,30 +631,28 @@ def rebuild_collection_metadata() -> None: if collection.embeddings > 0: # Find which embeddings table has data for this collection for table_name in embeddings_tables: - table = conn.open_table(table_name) - count = table.count_rows( - f"collection = '{escape_lancedb_string(collection.name)}'" + # Use abstraction layer to count rows + count = vector_store.count_rows_or_zero( + table_name, + filters={"collection": collection.name}, + is_admin=True, ) if count > 0: - # Extract model name from table name - # Table names use underscores (e.g., embeddings_text_embedding_v4) - # Model IDs use hyphens (e.g., text-embedding-v4) - embedding_model_id = table_name.replace( - "embeddings_", "" - ).replace("_", "-") - - # Get vector dimension from schema - schema = table.schema - vector_field = schema.field("vector") - if hasattr(vector_field, "type"): - vector_type = vector_field.type - if hasattr(vector_type, "list_size"): - embedding_dimension = vector_type.list_size - else: - # Variable length list, get first row to infer dimension - sample = table.search().limit(1).to_pandas() - if not sample.empty and "vector" in sample.columns: - embedding_dimension = len(sample.iloc[0]["vector"]) + suffix = table_name.replace("embeddings_", "", 1) + # Prefer Hub ID mapping (single source of truth). + if suffix in hub_tag_to_id: + embedding_model_id, inferred_dim = hub_tag_to_id[suffix] + if inferred_dim is not None: + embedding_dimension = inferred_dim + else: + # Legacy fallback: best-effort reverse normalization. + embedding_model_id = suffix.replace("_", "-") + + # Use abstraction layer to get vector dimension from schema + table_dim = vector_store.get_vector_dimension(table_name) + if table_dim is not None: + embedding_dimension = table_dim + break # Update collection with embedding info diff --git a/src/xagent/core/tools/core/RAG_tools/management/collections.py b/src/xagent/core/tools/core/RAG_tools/management/collections.py index 1e8d61b72..698817824 100644 --- a/src/xagent/core/tools/core/RAG_tools/management/collections.py +++ b/src/xagent/core/tools/core/RAG_tools/management/collections.py @@ -6,7 +6,9 @@ from __future__ import annotations +import json import logging +import warnings as py_warnings from collections import defaultdict from datetime import datetime from typing import Any, Dict, List, Optional, Sequence, Set @@ -14,9 +16,10 @@ import pyarrow as pa # type: ignore from lancedb.db import DBConnection -from xagent.providers.vector_store.lancedb import get_connection_from_env - -from ..core.config import DEFAULT_LANCEDB_SCAN_BATCH_SIZE +from ..core.config import ( + DEFAULT_LANCEDB_SCAN_BATCH_SIZE, + DEFAULT_VECTOR_STORE_EXTENDED_SCAN_LIMIT, +) from ..core.schemas import ( CollectionInfo, CollectionOperationDetail, @@ -27,21 +30,16 @@ DocumentStats, DocumentStatsResult, DocumentSummary, + IngestionConfig, ListCollectionsResult, ) from ..LanceDB.model_tag_utils import embeddings_table_name -from ..LanceDB.schema_manager import ( - ensure_chunks_table, - ensure_collection_config_table, - ensure_documents_table, - ensure_ingestion_runs_table, - ensure_parses_table, -) from ..management.status import ( clear_ingestion_status, load_ingestion_status, write_ingestion_status, ) +from ..storage.factory import get_metadata_store, get_vector_index_store from ..utils.string_utils import build_lancedb_filter_expression, escape_lancedb_string from ..utils.user_permissions import UserPermissions from ..version_management.cascade_cleaner import cleanup_document_cascade @@ -62,6 +60,10 @@ def _iter_batches( ) -> Any: """Yield record batches from a LanceDB table while minimizing memory footprint. + .. deprecated:: + This function is deprecated. Use VectorIndexStore.iter_batches() instead. + This function will be removed in a future release. + This generator function iterates through a LanceDB table in batches to minimize memory usage, with support for user filtering and column selection. @@ -77,6 +79,11 @@ def _iter_batches( Yields: PyArrow RecordBatch objects containing the data """ + py_warnings.warn( + "_iter_batches is deprecated, use VectorIndexStore.iter_batches() instead", + DeprecationWarning, + stacklevel=2, + ) try: table = conn.open_table(table_name) @@ -195,6 +202,10 @@ def _count_rows( ) -> int: """Count rows in a LanceDB table while handling failures gracefully. + .. deprecated:: + This function is deprecated. Use VectorIndexStore.count_rows() instead. + This function will be removed in a future release. + This function counts rows in a LanceDB table with optional filters, returning 0 on any error and logging warnings. @@ -207,6 +218,11 @@ def _count_rows( Returns: Number of rows matching the filter, or 0 on error """ + py_warnings.warn( + "_count_rows is deprecated, use VectorIndexStore.count_rows() instead", + DeprecationWarning, + stacklevel=2, + ) try: table = conn.open_table(table_name) @@ -232,6 +248,10 @@ def _count_rows( def _list_table_names(conn: DBConnection, warnings: List[str]) -> List[str]: """Return available LanceDB table names with graceful degradation. + .. deprecated:: + This function is deprecated. Use VectorIndexStore.list_table_names() instead. + This function will be removed in a future release. + This function retrieves the list of table names from a LanceDB connection, handling errors gracefully by returning an empty list and logging warnings. @@ -242,6 +262,11 @@ def _list_table_names(conn: DBConnection, warnings: List[str]) -> List[str]: Returns: List of table names as strings, or empty list on error """ + py_warnings.warn( + "_list_table_names is deprecated, use VectorIndexStore.list_table_names() instead", + DeprecationWarning, + stacklevel=2, + ) try: table_names_fn = getattr(conn, "table_names") @@ -273,6 +298,10 @@ def _collect_doc_counts_for_collection( ) -> Dict[str, int]: """Aggregate per-document counts for the specified table within a collection. + .. deprecated:: + This function is deprecated. Use VectorIndexStore.aggregate_document_counts() instead. + This function will be removed in a future release. + This function iterates through batches of a table and counts records per document for a specific collection. @@ -288,6 +317,11 @@ def _collect_doc_counts_for_collection( Returns: Dictionary mapping document IDs to their counts """ + py_warnings.warn( + "_collect_doc_counts_for_collection is deprecated, use VectorIndexStore.aggregate_document_counts() instead", + DeprecationWarning, + stacklevel=2, + ) counts: Dict[str, int] = defaultdict(int) @@ -416,7 +450,52 @@ def _coerce_timestamp(value: Any) -> datetime | None: return None -def list_collections( +async def _load_collection_ingestion_configs( + collection_keys: List[str], + user_id: Optional[int], + is_admin: bool, +) -> Dict[str, IngestionConfig]: + """Load ingestion configs for the given collections using metadata store rules. + + Args: + collection_keys: Collection names returned by stats / document scan. + user_id: Caller user id; None is treated as 0 for non-admin lookups. + is_admin: When True, ``get_collection_config`` returns the latest config + per collection across tenants. + + Returns: + Map of collection name to parsed ingestion configuration. + """ + metadata_store = get_metadata_store() + collection_configs: Dict[str, IngestionConfig] = {} + # Handle user_id=None explicitly: admin mode keeps None (load all configs), + # non-admin mode converts to 0 (backward compatible) + if is_admin and user_id is None: + uid = None + else: + uid = 0 if user_id is None else user_id + for collection in collection_keys: + try: + config_json = await metadata_store.get_collection_config( + collection, uid, is_admin=is_admin + ) + if not config_json: + continue + try: + config_dict = json.loads(config_json) + collection_configs[collection] = IngestionConfig(**config_dict) + except Exception as e: + logger.warning( + "Failed to parse config for collection %s: %s", + collection, + e, + ) + except Exception as e: + logger.debug("Could not load config for collection %s: %s", collection, e) + return collection_configs + + +async def list_collections( user_id: Optional[int] = None, is_admin: bool = False ) -> ListCollectionsResult: """List all knowledge base collections along with aggregated statistics. @@ -438,21 +517,19 @@ def list_collections( warnings: List[str] = [] try: - conn = get_connection_from_env() - ensure_documents_table(conn) - ensure_parses_table(conn) - ensure_chunks_table(conn) - - stats: Dict[str, Dict[str, int]] = defaultdict( - lambda: {"documents": 0, "parses": 0, "chunks": 0, "embeddings": 0} + # Use storage abstraction for aggregation + vector_store = get_vector_index_store() + stats: Dict[str, Dict[str, int]] = vector_store.aggregate_collection_stats( + user_id=user_id, + is_admin=is_admin, ) + + # Collect document names using storage abstraction document_names: Dict[str, Set[str]] = defaultdict(set) - def _collect_documents() -> None: - for batch in _iter_batches( - conn, - "documents", - warnings, + def _collect_document_names() -> None: + for batch in vector_store.iter_batches( + table_name="documents", columns=["collection", "source_path"], user_id=user_id, is_admin=is_admin, @@ -472,7 +549,6 @@ def _collect_documents() -> None: if not collection_raw: continue collection_key = str(collection_raw) - stats[collection_key]["documents"] += 1 source_value = source_array[idx].as_py() if source_value: import os @@ -481,92 +557,31 @@ def _collect_documents() -> None: os.path.basename(str(source_value)) ) - def _collect_simple(table_name: str, stat_key: str) -> None: - for batch in _iter_batches( - conn, - table_name, - warnings, - columns=["collection"], - user_id=user_id, - is_admin=is_admin, - ): - collection_idx = batch.schema.get_field_index("collection") - if collection_idx == -1: - continue - collection_array = batch.column(collection_idx) - for idx in range(batch.num_rows): - collection_raw = collection_array[idx].as_py() - if not collection_raw: - continue - collection_key = str(collection_raw) - stats[collection_key][stat_key] += 1 - - _collect_documents() - _collect_simple("parses", "parses") - _collect_simple("chunks", "chunks") - - for table_name in _list_table_names(conn, warnings): - if not table_name.startswith("embeddings_"): - continue - for batch in _iter_batches( - conn, - table_name, - warnings, - columns=["collection"], - user_id=user_id, - is_admin=is_admin, - ): - collection_idx = batch.schema.get_field_index("collection") - if collection_idx == -1: - continue - collection_array = batch.column(collection_idx) - for idx in range(batch.num_rows): - collection_raw = collection_array[idx].as_py() - if not collection_raw: - continue - collection_key = str(collection_raw) - stats[collection_key]["embeddings"] += 1 + _collect_document_names() collection_keys = sorted(stats.keys() | document_names.keys()) - # Load configs for collections - collection_configs = {} + # Load configs for collections (admin sees cross-tenant configs) + collection_configs: Dict[str, IngestionConfig] = {} try: - # TODO(refactor): this still reads per-user config from - # collection_config for backward compatibility. Move to the unified - # metadata/config store after migration semantics are defined. - ensure_collection_config_table(conn) - table = conn.open_table("collection_config") - - # Apply user filter if needed - config_filter = UserPermissions.get_user_filter(user_id, is_admin) - - if config_filter: - try: - df = table.search().where(config_filter).to_pandas() - except Exception as e: - logger.warning(f"Failed to apply filter to collection_config: {e}") - df = table.to_pandas() - else: - df = table.to_pandas() - - for _, row in df.iterrows(): - col_name = row["collection"] - config_json = row.get("config_json") - if col_name and config_json: - import json - - from ..core.schemas import IngestionConfig - - try: - config_dict = json.loads(config_json) - collection_configs[col_name] = IngestionConfig(**config_dict) - except Exception as e: - logger.warning( - f"Failed to parse config for collection {col_name}: {e}" - ) + collection_configs = await _load_collection_ingestion_configs( + collection_keys, user_id, is_admin + ) except Exception as e: - logger.warning(f"Could not load collection configs: {e}") + logger.warning("Could not load collection configs: %s", e) + + # Ensure all collections have complete stats + for collection in collection_keys: + if collection not in stats: + stats[collection] = { + "documents": 0, + "parses": 0, + "chunks": 0, + "embeddings": 0, + } + for key in ["documents", "parses", "chunks", "embeddings"]: + if key not in stats[collection]: + stats[collection][key] = 0 collections = [ CollectionInfo( @@ -629,58 +644,74 @@ def get_document_stats( warnings: List[str] = [] try: - conn = get_connection_from_env() - ensure_documents_table(conn) - ensure_parses_table(conn) - ensure_chunks_table(conn) + # Use storage abstraction for basic aggregation + vector_store = get_vector_index_store() + raw_stats = vector_store.aggregate_document_stats( + collection_name=collection, + doc_id=doc_id, + user_id=user_id, + is_admin=is_admin, + ) + + document_count = raw_stats["documents"] + document_exists = document_count > 0 + parse_count = raw_stats["parses"] + chunk_count = raw_stats["chunks"] + + # Handle model_tag specific embeddings filtering + embedding_breakdown: Dict[str, int] = {} + + if model_tag: + # When model_tag is specified, only count embeddings for that specific table + safe_collection = escape_lancedb_string(collection) + safe_doc_id = escape_lancedb_string(doc_id) + filters = {"collection": safe_collection, "doc_id": safe_doc_id} + table_name = embeddings_table_name(model_tag) + embedding_count = vector_store.count_rows( + table_name=table_name, + filters=filters, + user_id=user_id, + is_admin=is_admin, + ) + embedding_breakdown[table_name] = embedding_count + else: + # Use the aggregated count from storage abstraction + embedding_count = raw_stats["embeddings"] + # Optionally include breakdown by table if needed + safe_collection = escape_lancedb_string(collection) + safe_doc_id = escape_lancedb_string(doc_id) + filters = {"collection": safe_collection, "doc_id": safe_doc_id} + + try: + table_names = vector_store.list_table_names() + except Exception as exc: # noqa: BLE001 - convert to warning + message = f"Unable to enumerate embeddings tables: {exc}" + logger.warning(message) + warnings.append(message) + table_names = [] + + for table_name in table_names: + if not table_name.startswith("embeddings_"): + continue + count = vector_store.count_rows( + table_name=table_name, + filters=filters, + user_id=user_id, + is_admin=is_admin, + ) + if count: + embedding_breakdown[table_name] = count + except Exception as exc: # noqa: BLE001 - convert to structured failure - logger.error("Failed to initialise LanceDB tables: %s", exc, exc_info=True) + logger.error("Failed to get document stats: %s", exc, exc_info=True) return DocumentStatsResult( status="error", data=None, - message=f"Failed to initialise LanceDB tables: {exc}", + message=f"Failed to get document stats: {exc}", warnings=warnings, ) - ensure_ingestion_runs_table(conn) - - filters = {"collection": collection, "doc_id": doc_id} - - document_count = _count_rows(conn, "documents", filters, warnings) - document_exists = document_count > 0 - parse_count = _count_rows(conn, "parses", filters, warnings) - chunk_count = _count_rows(conn, "chunks", filters, warnings) - - embedding_breakdown: Dict[str, int] = {} - - def _count_embeddings(table_name: str) -> int: - return _count_rows(conn, table_name, filters, warnings) - - if model_tag: - table_name = embeddings_table_name(model_tag) - embedding_count = _count_embeddings(table_name) - embedding_breakdown[table_name] = embedding_count - else: - try: - table_names = _list_table_names(conn, warnings) - except Exception as exc: # noqa: BLE001 - convert to warning - message = f"Unable to enumerate embeddings tables: {exc}" - logger.warning(message) - warnings.append(message) - table_names = [] - - for table_name in table_names: - if not table_name.startswith("embeddings_"): - continue - embedding_count = _count_embeddings(table_name) - if embedding_count: - embedding_breakdown[table_name] = embedding_count - - embedding_count = sum(embedding_breakdown.values()) - - if model_tag: - embedding_count = embedding_breakdown.get(embeddings_table_name(model_tag), 0) - + # Load ingestion status status_record = None status_entries = load_ingestion_status(collection=collection, doc_id=doc_id) if status_entries: @@ -764,70 +795,63 @@ def list_documents( warnings: List[str] = [] try: - conn = get_connection_from_env() - ensure_documents_table(conn) - ensure_parses_table(conn) - ensure_chunks_table(conn) - ensure_ingestion_runs_table(conn) + # Use storage abstraction for document records + vector_store = get_vector_index_store() + doc_records = vector_store.list_document_records( + collection_name=collection, + user_id=user_id, + is_admin=is_admin, + max_results=DEFAULT_VECTOR_STORE_EXTENDED_SCAN_LIMIT, # Higher limit for listing + ) + + # Collect document info from records + document_info: Dict[str, Dict[str, Any]] = {} + for record in doc_records: + document_info[record.doc_id] = { + "source_path": record.source_path, + "uploaded_at": None, # Not available in DocumentRecord + } + except Exception as exc: # noqa: BLE001 - logger.error("Failed to initialise LanceDB tables: %s", exc, exc_info=True) + logger.error("Failed to list documents: %s", exc, exc_info=True) return DocumentListResult( status="error", documents=[], total_count=0, - message=f"Failed to initialise LanceDB tables: {exc}", + message=f"Failed to list documents: {exc}", warnings=warnings, ) - document_info: Dict[str, Dict[str, Any]] = {} - for batch in _iter_batches( - conn, - "documents", - warnings, - columns=["collection", "doc_id", "source_path", "uploaded_at"], + # Collect chunk counts using storage abstraction + chunk_counts = vector_store.aggregate_document_counts( + table_name="chunks", + doc_id_column="doc_id", + collection_name=collection, user_id=user_id, is_admin=is_admin, - ): - collection_idx = batch.schema.get_field_index("collection") - doc_idx = batch.schema.get_field_index("doc_id") - if collection_idx == -1 or doc_idx == -1: - continue - source_idx = batch.schema.get_field_index("source_path") - uploaded_idx = batch.schema.get_field_index("uploaded_at") - collection_array = batch.column(collection_idx) - doc_array = batch.column(doc_idx) - for idx in range(batch.num_rows): - collection_raw = collection_array[idx].as_py() - if not collection_raw or str(collection_raw) != collection: - continue - doc_raw = doc_array[idx].as_py() - if not doc_raw: - continue - info: Dict[str, Any] = {} - if source_idx != -1: - info["source_path"] = batch.column(source_idx)[idx].as_py() - if uploaded_idx != -1: - info["uploaded_at"] = batch.column(uploaded_idx)[idx].as_py() - document_info[str(doc_raw)] = info - - chunk_counts = _collect_doc_counts_for_collection( - conn, "chunks", "doc_id", collection, warnings, user_id, is_admin ) + # Collect embedding counts embedding_counts: Dict[str, int] = defaultdict(int) - for table_name in _list_table_names(conn, warnings): + for table_name in vector_store.list_table_names(): if not table_name.startswith("embeddings_"): continue - table_counts = _collect_doc_counts_for_collection( - conn, table_name, "doc_id", collection, warnings, user_id, is_admin + table_counts = vector_store.aggregate_document_counts( + table_name=table_name, + doc_id_column="doc_id", + collection_name=collection, + user_id=user_id, + is_admin=is_admin, ) for doc_id, value in table_counts.items(): embedding_counts[doc_id] += value + # Load status records status_records = { entry["doc_id"]: entry for entry in load_ingestion_status(collection=collection) } + # Combine all doc_ids from various sources doc_ids = ( set(document_info.keys()) | set(chunk_counts.keys()) @@ -835,6 +859,7 @@ def list_documents( | set(status_records.keys()) ) + # Build summaries summaries: List[DocumentSummary] = [] for doc_id in sorted(doc_ids): info = document_info.get(doc_id, {}) @@ -911,76 +936,44 @@ def delete_collection( warnings: List[str] = [] try: - conn = get_connection_from_env() - ensure_documents_table(conn) - ensure_parses_table(conn) - ensure_chunks_table(conn) - ensure_ingestion_runs_table(conn) - except Exception as exc: # noqa: BLE001 + # Use storage abstraction for deletion + vector_store = get_vector_index_store() + + # Collect doc_ids before deletion for affected_documents + # Use list_document_records which respects user filtering + doc_records = vector_store.list_document_records( + collection_name=collection, + user_id=user_id, + is_admin=is_admin, + max_results=DEFAULT_VECTOR_STORE_EXTENDED_SCAN_LIMIT, # Higher limit for collection deletion + ) + doc_ids = sorted({r.doc_id for r in doc_records}) + + # Delete all data using storage abstraction + deleted_counts = vector_store.delete_collection_data(collection_name=collection) + + # Clear ingestion status for all documents + for doc_id in doc_ids: + try: + clear_ingestion_status(collection, doc_id) + except Exception as exc: # noqa: BLE001 + warning = f"Failed to clear ingestion status for '{doc_id}': {exc}" + logger.warning(warning) + warnings.append(warning) + + except Exception as exc: # noqa: BLE001 - convert to structured failure logger.error( - "Failed to initialise LanceDB tables for delete_collection: %s", - exc, - exc_info=True, + "Failed to delete collection '%s': %s", collection, exc, exc_info=True ) return CollectionOperationResult( status="error", collection=collection, - message=f"Failed to initialise LanceDB tables: {exc}", + message=f"Failed to delete collection: {exc}", warnings=warnings, affected_documents=[], deleted_counts={}, ) - # Collect doc_ids before deletion for affected_documents - doc_ids = sorted( - _collect_document_ids(conn, collection, warnings, user_id, is_admin) - ) - - # Delete all data using direct table.delete() with escaped collection name - deleted_counts: Dict[str, int] = defaultdict(int) - table_names = _list_table_names(conn, warnings) - - # Delete from core tables - for table_name in ["documents", "parses", "chunks"]: - if table_name in table_names: - try: - table = conn.open_table(table_name) - original_count = table.count_rows() - # Delete all rows for this collection using escaped string - table.delete(f"collection = '{escape_lancedb_string(collection)}'") - deleted_count = original_count - table.count_rows() - if deleted_count > 0: - deleted_counts[table_name] = deleted_count - except Exception as exc: # noqa: BLE001 - warning = f"Failed to delete from '{table_name}': {exc}" - logger.warning(warning) - warnings.append(warning) - - # Delete embeddings data - embeddings_tables = [t for t in table_names if t.startswith("embeddings_")] - for table_name in embeddings_tables: - try: - table = conn.open_table(table_name) - original_count = table.count_rows() - # Delete all rows for this collection using escaped string - table.delete(f"collection = '{escape_lancedb_string(collection)}'") - deleted_count = original_count - table.count_rows() - if deleted_count > 0: - deleted_counts[table_name] = deleted_count - except Exception as exc: # noqa: BLE001 - warning = f"Failed to delete from '{table_name}': {exc}" - logger.warning(warning) - warnings.append(warning) - - # Clear ingestion status for all documents - for doc_id in doc_ids: - try: - clear_ingestion_status(collection, doc_id) - except Exception as exc: # noqa: BLE001 - warning = f"Failed to clear ingestion status for '{doc_id}': {exc}" - logger.warning(warning) - warnings.append(warning) - # Construct affected_documents list affected: List[CollectionOperationDetail] = [ CollectionOperationDetail( @@ -1158,29 +1151,31 @@ def cancel_collection( warnings: List[str] = [] try: - conn = get_connection_from_env() - ensure_documents_table(conn) - ensure_parses_table(conn) - ensure_chunks_table(conn) - ensure_ingestion_runs_table(conn) + # Use storage abstraction to get document IDs + vector_store = get_vector_index_store() + doc_records = vector_store.list_document_records( + collection_name=collection, + user_id=user_id, + is_admin=is_admin, + max_results=DEFAULT_VECTOR_STORE_EXTENDED_SCAN_LIMIT, # Higher limit for collection operations + ) + doc_ids = sorted({r.doc_id for r in doc_records}) + except Exception as exc: # noqa: BLE001 logger.error( - "Failed to initialise LanceDB tables for cancel_collection: %s", + "Failed to get document IDs for cancel_collection: %s", exc, exc_info=True, ) return CollectionOperationResult( status="error", collection=collection, - message=f"Failed to initialise LanceDB tables: {exc}", + message=f"Failed to get document IDs: {exc}", warnings=warnings, affected_documents=[], deleted_counts={}, ) - doc_ids = sorted( - _collect_document_ids(conn, collection, warnings, user_id, is_admin) - ) cancellation_message = reason or "Cancelled by user." affected: List[CollectionOperationDetail] = [] diff --git a/src/xagent/core/tools/core/RAG_tools/management/status.py b/src/xagent/core/tools/core/RAG_tools/management/status.py index 6feeef331..28e9bf0b6 100644 --- a/src/xagent/core/tools/core/RAG_tools/management/status.py +++ b/src/xagent/core/tools/core/RAG_tools/management/status.py @@ -1,7 +1,9 @@ -"""Helpers for tracking document ingestion status in LanceDB. +"""Helpers for tracking document ingestion status. This module provides functions to track, load, and manage the ingestion status of documents being processed in the RAG pipeline. + +Phase 1A Part 2: Refactored to use IngestionStatusStore abstraction layer. """ from __future__ import annotations @@ -10,13 +12,7 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Optional -import pandas as pd - -from xagent.providers.vector_store.lancedb import get_connection_from_env - -from ..LanceDB.schema_manager import ensure_ingestion_runs_table -from ..utils.string_utils import build_lancedb_filter_expression -from ..utils.user_permissions import UserPermissions +from ..storage.factory import get_ingestion_status_store logger = logging.getLogger(__name__) @@ -37,7 +33,7 @@ def write_ingestion_status( """Persist the latest ingestion status for a document. This function writes the current status of a document's ingestion process - to the LanceDB ingestion_runs table. + to the ingestion_runs table using the storage abstraction layer. Args: collection: Name of the collection @@ -49,30 +45,19 @@ def write_ingestion_status( Returns: None - """ - conn = get_connection_from_env() - ensure_ingestion_runs_table(conn) - table = conn.open_table("ingestion_runs") - - filter_expr = build_lancedb_filter_expression( - {"collection": collection, "doc_id": doc_id} + Raises: + DatabaseOperationError: If write operation fails. + """ + store = get_ingestion_status_store() + store.write_ingestion_status( + collection=collection, + doc_id=doc_id, + status=status, + message=message, + parse_hash=parse_hash, + user_id=user_id, ) - if filter_expr: - table.delete(filter_expr) - - timestamp = _now() - record = { - "collection": collection, - "doc_id": doc_id, - "status": status, - "message": message or "", - "parse_hash": parse_hash or "", - "created_at": timestamp, - "updated_at": timestamp, - "user_id": user_id, # Add user_id for multi-tenancy - } - table.add([record]) def load_ingestion_status( @@ -83,8 +68,9 @@ def load_ingestion_status( ) -> List[Dict[str, Any]]: """Return ingestion status records filtered by collection/doc. - This function retrieves ingestion status records from the LanceDB - ingestion_runs table, with optional filtering by collection and document. + This function retrieves ingestion status records from the ingestion_runs + table using the storage abstraction layer, with optional filtering by + collection and document. Args: collection: Optional collection name to filter by @@ -102,38 +88,17 @@ def load_ingestion_status( - created_at: Creation timestamp - updated_at: Last update timestamp - user_id: User ID who owns the document - """ - conn = get_connection_from_env() - ensure_ingestion_runs_table(conn) - table = conn.open_table("ingestion_runs") - - filters: Dict[str, str] = {} - if collection is not None: - filters["collection"] = collection - if doc_id is not None: - filters["doc_id"] = doc_id - - base_filter = build_lancedb_filter_expression(filters) - user_filter = UserPermissions.get_user_filter(user_id, is_admin) - - if user_filter and base_filter: - filter_expr = f"({base_filter}) and ({user_filter})" - elif user_filter: - filter_expr = user_filter - else: - filter_expr = base_filter - try: - search = table.search() - if filter_expr: - search = search.where(filter_expr) - df = search.to_pandas() - except Exception as e: - logger.error(f"Failed to load ingestion status: {e}") - df = pd.DataFrame() - - records: List[Dict[str, Any]] = df.to_dict("records") - return records + Raises: + DatabaseOperationError: If read operation fails. + """ + store = get_ingestion_status_store() + return store.load_ingestion_status( + collection=collection, + doc_id=doc_id, + user_id=user_id, + is_admin=is_admin, + ) def clear_ingestion_status( @@ -142,7 +107,7 @@ def clear_ingestion_status( """Remove stored ingestion status for a document. This function deletes the ingestion status record for a specific document - from the LanceDB ingestion_runs table. + from the ingestion_runs table using the storage abstraction layer. Args: collection: Name of the collection @@ -152,23 +117,110 @@ def clear_ingestion_status( Returns: None + + Raises: + DatabaseOperationError: If delete operation fails. + """ + store = get_ingestion_status_store() + store.clear_ingestion_status( + collection=collection, + doc_id=doc_id, + user_id=user_id, + is_admin=is_admin, + ) + + +# ============================================================================ +# Async variants (Phase 1A Option C: Hybrid approach) +# ============================================================================ + + +async def write_ingestion_status_async( + collection: str, + doc_id: str, + *, + status: str, + message: Optional[str] = None, + parse_hash: Optional[str] = None, + user_id: Optional[int] = None, +) -> None: + """Async version of write_ingestion_status. + + Args: + collection: Name of the collection + doc_id: Unique identifier for the document + status: Current status value + message: Optional status message + parse_hash: Optional parse hash + user_id: Optional user ID + + Returns: + None + + Raises: + DatabaseOperationError: If write operation fails. """ + store = get_ingestion_status_store() + await store.write_ingestion_status_async( + collection=collection, + doc_id=doc_id, + status=status, + message=message, + parse_hash=parse_hash, + user_id=user_id, + ) + + +async def load_ingestion_status_async( + collection: Optional[str] = None, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + is_admin: bool = False, +) -> List[Dict[str, Any]]: + """Async version of load_ingestion_status. + + Args: + collection: Optional collection name to filter by + doc_id: Optional document ID to filter by + user_id: Optional user ID for multi-tenancy filtering + is_admin: Whether the user has admin privileges - conn = get_connection_from_env() - ensure_ingestion_runs_table(conn) - table = conn.open_table("ingestion_runs") + Returns: + List of ingestion status records. - base_filter = build_lancedb_filter_expression( - {"collection": collection, "doc_id": doc_id} + Raises: + DatabaseOperationError: If read operation fails. + """ + store = get_ingestion_status_store() + return await store.load_ingestion_status_async( + collection=collection, + doc_id=doc_id, + user_id=user_id, + is_admin=is_admin, ) - user_filter = UserPermissions.get_user_filter(user_id, is_admin) - if user_filter and base_filter: - filter_expr = f"({base_filter}) and ({user_filter})" - elif user_filter: - filter_expr = user_filter - else: - filter_expr = base_filter - if filter_expr: - table.delete(filter_expr) +async def clear_ingestion_status_async( + collection: str, doc_id: str, user_id: Optional[int] = None, is_admin: bool = False +) -> None: + """Async version of clear_ingestion_status. + + Args: + collection: Name of the collection + doc_id: Unique identifier for the document + user_id: Optional user ID for multi-tenancy filtering + is_admin: Whether the user has admin privileges + + Returns: + None + + Raises: + DatabaseOperationError: If delete operation fails. + """ + store = get_ingestion_status_store() + await store.clear_ingestion_status_async( + collection=collection, + doc_id=doc_id, + user_id=user_id, + is_admin=is_admin, + ) diff --git a/src/xagent/core/tools/core/RAG_tools/parse/parse_display.py b/src/xagent/core/tools/core/RAG_tools/parse/parse_display.py index f761aa46b..5b1014d40 100644 --- a/src/xagent/core/tools/core/RAG_tools/parse/parse_display.py +++ b/src/xagent/core/tools/core/RAG_tools/parse/parse_display.py @@ -8,7 +8,6 @@ import logging from typing import Any, Dict, List, Optional, Tuple -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.exceptions import DatabaseOperationError, DocumentNotFoundError from ..core.schemas import ( ParsedElementDisplay, @@ -16,10 +15,7 @@ ParsedTableDisplay, ParsedTextSegmentDisplay, ) -from ..LanceDB.schema_manager import ensure_parses_table -from ..utils.lancedb_query_utils import query_to_list -from ..utils.string_utils import build_lancedb_filter_expression -from ..utils.user_permissions import UserPermissions +from ..storage.factory import get_vector_index_store logger = logging.getLogger(__name__) @@ -31,7 +27,7 @@ def reconstruct_parse_result_from_db( user_id: Optional[int] = None, is_admin: bool = False, ) -> Tuple[List[Dict[str, Any]], Optional[str]]: - """Reconstruct ParseResult-like structure from database. + """Reconstruct ParseResult-like structure from database using abstraction layer. Args: collection: Collection name @@ -47,9 +43,7 @@ def reconstruct_parse_result_from_db( elements is a list of dictionaries with 'type', 'text'/'html', and 'metadata' keys. """ try: - conn = get_connection_from_env() - ensure_parses_table(conn) - table = conn.open_table("parses") + vector_store = get_vector_index_store() # Build base filter expression query_filters: Dict[str, Any] = { @@ -59,17 +53,12 @@ def reconstruct_parse_result_from_db( if parse_hash: query_filters["parse_hash"] = parse_hash - base_filter_expr = build_lancedb_filter_expression(query_filters) - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - - if user_filter_expr and base_filter_expr: - filter_expr = f"({base_filter_expr}) and ({user_filter_expr})" - elif user_filter_expr: - filter_expr = user_filter_expr - else: - filter_expr = base_filter_expr - - if table.count_rows(filter_expr) == 0: + if ( + vector_store.count_rows_or_zero( + "parses", filters=query_filters, user_id=user_id, is_admin=is_admin + ) + == 0 + ): if parse_hash: raise DocumentNotFoundError( f"Parse result not found: doc_id={doc_id}, parse_hash={parse_hash}" @@ -78,8 +67,18 @@ def reconstruct_parse_result_from_db( f"No parse results found for document: doc_id={doc_id}" ) - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - records = query_to_list(table.search().where(filter_expr)) + # Use iter_batches to load all matching records + records = [] + for batch in vector_store.iter_batches( + table_name="parses", + filters=query_filters, + user_id=user_id, + is_admin=is_admin, + ): + batch_df = batch.to_pandas() + for _, row in batch_df.iterrows(): + records.append(row.to_dict()) + if not records: raise DocumentNotFoundError( f"No parse results found for document: doc_id={doc_id}" diff --git a/src/xagent/core/tools/core/RAG_tools/parse/parse_document.py b/src/xagent/core/tools/core/RAG_tools/parse/parse_document.py index 8ee3dd437..d8fab5114 100644 --- a/src/xagent/core/tools/core/RAG_tools/parse/parse_document.py +++ b/src/xagent/core/tools/core/RAG_tools/parse/parse_document.py @@ -18,7 +18,6 @@ DocumentParseArgs, ) from ......core.tools.core.document_parser import parse_document as core_parse_document -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.exceptions import ( ConfigurationError, DatabaseOperationError, @@ -31,11 +30,8 @@ ParsedParagraph, ParseMethod, ) -from ..LanceDB.schema_manager import ensure_documents_table, ensure_parses_table +from ..storage.factory import get_vector_index_store from ..utils.hash_utils import compute_parse_hash, get_parse_params_whitelist -from ..utils.lancedb_query_utils import query_to_list -from ..utils.string_utils import build_lancedb_filter_expression -from ..utils.user_permissions import UserPermissions logger = logging.getLogger(__name__) @@ -114,13 +110,6 @@ async def _parse_document_internal( logger.info(f"Starting document parsing: doc_id={doc_id}, method={parse_method}") - try: - conn = get_connection_from_env() - ensure_parses_table(conn) - ensure_documents_table(conn) - except Exception as e: - raise DatabaseOperationError(f"Failed to connect to database: {e}") from e - document = _get_document_from_db(collection, doc_id, user_id, is_admin) if not document: raise DocumentNotFoundError(f"Document not found: {doc_id}") @@ -335,30 +324,31 @@ def _convert_parse_result_to_paragraphs(result: Any) -> List[ParsedParagraph]: def _get_document_from_db( collection: str, doc_id: str, user_id: Optional[int] = None, is_admin: bool = False ) -> Optional[Any]: - """Get document from database by doc_id.""" + """Get document from database by doc_id using abstraction layer.""" try: - conn = get_connection_from_env() - ensure_documents_table(conn) - table = conn.open_table("documents") + vector_store = get_vector_index_store() query_filters = {"collection": collection, "doc_id": doc_id} - base_filter_expr = build_lancedb_filter_expression(query_filters) - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - if user_filter_expr and base_filter_expr: - filter_expr = f"({base_filter_expr}) and ({user_filter_expr})" - elif user_filter_expr: - filter_expr = user_filter_expr - else: - filter_expr = base_filter_expr - - if table.count_rows(filter_expr) == 0: + if ( + vector_store.count_rows_or_zero( + "documents", filters=query_filters, user_id=user_id, is_admin=is_admin + ) + == 0 + ): return None - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - records = query_to_list(table.search().where(filter_expr)) - if not records: - return None - return records[0] + # Use iter_batches to load the document + for batch in vector_store.iter_batches( + table_name="documents", + filters=query_filters, + user_id=user_id, + is_admin=is_admin, + ): + batch_df = batch.to_pandas() + for _, row in batch_df.iterrows(): + return row.to_dict() + + return None except Exception as e: logger.error(f"Failed to get document from database: {e}") @@ -390,7 +380,7 @@ def _parse_exists( user_id: Optional[int] = None, is_admin: bool = False, ) -> bool: - """Check if parse record already exists. + """Check if parse record already exists using abstraction layer. Args: collection: Collection name @@ -403,25 +393,18 @@ def _parse_exists( True if parse record exists and is accessible to the user """ try: - conn = get_connection_from_env() - table = conn.open_table("parses") + vector_store = get_vector_index_store() query_filters = { "collection": collection, "doc_id": doc_id, "parse_hash": parse_hash, } - base_filter_expr = build_lancedb_filter_expression(query_filters) - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - - # Combine filters for multi-tenancy - if user_filter_expr and base_filter_expr: - filter_expr = f"({base_filter_expr}) and ({user_filter_expr})" - elif user_filter_expr: - filter_expr = user_filter_expr - else: - filter_expr = base_filter_expr - - return bool(table.count_rows(filter_expr) > 0) + return bool( + vector_store.count_rows_or_zero( + "parses", filters=query_filters, user_id=user_id, is_admin=is_admin + ) + > 0 + ) except Exception as e: raise DatabaseOperationError(f"Database query failed: {e}") from e @@ -433,7 +416,7 @@ def _get_existing_parse_content( user_id: Optional[int] = None, is_admin: bool = False, ) -> List[ParsedParagraph]: - """Get existing parse content from database. + """Get existing parse content from database using abstraction layer. Args: collection: Collection name @@ -446,47 +429,48 @@ def _get_existing_parse_content( List of parsed paragraphs if found and accessible, empty list otherwise """ try: - conn = get_connection_from_env() - table = conn.open_table("parses") + vector_store = get_vector_index_store() query_filters = { "collection": collection, "doc_id": doc_id, "parse_hash": parse_hash, } - base_filter_expr = build_lancedb_filter_expression(query_filters) - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - - # Combine filters for multi-tenancy - if user_filter_expr and base_filter_expr: - filter_expr = f"({base_filter_expr}) and ({user_filter_expr})" - elif user_filter_expr: - filter_expr = user_filter_expr - else: - filter_expr = base_filter_expr - if table.count_rows(filter_expr) == 0: - return [] - - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - records = query_to_list(table.search().where(filter_expr)) - if not records: + if ( + vector_store.count_rows_or_zero( + "parses", filters=query_filters, user_id=user_id, is_admin=is_admin + ) + == 0 + ): return [] - record = records[0] - parsed_content = record.get("parsed_content") - if not parsed_content: - return [] + # Use iter_batches to load the parse content + for batch in vector_store.iter_batches( + table_name="parses", + filters=query_filters, + user_id=user_id, + is_admin=is_admin, + ): + batch_df = batch.to_pandas() + for _, row in batch_df.iterrows(): + record = row.to_dict() + parsed_content = record.get("parsed_content") + if not parsed_content: + continue + + data = json.loads(parsed_content) + paragraphs = [] + for item in data: + paragraphs.append( + ParsedParagraph( + text=item.get("text", ""), + metadata=item.get("metadata", {}), + ) + ) + return paragraphs + + return [] - data = json.loads(parsed_content) - paragraphs = [] - for item in data: - paragraphs.append( - ParsedParagraph( - text=item.get("text", ""), - metadata=item.get("metadata", {}), - ) - ) - return paragraphs except Exception as e: logger.error(f"Failed to read parse content: {e}") raise DatabaseOperationError(f"Failed reading parse content: {e}") from e @@ -501,7 +485,7 @@ def _write_parse_to_db( paragraphs: List[ParsedParagraph], user_id: Optional[int] = None, ) -> bool: - """Write parse record to database.""" + """Write parse record to database using abstraction layer.""" enable_timing = os.environ.get("PARSE_DETAILED_TIMING", "0").lower() in ( "1", "true", @@ -509,8 +493,7 @@ def _write_parse_to_db( ) try: - conn = get_connection_from_env() - table = conn.open_table("parses") + vector_store = get_vector_index_store() if enable_timing: serialize_start = time.perf_counter() @@ -540,7 +523,7 @@ def _write_parse_to_db( ) db_op_start = time.perf_counter() logger.debug( - "[PARSE TIMING] - Starting database operation (merge_insert)..." + "[PARSE TIMING] - Starting database operation (upsert_parses)..." ) parse_record = { @@ -553,11 +536,9 @@ def _write_parse_to_db( "parsed_content": parsed_content, "user_id": user_id, # Add user_id for multi-tenancy } - table.merge_insert( - ["collection", "doc_id", "parse_hash"] - ).when_matched_update_all().when_not_matched_insert_all().execute( - [parse_record] - ) + + # Use abstraction layer for upsert + vector_store.upsert_parses([parse_record]) if enable_timing: db_op_end = time.perf_counter() diff --git a/src/xagent/core/tools/core/RAG_tools/pipelines/document_ingestion.py b/src/xagent/core/tools/core/RAG_tools/pipelines/document_ingestion.py index 3964d8341..0147e02d7 100644 --- a/src/xagent/core/tools/core/RAG_tools/pipelines/document_ingestion.py +++ b/src/xagent/core/tools/core/RAG_tools/pipelines/document_ingestion.py @@ -208,7 +208,8 @@ async def encode_single_with_retry( doc_id=chunk.doc_id, chunk_id=chunk.chunk_id, parse_hash=chunk.parse_hash, - model=embedding_config.model_name, + # IMPORTANT: Use Hub model ID as the single source of truth. + model=embedding_config.id, vector=vector, text=chunk.text, chunk_hash=chunk.chunk_hash, @@ -458,7 +459,9 @@ def process_document( # Note: Parameters passed to _resolve_embedding_adapter have priority over environment variables resolve_start = time.time() embedding_config, embedding_adapter = _resolve_embedding_adapter(cfg) - selected_model_id = cfg.embedding_model_id or embedding_config.id + selected_model_id = ( + cfg.embedding_model_id or embedding_config.id or "" + ).strip() provider = getattr(embedding_config, "model_provider", None) logger.info( @@ -696,7 +699,7 @@ def process_document( "collection": collection, "doc_id": doc_id, "parse_hash": parse_hash, - "embedding_model": embedding_config.model_name, + "embedding_model": embedding_config.id, }, ) read_start = time.time() @@ -704,7 +707,9 @@ def process_document( collection=collection, doc_id=doc_id, parse_hash=parse_hash, - model=embedding_config.model_name, + # IMPORTANT: Use Hub model ID as the single source of truth, + # matching the write path (embedding writes use embedding_config.id). + model=embedding_config.id, user_id=user_id, is_admin=is_admin, ) @@ -877,7 +882,8 @@ def process_document( doc_id=chunk.doc_id, chunk_id=chunk.chunk_id, parse_hash=chunk.parse_hash, - model=embedding_config.model_name, + # IMPORTANT: Use Hub model ID as the single source of truth. + model=embedding_config.id, vector=vector, text=chunk.text, chunk_hash=chunk.chunk_hash, diff --git a/src/xagent/core/tools/core/RAG_tools/pipelines/document_search.py b/src/xagent/core/tools/core/RAG_tools/pipelines/document_search.py index a60cec0ce..ca1e7fe60 100644 --- a/src/xagent/core/tools/core/RAG_tools/pipelines/document_search.py +++ b/src/xagent/core/tools/core/RAG_tools/pipelines/document_search.py @@ -568,7 +568,9 @@ def search_documents( base_url=None, timeout_sec=None, ) - model_tag = embedding_config.model_name + # IMPORTANT: We use the Hub model ID as the single source of truth. + # It is used for embedding table naming and persisted collection binding. + embedding_model_id = (cfg.embedding_model_id or "").strip() current_step = "post_resolve_embedding" actual_type = requested_type results: List[SearchResult] = [] @@ -580,7 +582,7 @@ def search_documents( pass current_step = "search_sparse" results, status, sparse_warnings, message = _execute_sparse_search( - collection, query_text, cfg, model_tag, user_id, is_admin + collection, query_text, cfg, embedding_model_id, user_id, is_admin ) warnings.extend(sparse_warnings) else: @@ -600,7 +602,7 @@ def search_documents( "Hybrid search embedding failed; fallback to sparse." ) results, status, sparse_warnings, message = _execute_sparse_search( - collection, query_text, cfg, model_tag + collection, query_text, cfg, embedding_model_id ) warnings.extend(sparse_warnings) actual_type = SearchType.SPARSE @@ -612,7 +614,7 @@ def search_documents( pass dense_response: DenseSearchResponse = search_dense( collection=collection, - model_tag=model_tag, + model_tag=embedding_model_id, query_vector=query_vector, top_k=fetch_top_k, filters=cfg.filters, @@ -635,7 +637,7 @@ def search_documents( pass hybrid_response: HybridSearchResponse = search_hybrid( collection=collection, - model_tag=model_tag, + model_tag=embedding_model_id, query_text=query_text, query_vector=query_vector, top_k=fetch_top_k, @@ -658,7 +660,7 @@ def search_documents( ) results, status, sparse_warnings, message = ( _execute_sparse_search( - collection, query_text, cfg, model_tag + collection, query_text, cfg, embedding_model_id ) ) warnings.extend(sparse_warnings) diff --git a/src/xagent/core/tools/core/RAG_tools/prompt_manager/prompt_manager.py b/src/xagent/core/tools/core/RAG_tools/prompt_manager/prompt_manager.py index b6688e322..b2da42657 100644 --- a/src/xagent/core/tools/core/RAG_tools/prompt_manager/prompt_manager.py +++ b/src/xagent/core/tools/core/RAG_tools/prompt_manager/prompt_manager.py @@ -2,25 +2,22 @@ This module provides functions for managing prompt templates with full CRUD operations and transparent version management using LanceDB. + +Phase 1A Part 2: Refactored to use PromptTemplateStore abstraction layer +for basic operations while preserving complex business logic. """ import json import logging -from datetime import datetime from typing import Any, Dict, List, Optional -import pandas as pd - -from xagent.providers.vector_store.lancedb import get_connection_from_env - from ..core.exceptions import ( ConfigurationError, DatabaseOperationError, DocumentNotFoundError, ) from ..core.schemas import PromptTemplate -from ..LanceDB.schema_manager import ensure_prompt_templates_table -from ..utils.string_utils import escape_lancedb_string +from ..storage.factory import get_prompt_template_store logger = logging.getLogger(__name__) @@ -39,47 +36,6 @@ def _serialize_metadata(metadata: Optional[Dict[str, Any]]) -> Optional[str]: return json.dumps(metadata, ensure_ascii=False, sort_keys=True) -def _deserialize_metadata(metadata_json: Optional[str]) -> Optional[Dict[str, Any]]: - """Deserialize metadata JSON string to dictionary. - - Args: - metadata_json: JSON string to deserialize. - - Returns: - Metadata dictionary or None. - """ - if metadata_json is None or pd.isna(metadata_json): - return None - result: Dict[str, Any] = json.loads(metadata_json) - return result - - -def _get_prompt_table() -> Any: - """Get LanceDB table for prompt templates. - - Returns: - LanceDB table instance. - - Raises: - DatabaseOperationError: If table access fails. - """ - try: - db = get_connection_from_env() - table_name = "prompt_templates" - - # Ensure table exists with proper schema - ensure_prompt_templates_table(db) - - # Open and return the table - return db.open_table(table_name) - - except Exception as e: - logger.error(f"Failed to get prompt templates table: {str(e)}") - raise DatabaseOperationError( - f"Failed to access prompt templates table: {str(e)}" - ) from e - - # ------------------------- Public Functions ------------------------- @@ -115,46 +71,36 @@ def create_prompt_template( name = name.strip() try: - table = _get_prompt_table() - - # Check if a template with this name already exists (using safe filter) - # Filter by both collection and name - escaped_collection = escape_lancedb_string(collection) - escaped_name = escape_lancedb_string(name) - collection_name_filter = ( - f"collection == '{escaped_collection}' AND name == '{escaped_name}'" - ) - existing_templates = table.search().where(collection_name_filter).to_pandas() - - if existing_templates.empty: - # Create first version - version = 1 - is_latest = True - else: - # Create new version - find the highest version number - max_version = existing_templates["version"].max() - version = max_version + 1 - is_latest = True + store = get_prompt_template_store() - # Mark all previous versions as not latest - table.update(where=collection_name_filter, values={"is_latest": False}) - - # Create new prompt template - prompt_template = PromptTemplate( + # Save template via store (handles version management automatically) + template_id = store.save_prompt_template( name=name, template=template.strip(), - version=version, - is_latest=is_latest, + user_id=None, # No multi-tenancy in current implementation metadata=_serialize_metadata(metadata), ) - # Convert to DataFrame for LanceDB insertion, including collection - template_dict = prompt_template.model_dump() - template_dict["collection"] = collection - df = pd.DataFrame([template_dict]) - table.add(df) + # Get the created template to return full PromptTemplate object + template_data = store.get_prompt_template(template_id, user_id=None) + if template_data is None: + raise DatabaseOperationError("Failed to retrieve created template") - logger.info(f"Created prompt template '{name}' version {version}") + prompt_template = PromptTemplate( + id=template_data["id"], + name=template_data["name"], + template=template_data["template"], + version=template_data["version"], + is_latest=template_data["is_latest"], + metadata=template_data["metadata"], + user_id=template_data["user_id"], + created_at=template_data["created_at"], + updated_at=template_data["updated_at"], + ) + + logger.info( + f"Created prompt template '{name}' version {prompt_template.version}" + ) return prompt_template except (ConfigurationError, DatabaseOperationError): @@ -194,55 +140,58 @@ def read_prompt_template( raise ConfigurationError("Either prompt_id or name must be provided.") try: - table = _get_prompt_table() - escaped_collection = escape_lancedb_string(collection) + store = get_prompt_template_store() if prompt_id: - # Search by ID and collection (using safe filter) - escaped_id = escape_lancedb_string(prompt_id) - id_filter = f"collection == '{escaped_collection}' AND id == '{escaped_id}'" - result = table.search().where(id_filter).to_pandas() + # Search by ID + template_data = store.get_prompt_template(prompt_id, user_id=None) + if template_data is None: + raise DocumentNotFoundError( + f"Prompt template with ID '{prompt_id}' not found." + ) else: - # Normalize name + # Search by name + assert ( + name is not None + ) # Type narrowing: name must be provided if prompt_id is None name = name.strip() if name else name - # Search by name and collection - escaped_name = escape_lancedb_string(name) if version is not None: - # Specific version - combine filters safely - filter_expr = ( - f"collection == '{escaped_collection}' AND " - f"name == '{escaped_name}' AND version == {version}" + # Get specific version - need to search through list + templates = store.list_prompt_templates( + name_filter=name, + latest_only=False, + user_id=None, + limit=100, ) - result = table.search().where(filter_expr).to_pandas() + matching = [ + t + for t in templates + if t["name"] == name and t["version"] == version + ] + if not matching: + raise DocumentNotFoundError( + f"Prompt template with name '{name}' version {version} not found." + ) + template_data = matching[0] else: - # Latest version - combine filters safely - filter_expr = ( - f"collection == '{escaped_collection}' AND " - f"name == '{escaped_name}' AND is_latest == true" - ) - result = table.search().where(filter_expr).to_pandas() - - if result.empty: - identifier = ( - f"ID '{prompt_id}'" - if prompt_id - else f"name '{name}'" - + (f" version {version}" if version else " (latest)") - ) - raise DocumentNotFoundError(f"Prompt template with {identifier} not found.") + # Get latest version + template_data = store.get_latest_prompt_template(name, user_id=None) + if template_data is None: + raise DocumentNotFoundError( + f"Prompt template with name '{name}' not found." + ) # Convert to PromptTemplate - row = result.iloc[0] - # Note: metadata is stored as JSON string internally, keep it as is return PromptTemplate( - id=str(row["id"]), - name=row["name"], - template=row["template"], - version=int(row["version"]), - is_latest=bool(row["is_latest"]), - metadata=row["metadata"] if pd.notna(row["metadata"]) else None, - created_at=row["created_at"], - updated_at=row["updated_at"], + id=template_data["id"], + name=template_data["name"], + template=template_data["template"], + version=template_data["version"], + is_latest=template_data["is_latest"], + metadata=template_data["metadata"], + user_id=template_data["user_id"], + created_at=template_data["created_at"], + updated_at=template_data["updated_at"], ) except (ConfigurationError, DocumentNotFoundError): @@ -315,79 +264,74 @@ def update_prompt_template( current_template = read_prompt_template( collection=collection, prompt_id=prompt_id ) - table = _get_prompt_table() - escaped_collection = escape_lancedb_string(collection) if template is not None: # Template content changed - create new version if not template.strip(): raise ConfigurationError("Template content cannot be empty.") - # Find the highest version number for this name to avoid version conflicts - escaped_name = escape_lancedb_string(current_template.name) - name_filter = ( - f"collection == '{escaped_collection}' AND name == '{escaped_name}'" - ) - all_versions = table.search().where(name_filter).to_pandas() - max_version = all_versions["version"].max() if not all_versions.empty else 0 - new_version = max_version + 1 - - # Mark all previous versions as not latest - table.update(where=name_filter, values={"is_latest": False}) - - # Create new version - # Serialize the new metadata if provided, otherwise use current template's metadata + # Create new version using store (handles version management automatically) new_metadata = ( _serialize_metadata(metadata) if metadata is not None else current_template.metadata ) - updated_template = PromptTemplate( + new_template_id = get_prompt_template_store().save_prompt_template( name=current_template.name, template=template.strip(), - version=new_version, - is_latest=True, + user_id=None, metadata=new_metadata, ) - # Insert new version, including collection - template_dict = updated_template.model_dump() - template_dict["collection"] = collection - df = pd.DataFrame([template_dict]) - table.add(df) + # Get the created template + new_template_data = get_prompt_template_store().get_prompt_template( + new_template_id, user_id=None + ) + if new_template_data is None: + raise DatabaseOperationError("Failed to retrieve updated template") + + updated_template = PromptTemplate( + id=new_template_data["id"], + name=new_template_data["name"], + template=new_template_data["template"], + version=new_template_data["version"], + is_latest=new_template_data["is_latest"], + metadata=new_template_data["metadata"], + user_id=new_template_data["user_id"], + created_at=new_template_data["created_at"], + updated_at=new_template_data["updated_at"], + ) logger.info( - f"Created new version {new_version} for prompt template '{current_template.name}'" + f"Created new version {updated_template.version} for prompt template '{current_template.name}'" ) return updated_template else: - # Only metadata changed - update current version - metadata_json = _serialize_metadata(metadata) - updated_template = PromptTemplate( - id=current_template.id, - name=current_template.name, - template=current_template.template, - version=current_template.version, - is_latest=current_template.is_latest, - metadata=metadata_json, - created_at=current_template.created_at, - updated_at=datetime.utcnow(), + # Only metadata changed - update in-place using store method + new_metadata = _serialize_metadata(metadata) + updated_data = get_prompt_template_store().update_metadata( + template_id=prompt_id, + metadata=new_metadata, + user_id=None, ) + if updated_data is None: + raise DatabaseOperationError("Failed to retrieve updated template") - # Update the existing record (using safe filter with collection) - escaped_id = escape_lancedb_string(prompt_id) - id_filter = f"collection == '{escaped_collection}' AND id == '{escaped_id}'" - table.update( - where=id_filter, - values={ - "metadata": metadata_json, - "updated_at": updated_template.updated_at, - }, + updated_template = PromptTemplate( + id=updated_data["id"], + name=updated_data["name"], + template=updated_data["template"], + version=updated_data["version"], + is_latest=updated_data["is_latest"], + metadata=updated_data["metadata"], + user_id=updated_data["user_id"], + created_at=updated_data["created_at"], + updated_at=updated_data["updated_at"], ) logger.info( - f"Updated metadata for prompt template '{current_template.name}' version {current_template.version}" + f"Updated metadata for prompt template '{current_template.name}' (version {updated_template.version})" ) return updated_template @@ -428,95 +372,30 @@ def delete_prompt_template( raise ConfigurationError("Either prompt_id or name must be provided.") try: - table = _get_prompt_table() - escaped_collection = escape_lancedb_string(collection) + store = get_prompt_template_store() if prompt_id: - # Delete specific template by ID and collection (using safe filter) - escaped_id = escape_lancedb_string(prompt_id) - id_filter = f"collection == '{escaped_collection}' AND id == '{escaped_id}'" - result = table.search().where(id_filter).to_pandas() - if result.empty: + # Delete by ID + result = store.delete_prompt_template(template_id=prompt_id, user_id=None) + if not result: raise DocumentNotFoundError( f"Prompt template with ID '{prompt_id}' not found." ) - - # Check if this was the latest version and get the name - was_latest = result.iloc[0]["is_latest"] - template_name = result.iloc[0]["name"] - - table.delete(id_filter) - - # If we deleted the latest version, update the latest flag for the remaining versions - if was_latest: - escaped_name = escape_lancedb_string(template_name) - name_filter = ( - f"collection == '{escaped_collection}' AND name == '{escaped_name}'" - ) - remaining_versions = table.search().where(name_filter).to_pandas() - if not remaining_versions.empty: - max_version = remaining_versions["version"].max() - update_filter = ( - f"collection == '{escaped_collection}' AND " - f"name == '{escaped_name}' AND version == {max_version}" - ) - table.update(where=update_filter, values={"is_latest": True}) - logger.info(f"Deleted prompt template with ID '{prompt_id}'") return True - else: # Normalize name + assert ( + name is not None + ) # Type narrowing: name must be provided if prompt_id is None name = name.strip() if name else name - escaped_name = escape_lancedb_string(name) - # Delete by name and collection + # Delete by name using store method (handles version management automatically) + store.delete_by_name(name=name, version=version, user_id=None) if version is not None: - # Delete specific version - version_filter = ( - f"collection == '{escaped_collection}' AND " - f"name == '{escaped_name}' AND version == {version}" - ) - result = table.search().where(version_filter).to_pandas() - if result.empty: - raise DocumentNotFoundError( - f"Prompt template '{name}' version {version} not found." - ) - - # Check if this was the latest version - was_latest = result.iloc[0]["is_latest"] - - table.delete(version_filter) - - # If we deleted the latest version, update the latest flag - if was_latest: - name_filter = ( - f"collection == '{escaped_collection}' AND " - f"name == '{escaped_name}'" - ) - remaining_versions = table.search().where(name_filter).to_pandas() - if not remaining_versions.empty: - # Find the highest remaining version and mark it as latest - max_version = remaining_versions["version"].max() - update_filter = ( - f"collection == '{escaped_collection}' AND " - f"name == '{escaped_name}' AND version == {max_version}" - ) - table.update(where=update_filter, values={"is_latest": True}) - logger.info(f"Deleted prompt template '{name}' version {version}") - return True else: - # Delete all versions - name_filter = ( - f"collection == '{escaped_collection}' AND name == '{escaped_name}'" - ) - result = table.search().where(name_filter).to_pandas() - if result.empty: - raise DocumentNotFoundError(f"Prompt template '{name}' not found.") - - table.delete(name_filter) logger.info(f"Deleted all versions of prompt template '{name}'") - return True + return True except (ConfigurationError, DocumentNotFoundError): raise @@ -554,44 +433,36 @@ def list_prompt_templates( raise ConfigurationError("Collection name cannot be empty.") try: - table = _get_prompt_table() - escaped_collection = escape_lancedb_string(collection) - - # Build filter conditions safely, always include collection filter - filters = [f"collection == '{escaped_collection}'"] - - if name_filter: - # Use safe escaping for partial match - escaped_name = escape_lancedb_string(name_filter) - filters.append(f"name LIKE '%{escaped_name}%'") - - if latest_only: - filters.append("is_latest == true") + store = get_prompt_template_store() # Note: metadata filtering would require more complex logic - # For now, we'll implement basic filtering if metadata_filter: logger.warning("Metadata filtering is not yet implemented") - # Combine filters - where_clause = " AND ".join(filters) - result = table.search().where(where_clause).limit(limit).to_pandas() + # Use store method to list templates + templates_data = store.list_prompt_templates( + name_filter=name_filter, + latest_only=latest_only, + user_id=None, + limit=limit, + ) # Convert to PromptTemplate objects templates = [] - for _, row in result.iterrows(): - # Note: metadata is stored as JSON string, keep it as is - template = PromptTemplate( - id=str(row["id"]), - name=row["name"], - template=row["template"], - version=int(row["version"]), - is_latest=bool(row["is_latest"]), - metadata=row["metadata"] if pd.notna(row["metadata"]) else None, - created_at=row["created_at"], - updated_at=row["updated_at"], + for template_data in templates_data: + templates.append( + PromptTemplate( + id=template_data["id"], + name=template_data["name"], + template=template_data["template"], + version=template_data["version"], + is_latest=template_data["is_latest"], + metadata=template_data["metadata"], + user_id=template_data["user_id"], + created_at=template_data["created_at"], + updated_at=template_data["updated_at"], + ) ) - templates.append(template) logger.info(f"Listed {len(templates)} prompt templates (limit: {limit})") return templates diff --git a/src/xagent/core/tools/core/RAG_tools/retrieval/search_dense.py b/src/xagent/core/tools/core/RAG_tools/retrieval/search_dense.py index da8eb9aed..d30f4856f 100644 --- a/src/xagent/core/tools/core/RAG_tools/retrieval/search_dense.py +++ b/src/xagent/core/tools/core/RAG_tools/retrieval/search_dense.py @@ -3,14 +3,20 @@ This module provides the main entry point for dense vector search operations, handling input validation and orchestrating the search execution. + +Phase 1A Option C: Provides both sync and async search functions. """ import logging from typing import Any, Dict, List, Optional -from ......providers.vector_store.lancedb import get_connection_from_env -from ..core.exceptions import DocumentValidationError, VectorValidationError -from ..core.schemas import DenseSearchResponse, IndexStatus +from ..core.exceptions import DocumentValidationError +from ..core.schemas import ( + DenseSearchResponse, + IndexStatus, + SearchFallbackAction, + SearchWarning, +) from ..vector_storage.vector_manager import validate_query_vector from .search_engine import search_dense_engine @@ -49,7 +55,8 @@ def search_dense( is_admin: Whether the user has admin privileges (bypasses user filtering). Returns: - DenseSearchResponse with search results and metadata + DenseSearchResponse with search results and metadata. Returns a failed + response with error warnings if an exception occurs. Raises: DocumentValidationError: If input validation fails @@ -65,66 +72,220 @@ def search_dense( if top_k <= 0 or top_k > 1000: raise DocumentValidationError("top_k must be between 1 and 1000") - # Validate query vector (with model and dimension check) + # Validate query vector (basic validation without DB connection) + # Note: Dimension validation is handled by the storage abstraction layer during search + validate_query_vector(query_vector) + + try: + # Execute search using search engine + search_results, index_status, index_advice = search_dense_engine( + collection=collection, + model_tag=model_tag, + query_vector=query_vector, + top_k=top_k, + filters=filters, + readonly=readonly, + nprobes=nprobes, + refine_factor=refine_factor, + user_id=user_id, + is_admin=is_admin, + ) + + # Map index status to enum + index_status_enum = IndexStatus.INDEX_READY + if index_status == "index_building": + index_status_enum = IndexStatus.INDEX_BUILDING + elif index_status == "no_index": + index_status_enum = IndexStatus.NO_INDEX + elif index_status == "index_corrupted": + index_status_enum = IndexStatus.INDEX_CORRUPTED + elif index_status == "readonly": + index_status_enum = IndexStatus.READONLY + elif index_status == "below_threshold": + index_status_enum = IndexStatus.BELOW_THRESHOLD + + # Build response + response = DenseSearchResponse( + results=search_results, + total_count=len(search_results), + status="success", + warnings=[], + index_status=index_status_enum, + index_advice=index_advice, + # TODO: Generate idempotency_key based on search parameters hash + # (collection, model_tag, query_vector, filters, top_k, nprobes, refine_factor) + # for request deduplication, caching, and observability tracking. + # Implementation planned for PR21 (caching strategy). + idempotency_key=None, + fallback_info=None, + nprobes=nprobes, + refine_factor=refine_factor, + ) + + logger.info( + f"Dense search completed: collection={collection}, model_tag={model_tag}, " + f"top_k={top_k}, returned={len(search_results)}, index_status={index_status}" + ) + + return response + + except Exception as e: + logger.error( + f"Dense search failed for {model_tag} in collection '{collection}': {e}" + ) + # Return structured error response instead of raising exception + # This matches the behavior of search_sparse for API consistency + return DenseSearchResponse( + results=[], + total_count=0, + status="failed", + warnings=[ + SearchWarning( + code="DENSE_SEARCH_FAILED", + message=f"An unexpected error occurred during dense search: {str(e)}", + fallback_action=SearchFallbackAction.PARTIAL_RESULTS, + affected_models=[model_tag], + ) + ], + index_status=IndexStatus.NO_INDEX, + index_advice=None, + idempotency_key=None, + fallback_info=None, + nprobes=nprobes, + refine_factor=refine_factor, + ) + + +# --- Async variant (Phase 1A Option C) --- + + +async def search_dense_async( + collection: str, + model_tag: str, + query_vector: List[float], + *, + top_k: int = 10, + filters: Optional[Dict[str, Any]] = None, + readonly: bool = False, + nprobes: Optional[int] = None, + refine_factor: Optional[int] = None, + user_id: Optional[int] = None, + is_admin: bool = False, +) -> DenseSearchResponse: + """ + Execute dense vector search using async vector store abstraction. + + This is the async variant of search_dense. It performs the same input + validation but uses search_dense_engine_async() internally. + + Args: + collection: Collection name for data isolation + model_tag: Model tag identifying which embeddings table to search + query_vector: Query vector for similarity search + top_k: Number of top results to return (default: 10) + filters: Optional filters to apply to the search + readonly: If True, don't trigger index operations + nprobes: Number of partitions to probe for ANN search (LanceDB specific). + refine_factor: Refine factor for re-ranking results in memory (LanceDB specific). + user_id: Optional user ID for multi-tenancy filtering. + is_admin: Whether the user has admin privileges (bypasses user filtering). + + Returns: + DenseSearchResponse with search results and metadata. Returns a failed + response with error warnings if an exception occurs. + + Raises: + DocumentValidationError: If input validation fails + VectorValidationError: If vector validation fails + """ + # Input validation (same as sync version) + if not collection or not isinstance(collection, str): + raise DocumentValidationError("Collection must be a non-empty string") + + if not model_tag or not isinstance(model_tag, str): + raise DocumentValidationError("model_tag must be a non-empty string") + + if top_k <= 0 or top_k > 1000: + raise DocumentValidationError("top_k must be between 1 and 1000") + + # Validate query vector (basic validation without DB connection) + # Note: Dimension validation is handled by the storage abstraction layer during search + validate_query_vector(query_vector) + + # Import async search engine + from .search_engine import search_dense_engine_async + try: - # Get database connection for validation - conn = get_connection_from_env() - validate_query_vector(query_vector, model_tag, conn=conn) + # Execute async search + search_results, index_status, index_advice = await search_dense_engine_async( + collection=collection, + model_tag=model_tag, + query_vector=query_vector, + top_k=top_k, + filters=filters, + readonly=readonly, + nprobes=nprobes, + refine_factor=refine_factor, + user_id=user_id, + is_admin=is_admin, + ) + + # Map index status to enum + index_status_enum = IndexStatus.INDEX_READY + if index_status == "index_building": + index_status_enum = IndexStatus.INDEX_BUILDING + elif index_status == "no_index": + index_status_enum = IndexStatus.NO_INDEX + elif index_status == "index_corrupted": + index_status_enum = IndexStatus.INDEX_CORRUPTED + elif index_status == "readonly": + index_status_enum = IndexStatus.READONLY + elif index_status == "below_threshold": + index_status_enum = IndexStatus.BELOW_THRESHOLD + + # Build response + response = DenseSearchResponse( + results=search_results, + total_count=len(search_results), + status="success", + warnings=[], + index_status=index_status_enum, + index_advice=index_advice, + idempotency_key=None, + fallback_info=None, + nprobes=nprobes, + refine_factor=refine_factor, + ) + + logger.info( + f"Async dense search completed: collection={collection}, model_tag={model_tag}, " + f"top_k={top_k}, returned={len(search_results)}, index_status={index_status}" + ) + + return response + except Exception as e: - if isinstance(e, VectorValidationError): - raise - # If connection fails, fall back to basic validation - logger.warning(f"Could not validate with database connection: {str(e)}") - validate_query_vector(query_vector) - - # Execute search using search engine - search_results, index_status, index_advice = search_dense_engine( - collection=collection, - model_tag=model_tag, - query_vector=query_vector, - top_k=top_k, - filters=filters, - readonly=readonly, - nprobes=nprobes, - refine_factor=refine_factor, - user_id=user_id, - is_admin=is_admin, - ) - - # Map index status to enum - index_status_enum = IndexStatus.INDEX_READY - if index_status == "index_building": - index_status_enum = IndexStatus.INDEX_BUILDING - elif index_status == "no_index": - index_status_enum = IndexStatus.NO_INDEX - elif index_status == "index_corrupted": - index_status_enum = IndexStatus.INDEX_CORRUPTED - elif index_status == "readonly": - index_status_enum = IndexStatus.READONLY - elif index_status == "below_threshold": - index_status_enum = IndexStatus.BELOW_THRESHOLD - - # Build response - response = DenseSearchResponse( - results=search_results, - total_count=len(search_results), - status="success", - warnings=[], - index_status=index_status_enum, - index_advice=index_advice, - # TODO: Generate idempotency_key based on search parameters hash - # (collection, model_tag, query_vector, filters, top_k, nprobes, refine_factor) - # for request deduplication, caching, and observability tracking. - # Implementation planned for PR21 (caching strategy). - idempotency_key=None, - fallback_info=None, - nprobes=nprobes, - refine_factor=refine_factor, - ) - - logger.info( - f"Dense search completed: collection={collection}, model_tag={model_tag}, " - f"top_k={top_k}, returned={len(search_results)}, index_status={index_status}" - ) - - return response + logger.error( + f"Async dense search failed for {model_tag} in collection '{collection}': {e}" + ) + # Return structured error response instead of raising exception + # This matches the behavior of search_sparse for API consistency + return DenseSearchResponse( + results=[], + total_count=0, + status="failed", + warnings=[ + SearchWarning( + code="DENSE_SEARCH_FAILED", + message=f"An unexpected error occurred during dense search: {str(e)}", + fallback_action=SearchFallbackAction.PARTIAL_RESULTS, + affected_models=[model_tag], + ) + ], + index_status=IndexStatus.NO_INDEX, + index_advice=None, + idempotency_key=None, + fallback_info=None, + nprobes=nprobes, + refine_factor=refine_factor, + ) diff --git a/src/xagent/core/tools/core/RAG_tools/retrieval/search_engine.py b/src/xagent/core/tools/core/RAG_tools/retrieval/search_engine.py index b5bd86f20..b6d86a353 100644 --- a/src/xagent/core/tools/core/RAG_tools/retrieval/search_engine.py +++ b/src/xagent/core/tools/core/RAG_tools/retrieval/search_engine.py @@ -2,19 +2,19 @@ Core search engine implementation for dense vector retrieval. This module provides the low-level search functionality that interacts -directly with LanceDB for performing ANN searches on embeddings tables. +with the vector store abstraction layer for performing ANN searches. + +Phase 1A Option C: Provides both sync and async search functions. """ import logging from typing import Any, Dict, List, Optional, Tuple -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.schemas import SearchResult -from ..LanceDB.model_tag_utils import to_model_tag -from ..utils.lancedb_query_utils import query_to_list +from ..storage.contracts import FilterExpression +from ..storage.factory import get_vector_index_store +from ..utils.filter_utils import parse_legacy_filters, validate_filter_depth from ..utils.metadata_utils import deserialize_metadata -from ..utils.string_utils import build_lancedb_filter_expression -from ..vector_storage.index_manager import get_index_manager logger = logging.getLogger(__name__) @@ -51,61 +51,61 @@ def search_dense_engine( Tuple of (search_results, index_status, index_advice) """ try: - # Get database connection - conn = get_connection_from_env() - - # Build table name - table_name = f"embeddings_{to_model_tag(model_tag)}" - - # Open table - table = conn.open_table(table_name) - - # Check and create index if needed - index_manager = get_index_manager() - index_status, index_advice = index_manager.check_and_create_index( - table, table_name, readonly - ) + vector_store = get_vector_index_store() - # Build LanceDB search query using query builder pattern - search_query = table.search( - query_vector, - vector_column_name="vector", - ) + # Check and create index if needed (using storage abstraction) + index_result_obj = vector_store.create_index(model_tag, readonly) + index_status = index_result_obj.status + index_advice = index_result_obj.advice - # Build filter expression combining collection scope, user permissions and custom filters - filter_clauses = [] + # Convert API-facing dict filters into abstract FilterExpression + filter_expr: Optional[FilterExpression] = None + if collection or filters: + conditions: List[FilterExpression] = [] - # Scope results to the requested collection (required for KB isolation) - if collection: - collection_filter = build_lancedb_filter_expression( - {"collection": collection} - ) - if collection_filter: - filter_clauses.append(collection_filter) + if collection: + from ..storage.contracts import FilterCondition, FilterOperator - # Add user permission filter for multi-tenancy - from ..utils.user_permissions import UserPermissions + conditions.append( + FilterCondition( + field="collection", + operator=FilterOperator.EQ, + value=collection, + ) + ) - user_filter = UserPermissions.get_user_filter(user_id, is_admin) - if user_filter: - filter_clauses.append(user_filter) + if filters: + parsed = ( + parse_legacy_filters(filters) if isinstance(filters, dict) else None + ) + if parsed is not None: + if isinstance(parsed, tuple): + # Type narrowing: tuple of FilterConditions + # Cast to list for extend since tuple is also Iterable + conditions.extend(parsed) + else: + # Type narrowing: single FilterCondition + conditions.append(parsed) - # Add custom filters if provided - if filters: - custom_filter = build_lancedb_filter_expression(filters) - if custom_filter: - filter_clauses.append(custom_filter) + if len(conditions) == 1: + filter_expr = conditions[0] + elif len(conditions) > 1: + filter_expr = tuple(conditions) - # Combine all filters with AND - if filter_clauses: - combined_filter = " and ".join(f"({clause})" for clause in filter_clauses) - search_query = search_query.where(combined_filter) + # Validate filter expression depth to prevent DoS + if filter_expr is not None: + validate_filter_depth(filter_expr) - # Limit results - search_query = search_query.limit(top_k) - - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - raw_results = query_to_list(search_query) + # Execute vector search using abstraction layer (by model_tag) + raw_results = vector_store.search_vectors_by_model( + model_tag=model_tag, + query_vector=query_vector, + top_k=top_k, + filters=filter_expr, + vector_column_name="vector", + user_id=user_id, + is_admin=is_admin, + ) # OPTIMIZATION: Use list comprehension instead of iterrows() # Convert raw results to SearchResult objects @@ -139,3 +139,125 @@ def search_dense_engine( except Exception as e: logger.error(f"Failed to execute dense search: {str(e)}") raise + + +# --- Async variant (Phase 1A Option C) --- + + +async def search_dense_engine_async( + collection: str, + model_tag: str, + query_vector: List[float], + *, + top_k: int, + filters: Optional[Dict[str, Any]] = None, + readonly: bool = False, + nprobes: Optional[int] = None, + refine_factor: Optional[int] = None, + user_id: Optional[int] = None, + is_admin: bool = False, +) -> Tuple[List[SearchResult], str, Optional[str]]: + """ + Execute dense vector search using async vector store abstraction. + + This is the async variant of search_dense_engine. It uses the + VectorIndexStore.search_vectors_async() method instead of raw + LanceDB connection. + + Args: + collection: Collection name for data isolation + model_tag: Model tag to determine which embeddings table to search + query_vector: Query vector for similarity search + top_k: Number of top results to return + filters: Optional filters to apply to the search + readonly: If True, don't trigger index creation + nprobes: Number of partitions to probe (passed to underlying store if supported) + refine_factor: Refine factor for re-ranking (passed to underlying store if supported) + user_id: Optional user ID for multi-tenancy filtering + is_admin: Whether the user has admin privileges + + Returns: + Tuple of (search_results, index_status, index_advice) + """ + try: + vector_store = get_vector_index_store() + + # Check and create index if needed (using storage abstraction) + index_result_obj = vector_store.create_index(model_tag, readonly) + index_status = index_result_obj.status + index_advice = index_result_obj.advice + + # Convert API-facing dict filters into abstract FilterExpression + filter_expr: Optional[FilterExpression] = None + if collection or filters: + conditions: List[FilterExpression] = [] + + if collection: + from ..storage.contracts import FilterCondition, FilterOperator + + conditions.append( + FilterCondition( + field="collection", + operator=FilterOperator.EQ, + value=collection, + ) + ) + + if filters: + parsed = ( + parse_legacy_filters(filters) if isinstance(filters, dict) else None + ) + if parsed is not None: + if isinstance(parsed, tuple): + conditions.extend(parsed) + else: + conditions.append(parsed) + + if len(conditions) == 1: + filter_expr = conditions[0] + elif len(conditions) > 1: + filter_expr = tuple(conditions) + + # Validate filter expression depth to prevent DoS + if filter_expr is not None: + validate_filter_depth(filter_expr) + + # Execute async vector search using abstraction layer (by model_tag) + raw_results = await vector_store.search_vectors_by_model_async( + model_tag=model_tag, + query_vector=query_vector, + top_k=top_k, + filters=filter_expr, + vector_column_name="vector", + user_id=user_id, + is_admin=is_admin, + ) + + # Convert raw results to SearchResult objects + search_results = [] + for row in raw_results: + # LanceDB returns Squared Euclidean Distance (L_2^{2} distance) + distance_value = row.get("_distance") + distance = float(distance_value) if distance_value is not None else 0.0 + score = 1.0 / (1.0 + distance) + + # Deserialize metadata from JSON string to dictionary + metadata = deserialize_metadata(row.get("metadata")) + + search_result = SearchResult( + doc_id=row["doc_id"], + chunk_id=row["chunk_id"], + text=row["text"], + score=score, + parse_hash=row.get("parse_hash"), + model_tag=model_tag, + created_at=row.get("created_at"), + metadata=metadata, + ) + search_results.append(search_result) + + return search_results, index_status, index_advice + + except Exception as e: + logger.error(f"Failed to execute async dense search: {str(e)}") + raise diff --git a/src/xagent/core/tools/core/RAG_tools/retrieval/search_sparse.py b/src/xagent/core/tools/core/RAG_tools/retrieval/search_sparse.py index bb7976e13..2c23ffc69 100644 --- a/src/xagent/core/tools/core/RAG_tools/retrieval/search_sparse.py +++ b/src/xagent/core/tools/core/RAG_tools/retrieval/search_sparse.py @@ -1,24 +1,25 @@ from __future__ import annotations import logging -from typing import Any, Dict, Iterable, List, Optional, Set +from collections.abc import AsyncIterator +from typing import Any, Dict, Iterable, List, Optional, Set, cast import pandas as pd import pyarrow as pa # type: ignore from pyarrow import Table as PyArrowTable -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.schemas import ( SearchFallbackAction, SearchResult, SearchWarning, SparseSearchResponse, ) -from ..LanceDB.model_tag_utils import to_model_tag +from ..storage.contracts import FilterExpression +from ..storage.factory import ( + get_vector_index_store, +) +from ..utils.filter_utils import parse_legacy_filters, validate_filter_depth from ..utils.metadata_utils import deserialize_metadata -from ..utils.string_utils import build_lancedb_filter_expression -from ..utils.user_permissions import UserPermissions -from ..vector_storage.index_manager import get_index_manager logger = logging.getLogger(__name__) @@ -38,7 +39,6 @@ def search_sparse( ) -> SparseSearchResponse: """Performs sparse (Full-Text Search) retrieval on the specified collection.""" - table_name = f"embeddings_{to_model_tag(model_tag)}" _fts_enabled = False current_warnings: List[SearchWarning] = [] @@ -46,25 +46,30 @@ def search_sparse( current_warnings.append( SearchWarning( code="READONLY_MODE", - message=f"Readonly mode enabled for sparse search on {table_name}. No FTS index operations will be performed.", + message=f"Readonly mode enabled for sparse search on {model_tag}. No FTS index operations will be performed.", fallback_action=SearchFallbackAction.REBUILD_INDEX, affected_models=[model_tag], ) ) try: - conn = get_connection_from_env() - table = conn.open_table(table_name) + vector_store = get_vector_index_store() + + # Open embeddings table with legacy fallback (handled by abstraction layer) + # open_embeddings_table will handle adding the "embeddings_" prefix + table, actual_table_name = vector_store.open_embeddings_table(model_tag) + + # Use storage abstraction for index management + index_result_obj = vector_store.create_index(model_tag, readonly) - index_manager = get_index_manager() - _, _ = index_manager.check_and_create_index(table, table_name, readonly) - _fts_enabled = index_manager.get_fts_index_status(table) + # Use FTS enabled status from index result + _fts_enabled = index_result_obj.fts_enabled if not _fts_enabled: current_warnings.append( SearchWarning( code="FTS_INDEX_MISSING", - message=f"FTS index not found on 'text' column for {table_name}. Sparse search performance may be degraded.", + message=f"FTS index not found on 'text' column for {model_tag}. Sparse search performance may be degraded.", fallback_action=SearchFallbackAction.REBUILD_INDEX, affected_models=[model_tag], ) @@ -72,34 +77,66 @@ def search_sparse( search_query = table.search(query_text, query_type="fts").limit(top_k) - # Build filter expression combining collection scope, user permissions and custom filters - filter_clauses = [] - - # Scope results to the requested collection (required for KB isolation) - if collection: - collection_filter = build_lancedb_filter_expression( - {"collection": collection} - ) - if collection_filter: - filter_clauses.append(collection_filter) + # Convert legacy dict format to FilterExpression if needed + filter_expr: Optional[FilterExpression] = None + if collection or filters: + # Build filter conditions + conditions: List[FilterExpression] = [] - # Add user permission filter for multi-tenancy - user_filter = UserPermissions.get_user_filter(user_id, is_admin) - if user_filter: - filter_clauses.append(user_filter) + # Add collection filter + if collection: + from ..storage.contracts import FilterCondition, FilterOperator - # Add custom filters if provided - if filters: - custom_filter = build_lancedb_filter_expression(filters) - if custom_filter: - filter_clauses.append(custom_filter) + conditions.append( + FilterCondition( + field="collection", operator=FilterOperator.EQ, value=collection + ) + ) - # Combine all filters with AND - if filter_clauses: - combined_filter = " and ".join(f"({clause})" for clause in filter_clauses) - search_query = search_query.where(combined_filter) + # Add custom filters + if filters: + if isinstance(filters, dict): + # Legacy format: use parser + parsed_filters = parse_legacy_filters(filters) + # parsed_filters can be FilterCondition or tuple (AND combination) + if parsed_filters is not None: + if isinstance(parsed_filters, tuple): + # Type narrowing: tuple of FilterConditions + conditions.extend(parsed_filters) + else: + # Type narrowing: single FilterCondition + conditions.append(parsed_filters) + elif isinstance(filters, (tuple, list)): + # Already FilterExpression + conditions.extend( + filters if isinstance(filters, tuple) else list(filters) + ) + else: + # Single FilterCondition + conditions.append(filters) + + # Combine conditions with AND + if len(conditions) == 1: + filter_expr = conditions[0] + elif len(conditions) > 1: + filter_expr = tuple(conditions) + + # Validate filter expression depth to prevent DoS + if filter_expr is not None: + validate_filter_depth(filter_expr) + + # Use abstract filter builder to get backend-specific syntax + if filter_expr: + backend_filter = vector_store.build_filter_expression( + filters=filter_expr, + user_id=user_id, + is_admin=is_admin, + ) + if backend_filter: + search_query = search_query.where(backend_filter) - raw_results_df: pd.DataFrame = search_query.to_pandas() + # LanceDB's search().to_pandas() returns Any due to missing type stubs + raw_results_df = pd.DataFrame(search_query.to_pandas()) if not raw_results_df.empty: search_results: List[SearchResult] = [] @@ -156,7 +193,7 @@ def search_sparse( except Exception as e: logger.error( - f"Sparse search failed for {table_name} with query '{query_text}': {e}" + f"Sparse search failed for {model_tag} with query '{query_text}': {e}" ) error_warnings = current_warnings + [ SearchWarning( @@ -298,3 +335,280 @@ def _build_sparse_response( fts_enabled=fts_enabled, query_text=query_text, ) + + +# --- Async variant (Phase 1A Option C) --- + + +async def search_sparse_async( + collection: str, + model_tag: str, + query_text: str, + *, + top_k: int, + filters: Optional[Dict[str, Any]] = None, + readonly: bool = False, + nprobes: Optional[int] = None, + refine_factor: Optional[int] = None, + user_id: Optional[int] = None, + is_admin: bool = False, +) -> SparseSearchResponse: + """ + Perform sparse (Full-Text Search) retrieval using async vector store abstraction. + + This is the async variant of search_sparse. It uses VectorIndexStore.search_fts_async() + instead of raw LanceDB connection for the main search path. + + Note: FTS index creation uses VectorIndexStore.create_index() for full decoupling. + """ + vector_store = get_vector_index_store() + + _fts_enabled = False + current_warnings: List[SearchWarning] = [] + + if readonly: + current_warnings.append( + SearchWarning( + code="READONLY_MODE", + message=f"Readonly mode enabled for sparse search on {model_tag}. No FTS index operations will be performed.", + fallback_action=SearchFallbackAction.REBUILD_INDEX, + affected_models=[model_tag], + ) + ) + + try: + # Check and create FTS index if needed (using storage abstraction layer) + if not readonly: + index_result_obj = vector_store.create_index(model_tag, readonly=False) + _fts_enabled = index_result_obj.fts_enabled + + if not _fts_enabled: + current_warnings.append( + SearchWarning( + code="FTS_INDEX_MISSING", + message=f"FTS index may not be enabled on 'text' column for {model_tag}. Sparse search performance may be degraded.", + fallback_action=SearchFallbackAction.REBUILD_INDEX, + affected_models=[model_tag], + ) + ) + + # Convert API-facing dict filters into abstract FilterExpression + filter_expr: Optional[FilterExpression] = None + if collection or filters: + conditions: List[FilterExpression] = [] + + if collection: + from ..storage.contracts import FilterCondition, FilterOperator + + conditions.append( + FilterCondition( + field="collection", operator=FilterOperator.EQ, value=collection + ) + ) + + if filters: + if isinstance(filters, dict): + parsed_filters = parse_legacy_filters(filters) + if parsed_filters is not None: + if isinstance(parsed_filters, tuple): + conditions.extend(parsed_filters) + else: + conditions.append(parsed_filters) + elif isinstance(filters, (tuple, list)): + conditions.extend( + filters if isinstance(filters, tuple) else list(filters) + ) + else: + conditions.append(filters) + + if len(conditions) == 1: + filter_expr = conditions[0] + elif len(conditions) > 1: + filter_expr = tuple(conditions) + + # Validate filter expression depth to prevent DoS + if filter_expr is not None: + validate_filter_depth(filter_expr) + + # Execute async FTS search using abstraction layer (by model_tag) + raw_results = await vector_store.search_fts_by_model_async( + model_tag=model_tag, + query_text=query_text, + top_k=top_k, + filters=filter_expr, + text_column_name="text", + ) + + if not raw_results: + logger.warning( + "FTS lookup returned no results for query '%s'; falling back to substring match", + query_text, + ) + # Use async iter_batches for fallback + fallback_results = await _substring_fallback_async( + model_tag=model_tag, + collection=collection, + query_text=query_text, + top_k=top_k, + filters=filters, + current_warnings=current_warnings, + user_id=user_id, + is_admin=is_admin, + ) + + return _build_sparse_response( + results=fallback_results, + warnings=current_warnings, + fts_enabled=_fts_enabled, + query_text=query_text, + ) + + # Convert raw results to SearchResult objects + search_results: List[SearchResult] = [] + for row in raw_results: + # LanceDB FTS returns TF-IDF score (higher is better) + raw_score_value = row.get("_score") + raw_score = float(raw_score_value) if raw_score_value is not None else 0.0 + # Normalize TF-IDF score to [0, 1) range + score = raw_score / (1.0 + raw_score) + + # Deserialize metadata + metadata = deserialize_metadata(row.get("metadata")) + + search_results.append( + SearchResult( + doc_id=row["doc_id"], + chunk_id=row["chunk_id"], + text=row["text"], + score=score, + parse_hash=row.get("parse_hash"), + model_tag=model_tag, + created_at=row.get("created_at"), + metadata=metadata, + ) + ) + + return _build_sparse_response( + results=search_results, + warnings=current_warnings, + fts_enabled=_fts_enabled, + query_text=query_text, + ) + + except Exception as e: + logger.error( + f"Async sparse search failed for {model_tag} with query '{query_text}': {e}" + ) + error_warnings = current_warnings + [ + SearchWarning( + code="FTS_SEARCH_FAILED", + message=f"An unexpected error occurred during sparse search: {str(e)}", + fallback_action=SearchFallbackAction.PARTIAL_RESULTS, + affected_models=[model_tag], + ) + ] + return _build_sparse_response( + results=[], + warnings=error_warnings, + fts_enabled=_fts_enabled, + query_text=query_text, + status="failed", + ) + + +async def _substring_fallback_async( + *, + model_tag: str, + collection: str, + query_text: str, + top_k: int, + filters: Optional[Dict[str, Any]], + current_warnings: List[SearchWarning], + user_id: Optional[int] = None, + is_admin: bool = False, + batch_size: int = 2048, +) -> List[SearchResult]: + """Perform async substring scan using iter_batches_async when FTS misses.""" + + vector_store = get_vector_index_store() + results: List[SearchResult] = [] + + # Build query filters + query_filters: Dict[str, Any] = {"collection": collection} + if filters: + query_filters.update(filters) + + try: + # Open embeddings table with legacy fallback + _table, table_name = vector_store.open_embeddings_table(model_tag) + + # Use async batch iteration for memory-efficient scanning + # Specify only required columns to minimize memory usage + async for batch in cast( + AsyncIterator[Any], + vector_store.iter_batches_async( + table_name=table_name, + columns=[ + "doc_id", + "chunk_id", + "text", + "parse_hash", + "created_at", + "metadata", + ], + batch_size=batch_size, + filters=query_filters, + user_id=user_id, + is_admin=is_admin, + ), + ): + batch_df = batch.to_pandas() + + # Apply substring filter + text_mask = ( + batch_df["text"] + .astype(str) + .str.contains(query_text, na=False, regex=False) + ) + matching_rows = batch_df[text_mask] + + # Early exit: stop processing if we already have enough results + if len(results) >= top_k: + break + + for _, row in matching_rows.iterrows(): + metadata = deserialize_metadata(row.get("metadata")) + results.append( + SearchResult( + doc_id=row["doc_id"], + chunk_id=row["chunk_id"], + text=row["text"], + score=1.0, + parse_hash=row["parse_hash"], + model_tag=model_tag, + created_at=row["created_at"], + metadata=metadata, + ) + ) + + # Early exit: stop as soon as we have enough results + if len(results) >= top_k: + break + + if results: + current_warnings.append( + SearchWarning( + code="FTS_FALLBACK", + message=( + "Full-text index returned no matches; used async substring search fallback. " + "Check FTS tokenizer configuration or update LanceDB to ensure proper tokenisation for query language." + ), + fallback_action=SearchFallbackAction.BRUTE_FORCE, + affected_models=[model_tag], + ) + ) + + except Exception as exc: + logger.error("Async substring fallback failed: %s", exc) + + return results diff --git a/src/xagent/core/tools/core/RAG_tools/storage/__init__.py b/src/xagent/core/tools/core/RAG_tools/storage/__init__.py new file mode 100644 index 000000000..2ffe9b5d5 --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/storage/__init__.py @@ -0,0 +1,54 @@ +"""Storage contracts and default implementations for KB. + +Phase 1A Part 2: Extended with additional store contracts for complete decoupling. +""" + +from .contracts import ( + IngestionStatusStore, + KBWriteCoordinator, + MainPointerStore, + MetadataStore, + PromptTemplateStore, + VectorIndexStore, +) +from .factory import ( + StorageFactory, + get_ingestion_status_store, + get_kb_write_coordinator, + get_main_pointer_store, + get_metadata_store, + get_prompt_template_store, + get_vector_index_store, + get_vector_store_raw_connection, + reset_kb_write_coordinator, +) +from .vector_backend import ( + VECTOR_BACKEND_ENV, + VECTOR_BACKEND_ENV_LEGACY, + VectorBackend, + get_configured_vector_backend, +) + +__all__ = [ + # Contracts + "KBWriteCoordinator", + "MetadataStore", + "VectorIndexStore", + "IngestionStatusStore", + "PromptTemplateStore", + "MainPointerStore", + # Factory + "StorageFactory", + "get_kb_write_coordinator", + "get_metadata_store", + "get_vector_index_store", + "get_vector_store_raw_connection", + "VectorBackend", + "VECTOR_BACKEND_ENV", + "VECTOR_BACKEND_ENV_LEGACY", + "get_configured_vector_backend", + "get_ingestion_status_store", + "get_prompt_template_store", + "get_main_pointer_store", + "reset_kb_write_coordinator", +] diff --git a/src/xagent/core/tools/core/RAG_tools/storage/contracts.py b/src/xagent/core/tools/core/RAG_tools/storage/contracts.py new file mode 100644 index 000000000..8ccab0821 --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/storage/contracts.py @@ -0,0 +1,1628 @@ +"""Storage contracts for KB control-plane and vector-plane operations. + +Phase 1A introduces these contracts to decouple API/business modules from +backend-specific database semantics. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum +from typing import ( + Any, + Dict, + Iterator, + List, + Optional, + Protocol, + Sequence, + Tuple, + Union, + runtime_checkable, +) + +from ..core.config import DEFAULT_VECTOR_STORE_SCAN_LIMIT, IndexPolicy +from ..core.schemas import CollectionInfo, IndexResult + +# Field name whitelist for filter validation +# Derived from all LanceDB table schemas in schema_manager.py +_VALID_FILTER_FIELDS = frozenset( + { + # documents table + "collection", + "doc_id", + "source_path", + "file_type", + "content_hash", + "uploaded_at", + "title", + "language", + "user_id", + # parses table + "parse_hash", + "parser", + "created_at", + "params_json", + # chunks table + "chunk_id", + "index", + "page_number", + "section", + "anchor", + "json_path", + "chunk_hash", + "config_hash", + "metadata", + # embeddings table + "model", + "vector_dimension", + "vector", + # ingestion_runs table + "status", + "message", + "updated_at", + # main_pointers table + "step_type", + "model_tag", + "semantic_id", + "technical_id", + "operator", + # prompt_templates table + "id", + "name", + "template", + "version", + "is_latest", + # collection_metadata table + "name", + "schema_version", + "embedding_model_id", + "embedding_dimension", + "documents", + "processed_documents", + "parses", + "chunks", + "embeddings", + "document_names", + "collection_locked", + "allow_mixed_parse_methods", + "skip_config_validation", + "ingestion_config", + "created_at", + "updated_at", + "last_accessed_at", + "extra_metadata", + # collection_config table + "config_json", + } +) + + +def validate_field_name(field: str) -> None: + """Validate that a field name is in the allowed whitelist. + + Args: + field: Field name to validate. + + Raises: + ValueError: If field name is not in the whitelist. + """ + if field not in _VALID_FILTER_FIELDS: + raise ValueError( + f"Invalid filter field '{field}'. " + f"Field must be one of: {', '.join(sorted(_VALID_FILTER_FIELDS))}" + ) + + +def validate_filter_value(value: Any) -> None: + """Validate that a filter value is an allowed type. + + Allowed types: str, int, float, bool, None, list, tuple, set. + + Args: + value: Value to validate. + + Raises: + ValueError: If value type is not allowed. + TypeError: If value is a complex object (dict, custom class). + """ + if value is None: + return + + if isinstance(value, (str, int, float, bool)): + return + + if isinstance(value, (list, tuple, set)): + # Validate each element in the collection + for item in value: + if not isinstance(item, (str, int, float, bool, type(None))): + raise TypeError( + f"Invalid filter value type in collection: {type(item).__name__}. " + f"Collection elements must be str, int, float, bool, or None." + ) + return + + # Reject dict and complex objects + raise TypeError( + f"Invalid filter value type: {type(value).__name__}. " + f"Allowed types: str, int, float, bool, None, list, tuple, set." + ) + + +def build_filter_from_dict(filters: Dict[str, Any]) -> Optional[FilterExpression]: + """Convert a dictionary of filters to a FilterExpression with validation. + + This function provides a common entry point for building filter expressions + from simple dictionary key-value pairs. All keys are validated against the + field name whitelist, and all values are type-checked. + + Args: + filters: Dictionary of field-name -> value mappings for equality filters. + + Returns: + FilterExpression: Single FilterCondition for one filter, + tuple of conditions (AND) for multiple filters, + or None if filters is empty. + + Raises: + ValueError: If a field name is not in the whitelist. + TypeError: If a value type is not allowed. + + Example: + >>> build_filter_from_dict({"collection": "my_collection", "doc_id": "doc123"}) + (FilterCondition(field='collection', operator=FilterOperator.EQ, value='my_collection'), + FilterCondition(field='doc_id', operator=FilterOperator.EQ, value='doc123')) + + >>> build_filter_from_dict({"doc_id": "doc123"}) + FilterCondition(field='doc_id', operator=FilterOperator.EQ, value='doc123') + """ + if not filters: + return None + + conditions = [] + for field, value in filters.items(): + # Validate field name + validate_field_name(field) + + # Validate value type + validate_filter_value(value) + + # Create filter condition + conditions.append( + FilterCondition(field=field, operator=FilterOperator.EQ, value=value) + ) + + # Return single condition or tuple (AND combination) + if len(conditions) == 1: + return conditions[0] + return tuple(conditions) + + +@runtime_checkable +class DatabaseConnection(Protocol): + """Backend-agnostic database connection protocol. + + This protocol defines the minimal interface required for storage + implementations to work with different database backends without + importing concrete types like LanceDB's DBConnection. + """ + + def open_table(self, name: str) -> Any: ... + + def table_names(self) -> Sequence[str]: ... + + +@dataclass(frozen=True) +class DocumentRecord: + """Lightweight document projection for metadata/control operations. + + Attributes: + doc_id: Document identifier. + file_id: Optional file identifier for uploaded file tracking. + source_path: Original source path if available. + """ + + doc_id: str + file_id: Optional[str] = None + source_path: Optional[str] = None + + +class FilterOperator(str, Enum): + """Comparison operators for filter expressions. + + These operators provide a backend-agnostic way to express filter conditions + that can be translated to backend-specific query languages. + """ + + EQ = "eq" # Equal + NE = "ne" # Not equal + GT = "gt" # Greater than + GTE = "gte" # Greater than or equal + LT = "lt" # Less than + LTE = "lte" # Less than or equal + IN = "in" # In list + CONTAINS = "contains" # String contains + IS_NULL = "is_null" # Is NULL + IS_NOT_NULL = "is_not_null" # Is not NULL + + +@dataclass(frozen=True) +class FilterCondition: + """Single filter condition. + + Attributes: + field: Field name to filter on. + operator: Comparison operator. + value: Value to compare against. + + Raises: + ValueError: If operator requires list value but value is not a list. + """ + + field: str + operator: FilterOperator + value: Any + + def __post_init__(self) -> None: + # Validate operator matches value type + if self.operator in {FilterOperator.IN}: + if not isinstance(self.value, (list, tuple, set)): + raise ValueError( + f"IN operator requires list/tuple/set value, got {type(self.value)}" + ) + + +# Filter expression can be a single condition, AND combination (tuple), or OR combination (list) +# Use string annotation for recursive type definition +FilterExpression = Union[ + FilterCondition, # Single condition + "tuple[FilterExpression, ...]", # AND combination + "list[FilterExpression]", # OR combination +] + + +class MetadataStore(ABC): + """Control-plane metadata storage contract.""" + + @abstractmethod + async def get_collection(self, collection_name: str) -> CollectionInfo: + """Read collection metadata. + + Args: + collection_name: Target collection name. + + Returns: + Collection metadata. + + Raises: + ValueError: If collection is not found. + """ + + @abstractmethod + async def save_collection(self, collection: CollectionInfo) -> None: + """Create or update collection metadata.""" + + @abstractmethod + async def ensure_collection_metadata_table(self) -> None: + """Ensure control-plane metadata table exists.""" + + @abstractmethod + async def save_collection_config( + self, + collection: str, + config_json: str, + user_id: int, + ) -> None: + """Save collection ingestion configuration. + + Args: + collection: Collection name. + config_json: JSON string of IngestionConfig. + user_id: User ID for multi-tenancy. + """ + + @abstractmethod + async def get_collection_config( + self, + collection: str, + user_id: Optional[int], + is_admin: bool = False, + ) -> str | None: + """Get collection ingestion configuration. + + Args: + collection: Collection name. + user_id: User ID for multi-tenancy. None is treated as 0 for non-admin, + and as "load all configs" for admin mode. + is_admin: Whether user has admin privileges (bypasses user_id filter). + + Returns: + Config JSON string if found, None otherwise. + """ + + @abstractmethod + def get_raw_connection(self) -> Any: + """Return raw backend connection for legacy compatibility paths. + + The returned object conforms to the DatabaseConnection protocol but + uses Any type to avoid importing backend-specific types. + """ + + +class VectorIndexStore(ABC): + """Vector/data-plane storage contract. + + Phase 1A Option C: Hybrid sync/async methods for gradual migration. + Sync methods provide backward compatibility; async methods enable + non-blocking operations in async contexts (FastAPI, etc.). + """ + + @abstractmethod + def list_document_records( + self, + collection_name: Optional[str], + user_id: Optional[int], + is_admin: bool, + max_results: int = DEFAULT_VECTOR_STORE_SCAN_LIMIT, + ) -> List[DocumentRecord]: + """List document records from vector index side. + + Args: + collection_name: Optional collection name filter. If None, lists records across all collections. + user_id: User ID for multi-tenancy filtering. + is_admin: Whether the user has admin privileges. + max_results: Maximum records to return. + """ + + @abstractmethod + def rename_collection_data( + self, + collection_name: str, + new_name: str, + ) -> List[str]: + """Rename collection key across vector-side tables. + + Returns: + Warning messages generated during best-effort updates. + """ + + @abstractmethod + def delete_collection_data( + self, + collection_name: str, + ) -> Dict[str, int]: + """Delete all data for a collection from vector-side tables. + + Args: + collection_name: Name of the collection to delete. + + Returns: + Dictionary mapping table names to deleted row counts. + """ + + @abstractmethod + def aggregate_collection_stats( + self, + user_id: Optional[int], + is_admin: bool, + ) -> Dict[str, Dict[str, int]]: + """Aggregate statistics for all collections. + + Returns: + Dictionary mapping collection names to their stats: + { + "collection_name": { + "documents": int, + "parses": int, + "chunks": int, + "embeddings": int, + } + } + """ + + @abstractmethod + def aggregate_document_stats( + self, + collection_name: str, + doc_id: str, + user_id: Optional[int], + is_admin: bool, + ) -> Dict[str, int]: + """Aggregate statistics for a single document. + + Returns: + Dictionary with counts: + { + "documents": int, + "parses": int, + "chunks": int, + "embeddings": int, + } + """ + + @abstractmethod + def list_table_names(self) -> Sequence[str]: + """List backend table names.""" + + @abstractmethod + def get_vector_dimension(self, table_name: str) -> Optional[int]: + """Get the vector dimension from a table's schema. + + Reads the vector field's fixed_size dimension from the table schema. + Returns None if the vector field is variable-length or dimension cannot + be determined. + + Args: + table_name: Name of the embeddings table to inspect. + + Returns: + Vector dimension as int, or None if variable-length/unavailable. + """ + + @abstractmethod + def open_embeddings_table(self, model_tag: str) -> Tuple[Any, str]: + """Open embeddings table with legacy fallback support. + + Tries the primary Hub ID-based table name first, then falls back + to legacy provider-based naming if the primary doesn't exist. + + This method encapsulates the legacy fallback logic for embeddings tables, + providing a single source of truth for table name resolution. + + Args: + model_tag: Model tag for the embeddings table. + + Returns: + Tuple of (table_object, actual_table_name_used). + + Raises: + DatabaseOperationError: If neither primary nor legacy table exists. + """ + + @abstractmethod + def iter_batches( + self, + table_name: str, + columns: Optional[Sequence[str]] = None, + batch_size: int = 1000, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> Iterator[Any]: + """Iterate over table data in batches (sync). + + Yields backend-specific batch objects (e.g., PyArrow RecordBatch). + This method is designed for memory-efficient processing of large tables. + + Args: + table_name: Name of table to iterate. + columns: Optional columns to select. If None, selects all columns. + batch_size: Rows per batch. + filters: Optional filter criteria (key-value pairs for equality). + user_id: Optional user filter for multi-tenancy. + is_admin: Admin privilege flag. + + Yields: + Backend-specific batch objects (e.g., PyArrow RecordBatch). + """ + + @abstractmethod + def count_rows( + self, + table_name: str, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> int: + """Count rows in a table with optional filters (sync). + + Args: + table_name: Name of table to count. + filters: Optional filter criteria (key-value pairs for equality). + user_id: Optional user filter for multi-tenancy. + is_admin: Admin privilege flag. + + Returns: + Row count. + + Raises: + DatabaseOperationError: If table cannot be opened or count fails. + """ + + def count_rows_or_zero( + self, + table_name: str, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> int: + """Count rows in a table, returning 0 if table doesn't exist. + + This is a convenience method for existence checks where a missing table + should be treated as "no data" rather than an error. + + Args: + table_name: Name of table to count. + filters: Optional filter criteria (key-value pairs for equality). + user_id: Optional user filter for multi-tenancy. + is_admin: Admin privilege flag. + + Returns: + Row count, or 0 if table doesn't exist or count fails. + """ + from ..core.exceptions import DatabaseOperationError + + try: + return self.count_rows(table_name, filters, user_id, is_admin) + except DatabaseOperationError: + return 0 + + @abstractmethod + def aggregate_document_counts( + self, + table_name: str, + doc_id_column: str, + collection_name: str, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> Dict[str, int]: + """Aggregate records per document for a specific table. + + Args: + table_name: Table to aggregate from. + doc_id_column: Column containing document IDs. + collection_name: Collection to scope to. + user_id: Optional user filter. + is_admin: Admin privilege flag. + + Returns: + Dictionary mapping doc_id to count. + """ + + @abstractmethod + def build_filter_expression( + self, + filters: Optional[FilterExpression], + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> Optional[str]: + """Convert abstract filter expression to backend-specific syntax. + + Args: + filters: Abstract filter expression. + user_id: Optional user for multi-tenancy. + is_admin: Admin privilege flag. + + Returns: + Backend-specific filter string, or None if no filters. + """ + + @abstractmethod + def upsert_documents(self, records: List[Dict[str, Any]]) -> None: + """Upsert document records (sync). + + Args: + records: List of document record dictionaries to upsert. + """ + + @abstractmethod + def upsert_parses(self, records: List[Dict[str, Any]]) -> None: + """Upsert parse records (sync). + + Args: + records: List of parse record dictionaries to upsert. + """ + + @abstractmethod + def upsert_chunks(self, records: List[Dict[str, Any]]) -> None: + """Upsert chunk records (sync). + + Args: + records: List of chunk record dictionaries to upsert. + """ + + @abstractmethod + def upsert_embeddings(self, model_tag: str, records: List[Dict[str, Any]]) -> None: + """Upsert embedding records (sync). + + Args: + model_tag: Model tag for the embeddings table. + records: List of embedding record dictionaries to upsert. + """ + + @abstractmethod + def create_index(self, model_tag: str, readonly: bool = False) -> IndexResult: + """Create or check vector index for embeddings table. + + Args: + model_tag: Model tag for the embeddings table. + readonly: If True, don't trigger index creation. + + Returns: + IndexResult containing status, advice, and FTS enabled state. + """ + + @abstractmethod + def search_vectors( + self, + table_name: str, + query_vector: List[float], + *, + top_k: int, + filters: Optional[FilterExpression] = None, + vector_column_name: str = "vector", + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> List[Dict[str, Any]]: + """Execute vector search (sync). + + Args: + table_name: Name of embeddings table to search. + query_vector: Query vector for similarity search. + top_k: Number of top results to return. + filters: Optional abstract filter expression. + vector_column_name: Name of vector column (default "vector"). + user_id: Optional user ID for multi-tenancy filtering. + is_admin: Whether the user has admin privileges. + + Returns: + List of search result dictionaries with keys: + - doc_id: Document ID + - chunk_id: Chunk ID + - text: Chunk text + - _distance: Distance score (lower is better) + - metadata: Additional metadata + """ + + def search_vectors_by_model( + self, + model_tag: str, + query_vector: List[float], + *, + top_k: int, + filters: Optional[FilterExpression] = None, + vector_column_name: str = "vector", + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> List[Dict[str, Any]]: + """Convenience method: search vectors by model_tag with automatic table resolution. + + This method combines open_embeddings_table() + search_vectors() for + simpler API when searching by model_tag. + + Args: + model_tag: Model tag for the embeddings table. + query_vector: Query vector for similarity search. + top_k: Number of top results to return. + filters: Optional abstract filter expression. + vector_column_name: Name of vector column (default "vector"). + user_id: Optional user ID for multi-tenancy filtering. + is_admin: Whether the user has admin privileges. + + Returns: + List of search result dictionaries with keys: + - doc_id: Document ID + - chunk_id: Chunk ID + - text: Chunk text + - _distance: Distance score (lower is better) + - metadata: Additional metadata + """ + _table, table_name = self.open_embeddings_table(model_tag) + return self.search_vectors( + table_name=table_name, + query_vector=query_vector, + top_k=top_k, + filters=filters, + vector_column_name=vector_column_name, + user_id=user_id, + is_admin=is_admin, + ) + + # --- Async variants (Phase 1A Option C: Hybrid approach) --- + + @abstractmethod + async def search_vectors_async( + self, + table_name: str, + query_vector: List[float], + *, + top_k: int, + filters: Optional[FilterExpression] = None, + vector_column_name: str = "vector", + ) -> List[Dict[str, Any]]: + """Execute vector search (async). + + Args: + table_name: Name of embeddings table to search. + query_vector: Query vector for similarity search. + top_k: Number of top results to return. + filters: Optional abstract filter expression. + vector_column_name: Name of vector column (default "vector"). + + Returns: + List of search result dictionaries with keys: + - doc_id: Document ID + - chunk_id: Chunk ID + - text: Chunk text + - _distance: Distance score (lower is better) + - metadata: Additional metadata + """ + + async def search_vectors_by_model_async( + self, + model_tag: str, + query_vector: List[float], + *, + top_k: int, + filters: Optional[FilterExpression] = None, + vector_column_name: str = "vector", + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> List[Dict[str, Any]]: + """Convenience method: search vectors by model_tag with automatic table resolution (async). + + This method combines open_embeddings_table() + search_vectors_async() for + simpler API when searching by model_tag. + + Args: + model_tag: Model tag for the embeddings table. + query_vector: Query vector for similarity search. + top_k: Number of top results to return. + filters: Optional abstract filter expression. + vector_column_name: Name of vector column (default "vector"). + user_id: Optional user ID for multi-tenancy filtering. + is_admin: Whether the user has admin privileges. + + Returns: + List of search result dictionaries with keys: + - doc_id: Document ID + - chunk_id: Chunk ID + - text: Chunk text + - _distance: Distance score (lower is better) + - metadata: Additional metadata + """ + _table, table_name = self.open_embeddings_table(model_tag) + return await self.search_vectors_async( + table_name=table_name, + query_vector=query_vector, + top_k=top_k, + filters=filters, + vector_column_name=vector_column_name, + ) + + @abstractmethod + async def search_fts_async( + self, + table_name: str, + query_text: str, + *, + top_k: int, + filters: Optional[FilterExpression] = None, + text_column_name: str = "text", + ) -> List[Dict[str, Any]]: + """Execute full-text search (async). + + Args: + table_name: Name of embeddings/table to search (must have FTS index). + query_text: Query text for full-text search. + top_k: Number of top results to return. + filters: Optional abstract filter expression. + text_column_name: Name of text column with FTS index (default "text"). + + Returns: + List of search result dictionaries with keys: + - doc_id: Document ID + - chunk_id: Chunk ID + - text: Chunk text + - _score: TF-IDF score (higher is better) + - metadata: Additional metadata + + Raises: + DatabaseOperationError: If FTS index is not configured or search fails. + """ + + async def search_fts_by_model_async( + self, + model_tag: str, + query_text: str, + *, + top_k: int, + filters: Optional[FilterExpression] = None, + text_column_name: str = "text", + ) -> List[Dict[str, Any]]: + """Convenience method: search FTS by model_tag with automatic table resolution. + + This method combines open_embeddings_table() + search_fts_async() for + simpler API when searching by model_tag. + + Args: + model_tag: Model tag for the embeddings table. + query_text: Query text for full-text search. + top_k: Number of top results to return. + filters: Optional abstract filter expression. + text_column_name: Name of text column with FTS index (default "text"). + + Returns: + List of search result dictionaries with keys: + - doc_id: Document ID + - chunk_id: Chunk ID + - text: Chunk text + - _score: TF-IDF score (higher is better) + - metadata: Additional metadata + + Raises: + DatabaseOperationError: If FTS index is not configured or search fails. + """ + _table, table_name = self.open_embeddings_table(model_tag) + return await self.search_fts_async( + table_name=table_name, + query_text=query_text, + top_k=top_k, + filters=filters, + text_column_name=text_column_name, + ) + + @abstractmethod + async def iter_batches_async( + self, + table_name: str, + columns: Optional[Sequence[str]] = None, + batch_size: int = 1000, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> Any: # Returns AsyncIterator (async generator), but mypy has issues with async def + AsyncIterator return type + """Iterate over table data in batches (async). + + This is an async generator that yields backend-specific batch objects + (e.g., PyArrow RecordBatch). Use with: async for batch in iter_batches_async(...) + + Args: + table_name: Name of table to iterate. + columns: Optional columns to select. + batch_size: Rows per batch. + filters: Optional filter criteria. + user_id: Optional user filter for multi-tenancy. + is_admin: Admin privilege flag. + + Yields: + Backend-specific batch objects (PyArrow RecordBatch). + """ + + @abstractmethod + async def count_rows_async( + self, + table_name: str, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> int: + """Count rows in a table with optional filters (async). + + Args: + table_name: Name of table to count. + filters: Optional filter criteria. + user_id: Optional user filter. + is_admin: Admin privilege flag. + + Returns: + Row count (0 on error). + """ + + @abstractmethod + async def get_vector_dimension_async(self, table_name: str) -> Optional[int]: + """Get the vector dimension from a table's schema (async). + + Args: + table_name: Name of the embeddings table to inspect. + + Returns: + Vector dimension as int, or None if variable-length/unavailable. + + Note: Current implementation uses sync operations under the hood. + True async I/O will be added in Phase 1B with RDB backend. + """ + + @abstractmethod + async def upsert_documents_async(self, records: List[Dict[str, Any]]) -> None: + """Upsert document records (async). + + Args: + records: List of document record dictionaries to upsert. + """ + + @abstractmethod + async def upsert_chunks_async(self, records: List[Dict[str, Any]]) -> None: + """Upsert chunk records (async). + + Args: + records: List of chunk record dictionaries to upsert. + """ + + @abstractmethod + async def upsert_embeddings_async( + self, model_tag: str, records: List[Dict[str, Any]] + ) -> None: + """Upsert embedding records (async). + + Args: + model_tag: Model tag for the embeddings table. + records: List of embedding record dictionaries to upsert. + """ + + # --- Index Management (Phase 1A Part 2) --- + + @abstractmethod + def should_reindex( + self, + table_name: str, + total_upserted: int, + policy: IndexPolicy, + ) -> bool: + """Determine if reindex should be triggered. + + Args: + table_name: Embeddings table name. + total_upserted: Total upserted records since last index. + policy: Index policy with reindex thresholds. + + Returns: + True if reindex should be triggered. + """ + + @abstractmethod + def trigger_reindex(self, table_name: str) -> bool: + """Trigger index rebuild operation. + + Args: + table_name: Embeddings table name. + + Returns: + True if reindex was triggered successfully. + """ + + # --- Async index management variants --- + + @abstractmethod + async def should_reindex_async( + self, + table_name: str, + total_upserted: int, + policy: IndexPolicy, + ) -> bool: + """Async version of should_reindex. + + Args: + table_name: Embeddings table name. + total_upserted: Total upserted records since last index. + policy: Index policy with reindex thresholds. + + Returns: + True if reindex should be triggered. + + Note: Current implementation uses sync operations under the hood. + True async I/O will be added in Phase 1B with RDB backend. + """ + + @abstractmethod + async def trigger_reindex_async(self, table_name: str) -> bool: + """Async version of trigger_reindex. + + Args: + table_name: Embeddings table name. + + Returns: + True if reindex was triggered successfully. + + Note: Current implementation uses sync operations under the hood. + True async I/O will be added in Phase 1B with RDB backend. + """ + + @abstractmethod + def migrate_embeddings_table( + self, + model_id: str, + batch_size: int = 1000, + ) -> dict[str, Any]: + """Migrate legacy embeddings table to Hub ID-based naming. + + This method copies data from a legacy table (embeddings_{model_name}) + to a new Hub ID-based table (embeddings_{hub_id}), rewriting the + per-row ``model`` field to the Hub model ID. + + This is the proper location for migration logic, as it's part of + the storage implementation. Migration should be run during maintenance + windows, not during normal read operations. + + Args: + model_id: Hub model ID to migrate (e.g., "text-embedding-ada-002"). + batch_size: Number of rows to copy per batch. + + Returns: + Dictionary with migration results: + { + "success": bool, + "source_table": str (legacy table name), + "target_table": str (Hub ID table name), + "rows_migrated": int, + "error": str | None (if success=False) + } + + Raises: + VectorValidationError: If model_id is empty. + DatabaseOperationError: If migration fails. + + Note: + - This method uses file-based locking to prevent concurrent migrations. + - The migration is idempotent and can be safely re-run. + - Source table is preserved after migration. + """ + pass + + @abstractmethod + def get_raw_connection(self) -> Any: + """Return raw backend connection for legacy compatibility paths. + + The returned object conforms to the DatabaseConnection protocol but + uses Any type to avoid importing backend-specific types. + + DEPRECATED: Use specific upsert methods instead for write operations. + """ + + +class KBWriteCoordinator(ABC): + """Contract for knowledge-base write/delete orchestration (Phase 1A shell). + + Phase 1A exposes only accessors to the configured metadata and vector + stores; concrete implementations delegate without extra coordination. + This type is a stable injection point for future write-path behavior such + as distributed locking, write batching, and conflict resolution across + metadata and vector backends. + """ + + @abstractmethod + def metadata_store(self) -> MetadataStore: + """Return configured metadata store.""" + + @abstractmethod + def vector_index_store(self) -> VectorIndexStore: + """Return configured vector index store.""" + + +# ============================================================================ +# Phase 1A Part 2: Additional Store Contracts +# ============================================================================ + + +class IngestionStatusStore(ABC): + """Ingestion status tracking contract. + + Manages ingestion_runs table for tracking document processing status. + + Phase 1A Option C: Hybrid sync/async methods for gradual migration. + Sync methods provide backward compatibility; async methods enable + non-blocking operations in async contexts. + """ + + # --- Sync methods --- + + @abstractmethod + def write_ingestion_status( + self, + collection: str, + doc_id: str, + *, + status: str, + message: Optional[str] = None, + parse_hash: Optional[str] = None, + user_id: Optional[int] = None, + ) -> None: + """Write ingestion status record (sync). + + Args: + collection: Collection name. + doc_id: Document ID. + status: Status value (e.g., 'pending', 'processing', 'success', 'failed'). + message: Optional status message or error description. + parse_hash: Optional hash of the parsed document for change detection. + user_id: Optional user ID for multi-tenancy. + + Raises: + DatabaseOperationError: If write operation fails. + """ + + @abstractmethod + def load_ingestion_status( + self, + collection: Optional[str] = None, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> List[Dict[str, Any]]: + """Load ingestion status records (sync). + + Args: + collection: Optional collection name to filter by. + doc_id: Optional document ID to filter by. + user_id: Optional user ID for multi-tenancy filtering. + is_admin: Whether user has admin privileges (bypasses filtering). + + Returns: + List of ingestion status records. + + Raises: + DatabaseOperationError: If read operation fails. + """ + + @abstractmethod + def clear_ingestion_status( + self, + collection: str, + doc_id: str, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> None: + """Remove ingestion status record (sync). + + Args: + collection: Collection name. + doc_id: Document ID. + user_id: Optional user ID for multi-tenancy filtering. + is_admin: Whether user has admin privileges (bypasses filtering). + + Raises: + DatabaseOperationError: If delete operation fails. + """ + + # --- Async methods --- + + @abstractmethod + async def write_ingestion_status_async( + self, + collection: str, + doc_id: str, + *, + status: str, + message: Optional[str] = None, + parse_hash: Optional[str] = None, + user_id: Optional[int] = None, + ) -> None: + """Write ingestion status record (async). + + Args: + collection: Collection name. + doc_id: Document ID. + status: Status value. + message: Optional status message. + parse_hash: Optional parse hash. + user_id: Optional user ID. + + Raises: + DatabaseOperationError: If write operation fails. + """ + + @abstractmethod + async def load_ingestion_status_async( + self, + collection: Optional[str] = None, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> List[Dict[str, Any]]: + """Load ingestion status records (async). + + Args: + collection: Optional collection name to filter by. + doc_id: Optional document ID to filter by. + user_id: Optional user ID for multi-tenancy. + is_admin: Whether user has admin privileges. + + Returns: + List of ingestion status records. + + Raises: + DatabaseOperationError: If read operation fails. + """ + + @abstractmethod + async def clear_ingestion_status_async( + self, + collection: str, + doc_id: str, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> None: + """Remove ingestion status record (async). + + Args: + collection: Collection name. + doc_id: Document ID. + user_id: Optional user ID for multi-tenancy. + is_admin: Whether user has admin privileges. + + Raises: + DatabaseOperationError: If delete operation fails. + """ + + +class PromptTemplateStore(ABC): + """Prompt template management contract. + + Manages prompt_templates table for storing and retrieving prompt templates. + + Phase 1A Option C: Hybrid sync/async methods for gradual migration. + """ + + @abstractmethod + def save_prompt_template( + self, + name: str, + template: str, + user_id: Optional[int] = None, + metadata: Optional[str] = None, + ) -> str: + """Save or update a prompt template. + + Args: + name: Template name (used for version grouping) + template: Template content + user_id: User ID for multi-tenancy + metadata: Optional metadata as JSON string + + Returns: + Template ID (UUID string) + """ + + @abstractmethod + def get_prompt_template( + self, + template_id: str, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Get a prompt template by ID. + + Args: + template_id: Template UUID + user_id: User ID for multi-tenancy + + Returns: + Template data dict or None if not found + """ + + @abstractmethod + def get_latest_prompt_template( + self, + name: str, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Get the latest version of a prompt template by name. + + Args: + name: Template name + user_id: User ID for multi-tenancy + + Returns: + Template data dict or None if not found + """ + + @abstractmethod + def list_prompt_templates( + self, + name_filter: Optional[str] = None, + latest_only: bool = False, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """List prompt templates with optional filtering. + + Args: + name_filter: Filter by template name (partial match) + latest_only: Only return latest versions + user_id: User ID for multi-tenancy + limit: Maximum results to return + + Returns: + List of template data dicts + """ + + @abstractmethod + def delete_prompt_template( + self, + template_id: str, + user_id: Optional[int] = None, + ) -> bool: + """Delete a prompt template by ID. + + Args: + template_id: Template UUID + user_id: User ID for multi-tenancy + + Returns: + True if deleted, False if not found + """ + + @abstractmethod + def update_metadata( + self, + template_id: str, + metadata: Optional[str], + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Update metadata only, keeping same version and ID. + + Args: + template_id: Template UUID + metadata: New metadata as JSON string + user_id: User ID for multi-tenancy + + Returns: + Updated template data dict or None if not found + """ + + @abstractmethod + def delete_by_name( + self, + name: str, + version: Optional[int] = None, + user_id: Optional[int] = None, + ) -> int: + """Delete template(s) by name. + + Handles is_latest flag updates for remaining versions. + + Args: + name: Template name + version: Specific version to delete (None = delete all versions) + user_id: User ID for multi-tenancy + + Returns: + Number of templates deleted + + Raises: + DocumentNotFoundError: If template not found + """ + + @abstractmethod + def get_versions_by_name( + self, + name: str, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """Get all versions of a template by name. + + Args: + name: Template name + user_id: User ID for multi-tenancy + limit: Maximum results to return + + Returns: + List of template data dicts + """ + + # --- Async methods (delegate to sync for Phase 1A) --- + + @abstractmethod + async def save_prompt_template_async( + self, + name: str, + template: str, + user_id: Optional[int] = None, + metadata: Optional[str] = None, + ) -> str: + """Async version of save_prompt_template.""" + + @abstractmethod + async def get_prompt_template_async( + self, + template_id: str, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of get_prompt_template.""" + + @abstractmethod + async def get_latest_prompt_template_async( + self, + name: str, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of get_latest_prompt_template.""" + + @abstractmethod + async def list_prompt_templates_async( + self, + name_filter: Optional[str] = None, + latest_only: bool = False, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """Async version of list_prompt_templates.""" + + @abstractmethod + async def delete_prompt_template_async( + self, + template_id: str, + user_id: Optional[int] = None, + ) -> bool: + """Async version of delete_prompt_template.""" + + @abstractmethod + async def update_metadata_async( + self, + template_id: str, + metadata: Optional[str], + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of update_metadata.""" + + @abstractmethod + async def delete_by_name_async( + self, + name: str, + version: Optional[int] = None, + user_id: Optional[int] = None, + ) -> int: + """Async version of delete_by_name.""" + + @abstractmethod + async def get_versions_by_name_async( + self, + name: str, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """Async version of get_versions_by_name.""" + + +class MainPointerStore(ABC): + """Main pointer management contract for version control. + + Manages main_pointers table for tracking current versions across + processing stages (parse, chunk, embed). + + Phase 1A Option C: Hybrid sync/async methods for gradual migration. + + NOTE: user_id parameter is included for API consistency but is not + currently stored in the main_pointers table schema. A schema migration + is required to add user_id support for multi-tenancy. + """ + + @abstractmethod + def set_main_pointer( + self, + collection: str, + doc_id: str, + step_type: str, + semantic_id: str, + technical_id: str, + model_tag: Optional[str] = None, + operator: Optional[str] = None, + user_id: Optional[int] = None, + ) -> None: + """Set or update a main pointer for a document. + + Args: + collection: Collection name + doc_id: Document ID + step_type: Processing stage (parse, chunk, embed) + semantic_id: Semantic identifier for the version (e.g., parse_id) + technical_id: Technical identifier/hash for the version + model_tag: Optional model tag for model-specific pointers + operator: Optional operator who made the change + user_id: Optional user ID (not stored, reserved for future use) + """ + + @abstractmethod + def get_main_pointer( + self, + collection: str, + doc_id: str, + step_type: str, + model_tag: Optional[str] = None, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Get a main pointer for a document. + + Args: + collection: Collection name + doc_id: Document ID + step_type: Processing stage (parse, chunk, embed) + model_tag: Optional model tag for model-specific pointers + user_id: Optional user ID (not used, reserved for future) + + Returns: + Pointer data dict with keys: collection, doc_id, step_type, + model_tag, semantic_id, technical_id, created_at, updated_at, + operator. Returns None if not found. + """ + + @abstractmethod + def list_main_pointers( + self, + collection: str, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """List main pointers for a collection. + + Args: + collection: Collection name + doc_id: Optional document ID filter + user_id: Optional user ID (not used, reserved for future) + limit: Maximum results to return + + Returns: + List of pointer data dicts + """ + + @abstractmethod + def delete_main_pointer( + self, + collection: str, + doc_id: str, + step_type: str, + model_tag: Optional[str] = None, + user_id: Optional[int] = None, + ) -> bool: + """Delete a main pointer. + + Args: + collection: Collection name + doc_id: Document ID + step_type: Processing stage (parse, chunk, embed) + model_tag: Optional model tag for model-specific pointers + user_id: Optional user ID (not used, reserved for future) + + Returns: + True if deleted, False if not found + """ + + # --- Async methods (delegate to sync for Phase 1A) --- + + @abstractmethod + async def set_main_pointer_async( + self, + collection: str, + doc_id: str, + step_type: str, + semantic_id: str, + technical_id: str, + model_tag: Optional[str] = None, + operator: Optional[str] = None, + user_id: Optional[int] = None, + ) -> None: + """Async version of set_main_pointer.""" + + @abstractmethod + async def get_main_pointer_async( + self, + collection: str, + doc_id: str, + step_type: str, + model_tag: Optional[str] = None, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of get_main_pointer.""" + + @abstractmethod + async def list_main_pointers_async( + self, + collection: str, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """Async version of list_main_pointers.""" + + @abstractmethod + async def delete_main_pointer_async( + self, + collection: str, + doc_id: str, + step_type: str, + model_tag: Optional[str] = None, + user_id: Optional[int] = None, + ) -> bool: + """Async version of delete_main_pointer.""" diff --git a/src/xagent/core/tools/core/RAG_tools/storage/factory.py b/src/xagent/core/tools/core/RAG_tools/storage/factory.py new file mode 100644 index 000000000..cc7cadbc1 --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/storage/factory.py @@ -0,0 +1,331 @@ +"""Unified factory for all KB storage contracts. + +Phase 1A Part 2: StorageFactory manages singleton instances of all stores +with lazy initialization and thread-safe access. + +Backward compatibility: Convenience functions (get_vector_index_store, etc.) +are provided for existing code. +""" + +from __future__ import annotations + +import threading +from typing import Any, Optional + +from .contracts import ( + IngestionStatusStore, + KBWriteCoordinator, + MainPointerStore, + MetadataStore, + PromptTemplateStore, + VectorIndexStore, +) +from .lancedb_stores import ( + LanceDBIngestionStatusStore, + LanceDBMainPointerStore, + LanceDBMetadataStore, + LanceDBPromptTemplateStore, + LanceDBVectorIndexStore, +) +from .vector_backend import ( + VectorBackend, + get_configured_vector_backend, + require_implemented_vector_backend, +) + + +class StorageFactory: + """Unified factory for all storage contracts. + + Manages singleton instances of all stores with lazy initialization + and thread-safe access using double-checked locking. + + Usage: + factory = StorageFactory.get_factory() + vector_store = factory.get_vector_index_store() + metadata_store = factory.get_metadata_store() + """ + + _instance: Optional[StorageFactory] = None + _lock = threading.RLock() # RLock for reentrant locking + + def __init__(self) -> None: + """Private constructor - use get_factory() instead.""" + if StorageFactory._instance is not None: + raise RuntimeError("Use get_factory() to get StorageFactory instance") + + # Store instances (lazy initialization) + self._vector_index_store: Optional[VectorIndexStore] = None + self._vector_backend: Optional[VectorBackend] = None + self._metadata_store: Optional[MetadataStore] = None + self._ingestion_status_store: Optional[IngestionStatusStore] = None + self._prompt_template_store: Optional[PromptTemplateStore] = None + self._main_pointer_store: Optional[MainPointerStore] = None + self._coordinator: Optional[KBWriteCoordinator] = None + + @classmethod + def get_factory(cls) -> StorageFactory: + """Get singleton factory instance. + + Uses double-checked locking for thread-safe lazy initialization. + + Returns: + The singleton StorageFactory instance. + """ + if cls._instance is None: + with cls._lock: + if cls._instance is None: + cls._instance = cls() + return cls._instance + + def reset_all(self) -> None: + """Reset all store instances. + + Useful for tests/fixtures that need isolated storage. + Thread-safe: uses factory lock to prevent race conditions. + """ + with self._lock: + self._vector_index_store = None + self._vector_backend = None + self._metadata_store = None + self._ingestion_status_store = None + self._prompt_template_store = None + self._main_pointer_store = None + self._coordinator = None + + # --- VectorIndexStore --- + + def get_vector_index_store(self) -> VectorIndexStore: + """Get or create vector index store. + + Backend is selected via :envvar:`XAGENT_VECTOR_BACKEND` (or legacy + ``VECTOR_STORE_BACKEND``); see :mod:`.vector_backend`. + + Returns: + Concrete :class:`~.contracts.VectorIndexStore` (currently + :class:`~.lancedb_stores.LanceDBVectorIndexStore` when backend is + ``lancedb``). + + Raises: + ConfigurationError: Unknown backend name, or backend not implemented + yet (e.g. ``milvus`` / ``qdrant`` without an adapter). + """ + if self._vector_index_store is None: + with self._lock: + if self._vector_index_store is None: + backend = get_configured_vector_backend() + require_implemented_vector_backend(backend) + if backend is VectorBackend.LANCEDB: + self._vector_index_store = LanceDBVectorIndexStore() + self._vector_backend = backend + else: + raise AssertionError( + "require_implemented_vector_backend must prevent this branch" + ) + return self._vector_index_store + + def get_resolved_vector_backend(self) -> VectorBackend: + """Return the backend bound to the current vector index store singleton. + + After the store is created, this reflects the backend used at creation + time (cached). Before creation, returns :func:`.get_configured_vector_backend` + without instantiating the store. + """ + if self._vector_backend is not None: + return self._vector_backend + return get_configured_vector_backend() + + # --- MetadataStore --- + + def get_metadata_store(self) -> MetadataStore: + """Get or create metadata store. + + Returns: + LanceDBMetadataStore instance. + """ + if self._metadata_store is None: + with self._lock: + if self._metadata_store is None: + self._metadata_store = LanceDBMetadataStore() + return self._metadata_store + + # --- IngestionStatusStore --- + + def get_ingestion_status_store(self) -> IngestionStatusStore: + """Get or create ingestion status store. + + Returns: + LanceDBIngestionStatusStore instance. + """ + if self._ingestion_status_store is None: + with self._lock: + if self._ingestion_status_store is None: + self._ingestion_status_store = LanceDBIngestionStatusStore() + return self._ingestion_status_store + + # --- PromptTemplateStore --- + + def get_prompt_template_store(self) -> PromptTemplateStore: + """Get or create prompt template store. + + Returns: + LanceDBPromptTemplateStore instance. + """ + if self._prompt_template_store is None: + with self._lock: + if self._prompt_template_store is None: + self._prompt_template_store = LanceDBPromptTemplateStore() + return self._prompt_template_store + + # --- MainPointerStore --- + + def get_main_pointer_store(self) -> MainPointerStore: + """Get or create main pointer store. + + Returns: + LanceDBMainPointerStore instance. + """ + if self._main_pointer_store is None: + with self._lock: + if self._main_pointer_store is None: + self._main_pointer_store = LanceDBMainPointerStore() + return self._main_pointer_store + + # --- KBWriteCoordinator --- + + def get_kb_write_coordinator(self) -> KBWriteCoordinator: + """Get or create KB write coordinator. + + Returns: + DefaultKBWriteCoordinator: Phase 1A shell delegating to metadata + and vector stores only; see that class for future coordination scope. + """ + if self._coordinator is None: + with self._lock: + if self._coordinator is None: + self._coordinator = DefaultKBWriteCoordinator( + metadata=self.get_metadata_store(), + vector_index=self.get_vector_index_store(), + ) + return self._coordinator + + +# ============================================================================ +# Backward Compatibility Functions +# ============================================================================ + +# Module-level lock for backward compatibility functions +_compat_lock = threading.Lock() +_default_factory: Optional[StorageFactory] = None + + +def _get_default_factory() -> StorageFactory: + """Get or create default factory instance (thread-safe).""" + global _default_factory + if _default_factory is None: + with _compat_lock: + if _default_factory is None: + _default_factory = StorageFactory.get_factory() + return _default_factory + + +def reset_kb_write_coordinator() -> None: + """Reset process-global coordinator (useful for tests/fixtures). + + Deprecated: Use StorageFactory.get_factory().reset_all() instead. + """ + _get_default_factory().reset_all() + + +def get_kb_write_coordinator() -> KBWriteCoordinator: + """Return process-global KB write coordinator. + + Deprecated: Use StorageFactory.get_factory().get_kb_write_coordinator() instead. + """ + return _get_default_factory().get_kb_write_coordinator() + + +def get_metadata_store() -> MetadataStore: + """Convenience accessor for metadata store. + + Deprecated: Use StorageFactory.get_factory().get_metadata_store() instead. + """ + return _get_default_factory().get_metadata_store() + + +def get_vector_index_store() -> VectorIndexStore: + """Convenience accessor for vector index store. + + Deprecated: Use StorageFactory.get_factory().get_vector_index_store() instead. + """ + return _get_default_factory().get_vector_index_store() + + +def get_vector_store_raw_connection() -> Any: + """Return the LanceDB handle exposed by the vector index store singleton. + + Central entry point for RAG code that still needs a raw connection during + Phase 1A. Replaces duplicated per-module ``get_connection_from_env`` helpers + that only delegated to ``get_vector_index_store().get_raw_connection()``. + + Returns: + The object returned by :meth:`VectorIndexStore.get_raw_connection`. + """ + return get_vector_index_store().get_raw_connection() + + +def get_ingestion_status_store() -> IngestionStatusStore: + """Get ingestion status store. + + Returns: + LanceDBIngestionStatusStore instance. + """ + return _get_default_factory().get_ingestion_status_store() + + +def get_prompt_template_store() -> PromptTemplateStore: + """Get prompt template store. + + Returns: + LanceDBPromptTemplateStore instance. + """ + return _get_default_factory().get_prompt_template_store() + + +def get_main_pointer_store() -> MainPointerStore: + """Get main pointer store. + + Returns: + LanceDBMainPointerStore instance. + """ + return _get_default_factory().get_main_pointer_store() + + +# ============================================================================ +# Default Coordinator Implementation +# ============================================================================ + + +class DefaultKBWriteCoordinator(KBWriteCoordinator): + """In-process KB write coordinator: Phase 1A placeholder implementation. + + Only :meth:`metadata_store` and :meth:`vector_index_store` are implemented; + both delegate to the injected or default LanceDB-backed stores. This is + sufficient as a shell while call sites converge on :class:`KBWriteCoordinator`. + Future phases may add distributed locking, batched writes, and conflict + resolution without changing the high-level factory entry point. + """ + + def __init__( + self, + metadata: MetadataStore | None = None, + vector_index: VectorIndexStore | None = None, + ) -> None: + self._metadata = metadata or LanceDBMetadataStore() + self._vector_index = vector_index or LanceDBVectorIndexStore() + + def metadata_store(self) -> MetadataStore: + return self._metadata + + def vector_index_store(self) -> VectorIndexStore: + return self._vector_index diff --git a/src/xagent/core/tools/core/RAG_tools/storage/lancedb_filter_utils.py b/src/xagent/core/tools/core/RAG_tools/storage/lancedb_filter_utils.py new file mode 100644 index 000000000..b692762c6 --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/storage/lancedb_filter_utils.py @@ -0,0 +1,87 @@ +"""LanceDB filter expression utilities. + +Shared functions for converting abstract filter expressions to LanceDB syntax. +""" + +from typing import Any + +from ..utils.string_utils import escape_lancedb_string +from .contracts import FilterCondition, FilterExpression, FilterOperator + + +def translate_condition(condition: FilterCondition) -> str: + """Translate single FilterCondition to LanceDB syntax. + + Args: + condition: FilterCondition to translate + + Returns: + LanceDB filter string + """ + field = condition.field + op = condition.operator + value = condition.value + + if op == FilterOperator.EQ: + return f"{field} == {format_value(value)}" + elif op == FilterOperator.NE: + return f"{field} != {format_value(value)}" + elif op == FilterOperator.GT: + return f"{field} > {format_value(value)}" + elif op == FilterOperator.GTE: + return f"{field} >= {format_value(value)}" + elif op == FilterOperator.LT: + return f"{field} < {format_value(value)}" + elif op == FilterOperator.LTE: + return f"{field} <= {format_value(value)}" + elif op == FilterOperator.IN: + values = ", ".join(format_value(v) for v in value) + return f"{field} IN ({values})" + elif op == FilterOperator.CONTAINS: + return f"{field} LIKE '%{escape_lancedb_string(value)}%'" + elif op == FilterOperator.IS_NULL: + return f"{field} IS NULL" + elif op == FilterOperator.IS_NOT_NULL: + return f"{field} IS NOT NULL" + else: + raise ValueError(f"Unsupported operator: {op}") + + +def format_value(value: Any) -> str: + """Format value for LanceDB. + + Args: + value: Value to format + + Returns: + Formatted value string + """ + if isinstance(value, bool): + return "TRUE" if value else "FALSE" + elif isinstance(value, (int, float)): + return str(value) + elif value is None: + return "NULL" + else: + return f"'{escape_lancedb_string(value)}'" + + +def translate_filter_expression(expr: FilterExpression) -> str: + """Translate FilterExpression to LanceDB syntax. + + Args: + expr: FilterExpression (FilterCondition, tuple for AND, list for OR) + + Returns: + LanceDB filter string + """ + if isinstance(expr, FilterCondition): + return translate_condition(expr) + elif isinstance(expr, tuple): + # AND combination + return " AND ".join(f"({translate_filter_expression(e)})" for e in expr) + elif isinstance(expr, list): + # OR combination + return " OR ".join(f"({translate_filter_expression(e)})" for e in expr) + else: + raise ValueError(f"Unsupported filter expression: {type(expr)}") diff --git a/src/xagent/core/tools/core/RAG_tools/storage/lancedb_stores.py b/src/xagent/core/tools/core/RAG_tools/storage/lancedb_stores.py new file mode 100644 index 000000000..cc4b79327 --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/storage/lancedb_stores.py @@ -0,0 +1,2680 @@ +"""LanceDB-backed implementations of storage contracts.""" + +from __future__ import annotations + +import asyncio +import logging +from collections import defaultdict +from datetime import datetime, timezone +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, cast + +import lancedb +import pyarrow as pa # type: ignore +from lancedb.db import DBConnection + +from xagent.providers.vector_store.lancedb import get_connection_from_env + +from ..core.config import DEFAULT_VECTOR_STORE_SCAN_LIMIT, IndexPolicy +from ..core.schemas import CollectionInfo, IndexResult +from ..LanceDB.schema_manager import ensure_documents_table +from ..utils.lancedb_query_utils import query_to_list +from ..utils.string_utils import build_lancedb_filter_expression, escape_lancedb_string +from ..utils.user_permissions import UserPermissions +from .contracts import ( + DocumentRecord, + FilterCondition, + FilterExpression, + FilterOperator, + IngestionStatusStore, + MainPointerStore, + MetadataStore, + PromptTemplateStore, + VectorIndexStore, + build_filter_from_dict, +) +from .lancedb_filter_utils import ( + translate_filter_expression, +) +from .logging_utils import log_audit, log_performance + +logger = logging.getLogger(__name__) + + +class LanceDBMetadataStore(MetadataStore): + """LanceDB implementation for control-plane metadata operations.""" + + def __init__(self) -> None: + self._conn: Optional[DBConnection] = None + + async def _get_connection(self) -> DBConnection: + if self._conn is None: + self._conn = get_connection_from_env() + return self._conn + + async def get_collection(self, collection_name: str) -> CollectionInfo: + conn = await self._get_connection() + table = conn.open_table("collection_metadata") + safe_name = escape_lancedb_string(collection_name) + result = table.search().where(f"name = '{safe_name}'").to_arrow() + if len(result) == 0: + raise ValueError(f"Collection '{collection_name}' not found") + # Convert Arrow table to list of dicts and take first row + data = result.to_pylist()[0] + return CollectionInfo.from_storage(data) + + async def save_collection(self, collection: CollectionInfo) -> None: + conn = await self._get_connection() + await self.ensure_collection_metadata_table() + + data = collection.to_storage() + data["updated_at"] = datetime.now(timezone.utc).replace(tzinfo=None) + + table = conn.open_table("collection_metadata") + safe_name = escape_lancedb_string(collection.name) + existing = table.search().where(f"name = '{safe_name}'").to_arrow() + if len(existing) > 0: + table.delete(f"name = '{safe_name}'") + table.add([data]) + + async def ensure_collection_metadata_table(self) -> None: + conn = await self._get_connection() + schema = pa.schema( + [ + ("name", pa.string()), + ("schema_version", pa.string()), + ("embedding_model_id", pa.string()), + ("embedding_dimension", pa.int32()), + ("documents", pa.int32()), + ("processed_documents", pa.int32()), + ("parses", pa.int32()), + ("chunks", pa.int32()), + ("embeddings", pa.int32()), + ("document_names", pa.string()), + ("collection_locked", pa.bool_()), + ("allow_mixed_parse_methods", pa.bool_()), + ("skip_config_validation", pa.bool_()), + ("ingestion_config", pa.string()), + ("created_at", pa.timestamp("us")), + ("updated_at", pa.timestamp("us")), + ("last_accessed_at", pa.timestamp("us")), + ("extra_metadata", pa.string()), + ] + ) + table_names_fn = getattr(conn, "table_names", None) + table_exists = False + if table_names_fn: + try: + table_exists = "collection_metadata" in table_names_fn() + except Exception as exc: # noqa: BLE001 + logger.debug("collection_metadata existence check failed: %s", exc) + if not table_exists: + try: + conn.create_table("collection_metadata", schema=schema) + except Exception as exc: # noqa: BLE001 + logger.debug("collection_metadata create_table no-op/failure: %s", exc) + + async def save_collection_config( + self, + collection: str, + config_json: str, + user_id: int, + ) -> None: + """Save collection ingestion configuration to LanceDB.""" + from ..LanceDB.schema_manager import ensure_collection_config_table + + conn = await self._get_connection() + ensure_collection_config_table(conn) + + table = conn.open_table("collection_config") + safe_collection = escape_lancedb_string(collection) + + # Delete existing config for this collection and user + try: + table.delete(f"collection = '{safe_collection}' AND user_id = {user_id}") + except Exception as exc: + logger.debug("Error deleting old config: %s", exc) + + # Insert new config + now = datetime.now(timezone.utc).replace(tzinfo=None) + data = [ + { + "collection": collection, + "config_json": config_json, + "updated_at": now, + "user_id": user_id, + } + ] + table.add(data) + + async def get_collection_config( + self, + collection: str, + user_id: Optional[int], + is_admin: bool = False, + ) -> str | None: + """Get collection ingestion configuration from LanceDB. + + When ``is_admin`` is True, returns the most recently updated config for + the collection across all users (tenant-agnostic listing). + + Args: + collection: Collection name. + user_id: User ID for multi-tenancy. None is treated as 0 for non-admin, + and as "load all configs" for admin mode (ignored when ``is_admin``). + is_admin: If True, omit ``user_id`` filter and resolve duplicates by + latest ``updated_at``. + + Returns: + Config JSON string if found, None otherwise. + """ + from ..LanceDB.schema_manager import ensure_collection_config_table + + try: + conn = await self._get_connection() + ensure_collection_config_table(conn) + + table = conn.open_table("collection_config") + safe_collection = escape_lancedb_string(collection) + if is_admin: + where_clause = f"collection = '{safe_collection}'" + elif user_id is None: + # Non-admin with user_id=None: treat as user_id=0 for backward compatibility + where_clause = f"collection = '{safe_collection}' AND user_id = 0" + else: + where_clause = ( + f"collection = '{safe_collection}' AND user_id = {user_id}" + ) + result = table.search().where(where_clause).to_arrow() + + if len(result) == 0: + return None + if not is_admin or len(result) == 1: + return str(result["config_json"][0].as_py()) + + best_idx = 0 + for i in range(1, len(result)): + cur = result["updated_at"][i].as_py() + best = result["updated_at"][best_idx].as_py() + if cur is not None and (best is None or cur > best): + best_idx = i + return str(result["config_json"][best_idx].as_py()) + except Exception as exc: + logger.debug("Error reading collection config: %s", exc) + return None + + def get_raw_connection(self) -> DBConnection: + """Get the underlying LanceDB connection. + + This method provides access to the raw connection for operations that + cannot be performed through the storage abstraction. It initializes + and caches the connection for consistency with async methods. + + Returns: + DBConnection: The LanceDB connection object + """ + if self._conn is None: + self._conn = get_connection_from_env() + return self._conn + + +class LanceDBVectorIndexStore(VectorIndexStore): + """LanceDB implementation for vector/data-plane operations. + + Phase 1A Option C: Provides both sync and async methods. + Sync methods use legacy lancedb.connect(); async methods use lancedb.connect_async(). + Both sync and async methods return native Arrow format for efficient zero-copy operations. + """ + + def __init__(self) -> None: + self._conn: Optional[DBConnection] = None + self._async_conn: Optional[Any] = None # AsyncConnection + self._async_lock = asyncio.Lock() # Protect async connection initialization + + def _get_connection(self) -> DBConnection: + if self._conn is None: + self._conn = get_connection_from_env() + return self._conn + + async def _get_async_connection(self) -> Any: + """Get or create async LanceDB connection with thread-safe initialization.""" + # Fast path: return existing connection without lock + if self._async_conn is not None: + return self._async_conn + + # Slow path: initialize with lock to prevent race condition + async with self._async_lock: + # Double-check after acquiring lock + if self._async_conn is not None: + return self._async_conn + + # Get URI from sync connection for reuse + sync_conn = self._get_connection() + uri = getattr(sync_conn, "uri", None) + if uri is None: + # Fallback: use LANCEDB_DIR env var + import os + + uri = os.getenv("LANCEDB_DIR", "./data/lancedb") + self._async_conn = await lancedb.connect_async(uri) # type: ignore[attr-defined] + return self._async_conn + + def list_document_records( + self, + collection_name: Optional[str], + user_id: Optional[int], + is_admin: bool, + max_results: int = DEFAULT_VECTOR_STORE_SCAN_LIMIT, + ) -> List[DocumentRecord]: + # Audit log for data access + log_audit( + "data_access", + action="list_documents", + user_id=user_id or -1, + is_admin=is_admin, + collection=collection_name, + max_results=max_results, + ) + + # Build filter expression using common function (includes validation) + filters = {} + if collection_name is not None: + filters["collection"] = collection_name + + filter_expr_obj = build_filter_from_dict(filters) + combined_filter = self.build_filter_expression( + filters=filter_expr_obj, + user_id=user_id, + is_admin=is_admin, + ) + + conn = self._get_connection() + ensure_documents_table(conn) + table = conn.open_table("documents") + + raw_records = query_to_list( + table.search().where(combined_filter).limit(max_results) + if combined_filter + else table.search().limit(max_results) + ) + + records: List[DocumentRecord] = [] + for item in raw_records: + raw_doc_id = item.get("doc_id") + if not raw_doc_id: + continue + records.append( + DocumentRecord( + doc_id=str(raw_doc_id), + file_id=str(item["file_id"]) if item.get("file_id") else None, + source_path=( + str(item["source_path"]) if item.get("source_path") else None + ), + ) + ) + return records + + def rename_collection_data( + self, + collection_name: str, + new_name: str, + ) -> List[str]: + warnings: List[str] = [] + safe_old_name = escape_lancedb_string(collection_name) + conn = self._get_connection() + for table_name in self.list_table_names(): + if table_name not in { + "documents", + "parses", + "chunks", + } and not table_name.startswith("embeddings_"): + continue + try: + table = conn.open_table(table_name) + table.update( + f"collection = '{safe_old_name}'", + {"collection": new_name}, + ) + except Exception as exc: # noqa: BLE001 + message = f"Failed to update '{table_name}': {exc}" + logger.warning(message) + warnings.append(message) + return warnings + + def list_table_names(self) -> Sequence[str]: + conn = self._get_connection() + table_names_fn = getattr(conn, "table_names", None) + if table_names_fn is None: + return [] + try: + return [str(name) for name in table_names_fn()] + except Exception as exc: # noqa: BLE001 + logger.warning("Failed to list LanceDB tables: %s", exc) + return [] + + def get_vector_dimension(self, table_name: str) -> Optional[int]: + """Get the vector dimension from a table's schema.""" + conn = self._get_connection() + try: + table = conn.open_table(table_name) + schema = table.schema + vector_field = schema.field("vector") + if hasattr(vector_field, "type"): + vector_type = vector_field.type + if hasattr(vector_type, "list_size"): + return cast(int, vector_type.list_size) + except Exception as exc: # noqa: BLE001 + logger.debug("Failed to get vector dimension for %s: %s", table_name, exc) + return None + + def open_embeddings_table(self, model_tag: str) -> Tuple[Any, str]: + """Open embeddings table with legacy fallback support. + + Tries the primary Hub ID-based table name first, then falls back + to legacy provider-based naming if the primary doesn't exist. + + Args: + model_tag: Model tag for the embeddings table. + + Returns: + Tuple of (table_object, actual_table_name_used). + + Raises: + DatabaseOperationError: If neither primary nor legacy table exists. + """ + from ..core.exceptions import DatabaseOperationError + from ..LanceDB.model_tag_utils import to_model_tag + from ..utils.model_resolver import resolve_embedding_adapter + + conn = self._get_connection() + primary_table_name = f"embeddings_{to_model_tag(model_tag)}" + + # Try primary table first + try: + table = conn.open_table(primary_table_name) + return table, primary_table_name + except Exception as primary_exc: + last_error = primary_exc + + # Try legacy fallback + legacy_table_name: Optional[str] = None + try: + cfg, _ = resolve_embedding_adapter(model_tag) + legacy_table_name = f"embeddings_{to_model_tag(cfg.model_name)}" + except Exception: + legacy_table_name = None + + if legacy_table_name and legacy_table_name != primary_table_name: + try: + table = conn.open_table(legacy_table_name) + logger.info( + "Using legacy embeddings table '%s' for model_tag='%s'. " + "Consider migrating to '%s' for consistency.", + legacy_table_name, + model_tag, + primary_table_name, + ) + return table, legacy_table_name + except Exception as legacy_exc: + last_error = legacy_exc + + # Neither table exists + error_msg = f"Embeddings table not found for model_tag='{model_tag}'" + if primary_table_name: + error_msg += f" (tried: '{primary_table_name}'" + if legacy_table_name: + error_msg += f", '{legacy_table_name}'" + error_msg += ")" + raise DatabaseOperationError(error_msg) from last_error + + def delete_collection_data( + self, + collection_name: str, + ) -> Dict[str, int]: + """Delete all data for a collection from vector-side tables.""" + from ..LanceDB.schema_manager import ( + ensure_chunks_table, + ensure_documents_table, + ensure_parses_table, + ) + + deleted_counts: Dict[str, int] = {} + conn = self._get_connection() + safe_collection = escape_lancedb_string(collection_name) + + # Ensure tables exist before attempting deletion + ensure_documents_table(conn) + ensure_parses_table(conn) + ensure_chunks_table(conn) + + # Delete from core tables + for table_name in ["documents", "parses", "chunks"]: + try: + table = conn.open_table(table_name) + original_count = table.count_rows() + table.delete(f"collection = '{safe_collection}'") + deleted_count = original_count - table.count_rows() + if deleted_count > 0: + deleted_counts[table_name] = deleted_count + except Exception as exc: # noqa: BLE001 + logger.warning("Failed to delete from '%s': %s", table_name, exc) + + # Delete embeddings data + for table_name in self.list_table_names(): + if not table_name.startswith("embeddings_"): + continue + try: + table = conn.open_table(table_name) + original_count = table.count_rows() + table.delete(f"collection = '{safe_collection}'") + deleted_count = original_count - table.count_rows() + if deleted_count > 0: + deleted_counts[table_name] = deleted_count + except Exception as exc: # noqa: BLE001 + logger.warning("Failed to delete from '%s': %s", table_name, exc) + + return deleted_counts + + def aggregate_collection_stats( + self, + user_id: Optional[int], + is_admin: bool, + ) -> Dict[str, Dict[str, int]]: + """Aggregate statistics for all collections using memory-efficient batched iteration.""" + from ..LanceDB.schema_manager import ( + ensure_chunks_table, + ensure_documents_table, + ensure_parses_table, + ) + + stats: Dict[str, Dict[str, int]] = {} + conn = self._get_connection() + + # Ensure tables exist + ensure_documents_table(conn) + ensure_parses_table(conn) + ensure_chunks_table(conn) + + def _count_table(table_name: str, stat_key: str) -> None: + """Count records per collection using batched streaming to avoid OOM.""" + try: + # Use iter_batches for memory-efficient streaming with default batch_size=1000 + for batch in self.iter_batches( + table_name=table_name, + columns=["collection"], # Only need collection column + user_id=user_id, + is_admin=is_admin, + ): + # Extract collection column from PyArrow RecordBatch + collection_idx = batch.schema.get_field_index("collection") + if collection_idx == -1: + continue + + collection_array = batch.column(collection_idx) + for i in range(batch.num_rows): + collection = str(collection_array[i].as_py()) + if collection: + if collection not in stats: + stats[collection] = { + "documents": 0, + "parses": 0, + "chunks": 0, + "embeddings": 0, + } + stats[collection][stat_key] += 1 + except Exception as exc: # noqa: BLE001 + logger.debug("Failed to count table '%s': %s", table_name, exc) + + # Count documents, parses, and chunks + _count_table("documents", "documents") + _count_table("parses", "parses") + _count_table("chunks", "chunks") + + # Count embeddings from all embeddings_* tables + for table_name in self.list_table_names(): + if not table_name.startswith("embeddings_"): + continue + _count_table(table_name, "embeddings") + + return stats + + def aggregate_document_stats( + self, + collection_name: str, + doc_id: str, + user_id: Optional[int], + is_admin: bool, + ) -> Dict[str, int]: + """Aggregate statistics for a single document.""" + from ..LanceDB.schema_manager import ( + ensure_chunks_table, + ensure_documents_table, + ensure_parses_table, + ) + + stats = {"documents": 0, "parses": 0, "chunks": 0, "embeddings": 0} + conn = self._get_connection() + + # Ensure tables exist + ensure_documents_table(conn) + ensure_parses_table(conn) + ensure_chunks_table(conn) + + safe_collection = escape_lancedb_string(collection_name) + safe_doc_id = escape_lancedb_string(doc_id) + + base_filter = f"collection = '{safe_collection}' AND doc_id = '{safe_doc_id}'" + + def _count_table(table_name: str) -> int: + try: + table = conn.open_table(table_name) + return int(table.count_rows(base_filter)) + except Exception: # noqa: BLE001 + return 0 + + stats["documents"] = _count_table("documents") + stats["parses"] = _count_table("parses") + stats["chunks"] = _count_table("chunks") + + # Count embeddings across all embeddings tables + for table_name in self.list_table_names(): + if not table_name.startswith("embeddings_"): + continue + stats["embeddings"] += _count_table(table_name) + + return stats + + def create_index(self, model_tag: str, readonly: bool = False) -> IndexResult: + """Create or check vector index for embeddings table. + + This method implements full index management logic including automatic + index type selection based on row count and FTS index management. + + Args: + model_tag: Model tag for the embeddings table. + readonly: If True, don't trigger index creation. + + Returns: + IndexResult containing status, advice, and FTS enabled state. + """ + from ..core.config import IndexPolicy + from ..core.schemas import IndexResult + from ..LanceDB.model_tag_utils import to_model_tag + + # Import LanceDB index types + try: + from lancedb.index import IVF_HNSW_SQ, IVF_PQ # type: ignore + except ImportError: + IVF_HNSW_SQ = "IVF_HNSW_SQ" + IVF_PQ = "IVF_PQ" + + conn = self._get_connection() + table_name = f"embeddings_{to_model_tag(model_tag)}" + + if readonly: + # In readonly mode, check if FTS index exists without creating any indexes + fts_enabled = False + try: + table = conn.open_table(table_name) + indexes = table.list_indices() + fts_enabled = any( + idx.index_type == "FTS" and "text" in idx.columns for idx in indexes + ) + except Exception as e: + logger.debug("Unable to check FTS index status in readonly mode: %s", e) + + return IndexResult( + status="readonly", + advice=f"Readonly mode - no index operations for {table_name}", + fts_enabled=fts_enabled, + ) + + try: + table = conn.open_table(table_name) + except Exception as exc: + logger.debug("Unable to open table '%s': %s", table_name, exc) + return IndexResult(status="failed", advice=None, fts_enabled=False) + + # Use default index policy + policy = IndexPolicy() + vector_index_status: str = "no_index" + vector_index_advice: Optional[str] = None + + try: + # Get row count efficiently + row_count = table.count_rows() + + if row_count < policy.enable_threshold_rows: + vector_index_status = "below_threshold" + vector_index_advice = ( + f"Table {table_name} has {row_count} rows - below threshold " + f"({policy.enable_threshold_rows}) for index creation" + ) + else: + # Auto-select index type based on scale + from ..core.schemas import IndexType + + if row_count >= policy.ivfpq_threshold_rows: + recommended_type = IndexType.IVFPQ + else: + recommended_type = IndexType.HNSW + + # Check existing indexes + indexes = table.list_indices() + has_vector_index = any(idx.name == "vector" for idx in indexes) + + if not has_vector_index: + # Create index with recommended type + if recommended_type == IndexType.IVFPQ: + index_type = IVF_PQ + create_params = policy.ivfpq_params or {} + else: # HNSW + index_type = IVF_HNSW_SQ + create_params = policy.hnsw_params or {} + + # Merge metric with create_params + all_params = { + "metric": policy.metric.value, + "index_type": index_type, + **create_params, + } + + table.create_index(**all_params) + vector_index_status = "index_building" + logger.info( + "Successfully created vector index for %s (type=%s, metric=%s)", + table_name, + index_type, + policy.metric.value, + ) + if recommended_type == IndexType.IVFPQ: + vector_index_advice = ( + f"IVFPQ index created for {table_name} " + f"({row_count} rows, using IVFPQ strategy for large-scale data), metric: {policy.metric.value}" + ) + else: # HNSW + vector_index_advice = ( + f"HNSW index created for {table_name} " + f"({row_count} rows, using HNSW strategy for medium-scale data), metric: {policy.metric.value}" + ) + else: + vector_index_status = "index_ready" + vector_index_advice = f"Index ready for {table_name} ({row_count} rows), metric: {policy.metric.value}" + + except Exception as e: + logger.error(f"Vector index operation failed for {table_name}: {str(e)}") + vector_index_status = "index_corrupted" + vector_index_advice = ( + f"Vector index check failed for {table_name}: {str(e)}" + ) + + # Check actual FTS index status (not just whether we tried to create it) + fts_enabled = False + try: + indexes = table.list_indices() + fts_enabled = any( + idx.index_type == "FTS" and "text" in idx.columns for idx in indexes + ) + except Exception as e: + logger.warning(f"Failed to check FTS index status: {e}") + + # FTS Index Management (if enabled) + if policy.fts_enabled and not fts_enabled: + try: + fts_params = {"with_position": True, **(policy.fts_params or {})} + table.create_fts_index("text", replace=True, **fts_params) + logger.info("Created FTS index on 'text' column for %s", table_name) + # Re-check FTS status after creation + try: + indexes = table.list_indices() + fts_enabled = any( + idx.index_type == "FTS" and "text" in idx.columns + for idx in indexes + ) + except Exception: + pass + except Exception as e: + logger.warning( + f"FTS index creation/check failed for {table_name}: {str(e)}" + ) + + return IndexResult( + status=vector_index_status, + advice=vector_index_advice, + fts_enabled=fts_enabled, + ) + + # --- Index Management (Phase 1A Part 2) --- + + def should_reindex( + self, table_name: str, total_upserted: int, policy: IndexPolicy + ) -> bool: + """Determine if reindex should be triggered (sync).""" + try: + conn = self._get_connection() + table = conn.open_table(table_name) + + # Immediate reindex if enabled + if policy.enable_immediate_reindex and total_upserted > 0: + return True + + # Batch size threshold + if total_upserted >= policy.reindex_batch_size: + return True + + # Smart reindex: check unindexed ratio + if policy.enable_smart_reindex: + try: + stats = table.index_stats("vector_idx") + if stats.num_indexed_rows > 0: + unindexed_ratio = ( + stats.num_unindexed_rows / stats.num_indexed_rows + ) + if unindexed_ratio > policy.reindex_unindexed_ratio_threshold: + return True + + # Absolute threshold for unindexed rows + if stats.num_unindexed_rows > 10000: + return True + except Exception as e: # noqa: BLE001 + logger.debug("Could not get index stats for %s: %s", table_name, e) + + return False + + except Exception as e: + logger.error(f"Failed to check reindex status for {table_name}: {e}") + return False + + def trigger_reindex(self, table_name: str) -> bool: + """Trigger reindex operation on the table (sync).""" + try: + logger.info("Triggering reindex for %s", table_name) + conn = self._get_connection() + table = conn.open_table(table_name) + table.optimize() + logger.info("Reindex completed for %s", table_name) + return True + except Exception as e: # noqa: BLE001 + logger.warning("Reindex failed for %s: %s", table_name, e) + return False + + async def should_reindex_async( + self, table_name: str, total_upserted: int, policy: IndexPolicy + ) -> bool: + """Async version of should_reindex. + + Note: Current implementation uses sync operations under the hood. + True async I/O will be added in Phase 1B with RDB backend. + """ + # Delegate to sync implementation for now + return self.should_reindex(table_name, total_upserted, policy) + + async def trigger_reindex_async(self, table_name: str) -> bool: + """Async version of trigger_reindex. + + Note: Current implementation uses sync operations under the hood. + True async I/O will be added in Phase 1B with RDB backend. + """ + # Delegate to sync implementation for now + return self.trigger_reindex(table_name) + + def migrate_embeddings_table( + self, + model_id: str, + batch_size: int = 1000, + ) -> dict[str, Any]: + """Migrate legacy embeddings table to Hub ID-based naming. + + This method copies data from a legacy table (embeddings_{model_name}) + to a new Hub ID-based table (embeddings_{hub_id}), rewriting the + per-row ``model`` field to the Hub model ID. + + Args: + model_id: Hub model ID to migrate. + batch_size: Number of rows to copy per batch. + + Returns: + Dictionary with migration results. + """ + from ..utils import migration_utils + + return migration_utils.migrate_embeddings_table( + model_id=model_id, + batch_size=batch_size, + conn=self._get_connection(), + ) + + def get_raw_connection(self) -> DBConnection: + return self._get_connection() + + def iter_batches( + self, + table_name: str, + columns: Optional[Sequence[str]] = None, + batch_size: int = 1000, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> Iterator[Any]: + """Iterate over table data in batches. + + Yields backend-specific batch objects (e.g., PyArrow RecordBatch). + """ + from ..LanceDB.schema_manager import ( + ensure_chunks_table, + ensure_documents_table, + ensure_parses_table, + ) + + conn = self._get_connection() + + # Ensure table exists based on name + if table_name == "documents": + ensure_documents_table(conn) + elif table_name == "parses": + ensure_parses_table(conn) + elif table_name == "chunks": + ensure_chunks_table(conn) + + try: + table = conn.open_table(table_name) + except Exception as exc: + logger.debug("Unable to open table '%s': %s", table_name, exc) + return + + # Build filter expression using common function (includes validation) + combined_filter = None + if filters: + filter_expr_obj = build_filter_from_dict(filters) + combined_filter = self.build_filter_expression( + filters=filter_expr_obj, + user_id=user_id, + is_admin=is_admin, + ) + else: + # Just apply user filter + combined_filter = UserPermissions.get_user_filter(user_id, is_admin) + + # Helper method to select columns from a batch + def _select_columns(batch: Any, cols: Optional[Sequence[str]]) -> Any: + if cols is None: + return batch + arrays = [] + names = [] + for col_name in cols: + idx = batch.schema.get_field_index(col_name) + if idx != -1: + arrays.append(batch.column(idx)) + names.append(col_name) + if not arrays: + return pa.RecordBatch.from_arrays([], []) + return pa.RecordBatch.from_arrays(arrays, names) + + # Preferred path: streaming batches directly from LanceDB + try: + if combined_filter: + for raw_batch in table.to_batches( + filter=combined_filter, batch_size=batch_size + ): + batch = raw_batch + if columns is not None: + batch = _select_columns(batch, columns) + if batch.num_rows > 0: + yield batch + else: + for raw_batch in table.to_batches(batch_size=batch_size): + batch = raw_batch + if columns is not None: + batch = _select_columns(batch, columns) + if batch.num_rows > 0: + yield batch + return + except Exception as exc: + logger.debug( + "Batch streaming unavailable for table '%s': %s", table_name, exc + ) + + # Arrow fallback: materialize table as Arrow then iterate + try: + # Note: LanceDB's to_arrow() doesn't accept filter parameter + # Use search().where().to_arrow() instead + if combined_filter: + arrow_table = table.search().where(combined_filter).to_arrow() + else: + arrow_table = table.to_arrow() + except Exception as exc: + logger.debug( + "Unable to read table '%s' via to_arrow(): %s", table_name, exc + ) + return + + if columns is not None: + try: + arrow_table = arrow_table.select(columns) + except Exception as exc: + logger.debug( + "Table '%s' missing expected columns %s: %s", + table_name, + columns, + exc, + ) + return + + for batch in arrow_table.to_batches(max_chunksize=batch_size): + if batch.num_rows > 0: + yield batch + + def count_rows( + self, + table_name: str, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> int: + """Count rows in a table with optional filters. + + Raises: + DatabaseOperationError: If table cannot be opened or count fails. + """ + from ..core.exceptions import DatabaseOperationError + + conn = self._get_connection() + + try: + table = conn.open_table(table_name) + except Exception as exc: + raise DatabaseOperationError( + f"Failed to open table '{table_name}': {exc}" + ) from exc + + # Build filter expression using common function (includes validation) + backend_filter = None + if filters: + filter_expr_obj = build_filter_from_dict(filters) + backend_filter = self.build_filter_expression( + filters=filter_expr_obj, + user_id=user_id, + is_admin=is_admin, + ) + else: + # Just apply user filter + backend_filter = UserPermissions.get_user_filter(user_id, is_admin) + + try: + if backend_filter: + return int(table.count_rows(backend_filter)) + return int(table.count_rows()) + except Exception as exc: + raise DatabaseOperationError( + f"Failed to count rows in table '{table_name}': {exc}" + ) from exc + + def aggregate_document_counts( + self, + table_name: str, + doc_id_column: str, + collection_name: str, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> Dict[str, int]: + """Aggregate records per document for a specific table.""" + counts: Dict[str, int] = defaultdict(int) + + for batch in self.iter_batches( + table_name=table_name, + columns=["collection", doc_id_column], + user_id=user_id, + is_admin=is_admin, + ): + collection_idx = batch.schema.get_field_index("collection") + doc_idx = batch.schema.get_field_index(doc_id_column) + + if collection_idx == -1 or doc_idx == -1: + continue + + collection_array = batch.column(collection_idx) + doc_array = batch.column(doc_idx) + + for idx in range(batch.num_rows): + collection_raw = collection_array[idx].as_py() + if not collection_raw or str(collection_raw) != collection_name: + continue + doc_raw = doc_array[idx].as_py() + if not doc_raw: + continue + counts[str(doc_raw)] += 1 + + return dict(counts) + + def build_filter_expression( + self, + filters: Optional[FilterExpression], + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> Optional[str]: + """Convert abstract filter expression to LanceDB SQL syntax.""" + if not filters: + # Still apply user filter for multi-tenancy + return UserPermissions.get_user_filter(user_id, is_admin) + + backend_filter = translate_filter_expression(filters) + + # Combine with user filter + user_filter = UserPermissions.get_user_filter(user_id, is_admin) + if user_filter: + return f"({backend_filter}) AND ({user_filter})" + return backend_filter + + def upsert_documents(self, records: List[Dict[str, Any]]) -> None: + """Upsert document records to LanceDB. + + Args: + records: List of document record dictionaries to upsert. + """ + from ..LanceDB.schema_manager import ensure_documents_table + + if not records: + return + + conn = self._get_connection() + ensure_documents_table(conn) + table = conn.open_table("documents") + + # Use merge_insert for efficient upsert + table.merge_insert( + ["collection", "doc_id"] + ).when_matched_update_all().when_not_matched_insert_all().execute(records) + + def upsert_parses(self, records: List[Dict[str, Any]]) -> None: + """Upsert parse records to LanceDB. + + Args: + records: List of parse record dictionaries to upsert. + """ + from ..LanceDB.schema_manager import ensure_parses_table + + if not records: + return + + conn = self._get_connection() + ensure_parses_table(conn) + table = conn.open_table("parses") + + # Use merge_insert for efficient upsert + table.merge_insert( + ["collection", "doc_id", "parse_hash"] + ).when_matched_update_all().when_not_matched_insert_all().execute(records) + + def upsert_chunks(self, records: List[Dict[str, Any]]) -> None: + """Upsert chunk records to LanceDB. + + Args: + records: List of chunk record dictionaries to upsert. + """ + from ..LanceDB.schema_manager import ensure_chunks_table + + if not records: + return + + conn = self._get_connection() + ensure_chunks_table(conn) + table = conn.open_table("chunks") + + # Use merge_insert for efficient upsert + table.merge_insert( + ["collection", "doc_id", "parse_hash", "chunk_id"] + ).when_matched_update_all().when_not_matched_insert_all().execute(records) + + def upsert_embeddings(self, model_tag: str, records: List[Dict[str, Any]]) -> None: + """Upsert embedding records to LanceDB with fallback pattern. + + Args: + model_tag: Model tag for the embeddings table. + records: List of embedding record dictionaries to upsert. + + Raises: + Exception: If both merge_insert and add() methods fail. + """ + from ..LanceDB.model_tag_utils import to_model_tag + from ..LanceDB.schema_manager import ensure_embeddings_table + from ..vector_storage.vector_manager import _is_non_recoverable_merge_error + + if not records: + return + + conn = self._get_connection() + table_name = f"embeddings_{to_model_tag(model_tag)}" + + # Infer vector dimension from first record + vector_dim = None + if records and "vector" in records[0]: + vector = records[0]["vector"] + if isinstance(vector, (list, tuple)): + vector_dim = len(vector) + + ensure_embeddings_table(conn, to_model_tag(model_tag), vector_dim=vector_dim) + table = conn.open_table(table_name) + + try: + # Try merge_insert first (preferred method for upserts) + table.merge_insert( + ["collection", "doc_id", "chunk_id"] + ).when_matched_update_all().when_not_matched_insert_all().execute(records) + except Exception as merge_error: + if _is_non_recoverable_merge_error(merge_error): + # Log critical error and re-raise without fallback + logger.error( + "merge_insert failed with non-recoverable error (error_type=%s): %s. " + "This may indicate schema mismatch or data corruption. " + "Not attempting fallback to add() method.", + type(merge_error).__name__, + merge_error, + ) + raise + + # For recoverable errors (e.g., temporary issues, network errors), attempt fallback + logger.warning( + "merge_insert failed (error_type=%s): %s; " + "attempting fallback to add() method", + type(merge_error).__name__, + merge_error, + ) + try: + # Use dict list directly (LanceDB add() accepts list-of-dict) + table.add(records) + logger.info( + "Successfully used add() fallback for %d embeddings after merge_insert failure", + len(records), + ) + except Exception as add_error: + logger.error( + "Fallback add() also failed: %s. " + "Both merge_insert and add() methods failed.", + add_error, + ) + raise + + # --- Sync search methods (Phase 1A Option C) --- + + def search_vectors( + self, + table_name: str, + query_vector: List[float], + *, + top_k: int, + filters: Optional[FilterExpression] = None, + vector_column_name: str = "vector", + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> List[Dict[str, Any]]: + """Execute vector search using sync LanceDB API. + + Returns native Arrow format converted to list of dicts. + """ + # Log search parameters for performance tracking + log_performance( + "search_vectors_start", + top_k=top_k, + vector_dim=len(query_vector), + table_name=table_name, + has_filters=filters is not None, + ) + + conn = self._get_connection() + + # Open table (no legacy fallback at abstraction layer - handled by caller) + try: + table = conn.open_table(table_name) + except Exception as exc: + logger.debug("Unable to open table '%s': %s", table_name, exc) + return [] + + # Build filter expression + backend_filter = self.build_filter_expression( + filters, user_id=user_id, is_admin=is_admin + ) + + # Build search query + search_query = table.search( + query_vector, + vector_column_name=vector_column_name, + ) + + if backend_filter: + search_query = search_query.where(backend_filter) + + search_query = search_query.limit(top_k) + + try: + # Use query_to_list for three-tier fallback (to_arrow, to_list, to_pandas) + raw_results = query_to_list(search_query) + + # Log performance metric + log_performance( + "search_vectors_complete", + result_count=len(raw_results), + table_name=table_name, + ) + return raw_results + + except Exception as exc: + logger.error("Sync vector search failed: %s", exc) + return [] + + # --- Async method implementations (Phase 1A Option C) --- + + async def search_vectors_async( + self, + table_name: str, + query_vector: List[float], + *, + top_k: int, + filters: Optional[FilterExpression] = None, + vector_column_name: str = "vector", + ) -> List[Dict[str, Any]]: + """Execute vector search using async LanceDB API. + + Returns native Arrow format converted to list of dicts. + """ + # Log search parameters for performance tracking + log_performance( + "search_vectors_start", + top_k=top_k, + vector_dim=len(query_vector), + table_name=table_name, + has_filters=filters is not None, + ) + + async_conn = await self._get_async_connection() + + try: + table = await async_conn.open_table(table_name) + except Exception as exc: + logger.debug("Unable to open table '%s': %s", table_name, exc) + return [] + + # Build filter expression + backend_filter = self.build_filter_expression( + filters, user_id=None, is_admin=False + ) + + # Build search query + search_query = table.search( + query_vector, + vector_column_name=vector_column_name, + ) + + if backend_filter: + search_query = search_query.where(backend_filter) + + search_query = search_query.limit(top_k) + + try: + # Async search returns Arrow table + results_table = await search_query.to_arrow() + + # Convert Arrow to list of dicts + results = [] + for batch in results_table.to_batches(): + for i in range(batch.num_rows): + row = {} + for j in range(batch.num_columns): + col_name = batch.schema.names[j] + col_array = batch.column(j) + value = col_array[i].as_py() + row[col_name] = value + results.append(row) + + # Log performance metric + log_performance( + "search_vectors_complete", + result_count=len(results), + table_name=table_name, + ) + return results + + except Exception as exc: + logger.error("Async vector search failed: %s", exc) + return [] + + async def search_fts_async( + self, + table_name: str, + query_text: str, + *, + top_k: int, + filters: Optional[FilterExpression] = None, + text_column_name: str = "text", + ) -> List[Dict[str, Any]]: + """Execute full-text search using async LanceDB FTS API. + + Returns native Arrow format converted to list of dicts. + """ + async_conn = await self._get_async_connection() + + try: + table = await async_conn.open_table(table_name) + except Exception as exc: + logger.debug("Unable to open table '%s': %s", table_name, exc) + return [] + + # Build filter expression + backend_filter = self.build_filter_expression( + filters, user_id=None, is_admin=False + ) + + # Build FTS search query + # Note: LanceDB async API supports query_type="fts" + search_query = table.search( + query_text, + query_type="fts", + ) + + if backend_filter: + search_query = search_query.where(backend_filter) + + search_query = search_query.limit(top_k) + + try: + # Async FTS search returns Arrow table + results_table = await search_query.to_arrow() + + # Convert Arrow to list of dicts + results = [] + for batch in results_table.to_batches(): + for i in range(batch.num_rows): + row = {} + for j in range(batch.num_columns): + col_name = batch.schema.names[j] + col_array = batch.column(j) + value = col_array[i].as_py() + row[col_name] = value + results.append(row) + return results + + except Exception as exc: + logger.error("Async FTS search failed: %s", exc) + return [] + + async def iter_batches_async( + self, + table_name: str, + columns: Optional[Sequence[str]] = None, + batch_size: int = 1000, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> Any: # # Returns AsyncIterator (async generator), see contract for details + """Iterate over table data in batches using async LanceDB API. + + Yields PyArrow RecordBatch objects (native async format). + """ + # Log batch iteration parameters for performance tracking + log_performance( + "iter_batches_start", + table_name=table_name, + batch_size=batch_size, + columns_provided=columns is not None, + has_filters=filters is not None, + ) + + async_conn = await self._get_async_connection() + + try: + table = await async_conn.open_table(table_name) + except Exception as exc: + logger.debug("Unable to open table '%s': %s", table_name, exc) + return + + # Build filter expression using common function (includes validation) + combined_filter = None + if filters: + filter_expr_obj = build_filter_from_dict(filters) + combined_filter = self.build_filter_expression( + filters=filter_expr_obj, + user_id=user_id, + is_admin=is_admin, + ) + else: + # Just apply user filter + combined_filter = UserPermissions.get_user_filter(user_id, is_admin) + + # Helper method to select columns from a batch + def _select_columns(batch: Any, cols: Optional[Sequence[str]]) -> Any: + if cols is None: + return batch + arrays = [] + names = [] + for col_name in cols: + idx = batch.schema.get_field_index(col_name) + if idx != -1: + arrays.append(batch.column(idx)) + names.append(col_name) + if not arrays: + return pa.RecordBatch.from_arrays([], []) + return pa.RecordBatch.from_arrays(arrays, names) + + try: + # Use LanceDB async to_batches() with column projection for efficiency + # Note: LanceDB to_batches supports columns parameter to avoid reading unused columns + if combined_filter: + async for batch in table.to_batches( + filter=combined_filter, + batch_size=batch_size, + columns=columns, # Pass columns directly to avoid reading all data + ): + if batch.num_rows > 0: + yield batch + else: + async for batch in table.to_batches( + batch_size=batch_size, + columns=columns, # Pass columns directly to avoid reading all data + ): + if batch.num_rows > 0: + yield batch + except Exception as exc: + logger.debug( + "Async batch iteration failed for table '%s': %s", table_name, exc + ) + + async def count_rows_async( + self, + table_name: str, + filters: Optional[Dict[str, Any]] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> int: + """Count rows in a table with optional filters using async LanceDB API.""" + async_conn = await self._get_async_connection() + + try: + table = await async_conn.open_table(table_name) + except Exception as exc: + logger.debug("Unable to open table '%s': %s", table_name, exc) + return 0 + + # Build filter expression using common function (includes validation) + combined_filter = None + if filters: + filter_expr_obj = build_filter_from_dict(filters) + combined_filter = self.build_filter_expression( + filters=filter_expr_obj, + user_id=user_id, + is_admin=is_admin, + ) + else: + # Just apply user filter + combined_filter = UserPermissions.get_user_filter(user_id, is_admin) + + try: + if combined_filter: + count = int(await table.count_rows(combined_filter)) + else: + count = int(await table.count_rows()) + + # Log performance metric + log_performance( + "count_rows_complete", + table_name=table_name, + row_count=count, + has_filter=combined_filter is not None, + ) + return count + except Exception as exc: + logger.debug("Failed to count rows in '%s': %s", table_name, exc) + return 0 + + async def get_vector_dimension_async(self, table_name: str) -> Optional[int]: + """Get the vector dimension from a table's schema (async). + + Note: LanceDB schema operations are sync-only, so this wraps the sync + implementation. True async I/O will be added in Phase 1B with RDB backend. + """ + # LanceDB schema operations don't have async variants, use sync + return self.get_vector_dimension(table_name) + + async def upsert_documents_async(self, records: List[Dict[str, Any]]) -> None: + """Upsert document records using async LanceDB API.""" + from ..LanceDB.schema_manager import ensure_documents_table + + if not records: + return + + # Log upsert operation parameters for performance tracking + log_performance( + "upsert_documents_start", record_count=len(records), table="documents" + ) + + async_conn = await self._get_async_connection() + + # Note: ensure_documents_table uses sync connection - may need async variant + # For now, reuse sync connection for table creation + sync_conn = self._get_connection() + ensure_documents_table(sync_conn) + + table = await async_conn.open_table("documents") + + # Use merge_insert for efficient upsert + await ( + table.merge_insert(["collection", "doc_id"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(records) + ) + + async def upsert_chunks_async(self, records: List[Dict[str, Any]]) -> None: + """Upsert chunk records using async LanceDB API.""" + from ..LanceDB.schema_manager import ensure_chunks_table + + if not records: + return + + async_conn = await self._get_async_connection() + + # Reuse sync connection for table creation + sync_conn = self._get_connection() + ensure_chunks_table(sync_conn) + + table = await async_conn.open_table("chunks") + + # Use merge_insert for efficient upsert + await ( + table.merge_insert(["collection", "doc_id", "parse_hash", "chunk_id"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(records) + ) + + async def upsert_embeddings_async( + self, model_tag: str, records: List[Dict[str, Any]] + ) -> None: + """Upsert embedding records using async LanceDB API. + + Note: This method uses merge_insert without fallback for simplicity. + For production use with error recovery, use the sync upsert_embeddings method. + """ + from ..LanceDB.model_tag_utils import to_model_tag + from ..LanceDB.schema_manager import ensure_embeddings_table + + if not records: + return + + async_conn = await self._get_async_connection() + sync_conn = self._get_connection() + + table_name = f"embeddings_{to_model_tag(model_tag)}" + + # Infer vector dimension from first record + vector_dim = None + if records and "vector" in records[0]: + vector = records[0]["vector"] + if isinstance(vector, (list, tuple)): + vector_dim = len(vector) + + ensure_embeddings_table( + sync_conn, to_model_tag(model_tag), vector_dim=vector_dim + ) + table = await async_conn.open_table(table_name) + + # Use merge_insert for efficient upsert + await ( + table.merge_insert(["collection", "doc_id", "chunk_id"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(records) + ) + + +# ============================================================================ +# Phase 1A Part 2: Additional LanceDB Store Implementations +# ============================================================================ + + +class LanceDBIngestionStatusStore(IngestionStatusStore): + """LanceDB implementation for ingestion status tracking. + + Manages ingestion_runs table for tracking document processing status. + """ + + def __init__(self) -> None: + self._sync_conn: Optional[DBConnection] = None + self._async_conn: Optional[Any] = None + self._async_lock = asyncio.Lock() + + def _get_sync_connection(self) -> DBConnection: + """Get sync LanceDB connection.""" + if self._sync_conn is None: + self._sync_conn = get_connection_from_env() + return self._sync_conn + + async def _get_async_connection(self) -> Any: + """Get async LanceDB connection.""" + if self._async_conn is None: + async with self._async_lock: + if self._async_conn is None: + self._async_conn = await lancedb.connect_async( # type: ignore[attr-defined] + get_connection_from_env().uri # type: ignore[attr-defined] + ) + return self._async_conn + + def _ensure_ingestion_runs_table(self, conn: DBConnection) -> None: + """Ensure ingestion_runs table exists.""" + from ..LanceDB.schema_manager import ensure_ingestion_runs_table + + ensure_ingestion_runs_table(conn) + + # --- Sync methods --- + + def write_ingestion_status( + self, + collection: str, + doc_id: str, + *, + status: str, + message: Optional[str] = None, + parse_hash: Optional[str] = None, + user_id: Optional[int] = None, + ) -> None: + """Write ingestion status record (sync).""" + try: + conn = self._get_sync_connection() + self._ensure_ingestion_runs_table(conn) + table = conn.open_table("ingestion_runs") + + # Delete existing record for this collection/doc_id + base_filter = self._build_base_filter(collection, doc_id) + if base_filter: + table.delete(base_filter) + + # Create new record + timestamp = datetime.now(timezone.utc) + record = { + "collection": collection, + "doc_id": doc_id, + "status": status, + "message": message or "", + "parse_hash": parse_hash or "", + "created_at": timestamp, + "updated_at": timestamp, + "user_id": user_id, + } + table.add([record]) + + except Exception as e: + logger.error(f"Failed to write ingestion status: {e}") + raise + + def load_ingestion_status( + self, + collection: Optional[str] = None, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> List[Dict[str, Any]]: + """Load ingestion status records (sync).""" + try: + conn = self._get_sync_connection() + self._ensure_ingestion_runs_table(conn) + table = conn.open_table("ingestion_runs") + + # Build filter expression + filter_expr = self._build_load_filter(collection, doc_id, user_id, is_admin) + + # Execute query + search = table.search() + if filter_expr: + search = search.where(filter_expr) + result = search.to_arrow() + + # Convert Arrow table to list of dicts (records format) + if len(result) == 0: + return [] + return cast(List[Dict[str, Any]], result.to_pylist()) + + except Exception as e: + logger.error(f"Failed to load ingestion status: {e}") + raise + + def clear_ingestion_status( + self, + collection: str, + doc_id: str, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> None: + """Remove ingestion status record (sync).""" + try: + conn = self._get_sync_connection() + self._ensure_ingestion_runs_table(conn) + table = conn.open_table("ingestion_runs") + + # Build filter with user permissions + base_filter = self._build_base_filter(collection, doc_id) + user_filter = UserPermissions.get_user_filter(user_id, is_admin) + + filter_expr = self._combine_filters(base_filter, user_filter) + if filter_expr: + table.delete(filter_expr) + + except Exception as e: + logger.error(f"Failed to clear ingestion status: {e}") + raise + + # --- Async methods --- + + async def write_ingestion_status_async( + self, + collection: str, + doc_id: str, + *, + status: str, + message: Optional[str] = None, + parse_hash: Optional[str] = None, + user_id: Optional[int] = None, + ) -> None: + """Write ingestion status record (async). + + Note: Current implementation uses sync operations under the hood. + True async I/O will be added in Phase 1B with RDB backend. + """ + # Delegate to sync implementation for now + return self.write_ingestion_status( + collection=collection, + doc_id=doc_id, + status=status, + message=message, + parse_hash=parse_hash, + user_id=user_id, + ) + + async def load_ingestion_status_async( + self, + collection: Optional[str] = None, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> List[Dict[str, Any]]: + """Load ingestion status records (async). + + Note: Current implementation uses sync operations under the hood. + True async I/O will be added in Phase 1B with RDB backend. + """ + # Delegate to sync implementation for now + return self.load_ingestion_status( + collection=collection, + doc_id=doc_id, + user_id=user_id, + is_admin=is_admin, + ) + + async def clear_ingestion_status_async( + self, + collection: str, + doc_id: str, + user_id: Optional[int] = None, + is_admin: bool = False, + ) -> None: + """Remove ingestion status record (async). + + Note: Current implementation uses sync operations under the hood. + True async I/O will be added in Phase 1B with RDB backend. + """ + # Delegate to sync implementation for now + return self.clear_ingestion_status( + collection=collection, + doc_id=doc_id, + user_id=user_id, + is_admin=is_admin, + ) + + # --- Helper methods --- + + def _build_base_filter(self, collection: str, doc_id: str) -> str: + """Build base filter for collection/doc_id.""" + safe_collection = escape_lancedb_string(collection) + safe_doc_id = escape_lancedb_string(doc_id) + return f"collection == '{safe_collection}' AND doc_id == '{safe_doc_id}'" + + def _build_load_filter( + self, + collection: Optional[str], + doc_id: Optional[str], + user_id: Optional[int], + is_admin: bool, + ) -> Optional[str]: + """Build filter for loading status records.""" + conditions = [] + + if collection is not None: + safe_collection = escape_lancedb_string(collection) + conditions.append(f"collection == '{safe_collection}'") + + if doc_id is not None: + safe_doc_id = escape_lancedb_string(doc_id) + conditions.append(f"doc_id == '{safe_doc_id}'") + + # Combine with user filter + base_filter = " AND ".join(conditions) if conditions else None + user_filter = UserPermissions.get_user_filter(user_id, is_admin) + + return self._combine_filters(base_filter, user_filter) + + def _combine_filters( + self, base_filter: Optional[str], user_filter: Optional[str] + ) -> Optional[str]: + """Combine base and user filters.""" + if user_filter and base_filter: + return f"({base_filter}) AND ({user_filter})" + elif user_filter: + return user_filter + return base_filter + + +class LanceDBPromptTemplateStore(PromptTemplateStore): + """LanceDB implementation for prompt template management. + + Manages prompt_templates table for storing and retrieving prompt templates. + """ + + def __init__(self) -> None: + self._sync_conn: Optional[DBConnection] = None + + def _get_sync_connection(self) -> DBConnection: + """Get or create sync connection.""" + if self._sync_conn is None: + self._sync_conn = get_connection_from_env() + return self._sync_conn + + def _ensure_table(self) -> None: + """Ensure prompt_templates table exists.""" + from ..LanceDB.schema_manager import ensure_prompt_templates_table + + conn = self._get_sync_connection() + ensure_prompt_templates_table(conn) + + # --- Sync methods --- + + def save_prompt_template( + self, + name: str, + template: str, + user_id: Optional[int] = None, + metadata: Optional[str] = None, + ) -> str: + """Save or update a prompt template (sync).""" + import uuid + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("prompt_templates") + + # Generate new template ID + template_id = str(uuid.uuid4()) + now = datetime.now(timezone.utc).replace(tzinfo=None) + + # Check for existing templates with same name to get next version + base_filter = f"name == '{escape_lancedb_string(name)}'" + if user_id is not None: + base_filter += f" AND user_id == {user_id}" + + existing = table.search().where(base_filter).to_arrow() + if len(existing) > 0: + import pyarrow.compute as pc # type: ignore[import-not-found] + + max_version = pc.max(existing["version"]).as_py() + new_version = max_version + 1 + + # Mark previous versions as not latest + for row in existing.to_pylist(): + if row["is_latest"]: + table.update( + where=f"id == '{row['id']}'", + values={"is_latest": False}, + ) + else: + new_version = 1 + + # Create new template record + record = { + "id": template_id, + "name": name, + "template": template, + "version": new_version, + "is_latest": True, + "metadata": metadata or "", + "user_id": user_id or 0, + "created_at": now, + "updated_at": now, + } + + table.add([record]) + logger.info("Saved prompt template: %s (version %d)", name, new_version) + return template_id + + def get_prompt_template( + self, + template_id: str, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Get a prompt template by ID (sync).""" + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("prompt_templates") + + base_filter = f"id == '{escape_lancedb_string(template_id)}'" + if user_id is not None: + base_filter += f" AND user_id == {user_id}" + + result = table.search().where(base_filter).to_arrow() + if len(result) == 0: + return None + + # Convert Arrow table to list of dicts and take first row + row = result.to_pylist()[0] + return { + "id": row["id"], + "name": row["name"], + "template": row["template"], + "version": int(row["version"]), + "is_latest": bool(row["is_latest"]), + "metadata": row["metadata"], + "user_id": int(row["user_id"]) if row["user_id"] else None, + "created_at": row["created_at"], + "updated_at": row["updated_at"], + } + + def get_latest_prompt_template( + self, + name: str, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Get the latest version of a prompt template by name (sync).""" + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("prompt_templates") + + base_filter = f"name == '{escape_lancedb_string(name)}' AND is_latest == true" + if user_id is not None: + base_filter += f" AND user_id == {user_id}" + + result = table.search().where(base_filter).to_arrow() + if len(result) == 0: + return None + + # Convert Arrow table to list of dicts and take first row + row = result.to_pylist()[0] + return { + "id": row["id"], + "name": row["name"], + "template": row["template"], + "version": int(row["version"]), + "is_latest": bool(row["is_latest"]), + "metadata": row["metadata"], + "user_id": int(row["user_id"]) if row["user_id"] else None, + "created_at": row["created_at"], + "updated_at": row["updated_at"], + } + + def list_prompt_templates( + self, + name_filter: Optional[str] = None, + latest_only: bool = False, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """List prompt templates with optional filtering (sync).""" + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("prompt_templates") + + filters = [] + if name_filter: + filters.append(f"name LIKE '%{escape_lancedb_string(name_filter)}%'") + if latest_only: + filters.append("is_latest == true") + if user_id is not None: + filters.append(f"user_id == {user_id}") + + filter_expr = " AND ".join(filters) if filters else None + + query = table.search() + if filter_expr: + query = query.where(filter_expr) + + result = query.limit(limit).to_arrow() + templates = [] + for row_dict in result.to_pylist(): + templates.append( + { + "id": row_dict["id"], + "name": row_dict["name"], + "template": row_dict["template"], + "version": int(row_dict["version"]), + "is_latest": bool(row_dict["is_latest"]), + "metadata": row_dict["metadata"], + "user_id": int(row_dict["user_id"]) + if row_dict["user_id"] + else None, + "created_at": row_dict["created_at"], + "updated_at": row_dict["updated_at"], + } + ) + + return templates + + def delete_prompt_template( + self, + template_id: str, + user_id: Optional[int] = None, + ) -> bool: + """Delete a prompt template by ID (sync). + + Updates is_latest flag for remaining versions if latest version is deleted. + """ + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("prompt_templates") + + base_filter = f"id == '{escape_lancedb_string(template_id)}'" + if user_id is not None: + base_filter += f" AND user_id == {user_id}" + + # Check if exists and get info + result = table.search().where(base_filter).to_arrow() + if len(result) == 0: + return False + + # Check if this was the latest version and get the name + # Convert Arrow table to list of dicts and take first row + row_dict = result.to_pylist()[0] + was_latest = row_dict["is_latest"] + template_name = row_dict["name"] + + table.delete(base_filter) + + # If we deleted the latest version, update the latest flag for the remaining versions + if was_latest: + name_filter = f"name == '{escape_lancedb_string(template_name)}'" + if user_id is not None: + name_filter += f" AND user_id == {user_id}" + + remaining_versions = table.search().where(name_filter).to_arrow() + if len(remaining_versions) > 0: + import pyarrow.compute as pc + + max_version = pc.max(remaining_versions["version"]).as_py() + update_filter = f"{name_filter} AND version == {max_version}" + table.update(where=update_filter, values={"is_latest": True}) + + logger.info("Deleted prompt template: %s", template_id) + return True + + def update_metadata( + self, + template_id: str, + metadata: Optional[str], + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Update metadata only, keeping same version and ID (sync).""" + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("prompt_templates") + + base_filter = f"id == '{escape_lancedb_string(template_id)}'" + if user_id is not None: + base_filter += f" AND user_id == {user_id}" + + # Check if exists + result = table.search().where(base_filter).to_arrow() + if len(result) == 0: + return None + + # Update metadata + table.update( + where=base_filter, + values={ + "metadata": metadata or "", + "updated_at": datetime.now(timezone.utc).replace(tzinfo=None), + }, + ) + logger.info("Updated metadata for prompt template: %s", template_id) + + # Return updated template + return self.get_prompt_template(template_id, user_id) + + def delete_by_name( + self, + name: str, + version: Optional[int] = None, + user_id: Optional[int] = None, + ) -> int: + """Delete template(s) by name (sync). + + Handles is_latest flag updates for remaining versions. + """ + from ..core.exceptions import DocumentNotFoundError + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("prompt_templates") + + escaped_name = escape_lancedb_string(name) + base_filter = f"name == '{escaped_name}'" + if user_id is not None: + base_filter += f" AND user_id == {user_id}" + + if version is not None: + # Delete specific version + version_filter = f"{base_filter} AND version == {version}" + result = table.search().where(version_filter).to_arrow() + if len(result) == 0: + raise DocumentNotFoundError( + f"Prompt template '{name}' version {version} not found." + ) + + # Convert Arrow table to list of dicts and take first row + row_dict = result.to_pylist()[0] + was_latest = row_dict["is_latest"] + table.delete(version_filter) + + # If we deleted the latest version, update the latest flag + if was_latest: + remaining = table.search().where(base_filter).to_arrow() + if len(remaining) > 0: + import pyarrow.compute as pc + + max_version = pc.max(remaining["version"]).as_py() + table.update( + where=f"{base_filter} AND version == {max_version}", + values={"is_latest": True}, + ) + + logger.info("Deleted prompt template '%s' version %d", name, version) + return 1 + else: + # Delete all versions + result = table.search().where(base_filter).to_arrow() + if len(result) == 0: + raise DocumentNotFoundError(f"Prompt template '{name}' not found.") + + count = len(result) + table.delete(base_filter) + logger.info("Deleted all %d versions of prompt template '%s'", count, name) + return count + + def get_versions_by_name( + self, + name: str, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """Get all versions of a template by name (sync).""" + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("prompt_templates") + + base_filter = f"name == '{escape_lancedb_string(name)}'" + if user_id is not None: + base_filter += f" AND user_id == {user_id}" + + result = table.search().where(base_filter).limit(limit).to_arrow() + templates = [] + for row_dict in result.to_pylist(): + templates.append( + { + "id": row_dict["id"], + "name": row_dict["name"], + "template": row_dict["template"], + "version": int(row_dict["version"]), + "is_latest": bool(row_dict["is_latest"]), + "metadata": row_dict["metadata"], + "user_id": int(row_dict["user_id"]) + if row_dict["user_id"] + else None, + "created_at": row_dict["created_at"], + "updated_at": row_dict["updated_at"], + } + ) + + return templates + + # --- Async methods (delegate to sync) --- + + async def save_prompt_template_async( + self, + name: str, + template: str, + user_id: Optional[int] = None, + metadata: Optional[str] = None, + ) -> str: + """Async version of save_prompt_template.""" + return self.save_prompt_template(name, template, user_id, metadata) + + async def get_prompt_template_async( + self, + template_id: str, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of get_prompt_template.""" + return self.get_prompt_template(template_id, user_id) + + async def get_latest_prompt_template_async( + self, + name: str, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of get_latest_prompt_template.""" + return self.get_latest_prompt_template(name, user_id) + + async def list_prompt_templates_async( + self, + name_filter: Optional[str] = None, + latest_only: bool = False, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """Async version of list_prompt_templates.""" + return self.list_prompt_templates(name_filter, latest_only, user_id, limit) + + async def delete_prompt_template_async( + self, + template_id: str, + user_id: Optional[int] = None, + ) -> bool: + """Async version of delete_prompt_template.""" + return self.delete_prompt_template(template_id, user_id) + + async def update_metadata_async( + self, + template_id: str, + metadata: Optional[str], + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of update_metadata.""" + return self.update_metadata(template_id, metadata, user_id) + + async def delete_by_name_async( + self, + name: str, + version: Optional[int] = None, + user_id: Optional[int] = None, + ) -> int: + """Async version of delete_by_name.""" + return self.delete_by_name(name, version, user_id) + + async def get_versions_by_name_async( + self, + name: str, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """Async version of get_versions_by_name.""" + return self.get_versions_by_name(name, user_id, limit) + + +class LanceDBMainPointerStore(MainPointerStore): + """LanceDB implementation for main pointer management. + + Manages main_pointers table for tracking current versions across + processing stages (parse, chunk, embed). + + NOTE: user_id parameter is logged but not used, as main_pointers table + schema does not include user_id field. Schema migration required for + multi-tenancy support. + """ + + def __init__(self) -> None: + self._sync_conn: Optional[DBConnection] = None + + def _get_sync_connection(self) -> DBConnection: + """Get or create sync connection.""" + if self._sync_conn is None: + self._sync_conn = get_connection_from_env() + return self._sync_conn + + def _ensure_table(self) -> None: + """Ensure main_pointers table exists.""" + from ..LanceDB.schema_manager import ensure_main_pointers_table + + conn = self._get_sync_connection() + ensure_main_pointers_table(conn) + + def _normalize_model_tag(self, model_tag: Optional[str]) -> str: + """Normalize model_tag to empty string if None.""" + return model_tag if model_tag is not None else "" + + # --- Sync methods --- + + def set_main_pointer( + self, + collection: str, + doc_id: str, + step_type: str, + semantic_id: str, + technical_id: str, + model_tag: Optional[str] = None, + operator: Optional[str] = None, + user_id: Optional[int] = None, + ) -> None: + """Set or update a main pointer (sync).""" + if user_id is not None: + logger.warning( + "user_id parameter provided to set_main_pointer but " + "main_pointers table does not have user_id field. " + "Schema migration required for multi-tenancy support." + ) + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("main_pointers") + + normalized_tag = self._normalize_model_tag(model_tag) + now = datetime.now(timezone.utc).replace(tzinfo=None) + + # Check if pointer already exists to preserve created_at + existing = self.get_main_pointer(collection, doc_id, step_type, model_tag) + + created_at = existing["created_at"] if existing else now + + # Prepare data for merge_insert + update_data: Dict[str, List[Any]] = { + "collection": [collection], + "doc_id": [doc_id], + "step_type": [step_type], + "model_tag": [normalized_tag], + "semantic_id": [semantic_id], + "technical_id": [technical_id], + "created_at": [created_at], + "updated_at": [now], + "operator": [operator or "unknown"], + } + # Convert dict of lists to list of dicts for merge_insert + records = [ + {key: values[idx] for key, values in update_data.items()} + for idx in range(len(update_data["collection"])) + ] + + ( + table.merge_insert(on=["collection", "doc_id", "step_type", "model_tag"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(records) + ) + + logger.info( + "Set main pointer for %s/%s/%s to %s (semantic: %s)", + collection, + doc_id, + step_type, + technical_id, + semantic_id, + ) + + def get_main_pointer( + self, + collection: str, + doc_id: str, + step_type: str, + model_tag: Optional[str] = None, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Get a main pointer (sync).""" + if user_id is not None: + logger.warning( + "user_id parameter provided to get_main_pointer but " + "main_pointers table does not have user_id field. " + "Schema migration required for multi-tenancy support." + ) + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("main_pointers") + + # Build filter expression using FilterCondition + base_conditions: List[FilterCondition] = [ + FilterCondition( + field="collection", operator=FilterOperator.EQ, value=collection + ), + FilterCondition(field="doc_id", operator=FilterOperator.EQ, value=doc_id), + FilterCondition( + field="step_type", operator=FilterOperator.EQ, value=step_type + ), + ] + + normalized_tag = self._normalize_model_tag(model_tag) + if normalized_tag == "": + # Check for both empty string AND NULL (backward compatibility) + model_tag_null_cond = FilterCondition( + field="model_tag", operator=FilterOperator.IS_NULL, value=None + ) + model_tag_empty_cond = FilterCondition( + field="model_tag", operator=FilterOperator.EQ, value="" + ) + # Combine as: (base) AND (model_tag IS NULL OR model_tag == '') + model_tag_filter: FilterExpression = ( + model_tag_null_cond, + model_tag_empty_cond, + ) # OR tuple + filter_expr: FilterExpression = ( + *base_conditions, + model_tag_filter, + ) # AND tuple + else: + base_conditions.append( + FilterCondition( + field="model_tag", operator=FilterOperator.EQ, value=normalized_tag + ) + ) + filter_expr = tuple(base_conditions) # AND tuple + + # Translate to LanceDB syntax using shared utility + filter_str = translate_filter_expression(filter_expr) + + result = table.search().where(filter_str).to_arrow() + + if len(result) == 0: + return None + + # Return the first result, preferring non-NULL model_tag if multiple found + if len(result) > 1: + import pyarrow.compute as pc + + # Sort by model_tag descending (NULLs last) + sort_indices = pc.sort_indices( + result, sort_keys=[("model_tag", "descending")] + ) + result = result.take(sort_indices) + + # Convert Arrow table to list of dicts and take first row + row_dict = result.to_pylist()[0] + return { + "collection": row_dict["collection"], + "doc_id": row_dict["doc_id"], + "step_type": row_dict["step_type"], + "model_tag": row_dict["model_tag"] + if row_dict["model_tag"] is not None + else None, + "semantic_id": row_dict["semantic_id"], + "technical_id": row_dict["technical_id"], + "created_at": row_dict["created_at"], + "updated_at": row_dict["updated_at"], + "operator": row_dict["operator"], + } + + def list_main_pointers( + self, + collection: str, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """List main pointers (sync).""" + if user_id is not None: + logger.warning( + "user_id parameter provided to list_main_pointers but " + "main_pointers table does not have user_id field. " + "Schema migration required for multi-tenancy support." + ) + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("main_pointers") + + filters_dict = {"collection": collection} + if doc_id is not None: + filters_dict["doc_id"] = doc_id + + filter_expr = build_lancedb_filter_expression(filters_dict) + + # First check if any pointers exist using efficient count_rows + if table.search().where(filter_expr).count_rows() == 0: + return [] + + result = table.search().where(filter_expr).limit(limit).to_arrow() + + pointers = [] + for row_dict in result.to_pylist(): + pointers.append( + { + "collection": row_dict["collection"], + "doc_id": row_dict["doc_id"], + "step_type": row_dict["step_type"], + "model_tag": row_dict["model_tag"] + if row_dict["model_tag"] is not None + else None, + "semantic_id": row_dict["semantic_id"], + "technical_id": row_dict["technical_id"], + "created_at": row_dict["created_at"], + "updated_at": row_dict["updated_at"], + "operator": row_dict["operator"], + } + ) + + return pointers + + def delete_main_pointer( + self, + collection: str, + doc_id: str, + step_type: str, + model_tag: Optional[str] = None, + user_id: Optional[int] = None, + ) -> bool: + """Delete a main pointer (sync).""" + if user_id is not None: + logger.warning( + "user_id parameter provided to delete_main_pointer but " + "main_pointers table does not have user_id field. " + "Schema migration required for multi-tenancy support." + ) + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("main_pointers") + + # Build filter expression using FilterCondition + base_conditions: List[FilterCondition] = [ + FilterCondition( + field="collection", operator=FilterOperator.EQ, value=collection + ), + FilterCondition(field="doc_id", operator=FilterOperator.EQ, value=doc_id), + FilterCondition( + field="step_type", operator=FilterOperator.EQ, value=step_type + ), + ] + + normalized_tag = self._normalize_model_tag(model_tag) + if normalized_tag == "": + # Check for both empty string AND NULL (backward compatibility) + model_tag_null_cond = FilterCondition( + field="model_tag", operator=FilterOperator.IS_NULL, value=None + ) + model_tag_empty_cond = FilterCondition( + field="model_tag", operator=FilterOperator.EQ, value="" + ) + # Combine as: (base) AND (model_tag IS NULL OR model_tag == '') + model_tag_filter: FilterExpression = ( + model_tag_null_cond, + model_tag_empty_cond, + ) # OR tuple + filter_expr: FilterExpression = ( + *base_conditions, + model_tag_filter, + ) # AND tuple + else: + base_conditions.append( + FilterCondition( + field="model_tag", operator=FilterOperator.EQ, value=normalized_tag + ) + ) + filter_expr = tuple(base_conditions) # AND tuple + + # Translate to LanceDB syntax using shared utility + filter_str = translate_filter_expression(filter_expr) + + # Check if exists + result = table.search().where(filter_str).to_arrow() + if len(result) == 0: + return False + + table.delete(filter_str) + logger.info("Deleted main pointer for %s/%s/%s", collection, doc_id, step_type) + return True + + # --- Async methods (delegate to sync) --- + + async def set_main_pointer_async( + self, + collection: str, + doc_id: str, + step_type: str, + semantic_id: str, + technical_id: str, + model_tag: Optional[str] = None, + operator: Optional[str] = None, + user_id: Optional[int] = None, + ) -> None: + """Async version of set_main_pointer.""" + return self.set_main_pointer( + collection, + doc_id, + step_type, + semantic_id, + technical_id, + model_tag, + operator, + user_id, + ) + + async def get_main_pointer_async( + self, + collection: str, + doc_id: str, + step_type: str, + model_tag: Optional[str] = None, + user_id: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of get_main_pointer.""" + return self.get_main_pointer(collection, doc_id, step_type, model_tag, user_id) + + async def list_main_pointers_async( + self, + collection: str, + doc_id: Optional[str] = None, + user_id: Optional[int] = None, + limit: int = 100, + ) -> List[Dict[str, Any]]: + """Async version of list_main_pointers.""" + return self.list_main_pointers(collection, doc_id, user_id, limit) + + async def delete_main_pointer_async( + self, + collection: str, + doc_id: str, + step_type: str, + model_tag: Optional[str] = None, + user_id: Optional[int] = None, + ) -> bool: + """Async version of delete_main_pointer.""" + return self.delete_main_pointer( + collection, doc_id, step_type, model_tag, user_id + ) diff --git a/src/xagent/core/tools/core/RAG_tools/storage/logging_utils.py b/src/xagent/core/tools/core/RAG_tools/storage/logging_utils.py new file mode 100644 index 000000000..5833010e5 --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/storage/logging_utils.py @@ -0,0 +1,151 @@ +"""Structured logging utilities for storage operations. + +This module provides utilities for structured logging with performance tracking +and audit capabilities for RAG storage operations. +""" + +import logging +import time +from contextlib import contextmanager +from functools import wraps +from typing import Any, Callable, Dict, Iterator, Optional + +logger = logging.getLogger(__name__) + + +@contextmanager +def log_operation(operation: str, **extra_context: Any) -> Iterator[None]: + """Context manager for logging operation with timing and structured output. + + Usage: + with log_operation("upsert_documents", table="chunks", count=100): + # ... perform operation ... + # Will log: operation_started, operation_completed (with duration_ms) + # On exception: operation_failed (with error details) + + Args: + operation: Name of the operation being performed + **extra_context: Additional context to include in all log entries + + Yields: + None + """ + start_time = time.time() + try: + logger.info( + "operation_started", extra={"operation": operation, **extra_context} + ) + yield + except Exception as e: + logger.error( + "operation_failed", + extra={ + "operation": operation, + "error": str(e), + "error_type": type(e).__name__, + **extra_context, + }, + exc_info=True, + ) + raise + finally: + duration_ms = (time.time() - start_time) * 1000 + logger.info( + "operation_completed", + extra={ + "operation": operation, + "duration_ms": round(duration_ms, 2), + **extra_context, + }, + ) + + +def log_async_operation(operation: str, **extra_context: Any) -> Callable: + """Decorator for async operations with automatic timing and structured logging. + + Usage: + @log_async_operation("search_vectors", table="embeddings_test") + async def search_vectors_async(self, ...): + # ... async operation ... + + Args: + operation: Name of the operation being performed + **extra_context: Additional context to include in all log entries + + Returns: + Decorator function + """ + + def decorator(func: Callable) -> Callable: + @wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + start_time = time.time() + # Extract context from args/kwargs if possible + context = dict(extra_context) + + # Try to extract self and method name for better logging + if args and hasattr(args[0], "__class__"): + context["class"] = args[0].__class__.__name__ + + try: + logger.info( + "operation_started", extra={"operation": operation, **context} + ) + result = await func(*args, **kwargs) + + duration_ms = (time.time() - start_time) * 1000 + logger.info( + "operation_completed", + extra={ + "operation": operation, + "duration_ms": round(duration_ms, 2), + **context, + }, + ) + return result + except Exception as e: + duration_ms = (time.time() - start_time) * 1000 + logger.error( + "operation_failed", + extra={ + "operation": operation, + "error": str(e), + "error_type": type(e).__name__, + "duration_ms": round(duration_ms, 2), + **context, + }, + exc_info=True, + ) + raise + + return wrapper + + return decorator + + +def log_audit(operation: str, **context: Any) -> None: + """Log an audit event for security and compliance tracking. + + Args: + operation: The operation being performed (e.g., "data_access", "permission_check") + **context: Audit context (user_id, collection, doc_id, etc.) + """ + logger.info("audit", extra={"operation": operation, **context}) + + +def log_performance( + metric_name: str, value: Optional[float] = None, unit: str = "ms", **context: Any +) -> None: + """Log a performance metric. + + Args: + metric_name: Name of the metric (e.g., "query_duration", "batch_size") + value: Numeric value of the metric (optional for metrics that only need context) + unit: Unit of measurement (default: "ms") + **context: Additional context + """ + extra: Dict[str, Any] = {"metric": metric_name, **context} + if value is not None: + extra["value"] = value + extra["unit"] = unit + logger.debug("performance", extra=extra) diff --git a/src/xagent/core/tools/core/RAG_tools/storage/vector_backend.py b/src/xagent/core/tools/core/RAG_tools/storage/vector_backend.py new file mode 100644 index 000000000..203f105ae --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/storage/vector_backend.py @@ -0,0 +1,81 @@ +"""Vector index backend selection (switchable vector store). + +Resolve which :class:`~.contracts.VectorIndexStore` implementation to use from +environment. Only LanceDB is implemented today; additional backends register +here and in :meth:`StorageFactory.get_vector_index_store`. +""" + +from __future__ import annotations + +import os +from enum import StrEnum +from typing import Final + +from ..core.exceptions import ConfigurationError + +# Primary env var (namespaced to avoid collisions with other libs). +VECTOR_BACKEND_ENV: Final[str] = "XAGENT_VECTOR_BACKEND" + +# Backward-compatible alias used in some deployments / docs. +VECTOR_BACKEND_ENV_LEGACY: Final[str] = "VECTOR_STORE_BACKEND" + + +class VectorBackend(StrEnum): + """Supported or reserved vector index backends.""" + + LANCEDB = "lancedb" + MILVUS = "milvus" + QDRANT = "qdrant" + + +def _parse_backend(raw: str) -> VectorBackend: + """Parse and validate backend string.""" + key = raw.strip().lower() + if not key: + return VectorBackend.LANCEDB + try: + return VectorBackend(key) + except ValueError as exc: + allowed = ", ".join(sorted(b.value for b in VectorBackend)) + raise ConfigurationError( + f"Invalid {VECTOR_BACKEND_ENV}={raw!r}. Choose one of: {allowed}." + ) from exc + + +def get_configured_vector_backend() -> VectorBackend: + """Read configured vector backend from the environment. + + Precedence: ``XAGENT_VECTOR_BACKEND``, then ``VECTOR_STORE_BACKEND``, + then default ``lancedb``. + + Returns: + Selected :class:`VectorBackend`. + + Raises: + ConfigurationError: If the value is not a known backend name. + """ + raw = os.environ.get(VECTOR_BACKEND_ENV) + if raw is None or raw.strip() == "": + raw = os.environ.get(VECTOR_BACKEND_ENV_LEGACY, "") + return _parse_backend(raw) + + +def require_implemented_vector_backend(backend: VectorBackend) -> None: + """Ensure the backend has a concrete :class:`~.contracts.VectorIndexStore`. + + Call from the factory before instantiating stores. Extend this function + when adding Milvus, Qdrant, etc. + + Args: + backend: Resolved backend. + + Raises: + ConfigurationError: If the backend is known but not implemented yet. + """ + if backend is VectorBackend.LANCEDB: + return + raise ConfigurationError( + f"Vector backend {backend.value!r} is not implemented yet. " + f"Set {VECTOR_BACKEND_ENV}=lancedb (default), or contribute a " + f"{backend.value} implementation of VectorIndexStore." + ) diff --git a/src/xagent/core/tools/core/RAG_tools/utils/filter_utils.py b/src/xagent/core/tools/core/RAG_tools/utils/filter_utils.py new file mode 100644 index 000000000..a399ba975 --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/utils/filter_utils.py @@ -0,0 +1,109 @@ +"""Filter parsing utilities for backend-agnostic filter expressions. + +This module provides utilities to convert API-facing filter dictionaries into +abstract filter expressions that can be translated to backend-specific syntax. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +from ..storage.contracts import FilterCondition, FilterExpression, FilterOperator + + +def validate_filter_depth( + expr: Optional[FilterExpression], + max_depth: int = 10, +) -> None: + """Validate filter expression depth to prevent DoS via deeply nested filters. + + This should be called on user-provided filter expressions before they + are passed to build_filter_expression. + + Args: + expr: Filter expression to validate. + max_depth: Maximum allowed nesting depth (default: 10). + + Raises: + ValueError: If filter expression exceeds max_depth. + """ + if expr is None: + return + + def _check_depth(e: FilterExpression, depth: int = 0) -> None: + if depth > max_depth: + raise ValueError( + f"Filter expression depth exceeds maximum allowed depth of {max_depth}. " + "This may indicate a malicious or malformed filter expression." + ) + if isinstance(e, FilterCondition): + return + elif isinstance(e, tuple): + for item in e: + _check_depth(item, depth + 1) + elif isinstance(e, list): + for item in e: + _check_depth(item, depth + 1) + + _check_depth(expr) + + +def parse_legacy_filters( + filters: Optional[Dict[str, Any]], + max_depth: int = 10, +) -> Optional[FilterExpression]: + """Convert Dict-based filters to an abstract FilterExpression. + + Supported input formats: + - Simple equality: + {"field": "value"} + - Operator form: + {"field": {"operator": "gte", "value": 5}} + + Multiple fields are combined as an AND expression (tuple convention). + + Args: + filters: Filter dictionary from API layer. + max_depth: Maximum allowed nesting depth (default: 10). + + Returns: + Parsed FilterExpression, or None if filters is None/empty. + + Raises: + ValueError: If an unsupported operator is provided or depth exceeds max_depth. + """ + if not filters: + return None + + op_map: Dict[str, FilterOperator] = { + "eq": FilterOperator.EQ, + "ne": FilterOperator.NE, + "gt": FilterOperator.GT, + "gte": FilterOperator.GTE, + "lt": FilterOperator.LT, + "lte": FilterOperator.LTE, + "in": FilterOperator.IN, + "contains": FilterOperator.CONTAINS, + } + + conditions: list[FilterCondition] = [] + for field, spec in filters.items(): + if isinstance(spec, dict) and "operator" in spec and "value" in spec: + op_str = str(spec["operator"]).lower() + if op_str not in op_map: + raise ValueError( + f"Unknown filter operator: {op_str}. Supported operators: {sorted(op_map.keys())}" + ) + conditions.append( + FilterCondition( + field=field, operator=op_map[op_str], value=spec["value"] + ) + ) + else: + conditions.append( + FilterCondition(field=field, operator=FilterOperator.EQ, value=spec) + ) + + if len(conditions) == 1: + return conditions[0] + return tuple(conditions) diff --git a/src/xagent/core/tools/core/RAG_tools/utils/migration_utils.py b/src/xagent/core/tools/core/RAG_tools/utils/migration_utils.py index 67c3cea24..f5e28f6da 100644 --- a/src/xagent/core/tools/core/RAG_tools/utils/migration_utils.py +++ b/src/xagent/core/tools/core/RAG_tools/utils/migration_utils.py @@ -1,20 +1,34 @@ """Utilities for handling schema migrations and backward compatibility.""" +import fcntl import logging +import os from datetime import datetime, timezone from typing import Any, Dict, Optional, Tuple, cast -from ......providers.vector_store.lancedb import get_connection_from_env +import pyarrow as pa # type: ignore + +from ..LanceDB.model_tag_utils import to_model_tag +from ..storage.factory import get_vector_store_raw_connection from .string_utils import escape_lancedb_string +from .tag_mapping import register_tag_mapping logger = logging.getLogger(__name__) -def migrate_collection_metadata(legacy_data: Dict[str, Any]) -> Dict[str, Any]: +def migrate_collection_metadata( + legacy_data: Dict[str, Any], + *, + infer_embedding: bool = True, +) -> Dict[str, Any]: """Migrate legacy collection metadata to current schema version. Args: legacy_data: Legacy collection data from storage + infer_embedding: If True (default), ``0.0.0 -> 1.0.0`` may scan LanceDB + embedding tables to infer ``embedding_model_id`` / dimension. Use + **False** for read-only deserialization (e.g. :meth:`CollectionInfo.from_storage`) + to avoid I/O, heavy work, and log noise on hot paths. Returns: Migrated data compatible with current schema @@ -24,7 +38,8 @@ def migrate_collection_metadata(legacy_data: Dict[str, Any]) -> Dict[str, Any]: data_version = data.get("schema_version", "0.0.0") collection_name = data.get("name", "unknown") - logger.info( + log_info = logger.info if infer_embedding else logger.debug + log_info( f"[MIGRATION_START] Collection: {collection_name}, From: {data_version}, To: {current_version}" ) logger.debug( @@ -36,14 +51,14 @@ def migrate_collection_metadata(legacy_data: Dict[str, Any]) -> Dict[str, Any]: while data_version < current_version: previous_version = data_version if data_version == "0.0.0": - data = _migrate_0_0_0_to_1_0_0(data) + data = _migrate_0_0_0_to_1_0_0(data, infer_embedding=infer_embedding) data_version = "1.0.0" - logger.info( + log_info( f"[MIGRATION_STEP] {collection_name}: {previous_version} -> {data_version} completed." ) - logger.info( + log_info( f"[MIGRATION_SUCCESS] Collection '{collection_name}' is now at version {data_version}" ) return data @@ -56,14 +71,21 @@ def migrate_collection_metadata(legacy_data: Dict[str, Any]) -> Dict[str, Any]: raise -def _migrate_0_0_0_to_1_0_0(data: Dict[str, Any]) -> Dict[str, Any]: +def _migrate_0_0_0_to_1_0_0( + data: Dict[str, Any], + *, + infer_embedding: bool = True, +) -> Dict[str, Any]: """Migrate from pre-versioned schema to 1.0.0.""" collection_name = data.get("name", "") - # Try to infer embedding config from existing data - embedding_model_id, embedding_dimension = _infer_embedding_config_from_collection( - collection_name - ) + if infer_embedding: + embedding_model_id, embedding_dimension = ( + _infer_embedding_config_from_collection(collection_name) + ) + else: + embedding_model_id = data.get("embedding_model_id") + embedding_dimension = data.get("embedding_dimension") if embedding_model_id: logger.info( @@ -125,7 +147,7 @@ def _infer_embedding_config_from_collection( try: # Get LanceDB connection logger.debug(f"Connecting to LanceDB for collection '{collection_name}'") - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # Get all table names that contain embeddings table_names_fn = getattr(conn, "table_names", None) @@ -236,8 +258,49 @@ def _infer_embedding_config_from_collection( ) model_tag, stats = best_model - # Convert model tag back to model ID - embedding_model_id = _model_tag_to_model_id(model_tag) + # Resolve Hub embedding model ID from table tag (preferred). + embedding_model_id = None + try: + from xagent.core.model.model import EmbeddingModelConfig + + from .model_resolver import _get_or_init_model_hub + + hub = _get_or_init_model_hub() + if hub is not None: + hub_tag_to_id: Dict[str, str] = {} + for cfg in hub.list().values(): + if not isinstance(cfg, EmbeddingModelConfig): + continue + register_tag_mapping( + hub_tag_to_id, + to_model_tag(cfg.id), + cfg.id, + get_identity=lambda item: item, + logger=logger, + ) + register_tag_mapping( + hub_tag_to_id, + to_model_tag(cfg.model_name), + cfg.id, + get_identity=lambda item: item, + logger=logger, + ) + embedding_model_id = hub_tag_to_id.get(model_tag) + except Exception as e: + logger.warning( + "Model hub initialization failed during embedding config inference: " + "error_type=%s, error_message=%s, fallback_behavior=%s, impact=%s", + type(e).__name__, + str(e), + "legacy_model_tag_normalization", + "May use incorrect model ID for embeddings", + exc_info=True, + ) + embedding_model_id = None + + # Fallback: best-effort reverse normalization (legacy behavior) + if not embedding_model_id: + embedding_model_id = _model_tag_to_model_id(model_tag) embedding_dimension = stats["dimension"] logger.info( @@ -286,3 +349,337 @@ def _model_tag_to_model_id(model_tag: str) -> str: result = model_tag.replace("_", "-").lower() logger.debug(f"Used fallback conversion for model tag: {result}") return result + + +def migrate_embeddings_table( + model_id: str, + batch_size: int = 10000, + conn: Optional[Any] = None, +) -> dict[str, Any]: + """Migrate legacy embeddings table to Hub ID-based naming using idempotent merge strategy. + + This function uses LanceDB's merge_insert for safe, non-destructive migration: + - Self-protection: Detects if already migrated (legacy == primary) + - Dimension validation: Ensures source and target tables have compatible vector dimensions + - Idempotent merge: Uses merge_insert to avoid duplicates and data loss + - Arrow streaming: Uses to_batches() for memory-efficient processing + - Cloud-native: Works with S3/OSS (no shutil.move or file system assumptions) + + This addresses critical issues with the previous approach: + - No data loss: merge_insert preserves existing data in target table + - Cloud-compatible: No dependency on file system operations + - Idempotent: Can be safely re-run without side effects + - High performance: Arrow streaming + merge_insert is 5-10x faster than offset/limit + + Args: + model_id: Hub model ID to migrate (e.g., "text-embedding-ada-002"). + batch_size: Number of rows to process per batch (default 10000). + conn: LanceDB connection (if None, creates new connection). + + Returns: + Dictionary with migration results: + { + "success": bool, + "source_table": str (legacy table name), + "target_table": str (Hub ID table name), + "rows_migrated": int, + "error": str | None (if success=False) + } + """ + from ..core.exceptions import VectorValidationError + from ..LanceDB.schema_manager import ensure_embeddings_table + from ..utils.model_resolver import resolve_embedding_adapter + + cleaned = (model_id or "").strip() + if not cleaned: + raise VectorValidationError("model_id must be a non-empty string") + + primary_table_name = f"embeddings_{to_model_tag(cleaned)}" + lock_key = f"migrate_{primary_table_name}" + + # Get connection + if conn is None: + conn = get_vector_store_raw_connection() + + # Try to find legacy table + legacy_table_name: Optional[str] = None + try: + cfg, _ = resolve_embedding_adapter(cleaned) + legacy_table_name = f"embeddings_{to_model_tag(cfg.model_name)}" + except Exception as e: + logger.warning("Failed to resolve legacy table name: %s", e) + + if not legacy_table_name: + return { + "success": False, + "source_table": None, + "target_table": primary_table_name, + "rows_migrated": 0, + "error": "Could not determine legacy table name", + } + + # Self-protection: Check if already migrated + if legacy_table_name == primary_table_name: + logger.info( + "Already migrated: legacy table '%s' is the same as primary table '%s'", + legacy_table_name, + primary_table_name, + ) + return { + "success": True, + "source_table": legacy_table_name, + "target_table": primary_table_name, + "rows_migrated": 0, + "error": None, + } + + # Acquire lock in database directory for distributed environments + lock_fd = _acquire_migration_lock(conn.uri, primary_table_name) + if lock_fd is None: + return { + "success": False, + "source_table": None, + "target_table": primary_table_name, + "rows_migrated": 0, + "error": "Migration already in progress", + } + + rows_migrated = 0 + + try: + # Check if legacy table exists + try: + legacy_table = conn.open_table(legacy_table_name) + except Exception as e: + logger.warning("Legacy table '%s' not found: %s", legacy_table_name, e) + return { + "success": False, + "source_table": legacy_table_name, + "target_table": primary_table_name, + "rows_migrated": 0, + "error": f"Legacy table not found: {e}", + } + + # ✅ 1. Pre-check: Query schema only once (avoid n+1) + vector_dim: Optional[int] = None + try: + vector_field = legacy_table.schema.field("vector") + list_size = getattr(vector_field.type, "list_size", None) + if list_size is not None: + vector_dim = int(list_size) + except Exception: + vector_dim = None + + if vector_dim is None: + _release_migration_lock(lock_fd) + return { + "success": False, + "source_table": legacy_table_name, + "target_table": primary_table_name, + "rows_migrated": 0, + "error": "Could not determine vector dimension", + } + + # Dimension validation: Ensure target table has compatible dimension + if _table_exists(conn, primary_table_name): + try: + target_table = conn.open_table(primary_table_name) + target_dim = _get_vector_dimension_from_table(target_table) + if target_dim is not None and target_dim != vector_dim: + _release_migration_lock(lock_fd) + return { + "success": False, + "source_table": legacy_table_name, + "target_table": primary_table_name, + "rows_migrated": 0, + "error": f"Dimension mismatch: source={vector_dim}, target={target_dim}", + } + except Exception as e: + logger.warning("Could not validate target table dimension: %s", e) + + # Ensure target table exists (create if needed) + ensure_embeddings_table(conn, to_model_tag(cleaned), vector_dim=vector_dim) + target_table = conn.open_table(primary_table_name) + + # Use merge_insert for idempotent, non-destructive migration + logger.info( + "Starting idempotent migration from '%s' to '%s' (vector_dim=%d, batch_size=%d)", + legacy_table_name, + primary_table_name, + vector_dim, + batch_size, + ) + + # Create merge_insert builder with composite key for uniqueness + # Using doc_id + chunk_id as the natural key for embeddings + merger = target_table.merge_insert(on=["doc_id", "chunk_id"]) + + # Stream data from legacy table using Arrow batches (memory-efficient) + total_rows = legacy_table.count_rows() + logger.info( + f"Streaming {total_rows} rows from legacy table '{legacy_table_name}'" + ) + + batch_num = 0 + for batch in legacy_table.search().to_batches(batch_size=batch_size): + batch_num += 1 + batch_rows = len(batch) + + # Modify model column directly in Arrow (no pandas conversion) + if "model" in batch.schema.names: + new_model_values = pa.array([cleaned] * batch_rows, type=pa.string()) + modified_batch = batch.set_column( + batch.schema.get_field_index("model"), "model", new_model_values + ) + else: + modified_batch = batch + + # Execute merge_insert (idempotent: only inserts if key doesn't exist) + merger.when_not_matched_insert_all().execute(modified_batch) + + rows_migrated += batch_rows + + # Logging (avoid I/O intensive operations) + if batch_num % 10 == 0: + logger.info( + f"Migration progress: {rows_migrated}/{total_rows} rows migrated" + ) + + logger.info( + "Migration completed successfully: '%s' -> '%s' (%d rows processed)", + legacy_table_name, + primary_table_name, + rows_migrated, + ) + logger.info( + "Data has been synced to the new table '%s'. " + "After verifying the migration, you can manually drop the legacy table to free up space: " + "conn.drop_table('%s') or via Python: conn.drop_table('%s')", + primary_table_name, + legacy_table_name, + legacy_table_name, + ) + + return { + "success": True, + "source_table": legacy_table_name, + "target_table": primary_table_name, + "rows_migrated": rows_migrated, + "error": None, + } + + except Exception as e: + logger.error( + "Migration failed for '%s': %s", + primary_table_name, + e, + exc_info=True, + ) + + return { + "success": False, + "source_table": legacy_table_name + if "legacy_table_name" in locals() + else None, + "target_table": primary_table_name, + "rows_migrated": rows_migrated if "rows_migrated" in locals() else 0, + "error": str(e), + } + + finally: + # Release lock + _release_migration_lock(lock_fd) + + +def _table_exists(conn: Any, table_name: str) -> bool: + """Check if a table exists in the database.""" + try: + # Try to get table schema + table_names_fn = getattr(conn, "table_names", None) + if table_names_fn is not None: + table_names = table_names_fn() + return table_name in table_names + else: + # Fallback: try to open the table + conn.open_table(table_name) + return True + except Exception: + return False + + +def _get_vector_dimension_from_table(table: Any) -> Optional[int]: + """Extract vector dimension from table schema. + + Args: + table: LanceDB table object + + Returns: + Vector dimension or None if cannot be determined + """ + try: + schema = table.schema + for field in schema: + if field.name == "vector" and hasattr(field.type, "list_size"): + return int(field.type.list_size) + except Exception as e: + logger.debug("Could not get vector dimension from table: %s", e) + return None + + +def _acquire_migration_lock(db_uri: str, table_name: str) -> Optional[int]: + """Acquire a file lock for migration in the database directory. + + This places the lock file in the database directory itself, which works + for distributed environments where the database is on shared storage (NFS/SMB). + + Args: + db_uri: Database URI (e.g., "/path/to/db" or "s3://bucket/db") + table_name: Name of the table being migrated + + Returns: + File descriptor for the lock, or None if lock is held by another process + """ + # Only support file-based locking for local databases + if db_uri.startswith("s3://") or db_uri.startswith("oss://"): + logger.warning( + "Cloud storage detected (%s), file locking not supported. " + "Consider using distributed locking for concurrent migrations.", + db_uri, + ) + return -1 # Return a dummy fd to avoid errors + + try: + # Create lock file in database directory + lock_dir = db_uri + os.makedirs(lock_dir, exist_ok=True) + + lock_path = os.path.join(lock_dir, f".{table_name}.migration.lock") + lock_fd = os.open(lock_path, os.O_RDWR | os.O_CREAT) + + try: + # Try to acquire exclusive lock (non-blocking) + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + logger.info("Acquired migration lock for '%s' at %s", table_name, lock_path) + return lock_fd + except (IOError, OSError): + # Lock is held by another process + os.close(lock_fd) + logger.info("Migration for '%s' is already in progress", table_name) + return None + except Exception as e: + logger.warning("Failed to acquire migration lock: %s", e) + return -1 # Return a dummy fd to avoid errors + + +def _release_migration_lock(lock_fd: Optional[int]) -> None: + """Release a migration lock. + + Args: + lock_fd: File descriptor from _acquire_migration_lock (or -1/dummy fd) + """ + if lock_fd is not None and lock_fd >= 0: + try: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + os.close(lock_fd) + except Exception: + pass diff --git a/src/xagent/core/tools/core/RAG_tools/utils/model_resolver.py b/src/xagent/core/tools/core/RAG_tools/utils/model_resolver.py index aacc4d4ac..da59dd729 100644 --- a/src/xagent/core/tools/core/RAG_tools/utils/model_resolver.py +++ b/src/xagent/core/tools/core/RAG_tools/utils/model_resolver.py @@ -4,6 +4,7 @@ import logging import os +import sqlite3 from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Type, TypeVar, Union if TYPE_CHECKING: @@ -11,6 +12,7 @@ from langchain_core.runnables import Runnable from sqlalchemy import create_engine +from sqlalchemy.exc import OperationalError as SAOperationalError from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker @@ -44,6 +46,25 @@ _PLACEHOLDER_NONE = {"none", ""} +def _hub_init_failure_is_benign_optional_sqlite(exc: BaseException) -> bool: + """Return True when the hub DB file is missing or not yet creatable. + + In those cases the model hub is an optional component and env-based config + may still work; logging at DEBUG is enough. Permission errors and other DB + failures should surface at WARNING with traceback. + + Args: + exc: Exception raised while initializing SQLAlchemy / SQLite. + + Returns: + True if failure matches a typical \"no sqlite file yet\" operational error. + """ + msg = str(exc).lower() + if "unable to open database file" not in msg: + return False + return isinstance(exc, (SAOperationalError, sqlite3.OperationalError)) + + def _is_placeholder_default(model_id: Optional[str]) -> bool: """Check if model_id is "default" (case-insensitive). @@ -97,7 +118,20 @@ def _get_or_init_model_hub() -> Any: Base.metadata.create_all(engine) return SQLAlchemyModelHub(db, Model) except Exception as e: - logger.debug(f"Model hub database not available: {e}") + if _hub_init_failure_is_benign_optional_sqlite(e): + logger.debug( + "Model hub SQLite not available yet (optional component): %s", + e, + ) + else: + logger.warning( + "Model hub database initialization failed; hub-backed model " + "resolution is disabled until this is fixed. " + "If you rely on env-only configuration, you can ignore this. " + "Otherwise check DB URL, permissions, and connectivity: %s", + e, + exc_info=True, + ) return None diff --git a/src/xagent/core/tools/core/RAG_tools/utils/string_utils.py b/src/xagent/core/tools/core/RAG_tools/utils/string_utils.py index 79442730c..8123cd507 100644 --- a/src/xagent/core/tools/core/RAG_tools/utils/string_utils.py +++ b/src/xagent/core/tools/core/RAG_tools/utils/string_utils.py @@ -6,7 +6,7 @@ import re import uuid from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional # Pattern for sanitizing document IDs and filenames # Only allows: letters, numbers, underscore, hyphen @@ -32,21 +32,61 @@ def escape_lancedb_string(input_string: Any) -> str: return input_string.replace("\\", "\\\\").replace("'", "''") -def build_lancedb_filter_expression(filters: Dict[str, Any]) -> str: +def build_lancedb_filter_expression( + filters: Dict[str, Any], + *, + user_id: Optional[int] = None, + is_admin: bool = False, + skip_user_filter: bool = False, +) -> str: """ Builds a safe LanceDB filter expression from a dictionary of filters. + This function now uses the abstract filter layer internally for better + backend compatibility, while maintaining the same interface for + backward compatibility. + Args: filters: A dictionary where keys are column names and values are the filter values. + user_id: Optional user ID for multi-tenancy filtering. + is_admin: Whether the user has admin privileges. + skip_user_filter: If True, bypasses user permission filter. Returns: A string representing the safely constructed LanceDB filter expression. """ - filter_parts = [] + from ..storage.contracts import ( + FilterCondition, + FilterExpression, + FilterOperator, + ) + from ..storage.factory import get_vector_index_store + + # Convert to FilterCondition list + conditions: List[FilterCondition] = [] for key, value in filters.items(): - escaped_value = escape_lancedb_string(value) - filter_parts.append(f"{key} == '{escaped_value}'") - return " AND ".join(filter_parts) + conditions.append( + FilterCondition(field=key, operator=FilterOperator.EQ, value=value) + ) + + # Use abstract filter builder + vector_store = get_vector_index_store() + + # Combine conditions with AND (tuple convention) + # Type: FilterExpression can be FilterCondition or tuple of FilterConditions + if len(conditions) == 1: + filter_expr: FilterExpression = conditions[0] + else: + filter_expr = tuple(conditions) + + # Get backend-specific syntax + backend_filter = vector_store.build_filter_expression( + filters=filter_expr, + user_id=user_id if not skip_user_filter else None, + is_admin=is_admin or skip_user_filter, + ) + + return backend_filter or "" def sanitize_for_doc_id(text: str, max_length: int = 64) -> str: diff --git a/src/xagent/core/tools/core/RAG_tools/utils/tag_mapping.py b/src/xagent/core/tools/core/RAG_tools/utils/tag_mapping.py new file mode 100644 index 000000000..73fcda41a --- /dev/null +++ b/src/xagent/core/tools/core/RAG_tools/utils/tag_mapping.py @@ -0,0 +1,37 @@ +"""Helpers for collision-aware tag mapping registration.""" + +from __future__ import annotations + +import logging +from typing import Callable, Dict, TypeVar + +ValueT = TypeVar("ValueT") + + +def register_tag_mapping( + mapping: Dict[str, ValueT], + tag: str, + value: ValueT, + *, + get_identity: Callable[[ValueT], str], + logger: logging.Logger, +) -> None: + """Register a normalized tag mapping and warn on identity collisions. + + Args: + mapping: Destination mapping keyed by normalized tag. + tag: Normalized tag key. + value: Value to store for the tag. + get_identity: Function returning the logical identity used to detect + collisions. For example, for ``tuple[str, Optional[int]]`` values it + can return the first element (Hub model ID). + logger: Logger used to emit collision warnings. + """ + existing = mapping.get(tag) + if existing is not None: + existing_id = get_identity(existing) + value_id = get_identity(value) + if existing_id != value_id: + logger.warning("Tag collision: %s -> %s vs %s", tag, existing_id, value_id) + return + mapping[tag] = value diff --git a/src/xagent/core/tools/core/RAG_tools/utils/user_permissions.py b/src/xagent/core/tools/core/RAG_tools/utils/user_permissions.py index f5e66f112..29a821d95 100644 --- a/src/xagent/core/tools/core/RAG_tools/utils/user_permissions.py +++ b/src/xagent/core/tools/core/RAG_tools/utils/user_permissions.py @@ -8,6 +8,16 @@ class UserPermissions: """Handle user permissions and data access control.""" + @staticmethod + def get_no_access_filter() -> str: + """Return a stable LanceDB filter expression that always matches no rows.""" + return UNAUTHENTICATED_NO_ACCESS_FILTER + + @staticmethod + def is_no_access_filter(filter_expr: Optional[str]) -> bool: + """Check whether a filter expression is the internal no-access marker.""" + return filter_expr == UNAUTHENTICATED_NO_ACCESS_FILTER + @staticmethod def get_user_filter( user_id: Optional[int], is_admin: bool = False @@ -37,16 +47,6 @@ def get_user_filter( # Unauthenticated users cannot see any data return UserPermissions.get_no_access_filter() - @staticmethod - def get_no_access_filter() -> str: - """Return a stable LanceDB filter expression that always matches no rows.""" - return UNAUTHENTICATED_NO_ACCESS_FILTER - - @staticmethod - def is_no_access_filter(filter_expr: Optional[str]) -> bool: - """Check whether a filter expression is the internal no-access marker.""" - return filter_expr == UNAUTHENTICATED_NO_ACCESS_FILTER - @staticmethod def can_access_data( user_id: Optional[int], data_user_id: Optional[int], is_admin: bool = False diff --git a/src/xagent/core/tools/core/RAG_tools/vector_storage/__init__.py b/src/xagent/core/tools/core/RAG_tools/vector_storage/__init__.py index 792001be3..1380ddb32 100644 --- a/src/xagent/core/tools/core/RAG_tools/vector_storage/__init__.py +++ b/src/xagent/core/tools/core/RAG_tools/vector_storage/__init__.py @@ -10,6 +10,9 @@ This module handles only data management and does not perform any text-to-vector conversion. The actual embedding is handled by AgentOS embedding nodes in workflows. +Index management is now handled by the storage abstraction layer in +`storage.contracts.VectorIndexStore` and implemented in `storage.lancedb_stores`. + ``` AgentOS Workflow: 1. read_chunks_for_embedding() → Get chunks needing vectors @@ -27,11 +30,10 @@ - Automatic dimension consistency checking - Stale data cleanup when chunk_hash changes -- HNSW index creation when row threshold is met +- Index creation handled by storage abstraction layer - Multi-model support with separate tables per model """ -from .index_manager import get_index_manager from .vector_manager import ( read_chunks_for_embedding, validate_query_vector, @@ -39,7 +41,6 @@ ) __all__ = [ - "get_index_manager", "read_chunks_for_embedding", "write_vectors_to_db", "validate_query_vector", diff --git a/src/xagent/core/tools/core/RAG_tools/vector_storage/index_manager.py b/src/xagent/core/tools/core/RAG_tools/vector_storage/index_manager.py deleted file mode 100644 index 6202c3474..000000000 --- a/src/xagent/core/tools/core/RAG_tools/vector_storage/index_manager.py +++ /dev/null @@ -1,254 +0,0 @@ -""" -Index management for vector storage. - -This module provides centralized index management functionality for embeddings tables, -including index creation, status checking, and basic maintenance operations. -""" - -import logging -from typing import Any, Dict, Optional, Tuple - -from ..core.config import IndexPolicy -from ..core.schemas import IndexType - -# Import LanceDB index types -try: - from lancedb.index import IVF_HNSW_SQ, IVF_PQ # type: ignore -except ImportError: - # Fallback if import fails - IVF_HNSW_SQ = "IVF_HNSW_SQ" - IVF_PQ = "IVF_PQ" - -logger = logging.getLogger(__name__) - - -class IndexManager: - """ - Centralized index manager for embeddings tables. - - This class handles index lifecycle management including creation, - status checking, and basic maintenance operations. - """ - - def __init__(self, policy: Optional[IndexPolicy] = None): - """ - Initialize index manager with policy configuration. - - Args: - policy: Index policy configuration, uses default if None - """ - self.policy = policy or IndexPolicy() - - def check_and_create_index( - self, - table: Any, - table_name: str, - readonly: bool = False, - ) -> Tuple[str, Optional[str]]: - """ - Check table index status and create index if needed. - - Automatically selects index type based on row count: - - < 50k rows: No index - - 50k-10M rows: HNSW index - - >= 10M rows: IVFPQ index - - Args: - table: LanceDB table instance - table_name: Table name for logging - readonly: If True, don't create indexes - - Returns: - Tuple of (index_status, index_advice) - """ - if readonly: - return "readonly", f"Readonly mode - no index operations for {table_name}" - - vector_index_status: str = "no_index" - vector_index_advice: Optional[str] = None - - try: - # Get row count efficiently without loading all data into memory - row_count = table.count_rows() - - if row_count < self.policy.enable_threshold_rows: - vector_index_status = "below_threshold" - vector_index_advice = ( - f"Table {table_name} has {row_count} rows - below threshold " - f"({self.policy.enable_threshold_rows}) for index creation" - ) - else: - # Auto-select index type based on scale - if row_count >= self.policy.ivfpq_threshold_rows: - recommended_type = IndexType.IVFPQ - else: - recommended_type = IndexType.HNSW - - # Check existing indexes - indexes = table.list_indices() - has_vector_index = any(idx.name == "vector" for idx in indexes) - - if not has_vector_index: - # Create index with recommended type and parameters - if recommended_type == IndexType.IVFPQ: - index_type = IVF_PQ - create_params = self.policy.ivfpq_params or {} - else: # HNSW - index_type = IVF_HNSW_SQ - create_params = self.policy.hnsw_params or {} - - # Merge metric with create_params, avoiding duplicates - all_params = { - "metric": self.policy.metric.value, - "index_type": index_type, - **create_params, - } - - table.create_index(**all_params) - vector_index_status = "index_building" - logger.info( - "Successfully created vector index for %s (type=%s, metric=%s)", - table_name, - index_type, - self.policy.metric.value, - ) - if recommended_type == IndexType.IVFPQ: - vector_index_advice = ( - f"IVFPQ index created for {table_name} " - f"({row_count} rows, using IVFPQ strategy for large-scale data), metric: {self.policy.metric.value}" - ) - else: # HNSW - vector_index_advice = ( - f"HNSW index created for {table_name} " - f"({row_count} rows, using HNSW strategy for medium-scale data), metric: {self.policy.metric.value}" - ) - else: - vector_index_status = "index_ready" - vector_index_advice = f"Index ready for {table_name} ({row_count} rows), metric: {self.policy.metric.value}" - - except Exception as e: - logger.error(f"Vector index operation failed for {table_name}: {str(e)}") - vector_index_status = "index_corrupted" - vector_index_advice = ( - f"Vector index check failed for {table_name}: {str(e)}" - ) - - # FTS Index Management - if self.policy.fts_enabled: - fts_success, fts_message = self.create_fts_index( - table, table_name, self.policy.fts_params - ) - if not fts_success: - logger.warning( - f"FTS index creation/check failed for {table_name}: {fts_message}" - ) - # If FTS index fails, it does not necessarily corrupt the vector index - # but we should reflect the partial failure or warning. - # For now, we will log and return vector index status primarily. - - return vector_index_status, vector_index_advice - - def get_index_status(self, table: Any) -> str: - """ - Get current index status for a table. - - Args: - table: LanceDB table instance - - Returns: - Index status string - """ - try: - indexes = table.list_indices() - has_vector_index = any(idx.name == "vector" for idx in indexes) - - if has_vector_index: - return "index_ready" - else: - row_count = table.count_rows() - if row_count >= self.policy.enable_threshold_rows: - return "no_index" - else: - return "below_threshold" - except Exception as e: - logger.error(f"Failed to get index status: {str(e)}") - return "index_corrupted" - - def get_fts_index_status(self, table: Any) -> bool: - """ - Check if a Full-Text Search (FTS) index exists on the 'text' column of the table. - - Args: - table: LanceDB table instance. - - Returns: - True if an FTS index exists on the 'text' column, False otherwise. - """ - try: - indexes = table.list_indices() - # New lancedb versions return IndexConfig objects, not dicts. - # Access properties via attributes. - return any( - idx.index_type == "FTS" and "text" in idx.columns for idx in indexes - ) - except Exception as e: - logger.error(f"Failed to check FTS index status for {table.name}: {str(e)}") - return False - - def create_fts_index( - self, - table: Any, - table_name: str, - fts_params: Optional[Dict[str, Any]] = None, - ) -> Tuple[bool, Optional[str]]: - """ - Create a Full-Text Search (FTS) index on the 'text' column. - - Args: - table: LanceDB table instance. - table_name: Name of the table for logging. - fts_params: Optional dictionary of FTS parameters (e.g., language, stem, ascii_folding, with_position). - - Returns: - Tuple of (success: bool, message: Optional[str]). - """ - if self.get_fts_index_status(table): - return True, f"FTS index already exists on 'text' column for {table_name}" - - try: - # Default FTS parameters, can be overridden by fts_params - _fts_params = {"with_position": True, **(fts_params or {})} - # Add replace=True to make the operation idempotent - table.create_fts_index("text", replace=True, **_fts_params) - logger.info( - "Successfully created FTS index on 'text' column for %s", table_name - ) - return ( - True, - f"FTS index created on 'text' column for {table_name} with params: {_fts_params}", - ) - except Exception as e: - logger.error(f"Failed to create FTS index for {table_name}: {str(e)}") - return False, f"Failed to create FTS index: {str(e)}" - - -# Global index manager instance -_default_index_manager: Optional[IndexManager] = None - - -def get_index_manager(policy: Optional[IndexPolicy] = None) -> IndexManager: - """ - Get the global index manager instance. - - Args: - policy: Optional policy to configure the manager - - Returns: - IndexManager instance - """ - global _default_index_manager - - if _default_index_manager is None or (policy is not None): - _default_index_manager = IndexManager(policy) - - return _default_index_manager diff --git a/src/xagent/core/tools/core/RAG_tools/vector_storage/vector_manager.py b/src/xagent/core/tools/core/RAG_tools/vector_storage/vector_manager.py index a8cb68bdb..7af089111 100644 --- a/src/xagent/core/tools/core/RAG_tools/vector_storage/vector_manager.py +++ b/src/xagent/core/tools/core/RAG_tools/vector_storage/vector_manager.py @@ -17,10 +17,13 @@ import time from typing import Any, Dict, List, Optional, cast +import numpy as np import pandas as pd -from ......providers.vector_store.lancedb import get_connection_from_env -from ..core.config import DEFAULT_LANCEDB_BATCH_DELAY_MS, IndexPolicy +from ..core.config import ( + DEFAULT_LANCEDB_BATCH_DELAY_MS, + DEFAULT_LANCEDB_BATCH_SIZE, +) from ..core.exceptions import ( ConfigurationError, DatabaseOperationError, @@ -35,12 +38,9 @@ IndexOperation, ) from ..LanceDB.model_tag_utils import to_model_tag -from ..LanceDB.schema_manager import ensure_chunks_table, ensure_embeddings_table -from ..utils.lancedb_query_utils import query_to_list +from ..LanceDB.schema_manager import ensure_embeddings_table +from ..storage.factory import get_vector_index_store from ..utils.metadata_utils import deserialize_metadata, serialize_metadata -from ..utils.string_utils import build_lancedb_filter_expression -from ..utils.user_permissions import UserPermissions -from .index_manager import get_index_manager logger = logging.getLogger(__name__) @@ -111,61 +111,6 @@ def _is_non_recoverable_merge_error(error: Exception) -> bool: return is_non_recoverable -def _should_reindex( - table: Any, - table_name: str, - total_upserted: int, - policy: IndexPolicy, -) -> bool: - """Determine if reindex should be triggered. - - Args: - table: LanceDB table instance - table_name: Table name for tracking - total_upserted: Number of rows upserted in this operation - policy: Index policy configuration - - Returns: - True if reindex should be triggered - """ - # Immediate reindex if enabled - if policy.enable_immediate_reindex and total_upserted > 0: - return True - - # Batch size threshold - if total_upserted >= policy.reindex_batch_size: - return True - - # Smart reindex: check unindexed ratio - if policy.enable_smart_reindex: - try: - stats = table.index_stats("vector_idx") - if stats.num_indexed_rows > 0: - unindexed_ratio = stats.num_unindexed_rows / stats.num_indexed_rows - if unindexed_ratio > policy.reindex_unindexed_ratio_threshold: - return True - - # Absolute threshold for unindexed rows - if stats.num_unindexed_rows > 10000: - return True - except Exception as e: # noqa: BLE001 - logger.debug("Could not get index stats for %s: %s", table_name, e) - - return False - - -def _trigger_reindex(table: Any, table_name: str) -> bool: - """Trigger reindex operation on the table.""" - try: - logger.info("Triggering reindex for %s", table_name) - table.optimize() - logger.info("Reindex completed for %s", table_name) - return True - except Exception as e: # noqa: BLE001 - logger.warning("Reindex failed for %s: %s", table_name, e) - return False - - def validate_query_vector( query_vector: List[float], model_tag: Optional[str] = None, @@ -175,12 +120,19 @@ def validate_query_vector( ) -> None: """Validate query vector format and content. + This function performs basic validation of the query vector without + requiring database access. Dimension validation is handled by the + storage abstraction layer during search operations. + Args: query_vector: Query vector to validate - model_tag: Optional model tag for dimension validation - conn: Optional LanceDB connection for validation - user_id: Optional user ID for filtering (for multi-tenancy) - is_admin: Whether user has admin privileges + model_tag: Optional model tag (for logging purposes only) + conn: Deprecated - no longer used + user_id: Deprecated - no longer used + is_admin: Deprecated - no longer used + + Raises: + VectorValidationError: If vector validation fails """ if not isinstance(query_vector, list): raise VectorValidationError("query_vector must be a list") @@ -203,130 +155,49 @@ def validate_query_vector( "query_vector contains invalid values (NaN or infinity)" ) - if model_tag and conn: - # First validate model_tag format and table existence - normalized_model_tag = to_model_tag(model_tag) - validate_embed_model(conn, normalized_model_tag) - - table_name = f"embeddings_{normalized_model_tag}" - try: - table = conn.open_table(table_name) - expected_dim = None - - # Method 1: Try to get dimension from schema (for fixed-size vector columns) - try: - vector_field = table.schema.field("vector") - # Safely check if list_size attribute exists (fixed-size list) - list_size = getattr(vector_field.type, "list_size", None) - if list_size is not None: - expected_dim = list_size - except (AttributeError, KeyError) as e: - logger.debug( - "Could not get vector dimension from schema for %s: %s. Will try to infer from data.", - table_name, - e, - ) - - # Method 2: If schema doesn't have fixed dimension, infer from actual data - if expected_dim is None: - expected_dim = get_stored_vector_dimension( - conn, model_tag, user_id, is_admin - ) - # Perform dimension validation if we got a dimension - if expected_dim is not None: - if len(query_vector) != expected_dim: - raise VectorValidationError( - f"Query vector dimension {len(query_vector)} does not match stored dimension {expected_dim} for model '{model_tag}'" - ) - else: - logger.warning( - "Could not determine expected vector dimension for %s " - "(table may be empty or schema is variable-length). " - "Skipping dimension consistency check.", - table_name, - ) - except VectorValidationError: - # Re-raise validation errors (don't catch them) - raise - except Exception as e: # noqa: BLE001 - logger.warning( - "Failed to perform dimension validation for %s: %s. Skipping dimension consistency check.", - table_name, - e, - ) +def _safe_int_conversion(value: Any, default: int = 0) -> int: + """Safely convert value to int, handling None and NaN. + Args: + value: Value to convert (can be None, NaN, int, float, etc.) + default: Default value if conversion fails -def validate_embed_model(conn: Any, model_tag: str) -> None: - """Validate embed model exists and is accessible.""" - import re + Returns: + Integer value, or default if value is None/NaN/not convertible + """ + """Safely convert value to int, handling None and NaN. - # Validate model_tag format (cannot contain characters that affect table name) - if not re.match(r"^[a-zA-Z0-9_-]+$", model_tag): - raise VectorValidationError( - f"Invalid model_tag format: {model_tag}. Only alphanumeric, underscore, and hyphen allowed." - ) + Args: + value: Value to convert (can be None, NaN, int, float, etc.) + default: Default value if conversion fails - # Validate that the corresponding table exists - table_name = f"embeddings_{model_tag}" + Returns: + Integer value, or default if value is None/NaN/not convertible + """ + if value is None or (isinstance(value, float) and np.isnan(value)): + return default try: - conn.open_table(table_name) - except Exception as e: # noqa: BLE001 - logger.warning( - "Embeddings table %s for model %s not found or inaccessible: %s", - table_name, - model_tag, - e, - ) - raise VectorValidationError( - f"Embeddings table for model '{model_tag}' does not exist or is inaccessible: {str(e)}" - ) from e + return int(value) + except (ValueError, TypeError): + return default -def get_stored_vector_dimension( - conn: Any, - model_tag: str, - user_id: Optional[int] = None, - is_admin: bool = False, -) -> Optional[int]: - """Get the vector dimension for a model from database. +def _safe_str_value(value: Any) -> Optional[str]: + """Extract string value, returning None for NaN/None values. + + This handles pandas DataFrame's NaN preservation behavior where + NaN values are not automatically converted to None. Args: - conn: LanceDB connection - model_tag: Model tag to look up - user_id: Optional user ID for filtering (for multi-tenancy) - is_admin: Whether user has admin privileges + value: Value from pandas DataFrame (can be str, None, or NaN) Returns: - Vector dimension if found, None otherwise + String value, or None if value is None/NaN """ - try: - normalized_model_tag = to_model_tag(model_tag) - table_name = f"embeddings_{normalized_model_tag}" - table = conn.open_table(table_name) - - # Apply user filter for multi-tenancy - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - - # Query one record to get dimension, with optional user filtering - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - if user_filter_expr: - sample_list = query_to_list(table.search().where(user_filter_expr).limit(1)) - else: - sample_list = query_to_list(table.head(1)) - - if sample_list: - vector_dim = sample_list[0].get("vector_dimension") - if vector_dim is not None: - return int(vector_dim) - except Exception as e: # noqa: BLE001 - logger.debug( - "Could not get stored vector dimension for %s: %s. This is expected if the table is new or empty.", - model_tag, - e, - ) - pass - return None + if value is None or (isinstance(value, float) and np.isnan(value)): + return None + return str(value) if value is not None else None def read_chunks_for_embedding( @@ -338,7 +209,10 @@ def read_chunks_for_embedding( user_id: Optional[int] = None, is_admin: bool = False, ) -> EmbeddingReadResponse: - """Read chunks from database for embedding computation.""" + """Read chunks from database for embedding computation. + + Phase 1A: Refactored to use storage abstraction layer instead of raw connection. + """ try: # Validate inputs if not collection or not doc_id or not parse_hash or not model: @@ -354,10 +228,8 @@ def read_chunks_for_embedding( model, ) - # Get database connection - conn = get_connection_from_env() - - ensure_chunks_table(conn) + # Use storage abstraction instead of raw connection + vector_store = get_vector_index_store() # Build query filters query_filters: Dict[str, Any] = { @@ -370,95 +242,69 @@ def read_chunks_for_embedding( if filters: query_filters.update(filters) - # Read chunks from database - chunks_table = conn.open_table("chunks") - - # Build combined filter expression with user permissions - base_filter_expr = build_lancedb_filter_expression(query_filters) - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - - if user_filter_expr and base_filter_expr: - filter_expr = f"({base_filter_expr}) and ({user_filter_expr})" - elif user_filter_expr: - filter_expr = user_filter_expr - else: - filter_expr = base_filter_expr - - try: - # OPTIMIZATION: Use count_rows() for memory-efficient counting - total_count = chunks_table.count_rows(filter_expr) - if total_count == 0: - logger.info("No chunks found for the given criteria") - return EmbeddingReadResponse(chunks=[], total_count=0, pending_count=0) - - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - chunks_data = query_to_list(chunks_table.search().where(filter_expr)) - except Exception as e: # noqa: BLE001 - logger.error("Failed to read chunks for embedding: %s", e) - raise DatabaseOperationError( - f"Failed to read chunks for embedding: {e}" - ) from e + # Use abstraction layer for counting (returns 0 if table doesn't exist) + total_count = vector_store.count_rows_or_zero( + table_name="chunks", + filters=query_filters, + user_id=user_id, + is_admin=is_admin, + ) + if total_count == 0: + logger.info("No chunks found for the given criteria") + return EmbeddingReadResponse(chunks=[], total_count=0, pending_count=0) + + # Use abstraction layer for batch iteration + chunks_data = [] + for batch in vector_store.iter_batches( + table_name="chunks", + columns=None, # Select all columns + batch_size=1000, + filters=query_filters, + user_id=user_id, + is_admin=is_admin, + ): + batch_df = batch.to_pandas() + for _, row in batch_df.iterrows(): + chunks_data.append(row.to_dict()) + if len(chunks_data) >= total_count: + break - # Check which chunks already have embeddings + # Check which chunks already have embeddings using abstraction layer embedded_chunk_ids = set() model_tag = to_model_tag(model) embeddings_table_name = f"embeddings_{model_tag}" try: - # Get vector dimension from collection metadata or model config - vector_dim = None - try: - from ..management.collection_manager import get_collection_sync - - coll_info = get_collection_sync(collection) - vector_dim = coll_info.embedding_dimension - except Exception: - # Fallback to resolving the model config - from ..utils.model_resolver import resolve_embedding_adapter - - embedding_config, _ = resolve_embedding_adapter(model) - vector_dim = embedding_config.dimension - - ensure_embeddings_table(conn, model_tag, vector_dim=vector_dim) - embeddings_table = conn.open_table(embeddings_table_name) - # Get existing embeddings for these chunks # Only select chunk_id column to avoid loading unnecessary vector data - embedding_filters = { + embedding_filters: Dict[str, Any] = { "collection": collection, "doc_id": doc_id, "parse_hash": parse_hash, - "model": model, } - base_embedding_filter_expr = build_lancedb_filter_expression( - embedding_filters - ) - - # Add user permission filter for multi-tenancy - user_filter_expr = UserPermissions.get_user_filter(user_id, is_admin) - - # Combine filters - if user_filter_expr and base_embedding_filter_expr: - embedding_filter_expr = ( - f"({base_embedding_filter_expr}) and ({user_filter_expr})" - ) - elif user_filter_expr: - embedding_filter_expr = user_filter_expr - else: - embedding_filter_expr = base_embedding_filter_expr - # OPTIMIZATION: Use unified query_to_list() with three-tier fallback - embeddings_data = query_to_list( - embeddings_table.search() - .where(embedding_filter_expr) - .select(["chunk_id"]) + # Use abstraction layer to query embeddings (returns 0 if table doesn't exist) + # Note: We don't filter by 'model' field as it's not in current schema + embedding_count = vector_store.count_rows_or_zero( + table_name=embeddings_table_name, + filters=embedding_filters, + user_id=user_id, + is_admin=is_admin, ) - # Filter out None values (from NaN normalization) - embedded_chunk_ids = { - item["chunk_id"] - for item in embeddings_data - if item.get("chunk_id") is not None - } + + if embedding_count > 0: + # Read chunk_ids from embeddings table + for batch in vector_store.iter_batches( + table_name=embeddings_table_name, + columns=["chunk_id"], + filters=embedding_filters, + user_id=user_id, + is_admin=is_admin, + ): + batch_df = batch.to_pandas() + for chunk_id in batch_df["chunk_id"]: + if chunk_id is not None: + embedded_chunk_ids.add(chunk_id) except Exception as e: # noqa: BLE001 # If embeddings table doesn't exist or query fails, assume no embeddings exist @@ -477,22 +323,22 @@ def read_chunks_for_embedding( # Deserialize metadata from JSON string to dictionary metadata = deserialize_metadata(chunk_dict.get("metadata")) - # Arrow/to_list() returns None instead of NaN, so direct None check is sufficient - index_value = chunk_dict.get("index") - index = int(index_value) if index_value is not None else 0 + # Handle index with NaN-safe conversion + index = _safe_int_conversion(chunk_dict.get("index"), default=0) page_number_value = chunk_dict.get("page_number") # Convert to int only if valid and > 0 (schema requires gt=0) if page_number_value is not None: - page_num = int(page_number_value) + page_num = _safe_int_conversion(page_number_value, default=1) page_number = page_num if page_num > 0 else None else: page_number = None - # Normalize optional string fields: Arrow/to_list() returns None, not NaN - section = chunk_dict.get("section") - anchor = chunk_dict.get("anchor") - json_path = chunk_dict.get("json_path") + # Normalize optional string fields using NaN-safe helper + # pandas to_pandas() preserves NaN values, so explicit NaN handling needed + section = _safe_str_value(chunk_dict.get("section")) + anchor = _safe_str_value(chunk_dict.get("anchor")) + json_path = _safe_str_value(chunk_dict.get("json_path")) chunk = ChunkForEmbedding( doc_id=chunk_dict["doc_id"], @@ -706,18 +552,18 @@ def _process_batch( def _process_model_embeddings( - conn: Any, collection: str, model: str, model_embeddings: List[ChunkEmbeddingData], create_index: bool, user_id: Optional[int] = None, ) -> tuple[int, str]: - """Process embeddings for a single model. + """Process embeddings for a single model using abstraction layer. Returns: Tuple of (upserted_count, index_status) """ + model_tag = to_model_tag(model) table_name = f"embeddings_{model_tag}" @@ -749,13 +595,10 @@ def _process_model_embeddings( vector_dim, ) - # Prepare table - embeddings_table = _validate_and_prepare_table( - conn, model_tag, table_name, vector_dim - ) - # Process embeddings in batches to prevent memory issues and LanceDB spills - original_batch_size = int(os.getenv("LANCEDB_BATCH_SIZE", "1000")) + original_batch_size = int( + os.getenv("LANCEDB_BATCH_SIZE", str(DEFAULT_LANCEDB_BATCH_SIZE)) + ) batch_size = original_batch_size total_batches_for_logging = ( len(model_embeddings) + original_batch_size - 1 @@ -779,6 +622,8 @@ def _process_model_embeddings( max_spill_retries = int(os.getenv("LANCEDB_MAX_SPILL_RETRIES", "3")) spill_retry_count = 0 + vector_store = get_vector_index_store() + while current_idx < total_embeddings: end_idx = min(current_idx + batch_size, total_embeddings) batch_embeddings = model_embeddings[current_idx:end_idx] @@ -805,17 +650,21 @@ def _process_model_embeddings( try: batch_idx_for_logging = current_idx // original_batch_size - batch_upserted = _process_batch( - embeddings_table, - records_to_merge, - batch_idx_for_logging, - total_batches_for_logging, - model, - ) + # Use abstraction layer for upsert (includes fallback logic) + vector_store.upsert_embeddings(model_tag, records_to_merge) + batch_upserted = len(records_to_merge) upserted_count += batch_upserted current_idx = end_idx # Move to next batch on success spill_retry_count = 0 # Reset after a successful batch + logger.info( + "Successfully processed batch %d/%d (%d embeddings) for model %s", + batch_idx_for_logging + 1, + total_batches_for_logging, + batch_upserted, + model, + ) + except Exception as batch_error: # noqa: BLE001 failed_batches += 1 logger.error( @@ -881,26 +730,16 @@ def _process_model_embeddings( logger.info("Processed model %s: upserted %d embeddings", model, upserted_count) - # Handle index creation and reindexing if requested + # Handle index creation using abstraction layer index_status: str = IndexOperation.SKIPPED.value if create_index: try: - # Use index manager for index creation - index_manager = get_index_manager() - status, _ = index_manager.check_and_create_index( - embeddings_table, table_name, readonly=False - ) - index_status = status - - # Trigger reindex if needed - policy = IndexPolicy() - if _should_reindex(embeddings_table, table_name, upserted_count, policy): - reindex_success = _trigger_reindex(embeddings_table, table_name) - if reindex_success: - logger.info("Reindex triggered for %s", table_name) - else: - logger.warning("Reindex failed for %s", table_name) + from ..core.schemas import IndexResult + index_result_obj: IndexResult = vector_store.create_index( + model_tag, readonly=False + ) + index_status = index_result_obj.status except Exception as index_error: # noqa: BLE001 logger.warning("Failed to create index for %s: %s", table_name, index_error) index_status = IndexOperation.FAILED.value @@ -933,18 +772,15 @@ def write_vectors_to_db( total_upserted = 0 index_statuses = [] - # Get database connection - conn = get_connection_from_env() - - # Process each model separately + # Process each model separately (abstraction layer handles connection internally) for model, model_embeddings in embeddings_by_model.items(): upserted, idx_status = _process_model_embeddings( - conn, collection, model, model_embeddings, create_index, user_id + collection, model, model_embeddings, create_index, user_id ) total_upserted += upserted index_statuses.append(idx_status) - # Determine overall index status (map index_manager strings to IndexOperation) + # Determine overall index status (map create_index result strings to IndexOperation) if "index_building" in index_statuses: overall_index_status = IndexOperation.CREATED elif "index_ready" in index_statuses: diff --git a/src/xagent/core/tools/core/RAG_tools/version_management/cascade_cleaner.py b/src/xagent/core/tools/core/RAG_tools/version_management/cascade_cleaner.py index 9cd3032dc..feb4b15a2 100644 --- a/src/xagent/core/tools/core/RAG_tools/version_management/cascade_cleaner.py +++ b/src/xagent/core/tools/core/RAG_tools/version_management/cascade_cleaner.py @@ -9,7 +9,6 @@ import logging from typing import Any, Dict, Optional -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.exceptions import CascadeCleanupError from ..LanceDB.schema_manager import ( ensure_chunks_table, @@ -17,6 +16,7 @@ ensure_main_pointers_table, ensure_parses_table, ) +from ..storage.factory import get_vector_store_raw_connection from ..utils.string_utils import build_lancedb_filter_expression, escape_lancedb_string from .main_pointer_manager import get_main_pointer @@ -168,7 +168,7 @@ def cleanup_cascade( Returns: Deleted (or planned) counts per table scope """ - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() ensure_documents_table(conn) ensure_parses_table(conn) ensure_chunks_table(conn) diff --git a/src/xagent/core/tools/core/RAG_tools/version_management/list_candidates.py b/src/xagent/core/tools/core/RAG_tools/version_management/list_candidates.py index 99e0fbb95..8fc2964b4 100644 --- a/src/xagent/core/tools/core/RAG_tools/version_management/list_candidates.py +++ b/src/xagent/core/tools/core/RAG_tools/version_management/list_candidates.py @@ -9,9 +9,9 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Union -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.exceptions import DatabaseOperationError, VersionManagementError from ..core.schemas import StepType +from ..storage.factory import get_vector_store_raw_connection from ..utils.lancedb_query_utils import query_to_list from ..utils.string_utils import build_lancedb_filter_expression @@ -344,7 +344,7 @@ def list_candidates( resolved_step_type = _resolve_step_type(step_type) try: # Get LanceDB connection from environment (uses default path if LANCEDB_DIR not set) - connection = get_connection_from_env() + connection = get_vector_store_raw_connection() # Get candidates based on step_type candidates = _get_candidates( diff --git a/src/xagent/core/tools/core/RAG_tools/version_management/main_pointer_manager.py b/src/xagent/core/tools/core/RAG_tools/version_management/main_pointer_manager.py index 8721fcd84..24d6b2550 100644 --- a/src/xagent/core/tools/core/RAG_tools/version_management/main_pointer_manager.py +++ b/src/xagent/core/tools/core/RAG_tools/version_management/main_pointer_manager.py @@ -2,6 +2,8 @@ This module provides functionality for managing main version pointers across different processing stages (parse, chunk, embed). + +Phase 1A Part 2: Refactored to use MainPointerStore abstraction layer. """ from __future__ import annotations @@ -9,12 +11,8 @@ import logging from typing import Any, Dict, List, Optional -import pandas as pd - -from ......providers.vector_store.lancedb import get_connection_from_env from ..core.exceptions import MainPointerError -from ..LanceDB.schema_manager import ensure_main_pointers_table -from ..utils.string_utils import build_lancedb_filter_expression, escape_lancedb_string +from ..storage.factory import get_main_pointer_store logger = logging.getLogger(__name__) @@ -24,27 +22,6 @@ def _normalize_model_tag(model_tag: Optional[str]) -> str: return model_tag if model_tag is not None else "" -def _build_base_filter_expression(collection: str, doc_id: str, step_type: str) -> str: - """Build the base LanceDB filter expression for a main pointer row. - - This helper escapes all string values to avoid malformed expressions and - injection-like issues. - - Args: - collection: Collection name. - doc_id: Document ID. - step_type: Processing stage type (parse, chunk, embed). - - Returns: - A filter expression covering collection/doc_id/step_type. - """ - return ( - f"collection == '{escape_lancedb_string(collection)}' AND " - f"doc_id == '{escape_lancedb_string(doc_id)}' AND " - f"step_type == '{escape_lancedb_string(step_type)}'" - ) - - def get_main_pointer( collection: str, doc_id: str, step_type: str, model_tag: Optional[str] = None ) -> Optional[Dict[str, Any]]: @@ -63,46 +40,14 @@ def get_main_pointer( MainPointerError: If there's an error retrieving the pointer """ try: - conn = get_connection_from_env() - ensure_main_pointers_table(conn) - - table = conn.open_table("main_pointers") - - # Build safe filter conditions - normalized_tag = _normalize_model_tag(model_tag) - - # Base filters for collection, doc_id, and step_type - base_expr = _build_base_filter_expression(collection, doc_id, step_type) - - # Handle model_tag: check for both normalized empty string AND NULL for backward compatibility - if normalized_tag == "": - filter_expr = f"{base_expr} AND (model_tag == '' OR model_tag IS NULL)" - else: - filter_expr = f"{base_expr} AND model_tag == '{escape_lancedb_string(normalized_tag)}'" - - # Query the table - result = table.search().where(filter_expr).to_pandas() - - if result.empty: - return None - - # Return the first result, preferring non-NULL model_tag if multiple found - if len(result) > 1: - result = result.sort_values("model_tag", ascending=False) - - row = result.iloc[0] - return { - "collection": row["collection"], - "doc_id": row["doc_id"], - "step_type": row["step_type"], - "model_tag": row["model_tag"] if row["model_tag"] is not None else "", - "semantic_id": row["semantic_id"], - "technical_id": row["technical_id"], - "created_at": row["created_at"], - "updated_at": row["updated_at"], - "operator": row["operator"], - } - + store = get_main_pointer_store() + return store.get_main_pointer( + collection=collection, + doc_id=doc_id, + step_type=step_type, + model_tag=model_tag, + user_id=None, + ) except Exception as e: raise MainPointerError(f"Failed to get main pointer: {e}") @@ -119,11 +64,8 @@ def set_main_pointer( ) -> None: """Set or update the main pointer for a specific document and stage. - Uses merge_insert for atomicity and avoids 'delete-then-add' race conditions. - Normalizes None model_tag to empty string. - Args: - lancedb_dir: Directory for LanceDB (unused, using connection from env) + lancedb_dir: Directory for LanceDB (unused, kept for backward compatibility) collection: Collection name doc_id: Document ID step_type: Processing stage type (parse, chunk, embed) @@ -136,51 +78,17 @@ def set_main_pointer( MainPointerError: If there's an error setting the pointer """ try: - conn = get_connection_from_env() - ensure_main_pointers_table(conn) - - table = conn.open_table("main_pointers") - normalized_tag = _normalize_model_tag(model_tag) - now = pd.Timestamp.now(tz="UTC") - - # Check if pointer already exists to preserve created_at - existing = get_main_pointer(collection, doc_id, step_type, model_tag) - - if existing: - created_at = existing["created_at"] - - # Fix-up: normalize NULL model_tag to "" in DB before merge_insert to avoid duplicates - if normalized_tag == "": - base_expr = _build_base_filter_expression(collection, doc_id, step_type) - null_filter = f"{base_expr} AND model_tag IS NULL" - try: - table.update(where=null_filter, values={"model_tag": ""}) - except Exception as update_err: - logger.warning("Failed to normalize NULL model_tag: %s", update_err) - else: - created_at = now - - # Prepare data for merge_insert - update_data = { - "collection": [collection], - "doc_id": [doc_id], - "step_type": [step_type], - "model_tag": [normalized_tag], - "semantic_id": [semantic_id], - "technical_id": [technical_id], - "created_at": [created_at], - "updated_at": [now], - "operator": [operator or "unknown"], - } - df = pd.DataFrame(update_data) - - ( - table.merge_insert(on=["collection", "doc_id", "step_type", "model_tag"]) - .when_matched_update_all() - .when_not_matched_insert_all() - .execute(df) + store = get_main_pointer_store() + store.set_main_pointer( + collection=collection, + doc_id=doc_id, + step_type=step_type, + semantic_id=semantic_id, + technical_id=technical_id, + model_tag=model_tag, + operator=operator, + user_id=None, ) - logger.info( f"Set main pointer for {collection}/{doc_id}/{step_type} to {technical_id} (semantic: {semantic_id})" ) @@ -205,45 +113,13 @@ def list_main_pointers( MainPointerError: If there's an error listing pointers """ try: - conn = get_connection_from_env() - ensure_main_pointers_table(conn) - - table = conn.open_table("main_pointers") - - # Build safe filter conditions - filters_dict = {"collection": collection} - if doc_id is not None: - filters_dict["doc_id"] = doc_id - - filter_expr = build_lancedb_filter_expression(filters_dict) - - # First check if any pointers exist using efficient count_rows - if table.search().where(filter_expr).count_rows() == 0: - return [] - - # Only load data if pointers exist - result = table.search().where(filter_expr).to_pandas() - - # Convert to list of dictionaries - pointers = [] - for _, row in result.iterrows(): - pointers.append( - { - "collection": row["collection"], - "doc_id": row["doc_id"], - "step_type": row["step_type"], - "model_tag": row["model_tag"] - if row["model_tag"] is not None - else "", - "semantic_id": row["semantic_id"], - "technical_id": row["technical_id"], - "created_at": row["created_at"], - "updated_at": row["updated_at"], - "operator": row["operator"], - } - ) - - return pointers + store = get_main_pointer_store() + return store.list_main_pointers( + collection=collection, + doc_id=doc_id, + user_id=None, + limit=100, + ) except Exception as e: raise MainPointerError(f"Failed to list main pointers: {e}") @@ -271,29 +147,17 @@ def delete_main_pointer( MainPointerError: If there's an error deleting the pointer """ try: - conn = get_connection_from_env() - ensure_main_pointers_table(conn) - - table = conn.open_table("main_pointers") - - # Build safe filter conditions - normalized_tag = _normalize_model_tag(model_tag) - base_expr = _build_base_filter_expression(collection, doc_id, step_type) - - if normalized_tag == "": - filter_expr = f"{base_expr} AND (model_tag == '' OR model_tag IS NULL)" - else: - filter_expr = f"{base_expr} AND model_tag == '{escape_lancedb_string(normalized_tag)}'" - - # Check if pointer exists using count_rows for efficiency - count = table.search().where(filter_expr).count_rows() - if count == 0: - return False - - # Delete the pointer(s) - table.delete(filter_expr) - logger.info(f"Deleted main pointer for {collection}/{doc_id}/{step_type}") - return True + store = get_main_pointer_store() + result = store.delete_main_pointer( + collection=collection, + doc_id=doc_id, + step_type=step_type, + model_tag=model_tag, + user_id=None, + ) + if result: + logger.info(f"Deleted main pointer for {collection}/{doc_id}/{step_type}") + return result except Exception as e: raise MainPointerError(f"Failed to delete main pointer: {e}") diff --git a/src/xagent/core/tools/core/document_search.py b/src/xagent/core/tools/core/document_search.py index 3a619d058..b8ea5e1ef 100644 --- a/src/xagent/core/tools/core/document_search.py +++ b/src/xagent/core/tools/core/document_search.py @@ -89,7 +89,7 @@ async def list_knowledge_bases( RuntimeError: If listing knowledge bases fails """ try: - result = list_collections(user_id=user_id, is_admin=is_admin) + result = await list_collections(user_id=user_id, is_admin=is_admin) kb_list = [] for collection in result.collections: @@ -138,7 +138,7 @@ async def search_knowledge_base( """ try: # List all collections - collections_result = list_collections(user_id=user_id, is_admin=is_admin) + collections_result = await list_collections(user_id=user_id, is_admin=is_admin) if not collections_result.collections: return KnowledgeSearchResult( diff --git a/src/xagent/providers/vector_store/lancedb.py b/src/xagent/providers/vector_store/lancedb.py index 2ad72ab09..2dce27bef 100644 --- a/src/xagent/providers/vector_store/lancedb.py +++ b/src/xagent/providers/vector_store/lancedb.py @@ -27,6 +27,7 @@ __all__ = [ "LanceDBConnectionManager", "LanceDBVectorStore", + "clear_connection_cache", "get_connection", "get_connection_from_env", ] @@ -39,6 +40,16 @@ CONNECTION_TTL = int(os.getenv("LANCEDB_CONNECTION_TTL", "300")) +def clear_connection_cache() -> None: + """Clear the global LanceDB connection cache. + + This is primarily intended for test isolation to avoid reusing cached + connections across different `LANCEDB_DIR` values. + """ + with _cache_lock: + _connection_cache.clear() + + class LanceDBConnectionManager: """ LanceDB connection manager with caching and automatic cleanup. diff --git a/src/xagent/web/api/kb.py b/src/xagent/web/api/kb.py index 1557dc7e1..8df564bcd 100644 --- a/src/xagent/web/api/kb.py +++ b/src/xagent/web/api/kb.py @@ -24,6 +24,7 @@ from pydantic import BaseModel from sqlalchemy.orm import Session +from ...core.tools.core.RAG_tools.core.config import DEFAULT_VECTOR_STORE_SCAN_LIMIT from ...core.tools.core.RAG_tools.core.schemas import ( ChunkStrategy, CollectionOperationResult, @@ -53,7 +54,7 @@ from ...core.tools.core.RAG_tools.pipelines.document_search import run_document_search from ...core.tools.core.RAG_tools.pipelines.web_ingestion import run_web_ingestion from ...core.tools.core.RAG_tools.progress import get_progress_manager -from ...providers.vector_store.lancedb import get_connection_from_env +from ...core.tools.core.RAG_tools.storage.factory import get_vector_index_store from ..auth_dependencies import get_current_user from ..config import ( MAX_FILE_SIZE, @@ -77,9 +78,6 @@ from ..services.kb_file_service import ( get_document_record_file_id as _get_document_record_file_id, ) -from ..services.kb_file_service import ( - list_documents_for_user as _list_documents_for_user, -) from ..services.kb_file_service import ( resolve_document_filename as _resolve_document_filename, ) @@ -170,45 +168,17 @@ async def save_collection_config( _user: User = Depends(get_current_user), ) -> CollectionOperationResult: """Save ingestion configuration for a specific collection.""" - from datetime import datetime, timezone - - from ...core.tools.core.RAG_tools.LanceDB.schema_manager import ( - ensure_collection_config_table, - ) - from ...providers.vector_store.lancedb import get_connection_from_env + from ...core.tools.core.RAG_tools.storage.factory import get_metadata_store - def _save_config() -> None: - conn = get_connection_from_env() - # TODO(refactor): keep collection_config as a compatibility store for - # per-user ingestion settings; unify this with metadata-backed storage - # once config ownership and migration strategy are finalized. - ensure_collection_config_table(conn) - table = conn.open_table("collection_config") - - user_id_val = int(_user.id) - config_json = config.model_dump_json(exclude_unset=True) - now = datetime.now(timezone.utc).replace(tzinfo=None) - - try: - # Try to delete existing configuration for this collection and user - table.delete(f"collection = '{collection}' AND user_id = {user_id_val}") - except Exception as e: - logger.warning(f"Error deleting old config: {e}") - - # Insert new config - data = [ - { - "collection": collection, - "config_json": config_json, - "updated_at": now, - "user_id": user_id_val, - } - ] - - table.add(data) + config_json = config.model_dump_json(exclude_unset=True) try: - await asyncio.to_thread(_save_config) + metadata_store = get_metadata_store() + await metadata_store.save_collection_config( + collection=collection, + config_json=config_json, + user_id=int(_user.id), + ) return CollectionOperationResult( status="success", @@ -629,7 +599,7 @@ async def list_collections_api( try: result = await asyncio.wait_for( - asyncio.to_thread(list_collections, int(_user.id), bool(_user.is_admin)), + list_collections(user_id=int(_user.id), is_admin=bool(_user.is_admin)), timeout=kb_collections_timeout_seconds, ) return result @@ -670,7 +640,10 @@ async def search( ), filters: Optional[Dict[str, Any]] = Form( None, - description="Optional filters to apply during search (LanceDB format)", + description="Optional filters to apply during search. " + "Format: {field: value} for equality filters. " + "For advanced filters, use {field: {operator: str, value: Any}} " + "where operator can be: eq, ne, gt, gte, lt, lte, in, contains.", ), fusion_config: Optional[Dict[str, Any]] = Form( None, @@ -1084,10 +1057,11 @@ async def delete_collection_api( ), ) - collection_records = _list_documents_for_user( + vector_store = get_vector_index_store() + collection_records = vector_store.list_document_records( + collection_name=collection_name, user_id=int(_user.id), is_admin=bool(_user.is_admin), - collection_name=collection_name, ) collection_file_ids = { file_id @@ -1099,7 +1073,8 @@ async def delete_collection_api( result = delete_collection(collection_name, int(_user.id), bool(_user.is_admin)) - remaining_records = _list_documents_for_user( + remaining_records = vector_store.list_document_records( + collection_name=None, user_id=int(_user.id), is_admin=bool(_user.is_admin), ) @@ -1230,11 +1205,17 @@ async def check_documents_exist_api( if not requested: return {"existing_filenames": []} - records = _list_documents_for_user( + # Use storage abstraction layer to fetch document records + vector_store = get_vector_index_store() + records = vector_store.list_document_records( + collection_name=collection_name, user_id=int(_user.id), is_admin=False, - collection_name=collection_name, + max_results=DEFAULT_VECTOR_STORE_SCAN_LIMIT, ) + + # Build filename map from file_ids (for UploadedFile lookup) + # This preserves main branch's file_id -> filename resolution filename_map = _build_uploaded_filename_map( db, user_id=int(_user.id), @@ -1249,6 +1230,7 @@ async def check_documents_exist_api( existing_filenames = set() for record in records: + # Resolve filename using file_id first, then fallback to source_path basename resolved_filename = _resolve_document_filename(record, filename_map) if resolved_filename: existing_filenames.add(resolved_filename) @@ -1297,34 +1279,41 @@ async def delete_document_api( # NOTE: Exceptions are normalized by @handle_kb_exceptions for consistent API responses. from ...core.tools.core.RAG_tools.management.collections import delete_document - records = _list_documents_for_user( + # Use storage abstraction layer to fetch document records + vector_store = get_vector_index_store() + records = vector_store.list_document_records( + collection_name=collection_name, user_id=int(_user.id), is_admin=bool(_user.is_admin), - collection_name=collection_name, + max_results=DEFAULT_VECTOR_STORE_SCAN_LIMIT, ) + + # Build filename map from file_ids (for UploadedFile lookup and advanced matching) filename_map = _build_uploaded_filename_map( db, user_id=int(_user.id), file_ids=[ - current_file_id - for current_file_id in ( - _get_document_record_file_id(record) for record in records - ) - if current_file_id + file_id + for file_id in (_get_document_record_file_id(record) for record in records) + if file_id ], ) + # Find all matching documents (handle duplicates) matching_docs = [] for record in records: - current_doc_id = record.get("doc_id") + current_doc_id = record.doc_id current_file_id = _get_document_record_file_id(record) resolved_filename = _resolve_document_filename(record, filename_map) + + # Support filtering by doc_id, file_id, or filename (main branch feature) if doc_id and current_doc_id != doc_id: continue if file_id and current_file_id != file_id: continue if not doc_id and not file_id and resolved_filename != filename: continue + matching_docs.append( { "doc_id": current_doc_id, @@ -1342,7 +1331,9 @@ async def delete_document_api( deleted_doc_ids = [] deletion_errors = [] - remaining_records = _list_documents_for_user( + # Get remaining documents to check for orphaned UploadedFile records + remaining_records = vector_store.list_document_records( + collection_name=collection_name, user_id=int(_user.id), is_admin=bool(_user.is_admin), ) @@ -1431,19 +1422,14 @@ async def rename_collection_api( Returns: Success message """ - from ...core.tools.core.RAG_tools.management.collections import ( - _list_table_names, - ) from ...core.tools.core.RAG_tools.management.status import ( clear_ingestion_status, load_ingestion_status, write_ingestion_status, ) - from ...core.tools.core.RAG_tools.utils.string_utils import ( - escape_lancedb_string, - ) + from ...core.tools.core.RAG_tools.storage.factory import get_vector_index_store - conn = get_connection_from_env() + vector_store = get_vector_index_store() if not new_name or not new_name.strip(): raise HTTPException( @@ -1471,15 +1457,15 @@ async def rename_collection_api( physical_rename_error: Optional[str] = None old_collection_dir: Optional[Path] = None new_collection_dir: Optional[Path] = None + collection_records = vector_store.list_document_records( + collection_name=collection_name, + user_id=int(_user.id), + is_admin=bool(_user.is_admin), + ) collection_file_ids = { file_id for file_id in ( - _get_document_record_file_id(record) - for record in _list_documents_for_user( - user_id=int(_user.id), - is_admin=bool(_user.is_admin), - collection_name=collection_name, - ) + _get_document_record_file_id(record) for record in collection_records ) if file_id } @@ -1510,33 +1496,15 @@ async def rename_collection_api( ), ) - # Step 2: Update collection name in all tables - table_names = _list_table_names(conn, warnings) - - for table_name in ["documents", "parses", "chunks"]: - if table_name in table_names: - try: - table = conn.open_table(table_name) - table.update( - f"collection = '{escape_lancedb_string(collection_name)}'", - {"collection": new_name}, - ) - except Exception as e: - logger.warning("Failed to update '%s': %s", table_name, e) - warnings.append(f"Failed to update '{table_name}': {e}") - - for table_name in table_names: - if not table_name.startswith("embeddings_"): - continue - try: - table = conn.open_table(table_name) - table.update( - f"collection = '{escape_lancedb_string(collection_name)}'", - {"collection": new_name}, - ) - except Exception as e: - logger.warning("Failed to update embeddings table '%s': %s", table_name, e) - warnings.append(f"Failed to update '{table_name}': {e}") + # Step 2: Update collection name in all tables (documents, parses, chunks, embeddings) + # Use storage abstraction layer which handles all tables including embeddings + vector_store = get_vector_index_store() + warnings.extend( + vector_store.rename_collection_data( + collection_name=collection_name, + new_name=new_name, + ) + ) # Migrate ingestion status from old collection name to new try: diff --git a/src/xagent/web/services/kb_file_service.py b/src/xagent/web/services/kb_file_service.py index 6e366020d..09fd75b76 100644 --- a/src/xagent/web/services/kb_file_service.py +++ b/src/xagent/web/services/kb_file_service.py @@ -5,12 +5,13 @@ import logging import os from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union from sqlalchemy.orm import Session from ...config import get_uploads_dir from ...core.tools.core.RAG_tools.LanceDB.schema_manager import ensure_documents_table +from ...core.tools.core.RAG_tools.storage.contracts import DocumentRecord from ...core.tools.core.RAG_tools.utils.lancedb_query_utils import query_to_list from ...core.tools.core.RAG_tools.utils.string_utils import ( build_lancedb_filter_expression, @@ -104,9 +105,24 @@ def build_uploaded_filename_map( return {str(record.file_id): str(record.filename) for record in records} -def get_document_record_file_id(record: Dict[str, Any]) -> Optional[str]: - """Extract a normalized ``file_id`` from a KB document record.""" - raw_file_id = record.get("file_id") +def get_document_record_file_id( + record: Union[Dict[str, Any], DocumentRecord], +) -> Optional[str]: + """Extract a normalized ``file_id`` from a KB document record. + + Args: + record: Either a Dict[str, Any] or DocumentRecord dataclass. + + Returns: + Normalized file_id string or None. + """ + # Handle both Dict and DocumentRecord types + if isinstance(record, dict): + raw_file_id = record.get("file_id") + else: + # Assume DocumentRecord dataclass with file_id attribute + raw_file_id = getattr(record, "file_id", None) + if raw_file_id is None: return None file_id = str(raw_file_id).strip() @@ -114,15 +130,30 @@ def get_document_record_file_id(record: Dict[str, Any]) -> Optional[str]: def resolve_document_filename( - record: Dict[str, Any], filename_map: Dict[str, str] + record: Union[Dict[str, Any], DocumentRecord], filename_map: Dict[str, str] ) -> Optional[str]: - """Resolve a user-facing filename from ``file_id`` first, then legacy path.""" + """Resolve a user-facing filename from ``file_id`` first, then legacy path. + + Args: + record: Either a Dict[str, Any] or DocumentRecord dataclass. + filename_map: Mapping from file_id to filename. + + Returns: + Resolved filename or None. + """ file_id = get_document_record_file_id(record) if file_id and filename_map.get(file_id): return filename_map[file_id] - source_path = record.get("source_path") + + # Handle both Dict and DocumentRecord types for source_path + if isinstance(record, dict): + source_path = record.get("source_path") + else: + source_path = getattr(record, "source_path", None) + if source_path: return os.path.basename(str(source_path)) + return None @@ -133,7 +164,17 @@ def delete_uploaded_file_if_orphaned( user_id: int, remaining_file_ids: set[str], ) -> bool: - """Delete uploaded file row and local file when no documents still reference it.""" + """Delete uploaded file row and local file when no documents still reference it. + + Args: + db: Database session. + file_id: The ID of the file to check. + user_id: User ID for scoping. + remaining_file_ids: A set of all file_id values still referenced by other documents. + + Returns: + True if the file was deleted, False otherwise. + """ if not file_id or file_id in remaining_file_ids: return False @@ -161,6 +202,8 @@ def delete_uploaded_file_if_orphaned( else: if resolved_path.exists() and resolved_path.is_file(): resolved_path.unlink() + logger.info("Deleted orphaned physical file: %s", resolved_path) db.delete(file_record) + db.flush() return True diff --git a/tests/conftest.py b/tests/conftest.py index a4828b60d..0e845b53b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -22,6 +22,8 @@ from xagent.core.model import ChatModelConfig, EmbeddingModelConfig, RerankModelConfig from xagent.core.observability.langfuse_tracer import init_tracer, reset_tracer +from xagent.core.tools.core.RAG_tools.storage import reset_kb_write_coordinator +from xagent.providers.vector_store.lancedb import clear_connection_cache # YAML entrypoint has been removed, commenting out these imports # from xagent.entrypoint.yaml.parser import MigrationManager @@ -80,6 +82,13 @@ def pytest_collection_modifyitems(config, items): # ========================================== +def _security_test_subdir(tmp_path: Path, name: str) -> str: + """Create ``tmp_path / name`` and return its path as a string.""" + subdir = tmp_path / name + subdir.mkdir() + return str(subdir) + + @pytest.fixture def temp_dir(): """Provide a temporary directory for tests.""" @@ -87,28 +96,51 @@ def temp_dir(): yield temp_dir +@pytest.fixture(autouse=True, scope="function") +def isolate_lancedb_dir(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Isolate LanceDB and reset KB storage singletons for every test. + + By default, ``LANCEDB_DIR`` is set to a fresh directory under ``tmp_path`` + for each test. This avoids stale LanceDB schemas from a developer ``.env`` + or a fixed path, and matches CI-style ephemeral storage. Parallel workers + (pytest-xdist) each use their own process-local ``tmp_path``. + + If the environment sets ``XAGENT_PYTEST_RESPECT_LANCEDB_DIR=1``, the + existing ``LANCEDB_DIR`` from the environment is left unchanged (for CI or + local workflows that intentionally pin a path). + + Clears the LanceDB connection cache and resets the process-wide KB write + coordinator before and after each test. + """ + respect_env = os.environ.get("XAGENT_PYTEST_RESPECT_LANCEDB_DIR") == "1" + if not respect_env: + lancedb_dir = tmp_path / "lancedb" + lancedb_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("LANCEDB_DIR", str(lancedb_dir)) + + clear_connection_cache() + reset_kb_write_coordinator() + yield + reset_kb_write_coordinator() + clear_connection_cache() + + @pytest.fixture -def test_workspace_dir(tmp_path): - """Create test workspace directory for security testing.""" - workspace_dir = tmp_path / "test_workspace" - workspace_dir.mkdir() - return str(workspace_dir) +def test_workspace_dir(tmp_path: Path) -> str: + """Directory used as workspace root in ``test_service_security``.""" + return _security_test_subdir(tmp_path, "test_workspace") @pytest.fixture -def test_access_dir(tmp_path): - """Create test access directory for security testing.""" - access_dir = tmp_path / "test_access_restriction" - access_dir.mkdir() - return str(access_dir) +def test_access_dir(tmp_path: Path) -> str: + """Directory used for access-restriction scenarios in security tests.""" + return _security_test_subdir(tmp_path, "test_access_restriction") @pytest.fixture -def test_security_dir(tmp_path): - """Create test security directory for security testing.""" - security_dir = tmp_path / "test_security" - security_dir.mkdir() - return str(security_dir) +def test_security_dir(tmp_path: Path) -> str: + """Directory used for outside-access rejection scenarios in security tests.""" + return _security_test_subdir(tmp_path, "test_security") @pytest.fixture(autouse=True, scope="function") diff --git a/tests/core/test_config.py b/tests/core/test_config.py index 1c288799b..29ca6aed6 100644 --- a/tests/core/test_config.py +++ b/tests/core/test_config.py @@ -209,10 +209,11 @@ class TestGetLancedbPath: """Test get_lancedb_path() function.""" def test_default_lancedb_path(self, monkeypatch): - """Test default LanceDB path (relative to cwd).""" + """Test default LanceDB path (relative to storage root).""" monkeypatch.delenv(LANCEDB_PATH, raising=False) + monkeypatch.delenv(STORAGE_ROOT, raising=False) result = get_lancedb_path() - assert result == Path("data/lancedb") + assert result == Path.home() / ".xagent" / "data" / "lancedb" def test_lancedb_path_with_env_var(self, monkeypatch): """Test LanceDB path with environment variable.""" diff --git a/tests/core/tools/core/RAG_tools/LanceDB/test_schema_manager.py b/tests/core/tools/core/RAG_tools/LanceDB/test_schema_manager.py index 3e6fa2f3f..c5198cee9 100644 --- a/tests/core/tools/core/RAG_tools/LanceDB/test_schema_manager.py +++ b/tests/core/tools/core/RAG_tools/LanceDB/test_schema_manager.py @@ -12,13 +12,13 @@ ensure_embeddings_table, ensure_parses_table, ) -from xagent.providers.vector_store.lancedb import get_connection_from_env +from xagent.core.tools.core.RAG_tools.storage import get_vector_store_raw_connection def test_ensure_tables(tmp_path: Path, monkeypatch) -> None: db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() ensure_documents_table(conn) ensure_parses_table(conn) ensure_chunks_table(conn) @@ -40,7 +40,7 @@ def test_check_table_needs_migration_table_not_exists( """Test check_table_needs_migration when table doesn't exist.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # Table doesn't exist, should return False assert check_table_needs_migration(conn, "nonexistent_table") is False @@ -52,7 +52,7 @@ def test_check_table_needs_migration_table_without_user_id( """Test check_table_needs_migration when table exists but missing user_id field.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # Create a table without user_id field (old schema) old_schema = pa.schema( @@ -74,7 +74,7 @@ def test_check_table_needs_migration_table_with_user_id( """Test check_table_needs_migration when table exists and has user_id field.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # Create a table with user_id field (new schema) new_schema = pa.schema( @@ -97,7 +97,7 @@ def test_check_table_needs_migration_with_ensure_tables( """Test check_table_needs_migration with tables created by ensure_* functions.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # Create tables using ensure_* functions (which create tables with user_id) ensure_documents_table(conn) diff --git a/tests/core/tools/core/RAG_tools/LanceDB/test_schema_migration.py b/tests/core/tools/core/RAG_tools/LanceDB/test_schema_migration.py index afb77bd43..d938966ba 100644 --- a/tests/core/tools/core/RAG_tools/LanceDB/test_schema_migration.py +++ b/tests/core/tools/core/RAG_tools/LanceDB/test_schema_migration.py @@ -21,7 +21,7 @@ ensure_parses_table, ensure_prompt_templates_table, ) -from xagent.providers.vector_store.lancedb import get_connection_from_env +from xagent.core.tools.core.RAG_tools.storage import get_vector_store_raw_connection def test_get_sql_default_for_pa_type(): @@ -40,7 +40,7 @@ def test_auto_migration_adds_missing_columns(tmp_path: Path, monkeypatch): """Test that missing columns are automatically added with correct defaults.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # 1. Create a table with an OLD schema (missing 'language' and 'title') old_schema = pa.schema( @@ -82,7 +82,7 @@ def test_ensure_schema_fields_idempotency(tmp_path: Path, monkeypatch): """Test that calling migration on an up-to-date table is safe.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # Create table with FULL schema first ensure_collection_metadata_table(conn) @@ -136,7 +136,7 @@ def test_manual_migration_helper(tmp_path: Path, monkeypatch): """Test the low-level _ensure_schema_fields helper directly.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # Setup simple table conn.create_table("test_manual", schema=pa.schema([("a", pa.int32())])) @@ -163,7 +163,7 @@ def test_ensure_schema_fields_type_mismatch_keeps_existing_type( """Type mismatch should not rewrite existing column types.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() conn.create_table("test_type_mismatch", schema=pa.schema([("a", pa.int32())])) conn.open_table("test_type_mismatch").add([{"a": 7}]) @@ -187,7 +187,7 @@ def test_ensure_schema_fields_partial_failure_raises( """When add_columns fails, migration should raise instead of silently masking.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() conn.create_table("test_partial_failure", schema=pa.schema([("a", pa.int32())])) table = conn.open_table("test_partial_failure") @@ -235,7 +235,7 @@ def test_create_table_existing_with_schema_triggers_migration( """_create_table should migrate existing table when schema is provided.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() conn.create_table("create_table_migrate", schema=pa.schema([("a", pa.int32())])) target_schema = pa.schema([("a", pa.int32()), ("b", pa.string())]) @@ -252,7 +252,7 @@ def test_ensure_embeddings_table_with_fixed_vector_dim( """ensure_embeddings_table should use fixed-size list when vector_dim is set.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() ensure_embeddings_table(conn, "test_fixed", vector_dim=8) schema = conn.open_table("embeddings_test_fixed").schema @@ -266,7 +266,7 @@ def test_ensure_embeddings_table_with_variable_vector_dim( """ensure_embeddings_table should use variable list when vector_dim is None.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() ensure_embeddings_table(conn, "test_variable", vector_dim=None) schema = conn.open_table("embeddings_test_variable").schema @@ -280,7 +280,7 @@ def test_ensure_collection_config_table_create_and_idempotent( """ensure_collection_config_table should be creatable and idempotent.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() ensure_collection_config_table(conn) schema_before = conn.open_table("collection_config").schema @@ -298,7 +298,7 @@ def test_ensure_parses_table_migrates_missing_user_id( """ensure_parses_table should add user_id for legacy schema.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() old_schema = pa.schema( [ @@ -337,7 +337,7 @@ def test_ensure_prompt_templates_table_migrates_missing_user_id( """ensure_prompt_templates_table should add user_id for legacy schema.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() old_schema = pa.schema( [ @@ -380,7 +380,7 @@ def test_ensure_ingestion_runs_table_migrates_missing_user_id( """ensure_ingestion_runs_table should add user_id for legacy schema.""" db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() old_schema = pa.schema( [ @@ -416,16 +416,22 @@ def test_ensure_ingestion_runs_table_migrates_missing_user_id( def test_concurrent_ensure_collection_metadata_table_is_safe( tmp_path: Path, monkeypatch ) -> None: - """Concurrent ensure_collection_metadata_table calls should be safe.""" + """Concurrent ensure_collection_metadata_table calls should be safe. + + Note: This test verifies that the table creation logic is idempotent and safe + when called concurrently with different connections. Each thread uses its own + connection to avoid LanceDB connection threading issues. + """ db_dir = tmp_path / "db" monkeypatch.setenv("LANCEDB_DIR", str(db_dir)) - conn = get_connection_from_env() errors: list[Exception] = [] def _worker() -> None: try: - ensure_collection_metadata_table(conn) + # Each thread gets its own connection to avoid threading issues + worker_conn = get_vector_store_raw_connection() + ensure_collection_metadata_table(worker_conn) except Exception as exc: # noqa: BLE001 errors.append(exc) @@ -436,5 +442,7 @@ def _worker() -> None: t.join() assert errors == [] + # Verify the table was created successfully + conn = get_vector_store_raw_connection() schema = conn.open_table("collection_metadata").schema assert "ingestion_config" in schema.names diff --git a/tests/core/tools/core/RAG_tools/chunk/test_chunk_document.py b/tests/core/tools/core/RAG_tools/chunk/test_chunk_document.py index 4f697370a..393338c5d 100644 --- a/tests/core/tools/core/RAG_tools/chunk/test_chunk_document.py +++ b/tests/core/tools/core/RAG_tools/chunk/test_chunk_document.py @@ -497,9 +497,11 @@ def test_chunk_recursive_protected_content_keeps_code_block( ) assert chunk_result["created"] is True assert chunk_result["chunk_count"] > 0 - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -541,9 +543,11 @@ def test_chunk_markdown_with_headers_section_in_metadata( ) assert chunk_result["created"] is True assert chunk_result["chunk_count"] > 0 - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -594,9 +598,11 @@ def test_chunk_table_context_attached( ) assert chunk_result["created"] is True assert chunk_result["chunk_count"] > 0 - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -715,9 +721,11 @@ def test_chunk_config_hash_idempotency( assert chunk_result2["created"] is False # Should not write again # Verify database state - both should reference same config_hash - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -781,9 +789,11 @@ def test_chunk_separators_create_new_version( assert chunk_result2["created"] is True # Verify database has two different config_hash versions - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -836,9 +846,11 @@ def test_chunk_recursive_custom_separators_integration( assert chunk_result["created"] is True assert chunk_result["chunk_count"] >= 1 - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -896,9 +908,11 @@ def test_chunk_recursive_custom_separators_vs_default_different_result( assert chunk_default["created"] is True assert chunk_custom["created"] is True # Different separators must yield different config_hash (hence different version) - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -948,9 +962,11 @@ def test_chunk_row_level_hash_uniqueness( assert chunk_result["chunk_count"] > 1 # Need multiple chunks for this test # Step 3: Verify row-level chunk_hash uniqueness - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -1057,9 +1073,11 @@ def test_chunk_table_structure_validation( assert chunk_result["created"] is True # Step 3: Verify table structure - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") # Table should exist and be accessible @@ -1247,12 +1265,14 @@ def test_chunk_metadata_serialization_and_retrieval( assert chunk_result["created"] is True # Step 3: Verify metadata in database (should be serialized as JSON string) + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) from xagent.core.tools.core.RAG_tools.utils.metadata_utils import ( deserialize_metadata, ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() table = conn.open_table("chunks") df = ( table.search() @@ -1316,21 +1336,17 @@ def test_collection(self): def test_chunk_document_arrow_fallback_chain( self, temp_lancedb_dir, test_collection ) -> None: - """Test chunk_document uses to_arrow() -> to_list() -> to_pandas() fallback.""" + """Test chunk_document uses iter_batches with Arrow RecordBatch.""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.chunk.chunk_document import ( _get_existing_chunks, ) - mock_db_connection = MagicMock() - mock_table = MagicMock() - - def mock_open_table_func(table_name): - return mock_table - - mock_db_connection.open_table.side_effect = mock_open_table_func + # Mock the vector store + mock_vector_store = MagicMock() + # Create mock batch data (simulating Arrow RecordBatch) chunks_data = [ { "chunk_id": "chunk1", @@ -1341,26 +1357,25 @@ def mock_open_table_func(table_name): "index": 0, "created_at": pd.Timestamp.now(), "metadata": '{"key": "value"}', + "page_number": None, + "section": None, + "anchor": None, + "json_path": None, } ] - mock_arrow_table = MagicMock() - mock_arrow_table.to_pylist.return_value = chunks_data - - mock_search = MagicMock() - mock_where = MagicMock() - mock_table.search.return_value = mock_search - mock_search.where.return_value = mock_where - mock_where.to_arrow.return_value = mock_arrow_table - mock_table.count_rows.return_value = 1 - - with ( - patch( - "xagent.core.tools.core.RAG_tools.chunk.chunk_document.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.chunk.chunk_document.ensure_chunks_table" - ), + + # Create mock batch + mock_batch = MagicMock() + mock_batch.num_rows = 1 + mock_batch.to_pandas.return_value = pd.DataFrame([chunks_data[0]]) + + # Mock iter_batches to yield the mock batch + mock_vector_store.count_rows_or_zero.return_value = 1 + mock_vector_store.iter_batches.return_value = [mock_batch] + + with patch( + "xagent.core.tools.core.RAG_tools.chunk.chunk_document.get_vector_index_store", + return_value=mock_vector_store, ): result = _get_existing_chunks( collection=test_collection, @@ -1371,27 +1386,24 @@ def mock_open_table_func(table_name): assert len(result) == 1 assert result[0]["chunk_id"] == "chunk1" - # Verify to_arrow() was called - mock_where.to_arrow.assert_called_once() + # Verify count_rows_or_zero and iter_batches were called + mock_vector_store.count_rows_or_zero.assert_called_once() + mock_vector_store.iter_batches.assert_called_once() def test_chunk_document_fallback_to_list( self, temp_lancedb_dir, test_collection ) -> None: - """Test chunk_document fallback from to_arrow() to to_list().""" + """Test chunk_document handles batch data correctly.""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.chunk.chunk_document import ( _get_existing_chunks, ) - mock_db_connection = MagicMock() - mock_table = MagicMock() - - def mock_open_table_func(table_name): - return mock_table - - mock_db_connection.open_table.side_effect = mock_open_table_func + # Mock the vector store + mock_vector_store = MagicMock() + # Create mock batch data chunks_data = [ { "chunk_id": "chunk1", @@ -1402,26 +1414,25 @@ def mock_open_table_func(table_name): "index": 0, "created_at": pd.Timestamp.now(), "metadata": '{"key": "value"}', + "page_number": None, + "section": None, + "anchor": None, + "json_path": None, } ] - mock_search = MagicMock() - mock_where = MagicMock() - mock_table.search.return_value = mock_search - mock_search.where.return_value = mock_where - # to_arrow() fails, fallback to to_list() - mock_where.to_arrow.side_effect = AttributeError("to_arrow not available") - mock_where.to_list.return_value = chunks_data - mock_table.count_rows.return_value = 1 - - with ( - patch( - "xagent.core.tools.core.RAG_tools.chunk.chunk_document.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.chunk.chunk_document.ensure_chunks_table" - ), + # Create mock batch + mock_batch = MagicMock() + mock_batch.num_rows = 1 + mock_batch.to_pandas.return_value = pd.DataFrame([chunks_data[0]]) + + # Mock iter_batches to yield the mock batch + mock_vector_store.count_rows_or_zero.return_value = 1 + mock_vector_store.iter_batches.return_value = [mock_batch] + + with patch( + "xagent.core.tools.core.RAG_tools.chunk.chunk_document.get_vector_index_store", + return_value=mock_vector_store, ): result = _get_existing_chunks( collection=test_collection, @@ -1432,65 +1443,51 @@ def mock_open_table_func(table_name): assert len(result) == 1 assert result[0]["chunk_id"] == "chunk1" - # Verify fallback was used - mock_where.to_arrow.assert_called_once() - mock_where.to_list.assert_called_once() + # Verify methods were called + mock_vector_store.count_rows_or_zero.assert_called_once() + mock_vector_store.iter_batches.assert_called_once() def test_chunk_document_fallback_to_pandas_with_nan( self, temp_lancedb_dir, test_collection ) -> None: - """Test chunk_document fallback to to_pandas() and NaN normalization.""" + """Test chunk_document handles batch data correctly via iter_batches.""" from unittest.mock import MagicMock, patch - import numpy as np - from xagent.core.tools.core.RAG_tools.chunk.chunk_document import ( _get_existing_chunks, ) - mock_db_connection = MagicMock() - mock_table = MagicMock() - - def mock_open_table_func(table_name): - return mock_table - - mock_db_connection.open_table.side_effect = mock_open_table_func - - # Create DataFrame with NaN values - chunks_df = pd.DataFrame( - [ - { - "chunk_id": "chunk1", - "text": "test content", - "collection": test_collection, - "doc_id": "doc1", - "parse_hash": "hash1", - "index": 0, - "created_at": pd.Timestamp.now(), - "metadata": '{"key": "value"}', - "page_number": np.nan, # NaN value - } - ] - ) - - mock_search = MagicMock() - mock_where = MagicMock() - mock_table.search.return_value = mock_search - mock_search.where.return_value = mock_where - # Both to_arrow() and to_list() fail, fallback to to_pandas() - mock_where.to_arrow.side_effect = AttributeError("to_arrow not available") - mock_where.to_list.side_effect = AttributeError("to_list not available") - mock_where.to_pandas.return_value = chunks_df - mock_table.count_rows.return_value = 1 - - with ( - patch( - "xagent.core.tools.core.RAG_tools.chunk.chunk_document.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.chunk.chunk_document.ensure_chunks_table" - ), + # Mock the vector store + mock_vector_store = MagicMock() + + # Create mock batch data (without NaN - use None directly) + chunks_data = { + "chunk_id": "chunk1", + "text": "test content", + "collection": test_collection, + "doc_id": "doc1", + "parse_hash": "hash1", + "index": 0, + "created_at": pd.Timestamp.now(), + "metadata": '{"key": "value"}', + "page_number": None, + "section": None, + "anchor": None, + "json_path": None, + } + + # Create mock batch + mock_batch = MagicMock() + mock_batch.num_rows = 1 + mock_batch.to_pandas.return_value = pd.DataFrame([chunks_data]) + + # Mock iter_batches to yield the mock batch + mock_vector_store.count_rows_or_zero.return_value = 1 + mock_vector_store.iter_batches.return_value = [mock_batch] + + with patch( + "xagent.core.tools.core.RAG_tools.chunk.chunk_document.get_vector_index_store", + return_value=mock_vector_store, ): result = _get_existing_chunks( collection=test_collection, @@ -1501,9 +1498,5 @@ def mock_open_table_func(table_name): assert len(result) == 1 assert result[0]["chunk_id"] == "chunk1" - # Verify all fallbacks were attempted - mock_where.to_arrow.assert_called_once() - mock_where.to_list.assert_called_once() - mock_where.to_pandas.assert_called_once() - # Verify NaN was normalized to None - assert result[0].get("page_number") is None + # Verify None values are preserved + assert result[0]["page_number"] is None diff --git a/tests/core/tools/core/RAG_tools/core/test_factory_utils.py b/tests/core/tools/core/RAG_tools/core/test_factory_utils.py index 8e5f5219b..ed0b027fb 100644 --- a/tests/core/tools/core/RAG_tools/core/test_factory_utils.py +++ b/tests/core/tools/core/RAG_tools/core/test_factory_utils.py @@ -45,8 +45,7 @@ def test_get_default_index_policy() -> None: Note: This function returns static defaults only. The actual dynamic index type selection based on data scale (HNSW for 50k-10M rows, IVFPQ for >=10M rows) is - implemented in IndexManager.check_and_create_index() and comprehensively tested - in tests/vector_storage/test_index_manager.py. + implemented in storage.lancedb_stores.LanceDBVectorIndexStore.create_index(). """ threshold, index_type = get_default_index_policy() diff --git a/tests/core/tools/core/RAG_tools/file/test_register_document.py b/tests/core/tools/core/RAG_tools/file/test_register_document.py index c226fa897..f2142997e 100644 --- a/tests/core/tools/core/RAG_tools/file/test_register_document.py +++ b/tests/core/tools/core/RAG_tools/file/test_register_document.py @@ -247,10 +247,10 @@ def test_register_document_hash_computation_error( register_document(collection="test_collection", source_path=str(test_file)) @patch( - "xagent.core.tools.core.RAG_tools.file.register_document.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.file.register_document.get_vector_index_store" ) def test_register_document_configuration_error( - self, mock_get_db, tmp_path: Path + self, mock_get_store, tmp_path: Path ) -> None: """Test handling configuration errors.""" # Setup test file @@ -258,7 +258,11 @@ def test_register_document_configuration_error( test_file.write_text("Test content") # Mock database connection to raise configuration error - mock_get_db.side_effect = ConfigurationError("LANCEDB_DIR not configured") + mock_store = MagicMock() + mock_store.count_rows_or_zero.side_effect = ConfigurationError( + "LANCEDB_DIR not configured" + ) + mock_get_store.return_value = mock_store # Should propagate ConfigurationError with pytest.raises(ConfigurationError): @@ -285,10 +289,10 @@ def test_register_document_unsupported_file_type( register_document(collection=collection, source_path=str(unsupported_file)) @patch( - "xagent.core.tools.core.RAG_tools.file.register_document.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.file.register_document.get_vector_index_store" ) def test_register_document_database_operation_error( - self, mock_get_db, tmp_path: Path, monkeypatch + self, mock_get_store, tmp_path: Path, monkeypatch ) -> None: """Test handling database operation errors.""" # Setup environment variable @@ -299,15 +303,10 @@ def test_register_document_database_operation_error( test_file = tmp_path / "db_error_test.txt" test_file.write_text("Test content") - # Mock database connection to succeed, but table operations to fail - mock_db = MagicMock() - mock_get_db.return_value = mock_db - - # Mock ensure_documents_table to succeed - mock_db.ensure_documents_table = MagicMock() - - # Mock open_table to raise an error - mock_db.open_table.side_effect = Exception("Table access failed") + # Mock vector store to raise an error + mock_store = MagicMock() + mock_store.count_rows_or_zero.side_effect = Exception("Table access failed") + mock_get_store.return_value = mock_store # Should propagate DatabaseOperationError with pytest.raises(DatabaseOperationError, match="Table access failed"): diff --git a/tests/core/tools/core/RAG_tools/generate/test_format_generation_prompt.py b/tests/core/tools/core/RAG_tools/generate/test_format_generation_prompt.py index 21aacd5f7..caed905e9 100644 --- a/tests/core/tools/core/RAG_tools/generate/test_format_generation_prompt.py +++ b/tests/core/tools/core/RAG_tools/generate/test_format_generation_prompt.py @@ -11,44 +11,53 @@ @pytest.fixture -def sample_prompt_template() -> str: - """Provides a sample prompt template for testing.""" - return "Please summarize the following context:\n{context}" +def sample_prompt_template_placeholder() -> str: + """Provides a sample prompt template with placeholder.""" + return "Summarize this: {context}" @pytest.fixture -def sample_formatted_contexts() -> str: - """Provides sample formatted contexts for testing.""" - return "This is the first chunk.\n---\nThis is the second chunk." +def sample_prompt_template_plain() -> str: + """Provides a sample prompt template without placeholder.""" + return "Please summarize the following context:" @pytest.fixture -def expected_full_prompt( - sample_prompt_template: str, sample_formatted_contexts: str -) -> str: - """Provides the expected full prompt string.""" - return ( - f"{sample_prompt_template}\n\nContext:\n{sample_formatted_contexts}\n\nAnswer:" - ) +def sample_formatted_contexts() -> str: + """Provides sample formatted contexts for testing.""" + return "This is the first chunk.\n---\nThis is the second chunk." class TestFormatGenerationPrompt: """Tests for the format_generation_prompt core function.""" - def test_format_generation_prompt_success( + def test_format_generation_prompt_with_placeholder( + self, + sample_prompt_template_placeholder: str, + sample_formatted_contexts: str, + ) -> None: + """Test formatting when placeholder is present.""" + result = format_generation_prompt( + prompt_template=sample_prompt_template_placeholder, + formatted_contexts=sample_formatted_contexts, + ) + + expected = f"Summarize this: {sample_formatted_contexts}" + assert result == expected + + def test_format_generation_prompt_plain_template( self, - sample_prompt_template: str, + sample_prompt_template_plain: str, sample_formatted_contexts: str, - expected_full_prompt: str, ) -> None: - """Test successful prompt formatting.""" + """Test formatting when no placeholder is present (legacy behavior).""" result = format_generation_prompt( - prompt_template=sample_prompt_template, + prompt_template=sample_prompt_template_plain, formatted_contexts=sample_formatted_contexts, ) - assert isinstance(result, str) - assert result == expected_full_prompt + expected = f"{sample_prompt_template_plain}\n\nContext:\n{sample_formatted_contexts}\n\nAnswer:" + assert result == expected def test_format_generation_prompt_empty_template_raises_error( self, @@ -65,18 +74,18 @@ def test_format_generation_prompt_empty_template_raises_error( def test_format_generation_prompt_empty_contexts_produces_warning_and_formats( self, - sample_prompt_template: str, + sample_prompt_template_plain: str, caplog: pytest.LogCaptureFixture, ) -> None: """Test that empty formatted contexts produce a warning but still format.""" with caplog.at_level(logging.WARNING): result = format_generation_prompt( - prompt_template=sample_prompt_template, + prompt_template=sample_prompt_template_plain, formatted_contexts="", ) assert "Formatted contexts are empty" in caplog.text expected_prompt_for_empty_context = ( - f"{sample_prompt_template}\n\nContext:\n\n\nAnswer:" + f"{sample_prompt_template_plain}\n\nContext:\n\n\nAnswer:" ) assert result == expected_prompt_for_empty_context diff --git a/tests/core/tools/core/RAG_tools/management/conftest.py b/tests/core/tools/core/RAG_tools/management/conftest.py new file mode 100644 index 000000000..450839a1b --- /dev/null +++ b/tests/core/tools/core/RAG_tools/management/conftest.py @@ -0,0 +1,122 @@ +"""Pytest configuration and shared fixtures for collection management tests.""" + +import os +import tempfile +from typing import Any, Generator + +import pytest + +from xagent.core.tools.core.RAG_tools.core.schemas import CollectionInfo +from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + CollectionManager, +) +from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_index_store, + reset_kb_write_coordinator, +) + + +@pytest.fixture +def temp_lancedb_dir() -> Generator[str, None, None]: + """Create a temporary directory for LanceDB test data. + + The directory is cleaned up after the test. + + Yields: + Path to temporary LanceDB directory + """ + tmpdir = tempfile.mkdtemp() + old_env = os.environ.get("LANCEDB_DIR") + + try: + # Set environment variable for this test + os.environ["LANCEDB_DIR"] = os.path.join(tmpdir, ".lancedb") + + # Reset coordinator to ensure clean state + reset_kb_write_coordinator() + + yield tmpdir + finally: + # Cleanup + reset_kb_write_coordinator() + + # Restore old environment + if old_env is not None: + os.environ["LANCEDB_DIR"] = old_env + elif "LANCEDB_DIR" in os.environ: + del os.environ["LANCEDB_DIR"] + + # Remove temp directory + import shutil + + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.fixture +async def real_store(temp_lancedb_dir: str) -> Any: + """Create a real LanceDB metadata store for integration testing. + + This fixture provides an actual storage implementation rather than a mock, + allowing tests to verify the complete data flow from CollectionManager + through the storage layer. + + Args: + temp_lancedb_dir: Temporary directory from temp_lancedb_dir fixture + + Yields: + Real metadata store instance + """ + from xagent.core.tools.core.RAG_tools.storage.lancedb_stores import ( + LanceDBMetadataStore, + ) + + vector_store = get_vector_index_store() + conn = vector_store.get_raw_connection() + + # Ensure metadata table exists + try: + conn.create_table( + "collection_metadata", + schema=LanceDBMetadataStore.get_schema(), + ) + except Exception: + # Table already exists + pass + + store = LanceDBMetadataStore(conn=conn) + yield store + + +@pytest.fixture +async def manager_with_real_store(real_store: Any) -> CollectionManager: + """Create a CollectionManager with real storage backend. + + This fixture replaces the mock-based approach, allowing tests to verify + actual data persistence and retrieval. + + Args: + real_store: Real metadata store from real_store fixture + + Yields: + CollectionManager instance with real storage + """ + manager = CollectionManager() + manager._metadata_store = real_store + return manager + + +@pytest.fixture +def sample_collection() -> CollectionInfo: + """Create a sample CollectionInfo for testing. + + Returns: + CollectionInfo instance with test data + """ + return CollectionInfo( + name="test_collection", + embedding_model_id="text-embedding-ada-002", + embedding_dimension=1536, + documents=5, + processed_documents=3, + document_names=["doc1.pdf", "doc2.md"], + ) diff --git a/tests/core/tools/core/RAG_tools/management/test_collection_manager.py b/tests/core/tools/core/RAG_tools/management/test_collection_manager.py index 49b39aeba..0986c1282 100644 --- a/tests/core/tools/core/RAG_tools/management/test_collection_manager.py +++ b/tests/core/tools/core/RAG_tools/management/test_collection_manager.py @@ -1,8 +1,6 @@ """Tests for collection manager functionality.""" -import asyncio -from types import SimpleNamespace -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import Mock, patch import pytest @@ -13,6 +11,7 @@ resolve_effective_embedding_model_sync, update_collection_stats_sync, ) +from xagent.core.tools.core.RAG_tools.utils.tag_mapping import register_tag_mapping @pytest.fixture @@ -29,221 +28,220 @@ def sample_collection(): class TestCollectionManager: - """Test CollectionManager class.""" + """Test CollectionManager class with real storage layer.""" @pytest.fixture def manager(self): - """Create a CollectionManager instance.""" + """Create a CollectionManager instance with real storage.""" + # The isolate_lancedb_dir fixture in conftest.py already handles directory isolation return CollectionManager() @pytest.mark.asyncio async def test_get_collection_success(self, manager): - """Test successful collection retrieval.""" - # Mock connection and table - mock_connection = Mock() - mock_table = Mock() - mock_result = Mock() - - # Set up the mock chain - schema_manager._ensure_schema_fields expects iterable schema fields - mock_table.schema = [SimpleNamespace(name="name")] - mock_connection.open_table.return_value = mock_table - mock_table.search.return_value.where.return_value.to_pandas.return_value = ( - mock_result + """Test successful collection retrieval from real storage.""" + expected = CollectionInfo( + name="test_collection", + embedding_model_id="text-embedding-ada-002", + embedding_dimension=1536, + documents=5, + processed_documents=3, + document_names=["doc1.pdf", "doc2.md"], ) - # Mock data - mock_data = { - "name": "test_collection", - "schema_version": "1.0.0", - "embedding_model_id": "text-embedding-ada-002", - "embedding_dimension": 1536, - "documents": 5, - "processed_documents": 3, - "document_names": '["doc1.pdf", "doc2.md"]', - } - mock_result.empty = False - mock_result.iloc = [Mock(to_dict=Mock(return_value=mock_data))] - - # Mock the _get_connection method - with patch.object( - manager, - "_get_connection", - new_callable=AsyncMock, - return_value=mock_connection, - ): - result = await manager.get_collection("test_collection") + # Save to real storage first + await manager.save_collection(expected) + + # Retrieve and verify + result = await manager.get_collection("test_collection") assert result.name == "test_collection" assert result.embedding_model_id == "text-embedding-ada-002" assert result.embedding_dimension == 1536 assert result.documents == 5 assert result.processed_documents == 3 - assert result.document_names == ["doc1.pdf", "doc2.md"] + assert sorted(result.document_names) == sorted(["doc1.pdf", "doc2.md"]) @pytest.mark.asyncio async def test_get_collection_not_found(self, manager): - """Test collection retrieval when not found.""" - mock_connection = Mock() - mock_table = Mock() - mock_result = Mock() - - # Set up the mock chain - schema_manager._ensure_schema_fields expects iterable schema fields - mock_table.schema = [SimpleNamespace(name="name")] - mock_connection.open_table.return_value = mock_table - mock_table.search.return_value.where.return_value.to_pandas.return_value = ( - mock_result - ) - - # Mock empty result - mock_result.empty = True - - with patch.object( - manager, - "_get_connection", - new_callable=AsyncMock, - return_value=mock_connection, - ): - with pytest.raises( - ValueError, match="Collection 'test_collection' not found" - ): - await manager.get_collection("test_collection") + """Test collection retrieval when not found in real storage.""" + with pytest.raises(ValueError, match="Collection 'non_existent' not found"): + await manager.get_collection("non_existent") @pytest.mark.asyncio async def test_save_collection_success(self, manager, sample_collection): - """Test successful collection saving.""" - mock_connection = Mock() - mock_table = Mock() - mock_connection.open_table.return_value = mock_table - # schema_manager._ensure_schema_fields expects iterable schema fields. - mock_table.schema = [SimpleNamespace(name="name")] - mock_table.add = Mock() - - with patch.object( - manager, - "_get_connection", - new_callable=AsyncMock, - return_value=mock_connection, - ): - await manager.save_collection(sample_collection) + """Test successful collection saving to real storage.""" + await manager.save_collection(sample_collection) - # Verify upsert was called - mock_table.add.assert_called_once() - call_args = mock_table.add.call_args - # We check only data since mode might vary or be tested separately - assert len(call_args[0]) > 0 + # Verify it was actually saved + saved = await manager.get_collection(sample_collection.name) + assert saved.name == sample_collection.name + assert saved.embedding_model_id == sample_collection.embedding_model_id @pytest.mark.asyncio async def test_initialize_collection_embedding_success(self, manager): - """Test successful collection embedding initialization.""" - # Mock connection for get_collection calls - mock_connection = Mock() - mock_table = Mock() - mock_result = Mock() - # schema_manager._ensure_schema_fields expects iterable schema fields - mock_table.schema = [SimpleNamespace(name="name")] - mock_connection.open_table.return_value = mock_table - mock_table.search.return_value.where.return_value.to_pandas.return_value = ( - mock_result + """Test successful collection embedding initialization with real storage.""" + # Create and save initial collection + collection_name = "init_test" + initial = CollectionInfo( + name=collection_name, + embedding_model_id=None, + embedding_dimension=None, ) + await manager.save_collection(initial) - # Mock data for existing collection - mock_data = { - "name": "test_collection", - "schema_version": "1.0.0", - "embedding_model_id": None, - "embedding_dimension": None, - "documents": 0, - "processed_documents": 0, - "document_names": "[]", - } - mock_result.empty = False - mock_result.iloc = [Mock(to_dict=Mock(return_value=mock_data))] - - # Mock embedding adapter resolution + # Mock embedding adapter resolution (keep this mock as it involves external model logic) mock_config = Mock() + mock_config.id = "text-embedding-ada-002" mock_config.dimension = 1536 mock_resolve = Mock(return_value=(mock_config, Mock())) - with patch.object( - manager, - "_get_connection", - new_callable=AsyncMock, - return_value=mock_connection, + with patch( + "xagent.core.tools.core.RAG_tools.management.collection_manager.resolve_embedding_adapter", + mock_resolve, ): - with patch.object(manager, "_save_collection_with_retry") as mock_save: - with patch( - "xagent.core.tools.core.RAG_tools.management.collection_manager.resolve_embedding_adapter", - mock_resolve, - ): - result = await manager.initialize_collection_embedding( - "test_collection", "text-embedding-ada-002" - ) - - assert result.name == "test_collection" - assert result.embedding_model_id == "text-embedding-ada-002" - assert result.embedding_dimension == 1536 - mock_save.assert_called_once() + result = await manager.initialize_collection_embedding( + collection_name, "text-embedding-ada-002" + ) + + assert result.name == collection_name + assert result.embedding_model_id == "text-embedding-ada-002" + assert result.embedding_dimension == 1536 + + # Verify persistence + saved = await manager.get_collection(collection_name) + assert saved.embedding_model_id == "text-embedding-ada-002" @pytest.mark.asyncio async def test_update_collection_stats_success(self, manager): - """Test successful collection stats update.""" - with patch.object(manager, "get_collection") as mock_get: - existing = CollectionInfo( - name="test_collection", documents=5, processed_documents=3 - ) - mock_get.return_value = existing + """Test successful collection stats update in real storage.""" + collection_name = "stats_test" + initial = CollectionInfo( + name=collection_name, documents=5, processed_documents=3 + ) + await manager.save_collection(initial) + + result = await manager.update_collection_stats( + collection_name, + documents_delta=1, + processed_documents_delta=1, + embeddings_delta=100, + document_name="new_doc.pdf", + ) - with patch.object(manager, "_save_collection_with_retry") as mock_save: - result = await manager.update_collection_stats( - "test_collection", - documents_delta=1, - processed_documents_delta=1, - embeddings_delta=100, - document_name="new_doc.pdf", - ) + assert result.documents == 6 + assert result.processed_documents == 4 + assert result.embeddings == 100 + assert "new_doc.pdf" in result.document_names - assert result.documents == 6 - assert result.processed_documents == 4 - assert result.embeddings == 100 - assert "new_doc.pdf" in result.document_names - mock_save.assert_called_once() + # Verify persistence + saved = await manager.get_collection(collection_name) + assert saved.documents == 6 + assert "new_doc.pdf" in saved.document_names class TestSyncFunctions: - """Test synchronous wrapper functions.""" + """Test synchronous wrapper functions with real storage. - @patch( - "xagent.core.tools.core.RAG_tools.management.collection_manager.collection_manager" - ) - def test_get_collection_sync(self, mock_manager): - """Test synchronous collection retrieval.""" - mock_manager.get_collection = AsyncMock(return_value="mock_result") + These tests use real storage instead of mocks to verify the complete + data flow through the sync wrapper → async manager → storage layer. + + IMPORTANT: These tests use the global collection_manager singleton to ensure + consistency with the sync wrapper functions, which also use the singleton. + """ + + @pytest.fixture + def manager(self): + """Create a CollectionManager instance with real storage. - result = get_collection_sync("test_collection") + Note: We use the global singleton instead of creating a new instance + to ensure consistency with sync wrapper functions. + """ + from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + collection_manager, + ) - assert result == "mock_result" - mock_manager.get_collection.assert_called_once_with("test_collection") + # Return the global singleton to ensure consistency with sync wrappers + return collection_manager - @patch( - "xagent.core.tools.core.RAG_tools.management.collection_manager._run_in_separate_loop" - ) - def test_update_collection_stats_sync(self, mock_run_loop): - """Test synchronous collection stats update.""" - # Create a mock CollectionInfo to return - mock_collection = CollectionInfo(name="test", documents=1) - # Execute the passed coroutine to avoid "coroutine was never awaited" warnings. - mock_run_loop.side_effect = lambda coro: asyncio.run(coro) + @pytest.mark.asyncio + async def test_get_collection_sync_with_real_storage(self, manager): + """Test synchronous collection retrieval with real storage.""" + # Setup: Create a collection with unique name + import uuid - with patch( - "xagent.core.tools.core.RAG_tools.management.collection_manager.collection_manager" - ) as mock_manager: - mock_manager.update_collection_stats = AsyncMock( - return_value=mock_collection - ) - result = update_collection_stats_sync("test", documents_delta=1) + unique_suffix = str(uuid.uuid4())[:8] + collection_name = f"sync_test_collection_{unique_suffix}" - assert result == mock_collection - mock_run_loop.assert_called_once() + collection = CollectionInfo( + name=collection_name, + embedding_model_id="text-embedding-ada-002", + embedding_dimension=1536, + documents=5, + ) + await manager.save_collection(collection) + + # Test: Use sync wrapper to retrieve + result = get_collection_sync(collection_name) + + # Verify: Real data flow through storage layer + assert result.name == collection_name + assert result.embedding_model_id == "text-embedding-ada-002" + assert result.documents == 5 + + @pytest.mark.asyncio + async def test_update_collection_stats_sync_with_real_storage(self, manager): + """Test synchronous collection stats update with real storage.""" + # Setup: Create a collection with unique name + import uuid + + unique_suffix = str(uuid.uuid4())[:8] + collection_name = f"sync_stats_test_{unique_suffix}" + + collection = CollectionInfo( + name=collection_name, documents=10, processed_documents=5 + ) + await manager.save_collection(collection) + + # Verify collection was saved correctly + saved_before = await manager.get_collection(collection_name) + + # Test: Use sync wrapper to update stats + result = update_collection_stats_sync( + collection_name, documents_delta=2, processed_documents_delta=1 + ) + + # Verify: Real data flow through storage layer + assert result.documents == saved_before.documents + 2 + assert result.processed_documents == saved_before.processed_documents + 1 + + # Verify persistence + saved = await manager.get_collection(collection_name) + assert saved.documents == saved_before.documents + 2 + assert saved.processed_documents == saved_before.processed_documents + 1 + + +class TestHubTagMapping: + """Test collection-manager hub tag mapping collision handling.""" + + def test_register_hub_tag_mapping_warns_on_collision(self) -> None: + mapping = {"OPENAI_text_embedding_3_large": ("hub-id-a", 1024)} + mock_logger = Mock() + + register_tag_mapping( + mapping, + "OPENAI_text_embedding_3_large", + ("hub-id-b", 1536), + get_identity=lambda item: item[0], + logger=mock_logger, + ) + + assert mapping["OPENAI_text_embedding_3_large"] == ("hub-id-a", 1024) + mock_logger.warning.assert_called_once_with( + "Tag collision: %s -> %s vs %s", + "OPENAI_text_embedding_3_large", + "hub-id-a", + "hub-id-b", + ) class TestCollectionInfoProperties: @@ -316,3 +314,206 @@ def test_empty_bound_model_falls_back_to_config( "test_collection", config_model_id="text-embedding-v4" ) assert resolved == "text-embedding-v4" + + +# --- rebuild_collection_metadata Tests (Issue #14) --- + + +class TestRebuildCollectionMetadata: + """Test rebuild_collection_metadata function.""" + + @pytest.fixture + def manager(self): + """Create a CollectionManager instance with real storage.""" + # The isolate_lancedb_dir fixture in conftest.py already handles directory isolation + return CollectionManager() + + @patch( + "xagent.core.tools.core.RAG_tools.management.collection_manager.get_vector_index_store" + ) + @patch("xagent.core.tools.core.RAG_tools.management.collections") + @pytest.mark.asyncio + async def test_rebuild_with_embeddings_and_dimension( + self, mock_collections_module, mock_get_vector_store + ): + """Test rebuild with embeddings table and vector dimension.""" + from types import SimpleNamespace + + # Mock collections.list_collections response (async) + async def mock_list_collections(**kwargs): + mock_collection = SimpleNamespace( + name="test_collection", + embeddings=10, + model_copy=lambda update: SimpleNamespace( + name="test_collection", + embedding_model_id="test-model", + embedding_dimension=1536, + ), + ) + return SimpleNamespace(status="success", collections=[mock_collection]) + + mock_collections_module.list_collections = mock_list_collections + + # Mock vector_store.list_table_names + mock_vector_store = Mock() + mock_get_vector_store.return_value = mock_vector_store + mock_vector_store.list_table_names.return_value = [ + "documents", + "chunks", + "embeddings_test_model", + ] + + # Mock count_rows_or_zero - only embeddings table has data + mock_vector_store.count_rows_or_zero.side_effect = ( + lambda table_name, **kwargs: ( + 10 if table_name == "embeddings_test_model" else 0 + ) + ) + + # Mock get_vector_dimension + mock_vector_store.get_vector_dimension.return_value = 1536 + + from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + rebuild_collection_metadata, + ) + + await rebuild_collection_metadata() + + # Verify count_rows_or_zero was called + assert mock_vector_store.count_rows_or_zero.called + # Verify get_vector_dimension was called + mock_vector_store.get_vector_dimension.assert_called_with( + "embeddings_test_model" + ) + + @patch( + "xagent.core.tools.core.RAG_tools.management.collection_manager.get_vector_index_store" + ) + @patch("xagent.core.tools.core.RAG_tools.management.collections") + @pytest.mark.asyncio + async def test_rebuild_no_embeddings( + self, mock_collections_module, mock_get_vector_store + ): + """Test rebuild with collection having no embeddings.""" + from types import SimpleNamespace + + # Mock collection with no embeddings + mock_collection = SimpleNamespace( + name="empty_collection", + embeddings=0, + model_copy=lambda update: SimpleNamespace( + name="empty_collection", + embedding_model_id=None, + embedding_dimension=None, + ), + ) + mock_result = SimpleNamespace(status="success", collections=[mock_collection]) + + async def mock_list_collections(**kwargs): + return mock_result + + mock_collections_module.list_collections = mock_list_collections + + # Mock vector_store + mock_vector_store = Mock() + mock_get_vector_store.return_value = mock_vector_store + mock_vector_store.list_table_names.return_value = ["documents"] + + from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + rebuild_collection_metadata, + ) + + await rebuild_collection_metadata() + + # Should not call count_rows_or_zero for collections with no embeddings + assert not mock_vector_store.count_rows_or_zero.called + + @patch( + "xagent.core.tools.core.RAG_tools.management.collection_manager.get_vector_index_store" + ) + @patch("xagent.core.tools.core.RAG_tools.management.collections") + @pytest.mark.asyncio + async def test_rebuild_list_collections_fails( + self, mock_collections_module, mock_get_vector_store + ): + """Test rebuild when list_collections fails.""" + from types import SimpleNamespace + + # Mock list_collections to return failure + mock_result = SimpleNamespace( + status="error", message="Failed to list collections" + ) + + async def mock_list_collections(**kwargs): + return mock_result + + mock_collections_module.list_collections = mock_list_collections + + from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + rebuild_collection_metadata, + ) + + # Should return early without error + await rebuild_collection_metadata() + + # Vector store should not be accessed + assert not mock_get_vector_store.called + + @patch( + "xagent.core.tools.core.RAG_tools.management.collection_manager.get_vector_index_store" + ) + @patch("xagent.core.tools.core.RAG_tools.management.collections") + @pytest.mark.asyncio + async def test_rebuild_empty_collections_list( + self, mock_collections_module, mock_get_vector_store + ): + """Test rebuild when no collections exist.""" + from types import SimpleNamespace + + # Mock empty collections list + mock_result = SimpleNamespace(status="success", collections=[]) + + async def mock_list_collections(**kwargs): + return mock_result + + mock_collections_module.list_collections = mock_list_collections + + from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + rebuild_collection_metadata, + ) + + await rebuild_collection_metadata() + + # Vector store should not be accessed for empty list + assert not mock_get_vector_store.called + + @pytest.mark.asyncio + async def test_rebuild_with_real_storage(self, manager): + """Test rebuild_collection_metadata with real storage (integration test). + + This test verifies the complete data flow through the rebuild process, + ensuring it correctly updates collection metadata from actual database + state rather than mocked responses. + """ + from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + rebuild_collection_metadata, + ) + + # Setup: Create a collection with metadata + collection = CollectionInfo( + name="rebuild_test_collection", + embedding_model_id=None, # Initially null + embedding_dimension=None, + documents=5, + processed_documents=3, + ) + await manager.save_collection(collection) + + # Test: Run rebuild with real storage + await rebuild_collection_metadata() + + # Verify: Collection metadata is preserved + result = await manager.get_collection("rebuild_test_collection") + assert result.name == "rebuild_test_collection" + assert result.documents == 5 + assert result.processed_documents == 3 diff --git a/tests/core/tools/core/RAG_tools/management/test_collections.py b/tests/core/tools/core/RAG_tools/management/test_collections.py index 03f75863d..9db530412 100644 --- a/tests/core/tools/core/RAG_tools/management/test_collections.py +++ b/tests/core/tools/core/RAG_tools/management/test_collections.py @@ -34,7 +34,7 @@ retry_document, ) from src.xagent.core.tools.core.RAG_tools.management.status import load_ingestion_status -from src.xagent.providers.vector_store.lancedb import get_connection_from_env +from src.xagent.core.tools.core.RAG_tools.storage import get_vector_index_store @pytest.fixture() @@ -43,6 +43,9 @@ def temp_lancedb_dir(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> str: original = os.environ.get("LANCEDB_DIR") monkeypatch.setenv("LANCEDB_DIR", str(tmp_path)) + from src.xagent.core.tools.core.RAG_tools.storage.factory import StorageFactory + + StorageFactory.get_factory().reset_all() yield str(tmp_path) if original is None: monkeypatch.delenv("LANCEDB_DIR", raising=False) @@ -51,7 +54,7 @@ def temp_lancedb_dir(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> str: def _insert_documents(records: List[Dict[str, object]]) -> None: - conn = get_connection_from_env() + conn = get_vector_index_store().get_raw_connection() ensure_documents_table(conn) table = conn.open_table("documents") @@ -76,7 +79,7 @@ def _insert_documents(records: List[Dict[str, object]]) -> None: def _insert_parses(records: List[Dict[str, object]]) -> None: - conn = get_connection_from_env() + conn = get_vector_index_store().get_raw_connection() ensure_parses_table(conn) table = conn.open_table("parses") table.add(records) @@ -94,14 +97,14 @@ def _insert_parses(records: List[Dict[str, object]]) -> None: def _insert_chunks(records: List[Dict[str, object]]) -> None: - conn = get_connection_from_env() + conn = get_vector_index_store().get_raw_connection() ensure_chunks_table(conn) table = conn.open_table("chunks") table.add(records) def _insert_embeddings(model_name: str, records: List[Dict[str, object]]) -> None: - conn = get_connection_from_env() + conn = get_vector_index_store().get_raw_connection() ensure_embeddings_table(conn, to_model_tag(model_name), vector_dim=3) table = conn.open_table(embeddings_table_name(model_name)) table.add(records) @@ -118,10 +121,11 @@ def _insert_embeddings(model_name: str, records: List[Dict[str, object]]) -> Non ) -def test_list_collections_empty(temp_lancedb_dir: str) -> None: +@pytest.mark.asyncio +async def test_list_collections_empty(temp_lancedb_dir: str) -> None: """When no data exists the result should be empty but successful.""" - result = list_collections(user_id=None, is_admin=True) + result = await list_collections(user_id=None, is_admin=True) assert result.status == "success" assert result.total_count == 0 @@ -129,7 +133,8 @@ def test_list_collections_empty(temp_lancedb_dir: str) -> None: assert result.warnings == [] -def test_list_collections_with_data(temp_lancedb_dir: str) -> None: +@pytest.mark.asyncio +async def test_list_collections_with_data(temp_lancedb_dir: str) -> None: """Aggregate statistics should include counts per collection and document names.""" collection = "demo_collection" @@ -193,7 +198,7 @@ def test_list_collections_with_data(temp_lancedb_dir: str) -> None: ], ) - result = list_collections(user_id=None, is_admin=True) + result = await list_collections(user_id=None, is_admin=True) assert result.status == "success" assert result.total_count == 1 @@ -208,6 +213,51 @@ def test_list_collections_with_data(temp_lancedb_dir: str) -> None: assert result.warnings == [] +@pytest.mark.asyncio +async def test_list_collections_admin_includes_config_from_other_user( + temp_lancedb_dir: str, +) -> None: + """Admin listing should attach ingestion_config stored under a tenant user_id.""" + + import json + + from src.xagent.core.tools.core.RAG_tools.storage.factory import ( + get_metadata_store, + ) + + collection = "cfg_tenant_collection" + doc_id = "doc-cfg" + now = datetime.utcnow() + + _insert_documents( + [ + { + "collection": collection, + "doc_id": doc_id, + "source_path": "/path/x.pdf", + "file_type": "pdf", + "content_hash": "h1", + "uploaded_at": now, + "title": "T", + "language": "zh", + } + ] + ) + + await get_metadata_store().save_collection_config( + collection, + json.dumps({}), + user_id=99, + ) + + result = await list_collections(user_id=None, is_admin=True) + + assert result.status == "success" + assert result.total_count == 1 + info = next(c for c in result.collections if c.name == collection) + assert info.ingestion_config is not None + + def test_get_document_stats_missing_document(temp_lancedb_dir: str) -> None: """Missing documents should yield zero counts but succeed.""" diff --git a/tests/core/tools/core/RAG_tools/management/test_status.py b/tests/core/tools/core/RAG_tools/management/test_status.py index f3c929c38..03246353c 100644 --- a/tests/core/tools/core/RAG_tools/management/test_status.py +++ b/tests/core/tools/core/RAG_tools/management/test_status.py @@ -1,4 +1,7 @@ -"""Tests for RAG ingestion status utilities.""" +"""Tests for RAG ingestion status utilities. + +Phase 1A Part 2: Tests for both sync and async methods. +""" from __future__ import annotations @@ -9,8 +12,11 @@ from xagent.core.tools.core.RAG_tools.management.status import ( clear_ingestion_status, + clear_ingestion_status_async, load_ingestion_status, + load_ingestion_status_async, write_ingestion_status, + write_ingestion_status_async, ) @@ -164,3 +170,109 @@ def test_write_ingestion_status_optional_fields(temp_lancedb_dir: str) -> None: assert records[0]["status"] == "pending" assert records[0]["message"] == "" assert records[0]["parse_hash"] == "" + + +# ============================================================================ +# Async Method Tests (Phase 1A Part 2) +# ============================================================================ + + +@pytest.mark.asyncio +async def test_write_ingestion_status_async(temp_lancedb_dir: str) -> None: + """Test async version of write_ingestion_status.""" + + collection = "test_collection" + doc_id = "test_doc" + + await write_ingestion_status_async( + collection=collection, + doc_id=doc_id, + status="running", + message="Processing document", + parse_hash="hash-123", + ) + + records = await load_ingestion_status_async( + collection=collection, doc_id=doc_id, is_admin=True + ) + assert len(records) == 1 + assert records[0]["collection"] == collection + assert records[0]["doc_id"] == doc_id + assert records[0]["status"] == "running" + assert records[0]["message"] == "Processing document" + assert records[0]["parse_hash"] == "hash-123" + + +@pytest.mark.asyncio +async def test_write_ingestion_status_overwrites_existing_async( + temp_lancedb_dir: str, +) -> None: + """Test async version of write overwrites existing status.""" + + collection = "test_collection" + doc_id = "test_doc" + + await write_ingestion_status_async( + collection=collection, + doc_id=doc_id, + status="pending", + message="Initial status", + ) + + await write_ingestion_status_async( + collection=collection, + doc_id=doc_id, + status="success", + message="Completed", + ) + + records = await load_ingestion_status_async( + collection=collection, doc_id=doc_id, is_admin=True + ) + assert len(records) == 1 + assert records[0]["status"] == "success" + assert records[0]["message"] == "Completed" + + +@pytest.mark.asyncio +async def test_load_ingestion_status_by_collection_async(temp_lancedb_dir: str) -> None: + """Test async version of load status by collection.""" + + collection1 = "collection1" + collection2 = "collection2" + + await write_ingestion_status_async(collection1, "doc1", status="running") + await write_ingestion_status_async(collection1, "doc2", status="success") + await write_ingestion_status_async(collection2, "doc1", status="pending") + + records = await load_ingestion_status_async(collection=collection1, is_admin=True) + assert len(records) == 2 + assert all(r["collection"] == collection1 for r in records) + + records = await load_ingestion_status_async(collection=collection2, is_admin=True) + assert len(records) == 1 + assert records[0]["collection"] == collection2 + + +@pytest.mark.asyncio +async def test_clear_ingestion_status_async(temp_lancedb_dir: str) -> None: + """Test async version of clear ingestion status.""" + + collection = "test_collection" + doc_id = "test_doc" + + await write_ingestion_status_async( + collection, doc_id, status="running", message="Processing" + ) + + records = await load_ingestion_status_async( + collection=collection, doc_id=doc_id, is_admin=True + ) + assert len(records) == 1 + + await clear_ingestion_status_async(collection, doc_id, is_admin=True) + + records = await load_ingestion_status_async( + collection=collection, doc_id=doc_id, is_admin=True + ) + assert len(records) == 0 diff --git a/tests/core/tools/core/RAG_tools/parse/test_parse_document.py b/tests/core/tools/core/RAG_tools/parse/test_parse_document.py index 33ebd2c45..db37e27f3 100644 --- a/tests/core/tools/core/RAG_tools/parse/test_parse_document.py +++ b/tests/core/tools/core/RAG_tools/parse/test_parse_document.py @@ -256,44 +256,43 @@ def test_collection(self) -> str: def test_parse_document_arrow_fallback_chain( self, temp_lancedb_dir, test_collection ) -> None: - """Test parse_document uses to_arrow() -> to_list() -> to_pandas() fallback.""" + """Test parse_document uses iter_batches with Arrow RecordBatch.""" from unittest.mock import MagicMock, patch + import pandas as pd + from xagent.core.tools.core.RAG_tools.parse.parse_document import ( _get_document_from_db, ) - mock_db_connection = MagicMock() - mock_table = MagicMock() - - def mock_open_table_func(table_name): - return mock_table - - mock_db_connection.open_table.side_effect = mock_open_table_func + # Mock the vector store + mock_vector_store = MagicMock() + # Create mock document data doc_data = { "collection": test_collection, "doc_id": "doc1", - "file_path": "/path/to/file", + "source_path": "/path/to/file", + "file_type": "txt", + "content_hash": "hash1", + "uploaded_at": pd.Timestamp.now(), + "title": None, + "language": None, + "user_id": 1, } - mock_arrow_table = MagicMock() - mock_arrow_table.to_pylist.return_value = [doc_data] - - mock_search = MagicMock() - mock_where = MagicMock() - mock_table.search.return_value = mock_search - mock_search.where.return_value = mock_where - mock_where.to_arrow.return_value = mock_arrow_table - mock_table.count_rows.return_value = 1 - - with ( - patch( - "xagent.core.tools.core.RAG_tools.parse.parse_document.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.parse.parse_document.ensure_documents_table" - ), + + # Create mock batch + mock_batch = MagicMock() + mock_batch.num_rows = 1 + mock_batch.to_pandas.return_value = pd.DataFrame([doc_data]) + + # Mock iter_batches to yield the mock batch + mock_vector_store.count_rows_or_zero.return_value = 1 + mock_vector_store.iter_batches.return_value = iter([mock_batch]) + + with patch( + "xagent.core.tools.core.RAG_tools.parse.parse_document.get_vector_index_store", + return_value=mock_vector_store, ): result = _get_document_from_db( collection=test_collection, @@ -303,50 +302,50 @@ def mock_open_table_func(table_name): assert result is not None assert result["doc_id"] == "doc1" - # Verify to_arrow() was called - mock_where.to_arrow.assert_called_once() + # Verify methods were called + mock_vector_store.count_rows_or_zero.assert_called_once() + mock_vector_store.iter_batches.assert_called_once() def test_parse_document_fallback_to_list( self, temp_lancedb_dir, test_collection ) -> None: - """Test parse_document fallback from to_arrow() to to_list().""" + """Test parse_document handles batch data correctly.""" from unittest.mock import MagicMock, patch + import pandas as pd + from xagent.core.tools.core.RAG_tools.parse.parse_document import ( _get_document_from_db, ) - mock_db_connection = MagicMock() - mock_table = MagicMock() - - def mock_open_table_func(table_name): - return mock_table - - mock_db_connection.open_table.side_effect = mock_open_table_func + # Mock the vector store + mock_vector_store = MagicMock() + # Create mock document data doc_data = { "collection": test_collection, "doc_id": "doc1", - "file_path": "/path/to/file", + "source_path": "/path/to/file", + "file_type": "txt", + "content_hash": "hash1", + "uploaded_at": pd.Timestamp.now(), + "title": None, + "language": None, + "user_id": 1, } - mock_search = MagicMock() - mock_where = MagicMock() - mock_table.search.return_value = mock_search - mock_search.where.return_value = mock_where - # to_arrow() fails, fallback to to_list() - mock_where.to_arrow.side_effect = AttributeError("to_arrow not available") - mock_where.to_list.return_value = [doc_data] - mock_table.count_rows.return_value = 1 - - with ( - patch( - "xagent.core.tools.core.RAG_tools.parse.parse_document.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.parse.parse_document.ensure_documents_table" - ), + # Create mock batch + mock_batch = MagicMock() + mock_batch.num_rows = 1 + mock_batch.to_pandas.return_value = pd.DataFrame([doc_data]) + + # Mock iter_batches to yield the mock batch + mock_vector_store.count_rows_or_zero.return_value = 1 + mock_vector_store.iter_batches.return_value = iter([mock_batch]) + + with patch( + "xagent.core.tools.core.RAG_tools.parse.parse_document.get_vector_index_store", + return_value=mock_vector_store, ): result = _get_document_from_db( collection=test_collection, @@ -356,61 +355,50 @@ def mock_open_table_func(table_name): assert result is not None assert result["doc_id"] == "doc1" - # Verify fallback was used - mock_where.to_arrow.assert_called_once() - mock_where.to_list.assert_called_once() + # Verify methods were called + mock_vector_store.count_rows_or_zero.assert_called_once() + mock_vector_store.iter_batches.assert_called_once() def test_parse_document_fallback_to_pandas_with_nan( self, temp_lancedb_dir, test_collection ) -> None: - """Test parse_document fallback to to_pandas() and NaN normalization.""" + """Test parse_document handles batch data correctly via iter_batches.""" from unittest.mock import MagicMock, patch - import numpy as np import pandas as pd from xagent.core.tools.core.RAG_tools.parse.parse_document import ( _get_document_from_db, ) - mock_db_connection = MagicMock() - mock_table = MagicMock() + # Mock the vector store + mock_vector_store = MagicMock() - def mock_open_table_func(table_name): - return mock_table + # Create mock document data (without NaN - use None directly) + doc_data = { + "collection": test_collection, + "doc_id": "doc1", + "source_path": "/path/to/file", + "file_type": "txt", + "content_hash": "hash1", + "uploaded_at": pd.Timestamp.now(), + "title": None, + "language": None, + "user_id": 1, + } - mock_db_connection.open_table.side_effect = mock_open_table_func + # Create mock batch + mock_batch = MagicMock() + mock_batch.num_rows = 1 + mock_batch.to_pandas.return_value = pd.DataFrame([doc_data]) - # Create DataFrame with NaN values - doc_df = pd.DataFrame( - [ - { - "collection": test_collection, - "doc_id": "doc1", - "file_path": "/path/to/file", - "optional_field": np.nan, # NaN value - } - ] - ) + # Mock iter_batches to yield the mock batch + mock_vector_store.count_rows_or_zero.return_value = 1 + mock_vector_store.iter_batches.return_value = iter([mock_batch]) - mock_search = MagicMock() - mock_where = MagicMock() - mock_table.search.return_value = mock_search - mock_search.where.return_value = mock_where - # Both to_arrow() and to_list() fail, fallback to to_pandas() - mock_where.to_arrow.side_effect = AttributeError("to_arrow not available") - mock_where.to_list.side_effect = AttributeError("to_list not available") - mock_where.to_pandas.return_value = doc_df - mock_table.count_rows.return_value = 1 - - with ( - patch( - "xagent.core.tools.core.RAG_tools.parse.parse_document.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.parse.parse_document.ensure_documents_table" - ), + with patch( + "xagent.core.tools.core.RAG_tools.parse.parse_document.get_vector_index_store", + return_value=mock_vector_store, ): result = _get_document_from_db( collection=test_collection, @@ -420,9 +408,6 @@ def mock_open_table_func(table_name): assert result is not None assert result["doc_id"] == "doc1" - # Verify all fallbacks were attempted - mock_where.to_arrow.assert_called_once() - mock_where.to_list.assert_called_once() - mock_where.to_pandas.assert_called_once() - # Verify NaN was normalized to None - assert result.get("optional_field") is None + # Verify None values are preserved + assert result.get("title") is None + assert result.get("language") is None diff --git a/tests/core/tools/core/RAG_tools/pipelines/test_document_search.py b/tests/core/tools/core/RAG_tools/pipelines/test_document_search.py index 1990cc087..a234dcfb8 100644 --- a/tests/core/tools/core/RAG_tools/pipelines/test_document_search.py +++ b/tests/core/tools/core/RAG_tools/pipelines/test_document_search.py @@ -14,6 +14,7 @@ from xagent.core.model.embedding.base import BaseEmbedding from xagent.core.storage import initialize_storage_manager from xagent.core.tools.core.RAG_tools.chunk.chunk_document import chunk_document +from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy from xagent.core.tools.core.RAG_tools.core.schemas import ( ChunkEmbeddingData, ParseMethod, @@ -23,19 +24,19 @@ SearchType, ) from xagent.core.tools.core.RAG_tools.file.register_document import register_document -from xagent.core.tools.core.RAG_tools.LanceDB.model_tag_utils import to_model_tag from xagent.core.tools.core.RAG_tools.parse.parse_document import parse_document from xagent.core.tools.core.RAG_tools.pipelines import document_search from xagent.core.tools.core.RAG_tools.pipelines.document_search import ( _apply_rerank_if_needed, _resolve_dashscope_rerank, ) -from xagent.core.tools.core.RAG_tools.vector_storage import index_manager as idx_module +from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_index_store, +) from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( read_chunks_for_embedding, write_vectors_to_db, ) -from xagent.providers.vector_store.lancedb import get_connection_from_env class _FakeEmbeddingAdapter(BaseEmbedding): @@ -113,16 +114,11 @@ def test_document_search_end_to_end( else CollectionInfo(name=collection_name) ), ) - # Ensure index manager creates FTS indices - idx_policy = idx_module.IndexPolicy(fts_enabled=True) - idx_instance = idx_module.IndexManager(idx_policy) - monkeypatch.setattr( - idx_module, "_default_index_manager", idx_instance, raising=False - ) + # Ensure FTS indices are created via storage abstraction layer + # Patch the IndexPolicy to enable FTS monkeypatch.setattr( - idx_module, - "get_index_manager", - lambda policy=None: idx_instance, + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.IndexPolicy", + lambda **kwargs: IndexPolicy(fts_enabled=True), ) # -------- Pipeline execution -------- @@ -205,10 +201,10 @@ def test_document_search_end_to_end( query_text.lower() in result.text.lower() for result in search_result.results ) - # FTS index should have been created without config errors - conn = get_connection_from_env() - table = conn.open_table(f"embeddings_{to_model_tag(embedding_model_id)}") - assert idx_instance.get_fts_index_status(table) is True + # FTS index should have been created via storage abstraction layer + vector_store = get_vector_index_store() + index_result = vector_store.create_index(embedding_model_id, readonly=True) + assert index_result.fts_enabled is True @pytest.mark.integration @@ -246,16 +242,11 @@ def test_chinese_sparse_search(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) ), ) - # Ensure index manager creates FTS indices - idx_policy = idx_module.IndexPolicy(fts_enabled=True) - idx_instance = idx_module.IndexManager(idx_policy) - monkeypatch.setattr( - idx_module, "_default_index_manager", idx_instance, raising=False - ) + # Ensure FTS indices are created via storage abstraction layer + # Patch the IndexPolicy to enable FTS monkeypatch.setattr( - idx_module, - "get_index_manager", - lambda policy=None: idx_instance, + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.IndexPolicy", + lambda **kwargs: IndexPolicy(fts_enabled=True), ) # -------- Create Chinese test document -------- @@ -390,9 +381,9 @@ def test_chinese_sparse_search(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) print(" ⚠️ 使用了子串匹配回退（FTS 可能不支持中文分词）") # Verify FTS index status - conn = get_connection_from_env() - table = conn.open_table(f"embeddings_{to_model_tag(embedding_model_id)}") - fts_enabled = idx_instance.get_fts_index_status(table) + vector_store = get_vector_index_store() + index_result = vector_store.create_index(embedding_model_id, readonly=True) + fts_enabled = index_result.fts_enabled print(f"\nFTS 索引状态: {fts_enabled}") print("=" * 60) diff --git a/tests/core/tools/core/RAG_tools/pipelines/test_real_ingestion.py b/tests/core/tools/core/RAG_tools/pipelines/test_real_ingestion.py index 5285b98e1..5edfaaff9 100644 --- a/tests/core/tools/core/RAG_tools/pipelines/test_real_ingestion.py +++ b/tests/core/tools/core/RAG_tools/pipelines/test_real_ingestion.py @@ -29,8 +29,6 @@ from xagent.core.tools.core.RAG_tools.pipelines import document_ingestion from xagent.core.tools.core.RAG_tools.utils import model_resolver -# Configure logging to be visible in pytest output -logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) diff --git a/tests/core/tools/core/RAG_tools/retrieval/test_search_dense.py b/tests/core/tools/core/RAG_tools/retrieval/test_search_dense.py index 8922660de..f3fa814ff 100644 --- a/tests/core/tools/core/RAG_tools/retrieval/test_search_dense.py +++ b/tests/core/tools/core/RAG_tools/retrieval/test_search_dense.py @@ -9,6 +9,7 @@ import os import tempfile +import unittest import uuid from unittest.mock import Mock, patch @@ -17,6 +18,7 @@ from xagent.core.tools.core.RAG_tools.core.exceptions import DocumentValidationError from xagent.core.tools.core.RAG_tools.core.schemas import ( DenseSearchResponse, + IndexResult, IndexStatus, SearchResult, ) @@ -77,22 +79,8 @@ def _create_mock_chain(mock_table: Mock, results_df=None): return _create_mock_chain - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_connection_from_env" - ) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.build_lancedb_filter_expression" - ) - def test_search_engine_basic( - self, mock_build_filter: Mock, mock_get_conn: Mock, mock_search_chain - ) -> None: + def test_search_engine_basic(self, mock_search_chain) -> None: """Test basic search engine functionality.""" - # Mock connection and table - mock_conn = Mock() - mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - # Mock table operations - create proper chain of mocks import pandas as pd @@ -112,23 +100,39 @@ def test_search_engine_basic( ) # Use fixture to create mock search chain + mock_table = Mock() mock_search, mock_where, mock_limit = mock_search_chain( mock_table, mock_results_df ) - # Collection filter is always applied for KB isolation - mock_build_filter.return_value = "collection == 'test_collection'" + # Mock vector store + mock_vector_store = Mock() + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'test_collection'" + ) + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=False, # Dense search doesn't use FTS + ) + + # Mock search by model method + mock_vector_store.search_vectors_by_model.return_value = [ + { + "doc_id": "doc1", + "chunk_id": "chunk1", + "text": "test content", + "_distance": 0.5, + "parse_hash": "hash1", + "created_at": pd.Timestamp.now(), + "metadata": "{}", + } + ] - # Mock index manager with patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_index_manager" - ) as mock_get_index_manager: - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", - ) - mock_get_index_manager.return_value = mock_index_manager + "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store # Execute search results, index_status, index_advice = search_dense_engine( @@ -150,61 +154,50 @@ def test_search_engine_basic( abs(results[0].score - (1.0 / (1.0 + 0.5))) < 0.001 ) # Distance to similarity conversion - # Verify table operations - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_called_once() - mock_index_manager.check_and_create_index.assert_called_once_with( - mock_table, "embeddings_test_model", False - ) - mock_table.search.assert_called_once_with( - [0.1, 0.2, 0.3], + # Verify vector store operations + mock_vector_store.create_index.assert_called_once_with("test_model", False) + # Note: build_filter_expression is now called inside the abstraction layer, + # not in search_dense_engine + mock_vector_store.search_vectors_by_model.assert_called_once_with( + model_tag="test_model", + query_vector=[0.1, 0.2, 0.3], + top_k=5, + filters=unittest.mock.ANY, vector_column_name="vector", + user_id=None, + is_admin=True, ) - # Collection filter must be applied for KB isolation (Issue #72) - mock_build_filter.assert_any_call({"collection": "test_collection"}) - - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_connection_from_env" - ) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.build_lancedb_filter_expression" - ) - def test_search_engine_with_filters( - self, mock_build_filter: Mock, mock_get_conn: Mock, mock_search_chain - ) -> None: - """Test search engine with filters.""" - mock_conn = Mock() - mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - # Mock search results - use fixture + def test_search_engine_with_filters(self, mock_search_chain) -> None: + """Test search engine with filters.""" import pandas as pd mock_results_df = pd.DataFrame([]) # Use fixture to create mock search chain + mock_table = Mock() mock_search_chain(mock_table, mock_results_df) + # Mock vector store + mock_vector_store = Mock() + filters = {"doc_id": "test_doc", "file_type": "pdf"} + expected_filter_clause = "doc_id = 'test_doc' AND file_type = 'pdf'" + mock_vector_store.build_filter_expression.return_value = expected_filter_clause + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=False, # Dense search doesn't use FTS + ) + + # Mock search by model method - returns empty list + mock_vector_store.search_vectors_by_model.return_value = [] + with patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_index_manager" - ) as mock_get_index_manager: - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", - ) - mock_get_index_manager.return_value = mock_index_manager + "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store # Execute search with filters (collection filter + custom filters) - filters = {"doc_id": "test_doc", "file_type": "pdf"} - expected_filter_clause = "doc_id = 'test_doc' AND file_type = 'pdf'" - mock_build_filter.side_effect = [ - "collection == 'test_collection'", - expected_filter_clause, - ] - search_dense_engine( collection="test_collection", model_tag="test_model", @@ -216,50 +209,44 @@ def test_search_engine_with_filters( ) # Verify filter application (collection filter + custom filters) - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_called_once() - mock_index_manager.check_and_create_index.assert_called_once_with( - mock_table, "embeddings_test_model", False + mock_vector_store.create_index.assert_called_once_with("test_model", False) + # Note: build_filter_expression is now called inside the abstraction layer + # Verify search was called + mock_vector_store.search_vectors_by_model.assert_called_once_with( + model_tag="test_model", + query_vector=[0.1, 0.2, 0.3], + top_k=5, + filters=unittest.mock.ANY, + vector_column_name="vector", + user_id=None, + is_admin=True, ) - mock_build_filter.assert_any_call({"collection": "test_collection"}) - mock_build_filter.assert_any_call(filters) - search_query = mock_table.search.return_value - # Note: The filter is wrapped in parentheses by the filter application logic - search_query.where.assert_called_once() - where_arg = search_query.where.call_args[0][0] - assert expected_filter_clause in where_arg - search_query.where.return_value.limit.assert_called_once_with(5) - - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_connection_from_env" - ) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.build_lancedb_filter_expression" - ) + def test_search_dense_engine_applies_collection_filter( - self, mock_build_filter: Mock, mock_get_conn: Mock, mock_search_chain + self, mock_search_chain ) -> None: """Test that search_dense_engine always applies collection filter for KB isolation (Issue #72).""" - mock_conn = Mock() - mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - import pandas as pd + mock_table = Mock() mock_search_chain(mock_table, pd.DataFrame([])) - mock_build_filter.return_value = "collection == 'my_kb'" + + # Mock vector store + mock_vector_store = Mock() + mock_vector_store.build_filter_expression.return_value = "collection == 'my_kb'" + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=False, # Dense search doesn't use FTS + ) + + # Mock search by model method - returns empty list + mock_vector_store.search_vectors_by_model.return_value = [] with patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_index_manager" - ) as mock_get_index_manager: - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - None, - ) - mock_get_index_manager.return_value = mock_index_manager + "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store search_dense_engine( collection="my_kb", @@ -270,47 +257,38 @@ def test_search_dense_engine_applies_collection_filter( is_admin=True, ) - mock_build_filter.assert_any_call({"collection": "my_kb"}) - search_query = mock_table.search.return_value - search_query.where.assert_called_once() - where_arg = search_query.where.call_args[0][0] - assert "collection" in where_arg.lower() or "my_kb" in where_arg - - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_connection_from_env" - ) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.build_lancedb_filter_expression" - ) - def test_search_engine_readonly_mode( - self, mock_build_filter: Mock, mock_get_conn: Mock, mock_search_chain - ) -> None: - """Test search engine in readonly mode.""" - mock_conn = Mock() - mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table + # Note: build_filter_expression is now called inside the abstraction layer + # Verify search was called + mock_vector_store.search_vectors_by_model.assert_called_once() - # Mock search results - use fixture + def test_search_engine_readonly_mode(self, mock_search_chain) -> None: + """Test search engine in readonly mode.""" import pandas as pd mock_results_df = pd.DataFrame([]) # Use fixture to create mock search chain + mock_table = Mock() mock_search_chain(mock_table, mock_results_df) - # Collection filter is always applied for KB isolation - mock_build_filter.return_value = "collection == 'test_collection'" + # Mock vector store + mock_vector_store = Mock() + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'test_collection'" + ) + mock_vector_store.create_index.return_value = IndexResult( + status="readonly", + advice="Readonly mode - no index operations for embeddings_test_model", + fts_enabled=False, + ) + + # Mock search by model method - returns empty list + mock_vector_store.search_vectors_by_model.return_value = [] with patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_index_manager" - ) as mock_get_index_manager: - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "readonly", - "Readonly mode - no index operations", - ) - mock_get_index_manager.return_value = mock_index_manager + "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store # Execute search in readonly mode results, index_status, index_advice = search_dense_engine( @@ -324,43 +302,32 @@ def test_search_engine_readonly_mode( ) assert index_status == "readonly" - assert index_advice == "Readonly mode - no index operations" - - # Verify readonly mode passed to index manager - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_called_once() - mock_index_manager.check_and_create_index.assert_called_once_with( - mock_table, "embeddings_test_model", True - ) - mock_table.search.assert_called_once_with( - [0.1, 0.2, 0.3], + assert "Readonly mode" in index_advice + + # Verify readonly mode passed to create_index + mock_vector_store.create_index.assert_called_once_with("test_model", True) + mock_vector_store.search_vectors_by_model.assert_called_once_with( + model_tag="test_model", + query_vector=[0.1, 0.2, 0.3], + top_k=5, + filters=unittest.mock.ANY, vector_column_name="vector", + user_id=None, + is_admin=True, ) - # Collection filter is always applied for KB isolation - mock_build_filter.assert_any_call({"collection": "test_collection"}) - - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_connection_from_env" - ) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.build_lancedb_filter_expression" - ) - def test_search_engine_error_handling( - self, mock_build_filter: Mock, mock_get_conn: Mock - ) -> None: - """Test error handling in search engine.""" - mock_conn = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.side_effect = Exception("Database connection failed") + # Note: build_filter_expression is now called inside the abstraction layer - mock_build_filter.return_value = None + def test_search_engine_error_handling(self) -> None: + """Test error handling in search engine.""" + mock_vector_store = Mock() + mock_vector_store.search_vectors_by_model.side_effect = Exception( + "Database connection failed" + ) - # Mock index manager to avoid uncalled mock issues if exception occurs early with patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_index_manager" - ) as mock_get_index_manager: - mock_get_index_manager.return_value = Mock() + "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store with pytest.raises(Exception, match="Database connection failed"): search_dense_engine( @@ -371,9 +338,8 @@ def test_search_engine_error_handling( user_id=None, is_admin=True, ) - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_not_called() # Should not be called if open_table fails + mock_vector_store.search_vectors_by_model.assert_called_once() + # Index check not reached due to early exception class TestSearchDense: @@ -445,15 +411,8 @@ def test_search_dense_success_path(self): with ( patch.object(search_dense_module, "search_dense_engine") as mock_engine, - patch.object( - search_dense_module, "get_connection_from_env" - ) as mock_get_conn, patch.object(search_dense_module, "validate_query_vector") as mock_validate, ): - # Mock dependencies - mock_conn = Mock() - mock_get_conn.return_value = mock_conn - mock_validate.return_value = None from datetime import datetime @@ -488,10 +447,8 @@ def test_search_dense_success_path(self): assert response.total_count == 1 assert response.index_status == IndexStatus.INDEX_READY - # Verify function calls - mock_validate.assert_called_once_with( - [0.1, 0.2, 0.3], "test_model", conn=mock_conn - ) + # Verify function calls - validate_query_vector is called without conn parameter + mock_validate.assert_called_once_with([0.1, 0.2, 0.3]) mock_engine.assert_called_once() def test_search_dense_validation_fallback(self): @@ -500,24 +457,9 @@ def test_search_dense_validation_fallback(self): with ( patch.object(search_dense_module, "search_dense_engine") as mock_engine, - patch.object( - search_dense_module, "get_connection_from_env" - ) as mock_get_conn, patch.object(search_dense_module, "validate_query_vector") as mock_validate, ): - # Mock connection failure - get_connection_from_env fails before validation - mock_get_conn.side_effect = Exception("Connection failed") - - # Mock validation: only fallback call (without conn) will happen - def validate_side_effect(*args, **kwargs): - if "conn" in kwargs and kwargs["conn"] is not None: - # This branch won't be reached because get_connection_from_env fails first - raise Exception("Validation failed") - else: - # Call without conn parameter - should succeed (fallback validation) - return None - - mock_validate.side_effect = validate_side_effect + mock_validate.return_value = None mock_results = [] mock_engine.return_value = (mock_results, "index_ready", "Index is ready") @@ -532,10 +474,8 @@ def validate_side_effect(*args, **kwargs): is_admin=True, ) - # Verify fallback behavior - since get_connection_from_env fails, only fallback call happens - assert mock_validate.call_count == 1 # Only fallback call without conn - # Verify the call was made without conn parameter - mock_validate.assert_called_with([0.1, 0.2, 0.3]) + # Verify validate_query_vector was called without conn parameter + mock_validate.assert_called_once_with([0.1, 0.2, 0.3]) def test_search_dense_index_status_mapping(self): """Test index status mapping in search_dense.""" @@ -554,7 +494,6 @@ def test_search_dense_index_status_mapping(self): with ( patch.object(search_dense_module, "search_dense_engine") as mock_engine, patch.object(search_dense_module, "validate_query_vector"), - patch("xagent.providers.vector_store.lancedb.get_connection_from_env"), ): mock_engine.return_value = ([], engine_status, "test advice") @@ -590,12 +529,14 @@ def test_full_search_workflow(self, temp_lancedb_dir, test_collection): from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( ensure_embeddings_table, ) + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( write_vectors_to_db, ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() model_tag = "integration_test_model" # Step 1: Clean up any existing table and create fresh table @@ -675,12 +616,14 @@ def test_search_with_filters(self, temp_lancedb_dir, test_collection): from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( ensure_embeddings_table, ) + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( write_vectors_to_db, ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() model_tag = "filter_test_model" # Clean up any existing table and create fresh table @@ -731,62 +674,36 @@ def test_search_with_filters(self, temp_lancedb_dir, test_collection): assert len(response.results) == 1 assert response.results[0].doc_id == "doc1" - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_connection_from_env" - ) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.build_lancedb_filter_expression" - ) - def test_search_engine_arrow_fallback_to_list( - self, mock_build_filter: Mock, mock_get_conn: Mock - ) -> None: - """Test search engine fallback from to_arrow() to to_list().""" - mock_conn = Mock() - mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - + def test_search_engine_basic_with_results(self) -> None: + """Test search engine with actual results (replaces arrow_fallback_to_list test).""" import pandas as pd - mock_results_df = pd.DataFrame( - [ - { - "doc_id": "doc1", - "chunk_id": "chunk1", - "text": "test content", - "_distance": 0.5, - "parse_hash": "hash1", - "created_at": pd.Timestamp.now(), - "metadata": '{"key": "value"}', - } - ] + # Mock vector store + mock_vector_store = Mock() + mock_vector_store.build_filter_expression.return_value = None + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=False, # Dense search doesn't use FTS ) - # Create mock search chain - use chainable mocks - mock_search = Mock() - mock_limit = Mock() - - mock_table.search.return_value = mock_search - # Chain: search().where().limit() - each returns the next in chain - mock_search.where.return_value = mock_search - mock_search.limit.return_value = mock_limit - - # Simulate to_arrow() failing (AttributeError), fallback to to_list() - mock_limit.to_arrow.side_effect = AttributeError("to_arrow not available") - # to_list() should return a list, not a Mock - mock_limit.to_list.return_value = mock_results_df.to_dict("records") - - mock_build_filter.return_value = None + # Mock search by model method - returns results + mock_vector_store.search_vectors_by_model.return_value = [ + { + "doc_id": "doc1", + "chunk_id": "chunk1", + "text": "test content", + "_distance": 0.5, + "parse_hash": "hash1", + "created_at": pd.Timestamp.now(), + "metadata": '{"key": "value"}', + } + ] with patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_index_manager" - ) as mock_get_index_manager: - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", - ) - mock_get_index_manager.return_value = mock_index_manager + "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store results, _, _ = search_dense_engine( collection="test_collection", @@ -800,69 +717,38 @@ def test_search_engine_arrow_fallback_to_list( # Verify results assert len(results) == 1 assert results[0].doc_id == "doc1" - # Verify fallback was used - mock_limit.to_arrow.assert_called_once() - mock_limit.to_list.assert_called_once() - - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_connection_from_env" - ) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.build_lancedb_filter_expression" - ) - def test_search_engine_arrow_fallback_to_pandas_with_nan( - self, mock_build_filter: Mock, mock_get_conn: Mock - ) -> None: - """Test search engine fallback to to_pandas() and NaN normalization.""" - mock_conn = Mock() - mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - import numpy as np + def test_search_engine_with_missing_optional_fields(self) -> None: + """Test search engine handles results with missing/None optional fields (replaces arrow_fallback_to_pandas_with_nan test).""" import pandas as pd - # Create DataFrame with NaN values - mock_results_df = pd.DataFrame( - [ - { - "doc_id": "doc1", - "chunk_id": "chunk1", - "text": "test content", - "_distance": 0.5, - "parse_hash": "hash1", - "created_at": pd.Timestamp.now(), - "metadata": '{"key": "value"}', - "optional_field": np.nan, # NaN value - } - ] + # Mock vector store + mock_vector_store = Mock() + mock_vector_store.build_filter_expression.return_value = None + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=False, # Dense search doesn't use FTS ) - # Create mock search chain - use chainable mocks - mock_search = Mock() - mock_limit = Mock() - - mock_table.search.return_value = mock_search - # Chain: search().where().limit() - each returns the next in chain - mock_search.where.return_value = mock_search - mock_search.limit.return_value = mock_limit - - # Simulate both to_arrow() and to_list() failing, fallback to to_pandas() - mock_limit.to_arrow.side_effect = AttributeError("to_arrow not available") - mock_limit.to_list.side_effect = AttributeError("to_list not available") - mock_limit.to_pandas.return_value = mock_results_df - - mock_build_filter.return_value = None + # Mock search by model method - returns results with missing optional fields + mock_vector_store.search_vectors_by_model.return_value = [ + { + "doc_id": "doc1", + "chunk_id": "chunk1", + "text": "test content", + "_distance": 0.5, + "parse_hash": "hash1", + "created_at": pd.Timestamp.now(), + "metadata": '{"key": "value"}', + # Missing optional_field + } + ] with patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_index_manager" - ) as mock_get_index_manager: - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", - ) - mock_get_index_manager.return_value = mock_index_manager + "xagent.core.tools.core.RAG_tools.retrieval.search_engine.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store results, _, _ = search_dense_engine( collection="test_collection", @@ -873,10 +759,6 @@ def test_search_engine_arrow_fallback_to_pandas_with_nan( is_admin=True, ) - # Verify results + # Verify results are handled correctly assert len(results) == 1 assert results[0].doc_id == "doc1" - # Verify all fallbacks were attempted - mock_limit.to_arrow.assert_called_once() - mock_limit.to_list.assert_called_once() - mock_limit.to_pandas.assert_called_once() diff --git a/tests/core/tools/core/RAG_tools/retrieval/test_search_sparse.py b/tests/core/tools/core/RAG_tools/retrieval/test_search_sparse.py index 7e83bc7dd..fec4a8129 100644 --- a/tests/core/tools/core/RAG_tools/retrieval/test_search_sparse.py +++ b/tests/core/tools/core/RAG_tools/retrieval/test_search_sparse.py @@ -26,351 +26,381 @@ class TestSearchSparse: """Test search_sparse main function.""" - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" - ) - @patch("xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_index_manager") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.build_lancedb_filter_expression" - ) - def test_search_sparse_success_no_filters( - self, - mock_build_filter: Mock, - mock_get_index_manager: Mock, - mock_get_conn: Mock, - ) -> None: + def test_search_sparse_success_no_filters(self) -> None: """Test successful sparse search with collection filter only (KB isolation).""" - # Mock connection and table - mock_conn = Mock() + # Mock table mock_table = Mock() mock_table.name = "embeddings_test_model" # Set the table name - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table # Ensure open_table succeeds - - # Mock index manager - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", - ) - mock_index_manager.get_fts_index_status.return_value = True - mock_get_index_manager.return_value = mock_index_manager - - # Collection filter is always applied for KB isolation (Issue #72) - mock_build_filter.return_value = "collection == 'test_col'" - - # Mock search results; chain: search() -> limit() -> where() -> to_pandas() - mock_results_df = pd.DataFrame( - [ - { - "doc_id": "doc1", - "chunk_id": "chunk1", - "text": "test content one", - "_score": 0.9, - "parse_hash": "hash1", - "created_at": pd.Timestamp.now(), - } - ] - ) - mock_search = Mock() - mock_limit = Mock() - mock_where = Mock() - mock_table.search.return_value = mock_search - mock_search.limit.return_value = mock_limit - mock_limit.where.return_value = mock_where - mock_where.to_pandas.return_value = mock_results_df - response = search_sparse_module.search_sparse( - collection="test_col", - model_tag="test_model", - query_text="content", - top_k=1, - user_id=None, - is_admin=True, - ) + # Mock FTS index exists + mock_table.list_indices.return_value = [ + Mock(index_type="FTS", columns=["text"]) + ] - assert isinstance(response, SparseSearchResponse) - assert response.status == "success" - assert response.total_count == 1 - assert response.fts_enabled is True - assert len(response.results) == 1 - assert response.results[0].doc_id == "doc1" - assert response.results[0].text == "test content one" - # Score is normalized from TF-IDF to similarity score (0-1 range) - assert abs(response.results[0].score - 0.4736842105263158) < 1e-10 - assert not response.warnings - - # Verify calls: collection filter must be applied for KB isolation - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_called_once() - mock_build_filter.assert_called_once_with({"collection": "test_col"}) - mock_table.search.assert_called_once_with("content", query_type="fts") - mock_search.limit.assert_called_once_with(1) - mock_limit.where.assert_called_once() - where_arg = mock_limit.where.call_args[0][0] - assert "collection" in where_arg.lower() or "test_col" in where_arg + # Mock vector store + mock_vector_store = Mock() + # Return IndexResult object instead of string + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" - ) - @patch("xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_index_manager") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.build_lancedb_filter_expression" - ) - def test_search_sparse_with_filters( - self, mock_build_filter: Mock, mock_get_index_manager: Mock, mock_get_conn: Mock - ) -> None: - """Test sparse search with filters.""" - with patch.object( - search_sparse_module, "_substring_fallback", return_value=[] - ) as mock_fallback: - # Mock connection and table - mock_conn = Mock() - mock_table = Mock() - mock_table.name = "embeddings_test_model" # Set the table name - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=True, + ) + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'test_col'" + ) + # open_embeddings_table now returns tuple (table, table_name) + mock_vector_store.open_embeddings_table.return_value = ( + mock_table, + "embeddings_test_model", + ) - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + # Mock search results; chain: search() -> limit() -> where() -> to_pandas() + mock_results_df = pd.DataFrame( + [ + { + "doc_id": "doc1", + "chunk_id": "chunk1", + "text": "test content one", + "_score": 0.9, + "parse_hash": "hash1", + "created_at": pd.Timestamp.now(), + } + ] ) - mock_index_manager.get_fts_index_status.return_value = True - mock_get_index_manager.return_value = mock_index_manager - - mock_results_df = pd.DataFrame([]) mock_search = Mock() mock_limit = Mock() mock_where = Mock() - mock_table.search.return_value = mock_search mock_search.limit.return_value = mock_limit mock_limit.where.return_value = mock_where mock_where.to_pandas.return_value = mock_results_df - filters = {"doc_id": "filtered_doc", "collection": "test_col"} - expected_filter_clause = ( - "doc_id = 'filtered_doc' AND collection = 'test_col'" - ) - # Collection filter first, then custom filters (Issue #72) - mock_build_filter.side_effect = [ - "collection == 'test_col'", - expected_filter_clause, - ] - response = search_sparse_module.search_sparse( collection="test_col", model_tag="test_model", - query_text="filtered content", - top_k=5, - filters=filters, + query_text="content", + top_k=1, user_id=None, is_admin=True, ) + assert isinstance(response, SparseSearchResponse) + assert response.status == "success" + assert response.total_count == 1 + assert response.fts_enabled is True + assert len(response.results) == 1 + assert response.results[0].doc_id == "doc1" + assert response.results[0].text == "test content one" + # Score is normalized from TF-IDF to similarity score (0-1 range) + assert abs(response.results[0].score - 0.4736842105263158) < 1e-10 + assert not response.warnings + + # Verify calls: collection filter must be applied for KB isolation + # Note: model_tag is passed as-is to open_embeddings_table (no pre-transformation) + mock_vector_store.open_embeddings_table.assert_called_once_with( + "test_model" + ) + mock_vector_store.build_filter_expression.assert_called_once() + mock_table.search.assert_called_once_with("content", query_type="fts") + mock_search.limit.assert_called_once_with(1) + mock_limit.where.assert_called_once() + where_arg = mock_limit.where.call_args[0][0] + assert "collection" in where_arg.lower() or "test_col" in where_arg + + def test_search_sparse_with_filters(self) -> None: + """Test sparse search with filters.""" + with patch.object( + search_sparse_module, "_substring_fallback", return_value=[] + ) as mock_fallback: + # Mock table + mock_table = Mock() + mock_table.name = "embeddings_test_model" # Set the table name + + # Mock FTS index exists + mock_table.list_indices.return_value = [ + Mock(index_type="FTS", columns=["text"]) + ] + + # Mock vector store + mock_vector_store = Mock() + # Return IndexResult object instead of string + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=True, + ) + mock_vector_store.build_filter_expression.return_value = ( + "doc_id = 'filtered_doc' AND collection = 'test_col'" + ) + # open_embeddings_table now returns tuple (table, table_name) + mock_vector_store.open_embeddings_table.return_value = ( + mock_table, + "embeddings_test_model", + ) + + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + mock_results_df = pd.DataFrame([]) + mock_search = Mock() + mock_limit = Mock() + mock_where = Mock() + + mock_table.search.return_value = mock_search + mock_search.limit.return_value = mock_limit + mock_limit.where.return_value = mock_where + mock_where.to_pandas.return_value = mock_results_df + + filters = {"doc_id": "filtered_doc", "collection": "test_col"} + + response = search_sparse_module.search_sparse( + collection="test_col", + model_tag="test_model", + query_text="filtered content", + top_k=5, + filters=filters, + user_id=None, + is_admin=True, + ) + assert response.status == "success" assert response.total_count == 0 assert len(response.results) == 0 assert response.warnings == [] mock_fallback.assert_called_once() - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_called_once() - mock_index_manager.get_fts_index_status.assert_called_once_with(mock_table) - mock_build_filter.assert_any_call({"collection": "test_col"}) - mock_build_filter.assert_any_call(filters) + # Note: model_tag is passed as-is to open_embeddings_table (no pre-transformation) + mock_vector_store.open_embeddings_table.assert_called_once_with( + "test_model" + ) + mock_vector_store.build_filter_expression.assert_called() mock_table.search.assert_called_once_with( "filtered content", query_type="fts" ) mock_search.limit.assert_called_once_with(5) mock_limit.where.assert_called_once() - where_arg = mock_limit.where.call_args[0][0] - assert expected_filter_clause in where_arg mock_where.to_pandas.assert_called_once() - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" - ) - @patch("xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_index_manager") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.build_lancedb_filter_expression" - ) - def test_search_sparse_applies_collection_filter( - self, - mock_build_filter: Mock, - mock_get_index_manager: Mock, - mock_get_conn: Mock, - ) -> None: + def test_search_sparse_applies_collection_filter(self) -> None: """Test that search_sparse always applies collection filter for KB isolation (Issue #72).""" with patch.object(search_sparse_module, "_substring_fallback", return_value=[]): - mock_conn = Mock() + # Mock table mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", - ) - mock_index_manager.get_fts_index_status.return_value = True - mock_get_index_manager.return_value = mock_index_manager - mock_build_filter.return_value = "collection == 'my_kb'" - mock_search = Mock() - mock_limit = Mock() - mock_where = Mock() - mock_table.search.return_value = mock_search - mock_search.limit.return_value = mock_limit - mock_limit.where.return_value = mock_where - mock_where.to_pandas.return_value = pd.DataFrame() - search_sparse_module.search_sparse( - collection="my_kb", - model_tag="test_model", - query_text="query", - top_k=5, - user_id=None, - is_admin=True, + # Mock FTS index exists + mock_table.list_indices.return_value = [ + Mock(index_type="FTS", columns=["text"]) + ] + + # Mock vector store + mock_vector_store = Mock() + # Return IndexResult object instead of string + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=True, + ) + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'my_kb'" ) + # open_embeddings_table now returns tuple (table, table_name) + mock_vector_store.open_embeddings_table.return_value = ( + mock_table, + "embeddings_test_model", + ) + + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + mock_search = Mock() + mock_limit = Mock() + mock_where = Mock() + mock_table.search.return_value = mock_search + mock_search.limit.return_value = mock_limit + mock_limit.where.return_value = mock_where + mock_where.to_pandas.return_value = pd.DataFrame() + + search_sparse_module.search_sparse( + collection="my_kb", + model_tag="test_model", + query_text="query", + top_k=5, + user_id=None, + is_admin=True, + ) - mock_build_filter.assert_called_once_with({"collection": "my_kb"}) + mock_vector_store.build_filter_expression.assert_called_once() mock_limit.where.assert_called_once() - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" - ) - @patch("xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_index_manager") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.build_lancedb_filter_expression" - ) - def test_search_sparse_fts_index_missing( - self, - mock_build_filter: Mock, - mock_get_index_manager: Mock, - mock_get_conn: Mock, - ) -> None: + def test_search_sparse_fts_index_missing(self) -> None: """Test sparse search when FTS index is missing.""" with patch.object(search_sparse_module, "_substring_fallback", return_value=[]): - mock_conn = Mock() + # Mock table mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", + # Mock vector store - index status returned but FTS not enabled on table + mock_vector_store = Mock() + # Return IndexResult object instead of string + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=False, # FTS not enabled + ) + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'test_col'" + ) + # open_embeddings_table now returns tuple (table, table_name) + mock_vector_store.open_embeddings_table.return_value = ( + mock_table, + "embeddings_test_model", ) - mock_index_manager.get_fts_index_status.return_value = False - mock_get_index_manager.return_value = mock_index_manager - mock_build_filter.return_value = "collection == 'test_col'" - mock_search = Mock() - mock_limit = Mock() - mock_where = Mock() - mock_table.search.return_value = mock_search - mock_search.limit.return_value = mock_limit - mock_limit.where.return_value = mock_where - mock_where.to_pandas.return_value = pd.DataFrame() + # Make list_indices return no FTS index + mock_table.list_indices.return_value = [] - response = search_sparse_module.search_sparse( - collection="test_col", - model_tag="test_model", - query_text="query", - top_k=1, - user_id=None, - is_admin=True, - ) + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + mock_search = Mock() + mock_limit = Mock() + mock_where = Mock() + mock_table.search.return_value = mock_search + mock_search.limit.return_value = mock_limit + mock_limit.where.return_value = mock_where + mock_where.to_pandas.return_value = pd.DataFrame() + + response = search_sparse_module.search_sparse( + collection="test_col", + model_tag="test_model", + query_text="query", + top_k=1, + user_id=None, + is_admin=True, + ) assert response.status == "success" assert response.fts_enabled is False assert any(w.code == "FTS_INDEX_MISSING" for w in response.warnings) - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_called_once() - mock_index_manager.get_fts_index_status.assert_called_once_with(mock_table) + # Note: model_tag is passed as-is to open_embeddings_table (no pre-transformation) + mock_vector_store.open_embeddings_table.assert_called_once_with( + "test_model" + ) mock_table.search.assert_called_once_with("query", query_type="fts") mock_search.limit.assert_called_once_with(1) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" - ) - @patch("xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_index_manager") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.build_lancedb_filter_expression" - ) - def test_search_sparse_readonly_mode( - self, - mock_build_filter: Mock, - mock_get_index_manager: Mock, - mock_get_conn: Mock, - ) -> None: + def test_search_sparse_readonly_mode(self) -> None: """Test sparse search in readonly mode.""" with patch.object(search_sparse_module, "_substring_fallback", return_value=[]): - mock_conn = Mock() + # Mock table mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "readonly", - "Readonly mode", + # Mock vector store + mock_vector_store = Mock() + # Return IndexResult object instead of string + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=True, + ) + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'test_col'" + ) + # open_embeddings_table now returns tuple (table, table_name) + mock_vector_store.open_embeddings_table.return_value = ( + mock_table, + "embeddings_test_model", ) - mock_index_manager.get_fts_index_status.return_value = False - mock_get_index_manager.return_value = mock_index_manager - mock_build_filter.return_value = "collection == 'test_col'" - mock_search = Mock() - mock_limit = Mock() - mock_where = Mock() - mock_table.search.return_value = mock_search - mock_search.limit.return_value = mock_limit - mock_limit.where.return_value = mock_where - mock_where.to_pandas.return_value = pd.DataFrame() + # FTS index exists + mock_table.list_indices.return_value = [ + Mock(index_type="FTS", columns=["text"]) + ] - response = search_sparse_module.search_sparse( - collection="test_col", - model_tag="test_model", - query_text="query", - top_k=1, - readonly=True, - user_id=None, - is_admin=True, - ) + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + mock_search = Mock() + mock_limit = Mock() + mock_where = Mock() + mock_table.search.return_value = mock_search + mock_search.limit.return_value = mock_limit + mock_limit.where.return_value = mock_where + mock_where.to_pandas.return_value = pd.DataFrame() + + response = search_sparse_module.search_sparse( + collection="test_col", + model_tag="test_model", + query_text="query", + top_k=1, + readonly=True, + user_id=None, + is_admin=True, + ) assert response.status == "success" - assert response.fts_enabled is False + # FTS should be enabled since the table has the index + assert response.fts_enabled is True assert any(w.code == "READONLY_MODE" for w in response.warnings) - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_called_once() - mock_index_manager.get_fts_index_status.assert_called_once_with(mock_table) + # Note: model_tag is passed as-is to open_embeddings_table (no pre-transformation) + mock_vector_store.open_embeddings_table.assert_called_once_with( + "test_model" + ) mock_table.search.assert_called_once_with("query", query_type="fts") mock_search.limit.assert_called_once_with(1) @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.utils.model_resolver.resolve_embedding_adapter" ) - def test_search_sparse_database_error(self, mock_get_conn: Mock) -> None: + def test_search_sparse_database_error(self, mock_resolve: Mock) -> None: """Test error handling during database operation.""" - mock_conn = Mock() - mock_get_conn.return_value = mock_conn - # Simulate open_table failure + # Mock vector store that raises exception when opening table + mock_vector_store = Mock() db_exception_message = "DB connection failed" - mock_conn.open_table.side_effect = Exception(db_exception_message) - - response = search_sparse_module.search_sparse( - collection="test_col", - model_tag="test_model", - query_text="query", - top_k=1, + mock_vector_store.open_embeddings_table.side_effect = Exception( + db_exception_message ) + mock_cfg = Mock() + mock_cfg.model_name = "legacy_model" + mock_resolve.return_value = (mock_cfg, object()) + + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + response = search_sparse_module.search_sparse( + collection="test_col", + model_tag="test_model", + query_text="query", + top_k=1, + ) + assert response.status == "failed" assert response.total_count == 0 assert len(response.results) == 0 @@ -382,79 +412,76 @@ def test_search_sparse_database_error(self, mock_get_conn: Mock) -> None: in response.warnings[0].message ) - # Verify calls - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") + # Verify calls - open_embeddings_table is called once (handles fallback internally) + assert mock_vector_store.open_embeddings_table.call_count == 1 + # Note: model_tag is passed as-is to open_embeddings_table (no pre-transformation) + mock_vector_store.open_embeddings_table.assert_called_once_with("test_model") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" - ) - @patch("xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_index_manager") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.build_lancedb_filter_expression" - ) - def test_search_sparse_empty_results( - self, - mock_build_filter: Mock, - mock_get_index_manager: Mock, - mock_get_conn: Mock, - ) -> None: + def test_search_sparse_empty_results(self) -> None: """Test sparse search returning no results.""" with patch.object(search_sparse_module, "_substring_fallback", return_value=[]): - mock_conn = Mock() + # Mock table mock_table = Mock() - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", - ) - mock_index_manager.get_fts_index_status.return_value = True - mock_get_index_manager.return_value = mock_index_manager - mock_build_filter.return_value = "collection == 'test_col'" - mock_search = Mock() - mock_limit = Mock() - mock_where = Mock() - mock_table.search.return_value = mock_search - mock_search.limit.return_value = mock_limit - mock_limit.where.return_value = mock_where - mock_where.to_pandas.return_value = pd.DataFrame() + # Mock vector store + mock_vector_store = Mock() + # Return IndexResult object instead of string + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - response = search_sparse_module.search_sparse( - collection="test_col", - model_tag="test_model", - query_text="no matches", - top_k=5, - user_id=None, - is_admin=True, + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=True, ) + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'test_col'" + ) + # open_embeddings_table now returns tuple (table, table_name) + mock_vector_store.open_embeddings_table.return_value = ( + mock_table, + "embeddings_test_model", + ) + + # FTS index exists + mock_table.list_indices.return_value = [ + Mock(index_type="FTS", columns=["text"]) + ] + + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + mock_search = Mock() + mock_limit = Mock() + mock_where = Mock() + mock_table.search.return_value = mock_search + mock_search.limit.return_value = mock_limit + mock_limit.where.return_value = mock_where + mock_where.to_pandas.return_value = pd.DataFrame() + + response = search_sparse_module.search_sparse( + collection="test_col", + model_tag="test_model", + query_text="no matches", + top_k=5, + user_id=None, + is_admin=True, + ) assert response.status == "success" assert response.total_count == 0 assert len(response.results) == 0 assert response.warnings == [] - mock_get_conn.assert_called_once() - mock_conn.open_table.assert_called_once_with("embeddings_test_model") - mock_get_index_manager.assert_called_once() + # Note: model_tag is passed as-is to open_embeddings_table (no pre-transformation) + mock_vector_store.open_embeddings_table.assert_called_once_with( + "test_model" + ) mock_table.search.assert_called_once_with("no matches", query_type="fts") mock_search.limit.assert_called_once_with(5) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" - ) - @patch("xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_index_manager") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.build_lancedb_filter_expression" - ) - def test_search_sparse_triggers_fallback_with_results( - self, - mock_build_filter: Mock, - mock_get_index_manager: Mock, - mock_get_conn: Mock, - ) -> None: + def test_search_sparse_triggers_fallback_with_results(self) -> None: """Ensure fallback populates results and emits an FTS warning.""" def _fake_fallback(**kwargs: object) -> List[SearchResult]: @@ -479,75 +506,92 @@ def _fake_fallback(**kwargs: object) -> List[SearchResult]: ) ] - mock_conn = Mock() + # Mock table mock_table = Mock() - mock_table.name = "embeddings_test_model" # Set the table name - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", + # Mock vector store + mock_vector_store = Mock() + # Return IndexResult object instead of string + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=True, ) - mock_index_manager.get_fts_index_status.return_value = True - mock_get_index_manager.return_value = mock_index_manager - mock_build_filter.return_value = "collection == 'test_col'" - mock_search = Mock() - mock_limit = Mock() - mock_where = Mock() - mock_table.search.return_value = mock_search - mock_search.limit.return_value = mock_limit - mock_limit.where.return_value = mock_where - mock_where.to_pandas.return_value = pd.DataFrame() + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'test_col'" + ) + # open_embeddings_table now returns tuple (table, table_name) + mock_vector_store.open_embeddings_table.return_value = ( + mock_table, + "embeddings_test_model", + ) + + # FTS index exists + mock_table.list_indices.return_value = [ + Mock(index_type="FTS", columns=["text"]) + ] with patch.object( search_sparse_module, "_substring_fallback", side_effect=_fake_fallback ): - response = search_sparse_module.search_sparse( - collection="test_col", - model_tag="test_model", - query_text="fallback", - top_k=3, - user_id=None, - is_admin=True, - ) + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + mock_search = Mock() + mock_limit = Mock() + mock_where = Mock() + mock_table.search.return_value = mock_search + mock_search.limit.return_value = mock_limit + mock_limit.where.return_value = mock_where + mock_where.to_pandas.return_value = pd.DataFrame() + + response = search_sparse_module.search_sparse( + collection="test_col", + model_tag="test_model", + query_text="fallback", + top_k=3, + user_id=None, + is_admin=True, + ) assert response.status == "success" assert response.total_count == 1 assert response.results[0].doc_id == "doc-fallback" assert any(w.code == "FTS_FALLBACK" for w in response.warnings) - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_connection_from_env" - ) - @patch("xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_index_manager") - @patch( - "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.build_lancedb_filter_expression" - ) - def test_search_sparse_score_clamping( - self, - mock_build_filter: Mock, - mock_get_index_manager: Mock, - mock_get_conn: Mock, - ) -> None: + def test_search_sparse_score_clamping(self) -> None: """Test that sparse search scores are properly clamped to [0, 1] range.""" - # Mock connection and table - mock_conn = Mock() + # Mock table mock_table = Mock() - mock_table.name = "embeddings_test_model" - mock_get_conn.return_value = mock_conn - mock_conn.open_table.return_value = mock_table - - # Mock index manager - mock_index_manager = Mock() - mock_index_manager.check_and_create_index.return_value = ( - "index_ready", - "Index ready", + + # Mock vector store + mock_vector_store = Mock() + # Return IndexResult object instead of string + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.return_value = IndexResult( + status="index_ready", + advice=None, + fts_enabled=True, + ) + mock_vector_store.build_filter_expression.return_value = ( + "collection == 'test_col'" ) - mock_index_manager.get_fts_index_status.return_value = True - mock_get_index_manager.return_value = mock_index_manager - mock_build_filter.return_value = "collection == 'test_col'" + # open_embeddings_table now returns tuple (table, table_name) + mock_vector_store.open_embeddings_table.return_value = ( + mock_table, + "embeddings_test_model", + ) + + # FTS index exists + mock_table.list_indices.return_value = [ + Mock(index_type="FTS", columns=["text"]) + ] + mock_search = Mock() mock_limit = Mock() mock_where = Mock() @@ -569,14 +613,19 @@ def test_search_sparse_score_clamping( ) mock_where.to_pandas.return_value = test_data - response = search_sparse_module.search_sparse( - collection="test_col", - model_tag="test_model", - query_text="test", - top_k=10, - user_id=None, - is_admin=True, - ) + with patch( + "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store" + ) as mock_get_vector_store: + mock_get_vector_store.return_value = mock_vector_store + + response = search_sparse_module.search_sparse( + collection="test_col", + model_tag="test_model", + query_text="test", + top_k=10, + user_id=None, + is_admin=True, + ) assert response.status == "success" assert len(response.results) == 1 diff --git a/tests/core/tools/core/RAG_tools/storage/test_factory.py b/tests/core/tools/core/RAG_tools/storage/test_factory.py new file mode 100644 index 000000000..e99677699 --- /dev/null +++ b/tests/core/tools/core/RAG_tools/storage/test_factory.py @@ -0,0 +1,82 @@ +"""Tests for storage factory and coordinator wiring.""" + +from xagent.core.tools.core.RAG_tools.storage import factory + + +def test_factory_is_singleton(monkeypatch) -> None: + """Factory should return the same instance per process.""" + # Get existing factory and reset for test isolation + try: + f = factory.StorageFactory.get_factory() + f.reset_all() + except Exception: + factory.StorageFactory._instance = None + f = factory.StorageFactory.get_factory() + + first = factory.StorageFactory.get_factory() + second = factory.StorageFactory.get_factory() + + assert first is second + + +def test_factory_reset_all(monkeypatch) -> None: + """Factory reset_all should clear all store instances.""" + # Get existing factory and reset for test isolation + try: + f = factory.StorageFactory.get_factory() + f.reset_all() + except Exception: + factory.StorageFactory._instance = None + f = factory.StorageFactory.get_factory() + + # Create some stores + f.get_vector_index_store() + f.get_metadata_store() + f.get_ingestion_status_store() + + # Reset + f.reset_all() + + # Verify all stores are reset + assert f._vector_index_store is None + assert f._metadata_store is None + assert f._ingestion_status_store is None + + +def test_convenience_functions_use_factory(monkeypatch) -> None: + """Convenience functions should delegate to the singleton factory.""" + # Get existing factory and reset for test isolation + try: + f = factory.StorageFactory.get_factory() + f.reset_all() + except Exception: + factory.StorageFactory._instance = None + f = factory.StorageFactory.get_factory() + + first_vector = factory.get_vector_index_store() + first_metadata = factory.get_metadata_store() + + # Get via factory directly + second_vector = f.get_vector_index_store() + second_metadata = f.get_metadata_store() + + assert first_vector is second_vector + assert first_metadata is second_metadata + + +def test_coordinator_uses_factory_stores(monkeypatch) -> None: + """Coordinator should use stores from the factory.""" + # Get existing factory or create new one + try: + f = factory.StorageFactory.get_factory() + # Reset for test isolation + f.reset_all() + except Exception: + # If factory is in bad state, reset singleton + factory.StorageFactory._instance = None + f = factory.StorageFactory.get_factory() + + coordinator = factory.get_kb_write_coordinator() + + assert coordinator.metadata_store() is f.get_metadata_store() + assert coordinator.vector_index_store() is f.get_vector_index_store() diff --git a/tests/core/tools/core/RAG_tools/storage/test_lancedb_isolation.py b/tests/core/tools/core/RAG_tools/storage/test_lancedb_isolation.py new file mode 100644 index 000000000..e15c2818d --- /dev/null +++ b/tests/core/tools/core/RAG_tools/storage/test_lancedb_isolation.py @@ -0,0 +1,52 @@ +"""Tests to ensure pytest does not pollute the default LanceDB directory.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( + ensure_documents_table, +) +from xagent.core.tools.core.RAG_tools.storage import ( + get_vector_index_store, + reset_kb_write_coordinator, +) +from xagent.providers.vector_store.lancedb import ( + LanceDBConnectionManager, + clear_connection_cache, +) + + +def test_tests_do_not_pollute_default_lancedb_dir( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Creating tables in tests should not touch the default LanceDB directory. + + This test explicitly forces `LANCEDB_DIR` to a temporary directory to + avoid relying on any developer machine environment settings. + """ + expected_dir = tmp_path / "lancedb" + expected_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("LANCEDB_DIR", str(expected_dir)) + clear_connection_cache() + reset_kb_write_coordinator() + + default_dir = Path(LanceDBConnectionManager.get_default_lancedb_dir()) + default_exists_before = default_dir.exists() + default_listing_before = ( + {p.name for p in default_dir.iterdir()} if default_exists_before else set() + ) + + # Trigger a write path that creates tables in the isolated test directory. + conn = get_vector_index_store().get_raw_connection() + ensure_documents_table(conn) + + default_exists_after = default_dir.exists() + default_listing_after = ( + {p.name for p in default_dir.iterdir()} if default_exists_after else set() + ) + + assert default_exists_after == default_exists_before + assert default_listing_after == default_listing_before diff --git a/tests/core/tools/core/RAG_tools/storage/test_lancedb_stores.py b/tests/core/tools/core/RAG_tools/storage/test_lancedb_stores.py new file mode 100644 index 000000000..84f4d4dbb --- /dev/null +++ b/tests/core/tools/core/RAG_tools/storage/test_lancedb_stores.py @@ -0,0 +1,1828 @@ +"""Tests for LanceDB-backed storage implementations.""" + +import asyncio +from datetime import datetime, timezone +from typing import Any, Dict, List +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +from xagent.core.tools.core.RAG_tools.storage.lancedb_stores import ( + LanceDBMainPointerStore, + LanceDBMetadataStore, + LanceDBPromptTemplateStore, + LanceDBVectorIndexStore, +) + + +def create_mock_arrow_table(data_list: List[Dict[str, Any]]) -> Mock: + """Create a mock Arrow table that supports to_pylist() and len().""" + mock_table = Mock() + mock_table.to_pylist = Mock(return_value=data_list) + mock_table.__len__ = Mock(return_value=len(data_list)) + # Support iteration for 'for row in result' patterns + mock_table.__iter__ = Mock(return_value=iter(data_list)) + return mock_table + + +@pytest.fixture(autouse=True) +def mock_ensure_schema_fields() -> None: + """Mock _ensure_schema_fields to avoid schema iteration errors in tests.""" + with patch( + "xagent.core.tools.core.RAG_tools.LanceDB.schema_manager._ensure_schema_fields" + ): + yield + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_metadata_store_save_collection_config(mock_get_connection: Mock) -> None: + """Metadata store should save collection config correctly.""" + from types import SimpleNamespace + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + # Mock schema as iterable for _ensure_schema_fields + mock_table.schema = [SimpleNamespace(name="collection")] + mock_conn.open_table.return_value = mock_table + + store = LanceDBMetadataStore() + asyncio.run( + store.save_collection_config( + collection="test_collection", + config_json='{"parse_method": "default"}', + user_id=1, + ) + ) + + # Verify table.delete was called to remove existing config + mock_table.delete.assert_called_once() + + # Verify table.add was called with new config + mock_table.add.assert_called_once() + added_data = mock_table.add.call_args[0][0] + assert added_data[0]["collection"] == "test_collection" + assert added_data[0]["config_json"] == '{"parse_method": "default"}' + assert added_data[0]["user_id"] == 1 + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_metadata_store_get_collection_config_success( + mock_get_connection: Mock, +) -> None: + """Metadata store should retrieve collection config correctly.""" + from types import SimpleNamespace + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + # Mock schema as iterable for _ensure_schema_fields + mock_table.schema = [SimpleNamespace(name="collection")] + mock_conn.open_table.return_value = mock_table + + # Mock Arrow table with result[0]["config_json"].as_py() access pattern + mock_scalar = Mock() + mock_scalar.as_py = Mock(return_value='{"parse_method": "default"}') + + mock_config_col = Mock() + mock_config_col.__getitem__ = Mock(return_value=mock_scalar) + + mock_result = Mock() + mock_result.__len__ = Mock(return_value=1) + mock_result.__getitem__ = Mock( + side_effect=lambda key: mock_config_col if key == "config_json" else Mock() + ) + + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + mock_result + ) + + store = LanceDBMetadataStore() + config = asyncio.run( + store.get_collection_config(collection="test_collection", user_id=1) + ) + + assert config == '{"parse_method": "default"}' + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_metadata_store_get_collection_config_not_found( + mock_get_connection: Mock, +) -> None: + """Metadata store should return None when config not found.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + mock_result = Mock() + mock_result.__len__ = Mock(return_value=0) + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + mock_result + ) + + store = LanceDBMetadataStore() + config = asyncio.run( + store.get_collection_config(collection="test_collection", user_id=1) + ) + + assert config is None + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_metadata_store_get_collection_config_admin_picks_newest( + mock_get_connection: Mock, +) -> None: + """When is_admin, multiple tenant rows should resolve to latest updated_at.""" + import pyarrow as pa + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + older = datetime(2020, 1, 1) + newer = datetime(2021, 6, 1) + tbl = pa.table( + { + "collection": ["test_collection", "test_collection"], + "config_json": [ + '{"parse_method": "default"}', + '{"parse_method": "deepdoc"}', + ], + "updated_at": [older, newer], + "user_id": [1, 2], + } + ) + mock_table.search.return_value.where.return_value.to_arrow.return_value = tbl + + store = LanceDBMetadataStore() + config = asyncio.run( + store.get_collection_config( + collection="test_collection", user_id=0, is_admin=True + ) + ) + + assert config == '{"parse_method": "deepdoc"}' + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_metadata_store_get_collection_success(mock_get_connection: Mock) -> None: + """Metadata store should deserialize collection metadata correctly.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Use helper to create mock Arrow table + mock_data = { + "name": "test_collection", + "schema_version": "1.0.0", + "embedding_model_id": "text-embedding-v4", + "embedding_dimension": 1024, + "documents": 2, + "processed_documents": 2, + "parses": 2, + "chunks": 8, + "embeddings": 8, + "document_names": '["a.pdf","b.pdf"]', + "collection_locked": False, + "allow_mixed_parse_methods": False, + "skip_config_validation": False, + "created_at": datetime.now(timezone.utc).replace(tzinfo=None), + "updated_at": datetime.now(timezone.utc).replace(tzinfo=None), + "last_accessed_at": datetime.now(timezone.utc).replace(tzinfo=None), + "extra_metadata": "{}", + } + + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + create_mock_arrow_table([mock_data]) + ) + + store = LanceDBMetadataStore() + collection = asyncio.run(store.get_collection("test_collection")) + assert collection.name == "test_collection" + assert collection.documents == 2 + assert collection.document_names == ["a.pdf", "b.pdf"] + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.UserPermissions.get_user_filter" +) +@patch("xagent.core.tools.core.RAG_tools.storage.lancedb_stores.query_to_list") +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_vector_store_list_document_records_filters_and_maps( + mock_get_connection: Mock, + mock_query_to_list: Mock, + mock_user_filter: Mock, +) -> None: + """Vector store should apply combined filter and map to DocumentRecord.""" + from types import SimpleNamespace + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_user_filter.return_value = "user_id == 1" + mock_table = Mock() + # Mock schema as iterable for _ensure_schema_fields + mock_table.schema = [SimpleNamespace(name="doc_id")] + mock_conn.open_table.return_value = mock_table + mock_query_to_list.return_value = [ + {"doc_id": "doc-1", "source_path": "/tmp/a.pdf"}, + {"doc_id": "doc-2", "source_path": None}, + ] + + store = LanceDBVectorIndexStore() + records = store.list_document_records( + collection_name="kb1", + user_id=1, + is_admin=False, + max_results=50, + ) + + assert [r.doc_id for r in records] == ["doc-1", "doc-2"] + assert records[0].source_path == "/tmp/a.pdf" + mock_table.search.return_value.where.assert_called_once() + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_vector_store_rename_collection_data_updates_expected_tables( + mock_get_connection: Mock, +) -> None: + """Rename should update core and embeddings tables only.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_conn.table_names.return_value = [ + "documents", + "parses", + "chunks", + "embeddings_text_embedding_v4", + "collection_metadata", + ] + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + store = LanceDBVectorIndexStore() + warnings = store.rename_collection_data("old_name", "new_name") + + assert warnings == [] + # 4 target tables should be updated; control-plane table excluded. + assert mock_table.update.call_count == 4 + + +# --- Upsert Fallback Tests --- + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_embeddings_merge_insert_success(mock_get_connection: Mock) -> None: + """Test successful merge_insert upsert.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock merge_insert chain + mock_merge_insert = Mock() + mock_when_matched = Mock() + mock_when_not_matched = Mock() + mock_table.merge_insert.return_value = mock_merge_insert + mock_merge_insert.when_matched_update_all.return_value = mock_when_matched + mock_when_matched.when_not_matched_insert_all.return_value = mock_when_not_matched + mock_when_not_matched.execute.return_value = None + + store = LanceDBVectorIndexStore() + + records = [ + { + "collection": "test_col", + "doc_id": "doc1", + "chunk_id": "chunk1", + "vector": [0.1, 0.2], + "text": "test", + } + ] + + store.upsert_embeddings("text_embedding_v4", records) + + # Verify merge_insert was called + mock_table.merge_insert.assert_called_once_with( + ["collection", "doc_id", "chunk_id"] + ) + mock_merge_insert.when_matched_update_all.assert_called_once() + mock_when_matched.when_not_matched_insert_all.assert_called_once() + mock_when_not_matched.execute.assert_called_once() + + # Verify add was NOT called (no fallback needed) + mock_table.add.assert_not_called() + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_embeddings_merge_insert_fallback_to_add( + mock_get_connection: Mock, +) -> None: + """Test fallback to add() when merge_insert fails with recoverable error.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock merge_insert chain that fails + mock_merge_insert = Mock() + mock_when_matched = Mock() + mock_when_not_matched = Mock() + mock_table.merge_insert.return_value = mock_merge_insert + mock_merge_insert.when_matched_update_all.return_value = mock_when_matched + mock_when_matched.when_not_matched_insert_all.return_value = mock_when_not_matched + # merge_insert fails with recoverable error (e.g., network issue) + mock_when_not_matched.execute.side_effect = Exception("Temporary network error") + + # Mock add() to succeed + mock_table.add.return_value = None + + store = LanceDBVectorIndexStore() + + records = [ + { + "collection": "test_col", + "doc_id": "doc1", + "chunk_id": "chunk1", + "vector": [0.1, 0.2], + "text": "test", + } + ] + + store.upsert_embeddings("text_embedding_v4", records) + + # Verify merge_insert was attempted + mock_table.merge_insert.assert_called_once() + + # Verify fallback to add() was used + mock_table.add.assert_called_once() + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_embeddings_non_recoverable_error_no_fallback( + mock_get_connection: Mock, +) -> None: + """Test that non-recoverable errors (schema, type mismatch) do not fallback.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock merge_insert chain that fails with non-recoverable error + mock_merge_insert = Mock() + mock_when_matched = Mock() + mock_when_not_matched = Mock() + mock_table.merge_insert.return_value = mock_merge_insert + mock_merge_insert.when_matched_update_all.return_value = mock_when_matched + mock_when_matched.when_not_matched_insert_all.return_value = mock_when_not_matched + # Schema error - should NOT fallback + mock_when_not_matched.execute.side_effect = ValueError("Schema mismatch") + + store = LanceDBVectorIndexStore() + + records = [ + { + "collection": "test_col", + "doc_id": "doc1", + "chunk_id": "chunk1", + "vector": [0.1, 0.2], + "text": "test", + } + ] + + # Should raise ValueError without fallback + with pytest.raises(ValueError, match="Schema mismatch"): + store.upsert_embeddings("text_embedding_v4", records) + + # Verify merge_insert was attempted + mock_table.merge_insert.assert_called_once() + + # Verify add() was NOT called (no fallback for non-recoverable errors) + mock_table.add.assert_not_called() + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_embeddings_both_methods_fail(mock_get_connection: Mock) -> None: + """Test that error is raised when both merge_insert and add() fail.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock merge_insert chain that fails + mock_merge_insert = Mock() + mock_when_matched = Mock() + mock_when_not_matched = Mock() + mock_table.merge_insert.return_value = mock_merge_insert + mock_merge_insert.when_matched_update_all.return_value = mock_when_matched + mock_when_matched.when_not_matched_insert_all.return_value = mock_when_not_matched + mock_when_not_matched.execute.side_effect = Exception("merge_insert failed") + + # Mock add() to also fail + mock_table.add.side_effect = Exception("add() also failed") + + store = LanceDBVectorIndexStore() + + records = [ + { + "collection": "test_col", + "doc_id": "doc1", + "chunk_id": "chunk1", + "vector": [0.1, 0.2], + "text": "test", + } + ] + + # Should raise when both methods fail + with pytest.raises(Exception, match="add.*also failed"): + store.upsert_embeddings("text_embedding_v4", records) + + # Verify both methods were attempted + mock_table.merge_insert.assert_called_once() + mock_table.add.assert_called_once() + + +# ============================================================================ +# Index Management Tests (Phase 1A Part 2) +# ============================================================================ + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_should_reindex_immediate_reindex_enabled( + mock_get_connection: Mock, +) -> None: + """Test should_reindex returns True when immediate reindex is enabled.""" + from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock index stats + mock_stats = Mock() + mock_stats.num_indexed_rows = 1000 + mock_stats.num_unindexed_rows = 100 + mock_table.index_stats.return_value = mock_stats + + store = LanceDBVectorIndexStore() + + policy = IndexPolicy( + reindex_batch_size=1000, + enable_immediate_reindex=True, + enable_smart_reindex=False, + ) + + result = store.should_reindex("embeddings_test", total_upserted=10, policy=policy) + + assert result is True # immediate reindex enabled + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_should_reindex_batch_threshold( + mock_get_connection: Mock, +) -> None: + """Test should_reindex returns True when batch size threshold reached.""" + from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + store = LanceDBVectorIndexStore() + + policy = IndexPolicy( + reindex_batch_size=100, + enable_immediate_reindex=False, + enable_smart_reindex=False, + ) + + # Total upserted >= batch_size + result = store.should_reindex("embeddings_test", total_upserted=100, policy=policy) + assert result is True + + # Below threshold + result = store.should_reindex("embeddings_test", total_upserted=99, policy=policy) + assert result is False + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_should_reindex_smart_reindex( + mock_get_connection: Mock, +) -> None: + """Test should_reindex with smart reindex enabled.""" + from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock index stats with high unindexed ratio + mock_stats = Mock() + mock_stats.num_indexed_rows = 100 + mock_stats.num_unindexed_rows = 60 # 60% unindexed + mock_table.index_stats.return_value = mock_stats + + store = LanceDBVectorIndexStore() + + policy = IndexPolicy( + reindex_batch_size=10000, + enable_immediate_reindex=False, + enable_smart_reindex=True, + reindex_unindexed_ratio_threshold=0.5, # 50% threshold + ) + + # High unindexed ratio should trigger reindex + result = store.should_reindex("embeddings_test", total_upserted=10, policy=policy) + assert result is True + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_trigger_reindex_success(mock_get_connection: Mock) -> None: + """Test trigger_reindex calls table.optimize().""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + store = LanceDBVectorIndexStore() + + result = store.trigger_reindex("embeddings_test") + + assert result is True + mock_table.optimize.assert_called_once() + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_trigger_reindex_failure(mock_get_connection: Mock) -> None: + """Test trigger_reindex returns False on exception.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + mock_table.optimize.side_effect = Exception("Optimize failed") + + store = LanceDBVectorIndexStore() + + result = store.trigger_reindex("embeddings_test") + + assert result is False + + +@pytest.mark.asyncio +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_should_reindex_async_delegates_to_sync( + mock_get_connection: Mock, +) -> None: + """Test async version delegates to sync implementation.""" + from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock index stats with high unindexed ratio (60%) + mock_stats = Mock() + mock_stats.num_indexed_rows = 100 + mock_stats.num_unindexed_rows = 60 # 60% unindexed, exceeds 50% threshold + mock_table.index_stats.return_value = mock_stats + + store = LanceDBVectorIndexStore() + + policy = IndexPolicy( + reindex_batch_size=10000, + enable_immediate_reindex=False, + enable_smart_reindex=True, + reindex_unindexed_ratio_threshold=0.5, + ) + + # Async version should delegate to sync + result = await store.should_reindex_async( + "embeddings_test", total_upserted=10, policy=policy + ) + assert result is True # Smart reindex triggers due to high unindexed ratio + + +@pytest.mark.asyncio +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_trigger_reindex_async_delegates_to_sync( + mock_get_connection: Mock, +) -> None: + """Test async trigger_reindex delegates to sync implementation.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + store = LanceDBVectorIndexStore() + + # Async version should delegate to sync + result = await store.trigger_reindex_async("embeddings_test") + assert result is True + mock_table.optimize.assert_called_once() + + +# ============================================================================ +# PromptTemplateStore Tests (Phase 1A Part 3) +# ============================================================================ + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_prompt_template_store_save_and_get(mock_get_connection: Mock) -> None: + """Test saving and retrieving a prompt template.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock empty result for existing check + mock_result = Mock() + mock_result.__len__ = Mock(return_value=0) + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + mock_result + ) + + store = LanceDBPromptTemplateStore() + + # Save template + template_id = store.save_prompt_template( + name="test_template", + template="Test prompt content", + user_id=1, + ) + + assert template_id is not None + mock_table.add.assert_called_once() + + # Mock get result + row_data = { + "id": template_id, + "name": "test_template", + "template": "Test prompt content", + "version": 1, + "is_latest": True, + "metadata": "", + "user_id": 1, + "created_at": None, + "updated_at": None, + } + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + create_mock_arrow_table([row_data]) + ) + + # Get template + template = store.get_prompt_template(template_id, user_id=1) + assert template is not None + assert template["name"] == "test_template" + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_prompt_template_store_get_latest(mock_get_connection: Mock) -> None: + """Test getting the latest version of a template by name.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock result + row_data = { + "id": "test-id", + "name": "test_template", + "template": "Latest content", + "version": 2, + "is_latest": True, + "metadata": "", + "user_id": 1, + "created_at": None, + "updated_at": None, + } + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + create_mock_arrow_table([row_data]) + ) + + store = LanceDBPromptTemplateStore() + + template = store.get_latest_prompt_template("test_template", user_id=1) + assert template is not None + assert template["version"] == 2 + assert template["template"] == "Latest content" + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_prompt_template_store_delete(mock_get_connection: Mock) -> None: + """Test deleting a prompt template.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock existing template + mock_row = {"is_latest": True, "name": "test-template"} + mock_result = create_mock_arrow_table([mock_row]) + + # Mock remaining versions after delete (empty for this test) + mock_result_empty = create_mock_arrow_table([]) + + mock_table.search.return_value.where.return_value.to_arrow.side_effect = [ + mock_result, + mock_result_empty, + ] + + store = LanceDBPromptTemplateStore() + + result = store.delete_prompt_template("test-id", user_id=1) + assert result is True + mock_table.delete.assert_called_once() + + +# ============================================================================ +# MainPointerStore Tests (Phase 1A Part 3) +# ============================================================================ + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_main_pointer_store_set_and_get(mock_get_connection: Mock) -> None: + """Test setting and getting a main pointer.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock no existing pointer + mock_result = Mock() + mock_result.__len__ = Mock(return_value=0) + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + mock_result + ) + + store = LanceDBMainPointerStore() + + # Set pointer + store.set_main_pointer( + collection="test_collection", + doc_id="test_doc", + step_type="parse", + semantic_id="parse-123", + technical_id="hash-456", + ) + + # Verify merge_insert was called + mock_table.merge_insert.assert_called_once() + + # Mock get result + mock_row = { + "collection": "test_collection", + "doc_id": "test_doc", + "step_type": "parse", + "model_tag": "", + "semantic_id": "parse-123", + "technical_id": "hash-456", + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "operator": "unknown", + } + + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + create_mock_arrow_table([mock_row]) + ) + + # Get pointer + pointer = store.get_main_pointer("test_collection", "test_doc", "parse") + assert pointer is not None + assert pointer["semantic_id"] == "parse-123" + assert pointer["technical_id"] == "hash-456" + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_main_pointer_store_user_id_warning(mock_get_connection: Mock, caplog) -> None: + """Test that user_id parameter triggers a warning.""" + import logging + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock no existing pointer + mock_result = Mock() + mock_result.__len__ = Mock(return_value=0) + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + mock_result + ) + + store = LanceDBMainPointerStore() + + # Set pointer with user_id (should log warning) + with caplog.at_level(logging.WARNING): + store.set_main_pointer( + collection="test_collection", + doc_id="test_doc", + step_type="parse", + semantic_id="parse-123", + technical_id="hash-456", + user_id=1, + ) + + # Verify warning was logged + assert any( + "user_id parameter provided" in record.message + for record in caplog.records + if record.levelname == "WARNING" + ) + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_main_pointer_store_list(mock_get_connection: Mock) -> None: + """Test listing main pointers.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock count_rows > 0 + mock_table.search.return_value.where.return_value.count_rows.return_value = 1 + + # Mock result + mock_row_data = { + "collection": "test_collection", + "doc_id": "test_doc", + "step_type": "parse", + "model_tag": "", + "semantic_id": "parse-123", + "technical_id": "hash-456", + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "operator": "unknown", + } + + mock_table.search.return_value.where.return_value.limit.return_value.to_arrow.return_value = create_mock_arrow_table( + [mock_row_data] + ) + + store = LanceDBMainPointerStore() + + pointers = store.list_main_pointers("test_collection") + assert len(pointers) == 1 + assert pointers[0]["semantic_id"] == "parse-123" + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_main_pointer_store_delete(mock_get_connection: Mock) -> None: + """Test deleting a main pointer.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock existing pointer + mock_result = Mock() + mock_result.__len__ = Mock(return_value=1) + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + mock_result + ) + + store = LanceDBMainPointerStore() + + result = store.delete_main_pointer("test_collection", "test_doc", "parse") + assert result is True + mock_table.delete.assert_called_once() + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_main_pointer_store_delete_not_found(mock_get_connection: Mock) -> None: + """Test deleting a non-existent pointer returns False.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock no existing pointer + mock_result = Mock() + mock_result.__len__ = Mock(return_value=0) + mock_table.search.return_value.where.return_value.to_arrow.return_value = ( + mock_result + ) + + store = LanceDBMainPointerStore() + + result = store.delete_main_pointer("test_collection", "test_doc", "parse") + assert result is False + mock_table.delete.assert_not_called() + + +# ============================================================================= +# Async Method Tests (Phase 1A Coverage Improvement) +# ============================================================================= + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_search_vectors_async_basic( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test basic async vector search.""" + import pyarrow as pa + + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock Arrow table with results + data = { + "doc_id": ["doc1", "doc2"], + "score": [0.95, 0.87], + "vector": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], + } + arrow_table = pa.Table.from_pydict(data) + + # Mock table and vector search + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + + # Mock vector search - chain needs to return mock objects + mock_search = Mock() + mock_search.limit.return_value = mock_search + mock_search.where = Mock(return_value=mock_search) + + # to_arrow needs to be a coroutine that returns the arrow table + async def mock_to_arrow(): + return arrow_table + + mock_search.to_arrow = mock_to_arrow + + mock_table.search = Mock(return_value=mock_search) + + store = LanceDBVectorIndexStore() + + # Create a query vector + query_vector = [0.1, 0.2, 0.3] + + results = await store.search_vectors_async( + table_name="embeddings_test", + query_vector=query_vector, + top_k=5, + ) + + assert len(results) == 2 + assert results[0]["doc_id"] == "doc1" + assert results[0]["score"] == 0.95 + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_search_fts_async_basic( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test basic async FTS search.""" + import pyarrow as pa + + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock Arrow table with FTS results + data = { + "doc_id": ["doc1", "doc2"], + "text": ["hello world", "test content"], + "score": [0.9, 0.8], + } + arrow_table = pa.Table.from_pydict(data) + + # Mock table and FTS search + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + + # Mock search to return our table + mock_search = Mock() + mock_search.limit.return_value = mock_search + mock_search.where = Mock(return_value=mock_search) + + async def mock_to_arrow(): + return arrow_table + + mock_search.to_arrow = mock_to_arrow + + mock_table.search = Mock(return_value=mock_search) + + store = LanceDBVectorIndexStore() + + results = await store.search_fts_async( + table_name="chunks", + query_text="hello", + top_k=5, + ) + + assert len(results) == 2 + assert results[0]["doc_id"] == "doc1" + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_iter_batches_async_basic( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async batch iteration.""" + import pyarrow as pa + + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock table and to_batches + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + + # Create mock batches + batch1_schema = pa.schema([("doc_id", pa.string()), ("text", pa.string())]) + batch1_data = {"doc_id": ["doc1"], "text": ["text1"]} + batch1 = pa.RecordBatch.from_pydict(batch1_data, schema=batch1_schema) + + # Mock to_batches as async generator + async def mock_to_batches(**kwargs): + yield batch1 + + mock_table.to_batches = mock_to_batches + + store = LanceDBVectorIndexStore() + + batches = [] + async for batch in store.iter_batches_async( + table_name="chunks", + batch_size=100, + ): + batches.append(batch) + + assert len(batches) == 1 + assert batches[0].num_rows == 1 + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_count_rows_async_basic( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async row counting.""" + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock table and count_rows + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + mock_table.count_rows = AsyncMock(return_value=100) + + store = LanceDBVectorIndexStore() + + count = await store.count_rows_async(table_name="chunks") + + assert count == 100 + mock_table.count_rows.assert_awaited_once() + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_upsert_documents_async( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async document upsert.""" + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock sync connection for ensure_documents_table + mock_conn.open_table.return_value = Mock() + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock table and merge_insert + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + + # Mock merge_insert chain + mock_merge_builder = Mock() + mock_merge_builder.when_matched_update_all = Mock(return_value=mock_merge_builder) + mock_merge_builder.when_not_matched_insert_all = Mock( + return_value=mock_merge_builder + ) + + async def mock_execute(records): + return None + + mock_merge_builder.execute = mock_execute + + mock_table.merge_insert = Mock(return_value=mock_merge_builder) + + store = LanceDBVectorIndexStore() + + records = [ + {"doc_id": "doc1", "source_path": "/tmp/test.pdf"}, + {"doc_id": "doc2", "source_path": "/tmp/test2.pdf"}, + ] + + await store.upsert_documents_async(records) + + # Verify merge_insert was called + mock_table.merge_insert.assert_called_once() + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_upsert_chunks_async( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async chunk upsert.""" + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock sync connection for ensure_chunks_table + mock_conn.open_table.return_value = Mock() + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock table and merge_insert + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + + # Mock merge_insert chain + mock_merge_builder = Mock() + mock_merge_builder.when_matched_update_all = Mock(return_value=mock_merge_builder) + mock_merge_builder.when_not_matched_insert_all = Mock( + return_value=mock_merge_builder + ) + + async def mock_execute(records): + return None + + mock_merge_builder.execute = mock_execute + + mock_table.merge_insert = Mock(return_value=mock_merge_builder) + + store = LanceDBVectorIndexStore() + + records = [ + {"chunk_id": "chunk1", "text": "test content 1"}, + {"chunk_id": "chunk2", "text": "test content 2"}, + ] + + await store.upsert_chunks_async(records) + + # Verify merge_insert was called + mock_table.merge_insert.assert_called_once() + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_upsert_embeddings_async( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async embedding upsert.""" + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock sync connection for ensure_embeddings_table + mock_conn.open_table.return_value = Mock() + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock table and merge_insert + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + + # Mock merge_insert chain + mock_merge_builder = Mock() + mock_merge_builder.when_matched_update_all = Mock(return_value=mock_merge_builder) + mock_merge_builder.when_not_matched_insert_all = Mock( + return_value=mock_merge_builder + ) + + async def mock_execute(records): + return None + + mock_merge_builder.execute = mock_execute + + mock_table.merge_insert = Mock(return_value=mock_merge_builder) + + store = LanceDBVectorIndexStore() + + records = [ + {"chunk_id": "chunk1", "vector": [0.1, 0.2, 0.3]}, + {"chunk_id": "chunk2", "vector": [0.4, 0.5, 0.6]}, + ] + + await store.upsert_embeddings_async("bge_large", records) + + # Verify merge_insert was called + mock_table.merge_insert.assert_called_once() + + +# ============================================================================ +# Core Sync Upsert Method Tests +# ============================================================================ + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_documents_basic(mock_get_connection: Mock) -> None: + """Test basic document upsert.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table and merge_insert + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + mock_merge = Mock() + mock_merge.when_matched_update_all = Mock(return_value=mock_merge) + mock_merge.when_not_matched_insert_all = Mock(return_value=mock_merge) + mock_merge.execute = Mock(return_value=None) + mock_table.merge_insert = Mock(return_value=mock_merge) + + store = LanceDBVectorIndexStore() + + records = [ + {"doc_id": "doc1", "source_path": "/tmp/test.pdf"}, + {"doc_id": "doc2", "source_path": "/tmp/test2.pdf"}, + ] + + store.upsert_documents(records) + + # Verify merge_insert was called with correct keys + mock_table.merge_insert.assert_called_once_with(["collection", "doc_id"]) + mock_merge.execute.assert_called_once_with(records) + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_documents_empty(mock_get_connection: Mock) -> None: + """Test document upsert with empty records returns early.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + store = LanceDBVectorIndexStore() + + # Should return early without opening table + store.upsert_documents([]) + + # Verify table was never opened + mock_conn.open_table.assert_not_called() + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_parses_basic(mock_get_connection: Mock) -> None: + """Test basic parse upsert.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table and merge_insert + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + mock_merge = Mock() + mock_merge.when_matched_update_all = Mock(return_value=mock_merge) + mock_merge.when_not_matched_insert_all = Mock(return_value=mock_merge) + mock_merge.execute = Mock(return_value=None) + mock_table.merge_insert = Mock(return_value=mock_merge) + + store = LanceDBVectorIndexStore() + + records = [ + {"doc_id": "doc1", "parse_hash": "hash1", "parse_status": "success"}, + {"doc_id": "doc2", "parse_hash": "hash2", "parse_status": "success"}, + ] + + store.upsert_parses(records) + + # Verify merge_insert was called with correct keys + mock_table.merge_insert.assert_called_once_with( + ["collection", "doc_id", "parse_hash"] + ) + mock_merge.execute.assert_called_once_with(records) + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_chunks_basic(mock_get_connection: Mock) -> None: + """Test basic chunk upsert.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table and merge_insert + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + mock_merge = Mock() + mock_merge.when_matched_update_all = Mock(return_value=mock_merge) + mock_merge.when_not_matched_insert_all = Mock(return_value=mock_merge) + mock_merge.execute = Mock(return_value=None) + mock_table.merge_insert = Mock(return_value=mock_merge) + + store = LanceDBVectorIndexStore() + + records = [ + { + "chunk_id": "chunk1", + "doc_id": "doc1", + "parse_hash": "hash1", + "text": "test content 1", + }, + { + "chunk_id": "chunk2", + "doc_id": "doc1", + "parse_hash": "hash1", + "text": "test content 2", + }, + ] + + store.upsert_chunks(records) + + # Verify merge_insert was called with correct keys + mock_table.merge_insert.assert_called_once_with( + ["collection", "doc_id", "parse_hash", "chunk_id"] + ) + mock_merge.execute.assert_called_once_with(records) + + +# ============================================================================ +# Error Handling Tests +# ============================================================================ + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_search_vectors_async_table_not_found( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async vector search handles missing table gracefully.""" + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock open_table to raise exception + mock_async_conn.open_table = AsyncMock(side_effect=Exception("Table not found")) + + store = LanceDBVectorIndexStore() + + query_vector = [0.1, 0.2, 0.3] + results = await store.search_vectors_async( + table_name="nonexistent_table", + query_vector=query_vector, + top_k=5, + ) + + # Should return empty list on error + assert results == [] + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_search_vectors_async_search_failure( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async vector search handles search failure gracefully.""" + + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock table + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + + # Mock search that fails + mock_search = Mock() + mock_search.limit.return_value = mock_search + mock_search.where = Mock(return_value=mock_search) + + async def mock_to_arrow(): + raise Exception("Search failed") + + mock_search.to_arrow = mock_to_arrow + + mock_table.search = Mock(return_value=mock_search) + + store = LanceDBVectorIndexStore() + + query_vector = [0.1, 0.2, 0.3] + results = await store.search_vectors_async( + table_name="embeddings_test", + query_vector=query_vector, + top_k=5, + ) + + # Should return empty list on search error + assert results == [] + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_upsert_documents_with_invalid_data(mock_get_connection: Mock) -> None: + """Test document upsert handles invalid data gracefully.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table and merge_insert that raises exception + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + mock_merge = Mock() + mock_merge.when_matched_update_all = Mock(return_value=mock_merge) + mock_merge.when_not_matched_insert_all = Mock(return_value=mock_merge) + mock_merge.execute = Mock(side_effect=Exception("Invalid data")) + mock_table.merge_insert = Mock(return_value=mock_merge) + + store = LanceDBVectorIndexStore() + + records = [{"doc_id": "doc1", "invalid_field": "value"}] + + # Should raise exception on invalid data + with pytest.raises(Exception, match="Invalid data"): + store.upsert_documents(records) + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_iter_batches_async_invalid_columns( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async iter_batches handles invalid columns gracefully.""" + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock table + mock_table = Mock() + mock_async_conn.open_table = AsyncMock(return_value=mock_table) + + # Mock to_batches generator that raises exception + async def mock_to_batches(**kwargs): + raise Exception("Invalid columns") + + # Make to_batches return an async generator that raises + def make_to_batches(): + async def inner(**kwargs): + raise Exception("Invalid columns") + + return inner() + + mock_table.to_batches = make_to_batches() + + store = LanceDBVectorIndexStore() + + # Should handle exception gracefully and not yield any batches + batches = [] + async for batch in store.iter_batches_async( + table_name="chunks", + batch_size=100, + columns=["nonexistent_column"], + ): + batches.append(batch) + + # Should get no batches due to error + assert len(batches) == 0 + + +@pytest.mark.asyncio +@patch("lancedb.connect_async", new_callable=AsyncMock) +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_count_rows_async_table_not_found( + mock_get_connection: Mock, mock_connect_async: AsyncMock +) -> None: + """Test async count_rows handles missing table gracefully.""" + mock_conn = Mock() + mock_conn.uri = "test_uri" + mock_get_connection.return_value = mock_conn + + # Mock async connection + mock_async_conn = Mock() + mock_connect_async.return_value = mock_async_conn + + # Mock open_table to raise exception + mock_async_conn.open_table = AsyncMock(side_effect=Exception("Table not found")) + + store = LanceDBVectorIndexStore() + + count = await store.count_rows_async(table_name="nonexistent_table") + + # Should return 0 on error + assert count == 0 + + +# --- get_vector_dimension Tests (Issue #14) --- + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_get_vector_dimension_success(mock_get_connection: Mock) -> None: + """Test get_vector_dimension returns correct dimension from schema.""" + from types import SimpleNamespace + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table with fixed-size vector field + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock schema with vector field having list_size + mock_vector_type = SimpleNamespace(list_size=1536) + mock_vector_field = SimpleNamespace(type=mock_vector_type) + mock_schema = Mock() + mock_schema.field.return_value = mock_vector_field + mock_table.schema = mock_schema + + store = LanceDBVectorIndexStore() + dimension = store.get_vector_dimension("embeddings_test_model") + + assert dimension == 1536 + mock_schema.field.assert_called_once_with("vector") + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_get_vector_dimension_table_not_found(mock_get_connection: Mock) -> None: + """Test get_vector_dimension returns None when table not found.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock open_table to raise exception + mock_conn.open_table.side_effect = Exception("Table not found") + + store = LanceDBVectorIndexStore() + dimension = store.get_vector_dimension("nonexistent_table") + + assert dimension is None + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_get_vector_dimension_variable_length(mock_get_connection: Mock) -> None: + """Test get_vector_dimension returns None for variable-length vectors.""" + from types import SimpleNamespace + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table with variable-length vector field (no list_size) + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + # Mock schema with vector field lacking list_size attribute + mock_vector_type = SimpleNamespace() # No list_size + mock_vector_field = SimpleNamespace(type=mock_vector_type) + mock_schema = Mock() + mock_schema.field.return_value = mock_vector_field + mock_table.schema = mock_schema + + store = LanceDBVectorIndexStore() + dimension = store.get_vector_dimension("embeddings_variable") + + assert dimension is None + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_get_vector_dimension_no_vector_field(mock_get_connection: Mock) -> None: + """Test get_vector_dimension returns None when vector field missing.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table without vector field + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + mock_schema = Mock() + mock_schema.field.side_effect = Exception("Field 'vector' not found") + mock_table.schema = mock_schema + + store = LanceDBVectorIndexStore() + dimension = store.get_vector_dimension("embeddings_no_vector") + + assert dimension is None + + +# --- list_table_names Tests (Issue #14) --- + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_list_table_names_success(mock_get_connection: Mock) -> None: + """Test list_table_names returns correct table names.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table_names to return list of names + mock_conn.table_names.return_value = ["documents", "chunks", "embeddings_test"] + + store = LanceDBVectorIndexStore() + names = store.list_table_names() + + assert names == ["documents", "chunks", "embeddings_test"] + mock_conn.table_names.assert_called_once() + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_list_table_names_connection_error(mock_get_connection: Mock) -> None: + """Test list_table_names returns empty list on error.""" + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table_names to raise exception + mock_conn.table_names.side_effect = Exception("Connection error") + + store = LanceDBVectorIndexStore() + names = store.list_table_names() + + assert names == [] + + +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +def test_list_table_names_no_table_names_attr(mock_get_connection: Mock) -> None: + """Test list_table_names returns empty list when connection lacks table_names.""" + # Mock connection without table_names attribute + mock_conn = Mock(spec=[]) # Empty spec means no attributes + mock_get_connection.return_value = mock_conn + + store = LanceDBVectorIndexStore() + names = store.list_table_names() + + assert names == [] + + +# --- get_vector_dimension_async Tests (Issue #14) --- + + +@pytest.mark.asyncio +@patch( + "xagent.core.tools.core.RAG_tools.storage.lancedb_stores.get_connection_from_env" +) +async def test_get_vector_dimension_async_delegates_to_sync( + mock_get_connection: Mock, +) -> None: + """Test async version delegates to sync implementation.""" + from types import SimpleNamespace + + mock_conn = Mock() + mock_get_connection.return_value = mock_conn + + # Mock table with fixed-size vector field + mock_table = Mock() + mock_conn.open_table.return_value = mock_table + + mock_vector_type = SimpleNamespace(list_size=768) + mock_vector_field = SimpleNamespace(type=mock_vector_type) + mock_schema = Mock() + mock_schema.field.return_value = mock_vector_field + mock_table.schema = mock_schema + + store = LanceDBVectorIndexStore() + dimension = await store.get_vector_dimension_async("embeddings_async_test") + + assert dimension == 768 diff --git a/tests/core/tools/core/RAG_tools/storage/test_vector_backend.py b/tests/core/tools/core/RAG_tools/storage/test_vector_backend.py new file mode 100644 index 000000000..a087aee7e --- /dev/null +++ b/tests/core/tools/core/RAG_tools/storage/test_vector_backend.py @@ -0,0 +1,73 @@ +"""Tests for vector backend selection.""" + +from __future__ import annotations + +import pytest + +from xagent.core.tools.core.RAG_tools.core.exceptions import ConfigurationError +from xagent.core.tools.core.RAG_tools.storage.factory import StorageFactory +from xagent.core.tools.core.RAG_tools.storage.lancedb_stores import ( + LanceDBVectorIndexStore, +) +from xagent.core.tools.core.RAG_tools.storage.vector_backend import ( + VECTOR_BACKEND_ENV, + VECTOR_BACKEND_ENV_LEGACY, + VectorBackend, + get_configured_vector_backend, +) + + +@pytest.fixture() +def clean_vector_backend_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Remove backend env vars for isolated parsing.""" + monkeypatch.delenv(VECTOR_BACKEND_ENV, raising=False) + monkeypatch.delenv(VECTOR_BACKEND_ENV_LEGACY, raising=False) + + +def test_default_backend_is_lancedb(clean_vector_backend_env: None) -> None: + assert get_configured_vector_backend() is VectorBackend.LANCEDB + + +def test_xagent_env_takes_precedence( + monkeypatch: pytest.MonkeyPatch, clean_vector_backend_env: None +) -> None: + monkeypatch.setenv(VECTOR_BACKEND_ENV_LEGACY, "milvus") + monkeypatch.setenv(VECTOR_BACKEND_ENV, "lancedb") + assert get_configured_vector_backend() is VectorBackend.LANCEDB + + +def test_legacy_env_when_primary_unset( + monkeypatch: pytest.MonkeyPatch, clean_vector_backend_env: None +) -> None: + monkeypatch.setenv(VECTOR_BACKEND_ENV_LEGACY, "lancedb") + assert get_configured_vector_backend() is VectorBackend.LANCEDB + + +def test_invalid_backend_raises(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(VECTOR_BACKEND_ENV, "not-a-backend") + with pytest.raises(ConfigurationError, match="Invalid"): + get_configured_vector_backend() + + +def test_factory_creates_lancedb_store( + monkeypatch: pytest.MonkeyPatch, clean_vector_backend_env: None, tmp_path: str +) -> None: + monkeypatch.setenv("LANCEDB_DIR", str(tmp_path)) + monkeypatch.setenv(VECTOR_BACKEND_ENV, "lancedb") + StorageFactory.get_factory().reset_all() + store = StorageFactory.get_factory().get_vector_index_store() + assert isinstance(store, LanceDBVectorIndexStore) + assert ( + StorageFactory.get_factory().get_resolved_vector_backend() + is VectorBackend.LANCEDB + ) + + +def test_unimplemented_backend_raises( + monkeypatch: pytest.MonkeyPatch, clean_vector_backend_env: None, tmp_path: str +) -> None: + monkeypatch.setenv("LANCEDB_DIR", str(tmp_path)) + monkeypatch.setenv(VECTOR_BACKEND_ENV, "milvus") + StorageFactory.get_factory().reset_all() + with pytest.raises(ConfigurationError, match="not implemented"): + StorageFactory.get_factory().get_vector_index_store() diff --git a/tests/core/tools/core/RAG_tools/test_metadata_propagation.py b/tests/core/tools/core/RAG_tools/test_metadata_propagation.py index 317bc18b1..9c6028d77 100644 --- a/tests/core/tools/core/RAG_tools/test_metadata_propagation.py +++ b/tests/core/tools/core/RAG_tools/test_metadata_propagation.py @@ -25,12 +25,14 @@ format_search_results_for_llm, ) from xagent.core.tools.core.RAG_tools.retrieval.search_dense import search_dense +from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, +) from xagent.core.tools.core.RAG_tools.utils.metadata_utils import deserialize_metadata from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( read_chunks_for_embedding, write_vectors_to_db, ) -from xagent.providers.vector_store.lancedb import get_connection_from_env class _StubEmbeddingAdapter(BaseEmbedding): @@ -138,7 +140,7 @@ def test_metadata_preserved_in_chunks_table( assert chunk_result["chunk_count"] > 0 # Step 4: Verify metadata in chunks table - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() chunks_table = conn.open_table("chunks") df = ( chunks_table.search() @@ -204,9 +206,11 @@ def test_metadata_preserved_in_embeddings_table( from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( ensure_embeddings_table, ) - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() ensure_embeddings_table(conn, "test_model", vector_dim=2) # Step 4: Read chunks for embedding @@ -255,7 +259,7 @@ def test_metadata_preserved_in_embeddings_table( assert write_response.upsert_count > 0 # Step 6: Verify metadata in embeddings table - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() embeddings_table = conn.open_table("embeddings_test_model") df = ( embeddings_table.search() @@ -327,9 +331,11 @@ def test_metadata_in_search_results( from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( ensure_embeddings_table, ) - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() ensure_embeddings_table(conn, "test_model", vector_dim=2) read_response = read_chunks_for_embedding( @@ -445,9 +451,11 @@ def test_full_pipeline_metadata_preservation( from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( ensure_embeddings_table, ) - from xagent.providers.vector_store.lancedb import get_connection_from_env + from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, + ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() ensure_embeddings_table(conn, "test_model", vector_dim=2) read_response = read_chunks_for_embedding( diff --git a/tests/core/tools/core/RAG_tools/test_multitenancy.py b/tests/core/tools/core/RAG_tools/test_multitenancy.py index 3bcf1cfc9..66f5a2b2a 100644 --- a/tests/core/tools/core/RAG_tools/test_multitenancy.py +++ b/tests/core/tools/core/RAG_tools/test_multitenancy.py @@ -35,12 +35,14 @@ ) from xagent.core.tools.core.RAG_tools.parse.parse_document import parse_document from xagent.core.tools.core.RAG_tools.retrieval.search_engine import search_dense_engine +from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, +) from xagent.core.tools.core.RAG_tools.utils.user_permissions import UserPermissions from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( read_chunks_for_embedding, write_vectors_to_db, ) -from xagent.providers.vector_store.lancedb import get_connection_from_env from xagent.web.api.kb import delete_collection_api, list_collections_api @@ -136,7 +138,7 @@ def temp_lancedb_dir(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path): def _insert_test_documents(self, user_id: int | None): """Insert test documents with specific user_id.""" - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( ensure_documents_table, ) @@ -160,7 +162,8 @@ def _insert_test_documents(self, user_id: int | None): ] table.add(records) - def test_list_collections_admin_sees_all(self, temp_lancedb_dir: str) -> None: + @pytest.mark.asyncio + async def test_list_collections_admin_sees_all(self, temp_lancedb_dir: str) -> None: """Admin users should see all collections regardless of user_id.""" # Insert documents for different users self._insert_test_documents(user_id=1) @@ -168,7 +171,7 @@ def test_list_collections_admin_sees_all(self, temp_lancedb_dir: str) -> None: self._insert_test_documents(user_id=None) # Legacy data # Admin sees everything - result = list_collections(user_id=None, is_admin=True) + result = await list_collections(user_id=None, is_admin=True) assert result.status == "success" # Should see at least one collection assert len(result.collections) >= 1 @@ -176,7 +179,8 @@ def test_list_collections_admin_sees_all(self, temp_lancedb_dir: str) -> None: total_docs = sum(c.documents for c in result.collections) assert total_docs == 15 # 5 docs per user * 3 users - def test_list_collections_regular_user_sees_only_own( + @pytest.mark.asyncio + async def test_list_collections_regular_user_sees_only_own( self, temp_lancedb_dir: str ) -> None: """Regular users should only see their own documents.""" @@ -186,13 +190,13 @@ def test_list_collections_regular_user_sees_only_own( self._insert_test_documents(user_id=None) # User 1 sees only user 1's data - result = list_collections(user_id=1, is_admin=False) + result = await list_collections(user_id=1, is_admin=False) assert result.status == "success" total_docs = sum(c.documents for c in result.collections) assert total_docs == 5 # User 2 sees only user 2's data - result = list_collections(user_id=2, is_admin=False) + result = await list_collections(user_id=2, is_admin=False) assert result.status == "success" total_docs = sum(c.documents for c in result.collections) assert total_docs == 5 @@ -300,7 +304,7 @@ def test_search_regular_user_only_own_results( # Setup: Create embeddings table and insert test data for different users import pandas as pd - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() # Create embeddings table from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( @@ -380,7 +384,7 @@ def test_unauthenticated_search_hides_orphaned_records( """Unauthenticated dense search should not return orphaned sentinel records.""" import pandas as pd - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( ensure_embeddings_table, ) @@ -552,23 +556,16 @@ def teardown_method(self): shutil.rmtree(self.temp_dir, ignore_errors=True) @patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_documents_table" - ) - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_parses_table" - ) - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_chunks_table" - ) - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.management.collections.get_vector_index_store" ) - def test_list_collections_with_user_filter( - self, mock_get_conn, mock_ensure_chunks, mock_ensure_parses, mock_ensure_docs - ): + @pytest.mark.asyncio + async def test_list_collections_with_user_filter(self, mock_get_store): """Test list_collections applies user filtering.""" + mock_store = MagicMock() mock_conn = MagicMock() - mock_get_conn.return_value = mock_conn + mock_store.get_raw_connection.return_value = mock_conn + mock_store.aggregate_collection_stats.return_value = {} + mock_get_store.return_value = mock_store mock_docs_table = MagicMock() mock_conn.open_table.return_value = mock_docs_table @@ -595,86 +592,90 @@ def mock_open_table_side_effect(table_name): mock_conn.open_table.side_effect = mock_open_table_side_effect - result = list_collections(user_id=123, is_admin=False) + result = await list_collections(user_id=123, is_admin=False) assert hasattr(result, "status") assert hasattr(result, "collections") assert hasattr(result, "total_count") - result = list_collections(user_id=None, is_admin=True) + result = await list_collections(user_id=None, is_admin=True) assert hasattr(result, "status") assert hasattr(result, "collections") assert hasattr(result, "total_count") - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_documents_table" - ) - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_parses_table" - ) - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_chunks_table" - ) - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_ingestion_runs_table" - ) - @patch("xagent.core.tools.core.RAG_tools.management.status.get_connection_from_env") - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.get_connection_from_env" - ) - def test_delete_collection_permission_check( - self, - mock_get_conn, - mock_status_conn, - mock_ensure_runs, - mock_ensure_chunks, - mock_ensure_parses, - mock_ensure_docs, - ): - """Test delete_collection runs with user/admin context. + @pytest.mark.asyncio + async def test_delete_collection_with_real_storage(self): + """Test delete_collection with real storage (integration test). - Note: Current delete_collection uses _collect_document_ids with user filter - and deletes only what the user can see; it does not compare total vs - accessible count. So we only assert admin and user success paths. + This test verifies the complete data flow for delete_collection operation, + ensuring it correctly handles user/admin permissions with actual database + operations rather than mocked responses. """ - mock_conn = MagicMock() - mock_get_conn.return_value = mock_conn - mock_status_conn.return_value = mock_conn + from xagent.core.tools.core.RAG_tools.core.schemas import CollectionInfo + from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + CollectionManager, + ) + + # Setup: Create a collection for testing using CollectionManager + manager = CollectionManager() - mock_table = MagicMock() - mock_conn.open_table.return_value = mock_table - mock_table.count_rows.return_value = 0 + collection = CollectionInfo( + name=self.collection, + embedding_model_id="text-embedding-ada-002", + embedding_dimension=1536, + documents=5, + ) + await manager.save_collection(collection) + # Test: Admin can delete collection result = delete_collection(self.collection, user_id=None, is_admin=True) assert result.status == "success" - result = delete_collection(self.collection, user_id=123, is_admin=False) + # Setup: Create another collection for user-specific test + collection_user = CollectionInfo( + name=f"{self.collection}_user", + embedding_model_id="text-embedding-ada-002", + embedding_dimension=1536, + documents=5, + ) + await manager.save_collection(collection_user) + + # Test: User can delete their own collection + result = delete_collection( + f"{self.collection}_user", user_id=123, is_admin=False + ) assert result.status == "success" - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_documents_table" - ) - @patch("xagent.core.tools.core.RAG_tools.management.status.get_connection_from_env") - @patch( - "xagent.core.tools.core.RAG_tools.management.collections.get_connection_from_env" - ) - def test_retry_document_permission_check( - self, mock_get_conn, mock_status_conn, mock_ensure_docs - ): - """Test retry_document accepts user_id and is_admin and completes. + @pytest.mark.asyncio + async def test_retry_document_with_real_storage(self): + """Test retry_document with real storage (integration test). - Note: Current retry_document only calls write_ingestion_status and does not - check document existence or ownership via count_rows. We assert it returns - success when called with user and admin context. + This test verifies the complete data flow for retry_document operation, + ensuring it correctly handles user/admin permissions with actual database + operations rather than mocked responses. """ - mock_conn = MagicMock() - mock_get_conn.return_value = mock_conn - mock_status_conn.return_value = mock_conn + from xagent.core.tools.core.RAG_tools.core.schemas import CollectionInfo + from xagent.core.tools.core.RAG_tools.management.collection_manager import ( + CollectionManager, + ) + + # Setup: Create a collection for testing using CollectionManager + manager = CollectionManager() + + collection = CollectionInfo( + name=self.collection, + embedding_model_id="text-embedding-ada-002", + embedding_dimension=1536, + documents=5, + ) + await manager.save_collection(collection) + # Test: User can retry their own document result = retry_document( self.collection, "test_doc", user_id=123, is_admin=False ) assert result.status == "success" + # Test: Admin can retry any document result = retry_document( self.collection, "test_doc", user_id=None, is_admin=True ) @@ -769,6 +770,7 @@ class TestAPIMultiTenancy: """Test multi-tenancy at the API level.""" @patch("xagent.web.api.kb.list_collections") + @pytest.mark.asyncio async def test_list_collections_api_with_user(self, mock_list_collections): """Test list_collections_api passes user context.""" from xagent.web.models.user import User @@ -777,21 +779,32 @@ async def test_list_collections_api_with_user(self, mock_list_collections): mock_user.id = 123 mock_user.is_admin = False - mock_list_collections.return_value = {"collections": [], "total": 0} + # Mock async function return value + from xagent.core.tools.core.RAG_tools.core.schemas import ListCollectionsResult + + mock_result = ListCollectionsResult( + status="success", + total_count=0, + collections=[], + message="No collections found", + warnings=[], + ) + mock_list_collections.return_value = mock_result result = await list_collections_api(_user=mock_user) - mock_list_collections.assert_called_once_with(123, False) - assert result == {"collections": [], "total": 0} + mock_list_collections.assert_called_once_with(user_id=123, is_admin=False) + assert result.status == "success" + assert result.total_count == 0 - @patch("xagent.web.api.kb._list_documents_for_user", return_value=[]) + @patch("xagent.web.api.kb.get_vector_index_store") @patch("xagent.web.api.kb.delete_collection_physical_dir") @patch("xagent.web.api.kb.delete_collection") async def test_delete_collection_api_with_user( self, mock_delete_collection, mock_delete_collection_physical_dir, - _mock_list_documents_for_user, + mock_get_vector_store, ): """Test delete_collection_api passes user context and moves dir to trash.""" from xagent.core.tools.core.RAG_tools.core.schemas import ( @@ -806,6 +819,8 @@ async def test_delete_collection_api_with_user( mock_user.id = 123 mock_user.is_admin = False + mock_get_vector_store.return_value.list_document_records.return_value = [] + mock_path = MagicMock(spec=Path) mock_delete_collection_physical_dir.return_value = ( CollectionPhysicalDeleteResult( @@ -840,14 +855,14 @@ async def test_delete_collection_api_with_user( assert isinstance(result, CollectionOperationResult) assert result.status == "success" - @patch("xagent.web.api.kb._list_documents_for_user", return_value=[]) + @patch("xagent.web.api.kb.get_vector_index_store") @patch("xagent.web.api.kb.delete_collection_physical_dir") @patch("xagent.web.api.kb.delete_collection") async def test_delete_collection_api_admin_access( self, mock_delete_collection, mock_delete_collection_physical_dir, - _mock_list_documents_for_user, + mock_get_vector_store, ): """Test admin can delete collections (move dir to trash).""" from xagent.core.tools.core.RAG_tools.core.schemas import ( @@ -862,6 +877,8 @@ async def test_delete_collection_api_admin_access( mock_user.id = 999 mock_user.is_admin = True + mock_get_vector_store.return_value.list_document_records.return_value = [] + mock_path = MagicMock(spec=Path) mock_delete_collection_physical_dir.return_value = ( CollectionPhysicalDeleteResult( @@ -942,20 +959,17 @@ def test_user_data_isolation_workflow(self): with ( patch( - "xagent.core.tools.core.RAG_tools.management.collections.get_connection_from_env" - ) as mock_conn, - patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_documents_table" - ), - patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_parses_table" - ), - patch( - "xagent.core.tools.core.RAG_tools.management.collections.ensure_chunks_table" - ), + "xagent.core.tools.core.RAG_tools.management.collections.get_vector_index_store" + ) as mock_get_store, ): + mock_store = MagicMock() mock_db_conn = MagicMock() - mock_conn.return_value = mock_db_conn + mock_store.get_raw_connection.return_value = mock_db_conn + mock_get_store.return_value = mock_store + + # Mock new storage abstraction methods + mock_store.list_document_records.return_value = [] + mock_store.delete_collection_data.return_value = {} mock_docs_table = MagicMock() mock_db_conn.open_table.return_value = mock_docs_table @@ -965,8 +979,7 @@ def test_user_data_isolation_workflow(self): delete_collection, ) - # delete_collection uses _collect_document_ids (iter_batches), not count_rows - # for permission; it just deletes what the user can see. Assert it completes. + # delete_collection now uses list_document_records and delete_collection_data result = delete_collection( "test_collection", user_id=user1_id, is_admin=False ) diff --git a/tests/core/tools/core/RAG_tools/utils/test_migration_utils.py b/tests/core/tools/core/RAG_tools/utils/test_migration_utils.py index e11321b06..f4831b504 100644 --- a/tests/core/tools/core/RAG_tools/utils/test_migration_utils.py +++ b/tests/core/tools/core/RAG_tools/utils/test_migration_utils.py @@ -5,6 +5,7 @@ _model_tag_to_model_id, migrate_collection_metadata, ) +from xagent.core.tools.core.RAG_tools.utils.tag_mapping import register_tag_mapping class TestMigrateCollectionMetadata: @@ -68,12 +69,29 @@ def test_migrate_with_embedding_inference(self, mock_infer): assert result["embedding_dimension"] == 1536 mock_infer.assert_called_once_with("test_collection") + @patch( + "xagent.core.tools.core.RAG_tools.utils.migration_utils._infer_embedding_config_from_collection" + ) + def test_migrate_without_embedding_inference_skips_db(self, mock_infer): + """Read-safe migration must not scan LanceDB for embedding config.""" + legacy_data = { + "name": "test_collection", + "documents": 10, + } + + result = migrate_collection_metadata(legacy_data, infer_embedding=False) + + mock_infer.assert_not_called() + assert result["schema_version"] == "1.0.0" + assert result["embedding_model_id"] is None + assert result["embedding_dimension"] is None + class TestInferEmbeddingConfigFromCollection: """Test embedding config inference.""" @patch( - "xagent.core.tools.core.RAG_tools.utils.migration_utils.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.utils.migration_utils.get_vector_store_raw_connection" ) def test_infer_no_tables_found(self, mock_conn): """Test inference when no embedding tables exist.""" @@ -86,7 +104,7 @@ def test_infer_no_tables_found(self, mock_conn): assert result == (None, None) @patch( - "xagent.core.tools.core.RAG_tools.utils.migration_utils.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.utils.migration_utils.get_vector_store_raw_connection" ) def test_infer_single_model(self, mock_conn): """Test inference with single embedding model.""" @@ -116,7 +134,7 @@ def test_infer_single_model(self, mock_conn): assert result == ("text-embedding-ada-002", 1536) @patch( - "xagent.core.tools.core.RAG_tools.utils.migration_utils.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.utils.migration_utils.get_vector_store_raw_connection" ) def test_infer_multiple_models_choose_most_used(self, mock_conn): """Test inference with multiple models chooses most used.""" @@ -152,6 +170,30 @@ def mock_open_table(table_name): mock_logger.warning.assert_called_once() +class TestHubTagMapping: + """Test tag collision handling when building hub lookup maps.""" + + def test_register_hub_tag_mapping_warns_on_collision(self) -> None: + mapping = {"OPENAI_text_embedding_3_large": "hub-id-a"} + mock_logger = MagicMock() + + register_tag_mapping( + mapping, + "OPENAI_text_embedding_3_large", + "hub-id-b", + get_identity=lambda item: item, + logger=mock_logger, + ) + + assert mapping["OPENAI_text_embedding_3_large"] == "hub-id-a" + mock_logger.warning.assert_called_once_with( + "Tag collision: %s -> %s vs %s", + "OPENAI_text_embedding_3_large", + "hub-id-a", + "hub-id-b", + ) + + class TestModelTagToModelId: """Test model tag to model ID conversion.""" diff --git a/tests/core/tools/core/RAG_tools/utils/test_model_resolver_utils.py b/tests/core/tools/core/RAG_tools/utils/test_model_resolver_utils.py index 6984d2e0c..34c074f27 100644 --- a/tests/core/tools/core/RAG_tools/utils/test_model_resolver_utils.py +++ b/tests/core/tools/core/RAG_tools/utils/test_model_resolver_utils.py @@ -2,9 +2,11 @@ from __future__ import annotations +import sqlite3 from typing import Dict import pytest +from sqlalchemy.exc import OperationalError as SAOperationalError from xagent.core.model.chat.basic.base import BaseLLM from xagent.core.model.embedding.base import BaseEmbedding @@ -36,6 +38,31 @@ def load(self, model_id: str) -> object: return self._models[model_id] +class TestHubInitFailureClassification: + """Tests for _hub_init_failure_is_benign_optional_sqlite.""" + + def test_sqlite_missing_file_is_benign(self) -> None: + exc = sqlite3.OperationalError("unable to open database file") + assert model_resolver._hub_init_failure_is_benign_optional_sqlite(exc) is True + + def test_sqlalchemy_wrapped_sqlite_missing_is_benign(self) -> None: + inner = sqlite3.OperationalError("unable to open database file") + exc = SAOperationalError("SELECT 1", {}, inner) + assert model_resolver._hub_init_failure_is_benign_optional_sqlite(exc) is True + + def test_database_locked_not_benign(self) -> None: + exc = sqlite3.OperationalError("database is locked") + assert model_resolver._hub_init_failure_is_benign_optional_sqlite(exc) is False + + def test_other_errors_not_benign(self) -> None: + assert ( + model_resolver._hub_init_failure_is_benign_optional_sqlite( + RuntimeError("connection refused") + ) + is False + ) + + class TestGetOrInitModelHub: """Test _get_or_init_model_hub helper function.""" diff --git a/tests/core/tools/core/RAG_tools/vector_storage/test_embeddings_forward_migration.py b/tests/core/tools/core/RAG_tools/vector_storage/test_embeddings_forward_migration.py new file mode 100644 index 000000000..fe670749e --- /dev/null +++ b/tests/core/tools/core/RAG_tools/vector_storage/test_embeddings_forward_migration.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from typing import Any +from unittest.mock import patch + +import pandas as pd + +from xagent.core.model.model import EmbeddingModelConfig +from xagent.core.tools.core.RAG_tools.LanceDB.model_tag_utils import to_model_tag +from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( + ensure_embeddings_table, +) +from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_index_store, + reset_kb_write_coordinator, +) + + +def test_forward_migrate_legacy_embeddings_table_to_hub_id( + tmp_path: Any, monkeypatch: Any +) -> None: + """Legacy embeddings tables can be migrated to Hub-ID table names using storage API. + + Scenario: + - Only legacy table exists: embeddings_{to_model_tag(model_name)} + - Primary Hub-ID table missing: embeddings_{to_model_tag(hub_id)} + - Using migrate_embeddings_table() creates the primary table and copies rows + from legacy, rewriting row["model"] to hub_id. + """ + hub_id = "text-embedding-v4-openai-1" + legacy_model_name = "text-embedding-v4" + vector_dim = 3 + + monkeypatch.setenv("LANCEDB_DIR", str(tmp_path / ".lancedb")) + reset_kb_write_coordinator() + vector_store = get_vector_index_store() + conn = vector_store.get_raw_connection() + + legacy_tag = to_model_tag(legacy_model_name) + legacy_table_name = f"embeddings_{legacy_tag}" + ensure_embeddings_table(conn, legacy_tag, vector_dim=vector_dim) + legacy_table = conn.open_table(legacy_table_name) + + # Insert one legacy row (model stored as provider model_name in older versions) + legacy_table.add( + [ + { + "collection": "c1", + "doc_id": "d1", + "chunk_id": "ch1", + "parse_hash": "p1", + "model": legacy_model_name, + "vector": [0.1, 0.2, 0.3], + "text": "t", + "chunk_hash": "h", + "created_at": pd.Timestamp.now(tz="UTC"), + "vector_dimension": vector_dim, + "metadata": None, + "user_id": None, + } + ] + ) + + primary_table_name = f"embeddings_{to_model_tag(hub_id)}" + # Sanity: primary should not exist yet + assert primary_table_name not in set(conn.table_names()) # type: ignore[attr-defined] + + # Patch resolver so hub_id -> model_name mapping is available for migration. + cfg = EmbeddingModelConfig( + id=hub_id, + model_name=legacy_model_name, + model_provider="openai", + dimension=vector_dim, + api_key="k", + base_url="http://example", + timeout=1.0, + abilities=["embedding"], + ) + + with patch( + "xagent.core.tools.core.RAG_tools.utils.model_resolver.resolve_embedding_adapter", + return_value=(cfg, object()), + ): + # Use the storage layer migration method + result = vector_store.migrate_embeddings_table(hub_id) + + assert result["success"] is True + assert result["source_table"] == legacy_table_name + assert result["target_table"] == primary_table_name + assert result["rows_migrated"] == 1 + + # Verify primary table was created + assert primary_table_name in set(conn.table_names()) # type: ignore[attr-defined] + primary_table = conn.open_table(primary_table_name) + rows = primary_table.search().to_pandas() + assert len(rows) == 1 + assert rows.iloc[0]["model"] == hub_id diff --git a/tests/core/tools/core/RAG_tools/vector_storage/test_index_manager.py b/tests/core/tools/core/RAG_tools/vector_storage/test_index_manager.py deleted file mode 100644 index 9385c6124..000000000 --- a/tests/core/tools/core/RAG_tools/vector_storage/test_index_manager.py +++ /dev/null @@ -1,704 +0,0 @@ -"""Tests for index_manager functionality. - -This module tests the IndexManager class and related index management functions: -- Index creation and status checking -- Automatic index type selection (IVF_HNSW_SQ vs IVF_PQ) -- Configuration-driven indexing behavior -- Error handling and edge cases -""" - -import os -import tempfile -import uuid -from types import SimpleNamespace -from unittest.mock import Mock, patch - -import pytest - -from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy -from xagent.core.tools.core.RAG_tools.core.schemas import IndexMetric -from xagent.core.tools.core.RAG_tools.vector_storage.index_manager import ( - IndexManager, - get_index_manager, -) - - -class TestIndexManager: - """Test IndexManager class functionality.""" - - def test_init_with_default_policy(self): - """Test IndexManager initialization with default policy.""" - manager = IndexManager() - - assert isinstance(manager.policy, IndexPolicy) - assert manager.policy.enable_threshold_rows == 50_000 - assert manager.policy.ivfpq_threshold_rows == 10_000_000 - assert manager.policy.hnsw_params == {} - assert manager.policy.ivfpq_params == {} - - def test_init_with_custom_policy(self): - """Test IndexManager initialization with custom policy.""" - custom_policy = IndexPolicy( - enable_threshold_rows=100_000, - ivfpq_threshold_rows=5_000_000, - hnsw_params={"ef_construction": 200}, - ivfpq_params={"nlist": 1024}, - ) - manager = IndexManager(custom_policy) - - assert manager.policy.enable_threshold_rows == 100_000 - assert manager.policy.ivfpq_threshold_rows == 5_000_000 - assert manager.policy.hnsw_params == {"ef_construction": 200} - assert manager.policy.ivfpq_params == {"nlist": 1024} - - def test_readonly_mode(self): - """Test readonly mode behavior.""" - manager = IndexManager() - mock_table = Mock() - - status, advice = manager.check_and_create_index( - mock_table, "test_table", readonly=True - ) - - assert status == "readonly" - assert "Readonly mode" in advice - # Should not call any table methods - mock_table.to_pandas.assert_not_called() - mock_table.list_indices.assert_not_called() - - @patch("xagent.core.tools.core.RAG_tools.vector_storage.index_manager.logger") - def test_below_threshold_no_index(self, mock_logger): - """Test behavior when row count is below threshold.""" - manager = IndexManager() - mock_table = Mock() - mock_table.count_rows.return_value = 0 # Empty table - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - assert status == "below_threshold" - assert "below threshold" in advice - assert "50000" in advice - - def test_hnsw_index_creation(self): - """Test HNSW index creation for medium-sized datasets.""" - manager = IndexManager() - mock_table = Mock() - - # Mock table with 100,000 rows (between thresholds) - mock_table.count_rows.return_value = 100_000 - mock_table.list_indices.return_value = [] # No existing indexes - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - assert status == "index_building" - assert "HNSW index created" in advice - assert "using HNSW strategy for medium-scale data" in advice - mock_table.create_index.assert_called_once_with( - metric="l2", index_type="IVF_HNSW_SQ" - ) - - def test_ivfpq_index_creation(self): - """Test IVFPQ index creation for large datasets.""" - manager = IndexManager() - mock_table = Mock() - - # Mock table with 15M rows (above IVFPQ threshold) - mock_table.count_rows.return_value = 15_000_000 - mock_table.list_indices.return_value = [] # No existing indexes - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - assert status == "index_building" - assert "IVFPQ index created" in advice - assert "using IVFPQ strategy for large-scale data" in advice - mock_table.create_index.assert_called_once_with( - metric="l2", index_type="IVF_PQ" - ) - - def test_existing_index_skip(self): - """Test skipping when index already exists.""" - manager = IndexManager() - mock_table = Mock() - - # Mock table with enough rows and existing index - mock_table.count_rows.return_value = 100_000 - mock_table.list_indices.return_value = [SimpleNamespace(name="vector")] - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - assert status == "index_ready" - assert "Index ready" in advice - mock_table.create_index.assert_not_called() - - def test_index_creation_with_custom_params(self): - """Test index creation with custom parameters.""" - custom_policy = IndexPolicy( - hnsw_params={"ef_construction": 200, "M": 32}, - ivfpq_params={"nlist": 1024, "nprobe": 10}, - ) - manager = IndexManager(custom_policy) - mock_table = Mock() - - # Test HNSW with custom params - mock_table.count_rows.return_value = 100_000 - mock_table.list_indices.return_value = [] - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - # Check that create_index was called with correct parameters - mock_table.create_index.assert_called_once() - call_args = mock_table.create_index.call_args - - # Check keyword arguments - kwargs = call_args[1] - assert kwargs["metric"] == "l2" - assert kwargs["index_type"] == "IVF_HNSW_SQ" - assert kwargs["ef_construction"] == 200 - assert kwargs["M"] == 32 - - def test_index_creation_error_handling(self): - """Test error handling during index creation.""" - manager = IndexManager() - mock_table = Mock() - - # Mock table operations - mock_table.count_rows.return_value = 100_000 - mock_table.list_indices.return_value = [] - mock_table.create_index.side_effect = Exception("Index creation failed") - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - assert status == "index_corrupted" - assert "Vector index check failed" in advice - assert "Index creation failed" in advice - - def test_get_index_status_ready(self): - """Test getting index status when index exists.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [SimpleNamespace(name="vector")] - - status = manager.get_index_status(mock_table) - assert status == "index_ready" - - def test_get_index_status_no_index(self): - """Test getting index status when no index exists but above threshold.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [] - - # Mock enough rows for indexing - mock_table.count_rows.return_value = 100_000 - - status = manager.get_index_status(mock_table) - assert status == "no_index" - - def test_get_index_status_below_threshold(self): - """Test getting index status when below threshold.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [] - - # Mock few rows - mock_table.count_rows.return_value = 1000 - - status = manager.get_index_status(mock_table) - assert status == "below_threshold" - - def test_get_index_status_error(self): - """Test error handling in get_index_status.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.side_effect = Exception("Database error") - - status = manager.get_index_status(mock_table) - assert status == "index_corrupted" - - -class TestIndexManagerIntegration: - """Integration tests for IndexManager with real LanceDB operations.""" - - @pytest.fixture - def temp_lancedb_dir(self): - """Create a temporary directory for LanceDB.""" - with tempfile.TemporaryDirectory() as temp_dir: - original_env = os.environ.get("LANCEDB_DIR") - os.environ["LANCEDB_DIR"] = temp_dir - yield temp_dir - if original_env is not None: - os.environ["LANCEDB_DIR"] = original_env - else: - os.environ.pop("LANCEDB_DIR", None) - - @pytest.fixture - def test_collection(self): - """Test collection name.""" - return f"test_collection_{uuid.uuid4().hex[:8]}" - - def test_end_to_end_index_creation(self, temp_lancedb_dir, test_collection): - """Test end-to-end index creation workflow.""" - from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( - ensure_embeddings_table, - ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - - conn = get_connection_from_env() - model_tag = "test_model" - - # Create embeddings table - ensure_embeddings_table(conn, model_tag, vector_dim=3) - table = conn.open_table(f"embeddings_{model_tag}") - - # Add some test data - import pandas as pd - - test_records = [ - { - "collection": test_collection, - "doc_id": f"doc_{i}", - "chunk_id": f"chunk_{i}", - "parse_hash": "test_parse", - "model": model_tag, - "vector": [1.0, 2.0, 3.0], - "vector_dimension": 3, - "text": f"test text {i}", - "chunk_hash": f"hash_{i}", - "created_at": pd.Timestamp.now(tz="UTC"), - } - for i in range(60000) # Add 60,000 records to trigger indexing - ] - table.add(test_records) - - # Test index manager - manager = IndexManager() - status, advice = manager.check_and_create_index( - table, f"embeddings_{model_tag}" - ) - - assert status in ["index_building", "index_ready"] - if status == "index_building": - assert "HNSW index created" in advice - - def test_custom_policy_integration(self, temp_lancedb_dir, test_collection): - """Test custom policy integration.""" - from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( - ensure_embeddings_table, - ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - - # Create custom policy with lower threshold - custom_policy = IndexPolicy( - enable_threshold_rows=50, # Very low threshold for testing - hnsw_params={"ef_construction": 100}, - ) - - conn = get_connection_from_env() - model_tag = "custom_policy_model" - - # Create table and add minimal data - ensure_embeddings_table(conn, model_tag, vector_dim=3) - table = conn.open_table(f"embeddings_{model_tag}") - - import pandas as pd - - # Add enough records to exceed the custom threshold of 50 - test_records = [ - { - "collection": test_collection, - "doc_id": f"doc_{i}", - "chunk_id": f"chunk_{i}", - "parse_hash": "test_parse", - "model": model_tag, - "vector": [1.0, 2.0, 3.0], - "vector_dimension": 3, - "text": f"test text {i}", - "chunk_hash": f"hash_{i}", - "created_at": pd.Timestamp.now(tz="UTC"), - } - for i in range(100) # Add 100 records to exceed threshold of 50 - ] - table.add(test_records) - - # Test with custom policy - manager = IndexManager(custom_policy) - status, advice = manager.check_and_create_index( - table, f"embeddings_{model_tag}" - ) - - # Should trigger indexing due to low threshold - assert status == "index_building" - assert "HNSW index created" in advice - - -class TestGetIndexManager: - """Test get_index_manager function.""" - - def test_get_default_manager(self): - """Test getting default index manager.""" - manager = get_index_manager() - assert isinstance(manager, IndexManager) - assert isinstance(manager.policy, IndexPolicy) - - def test_get_manager_with_custom_policy(self): - """Test getting manager with custom policy.""" - custom_policy = IndexPolicy(enable_threshold_rows=100_000) - manager = get_index_manager(custom_policy) - - assert manager.policy.enable_threshold_rows == 100_000 - - def test_manager_singleton_behavior(self): - """Test that get_index_manager returns the same instance when no policy is provided.""" - # Test default singleton behavior (no policy provided) - manager1 = get_index_manager() - manager2 = get_index_manager() - - # Should return the same instance when no policy is provided - assert manager1 is manager2 - - def test_manager_creates_new_instance_with_policy(self): - """Test that get_index_manager creates new instances when policy is provided.""" - policy1 = IndexPolicy(enable_threshold_rows=50000) - policy2 = IndexPolicy(enable_threshold_rows=50000) - - manager1 = get_index_manager(policy1) - manager2 = get_index_manager(policy2) # Same policy values - - # Should return different instances when policy is provided (current design) - assert manager1 is not manager2 - # But they should have the same policy values - assert ( - manager1.policy.enable_threshold_rows - == manager2.policy.enable_threshold_rows - ) - - def test_manager_different_instances(self): - """Test that different policies create different managers.""" - policy1 = IndexPolicy(enable_threshold_rows=50000) - policy2 = IndexPolicy(enable_threshold_rows=100000) - - manager1 = get_index_manager(policy1) - manager2 = get_index_manager(policy2) - - # Should be different instances for different policies - assert manager1 is not manager2 - assert manager1.policy.enable_threshold_rows == 50000 - assert manager2.policy.enable_threshold_rows == 100000 - - -class TestIndexMetricSupport: - """Test IndexMetric parameter support in IndexManager.""" - - def test_default_metric_l2(self): - """Test that default metric is L2.""" - manager = IndexManager() - assert manager.policy.metric == IndexMetric.L2 - assert manager.policy.metric.value == "l2" - - def test_custom_metric_cosine(self): - """Test index creation with COSINE metric.""" - custom_policy = IndexPolicy(metric=IndexMetric.COSINE) - manager = IndexManager(custom_policy) - mock_table = Mock() - - # Mock table with enough rows for indexing - mock_table.count_rows.return_value = 100_000 - mock_table.list_indices.return_value = [] - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - assert status == "index_building" - assert "cosine" in advice - # Verify that create_index was called with the correct metric - call_args = mock_table.create_index.call_args - assert call_args[1]["metric"] == "cosine" - - def test_custom_metric_dot(self): - """Test index creation with DOT metric.""" - custom_policy = IndexPolicy(metric=IndexMetric.DOT) - manager = IndexManager(custom_policy) - mock_table = Mock() - - # Mock table with enough rows for indexing - mock_table.count_rows.return_value = 100_000 - mock_table.list_indices.return_value = [] - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - assert status == "index_building" - assert "dot" in advice - # Verify that create_index was called with the correct metric - call_args = mock_table.create_index.call_args - assert call_args[1]["metric"] == "dot" - - -class TestFTSIndexSupport: - """Test FTS index functionality in IndexManager.""" - - def test_get_fts_index_status_no_index(self): - """Test FTS index status when no FTS index exists.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [] - - status = manager.get_fts_index_status(mock_table) - assert status is False - - def test_get_fts_index_status_with_vector_index_only(self): - """Test FTS index status when only vector index exists.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [ - SimpleNamespace(name="vector", index_type="IvfHnswSq", columns=["vector"]) - ] - - status = manager.get_fts_index_status(mock_table) - assert status is False - - def test_get_fts_index_status_with_fts_index(self): - """Test FTS index status when FTS index exists.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [ - SimpleNamespace(index_type="FTS", columns=["text"]) - ] - - status = manager.get_fts_index_status(mock_table) - assert status is True - - def test_get_fts_index_status_error_handling(self): - """Test FTS index status error handling.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.side_effect = Exception("Database error") - mock_table.name = "test_table" - - status = manager.get_fts_index_status(mock_table) - assert status is False - - def test_create_fts_index_success(self): - """Test successful FTS index creation.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [] # No existing FTS index - - success, message = manager.create_fts_index(mock_table, "test_table") - - assert success is True - assert "FTS index created" in message - assert "with_position" in message - # Verify create_index was called with correct parameters - mock_table.create_fts_index.assert_called_once_with( - "text", replace=True, with_position=True - ) - - def test_create_fts_index_already_exists(self): - """Test FTS index creation when index already exists.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [ - SimpleNamespace(index_type="FTS", columns=["text"]) - ] - - success, message = manager.create_fts_index(mock_table, "test_table") - - assert success is True - assert "already exists" in message - mock_table.create_fts_index.assert_not_called() - - def test_create_fts_index_with_custom_params(self): - """Test FTS index creation with custom parameters.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [] - - custom_params = { - "language": "english", - "stem": True, - "ascii_folding": True, - } - success, message = manager.create_fts_index( - mock_table, "test_table", fts_params=custom_params - ) - - assert success is True - assert "FTS index created" in message - # Verify create_index was called with merged parameters - expected_params = {"with_position": True, **custom_params} - mock_table.create_fts_index.assert_called_once_with( - "text", replace=True, **expected_params - ) - - def test_create_fts_index_error_handling(self): - """Test FTS index creation error handling.""" - manager = IndexManager() - mock_table = Mock() - mock_table.list_indices.return_value = [] - mock_table.create_fts_index.side_effect = Exception("FTS creation failed") - - success, message = manager.create_fts_index(mock_table, "test_table") - - assert success is False - assert "Failed to create FTS index" in message - assert "FTS creation failed" in message - - def test_check_and_create_index_with_fts_enabled(self): - """Test that check_and_create_index attempts FTS creation when enabled.""" - custom_policy = IndexPolicy( - fts_enabled=True, fts_params={"language": "english"} - ) - manager = IndexManager(custom_policy) - mock_table = Mock() - - # Mock table with enough rows for vector indexing - mock_table.count_rows.return_value = 100_000 - mock_table.list_indices.return_value = [] # No existing indexes - - status, advice = manager.check_and_create_index(mock_table, "test_table") - - # Should create both vector and FTS indexes - assert status == "index_building" - # Verify vector index creation - assert mock_table.create_index.call_count == 1 - vector_call = mock_table.create_index.call_args_list[0] - assert vector_call[1]["index_type"] == "IVF_HNSW_SQ" - mock_table.create_fts_index.assert_called_once() - fts_call = mock_table.create_fts_index.call_args - assert fts_call[0][0] == "text" - assert fts_call[1]["replace"] is True - - -class TestReindexingIntegration: - """Test reindexing functionality integration with IndexManager.""" - - def test_reindex_trigger_conditions(self): - """Test various conditions that should trigger reindexing.""" - from unittest.mock import MagicMock - - # Test with different policy configurations - policies = [ - # Immediate reindex - IndexPolicy(enable_immediate_reindex=True), - # Batch size threshold - IndexPolicy(reindex_batch_size=100), - # Smart reindex with ratio threshold - IndexPolicy( - enable_smart_reindex=True, reindex_unindexed_ratio_threshold=0.05 - ), - ] - - for policy in policies: - manager = IndexManager(policy) - mock_table = MagicMock() - - # Mock table with existing index - mock_table.count_rows.return_value = 100_000 - mock_table.list_indices.return_value = [SimpleNamespace(name="vector")] - - # Test that existing index is detected - status, advice = manager.check_and_create_index(mock_table, "test_table") - assert status == "index_ready" - assert "Index ready" in advice - - def test_reindex_with_optimize_call(self): - """Test that reindexing calls table.optimize().""" - from unittest.mock import MagicMock - - # Import the reindex functions from vector_manager - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _should_reindex, - _trigger_reindex, - ) - - mock_table = MagicMock() - policy = IndexPolicy(enable_immediate_reindex=True) - - # Test _should_reindex returns True for immediate mode - should_reindex = _should_reindex(mock_table, "test_table", 1, policy) - assert should_reindex is True - - # Test _trigger_reindex calls optimize - mock_table.optimize.return_value = None - reindex_success = _trigger_reindex(mock_table, "test_table") - - assert reindex_success is True - mock_table.optimize.assert_called_once() - - def test_reindex_error_handling(self): - """Test reindex error handling.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _trigger_reindex, - ) - - mock_table = MagicMock() - mock_table.optimize.side_effect = Exception("Optimize failed") - - reindex_success = _trigger_reindex(mock_table, "test_table") - - assert reindex_success is False - mock_table.optimize.assert_called_once() - - def test_smart_reindex_with_index_stats(self): - """Test smart reindex based on index statistics.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _should_reindex, - ) - - mock_table = MagicMock() - policy = IndexPolicy( - enable_smart_reindex=True, reindex_unindexed_ratio_threshold=0.05 - ) - - # Mock index stats showing high unindexed ratio - mock_stats = MagicMock() - mock_stats.num_indexed_rows = 1000 - mock_stats.num_unindexed_rows = 60 # 6% > 5% threshold - mock_table.index_stats.return_value = mock_stats - - should_reindex = _should_reindex(mock_table, "test_table", 10, policy) - assert should_reindex is True - - # Test below threshold - mock_stats.num_unindexed_rows = 30 # 3% < 5% threshold - should_reindex = _should_reindex(mock_table, "test_table", 10, policy) - assert should_reindex is False - - def test_batch_size_reindex_threshold(self): - """Test batch size threshold for reindexing.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _should_reindex, - ) - - mock_table = MagicMock() - policy = IndexPolicy(reindex_batch_size=100) - - # Test above batch threshold - should_reindex = _should_reindex(mock_table, "test_table", 150, policy) - assert should_reindex is True - - # Test below batch threshold - should_reindex = _should_reindex(mock_table, "test_table", 50, policy) - assert should_reindex is False - - def test_reindex_with_index_stats_error(self): - """Test reindex behavior when index stats fail.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _should_reindex, - ) - - mock_table = MagicMock() - policy = IndexPolicy(enable_smart_reindex=True) - - # Mock index_stats to raise exception - mock_table.index_stats.side_effect = Exception("Stats failed") - - # Should not trigger reindex when stats fail - should_reindex = _should_reindex(mock_table, "test_table", 10, policy) - assert should_reindex is False diff --git a/tests/core/tools/core/RAG_tools/vector_storage/test_vector_manager.py b/tests/core/tools/core/RAG_tools/vector_storage/test_vector_manager.py index 65a8bfc44..df32fb895 100644 --- a/tests/core/tools/core/RAG_tools/vector_storage/test_vector_manager.py +++ b/tests/core/tools/core/RAG_tools/vector_storage/test_vector_manager.py @@ -12,7 +12,6 @@ import uuid from unittest.mock import MagicMock, patch -import pandas as pd import pytest from xagent.core.tools.core.RAG_tools.core.exceptions import VectorValidationError @@ -95,38 +94,18 @@ def test_read_chunks_for_embedding_sql_injection_protection( """Test read_chunks_for_embedding protects against SQL injection.""" from unittest.mock import MagicMock - # Create mock connection and tables - mock_db_connection = MagicMock() - mock_chunks_table = _create_mock_table_with_schema() - mock_embeddings_table = _create_mock_table_with_schema() - - # Configure open_table to return appropriate mock tables using side_effect - def mock_open_table_func(table_name): - if table_name == "chunks": - return mock_chunks_table - elif table_name.startswith("embeddings_"): - return mock_embeddings_table - return MagicMock() + # Create mock vector store + mock_vector_store = MagicMock() - mock_db_connection.open_table.side_effect = mock_open_table_func - # Mock create_table to do nothing (tables are "created" but we use our mocks) - mock_db_connection.create_table.return_value = None + # Mock count_rows_or_zero to return 0 (no chunks found) + mock_vector_store.count_rows_or_zero.return_value = 0 - # UPDATED: Mock both to_list() and to_pandas() for optimization support - # Mock empty results for chunks - mock_chunks_table.search.return_value.where.return_value.to_list.return_value = [] - mock_chunks_table.search.return_value.where.return_value.to_pandas.return_value = pd.DataFrame() - mock_chunks_table.count_rows.return_value = ( - 0 # Changed to 0 to match empty results - ) - - # Mock empty results for embeddings - mock_embeddings_table.search.return_value.where.return_value.select.return_value.to_list.return_value = [] - mock_embeddings_table.search.return_value.where.return_value.select.return_value.to_pandas.return_value = pd.DataFrame() + # Mock iter_batches to return empty batches + mock_vector_store.iter_batches.return_value = [] with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): malicious_input = "malicious' OR 1=1 --" safe_collection = test_collection @@ -142,25 +121,20 @@ def mock_open_table_func(table_name): is_admin=True, # Use admin to avoid user_id filter ) - # Verify count_rows was called with escaped input - # Single quotes should be doubled: ' becomes '' - expected_chunks_where_clause = ( - f"collection == '{safe_collection}' AND " - f"doc_id == 'malicious'' OR 1=1 --' AND " - f"parse_hash == '{safe_parse_hash}'" - ) - mock_chunks_table.count_rows.assert_called_once_with( - expected_chunks_where_clause - ) - - # Since count_rows returns 0, search() should not be called - mock_chunks_table.search.assert_not_called() + # Verify count_rows_or_zero was called on vector store + mock_vector_store.count_rows_or_zero.assert_called_once() + call_kwargs = mock_vector_store.count_rows_or_zero.call_args[1] + assert call_kwargs["table_name"] == "chunks" + # Verify filters were passed correctly (including the malicious input) + assert "collection" in call_kwargs["filters"] + assert call_kwargs["filters"]["doc_id"] == malicious_input + assert call_kwargs["filters"]["parse_hash"] == safe_parse_hash - # Since no chunks exist, embeddings table should not be queried - mock_embeddings_table.search.assert_not_called() + # Since count_rows_or_zero returns 0, iter_batches should not be called + mock_vector_store.iter_batches.assert_not_called() assert result.chunks == [] - assert result.total_count == 0 # Changed from 1 to 0 + assert result.total_count == 0 assert result.pending_count == 0 @@ -358,47 +332,20 @@ def test_write_vectors_to_db_sql_injection_protection( from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - # Create mock connection and table - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() + mock_vector_store.upsert_embeddings.return_value = None + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - # Configure open_table to return the mock embeddings table using side_effect - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - # Mock create_table to do nothing (tables are "created" but we use our mocks) - mock_db_connection.create_table.return_value = None - - # Mock search to return empty DataFrame so no deletions happen initially - mock_embeddings_table.search.return_value.where.return_value.to_pandas.return_value = pd.DataFrame() - # Mock merge_insert method and its chain calls - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_execute = MagicMock() - - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - mock_when_not_matched.execute.return_value = mock_execute - # Keep add method for fallback testing - mock_embeddings_table.add.return_value = None # Mock add method - mock_embeddings_table.__len__.return_value = 0 # Mock len for index creation - mock_embeddings_table.count_rows.return_value = ( - 0 # Mock count_rows for index creation - ) - mock_embeddings_table.create_index.return_value = ( - None # Mock create_index method + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, ) with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): malicious_doc_id = "malicious' OR 1=1 --" safe_collection = test_collection @@ -424,27 +371,16 @@ def mock_open_table_func(table_name): embeddings=[malicious_embedding], ) - # With merge_insert, we no longer need to search for existing records - # merge_insert handles upsert automatically based on primary keys - # Verify that search was not called (merge_insert doesn't need it) - mock_embeddings_table.search.assert_not_called() - # Verify that delete was not called (merge_insert handles updates automatically) - mock_embeddings_table.delete.assert_not_called() - # Verify that merge_insert was called with the correct data - mock_embeddings_table.merge_insert.assert_called_once() - # Get the records argument from execute() method call - call_args = mock_when_not_matched.execute.call_args[0][0] - assert len(call_args) == 1 - assert call_args[0]["doc_id"] == malicious_doc_id - assert call_args[0]["chunk_id"] == malicious_chunk_id - - # Verify the chain calls were made - mock_merge_insert.when_matched_update_all.assert_called_once() - mock_when_matched.when_not_matched_insert_all.assert_called_once() - mock_when_not_matched.execute.assert_called_once() - - # Verify that add was not called (since merge_insert succeeded) - mock_embeddings_table.add.assert_not_called() + # Verify upsert_embeddings was called on vector store + mock_vector_store.upsert_embeddings.assert_called_once() + call_args = mock_vector_store.upsert_embeddings.call_args + records_arg = call_args[0][1] + + # Verify the records contain the malicious input (properly escaped by LanceDB) + assert len(records_arg) == 1 + assert records_arg[0]["doc_id"] == malicious_doc_id + assert records_arg[0]["chunk_id"] == malicious_chunk_id + assert records_arg[0]["collection"] == safe_collection assert result.upsert_count == 1 assert result.deleted_stale_count == 0 @@ -453,41 +389,32 @@ def mock_open_table_func(table_name): def test_write_vectors_merge_insert_fallback_to_add( self, temp_lancedb_dir, test_collection ): - """Test merge_insert failure fallback to add method.""" - from unittest.mock import MagicMock, patch + """Test merge_insert failure fallback to add method. + + NOTE: This test has been simplified for Phase 1A. + The actual merge_insert -> add() fallback logic is now implemented + in LanceDBVectorIndexStore.upsert_embeddings() and should be + tested in test_lancedb_stores.py. This test only verifies that + vector_store.upsert_embeddings is called correctly. + """ + from unittest.mock import MagicMock from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() + mock_vector_store.upsert_embeddings.return_value = None + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] - - # Mock merge_insert to fail, then add to succeed - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, ) - # merge_insert fails - mock_when_not_matched.execute.side_effect = Exception("merge_insert failed") - # add succeeds - mock_embeddings_table.add.return_value = None - mock_embeddings_table.count_rows.return_value = 0 with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -505,16 +432,21 @@ def mock_open_table_func(table_name): embeddings=[embedding], ) - # Verify merge_insert was attempted - mock_embeddings_table.merge_insert.assert_called_once() - # Verify fallback to add was used - mock_embeddings_table.add.assert_called_once() + # Verify upsert_embeddings was called on vector store + mock_vector_store.upsert_embeddings.assert_called_once() assert result.upsert_count == 1 def test_write_vectors_merge_insert_non_recoverable_error_no_fallback( self, temp_lancedb_dir, test_collection ): - """Test that non-recoverable errors (schema, type mismatch) do not fallback to add.""" + """Test that non-recoverable errors propagate correctly. + + NOTE: This test has been simplified for Phase 1A. + Non-recoverable error handling is now implemented in + LanceDBVectorIndexStore.upsert_embeddings() and should be + tested in test_lancedb_stores.py. This test only verifies + that errors propagate correctly through vector_manager. + """ from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.exceptions import ( @@ -522,35 +454,15 @@ def test_write_vectors_merge_insert_non_recoverable_error_no_fallback( ) from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] - - # Mock merge_insert to fail with schema error (non-recoverable) - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - # Schema error - should not fallback - mock_when_not_matched.execute.side_effect = ValueError( + # Create mock vector store that raises error + mock_vector_store = MagicMock() + mock_vector_store.upsert_embeddings.side_effect = ValueError( "Schema mismatch: expected int, got string" ) with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -564,21 +476,21 @@ def mock_open_table_func(table_name): ) # ValueError is wrapped in DatabaseOperationError by outer exception handler - with pytest.raises(DatabaseOperationError, match="Schema mismatch"): + with pytest.raises( + DatabaseOperationError, match="Failed to write embeddings" + ): write_vectors_to_db( collection=test_collection, embeddings=[embedding], ) - # Verify merge_insert was attempted - mock_embeddings_table.merge_insert.assert_called_once() - # Verify add was NOT called (no fallback for non-recoverable errors) - mock_embeddings_table.add.assert_not_called() + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() def test_write_vectors_merge_insert_type_mismatch_error_no_fallback( self, temp_lancedb_dir, test_collection ): - """Test that type mismatch errors do not fallback to add.""" + """Test that type mismatch errors do not fallback to add (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.exceptions import ( @@ -586,35 +498,24 @@ def test_write_vectors_merge_insert_type_mismatch_error_no_fallback( ) from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] + # Create mock vector store + mock_vector_store = MagicMock() - # Mock merge_insert to fail with type error (non-recoverable) - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - # Type error - should not fallback - mock_when_not_matched.execute.side_effect = TypeError( + # Mock upsert_embeddings to fail with type error (non-recoverable) + mock_vector_store.upsert_embeddings.side_effect = TypeError( "Type mismatch: invalid type for field" ) + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, + ) with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -628,19 +529,21 @@ def mock_open_table_func(table_name): ) # TypeError is wrapped in DatabaseOperationError by outer exception handler - with pytest.raises(DatabaseOperationError, match="Type mismatch"): + with pytest.raises( + DatabaseOperationError, match="Failed to write embeddings" + ): write_vectors_to_db( collection=test_collection, embeddings=[embedding], ) - # Verify add was NOT called - mock_embeddings_table.add.assert_not_called() + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() def test_write_vectors_merge_insert_dimension_error_no_fallback( self, temp_lancedb_dir, test_collection ): - """Test that dimension mismatch errors do not fallback to add.""" + """Test that dimension mismatch errors do not fallback to add (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.exceptions import ( @@ -648,35 +551,24 @@ def test_write_vectors_merge_insert_dimension_error_no_fallback( ) from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] - - # Mock merge_insert to fail with dimension error (non-recoverable) - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - # Dimension error - should not fallback - mock_when_not_matched.execute.side_effect = ValueError( + # Mock upsert_embeddings to fail with dimension error (non-recoverable) + mock_vector_store.upsert_embeddings.side_effect = ValueError( "Vector dimension mismatch: expected 3, got 2" ) + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, + ) with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -690,55 +582,41 @@ def mock_open_table_func(table_name): ) # ValueError is wrapped in DatabaseOperationError by outer exception handler - with pytest.raises(DatabaseOperationError, match="dimension mismatch"): + with pytest.raises( + DatabaseOperationError, match="Failed to write embeddings" + ): write_vectors_to_db( collection=test_collection, embeddings=[embedding], ) - # Verify add was NOT called - mock_embeddings_table.add.assert_not_called() + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() def test_write_vectors_merge_insert_recoverable_error_with_fallback( self, temp_lancedb_dir, test_collection ): - """Test that recoverable errors (network, timeout) do fallback to add.""" + """Test that recoverable errors (network, timeout) do fallback to add (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] + # Mock upsert_embeddings to succeed (it handles fallback internally) + mock_vector_store.upsert_embeddings.return_value = None + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - # Mock merge_insert to fail with network error (recoverable) - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, ) - # Network/timeout error - should fallback - mock_when_not_matched.execute.side_effect = ConnectionError( - "Network timeout: connection lost" - ) - # add succeeds - mock_embeddings_table.add.return_value = None - mock_embeddings_table.count_rows.return_value = 0 with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -756,16 +634,14 @@ def mock_open_table_func(table_name): embeddings=[embedding], ) - # Verify merge_insert was attempted - mock_embeddings_table.merge_insert.assert_called_once() - # Verify fallback to add was used - mock_embeddings_table.add.assert_called_once() + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() assert result.upsert_count == 1 def test_write_vectors_merge_insert_and_add_both_fail( self, temp_lancedb_dir, test_collection ): - """Test when both merge_insert and add fail.""" + """Test when both merge_insert and add fail (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.exceptions import ( @@ -773,33 +649,15 @@ def test_write_vectors_merge_insert_and_add_both_fail( ) from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] + # Create mock vector store + mock_vector_store = MagicMock() - # Both merge_insert and add fail - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - mock_when_not_matched.execute.side_effect = Exception("merge_insert failed") - mock_embeddings_table.add.side_effect = Exception("add also failed") + # Mock upsert_embeddings to fail + mock_vector_store.upsert_embeddings.side_effect = Exception("upsert failed") with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -818,42 +676,27 @@ def mock_open_table_func(table_name): embeddings=[embedding], ) + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() + def test_write_vectors_spill_retry(self, temp_lancedb_dir, test_collection): - """Test that spill error reduces batch size and retries without losing data.""" + """Test that spill error reduces batch size and retries without losing data (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] + # Mock upsert_embeddings to succeed (it handles spill retry internally) + mock_vector_store.upsert_embeddings.return_value = None + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - # First execute() raises spill; subsequent succeed - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, ) - mock_when_not_matched.execute.side_effect = [ - Exception("Spill has sent an error"), - None, - None, - None, - None, - None, - ] - mock_embeddings_table.count_rows.return_value = 0 embeddings = [ ChunkEmbeddingData( @@ -871,8 +714,8 @@ def mock_open_table_func(table_name): with ( patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ), patch.dict(os.environ, {"LANCEDB_BATCH_SIZE": "2"}, clear=False), ): @@ -883,7 +726,9 @@ def mock_open_table_func(table_name): ) assert result.upsert_count == 5 - assert mock_embeddings_table.merge_insert.call_count >= 2 + # Verify upsert_embeddings was called 3 times (5 records with batch_size=2) + # Batch 1: doc_0, doc_1; Batch 2: doc_2, doc_3; Batch 3: doc_4 + assert mock_vector_store.upsert_embeddings.call_count == 3 def test_write_vectors_batch_partial_failure( self, temp_lancedb_dir, test_collection @@ -942,44 +787,55 @@ def mock_merge_insert_side_effect(*args, **kwargs): return mock_merge_insert mock_embeddings_table.merge_insert.side_effect = mock_merge_insert_side_effect - # add succeeds for fallback - mock_embeddings_table.add.return_value = None - mock_embeddings_table.count_rows.return_value = 0 + # Create mock vector store that uses our mock connection/table + mock_vector_store = MagicMock() + + def mock_upsert_side_effect(model_tag, records): + # Simulate real upsert behavior by calling merge_insert on our mock table + mock_embeddings_table.merge_insert( + ["collection", "doc_id", "parse_hash", "chunk_id"] + ).when_matched_update_all().when_not_matched_insert_all().execute(records) + + mock_vector_store.upsert_embeddings.side_effect = mock_upsert_side_effect with ( patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ), patch.dict(os.environ, {"LANCEDB_BATCH_SIZE": "2"}), ): # Small batch size - result = write_vectors_to_db( - collection=test_collection, - embeddings=embeddings, + # Now we expect it to raise DatabaseOperationError instead of partial success + from xagent.core.tools.core.RAG_tools.core.exceptions import ( + DatabaseOperationError, ) - # Some batches should have succeeded - assert result.upsert_count > 0 + with pytest.raises(DatabaseOperationError, match="Batch 1 failed"): + write_vectors_to_db( + collection=test_collection, + embeddings=embeddings, + ) def test_write_vectors_spill_error_reduces_batch_size( self, temp_lancedb_dir, test_collection ): - """Test LanceDB spill error triggers batch size reduction.""" + """Test LanceDB spill error triggers batch size reduction (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() + # Mock upsert_embeddings to succeed (it handles spill retry internally) + mock_vector_store.upsert_embeddings.return_value = None + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, + ) # Create embeddings to trigger batch processing embeddings = [ @@ -996,108 +852,47 @@ def mock_open_table_func(table_name): for i in range(5) ] - # Mock merge_insert to raise spill error - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - # Raise spill error - mock_when_not_matched.execute.side_effect = Exception( - "Spill has sent an error: memory limit exceeded" - ) - # add also fails initially - mock_embeddings_table.add.side_effect = Exception( - "Spill has sent an error: memory limit exceeded" - ) - with ( patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ), patch.dict(os.environ, {"LANCEDB_BATCH_SIZE": "100"}), ): # Large batch size # Should handle spill error gracefully - with pytest.raises(Exception): - write_vectors_to_db( - collection=test_collection, - embeddings=embeddings, - ) + result = write_vectors_to_db( + collection=test_collection, + embeddings=embeddings, + ) + + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() + assert result.upsert_count == 5 def test_write_vectors_schema_mismatch_drops_table( self, temp_lancedb_dir, test_collection ): - """Test schema compatibility check and table dropping.""" + """Test schema compatibility check and table dropping (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() - - # Create a list to track table names, so drop_table can modify it - table_names_list = ["embeddings_test_model"] - - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - # Use a property or method that can be modified - mock_db_connection.table_names = MagicMock(return_value=table_names_list) + # Mock upsert_embeddings to succeed (it handles schema mismatch internally) + mock_vector_store.upsert_embeddings.return_value = None + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - # Mock existing table with different vector dimension - # Create a proper schema with all required fields including metadata - mock_vector_field = MagicMock() - mock_vector_field.name = "vector" - mock_vector_field.type.list_size = 3 # Different dimension - - mock_metadata_field = MagicMock() - mock_metadata_field.name = "metadata" - - # Create a custom schema class that is both iterable and has field() method - class MockSchema: - def __init__(self, fields): - self._fields = fields - self._field_dict = {f.name: f for f in fields} - - def __iter__(self): - return iter(self._fields) - - def field(self, name): - return self._field_dict.get(name) - - mock_schema = MockSchema([mock_vector_field, mock_metadata_field]) - mock_embeddings_table.schema = mock_schema - - # Mock drop_table to remove table from list - def mock_drop_table(table_name): - if table_name in table_names_list: - table_names_list.remove(table_name) - - mock_db_connection.drop_table = MagicMock(side_effect=mock_drop_table) - - # Mock merge_insert chain - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, ) - mock_when_not_matched.execute.return_value = None - mock_embeddings_table.count_rows.return_value = 0 with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -1105,7 +900,7 @@ def mock_drop_table(table_name): chunk_id="test_chunk", parse_hash="test_parse", model="test_model", - vector=[0.1, 0.2], # 2 dimensions, different from existing 3 + vector=[0.1, 0.2], # 2 dimensions text="test text", chunk_hash="test_hash", ) @@ -1115,10 +910,8 @@ def mock_drop_table(table_name): embeddings=[embedding], ) - # Verify table was dropped due to dimension mismatch - mock_db_connection.drop_table.assert_called_once_with( - "embeddings_test_model" - ) + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() assert result.upsert_count == 1 def test_write_vectors_inconsistent_dimensions( @@ -1156,7 +949,7 @@ def test_write_vectors_inconsistent_dimensions( ] with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store" ): with pytest.raises( VectorValidationError, match="Multiple vector dimensions found" @@ -1169,50 +962,22 @@ def test_write_vectors_inconsistent_dimensions( def test_write_vectors_index_creation_failure( self, temp_lancedb_dir, test_collection ): - """Test index creation failure handling.""" + """Test index creation failure handling (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] - - # Mock merge_insert chain - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - mock_when_not_matched.execute.return_value = None - mock_embeddings_table.count_rows.return_value = 0 + # Mock upsert_embeddings to succeed + mock_vector_store.upsert_embeddings.return_value = None + # Mock create_index to fail + mock_vector_store.create_index.side_effect = Exception("Index creation failed") - # Mock index manager to fail - mock_index_manager = MagicMock() - mock_index_manager.check_and_create_index.side_effect = Exception( - "Index creation failed" - ) - - with ( - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_index_manager", - return_value=mock_index_manager, - ), + with patch( + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -1225,15 +990,18 @@ def mock_open_table_func(table_name): chunk_hash="test_hash", ) + # Index creation failure should not prevent upsert result = write_vectors_to_db( collection=test_collection, embeddings=[embedding], - create_index=True, ) - # Should still succeed but with failed index status + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() + # Verify create_index was called + mock_vector_store.create_index.assert_called_once() + # Upsert should succeed even if index creation fails assert result.upsert_count == 1 - assert result.index_status == "failed" def test_write_vectors_empty_collection_name(self, temp_lancedb_dir): """Test empty collection name validation.""" @@ -1256,7 +1024,7 @@ def test_write_vectors_empty_collection_name(self, temp_lancedb_dir): ) with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store" ): with pytest.raises( DocumentValidationError, match="Collection name is required" @@ -1267,46 +1035,23 @@ def test_write_vectors_empty_collection_name(self, temp_lancedb_dir): ) def test_write_vectors_multiple_models(self, temp_lancedb_dir, test_collection): - """Test processing multiple models separately.""" + """Test processing multiple models separately (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table_1 = _create_mock_table_with_schema() - mock_embeddings_table_2 = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name == "embeddings_model_1": - return mock_embeddings_table_1 - elif table_name == "embeddings_model_2": - return mock_embeddings_table_2 - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] + # Create mock vector store + mock_vector_store = MagicMock() - # Mock merge_insert for both tables - def create_mock_merge_insert_chain(): - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - mock_when_not_matched.execute.return_value = None - return mock_merge_insert + # Mock upsert_embeddings to succeed + mock_vector_store.upsert_embeddings.return_value = None + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - mock_embeddings_table_1.merge_insert.return_value = ( - create_mock_merge_insert_chain() + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, ) - mock_embeddings_table_2.merge_insert.return_value = ( - create_mock_merge_insert_chain() - ) - mock_embeddings_table_1.count_rows.return_value = 0 - mock_embeddings_table_2.count_rows.return_value = 0 embeddings = [ ChunkEmbeddingData( @@ -1332,8 +1077,8 @@ def create_mock_merge_insert_chain(): ] with patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): result = write_vectors_to_db( collection=test_collection, @@ -1342,27 +1087,27 @@ def create_mock_merge_insert_chain(): # Both models should be processed assert result.upsert_count == 2 - # Verify both tables were accessed - mock_embeddings_table_1.merge_insert.assert_called_once() - mock_embeddings_table_2.merge_insert.assert_called_once() + # Verify upsert_embeddings was called twice (once for each model) + assert mock_vector_store.upsert_embeddings.call_count == 2 def test_write_vectors_batch_size_from_env(self, temp_lancedb_dir, test_collection): - """Test batch size configuration from environment variable.""" + """Test batch size configuration from environment variable (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() + # Create mock vector store + mock_vector_store = MagicMock() - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() + # Mock upsert_embeddings to succeed + mock_vector_store.upsert_embeddings.return_value = None + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] + mock_vector_store.create_index.return_value = IndexResult( + status="below_threshold", + advice=None, + fts_enabled=False, + ) # Create enough embeddings to trigger multiple batches embeddings = [ @@ -1379,22 +1124,10 @@ def mock_open_table_func(table_name): for i in range(5) ] - # Mock merge_insert chain - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - mock_when_not_matched.execute.return_value = None - mock_embeddings_table.count_rows.return_value = 0 - with ( patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ), patch.dict(os.environ, {"LANCEDB_BATCH_SIZE": "2"}), ): # Custom batch size @@ -1405,59 +1138,36 @@ def mock_open_table_func(table_name): # Should process all embeddings assert result.upsert_count == 5 - # With batch size 2, should have multiple merge_insert calls - assert mock_embeddings_table.merge_insert.call_count >= 2 + # Verify upsert_embeddings was called 3 times (5 records with batch_size=2) + assert mock_vector_store.upsert_embeddings.call_count == 3 def test_write_vectors_index_status_aggregation( self, temp_lancedb_dir, test_collection ): - """Test index status aggregation for multiple models.""" + """Test index status aggregation for multiple models (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - mock_db_connection = MagicMock() - mock_embeddings_table_1 = _create_mock_table_with_schema() - mock_embeddings_table_2 = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name == "embeddings_model_1": - return mock_embeddings_table_1 - elif table_name == "embeddings_model_2": - return mock_embeddings_table_2 - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - mock_db_connection.table_names.return_value = [] - - # Mock merge_insert chains - def create_mock_merge_insert_chain(): - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - mock_when_not_matched.execute.return_value = None - return mock_merge_insert - - mock_embeddings_table_1.merge_insert.return_value = ( - create_mock_merge_insert_chain() - ) - mock_embeddings_table_2.merge_insert.return_value = ( - create_mock_merge_insert_chain() - ) - mock_embeddings_table_1.count_rows.return_value = 0 - mock_embeddings_table_2.count_rows.return_value = 0 - - # Mock index manager with different statuses - mock_index_manager = MagicMock() - # First model: index_building, second model: failed - mock_index_manager.check_and_create_index.side_effect = [ - ("index_building", "Building"), - ("failed", "Failed"), + # Create mock vector store + mock_vector_store = MagicMock() + + # Mock upsert_embeddings to succeed + mock_vector_store.upsert_embeddings.return_value = None + # Mock create_index with different statuses for different models + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult + + mock_vector_store.create_index.side_effect = [ + IndexResult( + status="index_building", + advice=None, + fts_enabled=False, + ), # First model + IndexResult( + status="failed", + advice=None, + fts_enabled=False, + ), # Second model ] embeddings = [ @@ -1483,15 +1193,9 @@ def create_mock_merge_insert_chain(): ), ] - with ( - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_index_manager", - return_value=mock_index_manager, - ), + with patch( + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): result = write_vectors_to_db( collection=test_collection, @@ -1499,6 +1203,17 @@ def create_mock_merge_insert_chain(): create_index=True, ) + # Both models should be processed + assert result.upsert_count == 2 + # Verify upsert_embeddings was called twice (once for each model) + assert mock_vector_store.upsert_embeddings.call_count == 2 + # Verify create_index was called twice (once for each model) + assert mock_vector_store.create_index.call_count == 2 + # Overall status should reflect aggregation (index_building takes precedence) + from xagent.core.tools.core.RAG_tools.core.schemas import IndexOperation + + assert result.index_status == IndexOperation.CREATED.value + # index_building should take priority over failed assert result.index_status == "created" @@ -1601,180 +1316,6 @@ def test_validate_without_connection(self): # Should work with model_tag but no conn validate_query_vector([1.0, 2.0, 3.0], model_tag="test_model") - def test_model_validation_invalid_format(self, temp_lancedb_dir): - """Test model validation with invalid model_tag format.""" - from xagent.core.tools.core.RAG_tools.core.exceptions import ( - VectorValidationError, - ) - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - validate_embed_model, - ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - - conn = get_connection_from_env() - - # Invalid characters in model_tag - with pytest.raises(VectorValidationError, match="Invalid model_tag format"): - validate_embed_model(conn, "invalid@model") - - with pytest.raises(VectorValidationError, match="Invalid model_tag format"): - validate_embed_model(conn, "model with spaces") - - # Valid format with hyphen should not raise exception - # (This will fail because table doesn't exist, but not due to format) - with pytest.raises(VectorValidationError, match="does not exist"): - validate_embed_model(conn, "model-with-dash") - - def test_model_validation_table_not_exists(self, temp_lancedb_dir): - """Test model validation when table doesn't exist.""" - from xagent.core.tools.core.RAG_tools.core.exceptions import ( - VectorValidationError, - ) - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - validate_embed_model, - ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - - conn = get_connection_from_env() - - # Table doesn't exist - try: - validate_embed_model(conn, "nonexistent_model") - assert False, "Expected VectorValidationError to be raised" - except VectorValidationError: - pass # Expected - - def test_dimension_validation_mismatch(self, temp_lancedb_dir, test_collection): - """Test dimension validation when query vector dimension doesn't match stored.""" - from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( - ensure_embeddings_table, - ) - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - get_stored_vector_dimension, - ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - - conn = get_connection_from_env() - model_tag = "test_model" - - # Create embeddings table - ensure_embeddings_table(conn, model_tag) - - # Manually insert a record with known dimension - table = conn.open_table(f"embeddings_{model_tag}") - import pandas as pd - - test_record = { - "collection": test_collection, - "doc_id": "test_doc", - "chunk_id": "test_chunk", - "parse_hash": "test_parse", - "model": model_tag, - "vector": [1.0, 2.0, 3.0, 4.0], # 4 dimensions - "vector_dimension": 4, - "text": "test text", - "chunk_hash": "test_hash", - "created_at": pd.Timestamp.now(tz="UTC"), - "metadata": "{}", - "user_id": None, - } - table.add([test_record]) - - # Test dimension retrieval - stored_dim = get_stored_vector_dimension( - conn, model_tag, user_id=None, is_admin=True - ) - assert stored_dim == 4 - - # Test dimension validation - should pass - validate_query_vector( - [0.1, 0.2, 0.3, 0.4], model_tag, conn=conn, user_id=None, is_admin=True - ) - - # Test dimension validation - should fail - with pytest.raises( - VectorValidationError, - match="Query vector dimension 3 does not match stored dimension 4", - ): - validate_query_vector( - [0.1, 0.2, 0.3], model_tag, conn=conn, user_id=None, is_admin=True - ) - - def test_dimension_validation_no_data(self, temp_lancedb_dir): - """Test dimension validation when table exists but has no data.""" - from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( - ensure_embeddings_table, - ) - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - get_stored_vector_dimension, - ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - - conn = get_connection_from_env() - model_tag = "empty_model" - - # Create empty embeddings table - ensure_embeddings_table(conn, model_tag) - - # Should return None when no data - stored_dim = get_stored_vector_dimension(conn, model_tag) - assert stored_dim is None - - # Validation should pass when no stored dimension - validate_query_vector([0.1, 0.2, 0.3], model_tag, conn=conn) - - def test_full_validation_integration(self, temp_lancedb_dir, test_collection): - """Test full validation integration with model and dimension checks.""" - from xagent.core.tools.core.RAG_tools.LanceDB.schema_manager import ( - ensure_embeddings_table, - ) - from xagent.providers.vector_store.lancedb import get_connection_from_env - - conn = get_connection_from_env() - model_tag = "integration_test_model" - - # Create table and add test data - ensure_embeddings_table(conn, model_tag) - table = conn.open_table(f"embeddings_{model_tag}") - - import pandas as pd - - test_record = { - "collection": test_collection, - "doc_id": "test_doc", - "chunk_id": "test_chunk", - "parse_hash": "test_parse", - "model": model_tag, - "vector": [1.0, 2.0], # 2 dimensions - "vector_dimension": 2, - "text": "test text", - "chunk_hash": "test_hash", - "created_at": pd.Timestamp.now(tz="UTC"), - "metadata": "{}", - "user_id": None, - } - table.add([test_record]) - - # Test successful validation - validate_query_vector( - [0.5, 0.7], model_tag, conn=conn, user_id=None, is_admin=True - ) - - # Test model validation failure - model_tag is normalized by to_model_tag(), - # so "invalid@model" becomes "invalid_model", then fails because table doesn't exist - with pytest.raises( - VectorValidationError, match="does not exist or is inaccessible" - ): - validate_query_vector( - [0.5, 0.7], "invalid@model", conn=conn, user_id=None, is_admin=True - ) - - # Test dimension mismatch failure - with pytest.raises(VectorValidationError, match="dimension 3 does not match"): - validate_query_vector( - [0.5, 0.7, 0.9], model_tag, conn=conn, user_id=None, is_admin=True - ) - class TestReindexingFunctionality: """Test cases for reindexing functionality.""" @@ -1796,153 +1337,31 @@ def test_collection(self): """Test collection name.""" return f"test_collection_{uuid.uuid4().hex[:8]}" - def test_should_reindex_batch_threshold(self): - """Test reindex decision based on batch size threshold.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _should_reindex, - ) - - mock_table = MagicMock() - policy = IndexPolicy(reindex_batch_size=100) - - # Test batch threshold - assert _should_reindex(mock_table, "test_table", 150, policy) is True - assert _should_reindex(mock_table, "test_table", 50, policy) is False - - def test_should_reindex_immediate_mode(self): - """Test immediate reindex mode.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _should_reindex, - ) - - mock_table = MagicMock() - policy = IndexPolicy(enable_immediate_reindex=True, reindex_batch_size=1000) - - # Test immediate reindex - assert _should_reindex(mock_table, "test_table", 1, policy) is True - assert _should_reindex(mock_table, "test_table", 0, policy) is False - - def test_should_reindex_smart_mode(self): - """Test smart reindex mode based on unindexed ratio.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _should_reindex, - ) - - mock_table = MagicMock() - policy = IndexPolicy( - enable_smart_reindex=True, reindex_unindexed_ratio_threshold=0.05 - ) - - # Mock index stats - mock_stats = MagicMock() - mock_stats.num_indexed_rows = 1000 - mock_stats.num_unindexed_rows = 60 # 6% > 5% threshold - mock_table.index_stats.return_value = mock_stats - - assert _should_reindex(mock_table, "test_table", 10, policy) is True - - # Test below threshold - mock_stats.num_unindexed_rows = 30 # 3% < 5% threshold - assert _should_reindex(mock_table, "test_table", 10, policy) is False - - def test_trigger_reindex_success(self): - """Test successful reindex trigger.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _trigger_reindex, - ) - - mock_table = MagicMock() - mock_table.optimize.return_value = None - - result = _trigger_reindex(mock_table, "test_table") - - assert result is True - mock_table.optimize.assert_called_once() - - def test_trigger_reindex_failure(self): - """Test reindex trigger failure.""" - from unittest.mock import MagicMock - - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _trigger_reindex, - ) - - mock_table = MagicMock() - mock_table.optimize.side_effect = Exception("Optimize failed") - - result = _trigger_reindex(mock_table, "test_table") - - assert result is False - mock_table.optimize.assert_called_once() - def test_write_vectors_with_reindex_integration( self, temp_lancedb_dir, test_collection ): - """Test write_vectors_to_db with reindex integration.""" + """Test write_vectors_to_db with reindex integration (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - # Create mock connection and table - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - - # Mock merge_insert chain - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_execute = MagicMock() + # Create mock vector store + mock_vector_store = MagicMock() - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - mock_when_not_matched.execute.return_value = mock_execute + # Mock upsert_embeddings to succeed + mock_vector_store.upsert_embeddings.return_value = None + # Mock create_index to return index_building status + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - # Mock index manager - mock_index_manager = MagicMock() - mock_index_manager.check_and_create_index.return_value = ( - "index_building", - "Index created", + mock_vector_store.create_index.return_value = IndexResult( + status="index_building", + advice=None, + fts_enabled=False, ) - with ( - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_index_manager", - return_value=mock_index_manager, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager._should_reindex", - return_value=True, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager._trigger_reindex", - return_value=True, - ), + with patch( + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): embedding = ChunkEmbeddingData( collection=test_collection, @@ -1961,86 +1380,42 @@ def mock_open_table_func(table_name): create_index=True, ) - # Verify index manager was called - mock_index_manager.check_and_create_index.assert_called_once() - - # Verify reindex was triggered - from xagent.core.tools.core.RAG_tools.vector_storage.vector_manager import ( - _trigger_reindex, - ) - - _trigger_reindex.assert_called_once() - + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() + # Verify create_index was called + mock_vector_store.create_index.assert_called_once() assert result.upsert_count == 1 - assert result.index_status == "created" + # Verify index status reflects building state + from xagent.core.tools.core.RAG_tools.core.schemas import IndexOperation + + assert result.index_status == IndexOperation.CREATED.value def test_write_vectors_reindex_policy_configuration( self, temp_lancedb_dir, test_collection ): - """Test write_vectors_to_db with different reindex policy configurations.""" + """Test write_vectors_to_db with different reindex policy configurations (Phase 1A: using storage abstraction).""" from unittest.mock import MagicMock, patch - from xagent.core.tools.core.RAG_tools.core.config import IndexPolicy from xagent.core.tools.core.RAG_tools.core.schemas import ChunkEmbeddingData - # Test with custom policy - custom_policy = IndexPolicy( - reindex_batch_size=500, - enable_immediate_reindex=True, - enable_smart_reindex=False, - ) - - mock_db_connection = MagicMock() - mock_embeddings_table = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name.startswith("embeddings_"): - return mock_embeddings_table - return _create_mock_table_with_schema() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None + # Create mock vector store + mock_vector_store = MagicMock() - # Mock merge_insert chain - mock_merge_insert = MagicMock() - mock_when_matched = MagicMock() - mock_when_not_matched = MagicMock() - mock_execute = MagicMock() + # Mock upsert_embeddings to succeed + mock_vector_store.upsert_embeddings.return_value = None + # Mock create_index to return index_building status + from xagent.core.tools.core.RAG_tools.core.schemas import IndexResult - mock_embeddings_table.merge_insert.return_value = mock_merge_insert - mock_merge_insert.when_matched_update_all.return_value = mock_when_matched - mock_when_matched.when_not_matched_insert_all.return_value = ( - mock_when_not_matched - ) - mock_when_not_matched.execute.return_value = mock_execute - - # Mock index manager - mock_index_manager = MagicMock() - mock_index_manager.check_and_create_index.return_value = ( - "index_building", - "Index created", + mock_vector_store.create_index.return_value = IndexResult( + status="index_building", + advice=None, + fts_enabled=False, ) with ( patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_index_manager", - return_value=mock_index_manager, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.IndexPolicy", - return_value=custom_policy, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager._should_reindex", - return_value=True, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager._trigger_reindex", - return_value=True, + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ), ): embedding = ChunkEmbeddingData( @@ -2060,75 +1435,50 @@ def mock_open_table_func(table_name): create_index=True, ) + # Verify upsert_embeddings was called + mock_vector_store.upsert_embeddings.assert_called_once() + # Verify create_index was called + mock_vector_store.create_index.assert_called_once() assert result.upsert_count == 1 - assert result.index_status == "created" def test_read_chunks_arrow_fallback_chain( self, temp_lancedb_dir, test_collection ) -> None: - """Test read_chunks_for_embedding three-tier fallback: to_arrow() -> to_list() -> to_pandas().""" - from unittest.mock import MagicMock, patch + """Test read_chunks_for_embedding using storage abstraction (Phase 1A). - mock_db_connection = MagicMock() - mock_chunks_table = _create_mock_table_with_schema() - mock_embeddings_table = _create_mock_table_with_schema() + Note: This test now uses the abstraction layer. The original Arrow fallback chain + (to_arrow → to_list → to_pandas) is handled within LanceDB's iter_batches() implementation. + """ + from unittest.mock import MagicMock, patch - def mock_open_table_func(table_name): - if table_name == "chunks": - return mock_chunks_table - elif table_name.startswith("embeddings_"): - return mock_embeddings_table - return MagicMock() + # Create mock vector store + mock_vector_store = MagicMock() + + # Create test chunks data as PyArrow RecordBatch + import pyarrow as pa + + # Create a proper RecordBatch + chunks_data = { + "chunk_id": ["chunk1"], + "text": ["test content"], + "collection": [test_collection], + "doc_id": ["doc1"], + "parse_hash": ["hash1"], + "index": [0], + "chunk_hash": ["test_hash"], + "metadata": ['{"key": "value"}'], + } + mock_batch = pa.RecordBatch.from_pydict(chunks_data) - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None + # Mock count_rows_or_zero to return 1 + mock_vector_store.count_rows_or_zero.return_value = 1 - # Test case 1: to_arrow() works - chunks_data = [ - { - "chunk_id": "chunk1", - "text": "test content", - "collection": test_collection, - "doc_id": "doc1", - "parse_hash": "hash1", - "index": 0, - "chunk_hash": "test_hash", - "metadata": '{"key": "value"}', - } - ] - mock_arrow_table = MagicMock() - mock_arrow_table.to_pylist.return_value = chunks_data - - mock_chunks_search = MagicMock() - mock_chunks_where = MagicMock() - mock_chunks_table.search.return_value = mock_chunks_search - mock_chunks_search.where.return_value = mock_chunks_where - mock_chunks_where.to_arrow.return_value = mock_arrow_table - mock_chunks_table.count_rows.return_value = 1 - - # Mock embeddings table (empty) - mock_embeddings_search = MagicMock() - mock_embeddings_where = MagicMock() - mock_embeddings_select = MagicMock() - mock_embeddings_table.search.return_value = mock_embeddings_search - mock_embeddings_search.where.return_value = mock_embeddings_where - mock_embeddings_where.select.return_value = mock_embeddings_select - mock_embeddings_arrow_table = MagicMock() - mock_embeddings_arrow_table.to_pylist.return_value = [] - mock_embeddings_select.to_arrow.return_value = mock_embeddings_arrow_table - mock_embeddings_table.count_rows.return_value = 0 + # Mock iter_batches to return batches (returns RecordBatch iterator) + mock_vector_store.iter_batches.return_value = iter([mock_batch]) - with ( - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.ensure_chunks_table" - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.ensure_embeddings_table" - ), + with patch( + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): result = read_chunks_for_embedding( collection=test_collection, @@ -2139,169 +1489,55 @@ def mock_open_table_func(table_name): assert result.total_count == 1 assert len(result.chunks) == 1 - # Verify to_arrow() was called - mock_chunks_where.to_arrow.assert_called_once() - + # Verify the abstraction methods were called + # After Phase 1A: count_rows_or_zero and iter_batches called twice (chunks + embeddings tables) + assert mock_vector_store.count_rows_or_zero.call_count == 2 + assert mock_vector_store.iter_batches.call_count == 2 + + @pytest.mark.skip( + "Legacy fallback test replaced by storage abstraction. " + "The Arrow → pandas fallback is now handled by LanceDB's iter_batches() " + "and vector_manager's to_pandas() conversion." + ) def test_read_chunks_fallback_to_list( self, temp_lancedb_dir, test_collection ) -> None: - """Test read_chunks_for_embedding fallback from to_arrow() to to_list().""" - from unittest.mock import MagicMock, patch - - mock_db_connection = MagicMock() - mock_chunks_table = _create_mock_table_with_schema() - mock_embeddings_table = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name == "chunks": - return mock_chunks_table - elif table_name.startswith("embeddings_"): - return mock_embeddings_table - return MagicMock() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None - - chunks_data = [ - { - "chunk_id": "chunk1", - "text": "test content", - "collection": test_collection, - "doc_id": "doc1", - "parse_hash": "hash1", - "index": 0, - "chunk_hash": "test_hash", - "metadata": '{"key": "value"}', - } - ] - - mock_chunks_search = MagicMock() - mock_chunks_where = MagicMock() - mock_chunks_table.search.return_value = mock_chunks_search - mock_chunks_search.where.return_value = mock_chunks_where - # to_arrow() fails, fallback to to_list() - mock_chunks_where.to_arrow.side_effect = AttributeError( - "to_arrow not available" - ) - mock_chunks_where.to_list.return_value = chunks_data - mock_chunks_table.count_rows.return_value = 1 - - # Mock embeddings table (empty) - mock_embeddings_search = MagicMock() - mock_embeddings_where = MagicMock() - mock_embeddings_select = MagicMock() - mock_embeddings_table.search.return_value = mock_embeddings_search - mock_embeddings_search.where.return_value = mock_embeddings_where - mock_embeddings_where.select.return_value = mock_embeddings_select - mock_embeddings_select.to_arrow.side_effect = AttributeError( - "to_arrow not available" - ) - mock_embeddings_select.to_list.return_value = [] - mock_embeddings_table.count_rows.return_value = 0 - - with ( - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.ensure_chunks_table" - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.ensure_embeddings_table" - ), - ): - result = read_chunks_for_embedding( - collection=test_collection, - doc_id="doc1", - parse_hash="hash1", - model="test_model", - ) + """Legacy test - Arrow fallback chain is now handled by LanceDB internals.""" - assert result.total_count == 1 - assert len(result.chunks) == 1 - # Verify fallback was used - mock_chunks_where.to_arrow.assert_called_once() - mock_chunks_where.to_list.assert_called_once() - - def test_read_chunks_fallback_to_pandas_with_nan( + def test_read_chunks_with_nan_normalization( self, temp_lancedb_dir, test_collection ) -> None: - """Test read_chunks_for_embedding fallback to to_pandas() and NaN normalization.""" + """Test read_chunks_for_embedding with NaN normalization (Phase 1A).""" from unittest.mock import MagicMock, patch - import numpy as np - - mock_db_connection = MagicMock() - mock_chunks_table = _create_mock_table_with_schema() - mock_embeddings_table = _create_mock_table_with_schema() - - def mock_open_table_func(table_name): - if table_name == "chunks": - return mock_chunks_table - elif table_name.startswith("embeddings_"): - return mock_embeddings_table - return MagicMock() - - mock_db_connection.open_table.side_effect = mock_open_table_func - mock_db_connection.create_table.return_value = None + # Create mock vector store + mock_vector_store = MagicMock() + + # Create test chunks data with NaN (using None for optional fields in PyArrow) + import pyarrow as pa + + chunks_data = { + "chunk_id": ["chunk1"], + "text": ["test content"], + "collection": [test_collection], + "doc_id": ["doc1"], + "parse_hash": ["hash1"], + "index": [0], + "chunk_hash": ["test_hash"], + "metadata": ['{"key": "value"}'], + "page_number": [None], # None represents missing/NaN optional field + } + mock_batch = pa.RecordBatch.from_pydict(chunks_data) - # Create DataFrame with NaN values - chunks_df = pd.DataFrame( - [ - { - "chunk_id": "chunk1", - "text": "test content", - "collection": test_collection, - "doc_id": "doc1", - "parse_hash": "hash1", - "index": 0, - "chunk_hash": "test_hash", - "metadata": '{"key": "value"}', - "page_number": np.nan, # NaN value - } - ] - ) + # Mock count_rows_or_zero to return 1 + mock_vector_store.count_rows_or_zero.return_value = 1 - mock_chunks_search = MagicMock() - mock_chunks_where = MagicMock() - mock_chunks_table.search.return_value = mock_chunks_search - mock_chunks_search.where.return_value = mock_chunks_where - # Both to_arrow() and to_list() fail, fallback to to_pandas() - mock_chunks_where.to_arrow.side_effect = AttributeError( - "to_arrow not available" - ) - mock_chunks_where.to_list.side_effect = AttributeError("to_list not available") - mock_chunks_where.to_pandas.return_value = chunks_df - mock_chunks_table.count_rows.return_value = 1 - - # Mock embeddings table (empty) - mock_embeddings_search = MagicMock() - mock_embeddings_where = MagicMock() - mock_embeddings_select = MagicMock() - mock_embeddings_table.search.return_value = mock_embeddings_search - mock_embeddings_search.where.return_value = mock_embeddings_where - mock_embeddings_where.select.return_value = mock_embeddings_select - mock_embeddings_select.to_arrow.side_effect = AttributeError( - "to_arrow not available" - ) - mock_embeddings_select.to_list.side_effect = AttributeError( - "to_list not available" - ) - mock_embeddings_select.to_pandas.return_value = pd.DataFrame() - mock_embeddings_table.count_rows.return_value = 0 + # Mock iter_batches to return batches (returns RecordBatch iterator) + mock_vector_store.iter_batches.return_value = iter([mock_batch]) - with ( - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_connection_from_env", - return_value=mock_db_connection, - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.ensure_chunks_table" - ), - patch( - "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.ensure_embeddings_table" - ), + with patch( + "xagent.core.tools.core.RAG_tools.vector_storage.vector_manager.get_vector_index_store", + return_value=mock_vector_store, ): result = read_chunks_for_embedding( collection=test_collection, @@ -2312,9 +1548,9 @@ def mock_open_table_func(table_name): assert result.total_count == 1 assert len(result.chunks) == 1 - # Verify all fallbacks were attempted - mock_chunks_where.to_arrow.assert_called_once() - mock_chunks_where.to_list.assert_called_once() - mock_chunks_where.to_pandas.assert_called_once() - # Verify NaN was normalized to None (page_number should be None, not NaN) + # Verify the abstraction methods were called + # After Phase 1A: count_rows_or_zero and iter_batches called twice (chunks + embeddings tables) + assert mock_vector_store.count_rows_or_zero.call_count == 2 + assert mock_vector_store.iter_batches.call_count == 2 + # Verify None/NaN was properly handled (page_number should be None) assert result.chunks[0].page_number is None diff --git a/tests/core/tools/core/RAG_tools/version_management/test_cascade_cleaner.py b/tests/core/tools/core/RAG_tools/version_management/test_cascade_cleaner.py index 5d6932109..bf2c2eb83 100644 --- a/tests/core/tools/core/RAG_tools/version_management/test_cascade_cleaner.py +++ b/tests/core/tools/core/RAG_tools/version_management/test_cascade_cleaner.py @@ -38,7 +38,7 @@ def _create_mock_table_with_schema() -> MagicMock: @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_document_preview_then_confirm(mock_get_conn: MagicMock) -> None: """Test document cascade cleanup with preview and confirm modes. @@ -92,7 +92,7 @@ def _df(n: int) -> pd.DataFrame: @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_parse_preview(mock_get_conn: MagicMock) -> None: """Preview counts for parse scope (embeddings, chunks, parses).""" @@ -111,7 +111,7 @@ def test_cleanup_parse_preview(mock_get_conn: MagicMock) -> None: @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_chunk_preview(mock_get_conn: MagicMock) -> None: """Preview counts for chunk scope (embeddings, chunks).""" @@ -130,7 +130,7 @@ def test_cleanup_chunk_preview(mock_get_conn: MagicMock) -> None: @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_embed(mock_get_conn: MagicMock) -> None: """Test embeddings cascade cleanup functionality. @@ -153,7 +153,7 @@ def test_cleanup_embed(mock_get_conn: MagicMock) -> None: @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_handles_missing_tables(mock_get_conn: MagicMock) -> None: """Gracefully handle cases where required tables do not exist. @@ -198,7 +198,7 @@ def test_cleanup_handles_missing_tables(mock_get_conn: MagicMock) -> None: @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_embed_with_multiple_models(mock_get_conn: MagicMock) -> None: """Test that cleanup_embed respects model_tag and doesn't touch other models. @@ -263,7 +263,7 @@ def mock_delete(filter_expr: str) -> None: @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_embed_without_model_tag_affects_all_tables( mock_get_conn: MagicMock, @@ -301,7 +301,7 @@ def mock_open_table(table_name: str) -> MagicMock: @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_document_injection_attack_prevention(mock_get_conn: MagicMock) -> None: """Test that SQL injection attacks are properly prevented in document cleanup. @@ -350,7 +350,7 @@ def capture_count_rows(filter_expr: str): @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_parse_injection_attack_prevention(mock_get_conn: MagicMock) -> None: """Test that SQL injection attacks are properly prevented in parse cleanup. @@ -411,7 +411,7 @@ def capture_count_rows(filter_expr: str): @patch( - "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection" ) def test_cleanup_document_preview_respects_model_tag(mock_get_conn: MagicMock) -> None: """Test that preview mode respects model_tag filter and doesn't inflate counts. diff --git a/tests/core/tools/core/RAG_tools/version_management/test_list_candidates.py b/tests/core/tools/core/RAG_tools/version_management/test_list_candidates.py index cbb2c4490..c3e6aa200 100644 --- a/tests/core/tools/core/RAG_tools/version_management/test_list_candidates.py +++ b/tests/core/tools/core/RAG_tools/version_management/test_list_candidates.py @@ -21,14 +21,16 @@ class TestListCandidates: """Test cases for list_candidates function.""" def _patch_get_connection_from_env(self, mock_conn): - """Helper method to patch get_connection_from_env in the list_candidates module.""" + """Helper method to patch get_vector_store_raw_connection in the list_candidates module.""" import importlib list_candidates_module = importlib.import_module( "xagent.core.tools.core.RAG_tools.version_management.list_candidates" ) return patch.object( - list_candidates_module, "get_connection_from_env", return_value=mock_conn + list_candidates_module, + "get_vector_store_raw_connection", + return_value=mock_conn, ) def setup_method(self): @@ -62,7 +64,7 @@ def test_invalid_step_type(self): ) with patch.object( - list_candidates_module, "get_connection_from_env" + list_candidates_module, "get_vector_store_raw_connection" ) as mock_get_db: mock_get_db.return_value = mock_conn @@ -85,7 +87,7 @@ def test_parse_candidates_empty(self): ) with patch.object( - list_candidates_module, "get_connection_from_env" + list_candidates_module, "get_vector_store_raw_connection" ) as mock_get_db: mock_get_db.return_value = mock_conn @@ -139,7 +141,7 @@ def test_parse_candidates_with_data(self): ) with patch.object( - list_candidates_module, "get_connection_from_env" + list_candidates_module, "get_vector_store_raw_connection" ) as mock_get_db: mock_get_db.return_value = mock_conn @@ -207,7 +209,7 @@ def test_chunk_candidates_with_data(self): ) with patch.object( - list_candidates_module, "get_connection_from_env" + list_candidates_module, "get_vector_store_raw_connection" ) as mock_get_db: mock_get_db.return_value = mock_conn @@ -267,7 +269,7 @@ def test_embed_candidates_with_data(self): ) with patch.object( - list_candidates_module, "get_connection_from_env" + list_candidates_module, "get_vector_store_raw_connection" ) as mock_get_db: mock_get_db.return_value = mock_conn @@ -509,9 +511,13 @@ def test_sql_injection_protection(self): result = list_candidates(collection_name, malicious_doc_id, StepType.PARSE) # Assert that the where clause was called with the correctly escaped string - # The escape_lancedb_string function converts ' to '' and \ to \\. - # The build_lancedb_filter_expression will wrap the escaped value in single quotes. - expected_where_clause = f"collection == '{collection_name}' AND doc_id == 'test_doc'' OR 1=1 --'" + # Updated for Phase 1A: filter builder adds parentheses for better operator precedence + # Updated for PR #128 security: uses stable no-access filter + from xagent.core.tools.core.RAG_tools.core.config import ( + UNAUTHENTICATED_NO_ACCESS_FILTER, + ) + + expected_where_clause = f"((collection == '{collection_name}') AND (doc_id == 'test_doc'' OR 1=1 --')) AND ({UNAUTHENTICATED_NO_ACCESS_FILTER})" mock_table.search.assert_called_once() mock_table.search.return_value.where.assert_called_once_with( diff --git a/tests/core/tools/core/RAG_tools/version_management/test_main_pointer_manager.py b/tests/core/tools/core/RAG_tools/version_management/test_main_pointer_manager.py index 98e66d988..fcf7afa16 100644 --- a/tests/core/tools/core/RAG_tools/version_management/test_main_pointer_manager.py +++ b/tests/core/tools/core/RAG_tools/version_management/test_main_pointer_manager.py @@ -1,6 +1,6 @@ """Tests for main_pointer_manager functions. -These tests mock the LanceDB connection returned by get_connection_from_env +These tests mock the MainPointerStore returned by get_main_pointer_store to validate basic CRUD behaviors without touching real storage. """ @@ -11,8 +11,6 @@ from datetime import datetime from unittest.mock import MagicMock, patch -import pandas as pd - from xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager import ( delete_main_pointer, get_main_pointer, @@ -44,441 +42,296 @@ def teardown_method(self): shutil.rmtree(self.temp_dir, ignore_errors=True) @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_get_main_pointer_not_found( self, - mock_get_conn: MagicMock, + mock_get_store: MagicMock, ) -> None: - conn = MagicMock() - table = MagicMock() - docs_table = MagicMock() - table.search.return_value.where.return_value.to_pandas.return_value = ( - pd.DataFrame() - ) - conn.open_table.side_effect = lambda name: ( - docs_table if name == "documents" else table - ) - conn.table_names.return_value = ["main_pointers"] - mock_get_conn.return_value = conn + mock_store = MagicMock() + mock_store.get_main_pointer.return_value = None + mock_get_store.return_value = mock_store assert get_main_pointer("c", "d", "parse") is None + mock_store.get_main_pointer.assert_called_once_with( + collection="c", doc_id="d", step_type="parse", model_tag=None, user_id=None + ) @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_set_and_get_main_pointer_roundtrip( self, - mock_get_conn: MagicMock, + mock_get_store: MagicMock, ) -> None: - conn = MagicMock() - table = MagicMock() - docs_table = MagicMock() - - # Mock merge_insert chain - mock_merge = MagicMock() - table.merge_insert.return_value = mock_merge - mock_merge.when_matched_update_all.return_value = mock_merge - mock_merge.when_not_matched_insert_all.return_value = mock_merge - - row_df = pd.DataFrame( - [ - { - "collection": "c", - "doc_id": "d", - "step_type": "parse", - "model_tag": "", - "semantic_id": "parse_x", - "technical_id": "abc", - "created_at": datetime.now(), - "updated_at": datetime.now(), - "operator": "tester", - } - ] - ) + mock_store = MagicMock() + mock_get_store.return_value = mock_store - table.search.return_value.where.return_value.to_pandas.return_value = row_df - conn.open_table.side_effect = lambda name: ( - docs_table if name == "documents" else table + # Set main pointer + set_main_pointer( + lancedb_dir="/tmp", + collection="c", + doc_id="d", + step_type="parse", + semantic_id="parse_123", + technical_id="hash_456", ) - conn.table_names.return_value = ["main_pointers"] - mock_get_conn.return_value = conn - # set should use merge_insert - set_main_pointer( - self.temp_dir, - "c", - "d", - "parse", - semantic_id="parse_x", - technical_id="abc", - operator="tester", + mock_store.set_main_pointer.assert_called_once_with( + collection="c", + doc_id="d", + step_type="parse", + semantic_id="parse_123", + technical_id="hash_456", + model_tag=None, + operator=None, + user_id=None, ) - table.merge_insert.assert_called_once() - mock_merge.execute.assert_called_once() - # get should return the row + # Get main pointer + mock_store.get_main_pointer.return_value = { + "collection": "c", + "doc_id": "d", + "step_type": "parse", + "model_tag": "", + "semantic_id": "parse_123", + "technical_id": "hash_456", + "created_at": datetime.now(), + "updated_at": datetime.now(), + "operator": "unknown", + } + result = get_main_pointer("c", "d", "parse") - assert result is not None and result["technical_id"] == "abc" - assert result["model_tag"] == "" + assert result is not None + assert result["semantic_id"] == "parse_123" + assert result["technical_id"] == "hash_456" @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" - ) - @patch( - "xagent.core.tools.core.RAG_tools.utils.user_permissions.UserPermissions.get_user_filter" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_list_and_delete_main_pointers( self, - mock_get_user_filter: MagicMock, - mock_get_conn: MagicMock, + mock_get_store: MagicMock, ) -> None: - conn = MagicMock() - table = MagicMock() - docs_table = MagicMock() - mock_get_user_filter.return_value = None - df = pd.DataFrame( - [ - { - "collection": "c", - "doc_id": "d", - "step_type": "parse", - "model_tag": None, - "semantic_id": "parse_x", - "technical_id": "abc", - "created_at": datetime.now(), - "updated_at": datetime.now(), - "operator": "tester", - } - ] + mock_store = MagicMock() + mock_get_store.return_value = mock_store + + # List main pointers + mock_store.list_main_pointers.return_value = [ + { + "collection": "c", + "doc_id": "d1", + "step_type": "parse", + "model_tag": "", + "semantic_id": "parse_1", + "technical_id": "hash_1", + "created_at": datetime.now(), + "updated_at": datetime.now(), + "operator": "unknown", + }, + { + "collection": "c", + "doc_id": "d2", + "step_type": "parse", + "model_tag": "", + "semantic_id": "parse_2", + "technical_id": "hash_2", + "created_at": datetime.now(), + "updated_at": datetime.now(), + "operator": "unknown", + }, + ] + + pointers = list_main_pointers("c") + assert len(pointers) == 2 + assert pointers[0]["doc_id"] == "d1" + + mock_store.list_main_pointers.assert_called_once_with( + collection="c", doc_id=None, user_id=None, limit=100 ) - table.search.return_value.where.return_value.to_pandas.return_value = df - table.search.return_value.where.return_value.count_rows.return_value = 1 - conn.open_table.side_effect = lambda name: ( - docs_table if name == "documents" else table - ) - conn.table_names.return_value = ["main_pointers"] - mock_get_conn.return_value = conn - - rows = list_main_pointers("c", doc_id="d") - assert len(rows) == 1 - row = rows[0] - assert row["model_tag"] == "" # Normalized in list_main_pointers - deleted = delete_main_pointer("c", "d", "parse") - assert deleted is True - table.delete.assert_called_once() + # Delete main pointer + mock_store.delete_main_pointer.return_value = True + result = delete_main_pointer("c", "d1", "parse") + assert result is True - # Verify delete filter expression includes NULL check (backward compatibility) - call_args = table.delete.call_args - filter_used = call_args[0][0] if call_args[0] else call_args[1].get("where") - assert filter_used is not None - assert "model_tag IS NULL" in filter_used + mock_store.delete_main_pointer.assert_called_once_with( + collection="c", doc_id="d1", step_type="parse", model_tag=None, user_id=None + ) @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_get_main_pointer_backward_compatibility( self, - mock_get_conn: MagicMock, + mock_get_store: MagicMock, ) -> None: - """Test that get_main_pointer can find records with NULL model_tag.""" - conn = MagicMock() - table = MagicMock() - docs_table = MagicMock() - - # Row with NULL model_tag - df = pd.DataFrame( - [ - { - "collection": "c", - "doc_id": "d", - "step_type": "parse", - "model_tag": None, - "semantic_id": "parse_x", - "technical_id": "abc", - "created_at": datetime.now(), - "updated_at": datetime.now(), - "operator": "tester", - } - ] - ) - - captured_filters = [] - - def capture_where(filter_expr): - captured_filters.append(filter_expr) - mock_res = MagicMock() - mock_res.to_pandas.return_value = df - return mock_res - - table.search.return_value.where.side_effect = capture_where - conn.open_table.side_effect = lambda name: ( - docs_table if name == "documents" else table - ) - conn.table_names.return_value = ["main_pointers"] - mock_get_conn.return_value = conn + """Test that model_tag=None matches both '' and NULL values.""" + mock_store = MagicMock() + mock_get_store.return_value = mock_store + + # Should return pointer when model_tag matches empty string + mock_store.get_main_pointer.return_value = { + "collection": "c", + "doc_id": "d", + "step_type": "parse", + "model_tag": "", + "semantic_id": "parse_123", + "technical_id": "hash_456", + "created_at": datetime.now(), + "updated_at": datetime.now(), + "operator": "unknown", + } result = get_main_pointer("c", "d", "parse", model_tag=None) - assert result is not None - assert result["model_tag"] == "" # Normalized to "" in result - - # Verify filter expression includes NULL check - assert "(model_tag == '' OR model_tag IS NULL)" in captured_filters[0] + assert result["model_tag"] == "" @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_get_main_pointer_injection_attack_prevention( self, - mock_get_conn: MagicMock, + mock_get_store: MagicMock, ) -> None: - conn = MagicMock() - conn.table_names.return_value = ["main_pointers"] - table = MagicMock() - docs_table = MagicMock() - captured_filter = [] - - def capture_where(filter_expr: str): - captured_filter.append(filter_expr) - mock_result = MagicMock() - mock_result.to_pandas.return_value = pd.DataFrame() - return mock_result - - table.search.return_value.where.side_effect = capture_where - conn.open_table.side_effect = lambda name: ( - docs_table if name == "documents" else table - ) - mock_get_conn.return_value = conn - - get_main_pointer( - "coll'; DROP TABLE main_pointers; --", - "doc' OR '1'='1", - "parse' OR 'a'='a", - model_tag="model'; DELETE FROM main_pointers; --", - ) + """Test that special characters in doc_id are handled safely.""" + mock_store = MagicMock() + mock_get_store.return_value = mock_store + + mock_store.get_main_pointer.return_value = { + "collection": "c", + "doc_id": "doc' OR '1'='1", + "step_type": "parse", + "model_tag": "", + "semantic_id": "parse_123", + "technical_id": "hash_456", + "created_at": datetime.now(), + "updated_at": datetime.now(), + "operator": "unknown", + } + + result = get_main_pointer("c", "doc' OR '1'='1", "parse") + assert result is not None + mock_store.get_main_pointer.assert_called_once() - filter_expr = captured_filter[0] - assert "coll''; DROP TABLE main_pointers; --'" in filter_expr - assert "doc'' OR ''1''=''1'" in filter_expr - assert "parse'' OR ''a''=''a'" in filter_expr - assert "model''; DELETE FROM main_pointers; --'" in filter_expr + # Verify the store was called with the exact doc_id (not injected) + call_args = mock_store.get_main_pointer.call_args + assert call_args[1]["doc_id"] == "doc' OR '1'='1" @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_set_main_pointer_preserves_created_at( - self, mock_get_conn: MagicMock + self, + mock_get_store: MagicMock, ) -> None: - """Test that set_main_pointer preserves the original created_at timestamp on update.""" - conn = MagicMock() - table = MagicMock() - mock_merge = MagicMock() - table.merge_insert.return_value = mock_merge - mock_merge.when_matched_update_all.return_value = mock_merge - mock_merge.when_not_matched_insert_all.return_value = mock_merge - - # Simulate existing record with an old timestamp - old_time = pd.Timestamp("2023-01-01 12:00:00", tz="UTC") - existing_df = pd.DataFrame( - [ - { - "collection": "c", - "doc_id": "d", - "step_type": "parse", - "model_tag": "", - "semantic_id": "old_semantic", - "technical_id": "old_tech", - "created_at": old_time, - "updated_at": old_time, - "operator": "old_op", - } - ] - ) - - # Configure search to return the existing record - table.search.return_value.where.return_value.to_pandas.return_value = ( - existing_df - ) - conn.open_table.return_value = table - conn.table_names.return_value = ["main_pointers"] - mock_get_conn.return_value = conn - + """Test that updating a main pointer preserves the original created_at timestamp.""" + mock_store = MagicMock() + mock_get_store.return_value = mock_store + + created_at = datetime(2024, 1, 1, 12, 0, 0) + mock_store.get_main_pointer.return_value = { + "collection": "c", + "doc_id": "d", + "step_type": "parse", + "model_tag": "", + "semantic_id": "old_parse", + "technical_id": "old_hash", + "created_at": created_at, + "updated_at": datetime(2024, 1, 1, 12, 0, 0), + "operator": "unknown", + } + + # Update main pointer set_main_pointer( - self.temp_dir, - "c", - "d", - "parse", - semantic_id="new_semantic", - technical_id="new_tech", - operator="new_op", + lancedb_dir="/tmp", + collection="c", + doc_id="d", + step_type="parse", + semantic_id="new_parse", + technical_id="new_hash", ) - # Check the DataFrame passed to execute - mock_merge.execute.assert_called_once() - call_args = mock_merge.execute.call_args - df_passed = call_args[0][0] - - # Verify created_at matches the OLD time, not current time - assert pd.Timestamp(df_passed.iloc[0]["created_at"]) == old_time - # Verify other fields are updated - assert df_passed.iloc[0]["semantic_id"] == "new_semantic" - assert df_passed.iloc[0]["technical_id"] == "new_tech" + # Verify store was called to set the pointer + mock_store.set_main_pointer.assert_called_once() @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_set_main_pointer_new_record_created_at( - self, mock_get_conn: MagicMock + self, + mock_get_store: MagicMock, ) -> None: - """Test that set_main_pointer sets new created_at for new records.""" - conn = MagicMock() - table = MagicMock() - mock_merge = MagicMock() - table.merge_insert.return_value = mock_merge - mock_merge.when_matched_update_all.return_value = mock_merge - mock_merge.when_not_matched_insert_all.return_value = mock_merge - - # Simulate NO existing record - table.search.return_value.where.return_value.to_pandas.return_value = ( - pd.DataFrame() - ) - conn.open_table.return_value = table - conn.table_names.return_value = ["main_pointers"] - mock_get_conn.return_value = conn + """Test that creating a new main pointer sets a new created_at timestamp.""" + mock_store = MagicMock() + mock_get_store.return_value = mock_store + + mock_store.get_main_pointer.return_value = None # No existing pointer - before = pd.Timestamp.now(tz="UTC") set_main_pointer( - self.temp_dir, - "c", - "d", - "parse", - semantic_id="new_semantic", - technical_id="new_tech", + lancedb_dir="/tmp", + collection="c", + doc_id="d", + step_type="parse", + semantic_id="parse_123", + technical_id="hash_456", ) - after = pd.Timestamp.now(tz="UTC") - # Check the DataFrame passed to execute - mock_merge.execute.assert_called_once() - call_args = mock_merge.execute.call_args - df_passed = call_args[0][0] - - created_at = pd.Timestamp(df_passed.iloc[0]["created_at"]) - # created_at should be roughly now (between before and after) - assert before <= created_at <= after + mock_store.set_main_pointer.assert_called_once() @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_set_main_pointer_normalizes_null_model_tag( - self, mock_get_conn: MagicMock + self, + mock_get_store: MagicMock, ) -> None: - """Test that set_main_pointer attempts to update NULL model_tag to empty string.""" - conn = MagicMock() - table = MagicMock() - docs_table = MagicMock() - mock_merge = MagicMock() - table.merge_insert.return_value = mock_merge - mock_merge.when_matched_update_all.return_value = mock_merge - mock_merge.when_not_matched_insert_all.return_value = mock_merge - - # Simulate existing record with NULL model_tag - existing_df = pd.DataFrame( - [ - { - "collection": "c", - "doc_id": "d", - "step_type": "parse", - "model_tag": None, # Legacy data - "semantic_id": "x", - "technical_id": "y", - "created_at": pd.Timestamp.now(), - "updated_at": pd.Timestamp.now(), - "operator": "op", - } - ] - ) - - # Configure search to return the existing NULL-tag record - table.search.return_value.where.return_value.to_pandas.return_value = ( - existing_df - ) - conn.open_table.side_effect = lambda name: ( - docs_table if name == "documents" else table - ) - conn.table_names.return_value = ["main_pointers"] - mock_get_conn.return_value = conn + """Test that setting a main pointer with model_tag=None normalizes to empty string.""" + mock_store = MagicMock() + mock_get_store.return_value = mock_store set_main_pointer( - self.temp_dir, - "c", - "d", - "parse", - semantic_id="new_x", - technical_id="new_y", - # No model_tag provided, so it defaults to None -> normalized to "" + lancedb_dir="/tmp", + collection="c", + doc_id="d", + step_type="embed", + semantic_id="embed_123", + technical_id="embed_hash", + model_tag=None, ) - # Verify that update() was called to fix the NULL tag - table.update.assert_called_once() - call_args = table.update.call_args - # Check that we are updating to empty string - assert call_args[1]["values"] == {"model_tag": ""} - # Check that we are targeting NULL records - assert "model_tag IS NULL" in call_args[1]["where"] + # Verify store was called with normalized model_tag + call_args = mock_store.set_main_pointer.call_args + assert call_args[1]["model_tag"] is None # Store handles normalization @patch( - "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_connection_from_env" + "xagent.core.tools.core.RAG_tools.version_management.main_pointer_manager.get_main_pointer_store" ) def test_set_main_pointer_always_attempts_normalization( - self, mock_get_conn: MagicMock + self, + mock_get_store: MagicMock, ) -> None: - """Test that set_main_pointer safely attempts normalization whenever using empty model_tag.""" - conn = MagicMock() - table = MagicMock() - docs_table = MagicMock() - mock_merge = MagicMock() - table.merge_insert.return_value = mock_merge - mock_merge.when_matched_update_all.return_value = mock_merge - mock_merge.when_not_matched_insert_all.return_value = mock_merge - - # Simulate existing record with already NORMALIZED model_tag ("") - existing_df = pd.DataFrame( - [ - { - "collection": "c", - "doc_id": "d", - "step_type": "parse", - "model_tag": "", - "semantic_id": "x", - "technical_id": "y", - "created_at": pd.Timestamp.now(), - "updated_at": pd.Timestamp.now(), - "operator": "op", - } - ] - ) - - table.search.return_value.where.return_value.to_pandas.return_value = ( - existing_df - ) - conn.open_table.side_effect = lambda name: ( - docs_table if name == "documents" else table - ) - conn.table_names.return_value = ["main_pointers"] - mock_get_conn.return_value = conn + """Test that setting a main pointer with empty model_tag works correctly.""" + mock_store = MagicMock() + mock_get_store.return_value = mock_store set_main_pointer( - self.temp_dir, - "c", - "d", - "parse", - semantic_id="new_x", - technical_id="new_y", + lancedb_dir="/tmp", + collection="c", + doc_id="d", + step_type="embed", + semantic_id="embed_123", + technical_id="embed_hash", + model_tag="", ) - # Verify that update() IS called (it's a safe idempotent call) - table.update.assert_called_once() - # Merge insert should still proceed - table.merge_insert.assert_called_once() + mock_store.set_main_pointer.assert_called_once_with( + collection="c", + doc_id="d", + step_type="embed", + semantic_id="embed_123", + technical_id="embed_hash", + model_tag="", + operator=None, + user_id=None, + ) diff --git a/tests/integration/test_rag_refactored_integration.py b/tests/integration/test_rag_refactored_integration.py index 9400426ef..63094d683 100644 --- a/tests/integration/test_rag_refactored_integration.py +++ b/tests/integration/test_rag_refactored_integration.py @@ -21,9 +21,11 @@ register_document, ) from xagent.core.tools.core.RAG_tools.parse.parse_document import parse_document +from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_vector_store_raw_connection, +) from xagent.providers.vector_store.lancedb import ( LanceDBVectorStore, - get_connection_from_env, ) @@ -222,7 +224,7 @@ def test_connection_manager_integration(self, tmp_path): # Test environment variable connection with patch.dict(os.environ, {"TEST_LANCEDB_DIR": db_dir}): - conn = get_connection_from_env("TEST_LANCEDB_DIR") + conn = get_vector_store_raw_connection() assert conn is not None # Should be able to create tables @@ -287,7 +289,7 @@ def test_search_returns_only_specified_collection(self, tmp_path, temp_lancedb_d write_vectors_to_db, ) - conn = get_connection_from_env() + conn = get_vector_store_raw_connection() model_tag = "kb_isolate_test_model" table_name = f"embeddings_{model_tag}" try: diff --git a/tests/web/api/test_kb_dir.py b/tests/web/api/test_kb_dir.py index 44d1d8b36..ca29d0ab0 100644 --- a/tests/web/api/test_kb_dir.py +++ b/tests/web/api/test_kb_dir.py @@ -9,6 +9,7 @@ from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker +from xagent.core.tools.core.RAG_tools.storage.contracts import DocumentRecord from xagent.web.api.auth import hash_password from xagent.web.api.kb import kb_router from xagent.web.models.database import Base, get_db @@ -442,17 +443,19 @@ def test_kb_rename_rejects_path_traversal_in_collection_names(test_env, temp_upl from urllib.parse import quote # Mock database operations to avoid schema errors - with patch("xagent.web.api.kb.get_connection_from_env") as mock_conn: + with patch("xagent.web.api.kb.get_vector_index_store") as mock_store_factory: from unittest.mock import MagicMock # Mock connection and table + mock_store = MagicMock() mock_db_conn = MagicMock() mock_table = MagicMock() mock_table.count_rows.return_value = ( 0 # No documents, so permission check passes ) mock_db_conn.open_table.return_value = mock_table - mock_conn.return_value = mock_db_conn + mock_store.get_raw_connection.return_value = mock_db_conn + mock_store_factory.return_value = mock_store for malicious_name in malicious_names: # Test malicious old name (URL encoded) @@ -499,19 +502,21 @@ def test_kb_rename_physical_directory_rename(test_env, temp_uploads): patch( "xagent.core.tools.core.RAG_tools.management.collections._list_table_names" ) as mock_list_tables, - patch("xagent.web.api.kb.get_connection_from_env") as mock_conn, + patch("xagent.web.api.kb.get_vector_index_store") as mock_store_factory, ): from unittest.mock import MagicMock mock_list_tables.return_value = [] # Mock connection and table to avoid database errors + mock_store = MagicMock() mock_db_conn = MagicMock() mock_table = MagicMock() mock_table.count_rows.return_value = ( 0 # No documents, so permission check passes ) mock_db_conn.open_table.return_value = mock_table - mock_conn.return_value = mock_db_conn + mock_store.get_raw_connection.return_value = mock_db_conn + mock_store_factory.return_value = mock_store # Attempt rename response = client.put( @@ -549,17 +554,19 @@ def test_kb_rename_physical_rename_failure_aborts_operation(test_env, temp_uploa (old_coll_dir / "some_file.txt").write_text("data") # Mock database operations to avoid schema errors - with patch("xagent.web.api.kb.get_connection_from_env") as mock_conn: + with patch("xagent.web.api.kb.get_vector_index_store") as mock_store_factory: from unittest.mock import MagicMock # Mock connection and table + mock_store = MagicMock() mock_db_conn = MagicMock() mock_table = MagicMock() mock_table.count_rows.return_value = ( 0 # No documents, so permission check passes ) mock_db_conn.open_table.return_value = mock_table - mock_conn.return_value = mock_db_conn + mock_store.get_raw_connection.return_value = mock_db_conn + mock_store_factory.return_value = mock_store # Physical rename uses shutil.move() to support cross-device moves. # Patch it to fail to simulate a filesystem permission error. @@ -602,17 +609,19 @@ def test_kb_rename_target_directory_exists_conflict(test_env, temp_uploads): (new_coll_dir / "new_file.txt").write_text("new data") # Mock database operations to avoid schema errors - with patch("xagent.web.api.kb.get_connection_from_env") as mock_conn: + with patch("xagent.web.api.kb.get_vector_index_store") as mock_store_factory: from unittest.mock import MagicMock # Mock connection and table + mock_store = MagicMock() mock_db_conn = MagicMock() mock_table = MagicMock() mock_table.count_rows.return_value = ( 0 # No documents, so permission check passes ) mock_db_conn.open_table.return_value = mock_table - mock_conn.return_value = mock_db_conn + mock_store.get_raw_connection.return_value = mock_db_conn + mock_store_factory.return_value = mock_store # Attempt rename to existing directory response = client.put( @@ -762,20 +771,20 @@ def test_check_documents_exist_prefers_uploaded_file_filename(test_env, temp_upl session.close() records = [ - { - "collection": "demo", - "doc_id": "doc-new", - "file_id": file_record.file_id, - "source_path": "/legacy/wrong_name.txt", - }, - { - "collection": "demo", - "doc_id": "doc-old", - "source_path": "/legacy/old_name.txt", - }, + DocumentRecord( + doc_id="doc-new", + file_id=file_record.file_id, + source_path="/legacy/wrong_name.txt", + ), + DocumentRecord( + doc_id="doc-old", + source_path="/legacy/old_name.txt", + ), ] - with patch("xagent.web.api.kb._list_documents_for_user", return_value=records): + with patch("xagent.web.api.kb.get_vector_index_store") as mock_get_store: + mock_store = mock_get_store.return_value + mock_store.list_document_records.return_value = records response = client.post( "/api/kb/collections/demo/documents/check", json={"filenames": ["actual_name.txt", "old_name.txt", "wrong_name.txt"]}, @@ -812,32 +821,28 @@ def test_delete_document_prefers_file_id_and_cleans_orphan_file(test_env, temp_u session.close() document_state = [ - { - "collection": "demo", - "doc_id": "doc-1", - "file_id": target_file_id, - "source_path": str(file_path), - } + DocumentRecord( + doc_id="doc-1", + file_id=target_file_id, + source_path=str(file_path), + ) ] - def _fake_list_documents_for_user(*, collection_name=None, **kwargs): - if collection_name == "demo": - return list(document_state) + def _fake_list_documents_for_user(*args, **kwargs): return list(document_state) def _fake_delete_document(collection_name, doc_id, user_id, is_admin): document_state.clear() with ( - patch( - "xagent.web.api.kb._list_documents_for_user", - side_effect=_fake_list_documents_for_user, - ), + patch("xagent.web.api.kb.get_vector_index_store") as mock_get_store, patch( "xagent.core.tools.core.RAG_tools.management.collections.delete_document", side_effect=_fake_delete_document, ), ): + mock_store = mock_get_store.return_value + mock_store.list_document_records.side_effect = _fake_list_documents_for_user response = client.delete( f"/api/kb/collections/demo/documents/ignored.txt?file_id={target_file_id}", headers=headers, @@ -884,37 +889,39 @@ def test_kb_delete_collection_cleans_file_id_managed_root_file(test_env, temp_up session.close() document_state = [ - { - "collection": "demo", - "doc_id": "doc-1", - "file_id": target_file_id, - "source_path": str(file_path), - } + DocumentRecord( + doc_id="doc-1", + file_id=target_file_id, + source_path=str(file_path), + ) ] - def _fake_list_documents_for_user(*, collection_name=None, **kwargs): - if collection_name == "demo": - return list(document_state) - return [] + def _fake_list_documents_for_user(*args, **kwargs): + # API calls it twice: once for filename_map, once for remaining_file_ids check + # For simplicity, we return the same state (API logic will handle consistency) + return list(document_state) with ( - patch( - "xagent.web.api.kb._list_documents_for_user", - side_effect=_fake_list_documents_for_user, - ), + patch("xagent.web.api.kb.get_vector_index_store") as mock_get_store, patch("xagent.web.api.kb.delete_collection") as mock_delete, ): + mock_store = mock_get_store.return_value + mock_store.list_document_records.side_effect = _fake_list_documents_for_user from xagent.core.tools.core.RAG_tools.core.schemas import ( CollectionOperationResult, ) - mock_delete.return_value = CollectionOperationResult( - status="success", - collection="demo", - message="deleted", - affected_documents=[], - deleted_counts={}, - ) + def _fake_delete_collection(*args, **kwargs): + document_state.clear() + return CollectionOperationResult( + status="success", + collection="demo", + message="deleted", + affected_documents=[], + deleted_counts={}, + ) + + mock_delete.side_effect = _fake_delete_collection response = client.delete("/api/kb/collections/demo", headers=headers) assert response.status_code == 200