diff --git a/backend/app/alembic/versions/0005_add_memory_system_tables.py b/backend/app/alembic/versions/0005_add_memory_system_tables.py index d1e64f7..e4cc106 100644 --- a/backend/app/alembic/versions/0005_add_memory_system_tables.py +++ b/backend/app/alembic/versions/0005_add_memory_system_tables.py @@ -361,13 +361,53 @@ def upgrade() -> None: unique=True, ) op.create_index("ix_procedures_project_name", "procedures", ["project_id", "name"]) - op.create_index("ix_procedures_agent_type", "procedures", ["agent_type_id"]) + # Note: agent_type_id already indexed via ix_procedures_agent_type_id (line 354) op.create_index( "ix_procedures_success_rate", "procedures", ["success_count", "failure_count"], ) + # ========================================================================= + # Add check constraints for data integrity + # ========================================================================= + + # Episode constraints + op.create_check_constraint( + "ck_episodes_importance_range", + "episodes", + "importance_score >= 0.0 AND importance_score <= 1.0", + ) + op.create_check_constraint( + "ck_episodes_duration_positive", + "episodes", + "duration_seconds >= 0.0", + ) + op.create_check_constraint( + "ck_episodes_tokens_positive", + "episodes", + "tokens_used >= 0", + ) + + # Fact constraints + op.create_check_constraint( + "ck_facts_confidence_range", + "facts", + "confidence >= 0.0 AND confidence <= 1.0", + ) + + # Procedure constraints + op.create_check_constraint( + "ck_procedures_success_positive", + "procedures", + "success_count >= 0", + ) + op.create_check_constraint( + "ck_procedures_failure_positive", + "procedures", + "failure_count >= 0", + ) + # ========================================================================= # Create memory_consolidation_log table # Tracks consolidation jobs @@ -432,6 +472,14 @@ def upgrade() -> None: def downgrade() -> None: """Drop Agent Memory System tables.""" + # Drop check constraints first + op.drop_constraint("ck_procedures_failure_positive", "procedures", type_="check") + op.drop_constraint("ck_procedures_success_positive", "procedures", type_="check") + op.drop_constraint("ck_facts_confidence_range", "facts", type_="check") + op.drop_constraint("ck_episodes_tokens_positive", "episodes", type_="check") + op.drop_constraint("ck_episodes_duration_positive", "episodes", type_="check") + op.drop_constraint("ck_episodes_importance_range", "episodes", type_="check") + # Drop tables in reverse order (dependencies first) op.drop_table("memory_consolidation_log") op.drop_table("procedures") diff --git a/backend/app/models/memory/episode.py b/backend/app/models/memory/episode.py index 49c322e..0dd4dfb 100644 --- a/backend/app/models/memory/episode.py +++ b/backend/app/models/memory/episode.py @@ -8,6 +8,7 @@ with context, actions, outcomes, and lessons learned. from sqlalchemy import ( BigInteger, + CheckConstraint, Column, DateTime, Enum, @@ -119,6 +120,19 @@ class Episode(Base, UUIDMixin, TimestampMixin): Index("ix_episodes_project_time", "project_id", "occurred_at"), # For importance-based pruning Index("ix_episodes_importance_time", "importance_score", "occurred_at"), + # Data integrity constraints + CheckConstraint( + "importance_score >= 0.0 AND importance_score <= 1.0", + name="ck_episodes_importance_range", + ), + CheckConstraint( + "duration_seconds >= 0.0", + name="ck_episodes_duration_positive", + ), + CheckConstraint( + "tokens_used >= 0", + name="ck_episodes_tokens_positive", + ), ) def __repr__(self) -> str: diff --git a/backend/app/models/memory/fact.py b/backend/app/models/memory/fact.py index b36299d..81b7045 100644 --- a/backend/app/models/memory/fact.py +++ b/backend/app/models/memory/fact.py @@ -7,6 +7,7 @@ triple format with confidence scores and source tracking. """ from sqlalchemy import ( + CheckConstraint, Column, DateTime, Float, @@ -15,6 +16,7 @@ from sqlalchemy import ( Integer, String, Text, + text, ) from sqlalchemy.dialects.postgresql import ( ARRAY, @@ -86,7 +88,7 @@ class Fact(Base, UUIDMixin, TimestampMixin): "predicate", "object", unique=True, - postgresql_where="project_id IS NOT NULL", + postgresql_where=text("project_id IS NOT NULL"), ), # Query patterns Index("ix_facts_subject_predicate", "subject", "predicate"), @@ -94,6 +96,11 @@ class Fact(Base, UUIDMixin, TimestampMixin): Index("ix_facts_confidence_time", "confidence", "last_reinforced"), # For finding facts by entity (subject or object) Index("ix_facts_subject", "subject"), + # Data integrity constraints + CheckConstraint( + "confidence >= 0.0 AND confidence <= 1.0", + name="ck_facts_confidence_range", + ), ) def __repr__(self) -> str: diff --git a/backend/app/models/memory/procedure.py b/backend/app/models/memory/procedure.py index 3342510..6c1b8d0 100644 --- a/backend/app/models/memory/procedure.py +++ b/backend/app/models/memory/procedure.py @@ -7,6 +7,7 @@ derived from successful task execution patterns. """ from sqlalchemy import ( + CheckConstraint, Column, DateTime, ForeignKey, @@ -91,22 +92,35 @@ class Procedure(Base, UUIDMixin, TimestampMixin): ), # Query patterns Index("ix_procedures_project_name", "project_id", "name"), - Index("ix_procedures_agent_type", "agent_type_id"), + # Note: agent_type_id already has index=True on Column definition # For finding best procedures Index("ix_procedures_success_rate", "success_count", "failure_count"), + # Data integrity constraints + CheckConstraint( + "success_count >= 0", + name="ck_procedures_success_positive", + ), + CheckConstraint( + "failure_count >= 0", + name="ck_procedures_failure_positive", + ), ) @property def success_rate(self) -> float: """Calculate the success rate of this procedure.""" - total = self.success_count + self.failure_count + # Snapshot values to avoid race conditions in concurrent access + success = self.success_count + failure = self.failure_count + total = success + failure if total == 0: return 0.0 - return self.success_count / total + return success / total @property def total_uses(self) -> int: """Get total number of times this procedure was used.""" + # Snapshot values for consistency return self.success_count + self.failure_count def __repr__(self) -> str: