feat(lancedb): automatic schema migration and collection_metadata centralization

- schema_manager: add _get_sql_default_for_pa_type() to map PyArrow types to
  LanceDB SQL default expressions
- Replace _validate_schema_fields with _ensure_schema_fields(): add missing
  columns to existing tables instead of raising; use add_columns() with
  type-appropriate defaults
- _create_table: when table exists and schema given, call _ensure_schema_fields
  for best-effort schema evolution
- ensure_chunks_table / ensure_embeddings_table: remove required_fields
  validation and long docstrings; rely on auto-migration
- Add ensure_collection_metadata_table(); trim docstrings for main_pointers,
  prompt_templates, ingestion_runs
- collection_manager: use ensure_collection_metadata_table(conn) in save path,
  remove _ensure_metadata_table() and pyarrow import
- Add tests/core/tools/core/RAG_tools/LanceDB/test_schema_migration.py for
  default mapping, auto-migration, idempotency, and _ensure_schema_fields

Migrated from FenixAOS fix/schema_adapter.

Made-with: Cursor

Author: sqhyz55

src/xagent/core/tools/core/RAG_tools/LanceDB/schema_manager.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from typing import Any
 
 import pyarrow as pa  # type: ignore
 from lancedb.db import DBConnection
@@ -15,6 +16,7 @@
     "ensure_main_pointers_table",
     "ensure_prompt_templates_table",
     "ensure_ingestion_runs_table",
+    "ensure_collection_metadata_table",
 ]
 
 
@@ -26,18 +28,27 @@ def _table_exists(conn: DBConnection, name: str) -> bool:
         return False
 
 
-def _validate_schema_fields(
-    conn: DBConnection, table_name: str, required_fields: list[str]
-) -> None:
-    """Validate that an existing table contains all required fields.
+def _get_sql_default_for_pa_type(pa_type: Any) -> str:
+    """Map PyArrow type to LanceDB SQL default value expression."""
+    if pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type):
+        return "''"
+    if pa.types.is_integer(pa_type):
+        return "0"
+    if pa.types.is_floating(pa_type):
+        return "0.0"
+    if pa.types.is_boolean(pa_type):
+        return "false"
+    if pa.types.is_timestamp(pa_type):
+        return "CAST(NULL AS TIMESTAMP)"
+    return "NULL"
+
 
-    Args:
-        conn: LanceDB connection
-        table_name: Name of the table to validate
-        required_fields: List of required field names
+def _ensure_schema_fields(
+    conn: DBConnection, table_name: str, target_schema: Any
+) -> None:
+    """Ensure an existing table matches the target schema by adding missing columns.
 
-    Raises:
-        ValueError: If the table exists but is missing required fields.
+    Only ADDS missing columns. Does not delete extra columns nor modify existing types.
     """
     if not _table_exists(conn, table_name):
         return
@@ -46,34 +57,37 @@ def _validate_schema_fields(
         table = conn.open_table(table_name)
         existing_schema = table.schema
         existing_field_names = {field.name for field in existing_schema}
-
-        missing_fields = [f for f in required_fields if f not in existing_field_names]
+        missing_fields = [
+            f for f in target_schema if f.name not in existing_field_names
+        ]
 
         if missing_fields:
-            error_msg = (
-                f"Table '{table_name}' exists but is missing required fields: {missing_fields}. "
-                f"This is likely due to a schema upgrade. "
-                f"Please delete the existing table or manually add the missing fields. "
-                f"Note: During development, we do not provide automatic migration scripts. "
-                f"To upgrade, you can either:\n"
-                f"1. Delete the table (data will be lost): conn.drop_table('{table_name}')\n"
-                f"2. Manually add the missing fields using LanceDB's schema update capabilities"
+            logger.info(
+                "Auto-migrating schema for table '%s'. Adding missing fields: %s",
+                table_name,
+                [f.name for f in missing_fields],
             )
-            logger.error(error_msg)
-            raise ValueError(error_msg)
-    except ValueError:
-        # Re-raise ValueError (our validation error)
-        raise
+            new_cols = {}
+            for field in missing_fields:
+                default_expr = _get_sql_default_for_pa_type(field.type)
+                new_cols[field.name] = default_expr
+            try:
+                table.add_columns(new_cols)
+                logger.info("Successfully migrated schema for table '%s'", table_name)
+            except Exception as e:
+                logger.error("Failed to add columns to table '%s': %s", table_name, e)
     except Exception as e:
-        # Log other errors but don't fail - schema validation is best-effort
         logger.warning(
-            f"Could not validate schema for table '{table_name}': {e}. "
-            f"Proceeding with table creation/usage."
+            "Could not validate or migrate schema for table '%s': %s",
+            table_name,
+            e,
         )
 
 
-def _create_table(conn: DBConnection, name: str, schema: object | None = None) -> None:
+def _create_table(conn: DBConnection, name: str, schema: Any | None = None) -> None:
     if _table_exists(conn, name):
+        if schema:
+            _ensure_schema_fields(conn, name, schema)
         return
     conn.create_table(name, schema=schema)
 
@@ -139,32 +153,9 @@ def ensure_parses_table(conn: DBConnection) -> None:
 def ensure_chunks_table(conn: DBConnection) -> None:
     """Ensure the chunks table exists with proper schema.
 
-    This function creates the table if it doesn't exist, and validates that
-    existing tables contain all required fields (especially 'metadata').
-
-    Args:
-        conn: LanceDB connection
-
-    Raises:
-        ValueError: If the table exists but is missing required fields.
-            This typically happens when an old table schema doesn't include
-            the 'metadata' field. During development, we do not provide
-            automatic migration scripts. Users must either delete the table
-            or manually add the missing fields.
-
-    Note:
-        There's no upgrade path for existing chunks tables. Any deployment
-        with an existing table will hit schema-mismatch errors once the pipeline
-        starts writing a column that doesn't exist. If you encounter this error,
-        you need to either delete the existing table or manually add the missing
-        'metadata' field.
+    If the table already exists, we attempt best-effort schema evolution by
+    adding any missing columns (see _ensure_schema_fields).
     """
-    # Required fields that must exist in the table (especially for schema validation)
-    required_fields = ["metadata"]
-
-    # Validate existing table schema before creating/using it
-    _validate_schema_fields(conn, "chunks", required_fields)
-
     schema = pa.schema(
         [
             pa.field("collection", pa.string()),
@@ -192,36 +183,11 @@ def ensure_embeddings_table(
 ) -> None:
     """Ensure the embeddings table exists with proper schema.
 
-    This function creates the table if it doesn't exist, and validates that
-    existing tables contain all required fields (especially 'metadata').
-
-    Args:
-        conn: LanceDB connection
-        model_tag: Model tag used to construct the table name (e.g., 'bge_large')
-        vector_dim: Optional vector dimension for fixed-size vectors
-
-    Raises:
-        ValueError: If the table exists but is missing required fields.
-            This typically happens when an old table schema doesn't include
-            the 'metadata' field. During development, we do not provide
-            automatic migration scripts. Users must either delete the table
-            or manually add the missing fields.
-
-    Note:
-        There's no upgrade path for existing embeddings tables. Any deployment
-        with an existing table will hit schema-mismatch errors once the pipeline
-        starts writing a column that doesn't exist. If you encounter this error,
-        you need to either delete the existing table or manually add the missing
-        'metadata' field.
+    If the table already exists, we attempt best-effort schema evolution by
+    adding any missing columns (see _ensure_schema_fields).
     """
     table_name = f"embeddings_{model_tag}"
 
-    # Required fields that must exist in the table (especially for schema validation)
-    required_fields = ["metadata"]
-
-    # Validate existing table schema before creating/using it
-    _validate_schema_fields(conn, table_name, required_fields)
-
     # Support dynamic vector dimension: if provided, create a FixedSizeList; otherwise allow variable-length
     vector_field_type = (
         pa.list_(pa.float32(), list_size=vector_dim)
@@ -252,11 +218,6 @@ def ensure_embeddings_table(
 
 
 def ensure_main_pointers_table(conn: DBConnection) -> None:
-    """Ensure the main_pointers table exists with proper schema.
-
-    Args:
-        conn: LanceDB connection
-    """
     schema = pa.schema(
         [
             pa.field("collection", pa.string()),
@@ -274,11 +235,6 @@ def ensure_main_pointers_table(conn: DBConnection) -> None:
 
 
 def ensure_prompt_templates_table(conn: DBConnection) -> None:
-    """Ensure the prompt_templates table exists with proper schema.
-
-    Args:
-        conn: LanceDB connection
-    """
     table_name = "prompt_templates"
     schema = pa.schema(
         [
@@ -311,13 +267,6 @@ def ensure_prompt_templates_table(conn: DBConnection) -> None:
 
 
 def ensure_ingestion_runs_table(conn: DBConnection) -> None:
-    """Ensure the ingestion_runs table exists with proper schema.
-
-    This table tracks the status of document ingestion processes.
-
-    Args:
-        conn: LanceDB connection
-    """
     schema = pa.schema(
         [
             pa.field("collection", pa.string()),
@@ -346,3 +295,29 @@ def ensure_ingestion_runs_table(conn: DBConnection) -> None:
             )
 
     _create_table(conn, "ingestion_runs", schema=schema)
+
+
+def ensure_collection_metadata_table(conn: DBConnection) -> None:
+    """Ensure the collection_metadata table exists with proper schema."""
+    schema = pa.schema(
+        [
+            pa.field("name", pa.string()),
+            pa.field("schema_version", pa.string()),
+            pa.field("embedding_model_id", pa.string()),
+            pa.field("embedding_dimension", pa.int32()),
+            pa.field("documents", pa.int32()),
+            pa.field("processed_documents", pa.int32()),
+            pa.field("parses", pa.int32()),
+            pa.field("chunks", pa.int32()),
+            pa.field("embeddings", pa.int32()),
+            pa.field("document_names", pa.string()),
+            pa.field("collection_locked", pa.bool_()),
+            pa.field("allow_mixed_parse_methods", pa.bool_()),
+            pa.field("skip_config_validation", pa.bool_()),
+            pa.field("created_at", pa.timestamp("us")),
+            pa.field("updated_at", pa.timestamp("us")),
+            pa.field("last_accessed_at", pa.timestamp("us")),
+            pa.field("extra_metadata", pa.string()),
+        ]
+    )
+    _create_table(conn, "collection_metadata", schema=schema)

src/xagent/core/tools/core/RAG_tools/management/collection_manager.py

@@ -12,11 +12,10 @@
 from functools import wraps
 from typing import Any, Awaitable, Callable, Optional, TypeVar
 
-import pyarrow as pa  # type: ignore
-
 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 ..utils.model_resolver import resolve_embedding_adapter
 from ..utils.string_utils import escape_lancedb_string
 
@@ -198,8 +197,8 @@ async def _save_collection_with_retry(
 
         for attempt in range(max_retries):
             try:
-                # Ensure collection_metadata table exists
-                await self._ensure_metadata_table()
+                # Ensure collection_metadata table exists and is up to date
+                ensure_collection_metadata_table(conn)
 
                 # Prepare data for storage
                 data = collection.to_storage()
@@ -235,53 +234,6 @@ async def _save_collection_with_retry(
                 )
                 await asyncio.sleep(wait_time)
 
-    async def _ensure_metadata_table(self) -> None:
-        """Ensure collection_metadata table exists in LanceDB.
-
-        Creates the table if it doesn't exist, otherwise does nothing.
-        """
-
-        conn = await self._get_connection()
-
-        schema = pa.schema(
-            [
-                ("name", pa.string()),
-                ("schema_version", pa.string()),
-                ("embedding_model_id", pa.string()),  # Nullable
-                ("embedding_dimension", pa.int32()),  # Nullable
-                ("documents", pa.int32()),
-                ("processed_documents", pa.int32()),
-                ("parses", pa.int32()),
-                ("chunks", pa.int32()),
-                ("embeddings", pa.int32()),
-                ("document_names", pa.string()),  # JSON string
-                ("collection_locked", pa.bool_()),
-                ("allow_mixed_parse_methods", pa.bool_()),
-                ("skip_config_validation", pa.bool_()),
-                ("created_at", pa.timestamp("us")),
-                ("updated_at", pa.timestamp("us")),
-                ("last_accessed_at", pa.timestamp("us")),
-                ("extra_metadata", pa.string()),  # JSON string
-            ]
-        )
-
-        # Check if table already exists
-        table_names_fn = getattr(conn, "table_names", None)
-        table_exists = False
-        if table_names_fn:
-            try:
-                existing_tables = table_names_fn()
-                table_exists = "collection_metadata" in existing_tables
-            except Exception as e:
-                logger.debug(f"Table names check failed: {e}")
-
-        if not table_exists:
-            try:
-                conn.create_table("collection_metadata", schema=schema)
-            except Exception as e:
-                logger.debug(f"Table creation failed (may already exist): {e}")
-                # Table might already exist, continue
-
     async def initialize_collection_embedding(
         self, collection_name: str, embedding_model_id: str
     ) -> "CollectionInfo":