diff --git a/migrations/versions/a7c3e9f1b2d4_add_document_sources_table.py b/migrations/versions/a7c3e9f1b2d4_add_document_sources_table.py new file mode 100644 index 00000000..09971e77 --- /dev/null +++ b/migrations/versions/a7c3e9f1b2d4_add_document_sources_table.py @@ -0,0 +1,129 @@ +"""add document_sources table, backfill from JSONB source_ids + +Normalize reasoning-tree linkage into a document_sources edge table. +Backfills from both the documents.source_ids column and the legacy +internal_metadata->'source_ids' location, dropping entries that are not +well-formed 21-char nanoids (they never resolved to documents anyway). + +The old source_ids column is kept (unwritten) for one release as a +rollback net; a follow-up migration drops it. + +Revision ID: a7c3e9f1b2d4 +Revises: e4eba9cfaa6f +Create Date: 2026-08-10 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +from migrations.utils import get_schema, index_exists, table_exists + +# revision identifiers, used by Alembic. +revision: str = "a7c3e9f1b2d4" +down_revision: str | None = "e4eba9cfaa6f" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +schema = get_schema() + + +def upgrade() -> None: + connection = op.get_bind() + inspector = sa.inspect(connection) + + if not table_exists("document_sources", inspector): + op.create_table( + "document_sources", + sa.Column("derived_id", sa.TEXT, nullable=False), + sa.Column("source_id", sa.TEXT, nullable=False), + sa.Column("position", sa.Integer, nullable=False, server_default="0"), + sa.Column("workspace_name", sa.TEXT, nullable=False), + sa.PrimaryKeyConstraint("derived_id", "source_id"), + sa.ForeignKeyConstraint( + ["derived_id"], + [f"{schema}.documents.id"], + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["workspace_name"], + [f"{schema}.workspaces.name"], + ), + sa.CheckConstraint("length(source_id) = 21", name="source_id_length"), + sa.CheckConstraint( + "source_id ~ '^[A-Za-z0-9_-]+$'", name="source_id_format" + ), + schema=schema, + ) + op.create_index( + "ix_document_sources_source_id", + "document_sources", + ["source_id", "workspace_name"], + schema=schema, + ) + + # Backfill: column takes precedence over legacy internal_metadata + # (mirrors the old resolved_source_ids property). Malformed entries + # are dropped; DISTINCT ON dedupes repeated IDs within one document. + op.execute(f""" + INSERT INTO {schema}.document_sources + (derived_id, source_id, position, workspace_name) + SELECT DISTINCT ON (d.id, s.value) + d.id, s.value, s.ord - 1, d.workspace_name + FROM {schema}.documents d + CROSS JOIN LATERAL jsonb_array_elements_text( + CASE + WHEN jsonb_typeof(d.source_ids) = 'array' + THEN d.source_ids + WHEN jsonb_typeof(d.internal_metadata->'source_ids') = 'array' + THEN d.internal_metadata->'source_ids' + WHEN jsonb_typeof(d.internal_metadata->'premise_ids') = 'array' + THEN d.internal_metadata->'premise_ids' + ELSE '[]'::jsonb + END + ) WITH ORDINALITY AS s(value, ord) + WHERE s.value ~ '^[A-Za-z0-9_-]{{21}}$' + ORDER BY d.id, s.value, s.ord + ON CONFLICT DO NOTHING + """) + + # The ORM no longer queries source_ids; drop its GIN index. The column + # itself stays for one release as a rollback net. + if index_exists("documents", "ix_documents_source_ids_gin", inspector): + op.drop_index( + "ix_documents_source_ids_gin", + table_name="documents", + schema=schema, + ) + + +def downgrade() -> None: + connection = op.get_bind() + inspector = sa.inspect(connection) + + # Reassemble JSONB arrays for rows written after the upgrade (the + # retained column still holds pre-upgrade data for older rows). + op.execute(f""" + UPDATE {schema}.documents d + SET source_ids = links.ids + FROM ( + SELECT derived_id, jsonb_agg(source_id ORDER BY position) AS ids + FROM {schema}.document_sources + GROUP BY derived_id + ) links + WHERE d.id = links.derived_id + """) + + if not index_exists("documents", "ix_documents_source_ids_gin", inspector): + op.create_index( + "ix_documents_source_ids_gin", + "documents", + ["source_ids"], + postgresql_using="gin", + schema=schema, + ) + + if table_exists("document_sources", inspector): + op.drop_table("document_sources", schema=schema) diff --git a/plans/source-ids-table.md b/plans/source-ids-table.md new file mode 100644 index 00000000..c9628dfa --- /dev/null +++ b/plans/source-ids-table.md @@ -0,0 +1,66 @@ +# Plan: normalize `source_ids` into a `document_sources` table + +Branch: `abigail/source-ids-table` (on top of `abigail/dev-2219`) + +## Goal + +Replace the JSONB `documents.source_ids` column with a proper edge table so +reasoning-tree linkage is queryable, indexed, and has one source of truth. +API shape is unchanged: `Conclusion.source_ids` stays `list[str] | None`. + +## Schema + +``` +document_sources + derived_id TEXT PK, FK -> documents.id ON DELETE CASCADE + source_id TEXT PK (NOT an FK: dreamer emits unresolvable IDs) + position INT (preserves premise order) + workspace_name TEXT FK -> workspaces.name + INDEX (source_id, workspace_name) -- reverse traversal, replaces GIN + CHECK length(source_id)=21, nanoid format +``` + +## Changes + +1. **models.py** — add `DocumentSource`; on `Document` drop the JSONB column + + GIN index, add `source_links` relationship (`lazy="selectin"` — required, + async lazy-load raises) and `source_ids` / `resolved_source_ids` properties. +2. **Migration** (one revision after `e4eba9cfaa6f`) — create table; backfill + from `source_ids` column AND legacy `internal_metadata->'source_ids'`, + dropping entries that don't match the nanoid regex; drop GIN index. + The old column is KEPT (unwritten) for one release as a rollback net — + follow-up migration drops it. +3. **crud/document.py** — insert sites build `DocumentSource` rows via + `build_source_links()` helper (dedupes + drops malformed IDs); + `get_child_observations` becomes a join instead of JSONB containment. +4. **utils/filter.py** — remove `source_ids` from `JSONB_COLUMNS`; special-case + it (and new alias `parent_id`) to an EXISTS subquery. Semantics parity: + scalar = membership, list = all present, `contains` = membership, + `in` = any present. Existing filter tests are the spec. +5. **utils/representation.py, utils/agent_tools.py** — delete + `internal_metadata.get("source_ids")` fallbacks (backfill retires them). + +## Decision points + +- **`/derived` endpoint**: `POST /conclusions/list` with + `{"filters": {"parent_id": ""}}` now covers it. Options: + (a) keep both, (b) drop `/derived` before it ships in a release. + Leaning (b) — one less route, avoids the `{conclusion_id}/derived` + path-capture footgun. SDK `derived()` helpers can wrap the filter. +- **Garbage IDs**: backfill and write path silently drop malformed entries. + They were already invisible to traversal; dry-run the backfill count on + real data before merging. +- **Old column retention**: kept unmapped for one release (see #2). + +## Costs + +- Every Document query gains one batched selectin SELECT. +- Write amplification: N link rows per conclusion. +- Backfill migration over all deductive/inductive documents. + +## Test plan + +- Existing `tests/routes/test_conclusions.py` filter + `/derived` tests pass + unchanged (parity spec). +- New: migration backfill test (pattern: `test_f1a2b3c4d5e6`), link + dedupe/malformed-drop unit tests, `parent_id` filter tests. diff --git a/src/crud/document.py b/src/crud/document.py index a7afee2e..02435e8c 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -1,4 +1,5 @@ import datetime +import re from collections.abc import Sequence from dataclasses import dataclass, field from enum import Enum @@ -33,6 +34,26 @@ from src.vector_store import ( logger = getLogger(__name__) +_SOURCE_ID_RE = re.compile(r"^[A-Za-z0-9_-]{21}$") + + +def build_source_links( + source_ids: list[str] | None, workspace_name: str +) -> list[models.DocumentSource]: + """Convert an LLM-provided source_ids list to DocumentSource rows. + + Dedupes (the PK is (derived_id, source_id)) and drops entries that are + not shaped like document IDs — the model occasionally emits timestamps + or numeric refs under schema pressure. + """ + if not source_ids: + return [] + return [ + models.DocumentSource(source_id=sid, position=i, workspace_name=workspace_name) + for i, sid in enumerate(dict.fromkeys(source_ids)) + if _SOURCE_ID_RE.match(sid) + ] + def get_all_documents( workspace_name: str, @@ -644,8 +665,7 @@ async def create_documents( internal_metadata=metadata_dict, session_name=doc.session_name, embedding=doc.embedding, - # Tree linkage column - source_ids=doc.source_ids, + source_links=build_source_links(doc.source_ids, workspace_name), ) else: new_doc = models.Document( @@ -657,8 +677,7 @@ async def create_documents( times_derived=doc.times_derived, internal_metadata=metadata_dict, session_name=doc.session_name, - # Tree linkage column - source_ids=doc.source_ids, + source_links=build_source_links(doc.source_ids, workspace_name), ) if doc.embedding: @@ -1365,10 +1384,9 @@ def get_child_observations( """ Get all observations that have this document as a source/premise. - Useful for traversing the reasoning tree upward (source -> derived observations). - Uses GIN index on source_ids for efficient lookups. Only matches linkage - stored in the source_ids column; legacy documents whose source_ids live in - internal_metadata are not found. + Useful for traversing the reasoning tree upward (source -> derived + observations). Joins through document_sources, which the backfill + migration populated from both current and legacy linkage storage. Args: workspace_name: Workspace identifier @@ -1381,13 +1399,17 @@ def get_child_observations( Select query for documents that reference this document as a source, for pagination support via apaginate() """ - # Find documents where source_ids contains the parent_id. The explicit - # jsonb_build_array keeps the bind typed as JSONB — SQLAlchemy would - # otherwise coerce the value to VARCHAR, which the @> operator rejects. - stmt = select(models.Document).where( - models.Document.workspace_name == workspace_name, - models.Document.source_ids.contains(func.jsonb_build_array(parent_id)), - models.Document.deleted_at.is_(None), + stmt = ( + select(models.Document) + .join( + models.DocumentSource, + models.DocumentSource.derived_id == models.Document.id, + ) + .where( + models.Document.workspace_name == workspace_name, + models.DocumentSource.source_id == parent_id, + models.Document.deleted_at.is_(None), + ) ) if observer: stmt = stmt.where(models.Document.observer == observer) diff --git a/src/models.py b/src/models.py index d46f1556..099147f3 100644 --- a/src/models.py +++ b/src/models.py @@ -390,9 +390,6 @@ class Document(Base): Integer, nullable=False, server_default=text("1") ) embedding: MappedColumn[Any] = mapped_column(Vector(_VECTOR_DIM), nullable=True) - source_ids: Mapped[list[str] | None] = mapped_column( - JSONB, nullable=True, server_default=text("NULL") - ) created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) @@ -420,10 +417,38 @@ class Document(Base): collection = relationship("Collection", back_populates="documents") + # selectin (not lazy) is required: async lazy-loads raise MissingGreenlet. + source_links: Mapped[list["DocumentSource"]] = relationship( + "DocumentSource", + order_by="DocumentSource.position", + cascade="all, delete-orphan", + passive_deletes=True, + lazy="selectin", + ) + + @property + def source_ids(self) -> list[str] | None: + """Parent conclusion IDs in original order; None when unlinked.""" + return [link.source_id for link in self.source_links] or None + + @source_ids.setter + def source_ids(self, value: list[str] | None) -> None: + # Constructor convenience: kwargs apply in order, so workspace_name + # must be passed before source_ids. Dedupes (PK is (derived_id, + # source_id)); crud.build_source_links additionally drops malformed + # IDs from LLM output before they reach this point. + self.source_links = [ + DocumentSource( + source_id=sid, position=i, workspace_name=self.workspace_name + ) + for i, sid in enumerate(dict.fromkeys(value or [])) + ] + @property def resolved_source_ids(self) -> list[str] | None: - """Source IDs, falling back to legacy internal_metadata storage.""" - return self.source_ids or (self.internal_metadata or {}).get("source_ids") + """Kept for API schema compatibility; legacy internal_metadata + storage is retired by the document_sources backfill migration.""" + return self.source_ids __table_args__ = ( CheckConstraint("length(id) = 21", name="id_length"), @@ -463,12 +488,6 @@ class Document(Base): "embedding": "vector_cosine_ops" }, # Cosine distance operator ), - # GIN index for efficient tree traversal (finding children by source IDs) - Index( - "ix_documents_source_ids_gin", - "source_ids", - postgresql_using="gin", - ), # Composite index for efficient reconciliation queries Index( "ix_documents_sync_state_last_sync_at", @@ -478,6 +497,33 @@ class Document(Base): ) +@final +class DocumentSource(Base): + """One reasoning-tree edge: derived_id was concluded from source_id.""" + + __tablename__: str = "document_sources" + + derived_id: Mapped[str] = mapped_column( + ForeignKey("documents.id", ondelete="CASCADE"), primary_key=True + ) + # Deliberately not an FK: the dreamer can emit IDs that never resolve, + # and sources may be deleted independently of their children. + source_id: Mapped[str] = mapped_column(TEXT, primary_key=True) + position: Mapped[int] = mapped_column( + Integer, nullable=False, server_default=text("0") + ) + workspace_name: Mapped[str] = mapped_column( + ForeignKey("workspaces.name"), nullable=False + ) + + __table_args__ = ( + # Reverse traversal ("who derived from me?") — replaces the old GIN index + Index("ix_document_sources_source_id", "source_id", "workspace_name"), + CheckConstraint("length(source_id) = 21", name="source_id_length"), + CheckConstraint("source_id ~ '^[A-Za-z0-9_-]+$'", name="source_id_format"), + ) + + @final class QueueItem(Base): __tablename__: str = "queue" diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index de011636..175cc6f2 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -971,12 +971,11 @@ async def create_observations( continue # Build metadata with level-specific fields + # source_ids intentionally omitted from metadata: linkage lives in + # the document_sources table (via DocumentCreate.source_ids below). metadata = schemas.DocumentMetadata( message_ids=message_ids, message_created_at=message_created_at, - source_ids=obs.source_ids - if obs.level in ("deductive", "inductive", "contradiction") - else None, premises=obs.premises if obs.level == "deductive" else None, sources=obs.sources if obs.level in ("inductive", "contradiction") diff --git a/src/utils/filter.py b/src/utils/filter.py index 995fd769..bc8fd61c 100644 --- a/src/utils/filter.py +++ b/src/utils/filter.py @@ -5,7 +5,17 @@ from logging import getLogger from typing import Any, TypeVar from typing import cast as typing_cast -from sqlalchemy import ColumnElement, Select, and_, case, cast, literal, not_, or_ +from sqlalchemy import ( + ColumnElement, + Select, + and_, + case, + cast, + literal, + not_, + or_, + select, +) from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.types import Numeric @@ -33,7 +43,7 @@ NUMERIC_OPERATORS = {"gte", "lte", "gt", "lt", "ne"} # JSONB columns keep containment semantics: bare lists are not membership # sugar, and dict values map to nested-metadata conditions rather than IN/Eq. -JSONB_COLUMNS = ("h_metadata", "configuration", "internal_metadata", "source_ids") +JSONB_COLUMNS = ("h_metadata", "configuration", "internal_metadata") ALLOWED_EXTERNAL_TO_INTERNAL_COLUMN_MAPPING = { "id": "name", @@ -62,6 +72,7 @@ ALLOWED_EXTERNAL_TO_INTERNAL_COLUMN_MAPPING_DOCUMENTS = { "observed_id": "observed", "level": "level", "source_ids": "source_ids", + "parent_id": "source_ids", "times_derived": "times_derived", "metadata": "internal_metadata", } @@ -300,6 +311,64 @@ def _build_filter_conditions( return and_(*conditions) +def _build_source_ids_condition( + value: Any, model_class: type[Any] +) -> ColumnElement[bool] | None: + """Filter documents by reasoning-tree linkage via document_sources. + + Preserves the old JSONB containment semantics: a scalar matches + membership, a bare list requires ALL entries present, {"contains": x} + matches membership, {"in": [...]} matches any entry present. + """ + from ..models import DocumentSource + + def _member(sid: Any) -> ColumnElement[bool]: + if not isinstance(sid, str) or not sid: + raise FilterError("source_ids filter entries must be non-empty strings") + return ( + select(literal(1)) + .where( + DocumentSource.derived_id == model_class.id, + DocumentSource.source_id == sid, + ) + .exists() + ) + + if value == "*": + return None + if isinstance(value, str): + return _member(value) + if isinstance(value, list | tuple | set): + entries = list(typing_cast(Sequence[Any], value)) + if "*" in entries: + return None + return _combine_conditions_with_and([_member(v) for v in entries]) + if isinstance(value, dict): + conditions: list[ColumnElement[bool]] = [] + for operator, op_value in typing_cast("dict[str, Any]", value).items(): + if op_value == "*": + continue + if operator == "contains": + conditions.append(_member(op_value)) + elif operator == "in": + if not isinstance(op_value, list | tuple | set): + raise FilterError( + f"Invalid value for 'in' operator: {op_value}. Expected an iterable" + ) + in_entries = list(typing_cast(Sequence[Any], op_value)) + if "*" in in_entries: + continue + members = [_member(v) for v in in_entries] + if members: + conditions.append(or_(*members)) + else: + raise FilterError( + f"Operator '{operator}' is not supported on source_ids" + ) + return _combine_conditions_with_and(conditions) + raise FilterError(f"Invalid source_ids filter value: {value}") + + def _build_field_condition( key: str, value: Any, model_class: type[Any] ) -> ColumnElement[bool] | None: @@ -336,6 +405,11 @@ def _build_field_condition( f"Column '{key}' is not allowed to be filtered on or does not exist on {model_class.__name__}" ) + # Reasoning-tree linkage lives in the document_sources table, not a + # column; translate to EXISTS subqueries before column resolution. + if model_class.__name__ == "Document" and column_name == "source_ids": + return _build_source_ids_condition(value, model_class) + # Check if the column exists on the model if not hasattr(model_class, column_name): raise FilterError(f"Column '{key}' does not exist on {model_class.__name__}") diff --git a/src/utils/representation.py b/src/utils/representation.py index 01e4b70a..6c059251 100644 --- a/src/utils/representation.py +++ b/src/utils/representation.py @@ -637,9 +637,7 @@ class Representation(BaseModel): doc.internal_metadata.get("message_ids", []) ), session_name=doc.session_name, - # Support both top-level and metadata locations for backward compatibility - source_ids=doc.source_ids - or doc.internal_metadata.get("premise_ids", []), + source_ids=doc.source_ids or [], premises=doc.internal_metadata.get("premises", []), ) for doc in documents @@ -654,9 +652,7 @@ class Representation(BaseModel): conclusion=doc.content, message_ids=doc.internal_metadata.get("message_ids", []), session_name=doc.session_name, - # Support both top-level and metadata locations for backward compatibility - source_ids=doc.source_ids - or doc.internal_metadata.get("source_ids", []), + source_ids=doc.source_ids or [], sources=doc.internal_metadata.get("sources", []), pattern_type=doc.internal_metadata.get("pattern_type", "pattern"), confidence=doc.internal_metadata.get("confidence", "medium"), @@ -673,9 +669,7 @@ class Representation(BaseModel): content=doc.content, message_ids=doc.internal_metadata.get("message_ids", []), session_name=doc.session_name, - # Support both top-level and metadata locations for backward compatibility - source_ids=doc.source_ids - or doc.internal_metadata.get("source_ids", []), + source_ids=doc.source_ids or [], sources=doc.internal_metadata.get("sources", []), ) for doc in documents diff --git a/tests/alembic/revisions/__init__.py b/tests/alembic/revisions/__init__.py index 77c15483..10ce184c 100644 --- a/tests/alembic/revisions/__init__.py +++ b/tests/alembic/revisions/__init__.py @@ -16,6 +16,7 @@ from . import ( test_564ba40505c5_add_session_name_column_to_documents, test_917195d9b5e9_add_messageembedding_table, test_a1b2c3d4e5f6_initial_schema, + test_a7c3e9f1b2d4_add_document_sources_table, test_b765d82110bd_change_metamessages_to_user_level_with_, test_b8183c5ffb48_codify_document_level_and_times_derived, test_baa22cad81e2_standardize_constraint_names, @@ -44,6 +45,7 @@ __all__ = [ "test_88b0fb10906f_add_webhooks_table", "test_917195d9b5e9_add_messageembedding_table", "test_a1b2c3d4e5f6_initial_schema", + "test_a7c3e9f1b2d4_add_document_sources_table", "test_b765d82110bd_change_metamessages_to_user_level_with_", "test_b8183c5ffb48_codify_document_level_and_times_derived", "test_baa22cad81e2_standardize_constraint_names", diff --git a/tests/alembic/revisions/test_a7c3e9f1b2d4_add_document_sources_table.py b/tests/alembic/revisions/test_a7c3e9f1b2d4_add_document_sources_table.py new file mode 100644 index 00000000..bdefb323 --- /dev/null +++ b/tests/alembic/revisions/test_a7c3e9f1b2d4_add_document_sources_table.py @@ -0,0 +1,105 @@ +"""Hooks for revision a7c3e9f1b2d4 (document_sources table + backfill).""" + +from __future__ import annotations + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +WORKSPACE_NAME = "sources_ws" +PEER_NAME = "sources_peer" + +# Documents covering every legacy linkage location plus malformed entries +DOC_COLUMN = generate_nanoid() # linkage in source_ids column +DOC_META = generate_nanoid() # linkage in internal_metadata.source_ids +DOC_PREMISE = generate_nanoid() # linkage in internal_metadata.premise_ids +DOC_GARBAGE = generate_nanoid() # column mixes valid + malformed entries + +SRC_A = generate_nanoid() +SRC_B = generate_nanoid() +SRC_C = generate_nanoid() +SRC_D = generate_nanoid() +SRC_E = generate_nanoid() + + +@register_before_upgrade("a7c3e9f1b2d4") +def prepare_document_sources(verifier: MigrationVerifier) -> None: + verifier.assert_table_exists("document_sources", exists=False) + + schema = verifier.schema + connection = verifier.conn + + connection.execute( + text(f'INSERT INTO "{schema}"."workspaces" ("id", "name") VALUES (:id, :n)'), + {"id": generate_nanoid(), "n": WORKSPACE_NAME}, + ) + connection.execute( + text( + f"""INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") + VALUES (:id, :n, :w)""" + ), + {"id": generate_nanoid(), "n": PEER_NAME, "w": WORKSPACE_NAME}, + ) + connection.execute( + text( + f"""INSERT INTO "{schema}"."collections" + ("id", "workspace_name", "observer", "observed") + VALUES (:id, :w, :p, :p)""" + ), + {"id": generate_nanoid(), "w": WORKSPACE_NAME, "p": PEER_NAME}, + ) + + def seed_doc(doc_id: str, source_ids: str | None, metadata: str) -> None: + connection.execute( + text( + f"""INSERT INTO "{schema}"."documents" + ("id", "workspace_name", "observer", "observed", "content", + "level", "source_ids", "internal_metadata") + VALUES (:id, :w, :p, :p, :c, 'deductive', + CAST(:s AS jsonb), CAST(:m AS jsonb))""" + ), + { + "id": doc_id, + "w": WORKSPACE_NAME, + "p": PEER_NAME, + "c": f"doc {doc_id}", + "s": source_ids, + "m": metadata, + }, + ) + + seed_doc(DOC_COLUMN, f'["{SRC_A}", "{SRC_B}"]', "{}") + seed_doc(DOC_META, None, f'{{"source_ids": ["{SRC_C}"]}}') + seed_doc(DOC_PREMISE, None, f'{{"premise_ids": ["{SRC_D}"]}}') + # Malformed entries (numeric ref, timestamp) must be dropped by backfill + seed_doc(DOC_GARBAGE, f'["{SRC_E}", "1234", "2024-01-01T00:00:00"]', "{}") + + +@register_after_upgrade("a7c3e9f1b2d4") +def verify_document_sources(verifier: MigrationVerifier) -> None: + verifier.assert_table_exists("document_sources") + verifier.assert_indexes_exist( + [("document_sources", "ix_document_sources_source_id")] + ) + verifier.assert_indexes_not_exist([("documents", "ix_documents_source_ids_gin")]) + + schema = verifier.schema + connection = verifier.conn + + def edges(doc_id: str) -> list[str]: + rows = connection.execute( + text( + f"""SELECT "source_id" FROM "{schema}"."document_sources" + WHERE "derived_id" = :d ORDER BY "position" + """ + ), + {"d": doc_id}, + ).all() + return [r.source_id for r in rows] + + assert edges(DOC_COLUMN) == [SRC_A, SRC_B] # column backfilled, order kept + assert edges(DOC_META) == [SRC_C] # legacy internal_metadata.source_ids + assert edges(DOC_PREMISE) == [SRC_D] # legacy internal_metadata.premise_ids + assert edges(DOC_GARBAGE) == [SRC_E] # malformed entries dropped diff --git a/tests/routes/test_conclusions.py b/tests/routes/test_conclusions.py index 8a1f91be..f1d0fc38 100644 --- a/tests/routes/test_conclusions.py +++ b/tests/routes/test_conclusions.py @@ -741,7 +741,9 @@ class TestConclusionRoutes: db_session: AsyncSession, sample_data: tuple[Workspace, Peer], ): - """Test that legacy internal_metadata.source_ids is coalesced into source_ids""" + """Legacy internal_metadata.source_ids is NOT read at runtime anymore: + the document_sources backfill migration (a7c3e9f1b2d4) moves it into + the edge table instead. A metadata-only doc therefore surfaces None.""" test_workspace, test_peer = sample_data # Create another peer @@ -775,7 +777,7 @@ class TestConclusionRoutes: assert response.status_code == 200 conclusion = response.json() - assert conclusion["source_ids"] == ["legacy_id_1", "legacy_id_2"] + assert conclusion["source_ids"] is None @pytest.mark.asyncio async def test_get_conclusion_not_found( @@ -858,6 +860,8 @@ class TestConclusionRoutes: db_session, test_workspace.name, test_peer.name, test_peer2.name ) + # document_sources check-constrains source_id to nanoid shape + source_ids = [str(generate_nanoid()), str(generate_nanoid())] doc = models.Document( workspace_name=test_workspace.name, observer=test_peer.name, @@ -865,7 +869,7 @@ class TestConclusionRoutes: content="Derived conclusion", embedding=[0.1] * 1536, level="inductive", - source_ids=["src_1", "src_2"], + source_ids=source_ids, times_derived=2, ) db_session.add(doc) @@ -884,7 +888,7 @@ class TestConclusionRoutes: assert response.status_code == 200 items = response.json()["items"] assert len(items) == 1 - assert items[0]["source_ids"] == ["src_1", "src_2"] + assert items[0]["source_ids"] == source_ids assert items[0]["times_derived"] == 2 @pytest.mark.asyncio diff --git a/tests/utils/test_agent_tools.py b/tests/utils/test_agent_tools.py index 8005fd4f..34857ec4 100644 --- a/tests/utils/test_agent_tools.py +++ b/tests/utils/test_agent_tools.py @@ -223,13 +223,15 @@ class TestCreateObservations: """Dialectic context (no current_messages) forces observations to be deductive.""" ctx = make_tool_context(current_messages=None) + # source links are check-constrained to nanoid shape + premise_ids = [str(generate_nanoid()), str(generate_nanoid())] result = await _handle_create_observations( ctx, { "observations": [ { "content": "Inferred preference for quiet spaces", - "source_ids": ["premise1", "premise2"], + "source_ids": premise_ids, "premises": [ "User mentioned working in libraries", "User avoids noisy cafes", @@ -249,7 +251,7 @@ class TestCreateObservations: doc = (await db_session.execute(stmt)).scalar_one_or_none() assert doc is not None assert doc.level == "deductive" - assert doc.source_ids == ["premise1", "premise2"] + assert doc.source_ids == premise_ids async def test_non_deriver_context_rejects_explicit( self, @@ -290,13 +292,14 @@ class TestCreateObservations: the prefix must be stripped so provenance links reference real IDs.""" ctx = make_tool_context(current_messages=None) + premise_ids = [str(generate_nanoid()), str(generate_nanoid())] result = await _handle_create_observations( ctx, { "observations": [ { "content": "Inferred preference for early mornings", - "source_ids": ["id:premise1", "ID:premise2"], + "source_ids": [f"id:{premise_ids[0]}", f"ID:{premise_ids[1]}"], "premises": [ "User schedules meetings before 9am", "User mentions waking at 5:30", @@ -313,7 +316,7 @@ class TestCreateObservations: ) doc = (await db_session.execute(stmt)).scalar_one_or_none() assert doc is not None - assert doc.source_ids == ["premise1", "premise2"] + assert doc.source_ids == premise_ids async def test_empty_observations_list_returns_error( self, make_tool_context: Callable[..., ToolContext]