hindsight-api 0.1.10__py3-none-any.whl → 0.1.12__py3-none-any.whl

hindsight_api/__init__.py

@@ -21,9 +21,11 @@ from .engine.search.trace import (
     WeightComponents,
 )
 from .engine.search.tracer import SearchTracer
+from .models import RequestContext
 
 __all__ = [
     "MemoryEngine",
+    "RequestContext",
     "HindsightConfig",
     "get_config",
     "SearchTrace",

hindsight_api/alembic/env.py

@@ -109,6 +109,9 @@ def run_migrations_online() -> None:
 
     get_database_url()  # Process and set the database URL in config
 
+    # Check if we're targeting a specific schema (for multi-tenant isolation)
+    target_schema = config.get_main_option("target_schema")
+
     connectable = engine_from_config(
         config.get_section(config.config_ini_section, {}),
         prefix="sqlalchemy.",
@@ -121,14 +124,34 @@ def run_migrations_online() -> None:
     def set_read_write_mode(dbapi_connection, connection_record):
         cursor = dbapi_connection.cursor()
         cursor.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE")
+        # If targeting a specific schema, set search_path
+        # Include public in search_path for access to shared extensions (pgvector)
+        if target_schema:
+            cursor.execute(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"')
+            cursor.execute(f'SET search_path TO "{target_schema}", public')
         cursor.close()
 
     with connectable.connect() as connection:
         # Also explicitly set read-write mode on this connection
         connection.execute(text("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE"))
+
+        # If targeting a specific schema, set search_path
+        # Include public in search_path for access to shared extensions (pgvector)
+        if target_schema:
+            connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"'))
+            connection.execute(text(f'SET search_path TO "{target_schema}", public'))
+
         connection.commit()  # Commit the SET command
 
-        context.configure(connection=connection, target_metadata=target_metadata)
+        # Configure context with version_table_schema if using a specific schema
+        context_opts = {
+            "connection": connection,
+            "target_metadata": target_metadata,
+        }
+        if target_schema:
+            context_opts["version_table_schema"] = target_schema
+
+        context.configure(**context_opts)
 
         with context.begin_transaction():
             context.run_migrations()

hindsight_api/alembic/versions/d9f6a3b4c5e2_rename_bank_to_interactions.py

@@ -6,7 +6,7 @@ Create Date: 2024-12-04 15:00:00.000000
 
 """
 
-from alembic import op
+from alembic import context, op
 
 # revision identifiers, used by Alembic.
 revision = "d9f6a3b4c5e2"
@@ -15,14 +15,22 @@ branch_labels = None
 depends_on = None
 
 
+def _get_schema_prefix() -> str:
+    """Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
+    schema = context.config.get_main_option("target_schema")
+    return f'"{schema}".' if schema else ""
+
+
 def upgrade():
+    schema = _get_schema_prefix()
+
     # Drop old check constraint FIRST (before updating data)
     op.drop_constraint("memory_units_fact_type_check", "memory_units", type_="check")
 
     # Update existing 'bank' values to 'experience'
-    op.execute("UPDATE memory_units SET fact_type = 'experience' WHERE fact_type = 'bank'")
+    op.execute(f"UPDATE {schema}memory_units SET fact_type = 'experience' WHERE fact_type = 'bank'")
     # Also update any 'interactions' values (in case of partial migration)
-    op.execute("UPDATE memory_units SET fact_type = 'experience' WHERE fact_type = 'interactions'")
+    op.execute(f"UPDATE {schema}memory_units SET fact_type = 'experience' WHERE fact_type = 'interactions'")
 
     # Create new check constraint with 'experience' instead of 'bank'
     op.create_check_constraint(
@@ -31,11 +39,13 @@ def upgrade():
 
 
 def downgrade():
+    schema = _get_schema_prefix()
+
     # Drop new check constraint FIRST
     op.drop_constraint("memory_units_fact_type_check", "memory_units", type_="check")
 
     # Update 'experience' back to 'bank'
-    op.execute("UPDATE memory_units SET fact_type = 'bank' WHERE fact_type = 'experience'")
+    op.execute(f"UPDATE {schema}memory_units SET fact_type = 'bank' WHERE fact_type = 'experience'")
 
     # Recreate old check constraint
     op.create_check_constraint(

hindsight_api/alembic/versions/e0a1b2c3d4e5_disposition_to_3_traits.py

@@ -12,7 +12,7 @@ system (skepticism, literalism, empathy with 1-5 integer values).
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-from alembic import op
+from alembic import context, op
 
 # revision identifiers, used by Alembic.
 revision: str = "e0a1b2c3d4e5"
@@ -21,9 +21,36 @@ branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 
 
+def _get_schema_prefix() -> str:
+    """Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
+    schema = context.config.get_main_option("target_schema")
+    return f'"{schema}".' if schema else ""
+
+
+def _get_target_schema() -> str:
+    """Get the target schema name (tenant schema or 'public')."""
+    schema = context.config.get_main_option("target_schema")
+    return schema if schema else "public"
+
+
 def upgrade() -> None:
     """Convert Big Five disposition to 3-trait disposition."""
     conn = op.get_bind()
+    schema = _get_schema_prefix()
+    target_schema = _get_target_schema()
+
+    # Check if disposition column exists (should have been created by previous migration)
+    result = conn.execute(
+        sa.text("""
+        SELECT column_name
+        FROM information_schema.columns
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
+    """),
+        {"schema": target_schema},
+    )
+    if not result.fetchone():
+        # Column doesn't exist yet (shouldn't happen but be safe)
+        return
 
     # Update all existing banks to use the new disposition format
     # Convert from old format to new format with reasonable mappings:
@@ -32,18 +59,18 @@ def upgrade() -> None:
     # - empathy: derived from agreeableness + inverse of neuroticism
     # Default all to 3 (neutral) for simplicity
     conn.execute(
-        sa.text("""
-        UPDATE banks
-        SET disposition = '{"skepticism": 3, "literalism": 3, "empathy": 3}'::jsonb
+        sa.text(f"""
+        UPDATE {schema}banks
+        SET disposition = '{{"skepticism": 3, "literalism": 3, "empathy": 3}}'::jsonb
         WHERE disposition IS NOT NULL
     """)
     )
 
     # Update the default for new banks
     conn.execute(
-        sa.text("""
-        ALTER TABLE banks
-        ALTER COLUMN disposition SET DEFAULT '{"skepticism": 3, "literalism": 3, "empathy": 3}'::jsonb
+        sa.text(f"""
+        ALTER TABLE {schema}banks
+        ALTER COLUMN disposition SET DEFAULT '{{"skepticism": 3, "literalism": 3, "empathy": 3}}'::jsonb
     """)
     )
 
@@ -51,20 +78,34 @@ def upgrade() -> None:
 def downgrade() -> None:
     """Convert back to Big Five disposition."""
     conn = op.get_bind()
+    schema = _get_schema_prefix()
+    target_schema = _get_target_schema()
+
+    # Check if disposition column exists
+    result = conn.execute(
+        sa.text("""
+        SELECT column_name
+        FROM information_schema.columns
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
+    """),
+        {"schema": target_schema},
+    )
+    if not result.fetchone():
+        return
 
     # Revert to Big Five format with default values
     conn.execute(
-        sa.text("""
-        UPDATE banks
-        SET disposition = '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb
+        sa.text(f"""
+        UPDATE {schema}banks
+        SET disposition = '{{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}}'::jsonb
         WHERE disposition IS NOT NULL
     """)
     )
 
     # Update the default for new banks
     conn.execute(
-        sa.text("""
-        ALTER TABLE banks
-        ALTER COLUMN disposition SET DEFAULT '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb
+        sa.text(f"""
+        ALTER TABLE {schema}banks
+        ALTER COLUMN disposition SET DEFAULT '{{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}}'::jsonb
     """)
     )

hindsight_api/alembic/versions/rename_personality_to_disposition.py

@@ -9,7 +9,7 @@ Create Date: 2024-12-04
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-from alembic import op
+from alembic import context, op
 from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
@@ -19,17 +19,25 @@ branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 
 
+def _get_target_schema() -> str:
+    """Get the target schema name (tenant schema or 'public')."""
+    schema = context.config.get_main_option("target_schema")
+    return schema if schema else "public"
+
+
 def upgrade() -> None:
     """Rename personality column to disposition in banks table (if it exists)."""
     conn = op.get_bind()
+    target_schema = _get_target_schema()
 
     # Check if 'personality' column exists (old database)
     result = conn.execute(
         sa.text("""
         SELECT column_name
         FROM information_schema.columns
-        WHERE table_name = 'banks' AND column_name = 'personality'
-    """)
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'personality'
+    """),
+        {"schema": target_schema},
     )
     has_personality = result.fetchone() is not None
 
@@ -38,8 +46,9 @@ def upgrade() -> None:
         sa.text("""
         SELECT column_name
         FROM information_schema.columns
-        WHERE table_name = 'banks' AND column_name = 'disposition'
-    """)
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
+    """),
+        {"schema": target_schema},
     )
     has_disposition = result.fetchone() is not None
 
@@ -63,12 +72,14 @@ def upgrade() -> None:
 def downgrade() -> None:
     """Revert disposition column back to personality."""
     conn = op.get_bind()
+    target_schema = _get_target_schema()
     result = conn.execute(
         sa.text("""
         SELECT column_name
         FROM information_schema.columns
-        WHERE table_name = 'banks' AND column_name = 'disposition'
-    """)
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
+    """),
+        {"schema": target_schema},
     )
     if result.fetchone():
         op.alter_column("banks", "disposition", new_column_name="personality")