From d7bdcc3bc1418687d4a32241effa23ef0bb47a2f Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Fri, 7 Nov 2025 13:22:24 -0500 Subject: [PATCH] feat: codify queue columns (#254) * feat: codify queue columns * fix: batch with python loop control * fix: cleanup merge * fix: down revision * fix: batch delete in migration * feat: only run alembic tests for changed migration / test (#264) * feat: only run alembic tests for changed migration / test * fix: Run full test suite if alembic testing infra changes * feat: codify times_derived + level on Document (#260) * feat: codify times_derived + level on Document * fix: CR comment * fix: CodeRabbit comments * fix: batch migrations; move types; remove fields from payload * fix: rm duplicate table args * fix: add messages.id FK --- .pre-commit-config.yaml | 6 +- ...codify_document_level_and_times_derived.py | 164 +++++++++ ...odify_workspace_name_and_message_id_in_.py | 252 +++++++++++++ scripts/run_alembic_tests.py | 141 ++++++++ src/crud/document.py | 2 + src/crud/representation.py | 6 +- src/deriver/consumer.py | 55 ++- src/deriver/enqueue.py | 33 +- src/deriver/queue_manager.py | 20 +- src/dreamer/dream_scheduler.py | 17 +- src/dreamer/dreamer.py | 23 +- src/models.py | 38 +- src/schemas.py | 17 +- src/utils/queue_payload.py | 31 +- src/utils/representation.py | 4 +- src/utils/types.py | 2 + src/utils/work_unit.py | 6 +- src/webhooks/events.py | 13 +- src/webhooks/webhook_delivery.py | 10 +- tests/alembic/revisions/__init__.py | 4 + ...codify_document_level_and_times_derived.py | 315 ++++++++++++++++ ...odify_workspace_name_and_message_id_in_.py | 336 ++++++++++++++++++ tests/crud/test_workspace.py | 1 + tests/deriver/conftest.py | 42 ++- tests/deriver/test_deriver_processing.py | 18 +- tests/deriver/test_queue_processing.py | 101 ++++-- tests/integration/test_representation.py | 14 +- tests/routes/test_queue_status.py | 11 +- tests/test_schema_validations.py | 7 +- 29 files changed, 1503 insertions(+), 186 deletions(-) create mode 100644 migrations/versions/b8183c5ffb48_codify_document_level_and_times_derived.py create mode 100644 migrations/versions/ec8f94139b02_codify_workspace_name_and_message_id_in_.py create mode 100755 scripts/run_alembic_tests.py create mode 100644 tests/alembic/revisions/test_b8183c5ffb48_codify_document_level_and_times_derived.py create mode 100644 tests/alembic/revisions/test_ec8f94139b02_codify_workspace_name_and_message_id_in_.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7cc0de63..6d5a3f6c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -74,11 +74,11 @@ repos: # Run Alembic tests only when migrations change - id: pytest-alembic name: pytest (alembic migrations) - entry: uv run pytest tests/alembic/ + entry: uv run python scripts/run_alembic_tests.py language: system - files: ^(migrations/.*\.py|tests/alembic/.*\.py)$ + files: ^(migrations/versions/.*\.py|tests/alembic/.*\.py)$ stages: [pre-push] - pass_filenames: false + pass_filenames: true require_serial: true # Ensure each alembic migration revision has a corresponding test file diff --git a/migrations/versions/b8183c5ffb48_codify_document_level_and_times_derived.py b/migrations/versions/b8183c5ffb48_codify_document_level_and_times_derived.py new file mode 100644 index 00000000..2e8c2a6b --- /dev/null +++ b/migrations/versions/b8183c5ffb48_codify_document_level_and_times_derived.py @@ -0,0 +1,164 @@ +"""codify_document_level_and_times_derived + +Revision ID: b8183c5ffb48 +Revises: ec8f94139b02 +Create Date: 2025-10-31 12:48:54.597269 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from sqlalchemy import text + +from migrations.utils import column_exists, constraint_exists, get_schema + +# revision identifiers, used by Alembic. +revision: str = "b8183c5ffb48" +down_revision: str | None = "ec8f94139b02" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +schema = get_schema() + + +def upgrade() -> None: + """Codify level and times_derived from internal_metadata into explicit columns.""" + connection = op.get_bind() + inspector = sa.inspect(connection) + + # Step 1: Add level column (nullable initially) + if not column_exists("documents", "level", inspector): + op.add_column( + "documents", + sa.Column( + "level", + sa.TEXT(), + nullable=True, + ), + schema=schema, + ) + + # Step 2: Add times_derived column (nullable initially) + if not column_exists("documents", "times_derived", inspector): + op.add_column( + "documents", + sa.Column("times_derived", sa.Integer(), nullable=True), + schema=schema, + ) + + # Step 3: Populate level and times_derived from internal_metadata in batches + # Default to 'explicit' for level and 1 for times_derived if not present in metadata + batch_size = 5000 + while True: + result = connection.execute( + text( + f""" + WITH batch AS ( + SELECT id + FROM {schema}.documents + WHERE level IS NULL OR times_derived IS NULL + LIMIT :batch_size + ) + UPDATE {schema}.documents d + SET + level = COALESCE( + d.internal_metadata->>'level', + 'explicit' + ), + times_derived = COALESCE( + (d.internal_metadata->>'times_derived')::integer, + 1 + ) + FROM batch + WHERE d.id = batch.id + """ + ), + {"batch_size": batch_size}, + ) + if result.rowcount == 0: + break + + # Step 4: Make level NOT NULL with server default + op.alter_column( + "documents", + "level", + nullable=False, + server_default=text("'explicit'"), + schema=schema, + ) + + # Step 5: Make times_derived NOT NULL with server default + op.alter_column( + "documents", + "times_derived", + nullable=False, + server_default=text("1"), + schema=schema, + ) + + # Step 6: Add CHECK constraint for level + if not constraint_exists("documents", "level_valid", "check", inspector): + op.create_check_constraint( + "level_valid", + "documents", + "level IN ('explicit', 'deductive')", + schema=schema, + ) + + +def downgrade() -> None: + """Restore level and times_derived to internal_metadata.""" + connection = op.get_bind() + inspector = sa.inspect(connection) + + # Step 1: Drop CHECK constraint for level + if constraint_exists("documents", "level_valid", "check", inspector): + op.drop_constraint( + "level_valid", + "documents", + type_="check", + schema=schema, + ) + + # Step 2: Copy level and times_derived back to internal_metadata in batches (optional, for safety) + batch_size = 5000 + while True: + result = connection.execute( + text( + f""" + WITH batch AS ( + SELECT id + FROM {schema}.documents + WHERE internal_metadata IS NULL + OR NOT (internal_metadata ? 'level') + OR NOT (internal_metadata ? 'times_derived') + LIMIT :batch_size + ) + UPDATE {schema}.documents d + SET internal_metadata = jsonb_set( + jsonb_set( + COALESCE(d.internal_metadata, '{{}}'::jsonb), + '{{level}}', + to_jsonb(d.level) + ), + '{{times_derived}}', + to_jsonb(d.times_derived) + ) + FROM batch + WHERE d.id = batch.id + """ + ), + {"batch_size": batch_size}, + ) + if result.rowcount == 0: + break + + # Step 3: Drop the level column + if column_exists("documents", "level", inspector): + op.drop_column("documents", "level", schema=schema) + + # Step 4: Drop the times_derived column + if column_exists("documents", "times_derived", inspector): + op.drop_column("documents", "times_derived", schema=schema) diff --git a/migrations/versions/ec8f94139b02_codify_workspace_name_and_message_id_in_.py b/migrations/versions/ec8f94139b02_codify_workspace_name_and_message_id_in_.py new file mode 100644 index 00000000..ae7210cb --- /dev/null +++ b/migrations/versions/ec8f94139b02_codify_workspace_name_and_message_id_in_.py @@ -0,0 +1,252 @@ +"""codify workspace_name and message_id in queue table + +Revision ID: ec8f94139b02 +Revises: e9b705f9adf9 +Create Date: 2025-10-28 17:39:51.778665 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +from migrations.utils import column_exists, fk_exists, get_schema, index_exists + +# revision identifiers, used by Alembic. +revision: str = "ec8f94139b02" +down_revision: str | None = "e9b705f9adf9" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +schema = get_schema() + + +def upgrade() -> None: + # Step 1: Drop rows where workspace_name is NULL in payload + # This removes invalid/corrupted queue items + conn = op.get_bind() + batch_size = 10000 + while True: + result = conn.execute( + sa.text( + f""" + DELETE FROM "{schema}".queue + WHERE id IN ( + SELECT id FROM "{schema}".queue + WHERE payload->>'workspace_name' IS NULL + LIMIT :batch_size + ) + """ + ), + {"batch_size": batch_size}, + ) + if result.rowcount == 0: + break + + # Step 2: Add workspace_name column (nullable initially for backfill) + op.add_column( + "queue", + sa.Column("workspace_name", sa.TEXT(), nullable=True), + schema=schema, + ) + + # Step 3: Add message_id column (nullable, as not all tasks have message_id) + op.add_column( + "queue", + sa.Column("message_id", sa.BigInteger(), nullable=True), + schema=schema, + ) + + # Step 4: Backfill workspace_name and message_id from payload in batches + batch_size = 5000 + while True: + result = conn.execute( + sa.text( + f""" + WITH batch AS ( + SELECT id + FROM "{schema}".queue + WHERE workspace_name IS NULL + OR (message_id IS NULL AND payload ? 'message_id' AND payload->>'message_id' IS NOT NULL) + LIMIT :batch_size + ) + UPDATE "{schema}".queue q + SET + workspace_name = COALESCE(q.workspace_name, q.payload->>'workspace_name'), + message_id = COALESCE( + q.message_id, + CASE + WHEN q.payload ? 'message_id' AND q.payload->>'message_id' IS NOT NULL + THEN (q.payload->>'message_id')::bigint + END + ) + FROM batch + WHERE q.id = batch.id + """ + ), + {"batch_size": batch_size}, + ) + if result.rowcount == 0: + break + + # Step 5: Remove workspace_name and message_id from JSONB payloads in batches + while True: + result = conn.execute( + sa.text( + f""" + WITH batch AS ( + SELECT id + FROM "{schema}".queue + WHERE payload ? 'workspace_name' OR payload ? 'message_id' + LIMIT :batch_size + ) + UPDATE "{schema}".queue q + SET payload = q.payload - 'workspace_name' - 'message_id' + FROM batch + WHERE q.id = batch.id + """ + ), + {"batch_size": batch_size}, + ) + if result.rowcount == 0: + break + + # Step 6: Make workspace_name non-nullable + op.alter_column("queue", "workspace_name", nullable=False, schema=schema) + + # Step 7: Add foreign key constraint on workspace_name -> workspaces.name + op.create_foreign_key( + "fk_queue_workspace_name", + "queue", + "workspaces", + ["workspace_name"], + ["name"], + source_schema=schema, + referent_schema=schema, + ) + + # Step 7b: Add foreign key constraint on message_id -> messages.id + op.create_foreign_key( + "fk_queue_message_id", + "queue", + "messages", + ["message_id"], + ["id"], + source_schema=schema, + referent_schema=schema, + ) + + # Step 8: Add index on workspace_name (for FK performance and filtering) + op.create_index( + op.f("ix_queue_workspace_name"), + "queue", + ["workspace_name"], + unique=False, + schema=schema, + ) + + # Step 9: Add partial index on message_id WHERE message_id IS NOT NULL + # This optimizes JOINs with the messages table + op.create_index( + "ix_queue_message_id_not_null", + "queue", + ["message_id"], + unique=False, + schema=schema, + postgresql_where=sa.text("message_id IS NOT NULL"), + ) + + # Step 10: Add composite index on (workspace_name, processed) + # This optimizes queries that filter unprocessed items by workspace + op.create_index( + "ix_queue_workspace_name_processed", + "queue", + ["workspace_name", "processed"], + unique=False, + schema=schema, + ) + + # Step 11: Add composite index on (work_unit_key, processed, id) + # This is critical for the hot path: "get next unprocessed item for this work unit" + # Covers: WHERE work_unit_key = ? AND NOT processed ORDER BY id + op.create_index( + "ix_queue_work_unit_key_processed_id", + "queue", + ["work_unit_key", "processed", "id"], + unique=False, + schema=schema, + ) + + +def downgrade() -> None: + inspector = sa.inspect(op.get_bind()) + + # Drop indexes + if index_exists("queue", "ix_queue_work_unit_key_processed_id", inspector): + op.drop_index( + "ix_queue_work_unit_key_processed_id", table_name="queue", schema=schema + ) + + if index_exists("queue", "ix_queue_workspace_name_processed", inspector): + op.drop_index( + "ix_queue_workspace_name_processed", table_name="queue", schema=schema + ) + + if index_exists("queue", "ix_queue_message_id_not_null", inspector): + op.drop_index("ix_queue_message_id_not_null", table_name="queue", schema=schema) + + if index_exists("queue", "ix_queue_workspace_name", inspector): + op.drop_index( + op.f("ix_queue_workspace_name"), table_name="queue", schema=schema + ) + + # Drop foreign key constraints + if fk_exists("queue", "fk_queue_message_id", inspector): + op.drop_constraint("fk_queue_message_id", "queue", schema=schema) + + if fk_exists("queue", "fk_queue_workspace_name", inspector): + op.drop_constraint("fk_queue_workspace_name", "queue", schema=schema) + + # Restore workspace_name and message_id to payload in batches + conn = op.get_bind() + batch_size = 5000 + + if column_exists("queue", "workspace_name", inspector) or column_exists( + "queue", "message_id", inspector + ): + while True: + result = conn.execute( + sa.text( + f""" + WITH batch AS ( + SELECT id + FROM "{schema}".queue + WHERE (workspace_name IS NOT NULL AND NOT (payload ? 'workspace_name')) + OR (message_id IS NOT NULL AND NOT (payload ? 'message_id')) + LIMIT :batch_size + ) + UPDATE "{schema}".queue q + SET payload = q.payload + || CASE WHEN q.workspace_name IS NOT NULL AND NOT (q.payload ? 'workspace_name') + THEN jsonb_build_object('workspace_name', q.workspace_name) + ELSE '{{}}'::jsonb END + || CASE WHEN q.message_id IS NOT NULL AND NOT (q.payload ? 'message_id') + THEN jsonb_build_object('message_id', q.message_id) + ELSE '{{}}'::jsonb END + FROM batch + WHERE q.id = batch.id + """ + ), + {"batch_size": batch_size}, + ) + if result.rowcount == 0: + break + + # Drop columns + if column_exists("queue", "message_id", inspector): + op.drop_column("queue", "message_id", schema=schema) + + if column_exists("queue", "workspace_name", inspector): + op.drop_column("queue", "workspace_name", schema=schema) diff --git a/scripts/run_alembic_tests.py b/scripts/run_alembic_tests.py new file mode 100755 index 00000000..8a2439e6 --- /dev/null +++ b/scripts/run_alembic_tests.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +""" +Run alembic tests selectively based on changed files. + +Note that this script is actively used within our precommit hooks and should not be removed. +If this script is moved, the corresponding precommit hook will need to be updated. + +This script determines which specific alembic tests to run based on: +1. If a test file changed, run the test for that revision +2. If a migration file changed, run the corresponding test + +The alembic test system uses a parameterized test in test_pipeline.py that runs +for each revision. We filter these tests using pytest's -k flag with the revision IDs. + +Usage: python scripts/run_alembic_tests.py ... +""" + +import re +import subprocess +import sys +from pathlib import Path + + +def extract_revision_id(filepath: Path) -> str | None: + """Extract the revision ID from a migration or test filename. + + Migration files: {revision_id}_{description}.py + Test files: test_{revision_id}_{description}.py + + Returns the revision_id (e.g., "05486ce795d5") or None if not found. + """ + filename = filepath.name + + # Remove .py extension + if not filename.endswith(".py"): + return None + + filename = filename[:-3] + + # Remove test_ prefix if present + if filename.startswith("test_"): + filename = filename[5:] + + # Extract revision ID (first part before underscore) + # Revision IDs are typically 12 characters of hex + match = re.match(r"^([a-f0-9]{12})_", filename) + if match: + return match.group(1) + + return None + + +def main(): + if len(sys.argv) < 2: + print("No files to check, skipping alembic tests") + sys.exit(0) + + changed_files = [Path(f) for f in sys.argv[1:]] + + # Paths + repo_root = Path(__file__).parent.parent + migrations_dir = repo_root / "migrations" / "versions" + tests_dir = repo_root / "tests" / "alembic" / "revisions" + alembic_tests_dir = repo_root / "tests" / "alembic" + + # Collect revision IDs to test + revision_ids: set[str] = set() + run_full_suite = False + + for filepath in changed_files: + filepath = Path(filepath).resolve() + + # Check if file is under tests/alembic (including subdirectories) + # If it's not a revision-specific test file, run full suite + if ( + filepath.parent == alembic_tests_dir + or alembic_tests_dir in filepath.parents + ) and not (filepath.parent == tests_dir and filepath.name.startswith("test_")): + run_full_suite = True + print( + f"Infrastructure file changed: {filepath.name} -> will run full test suite" + ) + continue + + # Case 1: Test file changed - extract its revision ID + if filepath.parent == tests_dir and filepath.name.startswith("test_"): + revision_id = extract_revision_id(filepath) + if revision_id: + revision_ids.add(revision_id) + print( + f"Test file changed: {filepath.name} -> testing revision {revision_id}" + ) + + # Case 2: Migration file changed - extract its revision ID + elif filepath.parent == migrations_dir: + revision_id = extract_revision_id(filepath) + if revision_id: + revision_ids.add(revision_id) + print( + f"Migration changed: {filepath.name} -> testing revision {revision_id}" + ) + + if run_full_suite: + # Run full test suite without -k filter + print("\nRunning full alembic test suite due to infrastructure file changes\n") + cmd = [ + "uv", + "run", + "pytest", + "tests/alembic/test_pipeline.py", + ] + elif revision_ids: + # Build a -k expression to filter tests by revision ID + # pytest -k "rev1 or rev2 or rev3" + k_expression = " or ".join(sorted(revision_ids)) + + print( + f"\nRunning tests for {len(revision_ids)} revision(s): {', '.join(sorted(revision_ids))}" + ) + print() + + # Run pytest on test_pipeline.py with -k filter + cmd = [ + "uv", + "run", + "pytest", + "tests/alembic/test_pipeline.py", + "-k", + k_expression, + ] + else: + print("No alembic tests to run") + sys.exit(0) + + result = subprocess.run(cmd, cwd=repo_root) + + sys.exit(result.returncode) + + +if __name__ == "__main__": + main() diff --git a/src/crud/document.py b/src/crud/document.py index 127b5d7b..012e4f73 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -127,6 +127,8 @@ async def create_documents( observer=observer, observed=observed, content=doc.content, + level=doc.level, + times_derived=doc.times_derived, internal_metadata=metadata_dict, embedding=doc.embedding, session_name=doc.session_name, diff --git a/src/crud/representation.py b/src/crud/representation.py index 13aefce7..4a41cdc6 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -145,7 +145,6 @@ class RepresentationManager: metadata: schemas.DocumentMetadata = schemas.DocumentMetadata( message_ids=[message_id_range], - level=obs_level, premises=obs_premises, message_created_at=format_datetime_utc(message_created_at), ) @@ -154,6 +153,7 @@ class RepresentationManager: schemas.DocumentCreate( content=obs_content, session_name=session_name, + level=obs_level, metadata=metadata, embedding=embedding, ) @@ -397,7 +397,7 @@ class RepresentationManager: models.Document.observer == self.observer, models.Document.observed == self.observed, ) - .order_by(models.Document.internal_metadata["times_derived"].desc()) + .order_by(models.Document.times_derived.desc()) ) result = await db.execute(stmt) @@ -454,7 +454,7 @@ class RepresentationManager: conditions: list[dict[str, Any]] = [] if level: - conditions.append({"internal_metadata": {"level": level}}) + conditions.append({"level": level}) if not conditions: return {} diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 6cafc0fc..9dbfb6bb 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -1,5 +1,4 @@ import logging -from typing import Any import sentry_sdk from pydantic import ValidationError @@ -26,8 +25,12 @@ logging.getLogger("sqlalchemy.engine.Engine").disabled = True console = Console(markup=True) -async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: +async def process_item(queue_item: models.QueueItem) -> None: """Process a single item from the queue.""" + task_type = queue_item.task_type + queue_payload = queue_item.payload + workspace_name = queue_item.workspace_name + if task_type == "webhook": try: validated = WebhookPayload(**queue_payload) @@ -39,7 +42,7 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: ) raise ValueError(f"Invalid payload structure: {str(e)}") from e async with tracked_db() as db: - await webhook_delivery.deliver_webhook(db, validated) + await webhook_delivery.deliver_webhook(db, validated, workspace_name) elif task_type == "summary": try: @@ -52,17 +55,21 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: ) raise ValueError(f"Invalid payload structure: {str(e)}") from e + # Use workspace_name and message_id from QueueItem columns + message_id = queue_item.message_id + + if message_id is None: + raise ValueError("Summary tasks require a message_id") + message_public_id = validated.message_public_id if not message_public_id: - logger.debug( - "Fetching message public ID for message %s", validated.message_id - ) + logger.debug("Fetching message public ID for message %s", message_id) async with tracked_db(operation_name="summary_fallback") as db: stmt = ( select(models.Message) - .where(models.Message.workspace_name == validated.workspace_name) + .where(models.Message.workspace_name == workspace_name) .where(models.Message.session_name == validated.session_name) - .where(models.Message.id == validated.message_id) + .where(models.Message.id == message_id) ) result = await db.execute(stmt) @@ -70,22 +77,20 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: if message is None: logger.error( "Failed to fetch message with ID %s for process_summary_task", - validated.message_id, + message_id, ) return message_public_id = message.public_id with sentry_sdk.start_transaction(name="process_summary_task", op="deriver"): await summarizer.summarize_if_needed( - validated.workspace_name, + workspace_name, validated.session_name, - validated.message_id, + message_id, validated.message_seq_in_session, message_public_id, ) - log_performance_metrics( - "summary", f"{validated.workspace_name}_{validated.message_id}" - ) + log_performance_metrics("summary", f"{workspace_name}_{message_id}") elif task_type == "dream": with sentry_sdk.start_transaction(name="process_dream_task", op="deriver"): @@ -98,7 +103,7 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: queue_payload, ) raise ValueError(f"Invalid payload structure: {str(e)}") from e - await process_dream(validated) + await process_dream(validated, workspace_name) else: raise ValueError(f"Invalid task type: {task_type}") @@ -109,29 +114,21 @@ async def process_representation_batch( observer: str | None, observed: str | None, ) -> None: - """Validate incoming queue payloads and dispatch to the appropriate handler. - - This function centralizes payload validation using a simple mapping from - task type to Pydantic model. After validation, routes the request to - the correct processor without repeating type checks elsewhere. - + """Prepares and processes a batch of messages for representation tasks. Args: - task_type: The type of task to process - queue_payloads: List of payload dictionaries to process - observed (optional): For representation tasks, the observed from work_unit_key - to identify which messages should be focused on - observer (optional): For representation tasks, the observer from work_unit_key - to identify which messages should be focused on + messages: List of messages to process + observer: The observer of the messages + observed: The observed of the messages """ if not messages or not messages[0]: - logger.debug("process_representation_batch received no payloads") + logger.debug("process_representation_batch received no messages") return if observed is None or observer is None: raise ValueError("observed and observer are required for representation tasks") logger.debug( - "process_representation_batch received %s payloads", + "process_representation_batch received %s messages", len(messages), ) diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py index 9731b91c..cf15c5d0 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -11,7 +11,7 @@ from src.dreamer.dream_scheduler import get_affected_dream_keys, get_dream_sched from src.exceptions import ValidationException from src.models import QueueItem from src.utils.queue_payload import create_payload -from src.utils.work_unit import get_work_unit_key +from src.utils.work_unit import construct_work_unit_key logger = logging.getLogger(__name__) @@ -159,8 +159,16 @@ def create_representation_record( session_id: Optional session ID Returns: - Queue record dictionary + Queue record dictionary with workspace_name and message_id as separate fields """ + workspace_name = message.get("workspace_name") + message_id = message.get("message_id") + + if not isinstance(workspace_name, str): + raise TypeError("workspace_name is required and must be a string") + if not isinstance(message_id, int): + raise TypeError("message_id is required and must be an integer") + processed_payload = create_payload( message=message, task_type="representation", @@ -168,10 +176,12 @@ def create_representation_record( observed=observed, ) return { - "work_unit_key": get_work_unit_key(processed_payload), + "work_unit_key": construct_work_unit_key(workspace_name, processed_payload), "payload": processed_payload, "session_id": session_id, "task_type": "representation", + "workspace_name": workspace_name, + "message_id": message_id, } @@ -185,23 +195,32 @@ def create_summary_record( Args: message: The message payload - observed: Name of the sender - observer: Name of the target session_id: Session ID + message_seq_in_session: The sequence number of the message in the session Returns: - Queue record dictionary + Queue record dictionary with workspace_name and message_id as separate fields """ + workspace_name = message.get("workspace_name") + message_id = message.get("message_id") + + if not isinstance(workspace_name, str): + raise ValueError("workspace_name is required and must be a string") + if not isinstance(message_id, int): + raise ValueError("message_id is required and must be an integer") + processed_payload = create_payload( message=message, task_type="summary", message_seq_in_session=message_seq_in_session, ) return { - "work_unit_key": get_work_unit_key(processed_payload), + "work_unit_key": construct_work_unit_key(workspace_name, processed_payload), "payload": processed_payload, "session_id": session_id, "task_type": "summary", + "workspace_name": workspace_name, + "message_id": message_id, } diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 953895a2..df2f1465 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -10,7 +10,7 @@ import sentry_sdk from dotenv import load_dotenv from nanoid import generate as generate_nanoid from sentry_sdk.integrations.asyncio import AsyncioIntegration -from sqlalchemy import BigInteger, and_, delete, select, update +from sqlalchemy import and_, delete, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import func @@ -464,9 +464,7 @@ class QueueManager: break try: - await process_item( - work_unit.task_type, queue_item.payload - ) + await process_item(queue_item) await self.mark_queue_items_as_processed( [queue_item], work_unit_key ) @@ -615,8 +613,7 @@ class QueueManager: .select_from(models.QueueItem) .join( models.Message, - func.cast(models.QueueItem.payload["message_id"].astext, BigInteger) - == models.Message.id, + models.QueueItem.message_id == models.Message.id, ) .where(~models.QueueItem.processed) .where(models.Message.session_name == parsed_key.session_name) @@ -662,10 +659,7 @@ class QueueManager: and_( models.QueueItem.work_unit_key == work_unit_key, ~models.QueueItem.processed, - func.cast( - models.QueueItem.payload["message_id"].astext, BigInteger - ) - == models.Message.id, + models.QueueItem.message_id == models.Message.id, ), ) .where(allowed_condition) @@ -690,7 +684,11 @@ class QueueManager: if items_to_process: max_queue_item_message_id = max( - [qi.payload["message_id"] for qi in items_to_process] + [ + qi.message_id + for qi in items_to_process + if qi.message_id is not None + ] ) messages_context = [ # remove any messages that are after the last message_id from queue items m for m in messages_context if m.id <= max_queue_item_message_id diff --git a/src/dreamer/dream_scheduler.py b/src/dreamer/dream_scheduler.py index a56edd3f..f28dead9 100644 --- a/src/dreamer/dream_scheduler.py +++ b/src/dreamer/dream_scheduler.py @@ -12,7 +12,7 @@ from src import models from src.config import settings from src.dependencies import tracked_db from src.utils.queue_payload import create_dream_payload -from src.utils.work_unit import get_work_unit_key, parse_work_unit_key +from src.utils.work_unit import construct_work_unit_key, parse_work_unit_key logger = getLogger(__name__) @@ -48,13 +48,13 @@ def get_affected_dream_keys(message: dict[str, Any]) -> list[str]: return [] # Generate dream work unit key for this peer's collection - dream_key = get_work_unit_key( + dream_key = construct_work_unit_key( + workspace_name, { "task_type": "dream", - "workspace_name": workspace_name, "observer": peer_name, "observed": peer_name, - } + }, ) return [dream_key] @@ -193,7 +193,6 @@ class DreamScheduler: ) -> None: """Execute the dream by enqueueing it and updating collection metadata.""" dream_payload = create_dream_payload( - workspace_name=workspace_name, dream_type="consolidate", observer=observer, observed=observed, @@ -205,6 +204,8 @@ class DreamScheduler: "payload": dream_payload, "session_id": None, "task_type": "dream", + "workspace_name": workspace_name, + "message_id": None, # Dreams don't have a message_id } await db.execute(insert(models.QueueItem), [dream_record]) @@ -325,13 +326,13 @@ async def check_and_schedule_dream( dream_scheduler = get_dream_scheduler() if dream_scheduler: - collection_work_unit_key = get_work_unit_key( + collection_work_unit_key = construct_work_unit_key( + collection.workspace_name, { "task_type": "dream", - "workspace_name": collection.workspace_name, "observer": collection.observer, "observed": collection.observed, - } + }, ) await dream_scheduler.schedule_dream( diff --git a/src/dreamer/dreamer.py b/src/dreamer/dreamer.py index 9ee43bb8..da73b86e 100644 --- a/src/dreamer/dreamer.py +++ b/src/dreamer/dreamer.py @@ -24,6 +24,7 @@ logger = logging.getLogger(__name__) @sentry_sdk.trace async def process_dream( payload: DreamPayload, + workspace_name: str, ) -> None: """ Process a dream task by performing collection maintenance operations. @@ -32,12 +33,12 @@ async def process_dream( payload: The dream task payload containing workspace, peer, and dream type information """ logger.info( - f"Processing dream task: {payload.dream_type} for {payload.workspace_name}/{payload.observer}/{payload.observed}" + f"Processing dream task: {payload.dream_type} for {workspace_name}/{payload.observer}/{payload.observed}" ) try: if payload.dream_type == "consolidate": - await _process_consolidate_dream(payload) + await _process_consolidate_dream(payload, workspace_name) ## TODO other dream types except Exception as e: @@ -50,7 +51,9 @@ async def process_dream( # Don't re-raise - we want to mark the dream task as processed even if it fails -async def _process_consolidate_dream(payload: DreamPayload) -> None: +async def _process_consolidate_dream( + payload: DreamPayload, workspace_name: str +) -> None: """ Process a consolidation dream task. @@ -63,7 +66,7 @@ async def _process_consolidate_dream(payload: DreamPayload) -> None: logger.info( f""" (っ- ‸ - ς)ᶻ z 𐰁 ᶻ z 𐰁 ᶻ z 𐰁\n -DREAM: consolidating documents for {payload.workspace_name}/{payload.observer}/{payload.observed}\n +DREAM: consolidating documents for {workspace_name}/{payload.observer}/{payload.observed}\n 𐰁 z ᶻ 𐰁 z ᶻ 𐰁 z ᶻ(っ- ‸ - ς)""" ) @@ -71,7 +74,7 @@ DREAM: consolidating documents for {payload.workspace_name}/{payload.observer}/{ async with tracked_db("dream_consolidate") as db: documents = await crud.get_all_documents( db, - payload.workspace_name, + workspace_name, observer=payload.observer, observed=payload.observed, ) @@ -88,7 +91,7 @@ DREAM: consolidating documents for {payload.workspace_name}/{payload.observer}/{ for cluster in clusters: await _consolidate_cluster( cluster, - payload.workspace_name, + workspace_name, db, observer=payload.observer, observed=payload.observed, @@ -118,9 +121,7 @@ async def _consolidate_cluster( logger.info("consolidated representation:\n%s", consolidated_representation) # TODO: less hacky preservation of times_derived - total_times_derived = sum( - doc.internal_metadata.get("times_derived", 1) for doc in cluster - ) + total_times_derived = sum(doc.times_derived for doc in cluster) new_documents = [ *consolidated_representation.explicit, @@ -141,10 +142,8 @@ async def _consolidate_cluster( # NOTE: other kinds of observations here in the future metadata = schemas.DocumentMetadata( - times_derived=total_times_derived, message_ids=obs.message_ids, message_created_at=format_datetime_utc(obs.created_at), - level=level, premises=premises, ) @@ -154,6 +153,8 @@ async def _consolidate_cluster( schemas.DocumentCreate( content=content, session_name=obs.session_name, + level=level, + times_derived=total_times_derived, metadata=metadata, embedding=embedding, ) diff --git a/src/models.py b/src/models.py index 2efad02f..f0042cd0 100644 --- a/src/models.py +++ b/src/models.py @@ -1,6 +1,6 @@ import datetime from logging import getLogger -from typing import Any, Literal, final +from typing import Any, final from dotenv import load_dotenv from nanoid import generate as generate_nanoid @@ -26,6 +26,8 @@ from sqlalchemy.orm.properties import MappedColumn from sqlalchemy.sql import func from typing_extensions import override +from src.utils.types import DocumentLevel, TaskType + from .db import Base load_dotenv(override=True) @@ -382,6 +384,12 @@ class Document(Base): "internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb") ) content: Mapped[str] = mapped_column(TEXT) + level: Mapped[DocumentLevel] = mapped_column( + TEXT, nullable=False, server_default="explicit" + ) + times_derived: Mapped[int] = mapped_column( + Integer, nullable=False, server_default=text("1") + ) embedding: MappedColumn[Any] = mapped_column(Vector(1536)) created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now() @@ -397,6 +405,7 @@ class Document(Base): CheckConstraint("length(id) = 21", name="id_length"), CheckConstraint("length(content) <= 65535", name="content_length"), CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"), + CheckConstraint("level IN ('explicit', 'deductive')", name="level_valid"), # Composite foreign key constraint for collections ForeignKeyConstraint( ["observer", "observed", "workspace_name"], @@ -439,9 +448,6 @@ class Document(Base): ) -TaskType = Literal["webhook", "summary", "representation", "dream"] - - @final class QueueItem(Base): __tablename__: str = "queue" @@ -460,14 +466,36 @@ class QueueItem(Base): created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now() ) + workspace_name: Mapped[str] = mapped_column( + ForeignKey("workspaces.name"), nullable=False + ) + message_id: Mapped[int | None] = mapped_column( + BigInteger, ForeignKey("messages.id"), nullable=True + ) __table_args__ = ( Index("ix_queue_created_at", "created_at"), Index("ix_queue_session_id", "session_id"), + Index( + "ix_queue_workspace_name", + "workspace_name", + ), + Index( + "ix_queue_message_id_not_null", + "message_id", + postgresql_where=text("message_id IS NOT NULL"), + ), + Index("ix_queue_workspace_name_processed", "workspace_name", "processed"), + Index( + "ix_queue_work_unit_key_processed_id", + "work_unit_key", + "processed", + "id", + ), ) def __repr__(self) -> str: - return f"QueueItem(id={self.id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed})" + return f"QueueItem(id={self.id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed}, workspace_name={self.workspace_name}, message_id={self.message_id})" @final diff --git a/src/schemas.py b/src/schemas.py index a1eac978..4d99dd42 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -288,20 +288,12 @@ class DocumentBase(BaseModel): class DocumentMetadata(BaseModel): - times_derived: int | None = Field( - default=None, - ge=1, - description="The number of times that a semantic duplicate document to this one has been derived", - ) message_ids: list[tuple[int, int]] = Field( description="The ID range(s) of the messages that this document was derived from. Acts as a link to the primary source of the document. Note that as a document gets deduplicated, additional ranges will be added, because the same document could be derived from completely separate message ranges." ) message_created_at: str = Field( description="The timestamp of the message that this document was derived from. Note that this is not the same as the created_at timestamp of the document. This timestamp is usually only saved with second-level precision." ) - level: Literal["explicit", "deductive"] = Field( - description="The level of the document (explicit or deductive)" - ) premises: list[str] | None = Field( default=None, description="The premises of the deduction -- only applicable for deductive observations", @@ -313,6 +305,15 @@ class DocumentCreate(DocumentBase): session_name: str = Field( description="The session from which the document was derived" ) + level: Literal["explicit", "deductive"] = Field( + default="explicit", + description="The level of the document (explicit or deductive)", + ) + times_derived: int = Field( + default=1, + ge=1, + description="The number of times that a semantic duplicate document to this one has been derived", + ) metadata: DocumentMetadata = Field() embedding: list[float] = Field() diff --git a/src/utils/queue_payload.py b/src/utils/queue_payload.py index e4d4a964..b267f04a 100644 --- a/src/utils/queue_payload.py +++ b/src/utils/queue_payload.py @@ -14,9 +14,7 @@ class RepresentationPayload(BasePayload): """Payload for representation tasks.""" task_type: Literal["representation"] = "representation" - workspace_name: str session_name: str - message_id: int content: str observer: str observed: str @@ -33,9 +31,7 @@ class SummaryPayload(BasePayload): """Payload for summary tasks.""" task_type: Literal["summary"] = "summary" - workspace_name: str session_name: str - message_id: int message_seq_in_session: int # Optional for backward compatibility with older queue items message_public_id: str | None = None @@ -45,7 +41,6 @@ class WebhookPayload(BasePayload): """Payload for webhook delivery tasks.""" task_type: Literal["webhook"] = "webhook" - workspace_name: str event_type: str data: dict[str, Any] @@ -54,35 +49,33 @@ class DreamPayload(BasePayload): """Payload for dream tasks.""" task_type: Literal["dream"] = "dream" - workspace_name: str dream_type: Literal["consolidate"] = "consolidate" observer: str observed: str def create_webhook_payload( - workspace_name: str, event_type: str, data: dict[str, Any], ) -> dict[str, Any]: - return WebhookPayload( - workspace_name=workspace_name, event_type=event_type, data=data - ).model_dump(mode="json") + """Create a webhook payload.""" + return WebhookPayload(event_type=event_type, data=data).model_dump( + mode="json", exclude_none=True + ) def create_dream_payload( - workspace_name: str, dream_type: Literal["consolidate"] = "consolidate", *, observer: str, observed: str, ) -> dict[str, Any]: + """Create a dream payload.""" return DreamPayload( - workspace_name=workspace_name, dream_type=dream_type, observer=observer, observed=observed, - ).model_dump(mode="json") + ).model_dump(mode="json", exclude_none=True) def create_payload( @@ -96,6 +89,10 @@ def create_payload( """ Create a processed payload from a message for queue processing. + Note: workspace_name and message_id are no longer included in the returned payload + as they are now stored in dedicated columns on the queue table. The caller is + responsible for extracting and passing these values separately. + Args: message: The original message dictionary task_type: Type of task ('representation' or 'summary') @@ -104,7 +101,7 @@ def create_payload( message_seq_in_session: Required for summary tasks, must be None for representation Returns: - Processed payload dictionary ready for queue processing + Processed payload dictionary ready for queue processing (without workspace_name and message_id) Raises: ValueError: If the payload doesn't match the expected schema @@ -142,9 +139,7 @@ def create_payload( validated_payload = RepresentationPayload( content=content, - workspace_name=workspace_name, session_name=session_name, - message_id=message_id, created_at=created_at, observer=observer, observed=observed, @@ -161,16 +156,14 @@ def create_payload( ) validated_payload = SummaryPayload( - workspace_name=workspace_name, session_name=session_name, - message_id=message_id, message_seq_in_session=message_seq_in_session, message_public_id=message_public_id, ) # Convert back to dict for compatibility with JSON serialization # mode='json' ensures datetime is converted to ISO string - payload = validated_payload.model_dump(mode="json") + payload = validated_payload.model_dump(mode="json", exclude_none=True) except Exception as e: raise ValueError(f"Failed to create valid payload: {str(e)}") from e diff --git a/src/utils/representation.py b/src/utils/representation.py index 30ebe71d..c0f499f0 100644 --- a/src/utils/representation.py +++ b/src/utils/representation.py @@ -271,7 +271,7 @@ class Representation(BaseModel): session_name=doc.session_name, ) for doc in documents - if doc.internal_metadata.get("level") == "explicit" + if doc.level == "explicit" ], deductive=[ DeductiveObservation( @@ -284,7 +284,7 @@ class Representation(BaseModel): premises=doc.internal_metadata.get("premises", []), ) for doc in documents - if doc.internal_metadata.get("level") == "deductive" + if doc.level == "deductive" ], ) diff --git a/src/utils/types.py b/src/utils/types.py index 317e545a..80dfc389 100644 --- a/src/utils/types.py +++ b/src/utils/types.py @@ -1,3 +1,5 @@ from typing import Literal SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom", "vllm"] +TaskType = Literal["webhook", "summary", "representation", "dream"] +DocumentLevel = Literal["explicit", "deductive"] diff --git a/src/utils/work_unit.py b/src/utils/work_unit.py index d7fd12bb..660baa16 100644 --- a/src/utils/work_unit.py +++ b/src/utils/work_unit.py @@ -15,11 +15,14 @@ class ParsedWorkUnit(BaseModel): observed: str | None -def get_work_unit_key(payload: dict[str, Any] | ParsedWorkUnit) -> str: +def construct_work_unit_key( + workspace_name: str, payload: dict[str, Any] | ParsedWorkUnit +) -> str: """ Generate a work unit key for a given task type, workspace name, and event type. Args: + workspace_name: The name of the workspace the work unit belongs to payload: Dictionary containing work unit information Returns: @@ -31,7 +34,6 @@ def get_work_unit_key(payload: dict[str, Any] | ParsedWorkUnit) -> str: if isinstance(payload, ParsedWorkUnit): payload = payload.model_dump() - workspace_name: str | None = payload.get("workspace_name") task_type: str | None = payload.get("task_type") if not workspace_name or not task_type: raise ValueError( diff --git a/src/webhooks/events.py b/src/webhooks/events.py index f942bb5d..a37eab89 100644 --- a/src/webhooks/events.py +++ b/src/webhooks/events.py @@ -7,7 +7,7 @@ from pydantic import BaseModel from src.dependencies import tracked_db from src.models import QueueItem from src.utils.queue_payload import create_webhook_payload -from src.utils.work_unit import get_work_unit_key +from src.utils.work_unit import construct_work_unit_key logger = logging.getLogger(__name__) @@ -51,23 +51,26 @@ async def publish_webhook_event(event: WebhookEvent) -> None: event: The webhook event to publish. """ try: + # Note: workspace_name is no longer included in the payload + # It's stored directly on the queue item payload = create_webhook_payload( - workspace_name=event.workspace_id, event_type=event.type.value, data=event.model_dump(mode="json", exclude={"type"}), ) async with tracked_db("publish_webhook_event") as db: queue_item = QueueItem( - work_unit_key=get_work_unit_key( + work_unit_key=construct_work_unit_key( + event.workspace_id, { "task_type": "webhook", - "workspace_name": event.workspace_id, - } + }, ), payload=payload, session_id=None, task_type="webhook", + workspace_name=event.workspace_id, + message_id=None, # Webhooks don't have a message_id ) db.add(queue_item) await db.commit() diff --git a/src/webhooks/webhook_delivery.py b/src/webhooks/webhook_delivery.py index 316fdccc..d26aa404 100644 --- a/src/webhooks/webhook_delivery.py +++ b/src/webhooks/webhook_delivery.py @@ -15,16 +15,18 @@ from src.utils.queue_payload import WebhookPayload logger = logging.getLogger(__name__) -async def deliver_webhook(db: AsyncSession, payload: WebhookPayload) -> None: +async def deliver_webhook( + db: AsyncSession, payload: WebhookPayload, workspace_name: str +) -> None: """ Deliver a single webhook event to its configured endpoints. """ async with httpx.AsyncClient(timeout=30.0) as client: try: - webhook_urls = await _get_webhook_urls(db, payload.workspace_name) + webhook_urls = await _get_webhook_urls(db, workspace_name) if not webhook_urls: logger.debug( - f"No webhook endpoints for workspace {payload.workspace_name}, skipping." + f"No webhook endpoints for workspace {workspace_name}, skipping." ) return @@ -72,7 +74,7 @@ async def deliver_webhook(db: AsyncSession, payload: WebhookPayload) -> None: ) except httpx.RequestError: - logger.exception(f"Error sending webhook for {payload.workspace_name}.") + logger.exception(f"Error sending webhook for {workspace_name}.") except Exception: logger.exception("Unexpected error delivering webhook.") diff --git a/tests/alembic/revisions/__init__.py b/tests/alembic/revisions/__init__.py index 988130db..fedc98c5 100644 --- a/tests/alembic/revisions/__init__.py +++ b/tests/alembic/revisions/__init__.py @@ -13,9 +13,11 @@ from . import ( test_917195d9b5e9_add_messageembedding_table, test_a1b2c3d4e5f6_initial_schema, test_b765d82110bd_change_metamessages_to_user_level_with_, + test_b8183c5ffb48_codify_document_level_and_times_derived, test_bb6fb3a7a643_add_message_seq_in_session_column, test_c3828084f472_add_indexes_for_messages_and_, test_d429de0e5338_adopt_peer_paradigm, + test_ec8f94139b02_codify_workspace_name_and_message_id_in_, ) __all__ = [ @@ -31,7 +33,9 @@ __all__ = [ "test_917195d9b5e9_add_messageembedding_table", "test_a1b2c3d4e5f6_initial_schema", "test_b765d82110bd_change_metamessages_to_user_level_with_", + "test_b8183c5ffb48_codify_document_level_and_times_derived", "test_bb6fb3a7a643_add_message_seq_in_session_column", "test_c3828084f472_add_indexes_for_messages_and_", "test_d429de0e5338_adopt_peer_paradigm", + "test_ec8f94139b02_codify_workspace_name_and_message_id_in_", ] diff --git a/tests/alembic/revisions/test_b8183c5ffb48_codify_document_level_and_times_derived.py b/tests/alembic/revisions/test_b8183c5ffb48_codify_document_level_and_times_derived.py new file mode 100644 index 00000000..c3e3d3a0 --- /dev/null +++ b/tests/alembic/revisions/test_b8183c5ffb48_codify_document_level_and_times_derived.py @@ -0,0 +1,315 @@ +"""Hooks for revision b8183c5ffb48 (codify_document_level_and_times_derived).""" + +from __future__ import annotations + +import json + +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 + +# Test data constants +WORKSPACE_NAME = "test-workspace" +OBSERVER_NAME = "test-observer" +OBSERVED_NAME = "test-observed" +SESSION_NAME = "test-session" + + +@register_before_upgrade("b8183c5ffb48") +def prepare_codify_document_level_and_times_derived( + verifier: MigrationVerifier, +) -> None: + """Seed state and assertions before upgrading to b8183c5ffb48.""" + # Verify columns don't exist yet + verifier.assert_column_exists("documents", "level", exists=False) + verifier.assert_column_exists("documents", "times_derived", exists=False) + + # Verify CHECK constraint doesn't exist yet + verifier.assert_constraint_exists("documents", "level_valid", "check", exists=False) + + conn = verifier.conn + schema = verifier.schema + + # Create workspace + conn.execute( + text( + f'INSERT INTO "{schema}"."workspaces" ("id", "name") VALUES (:ws_id, :ws_name)' + ), + {"ws_id": generate_nanoid(), "ws_name": WORKSPACE_NAME}, + ) + + # Create peers (observer and observed) + conn.execute( + text( + f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") ' + + "VALUES (:observer_id, :observer_name, :ws_name), " + + "(:observed_id, :observed_name, :ws_name)" + ), + { + "observer_id": generate_nanoid(), + "observer_name": OBSERVER_NAME, + "ws_name": WORKSPACE_NAME, + "observed_id": generate_nanoid(), + "observed_name": OBSERVED_NAME, + }, + ) + + # Create session + conn.execute( + text( + f'INSERT INTO "{schema}"."sessions" ("id", "name", "workspace_name") ' + + "VALUES (:session_id, :session_name, :ws_name)" + ), + { + "session_id": generate_nanoid(), + "session_name": SESSION_NAME, + "ws_name": WORKSPACE_NAME, + }, + ) + + # Create collection (required for documents) + conn.execute( + text( + f'INSERT INTO "{schema}"."collections" ("id", "observer", "observed", "workspace_name") ' + + "VALUES (:collection_id, :observer, :observed, :ws_name)" + ), + { + "collection_id": generate_nanoid(), + "observer": OBSERVER_NAME, + "observed": OBSERVED_NAME, + "ws_name": WORKSPACE_NAME, + }, + ) + + # Create documents with different scenarios + # Scenario 1: Document with both level and times_derived in internal_metadata + conn.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") ' + + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)" + ), + { + "id": generate_nanoid(), + "content": "Document with explicit level and times_derived=3", + "metadata": json.dumps({"level": "explicit", "times_derived": 3}), + "embedding": [0.1] * 1536, + "observer": OBSERVER_NAME, + "observed": OBSERVED_NAME, + "ws_name": WORKSPACE_NAME, + "session_name": SESSION_NAME, + }, + ) + + # Scenario 2: Document with deductive level in internal_metadata + conn.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") ' + + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)" + ), + { + "id": generate_nanoid(), + "content": "Document with deductive level and times_derived=5", + "metadata": json.dumps({"level": "deductive", "times_derived": 5}), + "embedding": [0.2] * 1536, + "observer": OBSERVER_NAME, + "observed": OBSERVED_NAME, + "ws_name": WORKSPACE_NAME, + "session_name": SESSION_NAME, + }, + ) + + # Scenario 3: Document without level or times_derived (should get defaults) + conn.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") ' + + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)" + ), + { + "id": generate_nanoid(), + "content": "Document without level or times_derived fields", + "metadata": json.dumps({"other_field": "value"}), + "embedding": [0.3] * 1536, + "observer": OBSERVER_NAME, + "observed": OBSERVED_NAME, + "ws_name": WORKSPACE_NAME, + "session_name": SESSION_NAME, + }, + ) + + # Scenario 4: Document with only level in metadata + conn.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") ' + + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)" + ), + { + "id": generate_nanoid(), + "content": "Document with only level field", + "metadata": json.dumps({"level": "explicit"}), + "embedding": [0.4] * 1536, + "observer": OBSERVER_NAME, + "observed": OBSERVED_NAME, + "ws_name": WORKSPACE_NAME, + "session_name": SESSION_NAME, + }, + ) + + # Scenario 5: Document with only times_derived in metadata + conn.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") ' + + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)" + ), + { + "id": generate_nanoid(), + "content": "Document with only times_derived field", + "metadata": json.dumps({"times_derived": 7}), + "embedding": [0.5] * 1536, + "observer": OBSERVER_NAME, + "observed": OBSERVED_NAME, + "ws_name": WORKSPACE_NAME, + "session_name": SESSION_NAME, + }, + ) + + # Verify we have exactly 5 documents + count = conn.execute(text(f'SELECT COUNT(*) FROM "{schema}"."documents"')).scalar() + assert count == 5, f"Expected 5 documents but found {count}" + + +@register_after_upgrade("b8183c5ffb48") +def verify_codify_document_level_and_times_derived(verifier: MigrationVerifier) -> None: + """Add assertions validating the effects of b8183c5ffb48.""" + # Verify columns were added with correct nullability + verifier.assert_column_exists("documents", "level", nullable=False) + verifier.assert_column_exists("documents", "times_derived", nullable=False) + + # Verify CHECK constraint exists + verifier.assert_constraint_exists("documents", "level_valid", "check") + + conn = verifier.conn + schema = verifier.schema + + # Verify all rows have non-null values after migration + verifier.assert_no_nulls("documents", "level") + verifier.assert_no_nulls("documents", "times_derived") + + # Verify data transformation: level extracted from internal_metadata + explicit_count = conn.execute( + text(f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE "level" = :level'), + {"level": "explicit"}, + ).scalar() + deductive_count = conn.execute( + text(f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE "level" = :level'), + {"level": "deductive"}, + ).scalar() + + # 3 documents should have explicit (scenarios 1, 3, 4) + # 1 document should have deductive (scenario 2) + # 1 document should have default explicit (scenario 5) + assert explicit_count == 4, f"Expected 4 explicit documents, got {explicit_count}" + assert deductive_count == 1, f"Expected 1 deductive document, got {deductive_count}" + + # Verify specific times_derived values were migrated correctly + times_derived_3 = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."documents" ' + + 'WHERE "times_derived" = 3 AND "content" LIKE :pattern' + ), + {"pattern": "%explicit level and times_derived=3%"}, + ).scalar() + assert ( + times_derived_3 == 1 + ), f"Expected 1 document with times_derived=3, got {times_derived_3}" + + times_derived_5 = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."documents" ' + + 'WHERE "times_derived" = 5 AND "content" LIKE :pattern' + ), + {"pattern": "%deductive level and times_derived=5%"}, + ).scalar() + assert ( + times_derived_5 == 1 + ), f"Expected 1 document with times_derived=5, got {times_derived_5}" + + times_derived_7 = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."documents" ' + + 'WHERE "times_derived" = 7 AND "content" LIKE :pattern' + ), + {"pattern": "%only times_derived field%"}, + ).scalar() + assert ( + times_derived_7 == 1 + ), f"Expected 1 document with times_derived=7, got {times_derived_7}" + + # Verify default times_derived=1 was applied to documents without it + times_derived_1 = conn.execute( + text(f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE "times_derived" = 1') + ).scalar() + assert ( + times_derived_1 == 2 + ), f"Expected 2 documents with times_derived=1, got {times_derived_1}" + + # Verify internal_metadata still contains the original data (NOT removed by migration) + level_in_metadata = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE internal_metadata ? \'level\'' + ) + ).scalar() + # 3 documents had level in metadata (scenarios 1, 2, 4) + assert ( + level_in_metadata == 3 + ), f"Expected 3 documents with level in metadata, got {level_in_metadata}" + + times_derived_in_metadata = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE internal_metadata ? \'times_derived\'' + ) + ).scalar() + # 3 documents had times_derived in metadata (scenarios 1, 2, 5) + assert ( + times_derived_in_metadata == 3 + ), f"Expected 3 documents with times_derived in metadata, got {times_derived_in_metadata}" + + # Verify server defaults work for new documents + new_doc_id = generate_nanoid() + conn.execute( + text( + f'INSERT INTO "{schema}"."documents" ' + + '("id", "content", "internal_metadata", "embedding", "observer", "observed", "workspace_name", "session_name") ' + + "VALUES (:id, :content, :metadata, :embedding, :observer, :observed, :ws_name, :session_name)" + ), + { + "id": new_doc_id, + "content": "New document after migration", + "metadata": json.dumps({}), + "embedding": [0.6] * 1536, + "observer": OBSERVER_NAME, + "observed": OBSERVED_NAME, + "ws_name": WORKSPACE_NAME, + "session_name": SESSION_NAME, + }, + ) + + # Verify the new document got default values + new_doc = conn.execute( + text( + f'SELECT "level", "times_derived" FROM "{schema}"."documents" WHERE "id" = :id' + ), + {"id": new_doc_id}, + ).one() + assert ( + new_doc.level == "explicit" + ), f"Expected new document to have level='explicit', got {new_doc.level}" + assert ( + new_doc.times_derived == 1 + ), f"Expected new document to have times_derived=1, got {new_doc.times_derived}" diff --git a/tests/alembic/revisions/test_ec8f94139b02_codify_workspace_name_and_message_id_in_.py b/tests/alembic/revisions/test_ec8f94139b02_codify_workspace_name_and_message_id_in_.py new file mode 100644 index 00000000..cbef748a --- /dev/null +++ b/tests/alembic/revisions/test_ec8f94139b02_codify_workspace_name_and_message_id_in_.py @@ -0,0 +1,336 @@ +"""Hooks for revision ec8f94139b02 (codify_workspace_name_and_message_id_in_).""" + +from __future__ import annotations + +import json + +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 + +# Test data constants +WORKSPACE_NAME_1 = "test-workspace-1" +WORKSPACE_NAME_2 = "test-workspace-2" +PEER_NAME_1 = "test-peer-1" +PEER_NAME_2 = "test-peer-2" +SESSION_NAME_1 = "test-session-1" +SESSION_NAME_2 = "test-session-2" + +# Indexes that should be created by the migration +_INDEXES = ( + ("queue", "ix_queue_workspace_name"), + ("queue", "ix_queue_message_id_not_null"), + ("queue", "ix_queue_workspace_name_processed"), + ("queue", "ix_queue_work_unit_key_processed_id"), +) + + +@register_before_upgrade("ec8f94139b02") +def prepare_codify_workspace_name_and_message_id_in( + verifier: MigrationVerifier, +) -> None: + """Seed state and assertions before upgrading to ec8f94139b02.""" + # Verify columns don't exist yet + verifier.assert_column_exists("queue", "workspace_name", exists=False) + verifier.assert_column_exists("queue", "message_id", exists=False) + + # Verify foreign key constraints don't exist yet + verifier.assert_constraint_exists( + "queue", "fk_queue_workspace_name", "foreign_key", exists=False + ) + verifier.assert_constraint_exists( + "queue", "fk_queue_message_id", "foreign_key", exists=False + ) + + # Verify indexes don't exist yet + verifier.assert_indexes_not_exist(_INDEXES) + + conn = verifier.conn + schema = verifier.schema + + # Create workspaces + conn.execute( + text( + f'INSERT INTO "{schema}"."workspaces" ("id", "name") ' + + "VALUES (:ws_id_1, :ws_name_1), (:ws_id_2, :ws_name_2)" + ), + { + "ws_id_1": generate_nanoid(), + "ws_name_1": WORKSPACE_NAME_1, + "ws_id_2": generate_nanoid(), + "ws_name_2": WORKSPACE_NAME_2, + }, + ) + + # Create peers + conn.execute( + text( + f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") ' + + "VALUES (:peer_id_1, :peer_name_1, :ws_name_1), " + + "(:peer_id_2, :peer_name_2, :ws_name_2)" + ), + { + "peer_id_1": generate_nanoid(), + "peer_name_1": PEER_NAME_1, + "ws_name_1": WORKSPACE_NAME_1, + "peer_id_2": generate_nanoid(), + "peer_name_2": PEER_NAME_2, + "ws_name_2": WORKSPACE_NAME_2, + }, + ) + + # Create sessions + conn.execute( + text( + f'INSERT INTO "{schema}"."sessions" ("id", "name", "workspace_name") ' + + "VALUES (:session_id_1, :session_name_1, :ws_name_1), " + + "(:session_id_2, :session_name_2, :ws_name_2)" + ), + { + "session_id_1": generate_nanoid(), + "session_name_1": SESSION_NAME_1, + "ws_name_1": WORKSPACE_NAME_1, + "session_id_2": generate_nanoid(), + "session_name_2": SESSION_NAME_2, + "ws_name_2": WORKSPACE_NAME_2, + }, + ) + + # Create some messages to reference in queue items + message_ids: list[str] = [] + for i in range(10): + message_id: str = generate_nanoid() + message_ids.append(message_id) + workspace = WORKSPACE_NAME_1 if i % 2 == 0 else WORKSPACE_NAME_2 + session = SESSION_NAME_1 if i % 2 == 0 else SESSION_NAME_2 + peer = PEER_NAME_1 if i % 2 == 0 else PEER_NAME_2 + + conn.execute( + text( + f'INSERT INTO "{schema}"."messages" ' + + '("public_id", "workspace_name", "session_name", "peer_name", "content", "seq_in_session") ' + + "VALUES (:msg_id, :ws_name, :session_name, :peer_name, :content, :seq)" + ), + { + "msg_id": message_id, + "ws_name": workspace, + "session_name": session, + "peer_name": peer, + "content": f"test message {i}", + "seq": i, + }, + ) + + # Get internal message IDs for queue references + message_db_ids: list[int] = [] + for msg_id in message_ids: + result = conn.execute( + text(f'SELECT "id" FROM "{schema}"."messages" WHERE "public_id" = :msg_id'), + {"msg_id": msg_id}, + ).one() + message_db_ids.append(result.id) + + # Bulk insert 100k queue items with workspace_name and message_id in payload + # Use efficient batch insert with generate_series + conn.execute(text("SET LOCAL synchronous_commit = OFF")) + + # Insert queue items in three categories: + # 1. Items with both workspace_name and message_id (60k) + # 2. Items with workspace_name but NO message_id (30k) + # 3. Items with workspace_name and NULL message_id value (10k) + + # Category 1: Both workspace_name and message_id (60k items) + conn.execute( + text( + f'INSERT INTO "{schema}"."queue" ' + + '("work_unit_key", "task_type", "payload", "processed") ' + + "SELECT " + + " 'work-unit-' || gs::text, " + + " 'representation', " + + " jsonb_build_object(" + + " 'workspace_name', CASE WHEN gs % 2 = 0 THEN :ws_name_1 ELSE :ws_name_2 END, " + + " 'message_id', :msg_db_id_0 + (gs % 10), " + + " 'other_field', 'value-' || gs::text" + + " ), " + + " false " + + "FROM generate_series(1, :n) AS gs" + ), + { + "ws_name_1": WORKSPACE_NAME_1, + "ws_name_2": WORKSPACE_NAME_2, + "msg_db_id_0": message_db_ids[0], + "n": 60_000, + }, + ) + + # Category 2: Only workspace_name, no message_id key (30k items) + conn.execute( + text( + f'INSERT INTO "{schema}"."queue" ' + + '("work_unit_key", "task_type", "payload", "processed") ' + + "SELECT " + + " 'work-unit-no-msg-' || gs::text, " + + " 'summary', " + + " jsonb_build_object(" + + " 'workspace_name', CASE WHEN gs % 2 = 0 THEN :ws_name_1 ELSE :ws_name_2 END, " + + " 'other_field', 'value-' || gs::text" + + " ), " + + " false " + + "FROM generate_series(60001, :n) AS gs" + ), + { + "ws_name_1": WORKSPACE_NAME_1, + "ws_name_2": WORKSPACE_NAME_2, + "n": 90_000, + }, + ) + + # Category 3: workspace_name with explicit NULL message_id (10k items) + conn.execute( + text( + f'INSERT INTO "{schema}"."queue" ' + + '("work_unit_key", "task_type", "payload", "processed") ' + + "SELECT " + + " 'work-unit-null-msg-' || gs::text, " + + " 'representation', " + + " jsonb_build_object(" + + " 'workspace_name', CASE WHEN gs % 2 = 0 THEN :ws_name_1 ELSE :ws_name_2 END, " + + " 'message_id', NULL::bigint, " + + " 'other_field', 'value-' || gs::text" + + " ), " + + " false " + + "FROM generate_series(90001, :n) AS gs" + ), + { + "ws_name_1": WORKSPACE_NAME_1, + "ws_name_2": WORKSPACE_NAME_2, + "n": 100_000, + }, + ) + + # Verify we have exactly 100k queue items + count = conn.execute(text(f'SELECT COUNT(*) FROM "{schema}"."queue"')).scalar() + assert count == 100_000, f"Expected 100k queue items but found {count}" + + +@register_after_upgrade("ec8f94139b02") +def verify_codify_workspace_name_and_message_id_in(verifier: MigrationVerifier) -> None: + """Add assertions validating the effects of ec8f94139b02.""" + # Verify columns were added with correct nullability + verifier.assert_column_exists("queue", "workspace_name", nullable=False) + verifier.assert_column_exists("queue", "message_id", nullable=True) + + # Verify all indexes were created + verifier.assert_indexes_exist(_INDEXES) + + # Verify foreign key constraints exist + verifier.assert_constraint_exists("queue", "fk_queue_workspace_name", "foreign_key") + verifier.assert_constraint_exists("queue", "fk_queue_message_id", "foreign_key") + + conn = verifier.conn + schema = verifier.schema + + # Verify all rows have non-null workspace_name after migration + verifier.assert_no_nulls("queue", "workspace_name") + + # Verify data transformation: workspace_name extracted from payload + ws1_count = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."queue" ' + + 'WHERE "workspace_name" = :ws_name' + ), + {"ws_name": WORKSPACE_NAME_1}, + ).scalar() + ws2_count = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."queue" ' + + 'WHERE "workspace_name" = :ws_name' + ), + {"ws_name": WORKSPACE_NAME_2}, + ).scalar() + + # Should be roughly 50/50 split (we alternate in the insert) + assert ( + ws1_count == 50_000 + ), f"Expected 50k items with workspace_name_1, got {ws1_count}" + assert ( + ws2_count == 50_000 + ), f"Expected 50k items with workspace_name_2, got {ws2_count}" + + # Verify data transformation: message_id extracted from payload where it exists + msg_id_count = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."queue" ' + + 'WHERE "message_id" IS NOT NULL' + ) + ).scalar() + + # Should be 60k items with message_id (category 1 only) + assert ( + msg_id_count == 60_000 + ), f"Expected 60k items with message_id, got {msg_id_count}" + + # Verify items without message_id in payload have NULL in column + null_msg_id_count = conn.execute( + text(f'SELECT COUNT(*) FROM "{schema}"."queue" ' + 'WHERE "message_id" IS NULL') + ).scalar() + + # Should be 40k items (30k without key + 10k with NULL value) + assert ( + null_msg_id_count == 40_000 + ), f"Expected 40k items with NULL message_id, got {null_msg_id_count}" + + # Verify workspace_name was removed from payload + ws_in_payload_count = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."queue" ' + + "WHERE payload ? 'workspace_name'" + ) + ).scalar() + assert ( + ws_in_payload_count == 0 + ), f"Found {ws_in_payload_count} items still with workspace_name in payload" + + # Verify message_id was removed from payload + msg_in_payload_count = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."queue" ' + "WHERE payload ? 'message_id'" + ) + ).scalar() + assert ( + msg_in_payload_count == 0 + ), f"Found {msg_in_payload_count} items still with message_id in payload" + + # Verify other_field remains in payload (data preservation) + other_field_count = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."queue" ' + + "WHERE payload ? 'other_field'" + ) + ).scalar() + assert ( + other_field_count == 100_000 + ), f"Expected all 100k items to retain other_field in payload, got {other_field_count}" + + # Spot check: verify a specific queue item was transformed correctly + sample_item = conn.execute( + text( + 'SELECT "workspace_name", "message_id", "payload" ' + + f'FROM "{schema}"."queue" ' + + "WHERE work_unit_key = 'work-unit-1' " + + "LIMIT 1" + ) + ).one() + + assert sample_item.workspace_name == WORKSPACE_NAME_2 # gs=1 is odd, so workspace 2 + assert sample_item.message_id is not None # Category 1 item + payload = ( + json.loads(sample_item.payload) + if isinstance(sample_item.payload, str) + else sample_item.payload + ) + assert "workspace_name" not in payload + assert "message_id" not in payload + assert payload.get("other_field") == "value-1" diff --git a/tests/crud/test_workspace.py b/tests/crud/test_workspace.py index 18601128..d8ea00a5 100644 --- a/tests/crud/test_workspace.py +++ b/tests/crud/test_workspace.py @@ -318,6 +318,7 @@ class TestWorkspaceCRUD: work_unit_key=f"representation:{test_workspace.name}:{session.name}:{test_peer.name}:{test_peer.name}", task_type="representation", payload={"test": "data"}, + workspace_name=test_workspace.name, ) db_session.add(queue_item) await db_session.flush() diff --git a/tests/deriver/conftest.py b/tests/deriver/conftest.py index 89410656..6ece4892 100644 --- a/tests/deriver/conftest.py +++ b/tests/deriver/conftest.py @@ -1,7 +1,7 @@ import asyncio -from collections.abc import Callable, Generator +from collections.abc import Awaitable, Callable, Generator, Sequence from datetime import datetime, timezone -from typing import Any, Literal +from typing import Any, Literal, TypeAlias, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -11,7 +11,10 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models, schemas from src.utils.queue_payload import create_payload -from src.utils.work_unit import get_work_unit_key +from src.utils.work_unit import construct_work_unit_key + +QueuePayload: TypeAlias = dict[str, Any] +QueuePayloadEntry: TypeAlias = QueuePayload | tuple[QueuePayload, int | None] @pytest.fixture @@ -155,18 +158,29 @@ def create_queue_payload() -> Callable[..., Any]: @pytest.fixture async def add_queue_items( db_session: AsyncSession, -) -> Callable[[list[dict[str, Any]], str], Any]: +) -> Callable[ + [Sequence[QueuePayloadEntry], str, str], Awaitable[list[models.QueueItem]] +]: """Helper function to add queue items to the database""" async def _add_items( - payloads: list[dict[str, Any]], session_id: str + payloads: Sequence[QueuePayloadEntry], + session_id: str, + workspace_name: str, ) -> list[models.QueueItem]: """Add queue items to the database and return them""" queue_items: list[models.QueueItem] = [] - for payload in payloads: + for payload_entry in payloads: + payload: QueuePayload + message_id: int | None + if isinstance(payload_entry, tuple): + payload, message_id = payload_entry + else: + payload = payload_entry + message_id = cast(int | None, payload.get("message_id")) # Generate work_unit_key from the payload - task_type = payload.get("task_type", "unknown") - work_unit_key = get_work_unit_key(payload) + task_type = cast(str, payload.get("task_type", "unknown")) + work_unit_key = construct_work_unit_key(workspace_name, payload) queue_item = models.QueueItem( session_id=session_id, @@ -174,6 +188,8 @@ async def add_queue_items( work_unit_key=work_unit_key, payload=payload, processed=False, + workspace_name=workspace_name, + message_id=message_id, ) db_session.add(queue_item) queue_items.append(queue_item) @@ -203,7 +219,7 @@ async def sample_queue_items( messages = sample_messages # Create various types of queue payloads - payloads: list[dict[str, Any]] = [] + payloads: list[tuple[dict[str, Any], int]] = [] # Create representation payloads for each message for message in messages: @@ -214,7 +230,7 @@ async def sample_queue_items( observer=message.peer_name, observed=message.peer_name, ) - payloads.append(payload1) + payloads.append((payload1, message.id)) # Representation for observer peer payload2 = create_queue_payload( @@ -223,7 +239,7 @@ async def sample_queue_items( observer=peer2.name, # peer2 observes others observed=message.peer_name, ) - payloads.append(payload2) + payloads.append((payload2, message.id)) # Create summary payloads for session for i, message in enumerate(messages): @@ -232,10 +248,10 @@ async def sample_queue_items( task_type="summary", message_seq_in_session=i + 1, ) - payloads.append(payload) + payloads.append((payload, message.id)) # Add all payloads as queue items - queue_items = await add_queue_items(payloads, session.id) + queue_items = await add_queue_items(payloads, session.id, session.workspace_name) return queue_items diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py index 002377df..4ebb86f7 100644 --- a/tests/deriver/test_deriver_processing.py +++ b/tests/deriver/test_deriver_processing.py @@ -9,6 +9,7 @@ import pytest from src import models from src.deriver.deriver import process_representation_tasks_batch from src.utils.representation import Representation +from src.utils.work_unit import construct_work_unit_key @pytest.mark.asyncio @@ -32,14 +33,12 @@ class TestDeriverProcessing: sample_session_with_peers: tuple[models.Session, list[models.Peer]], ): """Test that work unit keys are generated correctly""" - from src.utils.work_unit import get_work_unit_key session, peers = sample_session_with_peers peer1, peer2, _ = peers # Create a payload for representation task representation_payload = { - "workspace_name": "workspace1", "session_name": session.name, "observer": peer2.name, "observed": peer1.name, @@ -47,22 +46,25 @@ class TestDeriverProcessing: } # Generate work unit key for representation - work_unit_key = get_work_unit_key(representation_payload) - expected_key = ( - f"representation:workspace1:{session.name}:{peer2.name}:{peer1.name}" + work_unit_key = construct_work_unit_key( + session.workspace_name, representation_payload ) + expected_key = f"representation:{session.workspace_name}:{session.name}:{peer2.name}:{peer1.name}" assert work_unit_key == expected_key # Create a payload for summary task summary_payload = { - "workspace_name": "workspace1", "session_name": session.name, "task_type": "summary", } # Generate work unit key for summary - summary_work_unit_key = get_work_unit_key(summary_payload) - expected_summary_key = f"summary:workspace1:{session.name}:None:None" + summary_work_unit_key = construct_work_unit_key( + session.workspace_name, summary_payload + ) + expected_summary_key = ( + f"summary:{session.workspace_name}:{session.name}:None:None" + ) assert summary_work_unit_key == expected_summary_key async def test_mock_queue_manager( diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index d620392a..dfb85605 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -10,7 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import models from src.config import settings from src.deriver.queue_manager import QueueManager, WorkerOwnership -from src.utils.work_unit import get_work_unit_key +from src.utils.work_unit import construct_work_unit_key @pytest.mark.asyncio @@ -139,7 +139,7 @@ class TestQueueProcessing: for message in messages: await db_session.refresh(message) - payloads: list[Any] = [] + payloads: list[tuple[dict[str, Any], int]] = [] for message in messages: payload = create_queue_payload( # type: ignore[reportUnknownArgumentType] message=message, @@ -147,9 +147,9 @@ class TestQueueProcessing: observed=peer.name, observer=peer.name, ) - payloads.append(payload) + payloads.append((payload, message.id)) - items = await add_queue_items(payloads, session.id) + items = await add_queue_items(payloads, session.id, session.workspace_name) # Determine ascending order by DB id ordered = ( ( @@ -304,20 +304,23 @@ class TestQueueProcessing: await db_session.refresh(message) # Create queue items with token counts - payloads = [ - create_queue_payload( # type: ignore[reportUnknownArgumentType] - message=msg, - task_type="representation", - observed=peer.name, - observer=peer.name, + payload_entries = [ + ( + create_queue_payload( # type: ignore[reportUnknownArgumentType] + message=msg, + task_type="representation", + observed=peer.name, + observer=peer.name, + ), + msg, ) for msg in messages ] queue_items: list[models.QueueItem] = [] - for payload in payloads: + for payload, message in payload_entries: task_type = payload.get("task_type", "unknown") - work_unit_key = get_work_unit_key(payload) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) queue_item = models.QueueItem( session_id=session.id, @@ -325,6 +328,8 @@ class TestQueueProcessing: work_unit_key=work_unit_key, payload=payload, processed=False, + workspace_name=session.workspace_name, + message_id=message.id, ) db_session.add(queue_item) queue_items.append(queue_item) @@ -430,7 +435,7 @@ class TestQueueProcessing: observed=peer.name, observer=target.name, ) - work_unit_key = get_work_unit_key(payload) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) queue_item = models.QueueItem( session_id=session.id, @@ -438,6 +443,8 @@ class TestQueueProcessing: work_unit_key=work_unit_key, payload=payload, processed=False, + workspace_name=session.workspace_name, + message_id=message.id, ) db_session.add(queue_item) @@ -596,7 +603,7 @@ class TestQueueProcessing: observed=peer.name, observer=target.name, ) - work_unit_key = get_work_unit_key(payload) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) queue_item = models.QueueItem( session_id=session.id, @@ -604,6 +611,8 @@ class TestQueueProcessing: work_unit_key=work_unit_key, payload=payload, processed=False, + workspace_name=session.workspace_name, + message_id=message.id, ) db_session.add(queue_item) @@ -702,7 +711,6 @@ class TestQueueProcessing: token_counts = [500, 600] messages = [ models.Message( - id=999, session_name=session.name, workspace_name=session.workspace_name, peer_name=peer.name, @@ -711,7 +719,6 @@ class TestQueueProcessing: seq_in_session=1, ), models.Message( - id=1000, session_name=session.name, workspace_name=session.workspace_name, peer_name=peer.name, @@ -721,6 +728,15 @@ class TestQueueProcessing: ), ] + # Save messages to database first + for message in messages: + db_session.add(message) + await db_session.commit() + + # Refresh to get the actual IDs + for message in messages: + await db_session.refresh(message) + # Create payloads and queue items queue_items: list[models.QueueItem] = [] for i, message in enumerate(messages): @@ -729,7 +745,7 @@ class TestQueueProcessing: ) payload["token_count"] = token_counts[i] - work_unit_key = get_work_unit_key(payload) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) queue_item = models.QueueItem( session_id=session.id, @@ -737,6 +753,8 @@ class TestQueueProcessing: work_unit_key=work_unit_key, payload=payload, processed=False, + workspace_name=session.workspace_name, + message_id=message.id, ) db_session.add(queue_item) queue_items.append(queue_item) @@ -747,10 +765,11 @@ class TestQueueProcessing: processed_batches: list[dict[str, Any]] = [] async def mock_process_item( - task_type: str, - queue_payload: dict[str, Any], # pyright: ignore[reportUnusedParameter] + queue_item: models.QueueItem, ) -> None: - processed_batches.append({"task_type": task_type, "payload_count": 1}) + processed_batches.append( + {"task_type": queue_item.task_type, "payload_count": 1} + ) qm = QueueManager() work_unit_key = queue_items[0].work_unit_key @@ -838,21 +857,24 @@ class TestQueueProcessing: await db_session.refresh(message) # Create queue items - payloads = [ - create_queue_payload( # type: ignore[reportUnknownArgumentType] - message=msg, - task_type="representation", - observed=peer.name, - observer=peer.name, + payload_entries = [ + ( + create_queue_payload( # type: ignore[reportUnknownArgumentType] + message=msg, + task_type="representation", + observed=peer.name, + observer=peer.name, + ), + msg, ) for msg in messages ] # Add items to queue queue_items: list[models.QueueItem] = [] - for payload in payloads: + for payload, message in payload_entries: task_type = payload.get("task_type", "unknown") - work_unit_key = get_work_unit_key(payload) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) queue_item = models.QueueItem( session_id=session.id, @@ -860,6 +882,8 @@ class TestQueueProcessing: work_unit_key=work_unit_key, payload=payload, processed=False, + workspace_name=session.workspace_name, + message_id=message.id, ) db_session.add(queue_item) queue_items.append(queue_item) @@ -949,21 +973,24 @@ class TestQueueProcessing: await db_session.refresh(message) # Create queue items - payloads = [ - create_queue_payload( # type: ignore[reportUnknownArgumentType] - message=msg, - task_type="representation", - observed=peer.name, - observer=peer.name, + payload_entries = [ + ( + create_queue_payload( # type: ignore[reportUnknownArgumentType] + message=msg, + task_type="representation", + observed=peer.name, + observer=peer.name, + ), + msg, ) for msg in messages ] # Add items to queue queue_items: list[models.QueueItem] = [] - for payload in payloads: + for payload, message in payload_entries: task_type = payload.get("task_type", "unknown") - work_unit_key = get_work_unit_key(payload) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) queue_item = models.QueueItem( session_id=session.id, @@ -971,6 +998,8 @@ class TestQueueProcessing: work_unit_key=work_unit_key, payload=payload, processed=False, + workspace_name=session.workspace_name, + message_id=message.id, ) db_session.add(queue_item) queue_items.append(queue_item) diff --git a/tests/integration/test_representation.py b/tests/integration/test_representation.py index 15d8dd37..f3e112b5 100644 --- a/tests/integration/test_representation.py +++ b/tests/integration/test_representation.py @@ -222,8 +222,8 @@ class TestDocumentCreationWorkflow: observed=observed_peer.name, content="User likes dogs", session_name="test_session", + level="explicit", internal_metadata={ - "level": "explicit", "message_ids": [(1, 1)], "session_name": "test_session", }, @@ -274,7 +274,9 @@ class TestDocumentCreationWorkflow: observed=observed_peer.name, content="Highly derived observation", session_name=session.name, - internal_metadata={"level": "explicit", "times_derived": 5}, + level="explicit", + times_derived=5, + internal_metadata={}, embedding=[0.1] * 1536, ) db_session.add(highly_derived_doc) @@ -286,7 +288,9 @@ class TestDocumentCreationWorkflow: observed=observed_peer.name, content="Less derived observation", session_name=session.name, - internal_metadata={"level": "explicit", "times_derived": 2}, + level="explicit", + times_derived=2, + internal_metadata={}, embedding=[0.2] * 1536, ) db_session.add(less_derived_doc) @@ -315,8 +319,8 @@ class TestDocumentCreationWorkflow: observer="test_peer", observed="test_peer", content="User said they like programming", + level="explicit", internal_metadata={ - "level": "explicit", "message_ids": [(1, 1)], }, session_name="test_session", @@ -329,8 +333,8 @@ class TestDocumentCreationWorkflow: observer="test_peer", observed="test_peer", content="User is likely a software developer", + level="deductive", internal_metadata={ - "level": "deductive", "message_ids": [(1, 1)], "premises": ["User said they like programming"], }, diff --git a/tests/routes/test_queue_status.py b/tests/routes/test_queue_status.py index 85fd048f..bc50dcc7 100644 --- a/tests/routes/test_queue_status.py +++ b/tests/routes/test_queue_status.py @@ -3,7 +3,7 @@ from fastapi.testclient import TestClient from sqlalchemy.ext.asyncio import AsyncSession from src import models -from src.utils.work_unit import get_work_unit_key +from src.utils.work_unit import construct_work_unit_key @pytest.mark.asyncio @@ -157,9 +157,10 @@ class TestDeriverStatusEndpoint: queue_item = models.QueueItem( session_id=session.id, task_type="representation", - work_unit_key=get_work_unit_key(payload), + work_unit_key=construct_work_unit_key(workspace.name, payload), payload=payload, processed=False, + workspace_name=workspace.name, ) queue_items.append(queue_item) db_session.add_all(queue_items) @@ -231,9 +232,10 @@ class TestDeriverStatusEndpoint: queue_item = models.QueueItem( session_id=session.id, task_type="representation", - work_unit_key=get_work_unit_key(payload), + work_unit_key=construct_work_unit_key(workspace.name, payload), payload=payload, processed=False, + workspace_name=workspace.name, ) queue_items.append(queue_item) db_session.add_all(queue_items) @@ -292,9 +294,10 @@ class TestDeriverStatusEndpoint: queue_item = models.QueueItem( session_id=session.id, task_type="representation", - work_unit_key=get_work_unit_key(payload), + work_unit_key=construct_work_unit_key(workspace.name, payload), payload=payload, processed=False, + workspace_name=workspace.name, ) db_session.add(queue_item) await db_session.commit() diff --git a/tests/test_schema_validations.py b/tests/test_schema_validations.py index 86ad1b37..775c32b8 100644 --- a/tests/test_schema_validations.py +++ b/tests/test_schema_validations.py @@ -95,17 +95,18 @@ class TestDocumentValidations: def test_valid_document_create(self): metadata = DocumentMetadata( message_ids=[(1, 1)], - level="explicit", premises=[], message_created_at="2021-01-01T00:00:00Z", ) doc = DocumentCreate( content="test content", session_name="test", + level="explicit", metadata=metadata, embedding=[0.1, 0.2, 0.3], ) assert doc.content == "test content" + assert doc.level == "explicit" assert doc.metadata == metadata def test_document_content_too_short(self): @@ -113,9 +114,9 @@ class TestDocumentValidations: DocumentCreate( content="", session_name="test", + level="explicit", metadata=DocumentMetadata( message_ids=[(1, 1)], - level="explicit", premises=[], message_created_at="2021-01-01T00:00:00Z", ), @@ -129,9 +130,9 @@ class TestDocumentValidations: DocumentCreate( content="a" * 100001, session_name="test", + level="explicit", metadata=DocumentMetadata( message_ids=[(1, 1)], - level="explicit", premises=[], message_created_at="2021-01-01T00:00:00Z", ),