diff --git a/src/xagent/core/tools/core/RAG_tools/LanceDB/schema_manager.py b/src/xagent/core/tools/core/RAG_tools/LanceDB/schema_manager.py index f6d6341ff..8a879eacd 100644 --- a/src/xagent/core/tools/core/RAG_tools/LanceDB/schema_manager.py +++ b/src/xagent/core/tools/core/RAG_tools/LanceDB/schema_manager.py @@ -16,6 +16,7 @@ "ensure_parses_table", "ensure_chunks_table", "ensure_embeddings_table", + "ensure_active_generations_table", "ensure_main_pointers_table", "ensure_prompt_templates_table", "ensure_ingestion_runs_table", @@ -34,6 +35,27 @@ def _safe_close_table(table: Any) -> None: pass +def _add_nullable_string_columns( + conn: DBConnection, table_name: str, columns: dict[str, str] +) -> None: + """Add missing nullable string columns to an existing LanceDB table.""" + if not _table_exists(conn, table_name): + return + table = conn.open_table(table_name) + try: + existing = set(table.schema.names) + to_add = {name: expr for name, expr in columns.items() if name not in existing} + if to_add: + logger.info( + "Migrating '%s' table: adding columns %s", + table_name, + list(to_add.keys()), + ) + table.add_columns(to_add) + finally: + _safe_close_table(table) + + def _table_exists(conn: DBConnection, name: str) -> bool: try: return name in list_table_names(conn) @@ -352,6 +374,9 @@ def ensure_chunks_table(conn: DBConnection) -> None: _validate_schema_fields(conn, "chunks", required_fields) if _table_exists(conn, "chunks"): _migrate_table_user_id_to_int64(conn, "chunks") + _add_nullable_string_columns( + conn, "chunks", {"generation_id": "cast(null as string)"} + ) val_table = conn.open_table("chunks") try: @@ -373,6 +398,7 @@ def ensure_chunks_table(conn: DBConnection) -> None: pa.field("json_path", pa.string()), pa.field("chunk_hash", pa.string()), pa.field("config_hash", pa.string()), + pa.field("generation_id", pa.string()), pa.field("created_at", pa.timestamp("us")), pa.field("metadata", pa.string()), pa.field("user_id", pa.int64()), @@ -417,6 +443,14 @@ def ensure_embeddings_table( _validate_schema_fields(conn, table_name, required_fields) if _table_exists(conn, table_name): _migrate_table_user_id_to_int64(conn, table_name) + _add_nullable_string_columns( + conn, + table_name, + { + "generation_id": "cast(null as string)", + "config_hash": "cast(null as string)", + }, + ) val_table = conn.open_table(table_name) try: @@ -441,6 +475,8 @@ def ensure_embeddings_table( pa.field("vector_dimension", pa.int32()), pa.field("text", pa.large_string()), pa.field("chunk_hash", pa.string()), + pa.field("config_hash", pa.string()), + pa.field("generation_id", pa.string()), pa.field("created_at", pa.timestamp("us")), pa.field("metadata", pa.string()), pa.field("user_id", pa.int64()), @@ -453,6 +489,30 @@ def ensure_embeddings_table( ) +def ensure_active_generations_table(conn: DBConnection) -> None: + """Ensure the active_generations table exists with proper schema. + + Tracks the published searchable generation per document scope: + (collection, doc_id, parse_hash, user_id, model_tag). + """ + schema = pa.schema( + [ + pa.field("collection", pa.string()), + pa.field("doc_id", pa.string()), + pa.field("parse_hash", pa.string()), + pa.field("user_id", pa.int64()), + pa.field("model_tag", pa.string()), + pa.field("generation_id", pa.string()), + pa.field("config_hash", pa.string()), + pa.field("created_at", pa.timestamp("us")), + pa.field("updated_at", pa.timestamp("us")), + pa.field("published_at", pa.timestamp("us")), + pa.field("operator", pa.string()), + ] + ) + _create_table(conn, "active_generations", schema=schema) + + def ensure_main_pointers_table(conn: DBConnection) -> None: """Ensure the main_pointers table exists with proper schema. diff --git a/src/xagent/core/tools/core/RAG_tools/kb/storage_shim.py b/src/xagent/core/tools/core/RAG_tools/kb/storage_shim.py index 1958beac1..08776117b 100644 --- a/src/xagent/core/tools/core/RAG_tools/kb/storage_shim.py +++ b/src/xagent/core/tools/core/RAG_tools/kb/storage_shim.py @@ -6,6 +6,7 @@ if TYPE_CHECKING: from ..storage.contracts import ( + ActiveGenerationStore, IngestionStatusStore, KBWriteCoordinator, MainPointerStore, @@ -59,6 +60,10 @@ def get_main_pointer_store(self) -> MainPointerStore: """Return the legacy main pointer store singleton.""" return self._storage_factory.get_main_pointer_store() + def get_active_generation_store(self) -> ActiveGenerationStore: + """Return the legacy active generation store singleton.""" + return self._storage_factory.get_active_generation_store() + def reset_kb_write_coordinator(self) -> None: """Clear legacy storage singletons and coordinator-backed shim state.""" self._storage_factory.reset_all() diff --git a/src/xagent/core/tools/core/RAG_tools/storage/__init__.py b/src/xagent/core/tools/core/RAG_tools/storage/__init__.py index ee390fd4f..48045ebc4 100644 --- a/src/xagent/core/tools/core/RAG_tools/storage/__init__.py +++ b/src/xagent/core/tools/core/RAG_tools/storage/__init__.py @@ -4,6 +4,7 @@ """ from .contracts import ( + ActiveGenerationStore, IngestionStatusStore, KBWriteCoordinator, MainPointerStore, @@ -13,6 +14,7 @@ ) from .factory import ( StorageFactory, + get_active_generation_store, get_ingestion_status_store, get_kb_write_coordinator, get_main_pointer_store, @@ -38,6 +40,7 @@ "IngestionStatusStore", "PromptTemplateStore", "MainPointerStore", + "ActiveGenerationStore", # Factory "StorageFactory", "get_kb_write_coordinator", @@ -51,6 +54,7 @@ "get_ingestion_status_store", "get_prompt_template_store", "get_main_pointer_store", + "get_active_generation_store", "reset_kb_write_coordinator", "reset_rag_storage_for_tests", ] diff --git a/src/xagent/core/tools/core/RAG_tools/storage/contracts.py b/src/xagent/core/tools/core/RAG_tools/storage/contracts.py index d81be5abb..773b6500b 100644 --- a/src/xagent/core/tools/core/RAG_tools/storage/contracts.py +++ b/src/xagent/core/tools/core/RAG_tools/storage/contracts.py @@ -56,6 +56,7 @@ "json_path", "chunk_hash", "config_hash", + "generation_id", "metadata", # embeddings table "model", @@ -71,6 +72,8 @@ "semantic_id", "technical_id", "operator", + # active_generations table + "published_at", # prompt_templates table "id", "name", @@ -1794,3 +1797,93 @@ async def delete_main_pointer_async( user_id: Optional[int] = None, ) -> bool: """Async version of delete_main_pointer.""" + + +class _UnsetUserIdFilter: + """Sentinel type for omitted active-generation user_id filters.""" + + +USER_ID_FILTER_UNSET = _UnsetUserIdFilter() +UserIdFilter = Union[Optional[int], _UnsetUserIdFilter] + + +class ActiveGenerationStore(ABC): + """Published active-generation pointer contract for ingestion and retrieval. + + Each pointer identifies the single searchable generation for a document scope: + (collection, doc_id, parse_hash, user_id, model_tag). + """ + + @abstractmethod + def publish_active_generation( + self, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + model_tag: Optional[str], + generation_id: str, + config_hash: str, + operator: Optional[str] = None, + ) -> None: + """Publish or replace the active generation pointer for a scope.""" + + @abstractmethod + def get_active_generation( + self, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + model_tag: Optional[str] = None, + ) -> Optional[Dict[str, Any]]: + """Return the active generation pointer for a scope, if any.""" + + @abstractmethod + def list_active_generations( + self, + collection: str, + *, + doc_id: Optional[str] = None, + user_id: UserIdFilter = USER_ID_FILTER_UNSET, + model_tag: Optional[str] = None, + limit: int = 1000, + ) -> List[Dict[str, Any]]: + """List active generation pointers for a collection.""" + + @abstractmethod + async def publish_active_generation_async( + self, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + model_tag: Optional[str], + generation_id: str, + config_hash: str, + operator: Optional[str] = None, + ) -> None: + """Async version of publish_active_generation.""" + + @abstractmethod + async def get_active_generation_async( + self, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + model_tag: Optional[str] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of get_active_generation.""" + + @abstractmethod + async def list_active_generations_async( + self, + collection: str, + *, + doc_id: Optional[str] = None, + user_id: UserIdFilter = USER_ID_FILTER_UNSET, + model_tag: Optional[str] = None, + limit: int = 1000, + ) -> List[Dict[str, Any]]: + """Async version of list_active_generations.""" diff --git a/src/xagent/core/tools/core/RAG_tools/storage/factory.py b/src/xagent/core/tools/core/RAG_tools/storage/factory.py index 64fabe85d..2cf5096f2 100644 --- a/src/xagent/core/tools/core/RAG_tools/storage/factory.py +++ b/src/xagent/core/tools/core/RAG_tools/storage/factory.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Optional, cast from .contracts import ( + ActiveGenerationStore, IngestionStatusStore, KBWriteCoordinator, MainPointerStore, @@ -25,6 +26,7 @@ VectorIndexStore, ) from .lancedb_stores import ( + LanceDBActiveGenerationStore, LanceDBIngestionStatusStore, LanceDBMainPointerStore, LanceDBMetadataStore, @@ -96,6 +98,7 @@ def __init__(self) -> None: self._ingestion_status_store: Optional[IngestionStatusStore] = None self._prompt_template_store: Optional[PromptTemplateStore] = None self._main_pointer_store: Optional[MainPointerStore] = None + self._active_generation_store: Optional[ActiveGenerationStore] = None self._coordinator: Optional[KBWriteCoordinator] = None @classmethod @@ -126,6 +129,7 @@ def reset_all(self) -> None: self._ingestion_status_store = None self._prompt_template_store = None self._main_pointer_store = None + self._active_generation_store = None self._coordinator = None # Keep semantic KB coordinator state aligned with factory-backed stores. @@ -217,6 +221,16 @@ def get_prompt_template_store(self) -> PromptTemplateStore: self._prompt_template_store = LanceDBPromptTemplateStore() return self._prompt_template_store + # --- ActiveGenerationStore --- + + def get_active_generation_store(self) -> ActiveGenerationStore: + """Get or create active generation pointer store.""" + if self._active_generation_store is None: + with self._lock: + if self._active_generation_store is None: + self._active_generation_store = LanceDBActiveGenerationStore() + return self._active_generation_store + # --- MainPointerStore --- def get_main_pointer_store(self) -> MainPointerStore: @@ -359,6 +373,15 @@ def get_main_pointer_store() -> MainPointerStore: return _get_storage_shim().get_main_pointer_store() +def get_active_generation_store() -> ActiveGenerationStore: + """Get active generation pointer store. + + Returns: + LanceDBActiveGenerationStore instance. + """ + return _get_storage_shim().get_active_generation_store() + + # ============================================================================ # Default Coordinator Implementation # ============================================================================ 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 index 99c64f613..90ed8e39d 100644 --- a/src/xagent/core/tools/core/RAG_tools/storage/lancedb_stores.py +++ b/src/xagent/core/tools/core/RAG_tools/storage/lancedb_stores.py @@ -26,6 +26,8 @@ from ..utils.string_utils import build_lancedb_filter_expression, escape_lancedb_string from ..utils.user_permissions import UserPermissions from .contracts import ( + USER_ID_FILTER_UNSET, + ActiveGenerationStore, DocumentRecord, FilterCondition, FilterExpression, @@ -34,6 +36,7 @@ MainPointerStore, MetadataStore, PromptTemplateStore, + UserIdFilter, VectorIndexStore, build_filter_from_dict, ) @@ -3352,3 +3355,269 @@ async def delete_main_pointer_async( return self.delete_main_pointer( collection, doc_id, step_type, model_tag, user_id ) + + +class LanceDBActiveGenerationStore(ActiveGenerationStore): + """LanceDB implementation for active generation pointer management.""" + + LEGACY_USER_ID = -1 + + _MERGE_KEYS = ( + "collection", + "doc_id", + "parse_hash", + "user_id", + "model_tag", + ) + + def __init__(self) -> None: + self._sync_conn: Optional[DBConnection] = None + # Process-local guard so we only run schema setup once per store + # instance; subsequent publish/get/list calls skip the + # list_table_names round-trip on the hot path. + self._table_ensured: bool = False + + def _get_sync_connection(self) -> DBConnection: + if self._sync_conn is None: + self._sync_conn = get_connection_from_env() + return self._sync_conn + + def _ensure_table(self) -> None: + if self._table_ensured: + return + from ..LanceDB.schema_manager import ensure_active_generations_table + + ensure_active_generations_table(self._get_sync_connection()) + self._table_ensured = True + + @staticmethod + def _normalize_model_tag(model_tag: Optional[str]) -> str: + return model_tag if model_tag is not None else "" + + @classmethod + def _normalize_user_id(cls, user_id: Optional[int]) -> int: + """Map legacy user_id=None scope to a non-null merge key value.""" + return cls.LEGACY_USER_ID if user_id is None else user_id + + @classmethod + def _denormalize_record_user_id(cls, record: Dict[str, Any]) -> Dict[str, Any]: + """Hide the internal legacy-user sentinel from store callers.""" + if record.get("user_id") != cls.LEGACY_USER_ID: + return record + normalized_record = dict(record) + normalized_record["user_id"] = None + return normalized_record + + @classmethod + def _scope_filter_str( + cls, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + normalized_tag: str, + ) -> str: + """Build a SQL filter selecting exactly one pointer scope.""" + normalized_user_id = cls._normalize_user_id(user_id) + filter_expr: FilterExpression = ( + FilterCondition( + field="collection", operator=FilterOperator.EQ, value=collection + ), + FilterCondition(field="doc_id", operator=FilterOperator.EQ, value=doc_id), + FilterCondition( + field="parse_hash", operator=FilterOperator.EQ, value=parse_hash + ), + FilterCondition( + field="user_id", operator=FilterOperator.EQ, value=normalized_user_id + ), + FilterCondition( + field="model_tag", operator=FilterOperator.EQ, value=normalized_tag + ), + ) + return translate_filter_expression(filter_expr) + + def publish_active_generation( + self, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + model_tag: Optional[str], + generation_id: str, + config_hash: str, + operator: Optional[str] = None, + ) -> None: + from ..LanceDB.schema_manager import _safe_close_table + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("active_generations") + normalized_tag = self._normalize_model_tag(model_tag) + normalized_user_id = self._normalize_user_id(user_id) + + try: + now = datetime.now(timezone.utc).replace(tzinfo=None) + + # Look up created_at on the already-opened table to avoid the + # extra open_table/close round-trip that calling + # get_active_generation() would incur. + filter_str = self._scope_filter_str( + collection, doc_id, parse_hash, user_id, normalized_tag + ) + existing_rows = table.search().where(filter_str).to_arrow().to_pylist() + created_at = existing_rows[0]["created_at"] if existing_rows else now + + record = { + "collection": collection, + "doc_id": doc_id, + "parse_hash": parse_hash, + "user_id": normalized_user_id, + "model_tag": normalized_tag, + "generation_id": generation_id, + "config_hash": config_hash, + "created_at": created_at, + "updated_at": now, + "published_at": now, + "operator": operator or "unknown", + } + + ( + table.merge_insert(on=list(self._MERGE_KEYS)) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute([record]) + ) + finally: + _safe_close_table(table) + + def get_active_generation( + self, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + model_tag: Optional[str] = None, + ) -> Optional[Dict[str, Any]]: + from ..LanceDB.schema_manager import _safe_close_table + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("active_generations") + normalized_tag = self._normalize_model_tag(model_tag) + + try: + filter_str = self._scope_filter_str( + collection, doc_id, parse_hash, user_id, normalized_tag + ) + result = table.search().where(filter_str).to_arrow() + if len(result) == 0: + return None + rows = result.to_pylist() + return cast(Dict[str, Any], self._denormalize_record_user_id(rows[0])) + finally: + _safe_close_table(table) + + def list_active_generations( + self, + collection: str, + *, + doc_id: Optional[str] = None, + user_id: UserIdFilter = USER_ID_FILTER_UNSET, + model_tag: Optional[str] = None, + limit: int = 1000, + ) -> List[Dict[str, Any]]: + from ..LanceDB.schema_manager import _safe_close_table + + conn = self._get_sync_connection() + self._ensure_table() + table = conn.open_table("active_generations") + + try: + conditions: List[FilterExpression] = [ + FilterCondition( + field="collection", operator=FilterOperator.EQ, value=collection + ) + ] + if doc_id is not None: + conditions.append( + FilterCondition( + field="doc_id", operator=FilterOperator.EQ, value=doc_id + ) + ) + if user_id is not USER_ID_FILTER_UNSET: + conditions.append( + FilterCondition( + field="user_id", + operator=FilterOperator.EQ, + value=self._normalize_user_id(cast(Optional[int], user_id)), + ) + ) + if model_tag is not None: + conditions.append( + FilterCondition( + field="model_tag", + operator=FilterOperator.EQ, + value=self._normalize_model_tag(model_tag), + ) + ) + + filter_expr: FilterExpression = ( + conditions[0] if len(conditions) == 1 else tuple(conditions) + ) + filter_str = translate_filter_expression(filter_expr) + result = table.search().where(filter_str).limit(limit).to_arrow() + rows = result.to_pylist() + return [self._denormalize_record_user_id(row) for row in rows] + finally: + _safe_close_table(table) + + async def publish_active_generation_async( + self, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + model_tag: Optional[str], + generation_id: str, + config_hash: str, + operator: Optional[str] = None, + ) -> None: + return self.publish_active_generation( + collection, + doc_id, + parse_hash, + user_id, + model_tag, + generation_id, + config_hash, + operator, + ) + + async def get_active_generation_async( + self, + collection: str, + doc_id: str, + parse_hash: str, + user_id: Optional[int], + model_tag: Optional[str] = None, + ) -> Optional[Dict[str, Any]]: + return self.get_active_generation( + collection, doc_id, parse_hash, user_id, model_tag + ) + + async def list_active_generations_async( + self, + collection: str, + *, + doc_id: Optional[str] = None, + user_id: UserIdFilter = USER_ID_FILTER_UNSET, + model_tag: Optional[str] = None, + limit: int = 1000, + ) -> List[Dict[str, Any]]: + return self.list_active_generations( + collection, + doc_id=doc_id, + user_id=user_id, + model_tag=model_tag, + limit=limit, + ) diff --git a/tests/core/tools/core/RAG_tools/LanceDB/test_generation_schema.py b/tests/core/tools/core/RAG_tools/LanceDB/test_generation_schema.py new file mode 100644 index 000000000..227f5e85d --- /dev/null +++ b/tests/core/tools/core/RAG_tools/LanceDB/test_generation_schema.py @@ -0,0 +1,97 @@ +"""Schema tests for generation_id and active_generations (issue #438 PR1).""" + +from __future__ import annotations + +from pathlib import Path + +import pyarrow as pa + +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_active_generations_table, + ensure_chunks_table, + ensure_embeddings_table, +) +from xagent.core.tools.core.RAG_tools.storage import get_vector_store_raw_connection + + +def _schema_field_names(conn, table_name: str) -> set[str]: + table = conn.open_table(table_name) + try: + return set(table.schema.names) + finally: + if hasattr(table, "close"): + table.close() + + +def test_chunks_table_includes_generation_id(tmp_path: Path, monkeypatch) -> None: + """New and migrated chunks tables must expose generation_id.""" + monkeypatch.setenv("LANCEDB_DIR", str(tmp_path / "db")) + conn = get_vector_store_raw_connection() + ensure_chunks_table(conn) + assert "generation_id" in _schema_field_names(conn, "chunks") + + +def test_embeddings_table_includes_generation_id_and_config_hash( + tmp_path: Path, monkeypatch +) -> None: + """Embeddings tables must expose generation_id and config_hash.""" + monkeypatch.setenv("LANCEDB_DIR", str(tmp_path / "db")) + conn = get_vector_store_raw_connection() + model_tag = to_model_tag("BAAI/bge-large-zh-v1.5") + ensure_embeddings_table(conn, model_tag, vector_dim=4) + table_name = f"embeddings_{model_tag}" + names = _schema_field_names(conn, table_name) + assert "generation_id" in names + assert "config_hash" in names + + +def test_active_generations_table_schema(tmp_path: Path, monkeypatch) -> None: + """active_generations table must include tenant-scoped pointer fields.""" + monkeypatch.setenv("LANCEDB_DIR", str(tmp_path / "db")) + conn = get_vector_store_raw_connection() + ensure_active_generations_table(conn) + names = _schema_field_names(conn, "active_generations") + assert { + "collection", + "doc_id", + "parse_hash", + "user_id", + "model_tag", + "generation_id", + "config_hash", + "created_at", + "updated_at", + "published_at", + "operator", + }.issubset(names) + + +def test_chunks_table_adds_generation_id_column_for_legacy_schema( + tmp_path: Path, monkeypatch +) -> None: + """Existing chunks tables without generation_id receive a nullable column.""" + monkeypatch.setenv("LANCEDB_DIR", str(tmp_path / "db")) + conn = get_vector_store_raw_connection() + legacy_schema = pa.schema( + [ + pa.field("collection", pa.string()), + pa.field("doc_id", pa.string()), + pa.field("parse_hash", pa.string()), + pa.field("chunk_id", pa.string()), + pa.field("index", pa.int32()), + pa.field("text", pa.large_string()), + pa.field("page_number", pa.int32()), + pa.field("section", pa.string()), + pa.field("anchor", pa.string()), + pa.field("json_path", pa.string()), + pa.field("chunk_hash", pa.string()), + pa.field("config_hash", pa.string()), + pa.field("created_at", pa.timestamp("us")), + pa.field("metadata", pa.string()), + pa.field("user_id", pa.int64()), + ] + ) + conn.create_table("chunks", schema=legacy_schema) + ensure_chunks_table(conn) + assert "generation_id" in _schema_field_names(conn, "chunks") diff --git a/tests/core/tools/core/RAG_tools/storage/test_active_generation_store.py b/tests/core/tools/core/RAG_tools/storage/test_active_generation_store.py new file mode 100644 index 000000000..e26baf5da --- /dev/null +++ b/tests/core/tools/core/RAG_tools/storage/test_active_generation_store.py @@ -0,0 +1,305 @@ +"""Tests for ActiveGenerationStore and generation schema columns (issue #438 PR1).""" + +from __future__ import annotations + +import pytest + +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_active_generations_table, + ensure_chunks_table, + ensure_embeddings_table, +) +from xagent.core.tools.core.RAG_tools.storage import get_vector_store_raw_connection +from xagent.core.tools.core.RAG_tools.storage.factory import ( + get_active_generation_store, + reset_rag_storage_for_tests, +) +from xagent.core.tools.core.RAG_tools.storage.lancedb_stores import ( + LanceDBActiveGenerationStore, + LanceDBVectorIndexStore, +) + + +@pytest.fixture +def lancedb_conn(tmp_path, monkeypatch): + """Isolated LanceDB connection for integration tests.""" + monkeypatch.setenv("LANCEDB_DIR", str(tmp_path / "lancedb")) + reset_rag_storage_for_tests() + conn = get_vector_store_raw_connection() + yield conn + reset_rag_storage_for_tests() + + +def test_get_active_generation_store_returns_lancedb_impl() -> None: + store = get_active_generation_store() + assert store.__class__.__name__ == "LanceDBActiveGenerationStore" + + +def test_active_generation_publish_and_get_roundtrip(lancedb_conn) -> None: + ensure_active_generations_table(lancedb_conn) + store = get_active_generation_store() + + store.publish_active_generation( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=42, + model_tag="text_embedding_v4", + generation_id="gen_new", + config_hash="cfg_1", + operator="test", + ) + + pointer = store.get_active_generation( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=42, + model_tag="text_embedding_v4", + ) + assert pointer is not None + assert pointer["generation_id"] == "gen_new" + assert pointer["config_hash"] == "cfg_1" + assert pointer["user_id"] == 42 + + +def test_active_generation_scoped_by_user_id(lancedb_conn) -> None: + ensure_active_generations_table(lancedb_conn) + store = get_active_generation_store() + + store.publish_active_generation( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=1, + model_tag="", + generation_id="gen_user_1", + config_hash="cfg_1", + ) + store.publish_active_generation( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=2, + model_tag="", + generation_id="gen_user_2", + config_hash="cfg_2", + ) + + p1 = store.get_active_generation( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=1, + model_tag="", + ) + p2 = store.get_active_generation( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=2, + model_tag="", + ) + assert p1 is not None and p1["generation_id"] == "gen_user_1" + assert p2 is not None and p2["generation_id"] == "gen_user_2" + + +def test_active_generation_legacy_none_user_scope_republishes_one_row( + lancedb_conn, +) -> None: + """Legacy user_id=None scope should upsert one pointer and remain queryable.""" + ensure_active_generations_table(lancedb_conn) + store = get_active_generation_store() + + scope = dict( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=None, + model_tag="text_embedding_v4", + ) + + store.publish_active_generation( + **scope, + generation_id="gen_v1", + config_hash="cfg_v1", + operator="test", + ) + store.publish_active_generation( + **scope, + generation_id="gen_v2", + config_hash="cfg_v2", + operator="test", + ) + store.publish_active_generation( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=42, + model_tag="text_embedding_v4", + generation_id="gen_user_42", + config_hash="cfg_user_42", + operator="test", + ) + + pointer = store.get_active_generation(**scope) + assert pointer is not None + assert pointer["user_id"] is None + assert pointer["generation_id"] == "gen_v2" + assert pointer["config_hash"] == "cfg_v2" + + rows = store.list_active_generations(collection="col_a", user_id=None) + assert len(rows) == 1 + assert rows[0]["user_id"] is None + assert rows[0]["generation_id"] == "gen_v2" + + +def test_active_generation_republish_preserves_created_at(lancedb_conn) -> None: + """Republishing the same scope must keep created_at while bumping the rest.""" + ensure_active_generations_table(lancedb_conn) + store = get_active_generation_store() + + scope = dict( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=7, + model_tag="text_embedding_v4", + ) + + store.publish_active_generation( + **scope, + generation_id="gen_v1", + config_hash="cfg_v1", + operator="test", + ) + first = store.get_active_generation(**scope) + assert first is not None + original_created_at = first["created_at"] + + store.publish_active_generation( + **scope, + generation_id="gen_v2", + config_hash="cfg_v2", + operator="test", + ) + second = store.get_active_generation(**scope) + + assert second is not None + assert second["generation_id"] == "gen_v2" + assert second["config_hash"] == "cfg_v2" + # created_at must be carried over from the first publish even though + # publish does not perform a prior get_active_generation() round-trip. + assert second["created_at"] == original_created_at + assert second["updated_at"] >= original_created_at + assert second["published_at"] >= original_created_at + + +def test_active_generation_ensure_table_runs_once(monkeypatch, lancedb_conn) -> None: + """Publish/get hot paths must reuse the cached ensure_table result.""" + ensure_active_generations_table(lancedb_conn) + + from xagent.core.tools.core.RAG_tools.LanceDB import schema_manager as sm + + call_count = {"n": 0} + real_ensure = sm.ensure_active_generations_table + + def counting_ensure(conn): + call_count["n"] += 1 + return real_ensure(conn) + + monkeypatch.setattr(sm, "ensure_active_generations_table", counting_ensure) + + store = LanceDBActiveGenerationStore() + scope = dict( + collection="col_a", + doc_id="doc_1", + parse_hash="parse_abc", + user_id=11, + model_tag="", + ) + store.publish_active_generation( + **scope, + generation_id="gen_1", + config_hash="cfg", + ) + store.publish_active_generation( + **scope, + generation_id="gen_2", + config_hash="cfg", + ) + store.get_active_generation(**scope) + store.list_active_generations(collection="col_a") + + assert call_count["n"] == 1 + + +def test_chunks_table_persists_generation_id_column(lancedb_conn) -> None: + """PR1: schema accepts generation_id; merge-key isolation is covered in PR3.""" + ensure_chunks_table(lancedb_conn) + vector_store = LanceDBVectorIndexStore() + base = { + "collection": "col_a", + "doc_id": "doc_1", + "parse_hash": "parse_abc", + "index": 0, + "text": "hello", + "config_hash": "cfg_1", + "chunk_hash": "chash", + "created_at": None, + "metadata": None, + "user_id": 7, + } + vector_store.upsert_chunks( + [ + {**base, "chunk_id": "chunk_a", "generation_id": "gen_a"}, + {**base, "chunk_id": "chunk_b", "generation_id": "gen_b"}, + ] + ) + table = lancedb_conn.open_table("chunks") + try: + rows = table.search().where("doc_id = 'doc_1'").to_arrow().to_pylist() + finally: + if hasattr(table, "close"): + table.close() + generation_ids = {row["generation_id"] for row in rows} + assert generation_ids == {"gen_a", "gen_b"} + + +def test_embeddings_tables_accept_generation_id_and_config_hash( + lancedb_conn, +) -> None: + model_tag = to_model_tag("test-model") + ensure_embeddings_table(lancedb_conn, model_tag, vector_dim=2) + vector_store = LanceDBVectorIndexStore() + vector_store.upsert_embeddings( + "test-model", + [ + { + "collection": "col_a", + "doc_id": "doc_1", + "parse_hash": "parse_abc", + "generation_id": "gen_1", + "chunk_id": "chunk_1", + "model": "test-model", + "config_hash": "cfg_1", + "vector": [0.1, 0.2], + "vector_dimension": 2, + "text": "hello", + "chunk_hash": "chash", + "created_at": None, + "metadata": None, + "user_id": 3, + } + ], + ) + table_name = f"embeddings_{model_tag}" + table = lancedb_conn.open_table(table_name) + try: + row = table.search().where("chunk_id = 'chunk_1'").to_arrow().to_pylist()[0] + finally: + if hasattr(table, "close"): + table.close() + assert row["generation_id"] == "gen_1" + assert row["config_hash"] == "cfg_1"